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