Skip to content

Commit

Permalink
refactor(codypy): Restructure codypy module for better separation of …
Browse files Browse the repository at this point in the history 
…concerns

- Extract CodyServer class into separate server.py module
- Rename cody_py.py to agent.py and remove CodyServer class
- Move get_remote_repositories and receive_webviewmessage functions to utils.py
- Update imports and references in affected files
- Simplify CodyAgent initialization in main.py
- Rename "Claude 1.2 Instant" to "Claude Instant" in client_info.py

This refactor aims to improve the structure and organization of the codypy
module by separating the server-related functionality into its own module and
renaming files to better reflect their purpose. The changes also simplify the
usage of the CodyAgent class in the main script.
  • Loading branch information
@PriNova
PriNova committed May 21, 2024
1 parent 7359f85 commit ad2410f
Show file tree
Hide file tree
Showing 6 changed files with 225 additions and 234 deletions.
1 change: 0 additions & 1 deletion codypy/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
Models,
ModelSpec,
)
from .cody_py import CodyAgent, CodyServer
from .config import (
BLACK,
BLUE,
Expand Down
265 changes: 38 additions & 227 deletions codypy/cody_py.py → codypy/agent.py
View file
Original file line number Diff line number Diff line change
@@ -1,176 +1,33 @@
import asyncio
import os
import sys
from asyncio.subprocess import Process
from typing import Any, Self

from codypy.client_info import AgentSpecs, Models
from codypy.config import RED, RESET, debug_method_map
from codypy.logger import log_message
from codypy.utils import (
_check_for_binary_file,
_download_binary_to_path,
_format_binary_name,
)

from .config import BLUE, RED, RESET, YELLOW, Configs, debug_method_map
from .messaging import _send_jsonrpc_request, _show_last_message, request_response
from .server_info import CodyAgentSpecs


class CodyServer:
async def init(
binary_path: str,
version: str,
use_tcp: bool = False, # default because of ca-certificate verification
is_debugging: bool = False,
) -> Self:
cody_binary = ""
test_against_node_source: bool = (
False # Only for internal use to test against the Cody agent Node source
)
if not test_against_node_source:
has_agent_binary = await _check_for_binary_file(
binary_path, "cody-agent", version
)
if not has_agent_binary:
log_message(
"CodyServer: init:",
f"WARNING: The Cody Agent binary does not exist at the specified path: {binary_path}",
)
print(
f"{YELLOW}WARNING: The Cody Agent binary does not exist at the specified path: {binary_path}{RESET}"
)
print(
f"{YELLOW}WARNING: Start downloading the Cody Agent binary...{RESET}"
)
is_completed = await _download_binary_to_path(
binary_path, "cody-agent", version
)
if not is_completed:
log_message(
"CodyServer: init:",
"ERROR: Failed to download the Cody Agent binary.",
)
print(
f"{RED}ERROR: Failed to download the Cody Agent binary.{RESET}"
)
sys.exit(1)

cody_binary = os.path.join(
binary_path, await _format_binary_name("cody-agent", version)
)
from codypy.messaging import _show_last_message, request_response
from codypy.server import CodyServer
from codypy.server_info import CodyAgentSpecs

cody_agent = CodyServer(cody_binary, use_tcp, is_debugging)
await cody_agent._create_server_connection(test_against_node_source)
return cody_agent

def __init__(self, cody_binary: str, use_tcp: bool, is_debugging: bool) -> None:
self.cody_binary = cody_binary
self.use_tcp = use_tcp
self.is_debugging = is_debugging
self._process: Process | None = None
self._reader: asyncio.StreamReader | None = None
self._writer: asyncio.StreamWriter | None = None

async def _create_server_connection(
self, test_against_node_source: bool = False
class CodyAgent:
def __init__(
self,
cody_server: CodyServer,
agent_specs: AgentSpecs,
debug_method_map=debug_method_map,
) -> None:
"""
Asynchronously creates a connection to the Cody server.
If `cody_binary` is an empty string, it prints an error message and exits the program.
Sets the `CODY_AGENT_DEBUG_REMOTE` and `CODY_DEBUG` environment variables based on the `use_tcp` and `is_debugging` flags, respectively.
Creates a subprocess to run the Cody agent, either by executing the `bin/agent` binary or running the `index.js` file specified by `binary_path`.
Depending on the `use_tcp` flag, it either connects to the agent using stdio or opens a TCP connection to `localhost:3113`.
If the TCP connection fails after 5 retries, it prints an error message and exits the program.
Returns the reader and writer streams for the agent connection.
"""
if not test_against_node_source and self.cody_binary == "":
log_message(
"CodyServer: _create_server_connection:",
"ERROR: The Cody Agent binary path is empty.",
)
print(
f"{RED}You need to specify the BINARY_PATH to an absolute path to the agent binary or to the index.js file. Exiting...{RESET}"
)
sys.exit(1)

os.environ["CODY_AGENT_DEBUG_REMOTE"] = str(self.use_tcp).lower()
os.environ["CODY_DEBUG"] = str(self.is_debugging).lower()

args = []
binary = ""
if test_against_node_source:
binary = "node"
args.extend(
(
"--enable-source-maps",
"/home/prinova/CodeProjects/cody/agent/dist/index.js",
)
)
else:
binary = self.cody_binary
args.append("jsonrpc")
self._process: Process = await asyncio.create_subprocess_exec(
binary,
*args,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
env=os.environ,
)

self._reader = self._process.stdout
self._writer = self._process.stdin

if not self.use_tcp:
log_message(
"CodyServer: _create_server_connection:",
"Created a stdio connection to the Cody agent.",
)
if self.is_debugging:
print(f"{YELLOW}--- stdio connection ---{RESET}")
self._reader = self._process.stdout
self._writer = self._process.stdin

else:
log_message(
"CodyServer: _create_server_connection:",
"Created a TCP connection to the Cody agent.",
)
if self.is_debugging:
print(f"{YELLOW}--- TCP connection ---{RESET}")
retry: int = 0
retry_attempts: int = 5
for retry in range(retry_attempts):
try:
(self._reader, self._writer) = await asyncio.open_connection(
"localhost", 3113
)
if self._reader is not None and self._writer is not None:
log_message(
"CodyServer: _create_server_connection:",
"Connected to server: localhost:3113",
)
print(f"{YELLOW}Connected to server: localhost:3113{RESET}\n")
break

# return reader, writer, process
except ConnectionRefusedError:
await asyncio.sleep(1) # Retry after a short delay
retry += 1
if retry == retry_attempts:
log_message(
"CodyServer: _create_server_connection:",
"Could not connect to server. Exiting...",
)
print(f"{RED}Could not connect to server. Exiting...{RESET}")
sys.exit(1)
self._cody_server = cody_server
self.chat_id: str | None = None
self.repos: dict = {}
self.current_repo_context: list[str] = []
self.agent_specs = agent_specs
self.debug_method_map = debug_method_map

async def initialize_agent(
self,
agent_specs: AgentSpecs,
debug_method_map=debug_method_map,
is_debugging: bool = False,
) -> CodyAgentSpecs | None:
) -> None:
"""
Initializes the Cody agent by sending an "initialize" request to the agent and handling the response.
The method takes in agent specifications, a debug method map, and a boolean flag indicating whether debugging is enabled.
Expand All @@ -182,58 +39,33 @@ async def initialize_agent(
the debug method map, the reader and writer streams, the debugging flag, and the callback function.
"""

async def _handle_response(response: Any) -> None:
cody_agent_specs: CodyAgentSpecs = CodyAgentSpecs.model_validate(response)
log_message(
"CodyServer: initialize_agent:",
f"Agent Info: {cody_agent_specs}",
)
if is_debugging:
print(f"Agent Info: {cody_agent_specs}\n")
if not cody_agent_specs.authenticated:
log_message(
"CodyServer: initialize_agent:",
"Server is not authenticated.",
)
print(f"{RED}--- Server is not authenticated ---{RESET}")
await self.cleanup_server()
sys.exit(1)

response = await request_response(
"initialize",
agent_specs.model_dump(),
self.agent_specs.model_dump(),
debug_method_map,
self._reader,
self._writer,
self._cody_server._reader,
self._cody_server._writer,
is_debugging,
)

cody_agent_specs: CodyAgentSpecs = CodyAgentSpecs.model_validate(response)
log_message(
"CodyServer: initialize_agent:",
f"Agent Info: {cody_agent_specs}",
)
if is_debugging:
print(f"Agent Info: {cody_agent_specs}\n")
if cody_agent_specs.authenticated:
log_message(
"CodyServer: initialize_agent:",
"Server is authenticated.",
)
print(f"{YELLOW}--- Server is authenticated ---{RESET}")
else:
log_message(
"CodyServer: initialize_agent:",
"Server is not authenticated.",
)
print(f"{RED}--- Server is not authenticated ---{RESET}")
await self.cleanup_server()
sys.exit(1)
return await CodyAgent.init(self)

async def cleanup_server(self):
"""
Cleans up the server connection by sending a "shutdown" request to the server and terminating the server process if it is still running.
"""
log_message("CodyServer: cleanup_server:", "Cleanup Server...")
await _send_jsonrpc_request(self._writer, "shutdown", None)
if self._process.returncode is None:
self._process.terminate()
await self._process.wait()


class CodyAgent:
def __init__(self, cody_client: CodyServer) -> None:
self._cody_server = cody_client
self.chat_id: str | None = None
self.repos: dict = {}
self.current_repo_context: list[str] = []

async def init(cody_client: CodyServer):
return CodyAgent(cody_client)
await _handle_response(response)

async def new_chat(
self, debug_method_map=debug_method_map, is_debugging: bool = False
Expand Down Expand Up @@ -465,24 +297,3 @@ async def chat(
response,
context_files_response,
)


async def get_remote_repositories(
reader, writer, id: str, configs: Configs, debug_method_map
) -> Any:
return await request_response(
"chat/remoteRepos", id, debug_method_map, reader, writer, configs
)


async def receive_webviewmessage(
reader, writer, params, configs: Configs, debug_method_map
) -> Any:
return await request_response(
"webview/receiveMessage",
params,
debug_method_map,
reader,
writer,
configs,
)
2 changes: 1 addition & 1 deletion codypy/client_info.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ class Models(Enum):
model_id="anthropic/claude-2.1",
)
Claude1_2Instant = ModelSpec(
model_name="Claude 1.2 Instant",
model_name="Claude Instant",
model_id="anthropic/claude-instant-1.2",
)
Claude3Haiku = ModelSpec(
Expand Down
Loading

0 comments on commit ad2410f

Please sign in to comment.