File size: 28,426 Bytes
750b88d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
"""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)}