Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Commit 234727a

Browse files
aponcedeleonchaponcedeleonch
authored
Transform codegate-version and codegate-workspace to cli like codegate (#633)
* Transform `codegate-version` and `codegate-workspace` to cli like This PR transforms the existent shortcut commands for the pipleine and groups them in a single step. Now both commands can be accessed using: - `codegate version` - `codegate workspace` * Styling and test changes
1 parent 7e4b963 commit 234727a

13 files changed

Lines changed: 331 additions & 314 deletions

File tree

‎src/codegate/api/v1.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,10 @@
33
from fastapi.routing import APIRoute
44

55
from codegate.api import v1_models
6-
from codegate.pipeline.workspace import commands as wscmd
6+
from codegate.workspaces.crud import WorkspaceCrud
77

88
v1 = APIRouter()
9-
wscrud = wscmd.WorkspaceCrud()
9+
wscrud = WorkspaceCrud()
1010

1111

1212
def uniq_name(route: APIRoute):
@@ -61,9 +61,12 @@ async def create_workspace(request: v1_models.CreateWorkspaceRequest):
6161
return v1_models.Workspace(name=request.name)
6262

6363

64-
65-
@v1.delete("/workspaces/{workspace_name}", tags=["Workspaces"],
66-
generate_unique_id_function=uniq_name, status_code=204)
64+
@v1.delete(
65+
"/workspaces/{workspace_name}",
66+
tags=["Workspaces"],
67+
generate_unique_id_function=uniq_name,
68+
status_code=204,
69+
)
6770
async def delete_workspace(workspace_name: str):
6871
"""Delete a workspace by name."""
6972
raise NotImplementedError

‎src/codegate/api/v1_models.py‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,36 +9,44 @@ class Workspace(pydantic.BaseModel):
99
name: str
1010
is_active: bool
1111

12+
1213
class ActiveWorkspace(Workspace):
1314
# TODO: use a more specific type for last_updated
1415
last_updated: Any
1516

17+
1618
class ListWorkspacesResponse(pydantic.BaseModel):
1719
workspaces: list[Workspace]
1820

1921
@classmethod
2022
def from_db_workspaces(
21-
cls, db_workspaces: List[db_models.WorkspaceActive])-> "ListWorkspacesResponse":
22-
return cls(workspaces=[
23-
Workspace(name=ws.name, is_active=ws.active_workspace_id is not None)
24-
for ws in db_workspaces])
23+
cls, db_workspaces: List[db_models.WorkspaceActive]
24+
) -> "ListWorkspacesResponse":
25+
return cls(
26+
workspaces=[
27+
Workspace(name=ws.name, is_active=ws.active_workspace_id is not None)
28+
for ws in db_workspaces
29+
]
30+
)
31+
2532

2633
class ListActiveWorkspacesResponse(pydantic.BaseModel):
2734
workspaces: list[ActiveWorkspace]
2835

2936
@classmethod
3037
def from_db_workspaces(
31-
cls, ws: Optional[db_models.ActiveWorkspace]) -> "ListActiveWorkspacesResponse":
38+
cls, ws: Optional[db_models.ActiveWorkspace]
39+
) -> "ListActiveWorkspacesResponse":
3240
if ws is None:
3341
return cls(workspaces=[])
34-
return cls(workspaces=[
35-
ActiveWorkspace(name=ws.name,
36-
is_active=True,
37-
last_updated=ws.last_update)
38-
])
42+
return cls(
43+
workspaces=[ActiveWorkspace(name=ws.name, is_active=True, last_updated=ws.last_update)]
44+
)
45+
3946

4047
class CreateWorkspaceRequest(pydantic.BaseModel):
4148
name: str
4249

50+
4351
class ActivateWorkspaceRequest(pydantic.BaseModel):
4452
name: str

‎src/codegate/db/connection.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
alert_queue = asyncio.Queue()
3131
fim_cache = FimCache()
3232

33+
3334
class DbCodeGate:
3435
_instance = None
3536

