Type checking with ty, better protocol defs.

Author: tcdent

pyproject.toml

@@ -47,20 +47,19 @@ dev-dependencies = [
     "pytest-asyncio>=0.23.0",
     "pytest-cov>=4.1.0",
     "ruff>=0.5.0",
-    "mypy>=1.10.0",
+    "ty>=0.0.1a7",
     "fakeredis>=2.32.1",
 ]
 
 [tool.ruff]
 line-length = 100
 target-version = "py311"
 
-[tool.mypy]
-python_version = "3.11"
-warn_return_any = true
-warn_unused_configs = true
-disallow_untyped_defs = true
-disallow_incomplete_defs = true
+[tool.ty.environment]
+python-version = "3.11"
+
+[tool.ty.src]
+include = ["src/agentexec"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"

src/agentexec/__init__.py

@@ -59,7 +59,7 @@ async def search(agent_id: UUID, context: Input) -> Output:
 
 # OpenAI runner is only available if agents package is installed
 try:
-    from agentexec.runners import OpenAIRunner
+    from agentexec.runners import OpenAIRunner  # type: ignore[possibly-missing-import]
 
     __all__.append("OpenAIRunner")
 except ImportError:

src/agentexec/config.py

@@ -74,8 +74,8 @@ class Config(BaseSettings):
     )
 
     state_backend: str = Field(
-        default="redis",
-        description="State backend to use (redis, memory, etc.)",
+        default="agentexec.state.redis_backend",
+        description="State backend to use (fully-qualified module path)",
         validation_alias="AGENTEXEC_STATE_BACKEND",
     )
 

src/agentexec/core/task.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import asyncio
 import inspect
-from typing import Any, Protocol, TypedDict, get_type_hints
+from typing import Any, Protocol, TypeAlias, TypedDict, cast, get_type_hints
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_serializer
@@ -21,21 +20,38 @@ class TaskHandlerKwargs(TypedDict):
     context: BaseModel
 
 
-TaskResult = BaseModel  # Union[BaseModel, tuple[BaseModel, ...]]
+TaskResult: TypeAlias = BaseModel
 
 
-class TaskHandler(Protocol):
+class _SyncTaskHandler(Protocol):
     """Protocol for task handler functions."""
 
     __name__: str
 
     def __call__(
         self,
+        *,
+        agent_id: UUID,
+        context: BaseModel,
+    ) -> TaskResult: ...
+
+
+class _AsyncTaskHandler(Protocol):
+    """Protocol for async task handler functions."""
+
+    __name__: str
+
+    async def __call__(
+        self,
+        *,
         agent_id: UUID,
         context: BaseModel,
     ) -> TaskResult: ...
 
 
+TaskHandler: TypeAlias = _SyncTaskHandler | _AsyncTaskHandler
+
+
 class TaskDefinition:
     """Definition of a task type (created at registration time).
 
@@ -83,6 +99,15 @@ def __init__(
         self.context_type = context_type or self._infer_context_type(handler)
         self.result_type = result_type or self._infer_result_type(handler)
 
+    async def __call__(self, agent_id: UUID, context: BaseModel) -> TaskResult:
+        """Delegate calls to the handler function."""
+        if inspect.iscoroutinefunction(self.handler):
+            handler = cast(_AsyncTaskHandler, self.handler)
+            return await handler(agent_id=agent_id, context=context)
+        else:
+            handler = cast(_SyncTaskHandler, self.handler)
+            return handler(agent_id=agent_id, context=context)
+
     def _infer_context_type(self, handler: TaskHandler) -> type[BaseModel]:
         """Infer context class from handler's type annotations.
 
@@ -244,24 +269,19 @@ async def execute(self) -> TaskResult | None:
         )
 
         try:
