Skip to content

Commit

Permalink
test: remove llama test
Browse files Browse the repository at this point in the history
  • Loading branch information
@imotai
imotai committed Dec 13, 2023
1 parent 9956795 commit 46f957a
Show file tree
Hide file tree
Showing 9 changed files with 83 additions and 354 deletions.
5 changes: 4 additions & 1 deletion agent/src/og_agent/agent_api_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,10 +178,13 @@ class TaskRequest(BaseModel):


async def run_task(task: TaskRequest, key):
async for respond in agent_sdk.prompt(task.prompt, key, files=task.input_files, context_id=task.context_id):
async for respond in agent_sdk.prompt(
task.prompt, key, files=task.input_files, context_id=task.context_id
):
response = StepResponse.new_from(respond).model_dump(exclude_none=True)
yield "data: %s\n" % json.dumps(response)


@app.post("/process")
async def process_task(
task: TaskRequest,
Expand Down
6 changes: 4 additions & 2 deletions agent/src/og_agent/base_agent.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,14 +63,17 @@ class TypingState:
MESSAGE = 4
OTHER = 5


class BaseAgent:

def __init__(self, sdk):
self.kernel_sdk = sdk
self.model_name = ""
self.agent_memories = {}

def create_new_memory_with_default_prompt(self, user_name, user_id, actions = ACTIONS):
def create_new_memory_with_default_prompt(
self, user_name, user_id, actions=ACTIONS
):
"""
create a new memory for the user
"""
Expand Down Expand Up @@ -386,7 +389,6 @@ async def extract_message(
response_token_count + context_output_token_count
)
if is_json_format:

(
new_text_content,
new_code_content,
Expand Down
42 changes: 28 additions & 14 deletions agent/src/og_agent/llama_agent.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ async def handle_python_function(
state=task_context.to_context_state_proto(),
response_type=TaskResponse.OnStepActionStart,
on_step_action_start=OnStepActionStart(
input=tool_input, tool='execute'
input=tool_input, tool="execute"
),
)
)
Expand Down Expand Up @@ -131,10 +131,10 @@ async def arun(self, request, queue, context, task_opt):
context_id = (
request.context_id
if request.context_id
else self.create_new_memory_with_default_prompt("", "", actions=[FUNCTION_EXECUTE,
FUNCTION_DIRECT_MESSAGE])
else self.create_new_memory_with_default_prompt(
"", "", actions=[FUNCTION_EXECUTE, FUNCTION_DIRECT_MESSAGE]
)
)

if context_id not in self.agent_memories:
await queue.put(
TaskResponse(
Expand All @@ -145,7 +145,6 @@ async def arun(self, request, queue, context, task_opt):
)
)
return

agent_memory = self.agent_memories[context_id]
agent_memory.update_options(self.memory_option)
agent_memory.append_chat_message(
Expand Down Expand Up @@ -208,16 +207,23 @@ async def arun(self, request, queue, context, task_opt):
break
logger.debug(f" llama response {json_response}")
if (
'function_call'in json_response and json_response["function_call"] == "execute"
"function_call" in json_response
and json_response["function_call"] == "execute"
):
agent_memory.append_chat_message(message)
tools_mapping = {
"python": self.handle_python_function,
"bash": self.handle_bash_code,
}

function_result = await tools_mapping[json_response["arguments"]['language']](
json_response['arguments'], queue, context, task_context, task_opt
function_result = await tools_mapping[
json_response["arguments"]["language"]
](
json_response["arguments"],
queue,
context,
task_context,
task_opt,
)

logger.debug(f"the function result {function_result}")
Expand All @@ -242,23 +248,31 @@ async def arun(self, request, queue, context, task_opt):
"role": "user",
"content": f"{action_output} \n {function_result.console_stdout}",
})
agent_memory.append_chat_message({"role": "user", "content": current_question})
agent_memory.append_chat_message(
{"role": "user", "content": current_question}
)
elif function_result.has_error:
agent_memory.append_chat_message({
"role": "user",
"content": f"{action_output} \n {function_result.console_stderr}",
})
current_question = f"Generate a new step to fix the above error"
agent_memory.append_chat_message({"role": "user", "content": current_question})
agent_memory.append_chat_message(
{"role": "user", "content": current_question}
)
else:
agent_memory.append_chat_message({
"role": "user",
"content": f"{action_output} \n {function_result.console_stdout}",
})
agent_memory.append_chat_message({
"role": "user", "content": current_question})
elif 'function_call' in json_response and json_response["function_call"] == "direct_message":
message = json_response['arguments']['message']
agent_memory.append_chat_message(
{"role": "user", "content": current_question}
)
elif (
"function_call" in json_response
and json_response["function_call"] == "direct_message"
):
message = json_response["arguments"]["message"]
await queue.put(
TaskResponse(
state=task_context.to_context_state_proto(),
Expand Down
2 changes: 1 addition & 1 deletion agent/src/og_agent/llama_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ def __init__(self, endpoint, key, grammar):
super().__init__(endpoint + "/v1/chat/completions", key)
self.grammar = grammar

async def chat(self, messages, model, temperature=0, max_tokens=1024, stop=['\n']):
async def chat(self, messages, model, temperature=0, max_tokens=1024, stop=["\n"]):
data = {
"messages": messages,
"temperature": temperature,
Expand Down
80 changes: 39 additions & 41 deletions agent/src/og_agent/prompt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,52 +17,50 @@
"Use `execute` action to execute any code and `direct_message` action to send message to user",
]

FUNCTION_EXECUTE= ActionDesc(
name="execute",
desc="This action executes code in your programming environment and returns the output",
parameters=json.dumps({
"type": "object",
"properties": {
"explanation": {
"type": "string",
"description": "the explanation about the code parameters",
},
"code": {
"type": "string",
"description": "the bash code to be executed",
},
"language": {
"type": "string",
"description": "the language of the code, only python and bash are supported",
},
"saved_filenames": {
"type": "array",
"items": {"type": "string"},
"description": "A list of filenames that were created by the code",
},
FUNCTION_EXECUTE = ActionDesc(
name="execute",
desc="This action executes code in your programming environment and returns the output",
parameters=json.dumps({
"type": "object",
"properties": {
"explanation": {
"type": "string",
"description": "the explanation about the code parameters",
},
"required": ["explanation", "code", "language"],
}),
)
"code": {
"type": "string",
"description": "the bash code to be executed",
},
"language": {
"type": "string",
"description": "the language of the code, only python and bash are supported",
},
"saved_filenames": {
"type": "array",
"items": {"type": "string"},
"description": "A list of filenames that were created by the code",
},
},
"required": ["explanation", "code", "language"],
}),
)

FUNCTION_DIRECT_MESSAGE= ActionDesc(
name="direct_message",
desc="This action sends a direct message to user.",
parameters=json.dumps({
"type": "object",
"properties": {
"message": {
"type": "string",
"description": "the message will be sent to user",
},
FUNCTION_DIRECT_MESSAGE = ActionDesc(
name="direct_message",
desc="This action sends a direct message to user.",
parameters=json.dumps({
"type": "object",
"properties": {
"message": {
"type": "string",
"description": "the message will be sent to user",
},
"required": ["message"],
}),
},
"required": ["message"],
}),
)

ACTIONS = [
FUNCTION_EXECUTE
]
ACTIONS = [FUNCTION_EXECUTE]

OUTPUT_FORMAT = """The output format must be a JSON format with the following fields:
* function_call: The name of the action
Expand Down
Loading

0 comments on commit 46f957a

Please sign in to comment.