"""Letta (MemGPT) adapter. Each visit is one message to a Letta agent (split if longer than the chunk size); the agent manages its core, recall and archival memory with its own tools. Questions go to the same agent with the clinician answering prompt. After each answer the agent is restored to its state before the question: the question's messages are deleted, archival and core-memory edits are undone, and the system message gets its earlier text back. The system prompt is data/letta_memgpt_chat_0.7.0.txt. With `max_steps`, a chain at the cap is asked for its final reply. """ from __future__ import annotations import importlib import importlib.util import json import logging import os import shutil import sys import time import types from pathlib import Path from typing import Any, Dict, List from .base import BaseAgent, MemoryBuildResult, MemoryContext from ..llm_client import ENDPOINTS, MAX_ATTEMPTS, USAGE, api_key_for, http_timeout from ..prompts import ANSWER_NOCTX from .hipporag_agent import _unwrap_answer, _VISIT_HEADER_RE logger = logging.getLogger(__name__) _ACTIVE_LETTA_AGENT: "LettaAgent | None" = None LETTA_ROOT = (Path(__file__).resolve().parents[2] / "vendor" / "letta") _OPTIONAL_DEPS = { "composio": ("composio", "composio.constants", "composio.client", "composio.client.collections", "composio.exceptions", "composio_langchain"), "mcp": ("mcp", "mcp.client", "mcp.client.sse", "mcp.client.stdio", "mcp.types"), } def _stub_missing_optional_deps() -> None: """Register import stubs for composio and mcp when they are not installed; the agent never uses them.""" for root, names in _OPTIONAL_DEPS.items(): try: __import__(root) continue except ImportError: pass for name in names: if name in sys.modules: continue m = types.ModuleType(name) m.__path__ = [] m.__getattr__ = lambda attr, _n=name: type(attr, (), {}) m.DEFAULT_ENTITY_ID = "default" sys.modules[name] = m _NO_TEXT = object() def _message_kind(msg) -> str: """The message_type of a Letta message ('assistant_message', 'tool_call_message', ...), else its class name.""" kind = getattr(msg, "message_type", None) if kind is None: return type(msg).__name__ return str(kind.value) if hasattr(kind, "value") else str(kind) def _content_text(content): """Text of an assistant message's content: a string, a list of parts, or an object with `text` or `message`; _NO_TEXT when the content has none of these.""" if isinstance(content, str): return content if isinstance(content, list): pieces = [] for part in content: if hasattr(part, "text"): pieces.append(part.text) elif isinstance(part, str): pieces.append(part) elif hasattr(part, "content"): pieces.append(str(part.content)) return " ".join(pieces) for attr in ("text", "message"): if hasattr(content, attr): return getattr(content, attr) return _NO_TEXT def _tool_call(msg) -> tuple: """(name, arguments, tool_call_id) of a tool-call message; JSON-string arguments are decoded when possible.""" call = getattr(msg, "tool_call", None) if call: name, args = getattr(call, "name", ""), getattr(call, "arguments", "{}") else: name = getattr(msg, "function_name", None) or getattr(msg, "name", "") args = getattr(msg, "function_args", None) or getattr(msg, "arguments", {}) if isinstance(args, str): try: args = json.loads(args) except json.JSONDecodeError: pass return name, args, (getattr(call, "tool_call_id", None) if call else None) def _read_response(response) -> Dict[str, Any]: """The reply of one Letta step: the assistant message (or the message of a send_message call), the tool calls with their returns (truncated to 1,000 characters), the reasoning messages and the token usage.""" out = {"assistant_message": "", "function_calls": [], "internal_monologue": [], "usage": {}} usage = getattr(response, "usage", None) if usage: out["usage"] = {k: getattr(usage, k, 0) for k in ("completion_tokens", "prompt_tokens", "total_tokens")} for msg in getattr(response, "messages", None) or []: kind, cls = _message_kind(msg), str(type(msg)) if kind == "assistant_message" or "AssistantMessage" in cls: content = getattr(msg, "content", None) text = _content_text(content) if content else _NO_TEXT if text is not _NO_TEXT: out["assistant_message"] = text elif kind == "tool_call_message" or "ToolCall" in cls: name, args, call_id = _tool_call(msg) if not name: continue out["function_calls"].append({"name": name, "tool_call_id": call_id, "arguments": args if isinstance(args, dict) else str(args)}) if name == "send_message" and isinstance(args, dict): said = args.get("message", "") if said and not out["assistant_message"]: out["assistant_message"] = said elif kind in ("reasoning_message", "internal_monologue") or "Reasoning" in cls or "InternalMonologue" in cls: thought = (getattr(msg, "reasoning", None) or getattr(msg, "internal_monologue", None) or getattr(msg, "content", "")) if thought: out["internal_monologue"].append(thought) elif kind in ("tool_return_message", "function_return") or "ToolReturn" in cls or "FunctionReturn" in cls: returned = getattr(msg, "tool_return", None) or getattr(msg, "function_return", None) if not (returned and out["function_calls"]): continue call_id = getattr(msg, "tool_call_id", None) owner = next((c for c in reversed(out["function_calls"]) if call_id and c.get("tool_call_id") == call_id), out["function_calls"][-1]) owner["return"] = str(returned)[:1000] return out class LettaAgent(BaseAgent): ANSWERS_INTERNALLY = True MEMORIZE_WITH_LLM = True def __init__(self, *, model: str | None = None, temperature: float, max_tokens: int, retrieve_num: int, context_window: int, embedding_model: str | None = None, embedding_dim: int | None = None, embedding_chunk_size: int, memorize_chunk_tokens: int, memorize_chunk_overlap_tokens: int, memory_persona: str = ("I am an assistant helping with medical QA " "while preserving long-term memory."), memory_human: str = ("The user is a patient in a longitudinal " "medical dialogue setting."), internal_base_url: str | None = None, embedding_base_url: str | None = None, message_buffer_autoclear: bool, max_steps: int | None, working_dir: str | None = None, **params): super().__init__(**params) if not working_dir: raise ValueError("LettaAgent requires working_dir (the runner sets it per run)") self.model = model self.temperature = temperature self.max_tokens = max_tokens self.retrieve_num = retrieve_num self.context_window = context_window self.embedding_model = embedding_model self.embedding_dim = embedding_dim self.embedding_chunk_size = embedding_chunk_size self.memorize_chunk_tokens = memorize_chunk_tokens self.memorize_chunk_overlap_tokens = memorize_chunk_overlap_tokens self.max_question_tokens = context_window self.memory_persona = memory_persona self.memory_human = memory_human self.message_buffer_autoclear = bool(message_buffer_autoclear) self.max_steps = max_steps embedding_base_url = embedding_base_url or ENDPOINTS["embedding"] self.api_key = api_key_for(internal_base_url) if api_key_for(embedding_base_url) != self.api_key: raise ValueError("Letta sends one API key to its LLM and its embedding endpoint; give the memory (or " "answer) and embedding endpoints the same api_key_env in backbone.yaml") _default = os.environ.get("OPENAI_BASE_URL") or "https://api.openai.com/v1" self.llm_base_url = internal_base_url or _default self.embedding_base_url = embedding_base_url or _default self.persistence_root = (Path(working_dir).resolve() / "letta_db") if self.persistence_root.exists(): shutil.rmtree(self.persistence_root) self._agent_ids: Dict[str, str] = {} self._visits_memorized = 0 self._id2visit: Dict[str, int] = {} self._cur_visit = 0 self._last_hits: List[str] = [] self._answering = False self._answer_passages: List[str] = [] self._set_letta_dir() _stub_missing_optional_deps() self._import_letta() self._instrument_clients() self._connect() def _set_letta_dir(self) -> None: self.persistence_root.mkdir(parents=True, exist_ok=True) os.environ["LETTA_DIR"] = str(self.persistence_root) os.environ["LETTA_LETTA_DIR"] = str(self.persistence_root) def _import_letta(self) -> None: init_file = LETTA_ROOT / "__init__.py" if not init_file.exists(): raise ImportError(f"Letta source not found at {LETTA_ROOT}") for key in [k for k in sys.modules if k == "letta" or k.startswith("letta.")]: del sys.modules[key] spec = importlib.util.spec_from_file_location( "letta", str(init_file), submodule_search_locations=[str(LETTA_ROOT)]) module = importlib.util.module_from_spec(spec) sys.modules["letta"] = module spec.loader.exec_module(module) def _instrument_clients(self) -> None: """Add SDK retries and embedding-usage logging to the engine's clients, and hook passage inserts and searches to record the visit of each passage, the passages written while answering and the passages found.""" import openai as _openai usage_log = USAGE.log def _retrying_openai(**kw): kw["max_retries"] = MAX_ATTEMPTS - 1 kw.setdefault("timeout", http_timeout()) return _openai.OpenAI(**kw) llm_api = importlib.import_module("letta.llm_api.openai") llm_api.OpenAI = _retrying_openai oc = importlib.import_module("letta.llm_api.openai_client") oc.OpenAI = _retrying_openai class _TrackedEmbeddings: def __init__(self, real): self._real = real def create(self, **kw): resp = self._real.embeddings.create(**kw) u = getattr(resp, "usage", None) if u is not None: usage_log(kw.get("model", "?"), "embed", int(getattr(u, "prompt_tokens", 0) or 0), 0) return resp class _TrackedClient: def __init__(self, real): self.embeddings = _TrackedEmbeddings(real) def _tracked_openai(**kw): kw["max_retries"] = MAX_ATTEMPTS - 1 kw.setdefault("timeout", http_timeout()) return _TrackedClient(_openai.OpenAI(**kw)) pm = importlib.import_module("letta.services.passage_manager") pm.OpenAI = _tracked_openai emb = importlib.import_module("letta.embeddings") emb.OpenAI = _retrying_openai def _tracked_get_text_embedding(self, text: str): resp = self.client.embeddings.create(input=text, model=self.model) u = getattr(resp, "usage", None) if u is not None: usage_log(self.model, "embed", int(getattr(u, "prompt_tokens", 0) or 0), 0) return resp.data[0].embedding emb.OpenAIEmbeddings.get_text_embedding = _tracked_get_text_embedding global _ACTIVE_LETTA_AGENT _ACTIVE_LETTA_AGENT = self def _agent(): return _ACTIVE_LETTA_AGENT if not getattr(pm.PassageManager.insert_passage, "_memebench_hooked", False): _orig_insert = pm.PassageManager.insert_passage def _hooked_insert(pm_self, *a, **kw): out = _orig_insert(pm_self, *a, **kw) ids = [str(p.id) for p in (out or [])] for pid in ids: _agent()._id2visit[pid] = _agent()._cur_visit if _agent()._answering: _agent()._answer_passages.extend(ids) try: _agent()._trace("write_op", op="archival_insert", visit=_agent()._cur_visit, passage_ids=ids, text=kw.get("text")) except Exception: pass return out _hooked_insert._memebench_hooked = True pm.PassageManager.insert_passage = _hooked_insert am = importlib.import_module("letta.services.agent_manager") if not getattr(am.AgentManager.list_passages, "_memebench_hooked", False): _orig_list = am.AgentManager.list_passages def _hooked_list(am_self, *a, **kw): out = _orig_list(am_self, *a, **kw) if kw.get("query_text"): _agent()._last_hits = [str(p.id) for p in (out or [])] return out _hooked_list._memebench_hooked = True am.AgentManager.list_passages = _hooked_list def _connect(self) -> None: model_settings = importlib.import_module("letta.settings").model_settings if self.api_key: model_settings.openai_api_key = self.api_key model_settings.openai_api_base = self.embedding_base_url client_module = importlib.import_module("letta.client.client") schema = {name: importlib.import_module(f"letta.schemas.{name}") for name in ("llm_config", "embedding_config", "memory")} self._LLMConfig = schema["llm_config"].LLMConfig self._EmbeddingConfig = schema["embedding_config"].EmbeddingConfig self._ChatMemory = schema["memory"].ChatMemory self._client = client_module.create_client(base_url=None, token=None) if self.max_steps is not None: self._client.server.max_chaining_steps = int(self.max_steps) os.environ["MEMEBENCH_LETTA_FINAL_REPLY_AT_CAP"] = "1" logger.info("[Letta] max_steps=%d; the agent gives its final reply at the cap", int(self.max_steps)) def _llm_config(self): return self._LLMConfig(model=self.model, model_endpoint_type="openai", model_endpoint=self.llm_base_url, model_wrapper=None, context_window=self.context_window, temperature=self.temperature, max_tokens=self.max_tokens) def _embedding_config(self): return self._EmbeddingConfig(embedding_endpoint_type="openai", embedding_endpoint=self.embedding_base_url, embedding_model=self.embedding_model, embedding_dim=self.embedding_dim, embedding_chunk_size=self.embedding_chunk_size) def _system_prompt(self) -> str: return (Path(__file__).resolve().parent / "data" / "letta_memgpt_chat_0.7.0.txt").read_text(encoding="utf-8") def _agent_for(self, context_id: str) -> str: existing = self._agent_ids.get(context_id) if existing: return existing name = f"memebench_letta_ctx_{context_id}_{int(time.time())}" state = self._client.create_agent( name=name, system=self._system_prompt(), llm_config=self._llm_config(), embedding_config=self._embedding_config(), memory=self._ChatMemory(persona=self.memory_persona, human=self.memory_human), metadata={"retrieve_num": self.retrieve_num, "embedding_chunk_size": self.embedding_chunk_size}, message_buffer_autoclear=self.message_buffer_autoclear, ) self._agent_ids[context_id] = str(state.id) logger.info(f"[Letta] created agent {state.id} for context {context_id}") return str(state.id) def _context(self) -> str: return self._context_id if self._context_id is not None else "0" def _log_usage(self, parsed: Dict[str, Any]) -> None: u = parsed.get("usage") or {} p, c = int(u.get("prompt_tokens", 0) or 0), int(u.get("completion_tokens", 0) or 0) if p or c: USAGE.log(self.model, "chat", p, c) def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult: self._cur_visit = int(visit_number) t0 = time.time() agent_id = self._agent_for(self._context()) chunks = self.split_tokens(text, self.memorize_chunk_tokens, self.memorize_chunk_overlap_tokens) writes = 0 calls: List[str] = [] for i, chunk in enumerate(chunks): try: response = self._client.user_message(agent_id=agent_id, message=chunk) except Exception as e: logger.error(f"[Letta] user_message failed on v{visit_number} chunk {i}: {e}") return MemoryBuildResult(success=False, method="letta", action="user_message", visit_number=visit_number, chunk_count=self._visits_memorized, time_cost=time.time() - t0, extra={"error": str(e)[:200], "failed_chunk": i}) parsed = _read_response(response) self._log_usage(parsed) try: self._trace("tool_calls", stage="write", chunk=i, calls=[{"name": fc.get("name"), "arguments": fc.get("arguments")} for fc in parsed["function_calls"]], assistant_message=parsed.get("assistant_message")) except Exception: pass for fc in parsed["function_calls"]: calls.append(fc["name"]) if fc["name"] in ("archival_memory_insert", "core_memory_append", "core_memory_replace"): writes += 1 if fc["name"] in ("core_memory_append", "core_memory_replace"): a = fc.get("arguments") if isinstance(fc.get("arguments"), dict) else {} try: self._trace("write_op", op=("core_append" if fc["name"].endswith("append") else "core_replace"), visit=int(visit_number), block=a.get("label"), text=a.get("content") or a.get("new_content"), text_before=a.get("old_content"), passage_ids=[], chunk=i) except Exception: pass self._visits_memorized += 1 return MemoryBuildResult(success=True, method="letta", action="user_message", visit_number=visit_number, chunk_count=self._visits_memorized, added_chunks=writes, time_cost=time.time() - t0, extra={"memory_writes": writes, "tool_calls": calls}) def _state_before_question(self, agent_id: str) -> Dict[str, Any]: """The newest stored message, the in-context message ids, the core memory blocks and the system message (id and text) of the agent.""" srv, actor = self._client.server, self._client.user last = srv.message_manager.list_messages_for_agent(agent_id=agent_id, actor=actor, limit=1, ascending=False) agent = srv.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor) in_context = [str(m) for m in (agent.message_ids or [])] system = srv.message_manager.get_message_by_id(message_id=in_context[0], actor=actor) return {"last": str(last[0].id) if last else None, "in_context": in_context, "blocks": {b.label: b.value for b in agent.memory.blocks}, "system": (str(system.id), system.content[0].text)} def _forget_question(self, agent_id: str, before: Dict[str, Any]) -> Dict[str, int]: """Return the agent to `before`: delete the messages stored since, the archival passages written while answering, and restore the core memory blocks, the in-context message ids and the text of the system message. Letta rewrites the system message in place (same id) whenever core memory or archival memory changes, so restoring the ids alone would keep the text compiled while answering.""" srv, actor = self._client.server, self._client.user agent = srv.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor) if [str(m) for m in (agent.message_ids or [])] != before["in_context"]: srv.agent_manager.set_in_context_messages(agent_id=agent_id, message_ids=before["in_context"], actor=actor) new = srv.message_manager.list_messages_for_agent(agent_id=agent_id, actor=actor, after=before["last"], limit=None) for m in new: srv.message_manager.delete_message_by_id(message_id=m.id, actor=actor) for pid in self._answer_passages: srv.passage_manager.delete_passage_by_id(passage_id=pid, actor=actor) self._id2visit.pop(pid, None) n_passages = len(self._answer_passages) self._answer_passages = [] block_update = importlib.import_module("letta.schemas.block").BlockUpdate n_blocks = 0 for label, value in before["blocks"].items(): cur = srv.agent_manager.get_block_with_label(agent_id=agent_id, block_label=label, actor=actor) if cur.value != value: srv.block_manager.update_block(block_id=cur.id, block_update=block_update(value=value), actor=actor) n_blocks += 1 system_id, system_text = before["system"] n_system = 0 if srv.message_manager.get_message_by_id(message_id=system_id, actor=actor).content[0].text != system_text: message_update = importlib.import_module("letta.schemas.message").MessageUpdate srv.message_manager.update_message_by_id(message_id=system_id, message_update=message_update(content=system_text), actor=actor) n_system = 1 return {"messages": len(new), "passages": n_passages, "blocks": n_blocks, "system": n_system} def build_memory_context(self, question: str, as_of: str = "", display_question: str | None = None) -> MemoryContext: agent_id = self._agent_for(self._context()) before = self._state_before_question(agent_id) self._answering, self._answer_passages = True, [] try: ctx = self._answer(agent_id, question, as_of, display_question) finally: self._answering = False forgotten = self._forget_question(agent_id, before) ctx.extra = dict(ctx.extra or {}, forgotten=forgotten) return ctx def _answer(self, agent_id: str, question: str, as_of: str, display_question: str | None) -> MemoryContext: wrapped = ANSWER_NOCTX.format(as_of=as_of, qtext=display_question or question) wrapped = self.cut_tokens(wrapped, self.max_question_tokens) t0 = time.time() self._last_hits = [] try: response = self._client.user_message(agent_id=agent_id, message=wrapped) except Exception as e: kind = ("context_window_exceeded" if "ContextWindowExceeded" in type(e).__name__ else "agent_error") logger.error(f"[Letta] query failed ({kind}): {e}") return MemoryContext(mem="", n_tok=0, retrieved_count=0, error=True, answer=f"[err {kind}]", extra={"method": "letta", "error": str(e)[:200], "error_kind": kind}) parsed = _read_response(response) self._log_usage(parsed) searches = [fc for fc in parsed["function_calls"] if fc["name"] in ("archival_memory_search", "conversation_search") and fc.get("return")] blocks = [self.cut_tokens(fc["return"], self.context_window // 4) for fc in searches] mem = "\n\n".join(f"[Memory {i + 1}]\n{b}" for i, b in enumerate(blocks)) visits = sorted({int(v) for b in blocks for v in _VISIT_HEADER_RE.findall(b)} | {self._id2visit[pid] for pid in self._last_hits if pid in self._id2visit}) try: self._trace("tool_calls", stage="query", question=question, calls=[{"name": fc.get("name"), "arguments": fc.get("arguments")} for fc in parsed["function_calls"]], n_blocks=len(blocks)) except Exception: pass try: self._trace("query_retrieve", question=question, as_of=as_of, hits=[{"tool": fc.get("name"), "query": (fc.get("arguments") or {}).get("query") if isinstance(fc.get("arguments"), dict) else None, "tool_call_id": fc.get("tool_call_id"), "return": (fc.get("return") or "")[:2000]} for fc in searches], n_hits=len(searches), n_blocks=len(blocks), visits=visits, visit_source={"header_regex": sorted({int(v) for b in blocks for v in _VISIT_HEADER_RE.findall(b)}), "id_map": sorted({self._id2visit[pid] for pid in self._last_hits if pid in self._id2visit})}, anchored_ids=len(self._id2visit), retrieved_passage_ids=list(self._last_hits)) except Exception: pass return MemoryContext( mem=mem, n_tok=self.count_tokens(mem), retrieved_count=len(blocks), retrieved_visits=visits, retrieved_memories=[{"visit": None, "memory": b[:300]} for b in blocks], answer=_unwrap_answer(parsed["assistant_message"] or ""), extra={"method": "letta", "tool_calls": [fc["name"] for fc in parsed["function_calls"]], "query_time_s": round(time.time() - t0, 2)}, ) def dump_memory(self, out_dir): from .base import snapshot_tree db = Path(self.persistence_root) if not db.is_dir(): return None files = snapshot_tree(db, out_dir / "letta_db") (out_dir / "id2visit.json").write_text( json.dumps(self._id2visit, ensure_ascii=False), encoding="utf-8") return {"files": files, "anchored_ids": len(self._id2visit)}