"""MIRIX adapter. Visits are sent to MIRIX's meta agent, which routes them to the memory managers named in `memory_types`; a routing to a type outside the list is dropped. A routing that fails (a memory manager raised) reaches the meta agent as a tool error; the visit is then reported as not successful. Each run uses its own SQLite store, and episodic memories are stamped with the visit date. With `answer_agent: true`, each question is answered by a fresh MIRIX chat agent that loads memories into its system prompt, may call its memory-search tools for up to `answer_max_steps` steps, and replies through `send_message`. """ from __future__ import annotations import asyncio import json import logging import os import shutil import sys import time from pathlib import Path from typing import Any, Dict, List from .base import EMBED_ENC, BaseAgent, MemoryBuildResult, MemoryContext from ..llm_client import HTTP_TIMEOUT_S, MAX_ATTEMPTS, USAGE, http_timeout from ..prompts import ANSWER_NOCTX from .hipporag_agent import _unwrap_answer logger = logging.getLogger(__name__) def _log_embed_est(model: str, texts) -> None: USAGE.log(model, "embed", sum(len(EMBED_ENC.encode(str(t) or "")) for t in texts), 0) MIRIX_PATH = str(Path(__file__).resolve().parents[2] / "vendor" / "MIRIX") ALL_MEMORY_TYPES = ("core", "episodic", "semantic", "procedural", "resource", "knowledge_vault") class MIRIXAgent(BaseAgent): ANSWERS_INTERNALLY = False MEMORIZE_WITH_LLM = True def __init__(self, *, model: str | None = None, temperature: float, max_tokens: int, retrieve_num: int, memorize_chunk_tokens: int, memorize_chunk_overlap_tokens: int, query_memory_item_tokens: int, query_memory_context_tokens: int, max_input_tokens: int, context_window: int, embedding_max_input_tokens: int, embedding_model: str | None = None, embedding_dim: int | None = None, embedding_batch_size: int | None = None, memory_human: str = ("The user is a patient in a longitudinal " "medical dialogue setting."), memory_persona: str = ("I am an assistant helping with medical QA " "while preserving long-term memory."), internal_base_url: str | None = None, embedding_base_url: str | None = None, answer_agent: bool, answer_max_steps: int, memory_types: tuple | list, working_dir: str | None = None, **params): super().__init__(**params) if not working_dir: raise ValueError("MIRIXAgent 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.memorize_chunk_tokens = memorize_chunk_tokens self.memorize_chunk_overlap_tokens = memorize_chunk_overlap_tokens self.query_memory_item_tokens = query_memory_item_tokens self.query_memory_context_tokens = query_memory_context_tokens self.max_input_tokens = max_input_tokens self.context_window = context_window self.embedding_max_input_tokens = embedding_max_input_tokens self.embedding_model = embedding_model self.embedding_dim = embedding_dim self.embedding_batch_size = embedding_batch_size self.memory_human = memory_human self.memory_persona = memory_persona self.answer_agent = bool(answer_agent) self.answer_max_steps = int(answer_max_steps) self.memory_types = tuple(dict.fromkeys(memory_types)) if "core" not in self.memory_types or set(self.memory_types) - set(ALL_MEMORY_TYPES): raise ValueError(f"memory_types must contain 'core' and come from {ALL_MEMORY_TYPES}: {memory_types}") if self.answer_agent: self.ANSWERS_INTERNALLY = True self._chat_tool_ids: list = [] self._chat_llm_config = None self._chat_embedding_config = None self._chat_n = 0 self._answer_lookups: list = [] self._prompt_memories: dict | None = None _default = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") self._llm_base_url = internal_base_url or _default self._embedding_base_url = embedding_base_url or _default self._mirix_dir = mirix_dir = Path(working_dir).resolve() / "mirix" if mirix_dir.exists(): shutil.rmtree(mirix_dir) mirix_dir.mkdir(parents=True, exist_ok=True) self._preload_mirix_with_dir(str(mirix_dir)) self._occurred_at: str | None = None self._stage: str = "query" self._absorb_scope: list | None = None self._date2visit: Dict[str, int] = {} self._id2visit: Dict[str, set] = {} self._cur_visit: int | None = None self._visits_memorized = 0 self._update_failures = 0 self._update_error: str | None = None self._client = None self._meta_agent = None self._instrument_clients() self._set_retrieval_limits() self._start_engine() logging.getLogger("Mirix").setLevel(logging.WARNING) @staticmethod def _preload_mirix_with_dir(mirix_dir: str) -> None: """Load MIRIX with its data directory set to the run's working directory.""" import importlib.util import types if "mirix" in sys.modules: import mirix.constants as mc if mc.MIRIX_DIR != mirix_dir: raise RuntimeError("mirix already imported with a different MIRIX_DIR; " "one process per (method, patient) run is required") return if MIRIX_PATH not in sys.path: sys.path.insert(0, MIRIX_PATH) pkg = types.ModuleType("mirix") pkg.__path__ = [os.path.join(MIRIX_PATH, "mirix")] pkg.__version__ = "0.1.0" sys.modules["mirix"] = pkg spec = importlib.util.spec_from_file_location( "mirix.constants", os.path.join(MIRIX_PATH, "mirix", "constants.py")) mc = importlib.util.module_from_spec(spec) sys.modules["mirix.constants"] = mc spec.loader.exec_module(mc) mc.MIRIX_DIR = mirix_dir mc.MIRIX_DIR_TOOL_SANDBOX = os.path.join(mirix_dir, "tool_sandbox_dir") from mirix.schemas.embedding_config import EmbeddingConfig from mirix.schemas.llm_config import LLMConfig pkg.EmbeddingConfig = EmbeddingConfig pkg.LLMConfig = LLMConfig def _loop(self): try: loop = asyncio.get_event_loop() if loop.is_closed(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) return loop except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) return loop def _await(self, coro): return self._loop().run_until_complete(coro) def _instrument_clients(self) -> None: import openai as _openai usage_log = USAGE.log model_name = self.model import mirix.llm_api.openai_client as oc _RealAsync = _openai.AsyncOpenAI class _TrackedCompletions: def __init__(self, real): self._real = real async def create(self, **kw): resp = await self._real.chat.completions.create(**kw) u = getattr(resp, "usage", None) if u is not None: usage_log(kw.get("model", model_name), "chat", int(getattr(u, "prompt_tokens", 0) or 0), int(getattr(u, "completion_tokens", 0) or 0)) return resp class _TrackedChat: def __init__(self, real): self.completions = _TrackedCompletions(real) class _TrackedAsyncOpenAI: def __init__(self, **kw): kw.setdefault("max_retries", MAX_ATTEMPTS - 1) kw.setdefault("timeout", http_timeout()) self._real = _RealAsync(**kw) self.chat = _TrackedChat(self._real) def __getattr__(self, name): return getattr(self._real, name) oc.AsyncOpenAI = _TrackedAsyncOpenAI import llama_index.embeddings.openai as _lio _RealE = _lio.OpenAIEmbedding emb_model = self.embedding_model emb_batch = self.embedding_batch_size class _EmbWithModel(_RealE): def __init__(self, **kw): kw.setdefault("model", emb_model) kw.setdefault("max_retries", MAX_ATTEMPTS - 1) kw.setdefault("timeout", HTTP_TIMEOUT_S) if emb_batch: kw.setdefault("embed_batch_size", int(emb_batch)) super().__init__(**kw) def get_text_embedding(self, text): _log_embed_est(emb_model, [text]) return super().get_text_embedding(text) def get_text_embedding_batch(self, texts, **kw2): _log_embed_est(emb_model, list(texts)) return super().get_text_embedding_batch(texts, **kw2) _lio.OpenAIEmbedding = _EmbWithModel def _set_retrieval_limits(self) -> None: """Use retrieve_num as the number of memories per type in the agent's prompt, per search_in_memory call and per conversation_search page.""" import mirix.agent.agent as magent import mirix.constants as mconst k = int(self.retrieve_num) mconst.MAX_RETRIEVAL_LIMIT_IN_SYSTEM = mconst.RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE = k magent.MAX_RETRIEVAL_LIMIT_IN_SYSTEM = k def _install_engine_udf(self) -> None: from sqlalchemy import event import mirix.server.server as msrv from mirix.orm.sqlite_functions import cosine_distance eng = getattr(msrv, "engine", None) sync = getattr(eng, "sync_engine", None) if sync is None: return @event.listens_for(sync, "connect") def _reg(dbapi_conn, _rec): try: res = dbapi_conn.create_function("cosine_distance", 2, cosine_distance) if asyncio.iscoroutine(res): dbapi_conn.await_(res) except Exception as e: logger.warning(f"[MIRIX] cosine_distance registration failed: {e}") def _start_engine(self) -> None: from mirix.local_client.local_client import LocalClient from mirix.schemas.llm_config import LLMConfig from mirix.schemas.agent import CreateMetaAgent from mirix.server.server import ensure_tables_created self._install_engine_udf() self._await(ensure_tables_created()) llm_config = LLMConfig(model=self.model, model_endpoint_type="openai", model_endpoint=self._llm_base_url, context_window=self.context_window, temperature=self.temperature, max_tokens=self.max_tokens) embedding_config = self._embedding_config() async def create_client(): return await LocalClient.create(debug=False, default_llm_config=llm_config, default_embedding_config=embedding_config) self._client = self._await(create_client()) orig_send = self._client.server.send_messages agent_self = self async def send_messages_with_ts(*a, **kw): if kw.get("occurred_at") is None and agent_self._occurred_at is not None: kw["occurred_at"] = agent_self._occurred_at return await orig_send(*a, **kw) self._client.server.send_messages = send_messages_with_ts self._hook_memory_writes() self._hook_replace_tools() async def setup_meta_agent(): return await self._client.create_meta_agent( request=CreateMetaAgent( llm_config=llm_config, embedding_config=embedding_config, agents=[ {"core_memory_agent": {"blocks": [ {"label": "human", "value": self.memory_human}, {"label": "persona", "value": self.memory_persona}, ]}}, *(f"{t}_memory_agent" for t in ("resource", "semantic", "episodic", "procedural", "knowledge_vault") if t in self.memory_types), "meta_memory_agent", "reflexion_agent", "background_agent", ])) self._meta_agent = self._await(setup_meta_agent()) logger.info(f"[MIRIX] MetaAgent ready: {self._meta_agent.id}") self._ensure_core_blocks() if self.answer_agent: self._setup_answer_agent(llm_config, embedding_config) def _ensure_core_blocks(self) -> None: """Create the human and persona blocks for the runtime user.""" async def mk(): for label, value in (("human", self.memory_human), ("persona", self.memory_persona)): await self._client.create_block(label=label, value=value) try: self._await(mk()) logger.info("[MIRIX] core memory blocks (human, persona) created for the " "runtime user") except Exception as e: logger.error("[MIRIX] core block creation failed (%s); core_memory_append " "will raise on every call", e) def _embedding_config(self): from mirix.schemas.embedding_config import EmbeddingConfig return EmbeddingConfig(embedding_model=self.embedding_model, embedding_endpoint_type="openai", embedding_endpoint=self._embedding_base_url, embedding_dim=self.embedding_dim, embedding_chunk_size=self.embedding_max_input_tokens) _WRITE_MANAGERS = ( ("mirix.services.episodic_memory_manager", "EpisodicMemoryManager"), ("mirix.services.semantic_memory_manager", "SemanticMemoryManager"), ("mirix.services.knowledge_vault_manager", "KnowledgeVaultManager"), ("mirix.services.procedural_memory_manager", "ProceduralMemoryManager"), ("mirix.services.resource_memory_manager", "ResourceMemoryManager"), ) def _hook_memory_writes(self) -> None: """Wrap the create, update, lookup and delete methods of every memory manager to record which visit wrote each memory and what the agent searched.""" import importlib hooked = [] for mod_name, cls_name in self._WRITE_MANAGERS: try: cls = getattr(importlib.import_module(mod_name), cls_name) except Exception as e: logger.warning("[MIRIX] visit-tracking hook: cannot load %s.%s (%s)", mod_name, cls_name, e) continue if getattr(cls, "__memebench_hooked__", False): hooked.append(cls_name) continue n = 0 for fn_name in [x for x in dir(cls) if x.startswith("create_")]: orig = getattr(cls, fn_name, None) if callable(orig): setattr(cls, fn_name, self._wrap_create(orig)) n += 1 for fn_name in ("update_event", "update_item"): orig = getattr(cls, fn_name, None) if callable(orig): setattr(cls, fn_name, self._wrap_create(orig)) n += 1 lfn = self._LOOKUP_METHODS.get(cls_name) if lfn: orig = getattr(cls, lfn, None) if callable(orig): setattr(cls, lfn, self._wrap_lookup(orig)) n += 1 dfn = self._DELETE_METHODS.get(cls_name) if dfn: orig = getattr(cls, dfn, None) if callable(orig): setattr(cls, dfn, self._wrap_delete(orig)) n += 1 cls.__memebench_hooked__ = True hooked.append(f"{cls_name}({n})") logger.info("[MIRIX] visit-tracking hook installed on %s", ", ".join(hooked)) if not hooked: logger.error("[MIRIX] visit-tracking hook installed on no memory manager; " "the visits of retrieved memories cannot be recorded") _REPLACE_TOOLS = ("episodic_memory_replace", "semantic_memory_update", "resource_memory_update", "procedural_memory_update", "knowledge_vault_update") _DELETE_METHODS = { "EpisodicMemoryManager": "delete_event_by_id", "SemanticMemoryManager": "delete_semantic_item_by_id", "KnowledgeVaultManager": "delete_knowledge_by_id", "ProceduralMemoryManager": "delete_procedure_by_id", "ResourceMemoryManager": "delete_resource_by_id", } _STORE_OF = { "list_episodic_memory": "episodic", "list_semantic_items": "semantic", "list_knowledge": "knowledge_vault", "list_procedures": "procedural", "list_resources": "resource", "create_episodic_memory": "episodic", "update_event": "episodic", "delete_event_by_id": "episodic", "episodic_memory_replace": "episodic", "delete_semantic_item_by_id": "semantic", "semantic_memory_update": "semantic", "delete_knowledge_by_id": "knowledge_vault", "knowledge_vault_update": "knowledge_vault", "delete_procedure_by_id": "procedural", "procedural_memory_update": "procedural", "delete_resource_by_id": "resource", "resource_memory_update": "resource", } _PREFIX_STORE = {"ep": "episodic", "sem": "semantic", "kv": "knowledge_vault", "res": "resource", "proc": "procedural"} @classmethod def _store_for(cls, fn: str, memory_id: str = "") -> str | None: st = cls._STORE_OF.get(fn) if st: return st return cls._PREFIX_STORE.get(str(memory_id).split("_")[0]) if memory_id else None _LOOKUP_METHODS = { "EpisodicMemoryManager": "list_episodic_memory", "SemanticMemoryManager": "list_semantic_items", "KnowledgeVaultManager": "list_knowledge", "ProceduralMemoryManager": "list_procedures", "ResourceMemoryManager": "list_resources", } def _wrap_lookup(self, orig): """Record each memory search: its results while answering (for retrieved_visits), a trace entry while writing.""" agent_self = self async def wrapped(mgr_self, *a, **kw): out = await orig(mgr_self, *a, **kw) try: if agent_self._stage == "answer": rows = out if isinstance(out, (list, tuple)) else [] agent_self._answer_lookups.append({ "fn": orig.__name__, "store": agent_self._store_for(orig.__name__), "query": (str(kw.get("query"))[:300] if kw.get("query") else None), "search_method": kw.get("search_method"), "n_found": len(rows), "ids": [str(getattr(o, "id", "") or "") for o in rows[:50]]}) elif agent_self._stage == "write": rows = out if isinstance(out, (list, tuple)) else [] ast = kw.get("agent_state") agent_self._trace( "write_lookup", fn=orig.__name__, store=agent_self._store_for(orig.__name__), agent=getattr(ast, "name", None), query=(str(kw.get("query"))[:300] if kw.get("query") else None), search_field=kw.get("search_field") or None, visit=agent_self._cur_visit, embedded=kw.get("embedded_text") is not None, search_method=kw.get("search_method"), n_found=len(rows), candidates=[ {"id": str(getattr(o, "id", "") or ""), "visits": sorted(agent_self._id2visit.get( str(getattr(o, "id", "") or ""), ())), "text": str(getattr(o, "summary", None) or getattr(o, "name", None) or getattr(o, "caption", ""))[:200]} for o in rows[:20]], ) except Exception: pass return out return wrapped def _wrap_delete(self, orig): """Record what a delete removed, then drop it from the live map.""" import inspect try: params = list(inspect.signature(orig).parameters) id_param = params[1] if len(params) > 1 else None except Exception: id_param = None agent_self = self async def wrapped(mgr_self, *a, **kw): mid = "" try: mid = str(a[0] if a else kw.get(id_param, "") or "") except Exception: pass visits = sorted(agent_self._id2visit.get(mid, ())) if mid else [] try: agent_self._trace("write_op", op="delete", fn=orig.__name__, store=agent_self._store_for(orig.__name__, mid), memory_id=mid, visit=agent_self._cur_visit, visits_before=visits) except Exception: pass out = await orig(mgr_self, *a, **kw) try: if agent_self._absorb_scope is not None and mid: agent_self._absorb_scope.append({"memory_id": mid, "visits_before": visits}) agent_self._id2visit.pop(mid, None) except Exception: pass return out return wrapped def _hook_replace_tools(self) -> None: """For the tools that replace memories (delete, then insert), give the new memories the visits of the memories they replace.""" import functools, importlib try: mt = importlib.import_module("mirix.functions.function_sets.memory_tools") except Exception as e: logger.error("[MIRIX] replace hook: cannot load memory_tools (%s)", e) return agent_self = self hooked = [] for name in self._REPLACE_TOOLS: orig = getattr(mt, name, None) if not callable(orig) or getattr(orig, "__memebench_hooked__", False): continue def make(orig=orig, name=name): @functools.wraps(orig) async def wrapped(*a, **kw): prev = agent_self._absorb_scope agent_self._absorb_scope = [] before = set(agent_self._id2visit) try: out = await orig(*a, **kw) finally: removed = agent_self._absorb_scope agent_self._absorb_scope = prev absorbed = {v for r in removed for v in r["visits_before"]} created = sorted(set(agent_self._id2visit) - before) try: for nid in created: agent_self._id2visit[nid] |= absorbed if prev is not None: prev.extend(removed) except Exception: pass try: agent_self._trace("write_op", op="replace", fn=name, store=agent_self._store_for(name), visit=agent_self._cur_visit, removed=removed, created=created, absorbed_visits=sorted(absorbed)) except Exception: pass return out wrapped.__memebench_hooked__ = True return wrapped setattr(mt, name, make()) hooked.append(name) logger.info("[MIRIX] replace hook installed on %s", ", ".join(hooked) or "none") self._hook_router(mt) def _hook_router(self, mt) -> None: """Record which stores the meta agent routes to, dropping the stores outside `memory_types`, and count the routings that failed (MIRIX hands their error to the meta agent as the tool result).""" import functools agent_self = self orig = getattr(mt, "trigger_memory_update", None) if not callable(orig) or getattr(orig, "__memebench_hooked__", False): logger.error("[MIRIX] router hook: trigger_memory_update not hookable") return @functools.wraps(orig) async def wrapped(*a, **kw): types = kw.get("memory_types") if types is None and len(a) > 2: types = a[2] if isinstance(types, (list, tuple)): types = list(types) kept = [t for t in types if t in agent_self.memory_types] if "memory_types" in kw: kw["memory_types"] = kept elif len(a) > 2: a = a[:2] + (kept,) + a[3:] else: types, kept = [], [] try: agent_self._trace("route", fn="trigger_memory_update", visit=agent_self._cur_visit, memory_types=types, dispatched=list(dict.fromkeys(kept))) except Exception: pass try: return await orig(*a, **kw) except Exception as e: cause = e.__cause__ agent_self._update_failures += 1 agent_self._update_error = (f"{type(e).__name__}: {e}" + (f" ({type(cause).__name__}: {cause})" if cause else ""))[:300] raise wrapped.__memebench_hooked__ = True setattr(mt, "trigger_memory_update", wrapped) logger.info("[MIRIX] router hook installed on trigger_memory_update") def _wrap_create(self, orig): agent_self = self async def wrapped(mgr_self, *a, **kw): out = await orig(mgr_self, *a, **kw) try: is_upd = orig.__name__ in ("update_event", "update_item") for obj in (out if isinstance(out, (list, tuple)) else [out]): oid = str(getattr(obj, "id", "") or "") if not oid: continue agent_self._trace("write_op", op="update" if is_upd else "insert", fn=orig.__name__, memory_id=oid, store=agent_self._store_for(orig.__name__, oid), visit=agent_self._cur_visit, visits_before=sorted(agent_self._id2visit.get(oid, ())), text=str(getattr(obj, "summary", None) or getattr(obj, "details", None) or getattr(obj, "name", ""))[:300]) except Exception: pass try: v = agent_self._cur_visit if v is not None: for obj in (out if isinstance(out, (list, tuple)) else [out]): oid = getattr(obj, "id", None) if oid: agent_self._id2visit.setdefault(str(oid), set()).add(int(v)) except Exception: pass return out return wrapped def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult: """Send the visit to the meta agent. Not successful when send_message raises or a routing to the memory managers failed during the visit.""" t0 = time.time() failures_before = self._update_failures self._cur_visit = visit_number self._date2visit.setdefault(date, visit_number) bounded = self.cut_tokens(text, self.max_input_tokens) chunks = self.split_tokens(bounded, self.memorize_chunk_tokens, self.memorize_chunk_overlap_tokens) self._occurred_at = f"{date}T00:00:00" self._stage = "write" try: for chunk in chunks: self._await(self._client.send_message( agent_id=self._meta_agent.id, role="user", message=chunk)) except Exception as e: logger.error(f"[MIRIX] send_message failed on v{visit_number}: {e}") return MemoryBuildResult(success=False, method="mirix", action="send_message", visit_number=visit_number, chunk_count=self._visits_memorized, time_cost=time.time() - t0, extra={"error": str(e)[:200]}) finally: self._occurred_at = None self._stage = "query" self._visits_memorized += 1 extra = {"chunks": len(chunks)} update_failures = self._update_failures - failures_before if update_failures: logger.warning("[MIRIX] v%s: %d memory update(s) failed; this visit's content may be " "incomplete (%s)", visit_number, update_failures, self._update_error) extra.update(memory_update_failures=update_failures, error=self._update_error) return MemoryBuildResult(success=update_failures == 0, method="mirix", action="send_message", visit_number=visit_number, chunk_count=self._visits_memorized, added_chunks=len(chunks), time_cost=time.time() - t0, extra=extra) def _search_memory(self, query: str) -> List[Dict[str, Any]]: per_type_limit = max(1, (self.retrieve_num + 4) // 5) async def do_retrieve(): return await self._client.retrieve_memory( agent_id=self._meta_agent.id, query=query, memory_type="all", search_method="embedding", limit=per_type_limit) results = self._await(do_retrieve()) try: self._trace("query_retrieve", question=query, per_type_limit=per_type_limit, results=[{"memory_type": e.get("memory_type"), "entry": {k2: v2 for k2, v2 in e.items() if k2 != "embedding"}} for e in (results or {}).get("results", [])]) except Exception: pass memories = [] for entry in (results or {}).get("results", []): mt = entry.get("memory_type", "unknown") parts = [] if mt == "episodic": parts += [entry.get("summary"), entry.get("details")] elif mt == "semantic": if entry.get("name"): parts.append(f"[{entry['name']}]") parts += [entry.get("summary"), entry.get("details")] elif mt == "procedural": parts.append(entry.get("summary")) steps = entry.get("steps") if isinstance(steps, list): parts.append(" -> ".join(steps[:5])) elif steps: parts.append(str(steps)) elif mt == "resource": parts += [entry.get("summary"), (entry.get("content") or "")[:500] or None] elif mt == "knowledge_vault": parts += [entry.get("caption"), entry.get("secret_value")] else: for f in ("summary", "details", "content", "caption", "name"): if entry.get(f): parts.append(str(entry[f])) break content = " | ".join(p for p in parts if p) or str(entry) memories.append({"memory": self.cut_tokens(content, self.query_memory_item_tokens), "type": mt, "occurred_at": str(entry.get("occurred_at") or entry.get("timestamp") or "")}) return memories[:self.retrieve_num] def _setup_answer_agent(self, llm_config, embedding_config) -> None: """Register `send_message` and keep the configs a per-question chat agent needs.""" from mirix.schemas.tool import Tool as PydanticTool from mirix.schemas.enums import ToolType import mirix.constants as mc mc.EXTRAS_TOOLS[:] = [] async def reg(): return await self._client.server.tool_manager.create_or_update_tool( PydanticTool(name="send_message", tags=[ToolType.MIRIX_CORE.value], source_type="python", tool_type=ToolType.MIRIX_CORE), actor=self._client.client) tool = self._await(reg()) self._chat_tool_ids = [tool.id] self._chat_llm_config = llm_config self._chat_embedding_config = embedding_config self._hook_prompt_memories() logger.info("[MIRIX] answer agent enabled: send_message registered (%s); " "%d tool steps per question", tool.id, self.answer_max_steps) def _new_chat_agent(self): """A fresh chat agent for one question: it sees every memory written for this user and no earlier answer.""" from mirix.schemas.agent import AgentType from mirix.schemas.memory import Memory from mirix.schemas.tool_rule import TerminalToolRule self._chat_n += 1 async def mk(): return await self._client.create_agent( name=f"chat_agent_q{self._chat_n:04d}", agent_type=AgentType.chat_agent, memory=Memory(blocks=[]), llm_config=self._chat_llm_config, embedding_config=self._chat_embedding_config, tool_ids=list(self._chat_tool_ids), tool_rules=[TerminalToolRule(tool_name="send_message")], include_base_tools=True) return self._await(mk()) def _hook_prompt_memories(self) -> None: """Keep the memories the chat agent's system prompt was built from, for the row's memory_preview.""" from mirix.agent import agent as _ag from mirix.schemas.agent import AgentType cls = _ag.Agent if getattr(cls, "__memebench_prompt_hook__", False): return orig_build = cls.build_system_prompt_with_memories agent_self = self async def wrapped(self_, *a, **kw): out = await orig_build(self_, *a, **kw) try: if (agent_self._stage == "answer" and self_.agent_state.is_type(AgentType.chat_agent)): agent_self._prompt_memories = out[1] if isinstance(out, tuple) else None except Exception: pass return out cls.build_system_prompt_with_memories = wrapped cls.__memebench_prompt_hook__ = True @staticmethod def _render_prompt_memories(pm) -> str: """The memory block the chat agent was shown, as text.""" if not isinstance(pm, dict): return "" parts = [] for k, v in pm.items(): if k == "key_words": if v: parts.append(f"[topics] {v}") continue if isinstance(v, str): txt = v.strip() elif isinstance(v, dict): txt = "\n".join(str(x).strip() for kk, x in v.items() if isinstance(x, str) and x.strip()) else: txt = "" if txt: parts.append(f"[{k}]\n{txt}") return "\n\n".join(parts) def _answer_with_chat_agent(self, question: str, as_of: str, display_question: str | None) -> MemoryContext: """One question -> one fresh chat agent -> its send_message is the answer.""" t0 = time.time() raw_q = display_question or question wrapped = ANSWER_NOCTX.format(as_of=as_of, qtext=raw_q) server = self._client.server prev_cap = getattr(server, "max_chaining_steps", None) self._answer_lookups, self._prompt_memories = [], None self._stage = "answer" answer, tool_calls, n_assistant, error = "", [], 0, False try: chat = self._new_chat_agent() server.max_chaining_steps = self.answer_max_steps resp = self._await(self._client.send_message( agent_id=chat.id, role="user", message=wrapped)) sent = [] for m in getattr(resp, "messages", None) or []: mt = getattr(m, "message_type", None) mt = getattr(mt, "value", mt) if mt in ("tool_call_message", "function_call"): tc = getattr(m, "tool_call", None) or getattr(m, "function_call", None) name = getattr(tc, "name", None) args = str(getattr(tc, "arguments", "") or "") tool_calls.append({"name": name, "arguments": args[:300]}) if name == "send_message": try: sent.append(str((json.loads(args) or {}).get("message", ""))) except Exception: sent.append(args) elif mt == "assistant_message": n_assistant += 1 if not answer: c = getattr(m, "content", "") if isinstance(c, list): c = " ".join(str(getattr(x, "text", "") or "") for x in c) answer = str(c or "").strip() if not answer and sent: answer = next((x.strip() for x in sent if x and x.strip()), "") except Exception as e: logger.error(f"[MIRIX] answer agent failed: {e}") error, answer = True, f"[err {str(e)[:60]}]" finally: server.max_chaining_steps = prev_cap self._stage = "query" mem = self._render_prompt_memories(self._prompt_memories) ids = [i for lk in self._answer_lookups for i in lk.get("ids", []) if i] visits = sorted({v for i in ids for v in self._id2visit.get(i, ())}) search_steps = sum(1 for tc in tool_calls if tc.get("name") != "send_message") topics = (self._prompt_memories or {}).get("key_words") \ if isinstance(self._prompt_memories, dict) else None try: self._trace("answer_agent", question=raw_q[:300], topics=topics, tool_calls=tool_calls, search_steps=search_steps, assistant_messages=n_assistant, lookups=self._answer_lookups, answered=bool(answer) and not error, answer_chars=len(answer), retrieved_visits=visits, prompt_memory_chars=len(mem), time_s=round(time.time() - t0, 2)) except Exception: pass return MemoryContext( mem=mem, n_tok=self.count_tokens(mem), retrieved_count=len(set(ids)), retrieved_visits=visits, retrieved_memories=[{"visit": None, "memory": mem[:300]}] if mem else [], context_truncated=False, answer=(_unwrap_answer(answer) if answer and not error else answer), error=error, extra={"method": "mirix", "answer_agent": True, "topics": topics, "tool_calls": tool_calls, "search_steps": search_steps, "assistant_messages": n_assistant, "max_steps": self.answer_max_steps}, ) def dump_memory(self, out_dir): from .base import snapshot_tree if not self._mirix_dir.is_dir(): return None n = snapshot_tree(self._mirix_dir, out_dir / "mirix") (out_dir / "id2visit.json").write_text( json.dumps({k: sorted(v) for k, v in self._id2visit.items()}, ensure_ascii=False), encoding="utf-8") return {"files": n, "anchored_ids": len(self._id2visit)} def build_memory_context(self, question: str, as_of: str = "", display_question: str | None = None) -> MemoryContext: if self.answer_agent: return self._answer_with_chat_agent(question, as_of, display_question) try: memories = self._search_memory(question) except Exception as e: logger.error(f"[MIRIX] retrieval failed: {e}") return MemoryContext(mem="", n_tok=0, retrieved_count=0, error=True, extra={"method": "mirix", "error": str(e)[:200]}) blocks, visits, used = [], [], 0 for m in memories: t = f"[{m['type']}] {m['memory']}" n = self.count_tokens(t) if used + n > self.query_memory_context_tokens: break blocks.append(t) used += n occ = (m.get("occurred_at") or "")[:10] if occ in self._date2visit: visits.append(self._date2visit[occ]) mem = "\n\n".join(f"[Memory {i + 1}]\n{b}" for i, b in enumerate(blocks)) return MemoryContext( mem=mem, n_tok=self.count_tokens(mem), retrieved_count=len(blocks), retrieved_visits=sorted(set(visits)), retrieved_memories=[{"visit": None, "memory": b[:300]} for b in blocks], context_truncated=len(blocks) < len(memories), extra={"method": "mirix", "types": [m["type"] for m in memories]}, )