Merge pull request #83 from WorkflowAI/guillaume/add-uid-to-agent

Add uid to agent class

Author: guillaq

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "workflowai"
-version = "0.6.2"
+version = "0.6.3"
 description = "Python SDK for WorkflowAI"
 authors = ["Guillaume Aquilina <[email protected]>"]
 readme = "README.md"

tests/integration/conftest.py

@@ -34,7 +34,7 @@ def mock_register(self, schema_id: int = 1, task_id: str = "city-to-capital", va
         self.httpx_mock.add_response(
             method="POST",
             url=self.REGISTER_URL,
-            json={"schema_id": schema_id, "variant_id": variant_id, "id": task_id},
+            json={"schema_id": schema_id, "variant_id": variant_id, "id": task_id, "uid": 123},
         )
 
     def mock_response(

workflowai/core/client/_models.py

@@ -166,6 +166,7 @@ class CreateAgentRequest(BaseModel):
 class CreateAgentResponse(BaseModel):
     id: str
     schema_id: int
+    uid: int = 0
 
 
 class ModelMetadata(BaseModel):

workflowai/core/client/_types.py

@@ -54,6 +54,9 @@ class AgentInterface(_BaseObject, Generic[AgentInputContra, AgentOutput], Protoc
     __kwdefaults__: Optional[dict[str, Any]]
     __code__: Any
 
+    @property
+    def agent_uid(self) -> int: ...
+
     async def run(
         self,
         agent_input: AgentInputContra,

workflowai/core/client/agent.py

@@ -111,6 +111,8 @@ def __init__(
 
         self._default_validator = default_validator(output_cls)
         self._other_run_params = kwargs
+        # The UID of the agent. Set once the agent has been registered
+        self.agent_uid: int = 0
 
     @classmethod
     def build_tools(cls, tools: Iterable[Callable[..., Any]]):
@@ -267,6 +269,7 @@ async def register(self):
             returns=CreateAgentResponse,
         )
         self.schema_id = res.schema_id
+        self.agent_uid = res.uid
         return res.schema_id
 
     @classmethod

workflowai/core/client/agent_test.py

@@ -231,6 +231,7 @@ async def test_auto_register(self, httpx_mock: HTTPXMock, agent_no_schema: Agent
             json={
                 "id": "123",
                 "schema_id": 2,
+                "uid": 123,
             },
         )
         run_response = fixtures_json("task_run.json")
@@ -241,6 +242,7 @@ async def test_auto_register(self, httpx_mock: HTTPXMock, agent_no_schema: Agent
 
         out = await agent_no_schema.run(HelloTaskInput(name="Alice"))
         assert out.id == "8f635b73-f403-47ee-bff9-18320616c6cc"
+        assert agent_no_schema.agent_uid == 123
 
         run_response["id"] = "8f635b73-f403-47ee-bff9-18320616c6cc"
         # Try and run again
@@ -286,7 +288,7 @@ class AliasOutput(BaseModel):
 
         agent = Agent(agent_id="123", input_cls=AliasInput, output_cls=AliasOutput, api=api_client)
 
-        httpx_mock.add_response(url="http://localhost:8000/v1/_/agents", json={"id": "123", "schema_id": 2})
+        httpx_mock.add_response(url="http://localhost:8000/v1/_/agents", json={"id": "123", "schema_id": 2, "uid": 123})
 
         httpx_mock.add_response(
             url="http://localhost:8000/v1/_/agents/123/schemas/2/run",
@@ -805,7 +807,7 @@ async def test_list_models_registers_if_needed(
         # Mock the registration response
         httpx_mock.add_response(
             url="http://localhost:8000/v1/_/agents",
-            json={"id": "123", "schema_id": 2},
+            json={"id": "123", "schema_id": 2, "uid": 123},
         )
 
         # Mock the models response with the new structure