agentling/feature/agent

src/agentling/agent.py

@@ -0,0 +1,185 @@
+"""The agent loop.
+
+Agent assembles the framework's primitives — a Model, a set of Tools, and a
+Memory of typed steps — into an async ReAct loop. A single async generator
+(``_run_stream``) drives everything and yields Events; ``run()`` is a thin
+dispatcher that either hands back that stream (``stream=True``) or drains it and
+returns the final answer (``stream=False``).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import AsyncIterator, Callable, Sequence
+from typing import Any
+
+from .events import Event, FinalEvent, StepEvent, ToolCallEvent, ToolResultEvent
+from .memory import ActionStep, FinalStep, Memory, Step, TaskStep, ToolResult
+from .models import ChatMessage, Model, ToolCall
+from .tools import FinalAnswerTool, Tool
+
+DEFAULT_INSTRUCTIONS = (
+    "You are a helpful agent. Use the available tools to gather information or "
+    "take actions, thinking step by step. When you have the answer, call the "
+    "final_answer tool (or simply reply with plain text)."
+)
+
+
+class Agent:
+    """An async tool-calling agent.
+
+    Wires a Model, Tools (with final_answer always available), and a Memory into
+    a ReAct loop. Call run() to execute a task, optionally streaming Events.
+    """
+
+    def __init__(
+        self,
+        model: Model,
+        tools: Sequence[Tool] = (),
+        skills: Sequence[Any] = (),
+        instructions: str | None = None,
+        max_steps: int = 15,
+        step_callbacks: Sequence[Callable[[Step], None]] = (),
+        parallel_tools: bool = True,
+    ) -> None:
+        self.model = model
+        self.instructions = instructions or DEFAULT_INSTRUCTIONS
+        self.max_steps = max_steps
+        self.step_callbacks = list(step_callbacks)
+        self.parallel_tools = parallel_tools
+        self.skills = list(skills)  # wired in Day 6
+
+        self.memory = Memory()
+
+        # final_answer is always available so the model has an explicit way to stop.
+        all_tools = [*tools, FinalAnswerTool()]
+        self.tools: dict[str, Tool] = {tool.name: tool for tool in all_tools}
+        self._tool_schemas = [tool.to_schema() for tool in all_tools]
+
+    def run(
+        self,
+        task: str,
+        *,
+        stream: bool = False,
+        reset: bool = True,
+        max_steps: int | None = None,
+    ) -> Any:
+        """Run the agent on a task.
+
+        With stream=False (default) returns a coroutine resolving to the final
+        answer string:  ``answer = await agent.run(task)``.
+        With stream=True returns an async iterator of Events:
+        ``async for event in agent.run(task, stream=True): ...``.
+        """
+        events = self._run_stream(task, reset=reset, max_steps=max_steps)
+        if stream:
+            return events
+        return self._drain(events)
+
+    async def _drain(self, events: AsyncIterator[Event]) -> str:
+        """Consume the event stream and return the final answer."""
+        answer = ""
+        async for event in events:
+            if isinstance(event, FinalEvent):
+                answer = event.answer
+        return answer
+
+    async def _run_stream(
+        self, task: str, *, reset: bool = True, max_steps: int | None = None
+    ) -> AsyncIterator[Event]:
+        """The core loop: drive the model/tool cycle, yielding Events."""
+        if reset:
+            self.memory.reset()
+
+        self.memory.add(TaskStep(task=task))
+
+        limit = max_steps or self.max_steps
+        for _ in range(limit):
+            started = time.monotonic()
+            messages = self.memory.to_messages(self.instructions)
+            response = await self.model.generate(messages, tools=self._tool_schemas)
+
+            # Forgiving termination: a turn with no tool calls IS the final answer.
+            if not response.tool_calls:
+                self.memory.add(FinalStep(answer=response.content))
+                yield FinalEvent(answer=response.content, usage=response.usage)
+                return
+
+            # Announce every call, then run them (concurrently or in order).
+            for tool_call in response.tool_calls:
+                yield ToolCallEvent(tool_call=tool_call)
+
+            if self.parallel_tools:
+                results: list[ToolResult] = await asyncio.gather(
+                    *(self._execute_tool(tc) for tc in response.tool_calls)
+                )
+            else:
+                results = [await self._execute_tool(tc) for tc in response.tool_calls]
+
+            for result in results:
+                yield ToolResultEvent(result=result)
+
+            action = ActionStep(
+                model_message=response,
+                tool_results=results,
+                usage=response.usage,
+                duration=time.monotonic() - started,
+            )
+            self.memory.add(action)
+            for callback in self.step_callbacks:
+                callback(action)
+            yield StepEvent(step=action)
+
+            # Explicit termination: the model called final_answer (successfully).
+            final = next(
+                (r for r in results if r.name == "final_answer" and not r.is_error),
+                None,
+            )
+            if final is not None:
+                self.memory.add(FinalStep(answer=final.content))
+                yield FinalEvent(answer=final.content, usage=response.usage)
+                return
+
+        # Step limit reached: force one tool-free answer.
+        messages = self.memory.to_messages(self.instructions)
+        messages.append(
+            ChatMessage(
+                role="user",
+                content="Step limit reached. Give your best final answer now.",
+            )
+        )
+        response = await self.model.generate(messages)
+        self.memory.add(FinalStep(answer=response.content))
+        yield FinalEvent(answer=response.content, usage=response.usage)
+
+    async def _execute_tool(self, tool_call: ToolCall) -> ToolResult:
+        """Run one tool call, turning any failure into an error observation."""
+        tool = self.tools.get(tool_call.name)
+        if tool is None:
+            return ToolResult(
+                tool_call_id=tool_call.id,
+                name=tool_call.name,
+                content=(
+                    f"Unknown tool {tool_call.name!r}. "
+                    f"Available: {list(self.tools)}"
+                ),
+                is_error=True,
+            )
+        try:
+            output = await tool.call(tool_call.arguments)
+        except Exception as exc:
+            # Deliberate broad catch: a failing tool becomes an observation the
+            # model can recover from, never a crash. (Day 5 refines the layers.)
+            return ToolResult(
+                tool_call_id=tool_call.id,
+                name=tool_call.name,
+                content=f"{type(exc).__name__}: {exc}",
+                is_error=True,
+            )
+        return ToolResult(
+            tool_call_id=tool_call.id,
+            name=tool_call.name,
+            content=str(output),
+            is_error=False,
+        )

tests/test_agent.py

@@ -0,0 +1,191 @@
+from collections.abc import AsyncIterator, Sequence
+from typing import Any
+
+from agentling.agent import Agent
+from agentling.events import FinalEvent, StepEvent, ToolCallEvent, ToolResultEvent
+from agentling.memory import ActionStep, FinalStep, TaskStep
+from agentling.models import ChatMessage, Delta, ToolCall, Usage
+from agentling.tools import tool
+
+
+class FakeModel:
+    """A scripted model: returns pre-set ChatMessages in order, no network.
+
+    Records the messages it was called with so tests can assert on the loop.
+    """
+
+    def __init__(self, responses: Sequence[ChatMessage]) -> None:
+        self._responses = list(responses)
+        self._index = 0
+        self.calls: list[list[ChatMessage]] = []
+
+    async def generate(
+        self, messages: Sequence[ChatMessage], tools: Sequence[Any] | None = None
+    ) -> ChatMessage:
+        self.calls.append(list(messages))
+        response = self._responses[self._index]
+        self._index += 1
+        return response
+
+    def stream(
+        self, messages: Sequence[ChatMessage], tools: Sequence[Any] | None = None
+    ) -> AsyncIterator[Delta]:
+        raise NotImplementedError
+
+
+def _assistant(
+    content: str = "", tool_calls: list[ToolCall] | None = None
+) -> ChatMessage:
+    return ChatMessage(
+        role="assistant",
+        content=content,
+        tool_calls=tool_calls or [],
+        usage=Usage(1, 1),
+    )
+
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two integers.
+
+    Args:
+        a: The first number.
+        b: The second number.
+    """
+    return a + b
+
+
+# --------------------------------------------------------------------------- #
+# Termination paths
+# --------------------------------------------------------------------------- #
+async def test_forgiving_termination_returns_content() -> None:
+    model = FakeModel([_assistant(content="42")])
+    agent = Agent(model=model)
+
+    answer = await agent.run("What is 6 times 7?")
+
+    assert answer == "42"
+    assert isinstance(agent.memory.steps[0], TaskStep)
+    assert isinstance(agent.memory.steps[-1], FinalStep)
+
+
+async def test_explicit_final_answer() -> None:
+    model = FakeModel(
+        [
+            _assistant(
+                tool_calls=[
+                    ToolCall(id="c1", name="final_answer", arguments={"answer": "done"})
+                ]
+            )
+        ]
+    )
+    agent = Agent(model=model)
+
+    assert await agent.run("finish up") == "done"
+
+
+# --------------------------------------------------------------------------- #
+# Tool execution
+# --------------------------------------------------------------------------- #
+async def test_tool_call_then_answer() -> None:
+    model = FakeModel(
+        [
+            _assistant(
+                tool_calls=[ToolCall(id="c1", name="add", arguments={"a": 2, "b": 3})]
+            ),
+            _assistant(content="The sum is 5."),
+        ]
+    )
+    agent = Agent(model=model, tools=[add])
+
+    answer = await agent.run("add 2 and 3")
+
+    assert answer == "The sum is 5."
+    action = agent.memory.steps[1]
+    assert isinstance(action, ActionStep)
+    assert action.tool_results[0].content == "5"
+    assert action.tool_results[0].is_error is False
+
+
+async def test_multi_step_tool_task() -> None:
+    model = FakeModel(
+        [
+            _assistant(
+                tool_calls=[ToolCall(id="c1", name="add", arguments={"a": 2, "b": 3})]
+            ),
+            _assistant(
+                tool_calls=[ToolCall(id="c2", name="add", arguments={"a": 5, "b": 4})]
+            ),
+            _assistant(
+                tool_calls=[
+                    ToolCall(id="c3", name="final_answer", arguments={"answer": "9"})
+                ]
+            ),
+        ]
+    )
+    agent = Agent(model=model, tools=[add])
+
+    assert await agent.run("compute a couple of sums") == "9"
+    assert len(model.calls) == 3
+
+
+async def test_tool_error_becomes_observation() -> None:
+    @tool
+    def boom() -> str:
+        """Always fails."""
+        raise RuntimeError("kaboom")
+
+    model = FakeModel(
+        [
+            _assistant(tool_calls=[ToolCall(id="c1", name="boom", arguments={})]),
+            _assistant(content="recovered"),
+        ]
+    )
+    agent = Agent(model=model, tools=[boom])
+
+    assert await agent.run("try boom") == "recovered"
+    action = agent.memory.steps[1]
+    assert isinstance(action, ActionStep)
+    assert action.tool_results[0].is_error is True
+
+
+# --------------------------------------------------------------------------- #
+# Streaming
+# --------------------------------------------------------------------------- #
+async def test_streaming_yields_events() -> None:
+    model = FakeModel(
+        [
+            _assistant(
+                tool_calls=[ToolCall(id="c1", name="add", arguments={"a": 1, "b": 1})]
+            ),
+            _assistant(content="2"),
+        ]
+    )
+    agent = Agent(model=model, tools=[add])
+
+    events = [event async for event in agent.run("add", stream=True)]
+
+    assert any(isinstance(e, ToolCallEvent) for e in events)
+    assert any(isinstance(e, ToolResultEvent) for e in events)
+    assert any(isinstance(e, StepEvent) for e in events)
+    assert isinstance(events[-1], FinalEvent)
+    assert events[-1].answer == "2"
+
+
+# --------------------------------------------------------------------------- #
+# max_steps forced answer
+# --------------------------------------------------------------------------- #
+async def test_max_steps_forces_a_final_answer() -> None:
+    # The model never terminates on its own — it always calls a tool.
+    looping = [
+        _assistant(
+            tool_calls=[ToolCall(id=f"c{i}", name="add", arguments={"a": 1, "b": 1})]
+        )
+        for i in range(2)
+    ]
+    forced = _assistant(content="forced final")
+    model = FakeModel([*looping, forced])
+    agent = Agent(model=model, tools=[add], max_steps=2)
+
+    assert await agent.run("loop forever") == "forced final"
+    assert len(model.calls) == 3  # 2 loop steps + 1 forced tool-free answer