checkpoint

tensorzero · bglick13 · Jul 21, 2025 · Jul 28, 2025 · Jul 28, 2025 · Jul 28, 2025
commit 67f564da99f0e212e54a2250ba51f81c45ca1864
diff --git a/clients/python/tensorzero/agents.py b/clients/python/tensorzero/agents.py
@@ -1,7 +1,7 @@
 import contextvars
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Any, Optional, TypedDict
+from typing import Any, Generic, Optional, TypeVar
 
 import toml
 from pydantic import BaseModel
@@ -173,17 +173,23 @@ def chat(self) -> TensorZeroAgentsChat:
         return TensorZeroAgentsChat(self)
 
 
-class TensorZeroAgentsRunArgs(TypedDict):
-    system: BaseModel
-    user: BaseModel
+SystemModelT = TypeVar("SystemModelT", bound=BaseModel)
+UserModelT = TypeVar("UserModelT", bound=BaseModel)
 
 
-class TensorZeroAgentsRunner(agents.Runner):
+class TensorZeroAgentsRunArgs(BaseModel, Generic[SystemModelT, UserModelT]):
+    system: SystemModelT
+    user: UserModelT
+
+
+class TensorZeroAgentsRunner(agents.Runner, Generic[SystemModelT, UserModelT]):
     @classmethod
     async def run(
         cls,
         starting_agent: agents.Agent[agents.TContext],
-        input: str | list[agents.TResponseInputItem] | TensorZeroAgentsRunArgs,
+        input: str
+        | list[agents.TResponseInputItem]
+        | TensorZeroAgentsRunArgs[SystemModelT, UserModelT],
         *,
         context: agents.TContext | None = None,
         max_turns: int = agents.run.DEFAULT_MAX_TURNS,
@@ -197,8 +203,8 @@ async def run(
             "TensorZeroAgentsRunner must be called within a with_tensorzero_agents_patched context"
         )
 
-        if isinstance(input, dict):
-            args = input["system"].model_dump() | input["user"].model_dump()
+        if isinstance(input, TensorZeroAgentsRunArgs):
+            args = input.system.model_dump() | input.user.model_dump()
             args = {f"tensorzero::arguments::{k}": v for k, v in args.items()}
             existing_run_config = run_config or agents.run.RunConfig()
             existing_model_settings = (

diff --git a/clients/python/tensorzero/cli.py b/clients/python/tensorzero/cli.py
@@ -241,6 +241,7 @@ def main(argv: list[str] | None = None) -> None:
     # ---------------------------------------------------------------------
     identifiers = _collect_function_identifiers(cfg_dict)
     identifier_module_src = _render_identifier_module(identifiers)
+    output_dir.mkdir(parents=True, exist_ok=True)
 
     (output_dir / "__init__.py").touch(exist_ok=True)
     types_file = output_dir / "function_identifiers.py"

diff --git a/examples/rag-retrieval-augmented-generation/simple-agentic-rag/openai_v1.py b/examples/rag-retrieval-augmented-generation/simple-agentic-rag/openai_v1.py
@@ -2,7 +2,11 @@
 from datetime import datetime
 
 import dotenv
-from tensorzero.agents import TensorZeroAgentsRunner, with_tensorzero_agents_patched
+from tensorzero.agents import (
+    TensorZeroAgentsRunArgs,
+    TensorZeroAgentsRunner,
+    with_tensorzero_agents_patched,
+)
 from tools import answer_question, load_wikipedia_page, search_wikipedia, think
 
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
@@ -43,9 +47,12 @@ async def ask_question(question: str, verbose: bool = False) -> str:
     # Use the Agents SDK Runner - this handles the entire tool loop automatically
     system_model = SystemModel(date=datetime.now().strftime("%Y-%m-%d"))
     user_model = UserModel(question=question)
-    result = await TensorZeroAgentsRunner.run(
+    result = await TensorZeroAgentsRunner[SystemModel, UserModel].run(
         agent,
-        {"system": system_model, "user": user_model},
+        TensorZeroAgentsRunArgs(
+            system=system_model,
+            user=user_model,
+        ),
         max_turns=15,
     )
     if verbose: