diff --git a/pyproject.toml b/pyproject.toml index dafc105..499c1b5 100644 --- a/pyproject.toml +++ b/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 "] readme = "README.md" diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 55fd861..13ca5f7 100644 --- a/tests/integration/conftest.py +++ b/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( diff --git a/workflowai/core/client/_models.py b/workflowai/core/client/_models.py index 411c49c..b73cc73 100644 --- a/workflowai/core/client/_models.py +++ b/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): diff --git a/workflowai/core/client/_types.py b/workflowai/core/client/_types.py index ab5f541..7771c49 100644 --- a/workflowai/core/client/_types.py +++ b/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, diff --git a/workflowai/core/client/agent.py b/workflowai/core/client/agent.py index 651c4a2..7a8c86d 100644 --- a/workflowai/core/client/agent.py +++ b/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 diff --git a/workflowai/core/client/agent_test.py b/workflowai/core/client/agent_test.py index 8315f1e..031cda8 100644 --- a/workflowai/core/client/agent_test.py +++ b/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