diff --git a/README.md b/README.md index bc44295..91c59a4 100644 --- a/README.md +++ b/README.md @@ -67,3 +67,5 @@ Results are saved to `outputs/{dataset}/{memory}/{mode}/{domain}.json` and can b - Python ≥ 3.11 - `GEMINI_API_KEY` in `.env` or environment - For MemBench: set `MEMBENCH_DATA_PATH` to your local data directory +- For `--memory letta`: `LETTA_API_KEY` (Letta Cloud) or `LETTA_BASE_URL` (self-hosted server). + Optional: `LETTA_EMBEDDING_MODEL`, `LETTA_MODEL` and `LETTA_MAX_STEPS` (agent mode only) diff --git a/catalog.json b/catalog.json index c38afd8..2659ded 100644 --- a/catalog.json +++ b/catalog.json @@ -96,6 +96,13 @@ "link": "https://cognee.ai", "logo": "https://www.google.com/s2/favicons?sz=32&domain=cognee.ai" }, + "letta": { + "key": "letta", + "description": "Letta archival memory: documents are written as passages into a per-unit archive and retrieved by semantic search. Agent mode answers through a Letta agent that searches its own archival memory.", + "kind": "cloud", + "link": "https://letta.com", + "logo": "https://www.google.com/s2/favicons?sz=32&domain=letta.com" + }, "mastra": { "key": "mastra", "description": "Mastra semantic recall with LibSQL store and FastEmbed embeddings. topK=10.", diff --git a/pyproject.toml b/pyproject.toml index c7183a4..c0b47ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "sentence-transformers>=3.0", "python-dotenv>=1.0", "hindsight-all>=0.4", + "letta-client>=1.12.1", "supermemory>=0.1", "httpx>=0.27", "qdrant-client>=1.13", diff --git a/src/memory_bench/memory/__init__.py b/src/memory_bench/memory/__init__.py index 2b7e5e0..b08e97f 100644 --- a/src/memory_bench/memory/__init__.py +++ b/src/memory_bench/memory/__init__.py @@ -2,6 +2,7 @@ from .bm25 import BM25MemoryProvider from .cognee import CogneeMemoryProvider from .hindsight import HindsightCloudMemoryProvider, HindsightHTTPMemoryProvider, HindsightMemoryProvider +from .letta import LettaMemoryProvider from .mastra import MastraMemoryProvider from .mastra_om import MastraOMMemoryProvider from .mem0 import Mem0MemoryProvider @@ -21,6 +22,7 @@ "hindsight-cloud": HindsightCloudMemoryProvider, "hindsight-http": HindsightHTTPMemoryProvider, + "letta": LettaMemoryProvider, "mastra": MastraMemoryProvider, "mastra-om": MastraOMMemoryProvider, "mem0": Mem0MemoryProvider, diff --git a/src/memory_bench/memory/letta.py b/src/memory_bench/memory/letta.py new file mode 100644 index 0000000..a3db5cc --- /dev/null +++ b/src/memory_bench/memory/letta.py @@ -0,0 +1,230 @@ +"""Letta memory provider. + +Documents are stored as passages in a Letta archive (one archive per isolation +unit) and retrieved with Letta's semantic passage search. Agent mode answers +through a Letta agent that has the archive attached, so the agent decides for +itself when and what to search. + +Works against Letta Cloud (LETTA_API_KEY) or a self-hosted server +(LETTA_BASE_URL, e.g. http://localhost:8283). +""" + +import os +import threading +import uuid +from datetime import datetime, timezone +from pathlib import Path + +from ..models import Document +from .base import MemoryProvider + +_BATCH_SIZE = 50 + + +def _parse_iso_ts(ts: str | None) -> datetime | None: + """Best-effort ISO-8601 parse; returns a timezone-aware UTC datetime or None.""" + if not ts: + return None + try: + dt = datetime.fromisoformat(ts.replace("Z", "+00:00")) + except (ValueError, TypeError): + return None + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt.astimezone(timezone.utc) + + +def _message_text(content) -> str: + """Flatten Letta message content (a string or a list of text parts) into text.""" + if isinstance(content, str): + return content + return "".join(part.text for part in content or [] if getattr(part, "text", None)) + + +def _prefix_from_store_dir(store_dir: Path) -> str: + """Derive a stable archive-name prefix from the run's store directory.""" + parts = store_dir.parts + try: + idx = parts.index("_store") + return f"amb-{parts[idx - 2]}-{parts[idx + 1]}" + except (ValueError, IndexError): + return "amb-bench" + + +class LettaMemoryProvider(MemoryProvider): + name = "letta" + description = ( + "Letta archival memory: documents are written as passages into a per-unit archive " + "and retrieved by semantic search. Agent mode answers through a Letta agent that " + "searches its own archival memory." + ) + kind = "cloud" + link = "https://letta.com" + logo = "https://www.google.com/s2/favicons?sz=32&domain=letta.com" + + def __init__(self, k: int = 20): + self.k = k + self._client = None + self._prefix = "amb-bench" + self._per_unit = False + self._archive_ids: dict[str | None, str] = {} + self._agent_ids: dict[str | None, str] = {} + self._agent_locks: dict[str | None, threading.Lock] = {} + self._lock = threading.Lock() + self._embedding = os.environ.get("LETTA_EMBEDDING_MODEL", "openai/text-embedding-3-small") + self._model = os.environ.get("LETTA_MODEL", "openai/gpt-4.1") + self._max_steps = int(os.environ.get("LETTA_MAX_STEPS", "10")) + + def initialize(self) -> None: + from letta_client import Letta + + if not os.environ.get("LETTA_API_KEY") and not os.environ.get("LETTA_BASE_URL"): + raise RuntimeError( + "letta provider needs LETTA_API_KEY (Letta Cloud) or LETTA_BASE_URL (self-hosted server)" + ) + # api_key comes from LETTA_API_KEY; base_url from LETTA_BASE_URL when self-hosted. + self._client = Letta() + + def cleanup(self) -> None: + # Archives are kept (they hold the ingested corpus); the throwaway agents are not. + for agent_id in self._agent_ids.values(): + try: + self._client.agents.delete(agent_id) + except Exception: + pass + self._agent_ids.clear() + + def prepare(self, store_dir: Path, unit_ids: set[str] | None = None, reset: bool = True) -> None: + self._prefix = _prefix_from_store_dir(store_dir) + self._per_unit = unit_ids is not None + self._archive_ids.clear() + self._agent_ids.clear() + self._agent_locks.clear() + for unit in sorted(unit_ids) if unit_ids else [None]: + self._ensure_archive(unit, reset=reset) + + def _archive_name(self, unit: str | None) -> str: + return f"{self._prefix}-u{unit}" if unit is not None else self._prefix + + def _ensure_archive(self, unit: str | None, reset: bool = False) -> str: + with self._lock: + if unit in self._archive_ids: + return self._archive_ids[unit] + name = self._archive_name(unit) + existing = list(self._client.archives.list(name=name, limit=100)) + if reset: + for archive in existing: + self._client.archives.delete(archive.id) + existing = [] + archive = existing[0] if existing else self._client.archives.create( + name=name, + description="Agent Memory Benchmark run", + embedding=self._embedding, + ) + self._archive_ids[unit] = archive.id + return archive.id + + def _unit(self, user_id: str | None) -> str | None: + return user_id if self._per_unit else None + + @staticmethod + def _text(doc: Document) -> str: + if doc.timestamp: + return f"[Date: {doc.timestamp}]\n{doc.content}" + return doc.content + + def ingest(self, documents: list[Document]) -> None: + by_unit: dict[str | None, list[Document]] = {} + for doc in documents: + by_unit.setdefault(self._unit(doc.user_id), []).append(doc) + + for unit, docs in by_unit.items(): + archive_id = self._ensure_archive(unit) + passages = [ + {"text": self._text(doc), "metadata": {"doc_id": doc.id, "timestamp": doc.timestamp}} + for doc in docs + ] + for i in range(0, len(passages), _BATCH_SIZE): + self._client.archives.passages.create_many( + archive_id, passages=passages[i : i + _BATCH_SIZE] + ) + + def retrieve( + self, query: str, k: int = 10, user_id: str | None = None, query_timestamp: str | None = None + ) -> tuple[list[Document], dict | None]: + archive_id = self._ensure_archive(self._unit(user_id)) + k_eff = k or self.k + # Fetch a small buffer so a strict timestamp filter can still return up to k results. + results = self._client.passages.search( + archive_id=archive_id, query=query, limit=max(k_eff, 50) + ) + + query_dt = _parse_iso_ts(query_timestamp) + docs: list[Document] = [] + raw_results: list[dict] = [] + for i, r in enumerate(results): + passage = r.passage + if query_dt is not None: + doc_ts = _parse_iso_ts((passage.metadata or {}).get("timestamp")) + if doc_ts is not None and doc_ts > query_dt: + continue + docs.append(Document(id=passage.id or f"letta-{i}", content=passage.text)) + raw_results.append( + { + "id": passage.id, + "text": passage.text, + "score": r.score, + "tags": passage.tags, + "metadata": passage.metadata, + } + ) + if len(docs) >= k_eff: + break + return docs, {"results": raw_results} + + def _ensure_agent(self, unit: str | None) -> tuple[str, threading.Lock]: + archive_id = self._ensure_archive(unit) + with self._lock: + if unit not in self._agent_ids: + agent = self._client.agents.create( + name=f"{self._archive_name(unit)}-{uuid.uuid4().hex[:6]}", + model=self._model, + embedding=self._embedding, + include_base_tools=True, + message_buffer_autoclear=True, + ) + try: + self._client.agents.archives.attach(archive_id, agent_id=agent.id) + except Exception: + try: + self._client.agents.delete(agent.id) + except Exception: + pass + raise + self._agent_ids[unit] = agent.id + self._agent_locks[unit] = threading.Lock() + return self._agent_ids[unit], self._agent_locks[unit] + + def direct_answer( + self, query: str, user_id: str | None = None, query_timestamp: str | None = None + ) -> tuple[str, str, dict | None]: + unit = self._unit(user_id) + agent_id, lock = self._ensure_agent(unit) + # A Letta agent processes messages sequentially; concurrent sends interleave. + input_text = query + if query_timestamp: + input_text = f"[Question date: {query_timestamp} UTC]\n{query}" + with lock: + response = self._client.agents.messages.create( + agent_id, input=input_text, max_steps=self._max_steps + ) + + answers: list[str] = [] + context_parts: list[str] = [] + for message in response.messages: + if message.message_type == "assistant_message": + answers.append(_message_text(message.content)) + elif message.message_type == "tool_return_message": + context_parts.append(message.tool_return) + + return "\n".join(answers), "\n\n".join(context_parts), response.model_dump(mode="json") diff --git a/uv.lock b/uv.lock index 8f029f9..83f5f61 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.11" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -217,6 +217,7 @@ dependencies = [ { name = "groq" }, { name = "hindsight-all" }, { name = "httpx" }, + { name = "letta-client" }, { name = "mem0ai" }, { name = "python-dotenv" }, { name = "qdrant-client" }, @@ -239,6 +240,7 @@ requires-dist = [ { name = "groq", specifier = ">=1.1.1" }, { name = "hindsight-all", specifier = ">=0.4" }, { name = "httpx", specifier = ">=0.27" }, + { name = "letta-client", specifier = ">=1.12.1" }, { name = "mem0ai", specifier = ">=1.0.5" }, { name = "python-dotenv", specifier = ">=1.0" }, { name = "qdrant-client", specifier = ">=1.13" }, @@ -1006,7 +1008,7 @@ name = "cuda-bindings" version = "12.9.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder" }, + { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/45/e7/b47792cc2d01c7e1d37c32402182524774dadd2d26339bd224e0e913832e/cuda_bindings-12.9.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c912a3d9e6b6651853eed8eed96d6800d69c08e94052c292fec3f282c5a817c9", size = 12210593, upload-time = "2025-10-21T14:51:36.574Z" }, @@ -2686,6 +2688,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d0/a8/4202ca65561213ec84ca3800b1d4e5d37a1441cddeec533367ecbca7f408/langsmith-0.7.16-py3-none-any.whl", hash = "sha256:c84a7a06938025fe0aad992acc546dd75ce3f757ba8ee5b00ad914911d4fc02e", size = 347538, upload-time = "2026-03-09T21:11:15.02Z" }, ] +[[package]] +name = "letta-client" +version = "1.12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8e/55/34a347fb443f5797045cde591c34ee2926650e99f1f0d5f565dcefa99120/letta_client-1.12.1.tar.gz", hash = "sha256:3073adb6cc1ce3b906a40d8ce3b2b77a77ac1887ca2135770590cfe3cef6eba9", size = 396805, upload-time = "2026-06-02T00:31:12.848Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/af/b140a7117b4385b3e1b060df97d9d39703f0c54e156a4d889de8bb8dd48f/letta_client-1.12.1-py3-none-any.whl", hash = "sha256:6f554569a684d57e7f4e4513f8ce6792129b305df37393c650e50776f61baf99", size = 418856, upload-time = "2026-06-02T00:31:14.475Z" }, +] + [[package]] name = "limits" version = "4.8.0" @@ -3543,7 +3562,7 @@ name = "nvidia-cudnn-cu12" version = "9.10.2.21" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, @@ -3554,7 +3573,7 @@ name = "nvidia-cufft-cu12" version = "11.3.3.83" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, @@ -3581,9 +3600,9 @@ name = "nvidia-cusolver-cu12" version = "11.7.3.90" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12" }, - { name = "nvidia-cusparse-cu12" }, - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cusparse-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, @@ -3594,7 +3613,7 @@ name = "nvidia-cusparse-cu12" version = "12.5.8.93" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, @@ -5616,8 +5635,8 @@ name = "secretstorage" version = "3.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cryptography" }, - { name = "jeepney" }, + { name = "cryptography", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "jeepney", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } wheels = [