diff --git a/.claude/.codebase-info/.map-state.json b/.claude/.codebase-info/.map-state.json index b8367981..2d058f09 100644 --- a/.claude/.codebase-info/.map-state.json +++ b/.claude/.codebase-info/.map-state.json @@ -1,8 +1,8 @@ { "tool": "codebase-mapper", "version": "2.0.0", - "mappedAt": "2026-06-13", - "gitCommit": "3d33e62f2b6796b7259c1e3e5d77c4313a56f22d", + "mappedAt": "2026-07-05", + "gitCommit": "28099f5aa6eccb491c66be4bcd22532f15bb407f", "documents": [ "architecture.md", "tech-landscape.md", diff --git a/.claude/.codebase-info/architecture.md b/.claude/.codebase-info/architecture.md index e36947aa..d09f077a 100644 --- a/.claude/.codebase-info/architecture.md +++ b/.claude/.codebase-info/architecture.md @@ -1,6 +1,6 @@ # Architecture -*Last Updated: 2026-06-13* +*Last Updated: 2026-07-05* ## Summary @@ -47,7 +47,7 @@ Tool install (atomic-assembler TUI): |-----------|------|----------------| | Core agent | `atomic-agents/atomic_agents/agents/atomic_agent.py` | `AtomicAgent`, `AgentConfig`, run / stream / async methods | | Base contracts | `atomic-agents/atomic_agents/base/` | `BaseIOSchema`, `BaseTool`, `BaseToolConfig`, `BaseResource`, `BasePrompt` | -| Context | `atomic-agents/atomic_agents/context/` | `SystemPromptGenerator`, `BaseDynamicContextProvider`, `ChatHistory` | +| Context | `atomic-agents/atomic_agents/context/` | `SystemPromptGenerator`, `BaseDynamicContextProvider`, `BaseChatHistory` (pluggable memory contract) + `ChatHistory` | | Connectors | `atomic-agents/atomic_agents/connectors/mcp/` | Model Context Protocol tools / resources / prompts | | Utils | `atomic-agents/atomic_agents/utils/` | Token counting (LiteLLM), tool-message formatting | | Assembler (CLI) | `atomic-assembler/atomic_assembler/` | Textual TUI to fetch/install forge tools | diff --git a/.claude/.codebase-info/directory-structure.md b/.claude/.codebase-info/directory-structure.md index 09585ca6..6a82c412 100644 --- a/.claude/.codebase-info/directory-structure.md +++ b/.claude/.codebase-info/directory-structure.md @@ -1,6 +1,6 @@ # Directory Structure -*Last Updated: 2026-06-13* +*Last Updated: 2026-07-05* ## Root Layout @@ -10,7 +10,7 @@ atomic-agents/ # repo root (uv workspace) │ └── atomic_agents/ # import package │ ├── agents/ # AtomicAgent, AgentConfig │ ├── base/ # BaseIOSchema, BaseTool, BaseResource, BasePrompt -│ ├── context/ # SystemPromptGenerator, ChatHistory, context providers +│ ├── context/ # SystemPromptGenerator, BaseChatHistory/ChatHistory, context providers │ ├── connectors/mcp/ # Model Context Protocol integration │ └── utils/ # token counter, tool-message formatting │ └── tests/ # pytest suite (agents/, base/, context/, connectors/, utils/) @@ -19,7 +19,7 @@ atomic-agents/ # repo root (uv workspace) ├── atomic-forge/ # library of standalone tools (NOT a package) │ ├── tools// # one folder per tool: tool/.py, tests/, pyproject.toml │ └── guides/ # tool authoring guides (e.g. tool_structure.md) -├── atomic-examples/ # 15 runnable example apps (each its own project) +├── atomic-examples/ # 16 runnable example apps (each its own project) ├── docs/ # Sphinx + MyST documentation (api/, guides/, examples/) ├── guides/ # DEV_GUIDE.md and contributor guides ├── scripts/ # sync_version.py, generate_llms_files.py @@ -49,9 +49,9 @@ and copies a selected tool into the user's project. `pyproject.toml`/`requirements.txt`. Tools are copied into user projects, not pip-installed. ### `atomic-examples/` -15 standalone example apps (`quickstart`, `rag-chatbot`, `deep-research`, `web-search-agent`, -`mcp-agent`, `fastapi-memory`, `youtube-summarizer`, …), each with its own `pyproject.toml`. These -are excluded from the workspace build. +16 standalone example apps (`quickstart`, `rag-chatbot`, `deep-research`, `web-search-agent`, +`mcp-agent`, `fastapi-memory`, `persistent-memory`, `youtube-summarizer`, …), each with its own +`pyproject.toml`. These are excluded from the workspace build. ### `docs/` and `guides/` `docs/` is a Sphinx + MyST site (`api/` reference, `guides/`, `examples/`, `conf.py`), deployed to diff --git a/.claude/.codebase-info/entry-points.md b/.claude/.codebase-info/entry-points.md index 8dbf1948..e1129128 100644 --- a/.claude/.codebase-info/entry-points.md +++ b/.claude/.codebase-info/entry-points.md @@ -1,6 +1,6 @@ # Entry Points -*Last Updated: 2026-06-13* +*Last Updated: 2026-07-05* ## 1. Library API (the primary entry point) @@ -45,6 +45,7 @@ MCP server tools/resources/prompts become agent components via `atomic_agents.co | `mcp-agent` | MCP client/server, multiple transports | | `progressive-disclosure` | efficient MCP tool loading (3 servers, 24 tools) | | `fastapi-memory` | multi-user / multi-session memory behind a FastAPI service | +| `persistent-memory` | custom `BaseChatHistory` backend: cross-session recall via stdlib `sqlite3` (no extra deps) | | `hooks-example` | monitoring, error handling, retry via hooks | | `basic-multimodal` / `nested-multimodal` | images / PDFs in (nested) schemas | | `basic-pdf-analysis` | PDF analysis with a multimodal model | diff --git a/.claude/.codebase-info/modules.md b/.claude/.codebase-info/modules.md index b335683b..05be8b89 100644 --- a/.claude/.codebase-info/modules.md +++ b/.claude/.codebase-info/modules.md @@ -1,6 +1,6 @@ # Key Modules -*Last Updated: 2026-06-13* +*Last Updated: 2026-07-05* ## Core framework — `atomic-agents/atomic_agents/` @@ -24,8 +24,11 @@ - **Location:** `atomic-agents/atomic_agents/context/` - **Purpose:** System-prompt assembly and conversation memory. - **Key files:** `system_prompt_generator.py` (`SystemPromptGenerator`, `BaseDynamicContextProvider`), - `chat_history.py` (`ChatHistory`, `Message` — multimodal Image/Audio/PDF, turn grouping, - `dump()`/`load()` serialization). + `base_chat_history.py` (`BaseChatHistory` — interface-only ABC declaring the memory contract + `AtomicAgent` depends on; the pluggable seam for custom/persistent backends), + `chat_history.py` (`ChatHistory`, `Message` — the built-in in-memory implementation of + `BaseChatHistory`: multimodal Image/Audio/PDF, turn grouping, `dump()`/`load()` serialization). +- **Note:** `AgentConfig.history` is typed to `BaseChatHistory`, so any conforming backend drops in. ### connectors/mcp - **Location:** `atomic-agents/atomic_agents/connectors/mcp/` @@ -58,4 +61,4 @@ ### atomic-examples - **Location:** `atomic-examples/` -- **Purpose:** 15 runnable reference apps. Catalog in `entry-points.md`. +- **Purpose:** 16 runnable reference apps. Catalog in `entry-points.md`. diff --git a/.claude/.codebase-info/patterns.md b/.claude/.codebase-info/patterns.md index 7267dc02..d1797297 100644 --- a/.claude/.codebase-info/patterns.md +++ b/.claude/.codebase-info/patterns.md @@ -1,6 +1,6 @@ # Patterns & Conventions -*Last Updated: 2026-06-13* +*Last Updated: 2026-07-05* ## Atomicity Build with small, single-purpose, composable parts ("LEGO blocks"): each agent, tool, and context @@ -25,9 +25,15 @@ the next's input. See `atomic-examples/deep-research` and `orchestration-agent`. schema(s) → config → tool class + logic → example usage. ## Memory -- `ChatHistory` stores typed `Message`s grouped into turns (a user+assistant pair shares a `turn_id`), - is multimodal-aware (Image/Audio/PDF), and supports `dump()`/`load()`. `AtomicAgent` trims the - oldest whole turns to honor `max_context_tokens`. +- `BaseChatHistory` (`context/base_chat_history.py`) is an interface-only ABC declaring the memory + contract `AtomicAgent` relies on (`add_message`, `get_history`, turn handling, `dump()`/`load()`, + `copy()`, plus the `history`/`current_turn_id` attributes). It is the documented, dependency-free + seam for plugging in custom/persistent backends; `AgentConfig.history` is typed to it. +- `ChatHistory` is the built-in implementation: stores typed `Message`s grouped into turns (a + user+assistant pair shares a `turn_id`), is multimodal-aware (Image/Audio/PDF), and supports + `dump()`/`load()`. `AtomicAgent` trims the oldest whole turns to honor `max_context_tokens`. +- Custom backend pattern: subclass `ChatHistory` and override `add_message`/`load` to persist (see + the `persistent-memory` example and the "Writing a Custom Memory Backend" guide section). ## Error handling & retries - Lean on Instructor's validation/retry of structured outputs, plus hook events (`parse:error`, …) diff --git a/.coveragerc b/.coveragerc index d1c40623..510f0c8e 100644 --- a/.coveragerc +++ b/.coveragerc @@ -5,6 +5,8 @@ exclude_lines = @abc.abstractmethod if __name__ == .__main__.: if __name__ == '__main__': + if TYPE_CHECKING: + if typing.TYPE_CHECKING: pragma: no cover # pragma: no cover diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py index 3f77892c..d2ccfe6e 100644 --- a/atomic-agents/atomic_agents/agents/atomic_agent.py +++ b/atomic-agents/atomic_agents/agents/atomic_agent.py @@ -5,6 +5,7 @@ from typing import Optional, Type, Generator, AsyncGenerator, get_args, get_origin, Dict, List, Callable, Any import logging from atomic_agents.context.chat_history import ChatHistory +from atomic_agents.context.base_chat_history import BaseChatHistory from atomic_agents.context.system_prompt_generator import ( BaseDynamicContextProvider, SystemPromptGenerator, @@ -66,7 +67,7 @@ class BasicChatOutputSchema(BaseIOSchema): class AgentConfig(BaseModel): client: instructor.core.client.Instructor = Field(..., description="Client for interacting with the language model.") model: str = Field(default="gpt-5-mini", description="The model to use for generating responses.") - history: Optional[ChatHistory] = Field(default=None, description="History component for storing chat history.") + history: Optional[BaseChatHistory] = Field(default=None, description="History component for storing chat history.") system_prompt_generator: Optional[BaseSystemPromptGenerator] = Field( default=None, description=( @@ -117,7 +118,7 @@ class AtomicAgent[InputSchema: BaseIOSchema, OutputSchema: BaseIOSchema]: Attributes: client: Client for interacting with the language model. model (str): The model to use for generating responses. - history (ChatHistory): History component for storing chat history. + history (BaseChatHistory): History component for storing chat history. system_prompt_generator (BaseSystemPromptGenerator): Component for generating system prompts. system_role (Optional[str]): The role of the system in the conversation. None means no system prompt. assistant_role (str): The role of the assistant in the conversation. Use 'model' for Gemini, 'assistant' for OpenAI/Anthropic. @@ -189,7 +190,7 @@ def __init__(self, config: AgentConfig): """ self.client = config.client self.model = config.model - self.history = config.history or ChatHistory() + self.history = config.history if config.history is not None else ChatHistory() self.system_prompt_generator = config.system_prompt_generator or SystemPromptGenerator() self.system_role = config.system_role self.assistant_role = config.assistant_role diff --git a/atomic-agents/atomic_agents/context/__init__.py b/atomic-agents/atomic_agents/context/__init__.py index 487d54a6..47cb598b 100644 --- a/atomic-agents/atomic_agents/context/__init__.py +++ b/atomic-agents/atomic_agents/context/__init__.py @@ -1,3 +1,4 @@ +from .base_chat_history import BaseChatHistory from .chat_history import Message, ChatHistory from .system_prompt_generator import ( BaseDynamicContextProvider, @@ -7,6 +8,7 @@ __all__ = [ "Message", + "BaseChatHistory", "ChatHistory", "SystemPromptGenerator", "BaseDynamicContextProvider", diff --git a/atomic-agents/atomic_agents/context/base_chat_history.py b/atomic-agents/atomic_agents/context/base_chat_history.py new file mode 100644 index 00000000..359c5cfe --- /dev/null +++ b/atomic-agents/atomic_agents/context/base_chat_history.py @@ -0,0 +1,152 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Dict, List, Optional + +from atomic_agents.base.base_io_schema import BaseIOSchema + +if TYPE_CHECKING: + from atomic_agents.context.chat_history import Message + + +class BaseChatHistory(ABC): + """ + Declares the public contract of a chat history. + + This is the full method surface of the built-in ``ChatHistory``, of which ``AtomicAgent`` + itself uses a subset (``add_message``, ``get_history``, ``initialize_turn``, + ``delete_turn_id``, ``copy``, plus direct reads of the ``history`` / ``current_turn_id`` + attributes). The remaining methods (``dump`` / ``load`` / ``get_message_count`` / + ``get_current_turn_id``) round out the standard history surface that tooling and callers + rely on. + + This is an interface-only abstract base class: it defines no state and no behavior, + only the methods a chat history implementation must provide. ``ChatHistory`` is the + framework's built-in, in-memory implementation. Implement this interface directly (or + subclass ``ChatHistory``) to plug in a custom persistent or backend-specific memory store. + + Every implementation must maintain the following two attributes, since callers + (including ``AtomicAgent._trim_context``) read them directly: + + Attributes: + history (List[Message]): The in-memory working list of messages that make up the + chat history. ``AtomicAgent._trim_context`` reads this list directly (via + ``msg.turn_id``) to decide which turns to trim when a context-length limit + is exceeded. + current_turn_id (Optional[str]): The ID of the current turn, or ``None`` if no + turn has been started yet. Set by ``initialize_turn()``. + """ + + history: List["Message"] + current_turn_id: Optional[str] + + @abstractmethod + def initialize_turn(self) -> None: + """ + Starts a new turn, generating and storing a new turn ID. + + Implementations must set ``self.current_turn_id`` to a new, unique value each + time this is called. + """ + raise NotImplementedError + + @abstractmethod + def add_message(self, role: str, content: BaseIOSchema) -> None: + """ + Adds a message to the chat history, tagging it with the current turn ID. + + If no turn has been started yet (``current_turn_id`` is ``None``), implementations + should start one (e.g. by calling ``initialize_turn()``) before recording the message. + + Args: + role (str): The role of the message sender (e.g. 'user', 'system', 'assistant'). + content (BaseIOSchema): The content of the message. + """ + raise NotImplementedError + + @abstractmethod + def get_history(self) -> List[Dict]: + """ + Retrieves the chat history in the wire format expected by the LLM client. + + Returns: + List[Dict]: The list of messages as dictionaries with 'role' and 'content' keys, + suitable for passing to a chat completion call. + """ + raise NotImplementedError + + @abstractmethod + def get_current_turn_id(self) -> Optional[str]: + """ + Returns the current turn ID. + + Returns: + Optional[str]: The current turn ID, or ``None`` if no turn has been started. + """ + raise NotImplementedError + + @abstractmethod + def delete_turn_id(self, turn_id: str): + """ + Deletes all messages belonging to the given turn ID. + + Args: + turn_id (str): The turn ID whose messages should be removed. + + Raises: + ValueError: If the specified turn ID is not found in the history. + """ + raise NotImplementedError + + @abstractmethod + def get_message_count(self) -> int: + """ + Returns the number of messages currently in the chat history. + + Returns: + int: The number of messages. + """ + raise NotImplementedError + + @abstractmethod + def dump(self) -> str: + """ + Serializes the entire chat history to a JSON string. + + The returned string must be a valid argument to ``load()``, such that + ``instance.load(instance.dump())`` round-trips the full history state. + + Returns: + str: A JSON string representation of the chat history. + """ + raise NotImplementedError + + @abstractmethod + def load(self, serialized_data: str) -> None: + """ + Deserializes a JSON string produced by ``dump()`` and loads it into this instance, + replacing any existing history state. + + Args: + serialized_data (str): A JSON string representation of the chat history, + as produced by ``dump()``. + + Raises: + ValueError: If the serialized data is invalid or cannot be deserialized. + """ + raise NotImplementedError + + @abstractmethod + def copy(self) -> "BaseChatHistory": + """ + Creates a copy of the chat history, independent of the original. + + ``AtomicAgent`` calls ``copy()`` to snapshot ``initial_history`` at construction and + to restore it in ``reset_history()``. The built-in ``ChatHistory.copy()`` returns a + plain ``ChatHistory``; a subclass that carries extra state (a database handle, a + ``session_id``, a store reference) MUST override ``copy()`` to return its own type, + otherwise ``reset_history()`` silently replaces the backend with a plain in-memory + history and later writes stop reaching the store. + + Returns: + BaseChatHistory: A copy of the chat history. + """ + raise NotImplementedError diff --git a/atomic-agents/atomic_agents/context/chat_history.py b/atomic-agents/atomic_agents/context/chat_history.py index fb92513e..e4fbf91e 100644 --- a/atomic-agents/atomic_agents/context/chat_history.py +++ b/atomic-agents/atomic_agents/context/chat_history.py @@ -8,6 +8,7 @@ from pydantic import BaseModel, Field from atomic_agents.base.base_io_schema import BaseIOSchema +from atomic_agents.context.base_chat_history import BaseChatHistory INSTRUCTOR_MULTIMODAL_TYPES = (Image, Audio, PDF) @@ -28,7 +29,7 @@ class Message(BaseModel): turn_id: Optional[str] = None -class ChatHistory: +class ChatHistory(BaseChatHistory): """ Manages the chat history for an AI agent. diff --git a/atomic-agents/tests/context/test_base_chat_history.py b/atomic-agents/tests/context/test_base_chat_history.py new file mode 100644 index 00000000..358b8e1f --- /dev/null +++ b/atomic-agents/tests/context/test_base_chat_history.py @@ -0,0 +1,238 @@ +from unittest.mock import Mock + +import instructor +import pytest +from pydantic import Field + +from atomic_agents import ( + AgentConfig, + AtomicAgent, + BaseIOSchema, + BasicChatInputSchema, + BasicChatOutputSchema, +) +from atomic_agents.context import BaseChatHistory, ChatHistory, Message, SystemPromptGenerator + + +class InputSchema(BaseIOSchema): + """Test Input Schema""" + + test_field: str = Field(..., description="A test field") + + +def test_base_chat_history_cannot_be_instantiated_directly(): + with pytest.raises(TypeError): + BaseChatHistory() + + +def test_chat_history_is_subclass_and_instance_of_base(): + assert issubclass(ChatHistory, BaseChatHistory) + assert isinstance(ChatHistory(), BaseChatHistory) + + +def test_subclass_missing_abstractmethod_cannot_be_instantiated(): + class IncompleteHistory(BaseChatHistory): + """A subclass that forgets to implement `copy`.""" + + def initialize_turn(self) -> None: + pass + + def add_message(self, role, content): + pass + + def get_history(self): + return [] + + def get_current_turn_id(self): + return None + + def delete_turn_id(self, turn_id): + pass + + def get_message_count(self): + return 0 + + def dump(self): + return "{}" + + def load(self, serialized_data): + pass + + # `copy` intentionally omitted + + with pytest.raises(TypeError): + IncompleteHistory() + + +class RecordingChatHistory(ChatHistory): + """A custom persistent-memory-style backend used to exercise the extension seam.""" + + def __init__(self, max_messages=None): + super().__init__(max_messages=max_messages) + self.recorded_calls = 0 + + def add_message(self, role, content): + self.recorded_calls += 1 + super().add_message(role, content) + + +@pytest.fixture +def mock_instructor(): + mock = Mock(spec=instructor.Instructor) + mock.chat = Mock() + mock.chat.completions = Mock() + mock.chat.completions.create = Mock(return_value=BasicChatOutputSchema(chat_message="Test output")) + return mock + + +def test_custom_chat_history_subclass_used_by_agent(mock_instructor): + custom_history = RecordingChatHistory() + config = AgentConfig( + client=mock_instructor, + model="gpt-5-mini", + history=custom_history, + system_prompt_generator=SystemPromptGenerator(), + ) + agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) + + assert isinstance(agent.history, BaseChatHistory) + + result = agent.run(BasicChatInputSchema(chat_message="Hello")) + + assert result.chat_message == "Test output" + # One call for the user message, one for the assistant response. + assert custom_history.recorded_calls == 2 + assert custom_history.get_message_count() == 2 + + +class MinimalHistory(BaseChatHistory): + """A from-scratch BaseChatHistory implementation that does NOT subclass ChatHistory. + + Exercises the "own everything" path the docs describe, to prove a backend built purely + against the ABC drives AtomicAgent end to end. + """ + + def __init__(self, max_messages=None): + self.history = [] + self.current_turn_id = None + self.max_messages = max_messages + + def initialize_turn(self): + self.current_turn_id = "turn" + + def add_message(self, role, content): + if self.current_turn_id is None: + self.initialize_turn() + self.history.append(Message(role=role, content=content, turn_id=self.current_turn_id)) + + def get_history(self): + return [{"role": m.role, "content": m.content.model_dump_json()} for m in self.history] + + def get_current_turn_id(self): + return self.current_turn_id + + def delete_turn_id(self, turn_id): + self.history = [m for m in self.history if m.turn_id != turn_id] + + def get_message_count(self): + return len(self.history) + + def dump(self): + return "{}" + + def load(self, serialized_data): + pass + + def copy(self): + new = MinimalHistory(max_messages=self.max_messages) + new.history = list(self.history) + new.current_turn_id = self.current_turn_id + return new + + +def test_from_scratch_base_history_drives_agent(mock_instructor): + history = MinimalHistory() + config = AgentConfig( + client=mock_instructor, + model="gpt-5-mini", + history=history, + system_prompt_generator=SystemPromptGenerator(), + ) + agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) + + assert isinstance(agent.history, MinimalHistory) + + result = agent.run(BasicChatInputSchema(chat_message="Hello")) + + assert result.chat_message == "Test output" + assert agent.history.get_message_count() == 2 + + +class StatefulHistory(ChatHistory): + """A ChatHistory subclass carrying extra state, with copy() overridden to preserve it.""" + + def __init__(self, marker, max_messages=None): + super().__init__(max_messages=max_messages) + self.marker = marker + + def copy(self): + new = StatefulHistory(self.marker, max_messages=self.max_messages) + new.load(self.dump()) + new.current_turn_id = self.current_turn_id + return new + + +def test_reset_history_preserves_backend_when_copy_overridden(mock_instructor): + history = StatefulHistory(marker="db-handle") + config = AgentConfig( + client=mock_instructor, + model="gpt-5-mini", + history=history, + system_prompt_generator=SystemPromptGenerator(), + ) + agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) + + # The initial_history snapshot keeps the subclass type and its state. + assert isinstance(agent.initial_history, StatefulHistory) + assert agent.initial_history.marker == "db-handle" + + agent.run(BasicChatInputSchema(chat_message="Hi")) + agent.reset_history() + + # After reset, the backend is still the custom type, not a plain ChatHistory. + assert isinstance(agent.history, StatefulHistory) + assert agent.history.marker == "db-handle" + + +def test_subclass_without_copy_override_downgrades_to_chat_history(): + """Pins the known sharp edge: an un-overridden copy() returns a plain ChatHistory. + + This is exactly why subclasses carrying extra state must override copy(); documented in + the Memory guide and warned about in BaseChatHistory.copy's docstring. + """ + + class NaiveHistory(ChatHistory): + pass + + copied = NaiveHistory().copy() + assert type(copied) is ChatHistory + + +def test_falsy_custom_history_is_not_discarded(mock_instructor): + """AgentConfig must keep a provided history even if it is falsy (e.g. defines __len__).""" + + class FalsyWhenEmpty(ChatHistory): + def __len__(self): + return self.get_message_count() + + history = FalsyWhenEmpty() + assert not history # empty -> falsy via __len__ + config = AgentConfig( + client=mock_instructor, + model="gpt-5-mini", + history=history, + system_prompt_generator=SystemPromptGenerator(), + ) + agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](config) + + assert agent.history is history diff --git a/atomic-examples/persistent-memory/README.md b/atomic-examples/persistent-memory/README.md new file mode 100644 index 00000000..0dd30fde --- /dev/null +++ b/atomic-examples/persistent-memory/README.md @@ -0,0 +1,52 @@ +# Persistent Memory Example + +This example demonstrates how to plug a custom, persistent memory backend into Atomic Agents so a conversation survives across separate process runs (cross-session recall). + +It shows: + +- Subclassing `ChatHistory` to add persistence, following the pluggable `BaseChatHistory` extension point +- A `SQLiteChatHistory` backend built with **zero extra dependencies**, using only the Python standard library's `sqlite3` module +- Persisting on every `add_message` call and rehydrating the full history on startup via the inherited `dump()`/`load()` + +See the Memory guide's ["Writing a Custom Memory Backend"](../../docs/guides/memory.md#writing-a-custom-memory-backend) section for the full write-up of this pattern. + +## Getting Started + +1. Clone the main Atomic Agents repository: + + ```bash + git clone https://github.com/eigenwise/atomic-agents + ``` + +2. Navigate to this example's directory: + + ```bash + cd atomic-agents/atomic-examples/persistent-memory + ``` + +3. Install the dependencies using uv: + + ```bash + uv sync + ``` + +4. Set your `OPENAI_API_KEY` environment variable (or put it in a `.env` file in this directory). + +## Running the Example + +```bash +uv run python persistent_memory/main.py +``` + +Chat with the agent, then exit with `/exit`. A `chat_history.db` file is created in the current directory. + +**Run it a second time** and the agent picks up right where you left off: it prints the number of restored messages at startup and remembers everything from the previous run, because the conversation was persisted to `chat_history.db` under the `demo` session ID. + +## How It Works + +`SQLiteChatHistory` (in `persistent_memory/main.py`) subclasses `ChatHistory` instead of reimplementing the memory contract from scratch: + +- `__init__` opens (or creates) a SQLite database with a single `sessions(session_id TEXT PRIMARY KEY, data TEXT)` table, and loads any existing row for the given `session_id` via the inherited `load()`. +- `add_message` calls `super().add_message(...)` to reuse `ChatHistory`'s turn handling, then saves the entire history to SQLite via `self.dump()`. + +Because `dump()`/`load()` already handle serialization, the persistence layer only has to move a single JSON string in and out of storage. The same pattern works for any store (Redis, Postgres, a hosted memory service) by swapping out what `_save`/`__init__` talk to. diff --git a/atomic-examples/persistent-memory/persistent_memory/__init__.py b/atomic-examples/persistent-memory/persistent_memory/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/atomic-examples/persistent-memory/persistent_memory/main.py b/atomic-examples/persistent-memory/persistent_memory/main.py new file mode 100644 index 00000000..648b57f9 --- /dev/null +++ b/atomic-examples/persistent-memory/persistent_memory/main.py @@ -0,0 +1,137 @@ +""" +Persistent Memory Example +========================== + +Demonstrates a custom, persistent ``ChatHistory`` backend built with nothing but the +Python standard library's ``sqlite3`` module (zero extra dependencies). It follows the +recommended extension pattern from the Memory guide's "Writing a Custom Memory Backend" +section: subclass ``ChatHistory`` to inherit turn handling and serialization, and +override ``add_message`` to persist on every write. + +Run this script twice: the second run rehydrates the conversation from the first. +""" + +import contextlib +import os +import sqlite3 +from typing import Optional + +import instructor +import openai +from dotenv import load_dotenv +from rich.console import Console +from rich.text import Text + +from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema, BasicChatInputSchema, BasicChatOutputSchema +from atomic_agents.context import ChatHistory + +load_dotenv() + +# API Key setup +API_KEY = "" +if not API_KEY: + API_KEY = os.getenv("OPENAI_API_KEY") + +if not API_KEY: + raise ValueError( + "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY." + ) + + +class SQLiteChatHistory(ChatHistory): + """A ChatHistory that persists to a local SQLite database, keyed by session_id. + + Follows the recommended pattern from the Memory guide: subclass ChatHistory to inherit + turn handling and serialization, then override add_message to persist. copy() is also + overridden so the agent's initial_history snapshot and reset_history() keep this + persistent backend instead of silently degrading to a plain in-memory ChatHistory. + + A short-lived connection is opened per operation (via contextlib.closing) rather than + held open for the object's lifetime, so no handle is leaked and the backend is safe to + use from more than one thread. + """ + + def __init__(self, session_id: str, db_path: str = "chat_history.db", max_messages: Optional[int] = None): + super().__init__(max_messages=max_messages) + self.session_id = session_id + self.db_path = db_path + + with contextlib.closing(sqlite3.connect(self.db_path)) as conn: + conn.execute("CREATE TABLE IF NOT EXISTS sessions (session_id TEXT PRIMARY KEY, data TEXT)") + conn.commit() + row = conn.execute("SELECT data FROM sessions WHERE session_id = ?", (self.session_id,)).fetchone() + + if row is not None: + self.load(row[0]) + + def add_message(self, role: str, content: BaseIOSchema) -> None: + super().add_message(role, content) + self._save() + + def _save(self) -> None: + with contextlib.closing(sqlite3.connect(self.db_path)) as conn: + conn.execute( + "INSERT OR REPLACE INTO sessions (session_id, data) VALUES (?, ?)", + (self.session_id, self.dump()), + ) + conn.commit() + + def copy(self) -> "SQLiteChatHistory": + new_history = SQLiteChatHistory(self.session_id, db_path=self.db_path, max_messages=self.max_messages) + new_history.load(self.dump()) + new_history.current_turn_id = self.current_turn_id + return new_history + + +def main() -> None: + # Initialize a Rich Console for pretty console outputs + console = Console() + + # History setup - restores prior turns for this session_id if chat_history.db exists + history = SQLiteChatHistory(session_id="demo") + console.print(f"[bold yellow]Restored {history.get_message_count()} message(s) from {history.db_path}[/bold yellow]") + + if history.get_message_count() == 0: + initial_message = BasicChatOutputSchema(chat_message="Hello! How can I assist you today?") + history.add_message("assistant", initial_message) + + # OpenAI client setup using the Instructor library + client = instructor.from_openai(openai.OpenAI(api_key=API_KEY)) + + # Agent setup with specified configuration + agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( + config=AgentConfig( + client=client, + model="gpt-5-mini", + model_api_parameters={"reasoning_effort": "low"}, + history=history, + ) + ) + + # Display the last message so the user sees where the conversation left off + last_message = history.history[-1].content + console.print(Text("Agent:", style="bold green"), end=" ") + console.print(Text(last_message.chat_message, style="bold green")) + + # Start an infinite loop to handle user inputs and agent responses + while True: + # Prompt the user for input with a styled prompt + user_input = console.input("[bold blue]You:[/bold blue] ") + # Check if the user wants to exit the chat + if user_input.lower() in ["/exit", "/quit"]: + console.print( + f"Exiting chat... your conversation is saved in {history.db_path}. Run again to pick up where you left off." + ) + break + + # Process the user's input through the agent and get the response + input_schema = BasicChatInputSchema(chat_message=user_input) + response = agent.run(input_schema) + + agent_message = Text(response.chat_message, style="bold green") + console.print(Text("Agent:", style="bold green"), end=" ") + console.print(agent_message) + + +if __name__ == "__main__": + main() diff --git a/atomic-examples/persistent-memory/pyproject.toml b/atomic-examples/persistent-memory/pyproject.toml new file mode 100644 index 00000000..9de27ee5 --- /dev/null +++ b/atomic-examples/persistent-memory/pyproject.toml @@ -0,0 +1,25 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["persistent_memory"] + +[project] +name = "persistent-memory" +version = "1.0.0" +description = "Persistent memory backend example for Atomic Agents" +readme = "README.md" +authors = [ + { name = "Kenny Vaneetvelde", email = "kenny.vaneetvelde@gmail.com" } +] +requires-python = ">=3.12" +dependencies = [ + "atomic-agents", + "instructor[anthropic,groq,google-genai]==1.14.5", + "openai>=2.0.0,<3.0.0", + "python-dotenv>=1.0.1,<2.0.0", +] + +[tool.uv.sources] +atomic-agents = { workspace = true } diff --git a/docs/api/context.md b/docs/api/context.md index fa32400c..cd12cc0b 100644 --- a/docs/api/context.md +++ b/docs/api/context.md @@ -6,7 +6,7 @@ For a comprehensive guide on memory management, multi-agent patterns, and best p ## Agent History -The `ChatHistory` class manages conversation history and state for AI agents: +The `ChatHistory` class manages conversation history and state for AI agents. It implements `BaseChatHistory`, an interface-only abstract base class that declares the memory contract `AtomicAgent` depends on (the type `AgentConfig.history` accepts). Implement `BaseChatHistory` directly, or subclass `ChatHistory`, to plug in your own persistent or backend-specific memory store. See the [Memory guide's "Writing a Custom Memory Backend" section](/guides/memory#writing-a-custom-memory-backend) for the full contract and a recommended pattern. ```python from atomic_agents.context import ChatHistory @@ -260,6 +260,11 @@ For full API details: :undoc-members: :show-inheritance: +.. automodule:: atomic_agents.context.base_chat_history + :members: + :undoc-members: + :show-inheritance: + .. automodule:: atomic_agents.context.system_prompt_generator :members: :undoc-members: diff --git a/docs/guides/memory.md b/docs/guides/memory.md index 0c2e613e..d0bf563e 100644 --- a/docs/guides/memory.md +++ b/docs/guides/memory.md @@ -402,6 +402,79 @@ history = ChatHistory() # Create new instance --- +## Writing a Custom Memory Backend + +Memory is intentionally kept outside the framework core. `ChatHistory` is a solid in-memory default, but if you want conversations to survive a restart, live in Redis, Postgres, or some hosted memory service, that integration belongs in **your** codebase, not as a framework dependency. `BaseChatHistory` is the stable seam that makes this possible: build against it and `AtomicAgent` doesn't care what's underneath. + +### The Contract + +`BaseChatHistory` is an interface-only abstract base class (`ABC`). It declares no state and no behavior of its own, just the methods a history implementation must provide: + +- `initialize_turn(self) -> None` +- `add_message(self, role: str, content: BaseIOSchema) -> None` +- `get_history(self) -> List[Dict]` +- `get_current_turn_id(self) -> Optional[str]` +- `delete_turn_id(self, turn_id: str)` +- `get_message_count(self) -> int` +- `dump(self) -> str` +- `load(self, serialized_data: str) -> None` +- `copy(self) -> "BaseChatHistory"` + +Every implementation also has to maintain two attributes, because `AtomicAgent` reads them directly: + +- `history: List[Message]` - the in-memory working list of messages. `AtomicAgent._trim_context` reads this list (via each message's `turn_id`) to decide what to trim when a context-length limit is hit. +- `current_turn_id: Optional[str]` - the ID of the current turn, or `None` if no turn has started yet. Set by `initialize_turn()`. + +### Recommended Pattern: Subclass ChatHistory + +Don't implement `BaseChatHistory` from scratch unless you have a good reason to. Subclassing `ChatHistory` gets you turn handling and serialization for free, you only need to override the couple of methods that touch your storage layer: + +```python +from atomic_agents.context import ChatHistory + +class PersistentHistory(ChatHistory): + def __init__(self, session_id, store, max_messages=None): + super().__init__(max_messages=max_messages) + self.session_id = session_id + self.store = store + saved = store.get(session_id) + if saved: + self.load(saved) + + def add_message(self, role, content): + super().add_message(role, content) + self.store.put(self.session_id, self.dump()) + + def copy(self): + # AtomicAgent calls copy() for its initial_history snapshot and reset_history(). + # Without this override, copy() returns a plain ChatHistory and your backend is + # silently dropped after reset_history() (later writes stop reaching the store). + new_history = PersistentHistory(self.session_id, self.store, max_messages=self.max_messages) + new_history.load(self.dump()) + new_history.current_turn_id = self.current_turn_id + return new_history +``` + +The core of the pattern is small: load on init if there's saved state, persist on every `add_message`. `dump()`/`load()` already round-trip the full history, so there's no serialization logic to write yourself. + +```{important} +Override `copy()` too. `AtomicAgent` snapshots `initial_history` with `copy()` at construction and restores it in `reset_history()`. The built-in `ChatHistory.copy()` hard-codes a plain `ChatHistory`, so a subclass that carries extra state (a store handle, a `session_id`) that forgets to override `copy()` will be **silently replaced by a non-persistent in-memory history** the first time `reset_history()` runs. Any subclass with its own state must override `copy()` to return its own type. +``` + +```{note} +This pattern re-serializes and rewrites the **entire** history on every message (`self.dump()`), which is fine for a demo but O(n²) over a long conversation. A production backend against Redis/Postgres/etc. should persist incrementally (append the new message) rather than dumping the whole history each write. +``` + +### From-Scratch Alternative + +If you don't want the in-memory list at all (say your backend is a query, not a cache), implement `BaseChatHistory` directly. You give up the free turn tracking and `get_history()` serialization that `ChatHistory` provides, so you own all of it: generating turn IDs, building the wire-format list of dicts, and keeping `history` / `current_turn_id` in sync with whatever `AtomicAgent` expects to read. + +### Try It + +There's a full runnable example of this pattern (a dependency-free SQLite-backed history that survives across process runs) at [`atomic-examples/persistent-memory`](https://github.com/eigenwise/atomic-agents/tree/main/atomic-examples/persistent-memory). + +--- + ## Multimodal Content in History ChatHistory supports images, PDFs, and audio through Instructor's multimodal types. @@ -1175,6 +1248,24 @@ if history.get_message_count() > 40: | `load(data)` | Deserialize from JSON string | | `copy()` | Create deep copy | +### BaseChatHistory + +The abstract base class `ChatHistory` implements. Subclass it (directly, or via `ChatHistory`) to plug in a custom memory backend. See [Writing a Custom Memory Backend](#writing-a-custom-memory-backend). + +| Method | Description | +|--------|-------------| +| `initialize_turn()` | Start new turn with new UUID | +| `add_message(role, content)` | Add message to current turn | +| `get_history()` | Get all messages as list of dicts | +| `get_current_turn_id()` | Get current turn's UUID | +| `delete_turn_id(turn_id)` | Delete all messages in a turn | +| `get_message_count()` | Get number of messages | +| `dump()` | Serialize to JSON string | +| `load(data)` | Deserialize from JSON string | +| `copy()` | Create deep copy | + +Every implementation must also maintain a `history: List[Message]` attribute and a `current_turn_id: Optional[str]` attribute. + ### Message | Field | Type | Description | diff --git a/uv.lock b/uv.lock index d2b4b5d9..d0180678 100644 --- a/uv.lock +++ b/uv.lock @@ -26,6 +26,7 @@ members = [ "nested-multimodal", "orchestration-agent", "pdf-reader", + "persistent-memory", "progressive-disclosure", "quickstart", "rag-chatbot", @@ -2921,6 +2922,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/df/98/629f269c2bd91bdcac147aad5cf51ceb645c0196e23a41ee3c051125190f/pdoc3-0.11.6-py3-none-any.whl", hash = "sha256:8b72723767bd48d899812d2aec8375fc1c3476e179455db0b4575e6dccb44b93", size = 255188, upload-time = "2025-03-20T22:53:51.671Z" }, ] +[[package]] +name = "persistent-memory" +version = "1.0.0" +source = { editable = "atomic-examples/persistent-memory" } +dependencies = [ + { name = "atomic-agents" }, + { name = "instructor", extra = ["anthropic", "google-genai", "groq"] }, + { name = "openai" }, + { name = "python-dotenv" }, +] + +[package.metadata] +requires-dist = [ + { name = "atomic-agents", editable = "." }, + { name = "instructor", extras = ["anthropic", "groq", "google-genai"], specifier = "==1.14.5" }, + { name = "openai", specifier = ">=2.0.0,<3.0.0" }, + { name = "python-dotenv", specifier = ">=1.0.1,<2.0.0" }, +] + [[package]] name = "pillow" version = "12.0.0"