| """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 |
|
|
