agent-api

Author: jinhaiyang

tests/integration_tests/test_agents.py

@@ -17,7 +17,7 @@ def test_completions_sync(logging_conf):
         print(f"request_id:{request_id}")
         response = client.agents.invoke(
             request_id=request_id,
-            agent_id="TranslateAgent",
+            agent_id="general_translation",
             messages=[
                 {
                     "role": "user",
@@ -44,7 +44,7 @@ def test_completions_stream(logging_conf):
         print(f"request_id:{request_id}")
         response = client.agents.invoke(
             request_id=request_id,
-            agent_id="TranslateAgent",
+            agent_id="general_translation",
             messages=[
                 {
                     "role": "user",

zhipuai/core/_sse_client.py

@@ -50,9 +50,11 @@ def __stream__(self) -> Iterator[ResponseT]:
         for sse in iterator:
             if sse.data.startswith("[DONE]"):
                 break
-
             if sse.event is None:
                 data = sse.json_data()
+                if isinstance(data, Mapping) and data.get("agent_id"):
+                    yield self._data_process_func(data=data, cast_type=self._cast_type, response=self.response)
+                    break
                 if isinstance(data, Mapping) and data.get("error"):
                     raise APIResponseError(
                         message="An error occurred during streaming",

zhipuai/types/agents/agents_completion.py

@@ -19,12 +19,18 @@ class AgentsCompletionChoice(BaseModel):
     finish_reason: str
     message: AgentsCompletionMessage
 
+class AgentsError:
+    code: Optional[str] = None
+    message: Optional[str] = None
+
 
 class AgentsCompletion(BaseModel):
     agent_id: Optional[str] = None
-    created: Optional[int] = None
+    conversation_id: Optional[str] = None
+    status: Optional[str] = None
     choices: List[AgentsCompletionChoice]
     request_id: Optional[str] = None
     id: Optional[str] = None
-    usage: AgentsCompletionUsage
+    usage: Optional[AgentsCompletionUsage] = None
+    error: Optional[AgentsError] = None
 

zhipuai/types/agents/agents_completion_chunk.py

@@ -26,10 +26,15 @@ class AgentsCompletionUsage(BaseModel):
     completion_tokens: int
     total_tokens: int
 
+class AgentsError:
+    code: Optional[str] = None
+    message: Optional[str] = None
+
 
 class AgentsCompletionChunk(BaseModel):
     agent_id: Optional[str] = None
+    conversation_id: Optional[str] = None
     id: Optional[str] = None
     choices: List[AgentsChoice]
-    created: Optional[int] = None
-    usage: Optional[AgentsCompletionUsage] = None
+    usage: Optional[AgentsCompletionUsage] = None
+    error: Optional[AgentsError] = None