feat: improve agent example

Author: brian14708

examples/agent.py

@@ -2,31 +2,31 @@
 
 import argparse
 import asyncio
-from collections.abc import AsyncGenerator
+import os
+import random
+import readline
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Literal, cast
 from typing_extensions import override
 
-from openai import AsyncOpenAI
+from openai import AsyncOpenAI, pydantic_function_tool
 from openai.lib.streaming.chat import ChatCompletionStreamState
 from openai.types.chat import (
-    ChatCompletionChunk,
+    ChatCompletion,
     ChatCompletionMessageParam,
+    ChatCompletionMessageToolCallUnion,
 )
-from pydantic import TypeAdapter
+from pydantic import BaseModel, Field, TypeAdapter
+from rich.console import Console
 
 import duron
-from duron import RunOptions, op
+from duron import RunOptions
 from duron.codec import Codec
 from duron.contrib.storage import FileLogStorage
 
 if TYPE_CHECKING:
     from typing import Any
 
-    from openai.types.chat import (
-        ParsedChatCompletionMessage,
-    )
-
     from duron.codec import JSONValue
     from duron.typing import TypeHint
 
@@ -52,22 +52,90 @@ def decode_json(self, encoded: JSONValue, expected_type: TypeHint[Any]) -> objec
         return cast("object", TypeAdapter(expected_type).validate_python(encoded))
 
 
+@duron.op
+async def do_input() -> str:  # noqa: RUF029
+    try:
+        return input("> ")  # noqa: ASYNC250
+    except EOFError:
+        os._exit(0)
+    except KeyboardInterrupt:
+        os._exit(1)
+
+
 @duron.fn(codec=PydanticCodec())
 async def agent_fn(ctx: duron.Context) -> None:
-    result = await completion(
-        ctx,
-        messages=[
-            {
-                "role": "system",
-                "content": "You are a helpful assistant!",
-            },
-            {
-                "role": "user",
-                "content": "Say hello to Duron.",
-            },
-        ],
-    )
-    print(result.content)
+    console = Console()
+    history: list[ChatCompletionMessageParam] = [
+        {
+            "role": "system",
+            "content": "You are a helpful assistant!",
+        },
+    ]
+    while True:
+        msg = await ctx.run(do_input)
+        history.append({
+            "role": "user",
+            "content": msg,
+        })
+        console.print("[bold cyan]     USER[/bold cyan]", msg)
+        while True:
+            result = await completion(
+                ctx,
+                messages=history,
+            )
+            if result.choices[0].message.content:
+                console.print(
+                    "[bold red]ASSISTANT[/bold red] ", result.choices[0].message.content
+                )
+            history.append({
+                "role": "assistant",
+                "content": result.choices[0].message.content,
+                "tool_calls": [
+                    {
+                        "id": toolcall.id,
+                        "type": "function",
+                        "function": {
+                            "name": toolcall.function.name,
+                            "arguments": toolcall.function.arguments,
+                        },
+                    }
+                    for toolcall in result.choices[0].message.tool_calls or []
+                    if toolcall.type == "function"
+                ],
+            })
+            if not result.choices[0].message.tool_calls:
+                break
+
+            tasks: list[asyncio.Task[tuple[str, str]]] = []
+            for tool_call in result.choices[0].message.tool_calls:
+                console.print("[bold yellow]     CALL[/bold yellow]", tool_call.id)
+                console.print(tool_call.model_dump_json())
+                tasks.append(asyncio.create_task(ctx.run(call_tool, None, tool_call)))
+            for id_, tool_result in await asyncio.gather(*tasks):
+                console.print("[bold cyan]     TOOL[/bold cyan]", id_)
+                console.print(tool_result)
+                history.append({
+                    "role": "tool",
+                    "tool_call_id": id_,
+                    "content": tool_result,
+                })
+
+
+@duron.op
+async def call_tool(params: ChatCompletionMessageToolCallUnion) -> tuple[str, str]:  # noqa: RUF029
+    if params.type != "function" or not params.function.name:
+        return params.id, '{"status": "error", "message": "Invalid tool call"}'
+    tool_name = params.function.name
+
+    if tool_name == "get_temperature":
+        return params.id, get_temperature(
+            TemperatureInput.model_validate_json(params.function.arguments or "{}")
+        ).model_dump_json()
+    if tool_name == "get_forecast":
+        return params.id, get_forecast(
+            ForecastInput.model_validate_json(params.function.arguments or "{}")
+        ).model_dump_json()
+    return params.id, '{"status": "error", "message": "Unknown tool"}'
 
 
 async def main() -> None:
@@ -80,7 +148,7 @@ async def main() -> None:
     )
     args = parser.parse_args()
 
-    log_storage = FileLogStorage(Path("logs") / f"{args.session_id}.json")
+    log_storage = FileLogStorage(Path("logs") / f"{args.session_id}.jsonl")
     async with agent_fn.invoke(log_storage) as job:
         await job.start()
         await job.wait()
