diff --git a/src/agentling/events.py b/src/agentling/events.py new file mode 100644 index 0000000..301e659 --- /dev/null +++ b/src/agentling/events.py @@ -0,0 +1,56 @@ +"""Streaming events emitted by the agent loop. + +Memory is the durable, structured record of a run. Events are runtime +notifications emitted as the run progresses, letting callers render output, +inspect tool execution, and react to completion in real time. + +StepEvent connects the live event stream to memory by carrying the step that +was just recorded. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .memory import Step, ToolResult +from .models import ToolCall, Usage + + +@dataclass(frozen=True, slots=True) +class TextDelta: + """A chunk of streamed assistant text.""" + + text: str + + +@dataclass(frozen=True, slots=True) +class ToolCallEvent: + """Emitted before the agent executes a tool call.""" + + tool_call: ToolCall + + +@dataclass(frozen=True, slots=True) +class ToolResultEvent: + """Emitted after a tool call completes, successfully or with an error.""" + + result: ToolResult + + +@dataclass(frozen=True, slots=True) +class StepEvent: + """Emitted after a step has been recorded to memory.""" + + step: Step + + +@dataclass(frozen=True, slots=True) +class FinalEvent: + """Emitted once when the run completes, carrying the answer and total usage.""" + + answer: str + usage: Usage | None = None + + +# Public union type for the values yielded by the agent's event stream. +Event = TextDelta | ToolCallEvent | ToolResultEvent | StepEvent | FinalEvent diff --git a/src/agentling/memory.py b/src/agentling/memory.py new file mode 100644 index 0000000..415e72b --- /dev/null +++ b/src/agentling/memory.py @@ -0,0 +1,202 @@ +"""Typed step records and the agent's memory. + +The agent loop records each turn as a typed Step (TaskStep, ActionStep, +FinalStep) rather than as raw messages. Steps carry metadata — token usage, +timing, per-tool results and errors — and know how to render themselves back +into the ChatMessage list the model sees. Memory holds the steps, renders the +full conversation via to_messages(), and round-trips to JSON for persistence +and replay. +""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass, field +from typing import Any, ClassVar + +from .models import ChatMessage, ToolCall, Usage + + +@dataclass +class ToolResult: + """The outcome of a single executed tool call.""" + + tool_call_id: str + name: str + content: str + is_error: bool = False + + +@dataclass +class TaskStep: + """A task that started (or, in multi-turn, continued) the run.""" + + task: str + + def to_messages(self) -> list[ChatMessage]: + """Render as a single user message.""" + + return [ChatMessage(role="user", content=self.task)] + + +@dataclass +class ActionStep: + """One agent loop iteration. + + Stores the assistant message the model produced, any tool results from that + message, and optional execution metadata (error, usage, timing). + """ + + model_message: ChatMessage + tool_results: list[ToolResult] = field(default_factory=list) + error: str | None = None + usage: Usage | None = None + duration: float | None = None + + def to_messages(self) -> list[ChatMessage]: + """Render the assistant turn, then one tool message per result.""" + + # model_message is already role="assistant"; reuse it as-is. + messages: list[ChatMessage] = [self.model_message] + + for result in self.tool_results: + content = result.content + # Failed tool calls are returned to the model as observations so the + # next turn can recover. The tool name is included because a single + # step may contain several calls. + if result.is_error: + content = ( + f"Error from {result.name!r}: {result.content}. " + "Fix the arguments and try again." + ) + messages.append( + ChatMessage( + role="tool", content=content, tool_call_id=result.tool_call_id + ) + ) + + return messages + + +@dataclass +class FinalStep: + """The terminal answer produced by the run.""" + + answer: str + + def to_messages(self) -> list[ChatMessage]: + """Render as the assistant's final message (keeps multi-turn history whole).""" + + return [ChatMessage(role="assistant", content=self.answer)] + + +# A run's history is a sequence of these three step kinds. +Step = TaskStep | ActionStep | FinalStep + + +@dataclass +class Memory: + """Ordered history of an agent run. + + Stores typed steps rather than raw provider messages, renders them into the + ChatMessage list the model sees, and serializes them for persistence/replay. + """ + + # Discriminator tag written into each serialized step so from_dict() can + # rebuild the right class. ClassVar keeps it class-level config, not a field. + _KIND: ClassVar[dict[type, str]] = { + TaskStep: "task", + ActionStep: "action", + FinalStep: "final", + } + + steps: list[Step] = field(default_factory=list) + + def add(self, step: Step) -> None: + """Append a step to the run's history.""" + + self.steps.append(step) + + def reset(self) -> None: + """Clear all steps (used by run(reset=True) to start fresh).""" + + self.steps.clear() + + def to_messages(self, system_prompt: str) -> list[ChatMessage]: + """Render the full conversation: system prompt, then every step.""" + + # The system prompt is runtime configuration, not run history. Injecting + # it here means it is always first and never duplicated in stored steps. + messages: list[ChatMessage] = [ + ChatMessage(role="system", content=system_prompt) + ] + for step in self.steps: + messages.extend(step.to_messages()) + return messages + + def to_dict(self) -> dict[str, Any]: + """Serialize to a JSON-safe dict, tagging each step with its kind.""" + + # asdict() recurses the whole nested tree (ChatMessage -> ToolCalls -> + # Usage, ToolResults) into plain dicts but drops type information, which + # the load path restores from each step's "type" tag. The version field + # lets the format evolve without breaking older saved runs. + return { + "version": 1, + "steps": [ + {"type": self._KIND[type(step)], "data": asdict(step)} + for step in self.steps + ], + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Memory: + """Rebuild a Memory from the dict produced by to_dict().""" + + return cls(steps=[_step_from_dict(entry) for entry in data.get("steps", [])]) + + def dump_json(self) -> str: + """Serialize the memory to a JSON string.""" + + return json.dumps(self.to_dict()) + + @classmethod + def load_json(cls, raw: str) -> Memory: + """Rebuild a Memory from a JSON string produced by dump_json().""" + + return cls.from_dict(json.loads(raw)) + + +# dataclasses.asdict() converts nested dataclasses to dicts but does not record +# their types, so the load path rebuilds each nested object explicitly. +def _chat_message_from_dict(d: dict[str, Any]) -> ChatMessage: + """Reconstruct a ChatMessage (and its nested ToolCalls/Usage) from a dict.""" + + return ChatMessage( + role=d["role"], + content=d["content"], + tool_calls=[ToolCall(**tc) for tc in d["tool_calls"]], + tool_call_id=d["tool_call_id"], + usage=Usage(**d["usage"]) if d["usage"] else None, + ) + + +def _step_from_dict(entry: dict[str, Any]) -> Step: + """Reconstruct a Step from a {"type", "data"} entry, dispatching on the tag.""" + + kind, data = entry["type"], entry["data"] + + if kind == "task": + return TaskStep(**data) + if kind == "final": + return FinalStep(**data) + if kind == "action": + return ActionStep( + model_message=_chat_message_from_dict(data["model_message"]), + tool_results=[ToolResult(**tr) for tr in data["tool_results"]], + error=data["error"], + usage=Usage(**data["usage"]) if data["usage"] else None, + duration=data["duration"], + ) + + raise ValueError(f"Unknown step type: {kind!r}") diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 0000000..7c48867 --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,152 @@ +import json + +import pytest + +from agentling.memory import ( + ActionStep, + FinalStep, + Memory, + TaskStep, + ToolResult, +) +from agentling.models import ChatMessage, ToolCall, Usage + + +def _action_step() -> ActionStep: + """A representative ActionStep: one tool call, one successful result.""" + return ActionStep( + model_message=ChatMessage( + role="assistant", + content="", + tool_calls=[ + ToolCall(id="c1", name="get_weather", arguments={"city": "Paris"}) + ], + usage=Usage(10, 5), + ), + tool_results=[ + ToolResult(tool_call_id="c1", name="get_weather", content="Sunny") + ], + usage=Usage(10, 5), + duration=1.2, + ) + + +# --------------------------------------------------------------------------- # +# Per-step rendering +# --------------------------------------------------------------------------- # +def test_task_step_renders_user_message() -> None: + msgs = TaskStep(task="hi").to_messages() + assert [m.role for m in msgs] == ["user"] + assert msgs[0].content == "hi" + + +def test_action_step_renders_assistant_then_tool() -> None: + msgs = _action_step().to_messages() + assert [m.role for m in msgs] == ["assistant", "tool"] + assert msgs[1].content == "Sunny" + assert msgs[1].tool_call_id == "c1" + + +def test_action_step_renders_error_observation() -> None: + step = ActionStep( + model_message=ChatMessage( + role="assistant", + tool_calls=[ToolCall(id="c1", name="f", arguments={})], + ), + tool_results=[ + ToolResult(tool_call_id="c1", name="f", content="boom", is_error=True) + ], + ) + tool_msg = step.to_messages()[1] + assert tool_msg.content == "Error from 'f': boom. Fix the arguments and try again." + + +def test_final_step_renders_assistant_message() -> None: + msgs = FinalStep(answer="done").to_messages() + assert [m.role for m in msgs] == ["assistant"] + assert msgs[0].content == "done" + + +# --------------------------------------------------------------------------- # +# Memory.to_messages +# --------------------------------------------------------------------------- # +def test_to_messages_full_shape() -> None: + mem = Memory() + mem.add(TaskStep(task="weather?")) + mem.add(_action_step()) + mem.add(FinalStep(answer="Sunny.")) + + msgs = mem.to_messages("You are helpful.") + assert [m.role for m in msgs] == [ + "system", + "user", + "assistant", + "tool", + "assistant", + ] + assert msgs[0].content == "You are helpful." + + +def test_to_messages_empty_is_just_the_system_prompt() -> None: + msgs = Memory().to_messages("sys") + assert [m.role for m in msgs] == ["system"] + + +# --------------------------------------------------------------------------- # +# add / reset (multi-turn support) +# --------------------------------------------------------------------------- # +def test_add_and_reset() -> None: + mem = Memory() + mem.add(TaskStep(task="a")) + mem.add(TaskStep(task="b")) + assert len(mem.steps) == 2 + + mem.reset() + assert mem.steps == [] + + +# --------------------------------------------------------------------------- # +# JSON round-trip (persistence / replay) +# --------------------------------------------------------------------------- # +def test_json_round_trip_preserves_all_steps() -> None: + mem = Memory() + mem.add(TaskStep(task="weather?")) + mem.add(_action_step()) + mem.add(FinalStep(answer="Sunny.")) + + restored = Memory.load_json(mem.dump_json()) + assert restored.steps == mem.steps + + +def test_dump_json_tags_each_step_with_its_kind() -> None: + mem = Memory() + mem.add(TaskStep(task="hi")) + mem.add(_action_step()) + mem.add(FinalStep(answer="bye")) + + data = json.loads(mem.dump_json()) + assert [entry["type"] for entry in data["steps"]] == ["task", "action", "final"] + + +def test_round_trip_preserves_error_flag() -> None: + mem = Memory() + mem.add( + ActionStep( + model_message=ChatMessage( + role="assistant", + tool_calls=[ToolCall(id="c1", name="f", arguments={})], + ), + tool_results=[ + ToolResult(tool_call_id="c1", name="f", content="boom", is_error=True) + ], + ) + ) + + action = Memory.load_json(mem.dump_json()).steps[0] + assert isinstance(action, ActionStep) + assert action.tool_results[0].is_error is True + + +def test_from_dict_rejects_unknown_step_type() -> None: + with pytest.raises(ValueError, match="Unknown step type"): + Memory.from_dict({"steps": [{"type": "bogus", "data": {}}]})