-            result: TaskResult
-            if asyncio.iscoroutinefunction(self._definition.handler):
-                result = await self._definition.handler(
-                    agent_id=self.agent_id,
-                    context=self.context,
-                )
-            else:
-                result = self._definition.handler(
-                    agent_id=self.agent_id,
-                    context=self.context,
-                )
-
-            await state.aset_result(
-                self.agent_id,
-                result,
-                ttl_seconds=CONF.result_ttl,
+            result = await self._definition(
+                agent_id=self.agent_id,
+                context=self.context,
             )
 
+            # TODO ensure we are properly supporting None return values
+            if isinstance(result, BaseModel):
+                await state.aset_result(
+                    self.agent_id,
+                    result,
+                    ttl_seconds=CONF.result_ttl,
+                )
+
             activity.update(
                 agent_id=self.agent_id,
                 message=CONF.activity_message_complete,

src/agentexec/pipeline.py

@@ -8,8 +8,10 @@
     Any,
     Callable,
     ClassVar,
-    Coroutine,
+    Protocol,
+    TypeAlias,
     Union,
+    cast,
     get_type_hints,
     get_origin,
     get_args,
@@ -20,13 +22,47 @@
 
 from agentexec import activity
 from agentexec.core import queue
-from agentexec.core.task import Task
+from agentexec.core.task import Task, TaskResult
 
 if TYPE_CHECKING:
     from agentexec.worker import WorkerPool
 
-StepResult = Union[BaseModel, tuple[BaseModel, ...]]
-StepHandler = Callable[..., Union[StepResult, Coroutine[Any, Any, StepResult]]]
+StepResult: TypeAlias = Union[TaskResult, tuple[TaskResult, ...]]
+
+
+class _SyncStepHandler(Protocol):
+    """Protocol for pipeline step handler methods.
+
+    Step handlers are methods on a pipeline class that receive
+    one or more BaseModel context arguments from the previous step.
+    """
+
+    __name__: str
+
+    def __call__(
+        self,
+        instance: _PipelineBase,
+        **kwargs: BaseModel,
+    ) -> StepResult: ...
+
+
+class _AsyncStepHandler(Protocol):
+    """Protocol for async pipeline step handler methods.
+
+    Step handlers are methods on a pipeline class that receive
+    one or more BaseModel context arguments from the previous step.
+    """
+
+    __name__: str
+
+    async def __call__(
+        self,
+        instance: _PipelineBase,
+        **kwargs: BaseModel,
+    ) -> StepResult: ...
+
+
+StepHandler: TypeAlias = _SyncStepHandler | _AsyncStepHandler
 
 
 def _format_pipeline_name(cls: type) -> str:
@@ -38,7 +74,7 @@ class _PipelineBaseMeta(type):
     """Metaclass that registers pipeline subclasses with their bound pipeline."""
 
     @classmethod
-    def bind_pipeline(mcs, pipeline: Pipeline) -> type:
+    def bind_pipeline(mcs, pipeline: Pipeline) -> type[_PipelineBase]:
         """Create a new PipelineBase class bound to the given pipeline.
 
         Args:
@@ -66,8 +102,8 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
         if cls.__name__ == "PipelineBase":
             return
 
-        cls._pipeline.bind_user_pipeline(cls)
-        cls._pipeline.register_task()
+        cls._pipeline._bind_user_pipeline(cls)
+        cls._pipeline._register_task()
 
 
 class StepDefinition:
@@ -76,17 +112,17 @@ class StepDefinition:
     name: str
     order: Any
     handler: StepHandler
-    return_type: type | None
-    param_types: dict[str, type | None]
+    return_type: type[BaseModel] | None
+    param_types: dict[str, type[BaseModel] | None]
     _description: str | None = None
 
     def __init__(
         self,
         name: str,
         order: Any,
         handler: StepHandler,
-        return_type: type | None,
-        param_types: dict[str, type | None],
+        return_type: type[BaseModel] | None,
+        param_types: dict[str, type[BaseModel] | None],
         description: str | None = None,
     ) -> None:
         self.name = name
@@ -103,6 +139,15 @@ def description(self) -> str:
             return self._description
         return self.handler.__name__
 
+    async def __call__(self, instance: _PipelineBase, **kwargs: BaseModel) -> StepResult:
+        """Invoke the step handler."""
+        if asyncio.iscoroutinefunction(self.handler):
+            handler = cast(_AsyncStepHandler, self.handler)
+            return await handler(instance, **kwargs)
+        else:
+            handler = cast(_SyncStepHandler, self.handler)
+            return handler(instance, **kwargs)
+
 
 class Pipeline:
     """Orchestrates multi-step task workflows.
@@ -171,15 +216,15 @@ def Base(self) -> type[_PipelineBase]:
         """
         return _PipelineBaseMeta.bind_pipeline(self)
 
-    def bind_user_pipeline(self, cls: type[_PipelineBase]) -> None:
+    def _bind_user_pipeline(self, cls: type[_PipelineBase]) -> None:
         """Manually set the pipeline implementation class.
 
         Args:
             cls: Pipeline implementation class
         """
         self._user_pipeline_class = cls
 
-    def register_task(self) -> None:
+    def _register_task(self) -> None:
         """
         Register Task handler with the pipeline's worker pool
         """
@@ -190,7 +235,7 @@ def register_task(self) -> None:
 
         self._pool._add_task(
             name=self.name,
-            func=self._run_task,  # type: ignore[arg-type]
+            func=self._run_task,
             context_type=self._input_type,
             result_type=self._output_type,
         )
@@ -306,9 +351,10 @@ async def run(self, context: BaseModel) -> StepResult:
 
     async def _run_task(
         self,
+        *,
         agent_id: UUID,
         context: BaseModel,
-    ) -> StepResult:
+    ) -> TaskResult:
         """Run the pipeline as a task handler.
 
         Args:
@@ -330,6 +376,7 @@ async def _run_task(
             )
             _context = await self._run_step(step, _context)
 
+        assert isinstance(_context, TaskResult), "Final step must return BaseModel"
         return _context
 
     async def _run_step(
@@ -355,9 +402,7 @@ async def _run_step(
         else:
             kwargs = {}
 
-        if asyncio.iscoroutinefunction(step.handler):
-            return await step.handler(instance, **kwargs)
-        return step.handler(instance, **kwargs)
+        return await step(instance, **kwargs)
 
     def _validate_type_flow(self) -> None:
         """Verify that step return types match next step's parameters.

src/agentexec/runners/openai.py

@@ -5,6 +5,7 @@
 from agents import Agent, MaxTurnsExceeded, Runner, function_tool
 from agents.items import TResponseInputItem
 from agents.result import RunResult, RunResultStreaming
+from openai.types.responses.easy_input_message_param import EasyInputMessageParam
 
 from agentexec.runners.base import BaseAgentRunner, _RunnerTools
 
@@ -29,7 +30,7 @@ def _extract_input(e: MaxTurnsExceeded) -> list[TResponseInputItem]:
     final_input: list[TResponseInputItem] = (
         list(e.run_data.input)
         if isinstance(e.run_data.input, list)
-        else [{"role": "user", "content": e.run_data.input}]
+        else [EasyInputMessageParam(role="user", content=e.run_data.input)]
     )
 
     # Add all the conversation items that were generated
@@ -107,6 +108,12 @@ def __init__(
         # Override with OpenAI-specific tools
         self.tools = _OpenAIRunnerTools(self.agent_id)
 
+    def _wrap_up_prompt(self) -> EasyInputMessageParam:
+        return EasyInputMessageParam(
+            role="system",
+            content=self.prompts.wrap_up,
+        )
+
     async def run(
         self,
         agent: Agent[Any],
@@ -142,12 +149,7 @@ async def run(
 
             logger.info("Max turns exceeded, attempting recovery")
             final_input = _extract_input(e)
-            final_input.append(
-                {
-                    "role": "user",
-                    "content": self.prompts.wrap_up,
-                }
-            )
+            final_input.append(self._wrap_up_prompt())
             result = await Runner.run(
                 agent,
                 final_input,
@@ -212,12 +214,7 @@ async def run_streamed(
 
             logger.info("Max turns exceeded, attempting recovery")
             final_input = _extract_input(e)
-            final_input.append(
-                {
-                    "role": "user",
-                    "content": self.prompts.wrap_up,
-                }
-            )
+            final_input.append(self._wrap_up_prompt())
             result = Runner.run_streamed(
                 agent,
                 final_input,