fix: fix the agent api test case

dbpunk-labs · Nov 5, 2023 · fc6cf7d · fc6cf7d
1 parent c1e0de3
commit fc6cf7d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/agent/src/og_agent/agent_api_server.py b/agent/src/og_agent/agent_api_server.py
@@ -178,11 +178,10 @@ class TaskRequest(BaseModel):
 
 
 async def run_task(task: TaskRequest, key):
-    async for respond in agent_sdk.prompt(task.prompt, key, files=task.input_files):
+    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,

diff --git a/sdk/src/og_sdk/agent_sdk.py b/sdk/src/og_sdk/agent_sdk.py
@@ -182,11 +182,11 @@ async def add_kernel(self, key, endpoint, api_key):
         response = await self.stub.add_kernel(request, metadata=metadata)
         return response
 
-    async def prompt(self, prompt, api_key, files=[]):
+    async def prompt(self, prompt, api_key, files=[], context_id=None):
         metadata = aio.Metadata(
             ("api_key", api_key),
         )
-        request = agent_server_pb2.ProcessTaskRequest(task=prompt, input_files=files)
+        request = agent_server_pb2.ProcessTaskRequest(task=prompt, input_files=files, context_id=context_id)
         async for respond in self.stub.process_task(request, metadata=metadata):
             yield respond