arc53--docsgpt/application/retriever/graph_rag.py

"""GraphRAG local retriever — Personalized PageRank over a per-source graph.

Rephrased query -> entity-name NN seeds -> bounded 1-2-hop fetch -> networkx
Personalized PageRank (IDF-down-weighted hubs) -> chunks ranked by landed PPR
mass -> shared token budget. No LLM call at query time beyond the (optional,
reused) rephrase.

Composes :class:`ClassicRAG` rather than subclassing: PPR doesn't fit the
``_fetch_candidates`` hook, but the composed instance supplies the rephrase, the
token-budget loop, and the per-source fallback when a source has no graph.
"""

from __future__ import annotations

import logging
import math
from typing import Any, Dict, List

import networkx as nx

from application.core.settings import settings
from application.graphrag import graphrag_available
from application.graphrag.store import GraphStore
from application.retriever.base import BaseRetriever
from application.retriever.classic_rag import ClassicRAG
from application.retriever.labels import labels_from_metadata
from application.utils import num_tokens_from_string
from application.vectorstore.base import EmbeddingsSingleton

SEED_NODES = 10
SUBGRAPH_HOPS = 1


def _idf(doc_freq: Any) -> float:
    """Node-specificity weight: rarer entities (low ``doc_freq``) score higher."""
    return 1.0 / math.log(1.0 + max(int(doc_freq or 0), 0) + 1.0)


