"""GraphRAG retriever for the GraphRAG-vs-classical-RAG experiment. This is the second retrieval backend behind ``retrieve_tutor_context`` (the default is the hybrid ``LocalChromaRetriever`` in ``chroma_rag.py``). It reads a prebuilt Microsoft GraphRAG index (entities, relationships, text units, and community reports) produced offline by ``graphrag index`` over our corpus, and exposes the SAME interface as ``LocalChromaRetriever``:: retriever.search(query, allowed_sources=..., token_budget=...) -> list[SearchResult] so it drops straight into the agent tool and grades with the existing eval metrics (the matches it returns carry the real ``source``/``url`` recovered from ``corpus_manifest`` via each text unit's ``document_id``). Design constraints for the experiment (see evals.md / branch experiment/graphrag-vs-rag): * **Context provider only, no generation.** GraphRAG normally runs an LLM to synthesize an answer in local/global search. We must NOT do that: the eval's only generation model is the agent's 3.5 Flash. So this assembles the GraphRAG *context* (entity-linked text units + relevant community reports) and returns it as ``SearchResult`` chunks; the agent writes the answer. * **Same rerank/budget as classical RAG** so the comparison is fair: candidate text units are reranked with Cohere and capped by the same token budget. * The graph index is **local-only** (``data/graphrag/output``); it is never part of the HF vector-db bundle, so this backend is unavailable unless the index has been built on the machine. The index is built with Gemini 2.5 Flash; this module only reads its artifacts. """ from __future__ import annotations import logging import os from functools import lru_cache from pathlib import Path from typing import Any from .chroma_rag import ( DEFAULT_CONTEXT_TOKEN_BUDGET, DEFAULT_RERANK_MODEL, RERANK_SCORE_FLOOR, SearchResult, get_token_encoding, rerank_results, ) logger = logging.getLogger(__name__) GRAPHRAG_DIR = os.getenv("AI_TUTOR_GRAPHRAG_DIR", "data/graphrag") GRAPHRAG_OUTPUT_DIR = f"{GRAPHRAG_DIR}/output" GRAPHRAG_LANCEDB_DIR = f"{GRAPHRAG_OUTPUT_DIR}/lancedb" # Embedding model used at INDEX time (settings.yaml embedding_models). The query # must be embedded with the same model for entity similarity to be meaningful. GRAPHRAG_EMBED_MODEL_LITELLM = "openai/text-embedding-3-small" # Synthetic source key for community-report context (graph-level synthesis with # no single backing document). It is intentionally NOT a real corpus source, so # recall@source can never falsely credit it. GRAPHRAG_COMMUNITY_SOURCE = "graphrag_community" # Local-search-ish knobs. DEFAULT_ENTITY_TOP_K = 20 DEFAULT_COMMUNITY_TOP_K = 3 DEFAULT_RERANK_TOP_K = 5 # Candidate pool for query-relevant community-report selection: only the top # reports by the index's static rank are reranked against the query, so the # per-query cost stays bounded (one rerank call over at most this many docs). COMMUNITY_RERANK_CANDIDATES = 30 class GraphRAGIndexNotBuilt(RuntimeError): """Raised when the GraphRAG output artifacts are missing. Build them with:: uv run -m data.scraping_scripts.graphrag_prep_input uv run --env-file .env graphrag index --root data/graphrag """ def graphrag_index_exists(output_dir: str = GRAPHRAG_OUTPUT_DIR) -> bool: out = Path(output_dir) return out.is_dir() and any(out.glob("entities.parquet")) def _load_manifest_lookup(manifest_path: str) -> dict[str, dict[str, Any]]: """doc_id -> {source, url, title} from corpus_manifest.jsonl.""" import json lookup: dict[str, dict[str, Any]] = {} if not os.path.exists(manifest_path): return lookup with open(manifest_path, encoding="utf-8") as handle: for line in handle: line = line.strip() if not line: continue record = json.loads(line) doc_id = record.get("doc_id") if doc_id: lookup[doc_id] = { "source": record.get("source", ""), "url": record.get("url", ""), "title": record.get("title", ""), } return lookup class GraphRAGRetriever: """Local-search-style context assembly over a prebuilt GraphRAG index. NOTE: the exact LanceDB table name and parquet column names are validated against a real ``graphrag index`` run (the smoke index) before this backend is used in an eval; the loaders below discover them defensively. """ def __init__( self, *, cohere_api_key: str, output_dir: str = GRAPHRAG_OUTPUT_DIR, lancedb_dir: str = GRAPHRAG_LANCEDB_DIR, manifest_path: str | None = None, rerank_model: str = DEFAULT_RERANK_MODEL, embed_model_litellm: str = GRAPHRAG_EMBED_MODEL_LITELLM, entity_top_k: int = DEFAULT_ENTITY_TOP_K, community_top_k: int = DEFAULT_COMMUNITY_TOP_K, rerank_top_k: int = DEFAULT_RERANK_TOP_K, token_budget: int = DEFAULT_CONTEXT_TOKEN_BUDGET, answer_model_name: str | None = None, ) -> None: if not graphrag_index_exists(output_dir): raise GraphRAGIndexNotBuilt( f"GraphRAG index not found at {output_dir}. Build it with " "`uv run -m data.scraping_scripts.graphrag_prep_input` then " "`uv run --env-file .env graphrag index --root data/graphrag`." ) import cohere from .config import KB_MANIFEST_PATH self._output_dir = output_dir self._lancedb_dir = lancedb_dir self._rerank_model = rerank_model self._embed_model_litellm = embed_model_litellm self._entity_top_k = entity_top_k self._community_top_k = community_top_k self._rerank_top_k = rerank_top_k self._token_budget = token_budget self._encoding = get_token_encoding(answer_model_name) self._cohere = cohere.ClientV2(api_key=cohere_api_key) self._manifest = _load_manifest_lookup(manifest_path or KB_MANIFEST_PATH) self._load_tables() def _load_tables(self) -> None: """Load GraphRAG output parquets and open the entity-embedding store. Column/table names follow the graphrag 3.1 data model (TextUnit.document_id, Entity.text_unit_ids/community_ids); confirmed against the smoke index output before first eval use. """ import pandas as pd out = Path(self._output_dir) self._entities = pd.read_parquet(out / "entities.parquet") self._text_units = pd.read_parquet(out / "text_units.parquet") reports_path = out / "community_reports.parquet" self._reports = pd.read_parquet(reports_path) if reports_path.exists() else None # entity id -> its text unit ids, and text unit id -> row, both O(1) so a # query only touches the top-k nearest entities (not every entity/chunk). self._entity_text_units = { str(eid): list(tus) if tus is not None else [] for eid, tus in zip(self._entities["id"], self._entities["text_unit_ids"]) } self._text_unit_by_id = { str(row["id"]): row for _, row in self._text_units.iterrows() } self._entity_table = self._open_entity_vector_table() def _open_entity_vector_table(self): """Open the LanceDB table holding entity-description embeddings. graphrag names it like ``default-entity-description``; discover it by pattern so a version bump doesn't silently break retrieval. """ try: import lancedb db = lancedb.connect(self._lancedb_dir) names = list(db.table_names()) except Exception as exc: # pragma: no cover - exercised post-index logger.warning("GraphRAG LanceDB open failed (%s); entity search off.", exc) return None preferred = [n for n in names if "entity" in n and "description" in n] chosen = preferred or [n for n in names if "entity" in n] or names if not chosen: logger.warning("No GraphRAG LanceDB tables found in %s.", self._lancedb_dir) return None return db.open_table(chosen[0]) def _embed_query(self, query: str) -> list[float]: import litellm response = litellm.embedding(model=self._embed_model_litellm, input=[query]) return list(response.data[0]["embedding"]) def search( self, query: str, *, allowed_sources: list[str] | None = None, token_budget: int | None = None, ) -> list[SearchResult]: """Assemble GraphRAG context as reranked SearchResult chunks. Steps: embed query -> nearest entities (LanceDB) -> their text units (mapped to real source/url via document_id) + query-relevant community reports -> Cohere rerank -> token budget. Returns [] on any failure so the agent degrades to KB browsing, matching the classical backend's contract. """ try: text_unit_hits = self._entity_linked_text_units(query, allowed_sources) reranked = self._rerank(query, text_unit_hits) community_hits = self._community_context(query) return self._apply_token_budget(reranked + community_hits, token_budget) except Exception as exc: # pragma: no cover - exercised post-index logger.warning("GraphRAG search failed; returning no matches. %s", exc) return [] # --- internals (finalized against the real index) --------------------- def _entity_linked_text_units( self, query: str, allowed_sources: list[str] | None ) -> list[SearchResult]: if self._entity_table is None: return [] query_vec = self._embed_query(query) hits = ( self._entity_table.search(query_vec).limit(self._entity_top_k).to_pandas() ) # Walk the nearest entities in similarity order, collecting the text # units they were extracted from (deduped, order preserved). text_unit_ids: list[str] = [] seen: set[str] = set() for entity_id in hits.get("id", []): for tu_id in self._entity_text_units.get(str(entity_id), []): tu_id = str(tu_id) if tu_id not in seen: seen.add(tu_id) text_unit_ids.append(tu_id) results: list[SearchResult] = [] allowed = set(allowed_sources or []) for tu_id in text_unit_ids: row = self._text_unit_by_id.get(tu_id) if row is None: continue doc_id = str(row.get("document_id") or "") meta = self._manifest.get(doc_id, {}) source = meta.get("source", "") if allowed and source not in allowed: continue text = str(row.get("text") or "") results.append( SearchResult( chunk_id=tu_id, doc_id=doc_id, title=meta.get("title", "") or doc_id, url=meta.get("url", ""), source=source or "unknown", retrieve_doc=False, tokens=int(row.get("n_tokens") or 0), score=0.0, content=text, chunk_content=text, heading_path="", retrieval_method="graphrag", ) ) return results def _rerank(self, query: str, results: list[SearchResult]) -> list[SearchResult]: if not results: return [] return rerank_results( self._cohere, query, results, model=self._rerank_model, top_n=self._rerank_top_k, ) def _community_context(self, query: str) -> list[SearchResult]: """Query-relevant community reports as context-only chunks (synthetic source, no url -> no recall/citation credit, by design). Candidates are the ``COMMUNITY_RERANK_CANDIDATES`` highest-ranked reports (the index's static rank); the ``community_top_k`` most relevant to the query are selected with the same Cohere rerank used for text units. Scores are pinned back to 0.0 so the reports stay context-only chunks (and stay exempt from the rerank-score floor in ``_apply_token_budget``). """ if self._reports is None or self._community_top_k <= 0: return [] reports = self._reports if "rank" in reports.columns: reports = reports.sort_values("rank", ascending=False) candidates: list[SearchResult] = [] for _, row in reports.head(COMMUNITY_RERANK_CANDIDATES).iterrows(): content = str(row.get("full_content") or row.get("summary") or "") if not content: continue candidates.append( SearchResult( chunk_id=f"community:{row.get('community', '')}", doc_id="", title=str(row.get("title") or "Community report"), url="", source=GRAPHRAG_COMMUNITY_SOURCE, retrieve_doc=False, tokens=0, score=0.0, content=content, chunk_content=content, heading_path="", retrieval_method="graphrag_community", ) ) if len(candidates) <= self._community_top_k: return candidates selected = rerank_results( self._cohere, query, candidates, model=self._rerank_model, top_n=self._community_top_k, ) for result in selected: result.score = 0.0 return selected def _apply_token_budget( self, results: list[SearchResult], token_budget: int | None ) -> list[SearchResult]: budget = self._token_budget if token_budget is None else token_budget filtered: list[SearchResult] = [] total = 0 for result in results: # Same low-relevance floor as the classical retriever, so both eval # arms drop weak reranked hits (the module's fairness contract). # Community reports are context-only chunks pinned at score 0.0 and # are exempt, or the floor would drop them all. if ( result.source != GRAPHRAG_COMMUNITY_SOURCE and result.score < RERANK_SCORE_FLOOR ): continue result_tokens = len( self._encoding.encode(result.content, disallowed_special=()) ) if total + result_tokens > budget: # Skip an oversized result instead of cutting off the whole # list; keep filling the budget with later results that fit. continue total += result_tokens filtered.append(result) return filtered @lru_cache(maxsize=1) def _build_graphrag_retriever() -> GraphRAGRetriever: cohere_api_key = os.environ["COHERE_API_KEY"] return GraphRAGRetriever(cohere_api_key=cohere_api_key) def get_graphrag_retriever() -> GraphRAGRetriever: return _build_graphrag_retriever()