"""HippoRAG 2 adapter. Visits accumulate as passages. When a question arrives after new visits, the graph is extended over them (OpenIE, embeddings, fact and synonymy edges), so memory at a checkpoint covers exactly the visits up to it and accumulates across the whole stream. If the extension fails, the visits stay pending for the next question and the current question is returned as an error. Answers come from HippoRAG's own reader (retrieve, then qa), with the clinician answering prompt as the question. """ from __future__ import annotations import json import logging import os import re import shutil import time import traceback from dataclasses import dataclass, field from pathlib import Path from typing import Any, Optional from .base import ENC, BaseAgent, MemoryBuildResult, MemoryContext, as_list from .engine_bridge import (EngineChat, EngineEmbedder, add_engine_source, stub_optional_imports, read_openie_docs) from ..llm_client import USAGE, create_llm_client, extract_json_object from ..prompts import ANSWER_NOCTX logger = logging.getLogger(__name__) _VISIT_HEADER_RE = re.compile(r"=== Visit v(\d+) ยท ") _VISIT_START = "=== Visit v" _READER_FRAMING = "Wikipedia Title: \n\n" """Text HippoRAG's reader puts around each passage; charged to every passage when fitting the budget.""" def _graph_nodes(hipporag) -> set: """Node names already in the graph.""" g = hipporag.graph return set(g.vs["name"]) if "name" in g.vs.attribute_names() else set() def _wrap_fact_edges(hipporag) -> None: """Pass only new chunks to fact-edge construction and register entity links for earlier chunks.""" from hipporag.utils.misc_utils import compute_mdhash_id orig_add = hipporag.add_fact_edges def add_fact_edges(chunk_ids, chunk_triples): existing = _graph_nodes(hipporag) new = [(k, t) for k, t in zip(chunk_ids, chunk_triples) if k not in existing] orig_add([k for k, _ in new], [t for _, t in new]) for k, t in zip(chunk_ids, chunk_triples): if k in existing: for tr in t: tr = tuple(tr) for ent in (compute_mdhash_id(content=tr[0], prefix="entity-"), compute_mdhash_id(content=tr[2], prefix="entity-")): hipporag.ent_node_to_chunk_ids[ent] = hipporag.ent_node_to_chunk_ids.get( ent, set()).union({k}) hipporag.add_fact_edges = add_fact_edges def _wrap_synonymy_edges(hipporag) -> None: """Build synonymy edges over all nodes, dropping pairs that already exist between earlier nodes.""" orig_syn = hipporag.add_synonymy_edges def add_synonymy_edges(): existing = _graph_nodes(hipporag) before = dict(hipporag.node_to_node_stats) orig_syn() if existing: stats = hipporag.node_to_node_stats for edge in [e for e, w in stats.items() if (e not in before or before[e] != w) and e[0] in existing and e[1] in existing]: del stats[edge] hipporag.add_synonymy_edges = add_synonymy_edges def _install_ner_array_fallback(hipporag) -> None: """Accept a bare JSON array as a NER response, in addition to {"named_entities": [...]}.""" orig_ner = hipporag.openie.ner def ner(chunk_key, passage): out = orig_ner(chunk_key=chunk_key, passage=passage) if not out.unique_entities and out.response: txt = out.response.strip() if txt.startswith("```"): txt = txt.strip("`").removeprefix("json").strip() try: arr = json.loads(txt) if isinstance(arr, list) and all(isinstance(x, str) for x in arr): out.unique_entities = list(dict.fromkeys(arr)) except Exception: pass return out hipporag.openie.ner = ner def _unwrap_answer(text: str) -> str: """Return the `answer` field if `text` is the JSON object the answering prompt requests.""" s = (text or "").strip() if not s or "{" not in s: return s try: try: obj = json.loads(s, strict=False) except Exception: obj = extract_json_object(s) if isinstance(obj, dict) and "answer" in obj: return str(obj["answer"]).strip() except Exception: pass return s def _visit_passages(texts: list[str]) -> list[str]: """One passage per visit, from its header on; the memorize-template sentence in front of the first visit becomes one more passage (the store keeps it once, because passages are deduplicated).""" passages = [] for text in texts: body = text.strip() if body: at = body.find(_VISIT_START) passages.append(body[at:] if at > 0 else body) lead = texts[0].find(_VISIT_START) if lead > 0: passages.append(texts[0][:lead].strip()) return passages def _fit_passages(passages, budget: int, min_tail: int) -> list[str]: """Passages in rank order while they fit `budget` tokens (each also charged the reader's framing). The first one that does not fit is cut to the room left when that is more than `min_tail` tokens.""" framing = BaseAgent.count_tokens(_READER_FRAMING) kept, used = [], 0 for p in passages: text = p if isinstance(p, str) else str(p) size = BaseAgent.count_tokens(text) + framing if used + size <= budget: kept.append(text) used += size continue room = budget - used - framing if room > min_tail: kept.append(ENC.decode(ENC.encode(text)[:room])) break return kept @dataclass class _Patient: """Per-patient state: visits waiting for the next build and the visits already in the graph.""" texts: list = field(default_factory=list) visits: list = field(default_factory=list) indexed: list = field(default_factory=list) received: int = 0 built: bool = False engine: Any = None class HippoRAGAgent(BaseAgent): MEMORIZE_WITH_LLM = True ANSWERS_INTERNALLY = True def __init__( self, *, model: str | None = None, temperature: float, max_tokens: int, provider: str = "openai", api_key: Optional[str] = None, base_url: Optional[str] = None, internal_base_url: Optional[str] = None, openie_mode: str, is_directed_graph: bool, synonymy_edge_sim_threshold: float, synonymy_edge_topk: int, linking_top_k: int, retrieval_top_k: int, qa_top_k: int, damping: float, passage_node_weight: float, force_index_from_scratch: bool, force_openie_from_scratch: bool, save_openie: bool, embedding_provider: str, embedding_model: str | None = None, embedding_dim: Optional[int] = None, embedding_api_key: Optional[str] = None, embedding_base_url: Optional[str] = None, embedding_batch_size: int | None = None, embedding_max_seq_len: int, max_context_tokens: int, context_reserve_tokens: int, min_truncated_passage_tokens: int, working_dir: Optional[str] = None, **kwargs, ): super().__init__(**kwargs) if embedding_provider not in ("openai", "api"): raise ValueError(f"embedding_provider must be 'openai' or 'api', not {embedding_provider!r}") base_url = base_url or internal_base_url self.openie_mode = openie_mode self.qa_top_k = qa_top_k self.max_context_tokens = max_context_tokens self.context_reserve_tokens = context_reserve_tokens self.min_truncated_passage_tokens = min_truncated_passage_tokens self.working_dir = working_dir or str(Path(__file__).resolve().parents[2] / "outputs" / "_hipporag_store") self._store_cleared = False self._engine_settings = dict( dataset=None, llm_name=model, llm_base_url=base_url or os.environ.get("OPENAI_BASE_URL"), temperature=temperature, max_new_tokens=max_tokens, openie_mode=openie_mode, is_directed_graph=is_directed_graph, synonymy_edge_sim_threshold=synonymy_edge_sim_threshold, synonymy_edge_topk=synonymy_edge_topk, linking_top_k=linking_top_k, retrieval_top_k=retrieval_top_k, qa_top_k=qa_top_k, damping=damping, passage_node_weight=passage_node_weight, force_index_from_scratch=force_index_from_scratch, force_openie_from_scratch=force_openie_from_scratch, save_openie=save_openie, embedding_model_name=embedding_model, embedding_batch_size=embedding_batch_size, embedding_max_seq_len=embedding_max_seq_len, embedding_return_as_normalized=True, ) client = create_llm_client(provider=provider, model=model, temperature=temperature, max_tokens=max_tokens, api_key=api_key, base_url=base_url) self._chat = EngineChat(client, temperature=temperature, max_tokens=max_tokens) self._embedder = EngineEmbedder(model=embedding_model, batch_size=embedding_batch_size, dim=embedding_dim, api_key=embedding_api_key, base_url=embedding_base_url) self._patients: dict[Any, _Patient] = {} self._last_build: dict = {} add_engine_source("HippoRAG") stub_optional_imports() def _patient(self) -> tuple[Any, _Patient]: key = self._context_id if self._context_id is not None else 0 if key not in self._patients: self._patients[key] = _Patient() return key, self._patients[key] def _store_dir(self, key) -> str: """The patient's engine directory; at the first use in this run an earlier store is removed.""" path = os.path.join(self.working_dir, f"context_{key}") if not self._store_cleared: self._store_cleared = True if os.path.isdir(path): shutil.rmtree(path) logger.info("[HippoRAG] removed the store left at %s", path) os.makedirs(path, exist_ok=True) return path def _engine(self, key, state: _Patient): """The patient's HippoRAG instance, created on first use with the evaluation's chat and embedding clients.""" if state.engine is not None: return state.engine from hipporag.HippoRAG import HippoRAG from hipporag.utils.config_utils import BaseConfig config = BaseConfig(save_dir=self._store_dir(key), **self._engine_settings) engine = HippoRAG(global_config=config, embedding_model=self._embedder) for role in ("llm_model", "qa_llm", "extraction_llm"): setattr(engine, role, self._chat) _wrap_fact_edges(engine) _wrap_synonymy_edges(engine) if getattr(engine, "openie", None) is not None: engine.openie.llm_model = self._chat _install_ner_array_fallback(engine) if getattr(engine, "rerank_filter", None) is not None: engine.rerank_filter.llm_infer_fn = self._chat.infer engine.embedding_model = self._embedder for store in ("chunk_embedding_store", "entity_embedding_store", "fact_embedding_store"): if getattr(engine, store, None) is not None: getattr(engine, store).embedding_model = self._embedder logger.info("[HippoRAG] engine ready for %s (openie_mode=%s, store=%s)", key, self.openie_mode, config.save_dir) state.engine = engine return engine def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult: """Accumulate one visit; the graph is extended at the next question.""" USAGE.set_phase("memorize") t0 = time.time() key, state = self._patient() state.texts.append(text) state.visits.append(visit_number) state.received += 1 return MemoryBuildResult( success=True, method="hipporag", action="accumulate", visit_number=visit_number, chunk_count=state.received, added_chunks=1, time_cost=time.time() - t0, extra={"context_id": key, "session_number": state.received, "pending_sessions": len(state.texts), "total_accumulated_chars": sum(map(len, state.texts))}) def dump_memory(self, out_dir): """Per-checkpoint export: the indexed passages and a graph size summary.""" from .base import snapshot_jsonl_gz, snapshot_tree stats = {} for cid, state in self._patients.items(): hr = state.engine if hr is None: continue try: src = os.path.join(self.working_dir, f"context_{cid}") if os.path.isdir(src): stats[f"store_files_{cid}"] = snapshot_tree(src, out_dir / f"store_{cid}") except Exception as e: logger.error("[HippoRAG] store snapshot failed (run continues): %s", e) rows = [] store = getattr(hr, "chunk_embedding_store", None) if store is not None: try: mapping = getattr(store, "hash_id_to_row", None) if mapping is None and hasattr(store, "get_all_ids"): ids = list(store.get_all_ids()) mapping = store.get_rows(ids) if ids else {} for hid, row in (mapping or {}).items(): text = str((row or {}).get("content", "")) m = _VISIT_HEADER_RE.findall(text) rows.append({"hash_id": str(hid), "source_visit": int(m[0]) if m else None, "content": text[:300]}) except Exception as e: logger.error("[HippoRAG] passage export failed (run continues): %s", e) if not rows: logger.warning("[HippoRAG] passage export is empty for context %s; the " "store snapshot still holds the passages", cid) stats[f"passages_{cid}"] = snapshot_jsonl_gz( out_dir / f"passages_{cid}.jsonl.gz", rows) g = getattr(hr, "graph", None) if g is not None: try: stats[f"graph_{cid}"] = {"vertices": g.vcount(), "edges": g.ecount()} except Exception: pass return stats or None def _extend_graph(self, key, state: _Patient) -> dict: """Index the waiting visits into the patient's graph. Returns the build summary; on failure the waiting visits are kept, so the next question indexes them again, and the summary carries the error.""" t0 = time.time() waiting = len(state.texts) try: engine = self._engine(key, state) passages = _visit_passages(state.texts) expected = waiting + int(state.texts[0].find(_VISIT_START) > 0) engine.index(passages) engine.ready_to_retrieve = False graph = engine.get_graph_info() openie = read_openie_docs(engine.openie_results_path) except Exception as e: logger.error("[HippoRAG] graph build failed: %s", e) traceback.print_exc() return {"success": False, "session_count": waiting, "error": str(e), "time_cost": time.time() - t0} state.indexed.extend(state.visits) state.texts, state.visits = [], [] state.built = True logger.info("[HippoRAG] %d visits indexed as %d passages in %.1fs; graph %s", waiting, len(passages), time.time() - t0, graph) return {"success": True, "session_count": waiting, "num_documents": len(passages), "expected_documents": expected, "graph_info": graph, "openie_count": len(openie) if openie else 0, "time_cost": time.time() - t0} def build_memory_context(self, question: str, as_of: str = "", display_question: str | None = None) -> MemoryContext: """Index pending visits, retrieve, and answer with HippoRAG's reader. When indexing fails the question is returned as an error ("[err graph_build]") instead of being answered from the earlier graph.""" key, state = self._patient() if state.texts: USAGE.set_phase("memorize") build = self._extend_graph(key, state) graph = build.get("graph_info") or {} self._last_build = { "new_documents": build.get("num_documents"), "sessions_indexed": build.get("session_count"), "indexed_visits_total": len(state.indexed), "graph_nodes": graph.get("num_total_nodes"), "graph_triples": graph.get("num_total_triples"), "build_time_s": round(build.get("time_cost", 0.0), 1), "build_ok": build.get("success", False), } if not build.get("success"): return MemoryContext(mem="", n_tok=0, retrieved_count=0, error=True, answer="[err graph_build]", extra={"method": "hipporag", "error": "graph_build", "build_error": str(build.get("error"))[:300], **self._last_build}) if build["num_documents"] != build["expected_documents"]: raise RuntimeError(f"[HippoRAG] {build['num_documents']} documents for " f"{build['session_count']} pending visits") USAGE.set_phase("query") t0 = time.time() if not state.built: logger.warning("[HippoRAG] nothing is indexed for %s", key) return MemoryContext(mem="", n_tok=0, retrieved_count=0, answer="[err no_data_indexed]", error=True, extra={"method": "hipporag", "error": "no_data_indexed"}) prompt = ANSWER_NOCTX.format(as_of=as_of, qtext=display_question or question) answer, shown, failed, cut = "", [], False, False try: answer, shown, cut = self._retrieve_and_read(self._engine(key, state), question, prompt) except Exception as e: logger.error("[HippoRAG] query failed: %s", e) traceback.print_exc() answer, failed = f"[err {str(e)[:60]}]", True mem = "\n\n".join(f"[Memory {i}]\n{p['content']}" for i, p in enumerate(shown, 1)) return MemoryContext( mem=mem, n_tok=self.count_tokens(mem), retrieved_count=len(shown), retrieved_visits=sorted({int(v) for p in shown for v in _VISIT_HEADER_RE.findall(p["content"])}), retrieved_memories=[{"visit": None, "memory": p["content"][:300], "score": p.get("score")} for p in shown], context_truncated=cut, answer=_unwrap_answer(answer), error=failed, extra={"method": "hipporag", "openie_mode": self.openie_mode, "query_time_s": round(time.time() - t0, 2), **self._last_build}, ) def _retrieve_and_read(self, engine, question: str, prompt: str) -> tuple[str, list, bool]: """Retrieve with the question, fit the top passages to the context budget and let the reader answer `prompt`. Returns the answer, the passages it read (with scores) and whether the budget cut passages.""" if not engine.ready_to_retrieve: engine.prepare_retrieval_objects() found = engine.retrieve(queries=[question]) try: first = (found or [None])[0] self._trace("query_retrieve", question=question, docs=as_list(getattr(first, "docs", None))[:20], scores=[float(x) for x in as_list(getattr(first, "doc_scores", None))[:20]]) except Exception: pass if not found: return "", [], False solution = found[0] top = solution.docs[:self.qa_top_k] if solution.docs else [] cut = False if top and self.max_context_tokens > 0: budget = max(self.max_context_tokens - self.count_tokens(prompt) - self.context_reserve_tokens, 0) fitted = _fit_passages(top, budget, self.min_truncated_passage_tokens) cut = len(fitted) < len(top) top = fitted solution.docs = top solution.question = prompt read, _, _ = engine.qa([solution]) solution = read[0] if read else solution answer = solution.answer if hasattr(solution, "answer") else str(solution) shown = [] docs = getattr(solution, "docs", None) if docs: scores = getattr(solution, "doc_scores", None) for rank, doc in enumerate(docs[:self.qa_top_k]): entry = {"rank": rank + 1, "content": doc if isinstance(doc, str) else str(doc)} if scores is not None and len(scores) > rank: entry["score"] = float(scores[rank]) shown.append(entry) return answer, shown, cut