class GraphRAGRetriever(BaseRetriever):
    """Per-source PPR retriever; falls back to ClassicRAG when a source has no graph."""

    def __init__(
        self,
        source,
        chat_history=None,
        prompt="",
        chunks=2,
        doc_token_limit=50000,
        model_id="docsgpt-local",
        user_api_key=None,
        agent_id=None,
        llm_name=settings.LLM_PROVIDER,
        api_key=settings.API_KEY,
        decoded_token=None,
        model_user_id=None,
        defer_rephrase=False,
        request_id=None,
    ):
        self._classic = ClassicRAG(
            source=source,
            chat_history=chat_history,
            prompt=prompt,
            chunks=chunks,
            doc_token_limit=doc_token_limit,
            model_id=model_id,
            user_api_key=user_api_key,
            agent_id=agent_id,
            llm_name=llm_name,
            api_key=api_key,
            decoded_token=decoded_token,
            model_user_id=model_user_id,
            defer_rephrase=defer_rephrase,
            request_id=request_id,
        )
        self.original_question = self._classic.original_question
        self.chunks = self._classic.chunks
        self.doc_token_limit = doc_token_limit
        self.vectorstores = self._classic.vectorstores
        self.per_source_retrieval = {}

    def _embed_query(self, question: str) -> List[float]:
        embedding = EmbeddingsSingleton.get_instance(
            settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY
        )
        return embedding.embed_query(question)

    def _ppr_scores(self, subgraph, seeds) -> Dict[str, float]:
        """Run Personalized PageRank, then down-weight hub nodes by IDF.

        ``seeds`` maps seed node id -> personalization weight (seed similarity).
        After PPR, each node's mass is scaled by ``1/log(2 + doc_freq)`` so a
        high-degree hub contributes less than a specific entity at equal mass.
        """
        graph = nx.Graph()
        for node in subgraph.get("nodes", []):
            graph.add_node(node["id"], doc_freq=node.get("doc_freq", 0))
        for edge in subgraph.get("edges", []):
            src, dst = edge["src_node_id"], edge["dst_node_id"]
            if src in graph and dst in graph:
                weight = float(edge.get("weight") or 1.0)
                graph.add_edge(src, dst, weight=weight)
        if graph.number_of_nodes() == 0:
            return {}

        personalization = {n: seeds.get(n, 0.0) for n in graph.nodes}
        if not any(personalization.values()):
            personalization = None

        ranks = nx.pagerank(graph, personalization=personalization, weight="weight")
        return {
            node: rank * _idf(graph.nodes[node].get("doc_freq", 0))
            for node, rank in ranks.items()
        }

    def _rank_chunks(self, store, source_id, node_scores) -> List[str]:
        """Score chunks by summed (PPR mass x IDF) of their linked nodes; top candidates.

        Over-fetches beyond ``self.chunks`` so chunks with missing text don't drop
        the final count below the budget; the budget loop caps the real total.
        """
        node_ids = list(node_scores.keys())
        chunk_links = store.get_chunk_ids_for_nodes(source_id, node_ids)
        chunk_scores: Dict[str, float] = {}
        for node_id, chunk_ids in chunk_links.items():
            node_score = node_scores.get(node_id, 0.0)
            for chunk_id in chunk_ids:
                chunk_scores[chunk_id] = chunk_scores.get(chunk_id, 0.0) + node_score
        ranked = sorted(chunk_scores, key=lambda c: chunk_scores[c], reverse=True)
        candidates = max(self.chunks * 2, self.chunks + 5)
        return ranked[: max(1, candidates)]

    def _graph_docs_for_source(self, store, source_id) -> List[Dict[str, Any]]:
        """Local PPR retrieval for one source (caller guarantees it has a graph)."""
        question = self._classic._get_rephrased_question()
        query_embedding = self._embed_query(question)
        seed_rows = store.search_nodes_by_embedding(
            source_id, query_embedding, k=SEED_NODES
        )
        if not seed_rows:
            return []

        seed_ids = [row["id"] for row in seed_rows]
        # Clamp to >= 0: cosine distance can exceed 1 (negative similarity) for
        # some embedding backends, and networkx pagerank produces garbage on
        # negative personalization (and ZeroDivisionError when the weights sum
        # to ~0). All-zero collapses to uniform PPR via the None guard below.
        seeds = {
            row["id"]: max(0.0, 1.0 - float(row.get("distance") or 0.0))
            for row in seed_rows
        }

        subgraph = store.get_subgraph(source_id, seed_ids, hops=SUBGRAPH_HOPS)
        node_scores = self._ppr_scores(subgraph, seeds)
        if not node_scores:
            return []

        chunk_ids = self._rank_chunks(store, source_id, node_scores)
        chunk_data = store.get_chunk_texts(source_id, chunk_ids)

        docs: List[Dict[str, Any]] = []
        token_budget = max(int(self.doc_token_limit * 0.9), 100)
        cumulative_tokens = 0
        for chunk_id in chunk_ids:
            if len(docs) >= self.chunks:
                break
            chunk = chunk_data.get(chunk_id)
            text = chunk.get("text") if chunk else None
            if not text:
                continue
            labels = labels_from_metadata(chunk.get("metadata"), text, source_id)
            doc_tokens = num_tokens_from_string(f"{labels['filename']}\n{text}")
            if cumulative_tokens + doc_tokens >= token_budget:
                break
            docs.append({"text": text, **labels})
            cumulative_tokens += doc_tokens
        return docs

    def _classic_for_source(self, source_id) -> List[Dict[str, Any]]:
        """Reuse the composed ClassicRAG to retrieve one source's chunks."""
        original = self._classic.vectorstores
        original_overrides = self._classic.per_source_retrieval
        try:
            self._classic.vectorstores = [source_id]
            self._classic.per_source_retrieval = {
                k: v for k, v in self.per_source_retrieval.items() if k == source_id
            }
            return self._classic._get_data()
        finally:
            self._classic.vectorstores = original
            self._classic.per_source_retrieval = original_overrides

    def _get_data(self) -> List[Dict[str, Any]]:
        if not self.vectorstores:
            return []

        store = None
        if graphrag_available():
            try:
                store = GraphStore()
            except Exception as e:
                logging.error(f"GraphRAG store unavailable, falling back: {e}")
                store = None

        all_docs: List[Dict[str, Any]] = []
        for source_id in self.vectorstores:
            if not source_id:
                continue
            has_graph = False
            if store is not None:
                try:
                    has_graph = store.count_nodes(source_id) > 0
                except Exception as e:
                    logging.error(f"GraphRAG count_nodes failed for {source_id}: {e}")
                    has_graph = False

            if not has_graph:
                all_docs.extend(self._classic_for_source(source_id))
                continue

            try:
                all_docs.extend(self._graph_docs_for_source(store, source_id))
            except Exception as e:
                logging.error(
                    f"GraphRAG retrieval failed for {source_id}, falling back: {e}",
                    exc_info=True,
                )
                all_docs.extend(self._classic_for_source(source_id))
        return all_docs

    def search(self, query: str = "") -> List[Dict[str, Any]]:
        if query:
            self.original_question = query
            self._classic.original_question = query
            self._classic._rephrased_question = None
            self._classic.question = self._classic._rephrase_query()
            self._classic._rephrased_question = self._classic.question
        return self._get_data()