@@ -89,62 +157,116 @@ async def main() -> None:
 async def completion(
     ctx: duron.Context,
     messages: list[ChatCompletionMessageParam],
-) -> ParsedChatCompletionMessage[None]:
-    @op(
-        checkpoint=True,
-        action_type=ChatCompletionChunk,
-        return_type=ChatCompletionStreamState | None,
-        initial=lambda: None,
-        reducer=lambda a, b: (
-            s := a or ChatCompletionStreamState(),
-            s.handle_chunk(b),
-            s,
-        )[-1],
-    )
-    async def _completion_stream(
-        prev: ChatCompletionStreamState | None,
+) -> ChatCompletion:
+    @duron.op
+    async def _completion(
         messages: list[ChatCompletionMessageParam],
-    ) -> AsyncGenerator[ChatCompletionChunk, ChatCompletionStreamState | None]:
-        if prev:
-            msg = prev.current_completion_snapshot.choices[0].message
-            messages = [
-                *messages,
-                {
-                    "role": "assistant",
-                    "content": msg.content,
-                    "tool_calls": (
-                        {
-                            "id": call.id,
-                            "type": call.type,
-                            "function": {
-                                "name": call.function.name,
-                                "arguments": call.function.arguments,
-                            },
-                        }
-                        for call in msg.tool_calls
-                    )
-                    if msg.tool_calls
-                    else (),
-                },
-            ]
+    ) -> ChatCompletion:
+        state = ChatCompletionStreamState()
         async for chunk in await client.chat.completions.create(
             messages=messages,
+            tools=[
+                pydantic_function_tool(
+                    TemperatureInput,
+                    name="get_temperature",
+                    description="Get current temperature for a location",
+                ),
+                pydantic_function_tool(
+                    ForecastInput,
+                    name="get_forecast",
+                    description="Get weather forecast for a location",
+                ),
+            ],
             model=DEFAULT_MODEL,
             stream=True,
         ):
             if chunk.object:  # type: ignore[redundant-expr]
-                yield chunk
+                _ = state.handle_chunk(chunk)
+        return state.get_final_completion()
 
-    state = await ctx.run(
-        _completion_stream,
+    return await ctx.run(
+        _completion,
         RunOptions(
             metadata={"type": "chat.completions.create"},
         ),
         messages,
     )
-    assert state  # noqa: S101
-    return state.get_final_completion().choices[0].message
+
+
+# tools
+
+
+class TemperatureInput(BaseModel):
+    location: str = Field(..., description="Location to get weather for")
+    unit: Literal["celsius", "fahrenheit"] = Field(
+        default="celsius", description="Temperature unit"
+    )
+
+
+class TemperatureOutput(BaseModel):
+    location: str
+    temperature: float | None
+    unit: Literal["celsius", "fahrenheit"]
+    status: Literal["success", "error"]
+    message: str | None = None
+
+
+class ForecastInput(BaseModel):
+    location: str = Field(..., description="Location for forecast")
+    days: int = Field(default=3, ge=1, le=7, description="Number of days (1-7)")
+
+
+class ForecastDay(BaseModel):
+    day: int
+    high: float
+    low: float
+    humidity: int
+    wind_speed: int
+
+
+class ForecastOutput(BaseModel):
+    location: str
+    forecast: list[ForecastDay]
+    status: Literal["success", "error"]
+
+
+# Simplified tool implementations
+def get_temperature(input_data: TemperatureInput) -> TemperatureOutput:
+    # Generate random temperature based on realistic ranges
+    if input_data.unit == "celsius":
+        temp = round(random.uniform(0, 37), 1)
+    else:  # fahrenheit
+        temp = round(random.uniform(-4, 113), 1)
+
+    return TemperatureOutput(
+        location=input_data.location,
+        temperature=temp,
+        unit=input_data.unit,
+        status="success",
+    )
+
+
+def get_forecast(input_data: ForecastInput) -> ForecastOutput:
+    forecast: list[ForecastDay] = []
+    for i in range(input_data.days):
+        high = round(random.uniform(15, 35), 1)
+        low = round(random.uniform(5, high - 5), 1)  # Low is always less than high
+
+        forecast.append(
+            ForecastDay(
+                day=i + 1,
+                high=high,
+                low=low,
+                humidity=random.randint(30, 90),
+                wind_speed=random.randint(5, 25),
+            )
+        )
+
+    return ForecastOutput(
+        location=input_data.location, forecast=forecast, status="success"
+    )
 
 
 if __name__ == "__main__":
+    _ = readline
     asyncio.run(main())

pyproject.toml

@@ -36,6 +36,7 @@ docs = [
 examples = [
     "pydantic>=2.11.9",
     "openai>=2.2.0",
+    "rich>=14.1.0",
 ]
 
 [tool.ruff.lint]
@@ -51,7 +52,7 @@ future-annotations = true
 
 [tool.ruff.lint.per-file-ignores]
 "tests/**/*.py" = ["S101", "S311"]
-"examples/**/*.py" = ["T201"]
+"examples/**/*.py" = ["T201", "S311"]
 
 [tool.ruff.lint.isort]
 extra-standard-library = ["typing_extensions"]

src/duron/_core/context.py

@@ -108,13 +108,13 @@ async def run(
                 await stream.discard()
                 return await stream
 
-        return_type = (
-            fn.return_type
-            if isinstance(fn, Op) and fn.return_type
-            else self._task.codec.inspect_function(
-                cast("Callable[..., object]", fn),
-            ).return_type
-        )
+        if isinstance(fn, Op):
+            if fn.return_type:
+                return_type = fn.return_type
+            else:
+                return_type = self._task.codec.inspect_function(fn.fn).return_type
+        else:
+            return_type = self._task.codec.inspect_function(fn).return_type
 
         metadata = options.metadata if options else None
         op = create_op(