diff --git a/src/dynavec/__init__.py b/src/dynavec/__init__.py index b384daf..72d47e1 100644 --- a/src/dynavec/__init__.py +++ b/src/dynavec/__init__.py @@ -75,6 +75,15 @@ run_context, set_default_run_logger, ) +from .memory import ( + BaseMemory, + DynamoDBMemoryStore, + InMemoryMemoryStore, + MemoryPolicy, + SummaryMemoryPolicy, + TokenBudgetMemoryPolicy, + WindowMemoryPolicy, +) from .models import ( Document, ExplainedSearchResult, @@ -189,6 +198,14 @@ "BaseCheckpointer", "MemoryCheckpointer", "DynamoDBCheckpointer", + # Conversation Memory (dynaflow) + "BaseMemory", + "InMemoryMemoryStore", + "DynamoDBMemoryStore", + "MemoryPolicy", + "WindowMemoryPolicy", + "TokenBudgetMemoryPolicy", + "SummaryMemoryPolicy", # Ingestion "Record", "ingest", diff --git a/src/dynavec/agents/__init__.py b/src/dynavec/agents/__init__.py index 262cfa9..9ec7aef 100644 --- a/src/dynavec/agents/__init__.py +++ b/src/dynavec/agents/__init__.py @@ -8,6 +8,15 @@ DynamoDBCheckpointer, MemoryCheckpointer, ) +from ..memory import ( + BaseMemory, + DynamoDBMemoryStore, + InMemoryMemoryStore, + MemoryPolicy, + SummaryMemoryPolicy, + TokenBudgetMemoryPolicy, + WindowMemoryPolicy, +) from .base import ( AgentResult, AgentStep, @@ -45,4 +54,11 @@ "BaseCheckpointer", "MemoryCheckpointer", "DynamoDBCheckpointer", + "BaseMemory", + "InMemoryMemoryStore", + "DynamoDBMemoryStore", + "MemoryPolicy", + "WindowMemoryPolicy", + "TokenBudgetMemoryPolicy", + "SummaryMemoryPolicy", ] diff --git a/src/dynavec/agents/react.py b/src/dynavec/agents/react.py index 3bd9546..e02fcc6 100644 --- a/src/dynavec/agents/react.py +++ b/src/dynavec/agents/react.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from ..checkpoint import BaseCheckpointer + from ..memory import BaseMemory DEFAULT_REACT_SYSTEM_PROMPT = """You are a helpful and precise reasoning agent. You solve tasks step-by-step using a ReAct (Reason + Act) approach. @@ -76,6 +77,7 @@ def __init__( system_prompt: str | None = None, max_steps: int = 10, checkpointer: BaseCheckpointer | None = None, + memory: BaseMemory | None = None, ) -> None: self.model = model self.system_prompt = ( @@ -83,6 +85,7 @@ def __init__( ) self.max_steps = max(1, max_steps) self.checkpointer = checkpointer + self.memory = memory self._tools_map: dict[str, AgentTool] = {} self._chat_tools: list[Tool] = [] @@ -93,9 +96,7 @@ def __init__( self._chat_tools.append(t.to_chat_tool()) elif tools: for raw_tool in tools: - agent_tool = ( - raw_tool if isinstance(raw_tool, AgentTool) else AgentTool(raw_tool) - ) + agent_tool = raw_tool if isinstance(raw_tool, AgentTool) else AgentTool(raw_tool) self._tools_map[agent_tool.name] = agent_tool self._chat_tools.append(agent_tool.to_chat_tool()) @@ -118,9 +119,7 @@ async def _adispatch_tool_call(self, tc: ToolCall) -> str: return await tool_instance.aexecute(tc.arguments) return f"Error: Tool {tc.name!r} is not registered in available tools." - def _process_tool_calls( - self, tool_calls: list[ToolCall], messages: list[Message] - ) -> list[str]: + def _process_tool_calls(self, tool_calls: list[ToolCall], messages: list[Message]) -> list[str]: step_observations: list[str] = [] for tc in tool_calls: obs = self._dispatch_tool_call(tc) @@ -150,12 +149,9 @@ async def _aprocess_tool_calls( ) return step_observations - def _finalize_result( - self, steps: list[AgentStep], total_tool_calls: int - ) -> AgentResult: + def _finalize_result(self, steps: list[AgentStep], total_tool_calls: int) -> AgentResult: last_output = ( - steps[-1].thought - or (steps[-1].observations[-1] if steps[-1].observations else "") + steps[-1].thought or (steps[-1].observations[-1] if steps[-1].observations else "") if steps else "" ) @@ -194,7 +190,21 @@ def run( start_step = existing_cp.step + 1 if not messages: - messages = self._init_messages(goal) + if self.memory is not None and thread_id is not None: + existing = self.memory.get_messages(thread_id) + if not existing and self.system_prompt: + self.memory.append( + thread_id, + Message(role="system", content=self.system_prompt), + ) + self.memory.append( + thread_id, + Message(role="user", content=goal), + ) + messages = list(self.memory.get_messages(thread_id)) + else: + messages = self._init_messages(goal) + if self.checkpointer is not None and thread_id is not None: self.checkpointer.put( thread_id, @@ -225,6 +235,12 @@ def run( observations=[], ) steps.append(step) + out = msg.content or "" + if self.memory is not None and thread_id is not None: + self.memory.append( + thread_id, + Message(role="assistant", content=out), + ) if self.checkpointer is not None and thread_id is not None: self.checkpointer.put( thread_id, @@ -233,14 +249,14 @@ def run( "messages": [_message_to_dict(m) for m in messages], "steps": [_step_to_dict(s) for s in steps], "total_tool_calls": total_tool_calls, - "output": msg.content or "", + "output": out, "finished": True, }, node_id=f"step_{step_idx}", step=step_idx, ) return AgentResult( - output=msg.content or "", + output=out, steps=steps, finished=True, termination_reason="completed", @@ -274,6 +290,11 @@ def run( ) res = self._finalize_result(steps, total_tool_calls) + if self.memory is not None and thread_id is not None: + self.memory.append( + thread_id, + Message(role="assistant", content=res.output), + ) if self.checkpointer is not None and thread_id is not None: self.checkpointer.put( thread_id, @@ -316,7 +337,20 @@ async def arun( start_step = existing_cp.step + 1 if not messages: - messages = self._init_messages(goal) + if self.memory is not None and thread_id is not None: + existing = self.memory.get_messages(thread_id) + if not existing and self.system_prompt: + self.memory.append( + thread_id, + Message(role="system", content=self.system_prompt), + ) + self.memory.append( + thread_id, + Message(role="user", content=goal), + ) + messages = list(self.memory.get_messages(thread_id)) + else: + messages = self._init_messages(goal) if self.checkpointer is not None and thread_id is not None: self.checkpointer.put( thread_id, @@ -347,6 +381,12 @@ async def arun( observations=[], ) steps.append(step) + out = msg.content or "" + if self.memory is not None and thread_id is not None: + self.memory.append( + thread_id, + Message(role="assistant", content=out), + ) if self.checkpointer is not None and thread_id is not None: self.checkpointer.put( thread_id, @@ -355,14 +395,14 @@ async def arun( "messages": [_message_to_dict(m) for m in messages], "steps": [_step_to_dict(s) for s in steps], "total_tool_calls": total_tool_calls, - "output": msg.content or "", + "output": out, "finished": True, }, node_id=f"step_{step_idx}", step=step_idx, ) return AgentResult( - output=msg.content or "", + output=out, steps=steps, finished=True, termination_reason="completed", @@ -371,9 +411,7 @@ async def arun( ) total_tool_calls += len(msg.tool_calls) - step_observations = await self._aprocess_tool_calls( - list(msg.tool_calls), messages - ) + step_observations = await self._aprocess_tool_calls(list(msg.tool_calls), messages) step = AgentStep( step_number=step_idx, @@ -398,6 +436,11 @@ async def arun( ) res = self._finalize_result(steps, total_tool_calls) + if self.memory is not None and thread_id is not None: + self.memory.append( + thread_id, + Message(role="assistant", content=res.output), + ) if self.checkpointer is not None and thread_id is not None: self.checkpointer.put( thread_id, diff --git a/src/dynavec/memory.py b/src/dynavec/memory.py new file mode 100644 index 0000000..e02b288 --- /dev/null +++ b/src/dynavec/memory.py @@ -0,0 +1,422 @@ +"""Conversation and short-term memory store on DynamoDB (dynaflow). + +Complements dynavec's long-term semantic vector memory with short-term +conversation history and scratchpad persisted on DynamoDB, supporting +sliding window, token budget, and summary-on-overflow memory policies. +""" + +from __future__ import annotations + +import threading +from abc import ABC, abstractmethod +from collections.abc import Callable, Sequence +from datetime import datetime, timezone +from typing import Any + +from .chat.base import Message, Role, ToolCall +from .stores.dynamodb import _from_dynamo, _to_dynamo + + +def message_to_dict(msg: Message) -> dict[str, Any]: + """Serialize a Message to a dictionary.""" + return { + "role": msg.role, + "content": msg.content, + "tool_call_id": msg.tool_call_id, + "tool_calls": [ + {"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in msg.tool_calls + ], + } + + +def dict_to_message(d: dict[str, Any]) -> Message: + """Deserialize a dictionary to a Message.""" + tcs = [ + ToolCall(id=t["id"], name=t["name"], arguments=t["arguments"]) + for t in d.get("tool_calls", []) + ] + return Message( + role=d["role"], + content=d.get("content"), + tool_calls=tcs, + tool_call_id=d.get("tool_call_id"), + ) + + +# --- Memory Policies --- + + +class MemoryPolicy(ABC): + """Abstract policy governing message retention, trimming, and compression.""" + + @abstractmethod + def apply(self, messages: list[Message]) -> list[Message]: + """Apply policy to a list of messages, returning the retained messages.""" + + +class WindowMemoryPolicy(MemoryPolicy): + """Retains the most recent k non-system messages, optionally preserving system prompt.""" + + def __init__(self, k: int = 10, *, preserve_system: bool = True) -> None: + self.k = max(1, k) + self.preserve_system = preserve_system + + def apply(self, messages: list[Message]) -> list[Message]: + if not messages or len(messages) <= self.k: + return list(messages) + + system_msg: Message | None = None + other_messages: list[Message] = [] + + for m in messages: + if m.role == "system" and system_msg is None and self.preserve_system: + system_msg = m + else: + other_messages.append(m) + + trimmed = other_messages[-self.k :] + if system_msg is not None: + return [system_msg, *trimmed] + return trimmed + + +def default_token_counter(msg: Message) -> int: + """Approximate token count for a Message (1 token ~= 4 chars).""" + text = msg.content or "" + if msg.tool_calls: + for tc in msg.tool_calls: + text += f" {tc.name} {tc.arguments}" + return max(1, len(text) // 4) + + +class TokenBudgetMemoryPolicy(MemoryPolicy): + """Trims oldest non-system messages to stay within a maximum token budget.""" + + def __init__( + self, + max_tokens: int = 2048, + *, + token_counter: Callable[[Message], int] | None = None, + preserve_system: bool = True, + ) -> None: + self.max_tokens = max(1, max_tokens) + self.counter = token_counter or default_token_counter + self.preserve_system = preserve_system + + def apply(self, messages: list[Message]) -> list[Message]: + if not messages: + return [] + + system_msg: Message | None = None + remaining: list[Message] = [] + + for m in messages: + if m.role == "system" and system_msg is None and self.preserve_system: + system_msg = m + else: + remaining.append(m) + + sys_cost = self.counter(system_msg) if system_msg is not None else 0 + budget = max(0, self.max_tokens - sys_cost) + + retained: list[Message] = [] + current_tokens = 0 + + # Traverse backwards from newest to oldest + for m in reversed(remaining): + cost = self.counter(m) + if current_tokens + cost <= budget or not retained: + retained.append(m) + current_tokens += cost + else: + break + + retained.reverse() + if system_msg is not None: + return [system_msg, *retained] + return retained + + +class SummaryMemoryPolicy(MemoryPolicy): + """Compacts older messages into a summary message when exceeding max_messages.""" + + def __init__( + self, + max_messages: int = 6, + *, + summarizer: Callable[[list[Message]], str] | None = None, + preserve_system: bool = True, + ) -> None: + self.max_messages = max(2, max_messages) + self.summarizer = summarizer or self._default_summary + self.preserve_system = preserve_system + + def _default_summary(self, messages: list[Message]) -> str: + points = [] + for m in messages: + content = (m.content or "").strip() + if len(content) > 120: + content = content[:117] + "..." + points.append(f"- {m.role.capitalize()}: {content}") + return "\n".join(points) + + def apply(self, messages: list[Message]) -> list[Message]: + if len(messages) <= self.max_messages: + return list(messages) + + system_msg: Message | None = None + others: list[Message] = [] + + for m in messages: + if m.role == "system" and system_msg is None and self.preserve_system: + system_msg = m + else: + others.append(m) + + cutoff = len(others) - (self.max_messages // 2) + to_summarize = others[:cutoff] + recent = others[cutoff:] + + summary_text = f"Summary of prior conversation:\n{self.summarizer(to_summarize)}" + summary_msg = Message(role="system", content=summary_text) + + result: list[Message] = [] + if system_msg is not None: + result.append(system_msg) + result.append(summary_msg) + result.extend(recent) + return result + + +# --- Memory Stores --- + + +class BaseMemory(ABC): + """Abstract interface for managing agent short-term message memory.""" + + @abstractmethod + def append( + self, + thread_id: str, + messages: Message | Sequence[Message] | str, + *, + role: Role = "user", + ) -> list[Message]: + """Append one or more messages to the thread history.""" + + @abstractmethod + def get_messages(self, thread_id: str) -> list[Message]: + """Retrieve the current message history for a thread.""" + + @abstractmethod + def clear(self, thread_id: str) -> None: + """Clear message history and scratchpad for a thread.""" + + @abstractmethod + def trim(self, thread_id: str) -> list[Message]: + """Apply memory policy and save trimmed history.""" + + @abstractmethod + def get_scratchpad(self, thread_id: str) -> dict[str, Any]: + """Retrieve the thread scratchpad dictionary.""" + + @abstractmethod + def set_scratchpad(self, thread_id: str, scratchpad: dict[str, Any]) -> None: + """Set the thread scratchpad dictionary.""" + + +class InMemoryMemoryStore(BaseMemory): + """Thread-safe, in-memory conversation memory store for tests and local agents.""" + + def __init__(self, policy: MemoryPolicy | None = None) -> None: + self.policy = policy + self._messages: dict[str, list[Message]] = {} + self._scratchpads: dict[str, dict[str, Any]] = {} + self._lock = threading.Lock() + + def _normalize( + self, + messages: Message | Sequence[Message] | str, + role: Role, + ) -> list[Message]: + if isinstance(messages, str): + return [Message(role=role, content=messages)] + if isinstance(messages, Message): + return [messages] + return list(messages) + + def append( + self, + thread_id: str, + messages: Message | Sequence[Message] | str, + *, + role: Role = "user", + ) -> list[Message]: + new_msgs = self._normalize(messages, role) + with self._lock: + history = self._messages.setdefault(thread_id, []) + history.extend(new_msgs) + if self.policy is not None: + history = self.policy.apply(history) + self._messages[thread_id] = history + return list(history) + + def get_messages(self, thread_id: str) -> list[Message]: + with self._lock: + history = self._messages.get(thread_id, []) + if self.policy is not None: + return self.policy.apply(list(history)) + return list(history) + + def clear(self, thread_id: str) -> None: + with self._lock: + self._messages.pop(thread_id, None) + self._scratchpads.pop(thread_id, None) + + def trim(self, thread_id: str) -> list[Message]: + with self._lock: + history = self._messages.get(thread_id, []) + if self.policy is not None: + history = self.policy.apply(history) + self._messages[thread_id] = history + return list(history) + + def get_scratchpad(self, thread_id: str) -> dict[str, Any]: + with self._lock: + return dict(self._scratchpads.get(thread_id, {})) + + def set_scratchpad(self, thread_id: str, scratchpad: dict[str, Any]) -> None: + with self._lock: + self._scratchpads[thread_id] = dict(scratchpad) + + +class DynamoDBMemoryStore(BaseMemory): + """DynamoDB-backed conversation and scratchpad store. + + Stores conversation state per thread using single-table schema: + - pk: "{namespace}#memory#{thread_id}" + - messages: serialized message history + - scratchpad: agent working memory dict + - updated_at: ISO UTC timestamp + """ + + def __init__( + self, + table: Any, + *, + namespace: str = "default", + policy: MemoryPolicy | None = None, + ) -> None: + self.table = table + self.namespace = namespace + self.policy = policy + + def _pk(self, thread_id: str) -> str: + return f"{self.namespace}#memory#{thread_id}" + + def _normalize( + self, + messages: Message | Sequence[Message] | str, + role: Role, + ) -> list[Message]: + if isinstance(messages, str): + return [Message(role=role, content=messages)] + if isinstance(messages, Message): + return [messages] + return list(messages) + + def get_messages(self, thread_id: str) -> list[Message]: + """Fetch conversation message history for thread_id.""" + resp = self.table.get_item(Key={"pk": self._pk(thread_id)}) + item = resp.get("Item") + if not item or "messages" not in item: + return [] + + clean_data = _from_dynamo(item["messages"]) + messages = [dict_to_message(m) for m in clean_data] + if self.policy is not None: + return self.policy.apply(messages) + return messages + + def append( + self, + thread_id: str, + messages: Message | Sequence[Message] | str, + *, + role: Role = "user", + ) -> list[Message]: + """Append messages to thread history and persist to DynamoDB.""" + to_add = self._normalize(messages, role) + current = self.get_messages(thread_id) + current.extend(to_add) + + if self.policy is not None: + current = self.policy.apply(current) + + serialized = [message_to_dict(m) for m in current] + pk = self._pk(thread_id) + + item: dict[str, Any] = { + "pk": pk, + "ns": self.namespace, + "type": "conversation_memory", + "thread_id": thread_id, + "messages": serialized, + "updated_at": datetime.now(timezone.utc).isoformat(), + } + + # Preserve existing scratchpad if already written + try: + resp = self.table.get_item(Key={"pk": pk}) + existing = resp.get("Item") + if existing and "scratchpad" in existing: + item["scratchpad"] = existing["scratchpad"] + except Exception: + pass + + self.table.put_item(Item=_to_dynamo(item)) + return current + + def trim(self, thread_id: str) -> list[Message]: + """Apply memory policy and save trimmed history.""" + messages = self.get_messages(thread_id) + if self.policy is not None: + messages = self.policy.apply(messages) + + serialized = [message_to_dict(m) for m in messages] + pk = self._pk(thread_id) + + self.table.update_item( + Key={"pk": pk}, + UpdateExpression="SET #m = :m, updated_at = :ts", + ExpressionAttributeNames={"#m": "messages"}, + ExpressionAttributeValues={ + ":m": _to_dynamo(serialized), + ":ts": datetime.now(timezone.utc).isoformat(), + }, + ) + return messages + + def clear(self, thread_id: str) -> None: + """Clear memory for a thread.""" + self.table.delete_item(Key={"pk": self._pk(thread_id)}) + + def get_scratchpad(self, thread_id: str) -> dict[str, Any]: + """Retrieve scratchpad data for thread_id.""" + resp = self.table.get_item(Key={"pk": self._pk(thread_id)}) + item = resp.get("Item") + if not item or "scratchpad" not in item: + return {} + return dict(_from_dynamo(item["scratchpad"])) + + def set_scratchpad(self, thread_id: str, scratchpad: dict[str, Any]) -> None: + """Update scratchpad data for thread_id.""" + pk = self._pk(thread_id) + self.table.update_item( + Key={"pk": pk}, + UpdateExpression="SET scratchpad = :s, updated_at = :ts", + ExpressionAttributeValues={ + ":s": _to_dynamo(scratchpad), + ":ts": datetime.now(timezone.utc).isoformat(), + }, + ) diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 0000000..4bf3a05 --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,286 @@ +"""Unit tests for conversation and short-term memory store on DynamoDB (Issue #290).""" + +from __future__ import annotations + +import boto3 +from moto import mock_aws + +from dynavec.agents import ReActAgent +from dynavec.chat.base import ChatModel, ChatResult, Message, Tool, ToolCall +from dynavec.memory import ( + BaseMemory, + DynamoDBMemoryStore, + InMemoryMemoryStore, + SummaryMemoryPolicy, + TokenBudgetMemoryPolicy, + WindowMemoryPolicy, + dict_to_message, + message_to_dict, +) + + +class ScriptedChatModel(ChatModel): + """Deterministic test double for ChatModel.""" + + def __init__(self, responses: list[ChatResult | Message | str]) -> None: + self.responses = list(responses) + self.call_count = 0 + self.received_messages: list[list[Message]] = [] + + def invoke( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: object, + ) -> ChatResult: + self.received_messages.append(list(messages)) + if self.call_count >= len(self.responses): + return ChatResult( + message=Message(role="assistant", content="Done."), + finish_reason="stop", + ) + resp = self.responses[self.call_count] + self.call_count += 1 + if isinstance(resp, ChatResult): + return resp + if isinstance(resp, Message): + return ChatResult(message=resp, finish_reason="stop") + return ChatResult( + message=Message(role="assistant", content=resp), + finish_reason="stop", + ) + + async def ainvoke( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: object, + ) -> ChatResult: + return self.invoke(messages, tools, **kwargs) + + def stream( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: object, + ): + raise NotImplementedError + + async def astream( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: object, + ): + raise NotImplementedError + + +def test_message_serialization_round_trip(): + """Verify message dictionary serialization and deserialization.""" + msg = Message( + role="assistant", + content="Calculated result.", + tool_calls=[ToolCall(id="tc-1", name="add", arguments='{"a": 1, "b": 2}')], + tool_call_id=None, + ) + + d = message_to_dict(msg) + assert d["role"] == "assistant" + assert d["content"] == "Calculated result." + assert len(d["tool_calls"]) == 1 + assert d["tool_calls"][0]["name"] == "add" + + restored = dict_to_message(d) + assert restored.role == msg.role + assert restored.content == msg.content + assert len(restored.tool_calls) == 1 + assert restored.tool_calls[0].arguments == '{"a": 1, "b": 2}' + + +def test_window_memory_policy(): + """Verify WindowMemoryPolicy retains last k messages while preserving system message.""" + messages = [ + Message(role="system", content="System instruction"), + Message(role="user", content="msg 1"), + Message(role="assistant", content="msg 2"), + Message(role="user", content="msg 3"), + Message(role="assistant", content="msg 4"), + ] + + policy = WindowMemoryPolicy(k=2, preserve_system=True) + trimmed = policy.apply(messages) + + assert len(trimmed) == 3 + assert trimmed[0].role == "system" + assert trimmed[0].content == "System instruction" + assert trimmed[1].content == "msg 3" + assert trimmed[2].content == "msg 4" + + +def test_token_budget_memory_policy(): + """Verify TokenBudgetMemoryPolicy trims history to remain under token budget.""" + messages = [ + Message(role="system", content="System"), # ~1-2 tokens + Message(role="user", content="A" * 80), # ~20 tokens + Message(role="assistant", content="B" * 80), # ~20 tokens + Message(role="user", content="C" * 20), # ~5 tokens + ] + + # Limit budget so only newest user message and system fit + policy = TokenBudgetMemoryPolicy(max_tokens=15, preserve_system=True) + trimmed = policy.apply(messages) + + assert len(trimmed) == 2 + assert trimmed[0].role == "system" + assert trimmed[1].content == "C" * 20 + + +def test_summary_memory_policy(): + """Verify SummaryMemoryPolicy compacts overflow messages into a summary block.""" + messages = [ + Message(role="system", content="You are a helper."), + Message(role="user", content="What is DynamoDB?"), + Message(role="assistant", content="It is a NoSQL database."), + Message(role="user", content="What is S3?"), + Message(role="assistant", content="It is an object storage."), + Message(role="user", content="What is dynavec?"), + Message(role="assistant", content="A serverless vector database."), + ] + + policy = SummaryMemoryPolicy(max_messages=4, preserve_system=True) + compacted = policy.apply(messages) + + assert len(compacted) < len(messages) + assert compacted[0].role == "system" + assert compacted[0].content == "You are a helper." + # Second message is the summary + assert compacted[1].role == "system" + assert "Summary of prior conversation:" in compacted[1].content + assert "DynamoDB" in compacted[1].content + # Tail messages preserved + assert compacted[-1].content == "A serverless vector database." + + +def test_in_memory_memory_store(): + """Verify InMemoryMemoryStore append, retrieval, scratchpad, and clear.""" + store = InMemoryMemoryStore(policy=WindowMemoryPolicy(k=2)) + assert isinstance(store, BaseMemory) + + store.append("th-1", "Hello") + store.append("th-1", Message(role="assistant", content="Hi! How can I help?")) + store.append("th-1", "Tell me about AWS.") + + msgs = store.get_messages("th-1") + assert len(msgs) == 2 + assert msgs[0].content == "Hi! How can I help?" + assert msgs[1].content == "Tell me about AWS." + + # Scratchpad + store.set_scratchpad("th-1", {"current_step": 3, "intent": "cloud_query"}) + sp = store.get_scratchpad("th-1") + assert sp["intent"] == "cloud_query" + + store.clear("th-1") + assert store.get_messages("th-1") == [] + assert store.get_scratchpad("th-1") == {} + + +def test_dynamodb_memory_store_moto(): + """Verify DynamoDBMemoryStore with mocked AWS DynamoDB table.""" + with mock_aws(): + ddb = boto3.resource("dynamodb", region_name="us-east-1") + table = ddb.create_table( + TableName="test_agent_memory", + KeySchema=[{"AttributeName": "pk", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "pk", "AttributeType": "S"}], + BillingMode="PAY_PER_REQUEST", + ) + + store = DynamoDBMemoryStore( + table=table, + namespace="prod_app", + policy=WindowMemoryPolicy(k=3, preserve_system=True), + ) + + # Initially empty + assert store.get_messages("th-ddb-1") == [] + + # Append system and conversation messages + store.append("th-ddb-1", Message(role="system", content="You are a tutor.")) + store.append("th-ddb-1", "Explain vectors.") + store.append( + "th-ddb-1", + Message(role="assistant", content="Vectors are mathematical lists of numbers."), + ) + store.append("th-ddb-1", "How are they stored?") + store.append("th-ddb-1", Message(role="assistant", content="In S3 Vectors and DynamoDB.")) + + # Retrieve with window policy (k=3 non-system + 1 system) + messages = store.get_messages("th-ddb-1") + assert len(messages) == 4 + assert messages[0].role == "system" + assert messages[1].content == "Vectors are mathematical lists of numbers." + assert messages[2].content == "How are they stored?" + assert messages[3].content == "In S3 Vectors and DynamoDB." + + # Scratchpad + store.set_scratchpad("th-ddb-1", {"topic": "math", "progress": 0.5, "attempts": 2}) + sp = store.get_scratchpad("th-ddb-1") + assert sp["topic"] == "math" + assert sp["progress"] == 0.5 + + # Multi-thread isolation + store.append("th-ddb-2", "Different thread message") + assert len(store.get_messages("th-ddb-2")) == 1 + assert len(store.get_messages("th-ddb-1")) == 4 + + # Clear + store.clear("th-ddb-2") + assert store.get_messages("th-ddb-2") == [] + assert len(store.get_messages("th-ddb-1")) == 4 + + +def test_react_agent_memory_multi_turn(): + """Verify ReActAgent maintains multi-turn conversation memory across multiple runs.""" + model = ScriptedChatModel( + [ + # Turn 1 response + ChatResult(message=Message(role="assistant", content="Nice to meet you, Alice!")), + # Turn 2 response + ChatResult(message=Message(role="assistant", content="Your name is Alice.")), + ] + ) + + memory = InMemoryMemoryStore() + agent = ReActAgent( + model=model, + system_prompt="Helpful assistant.", + memory=memory, + ) + + # Turn 1 + res1 = agent.run("My name is Alice.", thread_id="thread-conv-1") + assert res1.output == "Nice to meet you, Alice!" + + # Verify memory recorded turn 1 (system + user + assistant) + msgs_t1 = memory.get_messages("thread-conv-1") + assert len(msgs_t1) == 3 + assert msgs_t1[0].role == "system" + assert msgs_t1[1].content == "My name is Alice." + assert msgs_t1[2].content == "Nice to meet you, Alice!" + + # Turn 2 + res2 = agent.run("What is my name?", thread_id="thread-conv-1") + assert res2.output == "Your name is Alice." + + # Verify agent received prior turn history in turn 2 + assert len(model.received_messages) == 2 + t2_messages = model.received_messages[1] + assert len(t2_messages) == 4 + assert t2_messages[1].content == "My name is Alice." + assert t2_messages[2].content == "Nice to meet you, Alice!" + assert t2_messages[3].content == "What is my name?" + + # Verify memory contains all 5 messages + final_msgs = memory.get_messages("thread-conv-1") + assert len(final_msgs) == 5