@@ -256,8 +257,7 @@ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
256257
"""
257258
)
258259
try:
259-
added_workspace = await self._execute_update_pydantic_model(
260-
workspace, sql)
260+
added_workspace = await self._execute_update_pydantic_model(workspace, sql)
261261
except Exception as e:
262262
logger.error(f"Failed to add workspace: {workspace_name}.", error=str(e))
263263
return None

‎src/codegate/pipeline/cli/cli.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
import shlex
2+
3+
from litellm import ChatCompletionRequest
4+
5+
from codegate.pipeline.base import (
6+
PipelineContext,
7+
PipelineResponse,
8+
PipelineResult,
9+
PipelineStep,
10+
)
11+
from codegate.pipeline.cli.commands import Version, Workspace
12+
13+
HELP_TEXT = """
14+
## CodeGate CLI\n
15+
**Usage**: `codegate [-h] <command> [args]`\n
16+
Check the help of each command by running `codegate <command> -h`\n
17+
Available commands:
18+
- `version`: Show the version of CodeGate
19+
- `workspace`: Perform different operations on workspaces
20+
"""
21+
22+
NOT_FOUND_TEXT = "Command not found. Use `codegate -h` to see available commands."
23+
24+
25+
async def codegate_cli(command):
26+
"""
27+
Process the 'codegate' command.
28+
"""
29+
if len(command) == 0:
30+
return HELP_TEXT
31+
32+
available_commands = {
33+
"version": Version().exec,
34+
"workspace": Workspace().exec,
35+
"-h": lambda _: HELP_TEXT,
36+
}
37+
out_func = available_commands.get(command[0])
38+
if out_func is None:
39+
return NOT_FOUND_TEXT
40+
41+
return await out_func(command[1:])
42+
43+
44+
class CodegateCli(PipelineStep):
45+
"""Pipeline step that handles codegate cli."""
46+
47+
@property
48+
def name(self) -> str:
49+
"""
50+
Returns the name of this pipeline step.
51+
52+
Returns:
53+
str: The identifier 'codegate-cli'
54+
"""
55+
return "codegate-cli"
56+
57+
async def process(
58+
self, request: ChatCompletionRequest, context: PipelineContext
59+
) -> PipelineResult:
60+
"""
61+
Checks if the last user message contains "codegate" and process the command.
62+
This short-circuits the pipeline if the message is found.
63+
64+
Args:
65+
request (ChatCompletionRequest): The chat completion request to process
66+
context (PipelineContext): The current pipeline context
67+
68+
Returns:
69+
PipelineResult: Contains the response if triggered, otherwise continues
70+
pipeline
71+
"""
72+
last_user_message = self.get_last_user_message(request)
73+
74+
if last_user_message is not None:
75+
last_user_message_str, _ = last_user_message
76+
splitted_message = last_user_message_str.lower().split(" ")
77+
# We expect codegate as the first word in the message
78+
if splitted_message[0] == "codegate":
79+
context.shortcut_response = True
80+
args = shlex.split(last_user_message_str)
81+
cmd_out = await codegate_cli(args[1:])
82+
return PipelineResult(
83+
response=PipelineResponse(
84+
step_name=self.name,
85+
content=cmd_out,
86+
model=request["model"],
87+
),
88+
context=context,
89+
)
90+
91+
# Fall through
92+
return PipelineResult(request=request, context=context)

‎src/codegate/pipeline/cli/commands.py‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
from abc import ABC, abstractmethod
2+
from typing import List
3+
4+
from codegate import __version__
5+
from codegate.workspaces.crud import WorkspaceCrud
6+
7+
8+
class CodegateCommand(ABC):
9+
@abstractmethod
10+
async def run(self, args: List[str]) -> str:
11+
pass
12+
13+
@property
14+
@abstractmethod
15+
def help(self) -> str:
16+
pass
17+
18+
async def exec(self, args: List[str]) -> str:
19+
if args and args[0] == "-h":
20+
return self.help
21+
return await self.run(args)
22+
23+
24+
class Version(CodegateCommand):
25+
async def run(self, args: List[str]) -> str:
26+
return f"CodeGate version: {__version__}"
27+
28+
@property
29+
def help(self) -> str:
30+
return (
31+
"### CodeGate Version\n\n"
32+
"Prints the version of CodeGate.\n\n"
33+
"**Usage**: `codegate version`\n\n"
34+
"*args*: None"
35+
)
36+
37+
38+
class Workspace(CodegateCommand):
39+
40+
def __init__(self):
41+
self.workspace_crud = WorkspaceCrud()
42+
self.commands = {
43+
"list": self._list_workspaces,
44+
"add": self._add_workspace,
45+
"activate": self._activate_workspace,
46+
}
47+
48+
async def _list_workspaces(self, *args: List[str]) -> str:
49+
"""
50+
List all workspaces
51+
"""
52+
workspaces = await self.workspace_crud.get_workspaces()
53+
respond_str = ""
54+
for workspace in workspaces:
55+
respond_str += f"- {workspace.name}"
56+
if workspace.active_workspace_id:
57+
respond_str += " **(active)**"
58+
respond_str += "\n"
59+
return respond_str
60+
61+
async def _add_workspace(self, args: List[str]) -> str:
62+
"""
63+
Add a workspace
64+
"""
65+
if args is None or len(args) == 0:
66+
return "Please provide a name. Use `codegate-workspace add your_workspace_name`"
67+
68+
new_workspace_name = args[0]
69+
if not new_workspace_name:
70+
return "Please provide a name. Use `codegate-workspace add your_workspace_name`"
71+
72+
workspace_created = await self.workspace_crud.add_workspace(new_workspace_name)
73+
if not workspace_created:
74+
return (
75+
"Something went wrong. Workspace could not be added.\n"
76+
"1. Check if the name is alphanumeric and only contains dashes, and underscores.\n"
77+
"2. Check if the workspace already exists."
78+
)
79+
return f"Workspace **{new_workspace_name}** has been added"
80+
81+
async def _activate_workspace(self, args: List[str]) -> str:
82+
"""
83+
Activate a workspace
84+
"""
85+
if args is None or len(args) == 0:
86+
return "Please provide a name. Use `codegate-workspace activate workspace_name`"
87+
88+
workspace_name = args[0]
89+
if not workspace_name:
90+
return "Please provide a name. Use `codegate-workspace activate workspace_name`"
91+
92+
was_activated = await self.workspace_crud.activate_workspace(workspace_name)
93+
if not was_activated:
94+
return (
95+
f"Workspace **{workspace_name}** does not exist or was already active. "
96+
f"Use `codegate-workspace add {workspace_name}` to add it"
97+
)
98+
return f"Workspace **{workspace_name}** has been activated"
99+
100+
async def run(self, args: List[str]) -> str:
101+
if not args:
102+
return "Please provide a command. Use `codegate workspace -h` to see available commands"
103+
command = args[0]
104+
command_to_execute = self.commands.get(command)
105+
if command_to_execute is not None:
106+
return await command_to_execute(args[1:])
107+
else:
108+
return "Command not found. Use `codegate workspace -h` to see available commands"
109+
110+
@property
111+
def help(self) -> str:
112+
return (
113+
"### CodeGate Workspace\n\n"
114+
"Manage workspaces.\n\n"
115+
"**Usage**: `codegate workspace <command> [args]`\n\n"
116+
"Available commands:\n\n"
117+
"- `list`: List all workspaces\n\n"
118+
" - *args*: None\n\n"
119+
"- `add`: Add a workspace\n\n"
120+
" - *args*:\n\n"
121+
" - `workspace_name`\n\n"
122+
"- `activate`: Activate a workspace\n\n"
123+
" - *args*:\n\n"
124+
" - `workspace_name`"
125+
)

‎src/codegate/pipeline/factory.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from codegate.config import Config
44
from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
5+
from codegate.pipeline.cli.cli import CodegateCli
56
from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
67
from codegate.pipeline.extract_snippets.extract_snippets import CodeSnippetExtractor
78
from codegate.pipeline.extract_snippets.output import CodeCommentStep
@@ -13,8 +14,6 @@
1314
SecretUnredactionStep,
1415
)
1516
from codegate.pipeline.system_prompt.codegate import SystemPrompt
16-
from codegate.pipeline.version.version import CodegateVersion
17-
from codegate.pipeline.workspace.workspace import CodegateWorkspace
1817

1918

2019
class PipelineFactory:
@@ -28,8 +27,7 @@ def create_input_pipeline(self) -> SequentialPipelineProcessor:
2827
# and without obfuscating the secrets, we'd leak the secrets during those
2928
# later steps
3029
CodegateSecrets(),
31-
CodegateVersion(),
32-
CodegateWorkspace(),
30+
CodegateCli(),
3331
CodeSnippetExtractor(),
3432
CodegateContextRetriever(),
3533
SystemPrompt(Config.get_config().prompts.default_chat),

0 commit comments

Comments
 (0)