fix: base agent protocol

Author: guillaq

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "workflowai"
-version = "0.6.0.dev18"
+version = "0.6.0.dev19"
 description = ""
 authors = ["Guillaume Aquilina <[email protected]>"]
 readme = "README.md"

workflowai/core/client/_types.py

@@ -11,6 +11,7 @@
 from typing_extensions import Unpack
 
 from workflowai.core._common_types import AgentInputContra, AgentOutputCov, RunParams
+from workflowai.core.client._models import ModelInfo
 from workflowai.core.domain.run import Run
 from workflowai.core.domain.task import AgentInput, AgentOutput
 from workflowai.core.domain.tool_call import ToolCallResult
@@ -75,6 +76,8 @@ async def reply(
         **kwargs: Unpack[RunParams[AgentOutput]],
     ): ...
 
+    async def list_models(self) -> list[ModelInfo]: ...
+
 
 class RunnableAgent(_BaseProtocol[AgentInputContra, AgentOutput], Protocol):
     async def __call__(

workflowai/core/domain/run.py

@@ -150,9 +150,6 @@ async def fetch_completions(self) -> list[Completion]:
 
 
 class _AgentBase(Protocol, Generic[AgentOutput]):
-    # TODO: fix circular dep
-    from workflowai.core.client._models import ModelInfo
-
     async def reply(
         self,
         run_id: str,
@@ -164,5 +161,3 @@ async def reply(
         ...
 
     async def fetch_completions(self, run_id: str) -> list[Completion]: ...
-
-    async def list_models(self) -> list[ModelInfo]: ...