File size: 42,946 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 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 | """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]},
)
|