| """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]}, |
| ) |
|
|
