updated code to include graph rag

backend/app/api/routes_chat.py

@@ -1,47 +1,45 @@
-"""Chat routes for hybrid RAG-based Q&A."""
+"""Chat routes for Graph-RAG."""
 
 from app.core.llm import llm_chat
 from app.core.prompts import build_rag_prompt
 from app.models.api import ChatRequest, ChatResponse, Citation
-from app.retrieval.hybrid import hybrid_search
+from app.retrieval.retrieve import hybrid_graph_search
 from fastapi import APIRouter
 
 router = APIRouter()
 
 
 @router.post("/ask", response_model=ChatResponse)
-def chat_hybrid(request: ChatRequest) -> ChatResponse:
-    """Hybrid RAG Q&A endpoint (vector + BM25)."""
-    # Hybrid retrieval
-    chunks = hybrid_search(request.query, top_k=1)
+def chat(request: ChatRequest) -> ChatResponse:
+    """Graph-augmented RAG endpoint."""
+    results = hybrid_graph_search(request.query, request.top_k)
 
-    if not chunks:
+    if not results:
         return ChatResponse(
             answer="I don't know based on the provided documents.",
             citations=[],
         )
 
-    chunk = chunks[0]
+    context = "\n\n".join(
+        f"(Pages {sc.chunk.page_start}-\
+{sc.chunk.page_end})\n{sc.chunk.text}"
+        for sc in results
+    )
 
-    # Build prompt from ONLY the best chunk
     messages = build_rag_prompt(
-        context=chunk.chunk.text,
+        context=context,
         question=request.query,
     )
 
-    # Ask LLM
     answer = llm_chat(messages=messages)
 
-    # Cite ONLY what was used
     citations = [
         Citation(
-            page_start=chunk.chunk.page_start,
-            page_end=chunk.chunk.page_end,
-            snippet=chunk.chunk.text[:300],
+            page_start=sc.chunk.page_start,
+            page_end=sc.chunk.page_end,
+            snippet=sc.chunk.text[:300],
         )
+        for sc in results
     ]
 
-    return ChatResponse(
-        answer=answer,
-        citations=citations,
-    )
+    return ChatResponse(answer=answer, citations=citations)

backend/app/ingestion/entities.py

@@ -0,0 +1,55 @@
+"""Load entities from chunk texts."""
+
+from typing import List, Set
+
+import spacy
+
+# Load once at module import
+NLP = spacy.load("en_core_web_sm")
+
+# Entity labels we accept.
+# Keep this BROAD on purpose.
+ALLOWED_ENTITY_LABELS = {
+    "PERSON",
+    "ORG",
+    "GPE",
+    "LOC",
+    "PRODUCT",
+    "EVENT",
+    "WORK_OF_ART",
+    "LAW",
+    "LANGUAGE",
+    "NORP",
+    "FAC",
+}
+
+
+def extract_entities(text: str) -> List[str]:
+    """Extract entities from text using spaCy.
+
+    Rules:
+    - Deterministic (no LLM)
+    - Preserve surface form
+    - Deduplicate
+    - Ignore very short / noisy entities
+    """
+    if not text.strip():
+        return []
+
+    doc = NLP(text)
+
+    entities: Set[str] = set()
+
+    for ent in doc.ents:
+        if ent.label_ not in ALLOWED_ENTITY_LABELS:
+            continue
+
+        value = ent.text.strip()
+
+        # Drop trivial junk
+        if len(value) < 3 or len(value.split()) > 5:
+            continue
+
+        entities.add(value)
+
+    return sorted(entities)

backend/app/ingestion/pipeline.py

@@ -5,26 +5,34 @@ from typing import List
 
 from app.ingestion.chunking import chunk_segments
 from app.ingestion.cleaning import clean_text
+from app.ingestion.entities import extract_entities
 from app.ingestion.indexing import index_chunks
 from app.ingestion.pdf_loader import extract_pages
 from app.models.ingestion import Chunk, RawSegment
+from app.retrieval.chunk_registry import register_chunks
+from app.retrieval.graph_utils import index_entities
 from app.retrieval.keyword_index import build_bm25_index
 
 
 def ingest_pdf(file_path: Path, doc_id: str) -> List[Chunk]:
-    """Ingest a PDF document into indexed chunks."""
-    raw_segments = extract_pages(file_path=file_path, doc_id=doc_id)
+    """Ingest a PDF document."""
+    raw_segments = extract_pages(file_path, doc_id)
     cleaned_segments = _clean_segments(raw_segments)
+
     chunks = chunk_segments(cleaned_segments)
 
+    for chunk in chunks:
+        chunk.entities = extract_entities(chunk.text)
+
+    register_chunks(chunks)
+    index_entities(chunks)
     index_chunks(chunks)
     build_bm25_index(chunks)
-
     return chunks
 
 
 def _clean_segments(segments: List[RawSegment]) -> List[RawSegment]:
-    """Apply text cleaning to raw segments."""
+    """Apply text cleaning."""
     return [
         RawSegment(
             doc_id=s.doc_id,

backend/app/models/ingestion.py

@@ -1,10 +1,12 @@
-"""Pydantic models for Ingestion artifacts."""
+"""Pydantic models for ingestion artifacts."""
 
-from pydantic import BaseModel
+from typing import List
+
+from pydantic import BaseModel, Field
 
 
 class RawSegment(BaseModel):
-    """Represents raw page-level text extracted from a PDF."""
+    """Represents raw page-level text extracted from a document."""
 
     doc_id: str
     page: int
@@ -19,3 +21,4 @@ class Chunk(BaseModel):
     page_start: int
     page_end: int
     text: str
+    entities: List[str] = Field(default_factory=list)

backend/app/retrieval/chunk_registry.py

@@ -0,0 +1,31 @@
+"""In-memory chunk registry.
+
+Single source of truth for all ingested chunks.
+Used by graph-based retrieval to map entities back to chunks.
+
+Note:
+- Ephemeral by design (non-persistent)
+- Rebuilt on each ingestion cycle
+"""
+
+from typing import Dict, List
+
+from app.models.ingestion import Chunk
+
+_CHUNKS: Dict[str, Chunk] = {}
+
+
+def register_chunks(chunks: List[Chunk]) -> None:
+    """Register chunks in memory."""
+    for chunk in chunks:
+        _CHUNKS[chunk.chunk_id] = chunk
+
+
+def get_chunks() -> List[Chunk]:
+    """Return all registered chunks."""
+    return list(_CHUNKS.values())
+
+
+def clear_chunks() -> None:
+    """Clear registry (useful for tests)."""
+    _CHUNKS.clear()

backend/app/retrieval/graph_utils.py

@@ -0,0 +1,102 @@
+"""Graph utilities for Graph-RAG.
+
+Responsibilities:
+- Build entity co-occurrence graph
+- Index entity → chunk mappings
+- Extract entities from queries
+- Expand entities via graph traversal
+- Recall chunks via entity relationships
+"""
+
+from collections import defaultdict
+from typing import Dict, Iterable, List, Set
+
+import networkx as nx
+from app.models.ingestion import Chunk
+
+# In-memory entity → chunk index
+_ENTITY_TO_CHUNKS: Dict[str, Set[str]] = defaultdict(set)
+
+
+def index_entities(chunks: List[Chunk]) -> None:
+    """Index entities to chunk IDs.
+
+    Called once during ingestion.
+    """
+    for chunk in chunks:
+        for entity in chunk.entities:
+            _ENTITY_TO_CHUNKS[entity].add(chunk.chunk_id)
+
+
+def build_graph(chunks: List[Chunk]) -> nx.Graph:
+    """Build an entity co-occurrence graph.
+
+    Nodes: entities
+    Edges: co-occurrence within the same chunk
+    """
+    graph = nx.Graph()
+
+    for chunk in chunks:
+        entities = chunk.entities
+
+        for entity in entities:
+            graph.add_node(entity)
+
+        for i, e1 in enumerate(entities):
+            for e2 in entities[i + 1 :]:
+                if graph.has_edge(e1, e2):
+                    graph[e1][e2]["weight"] += 1
+                else:
+                    graph.add_edge(e1, e2, weight=1)
+
+    return graph
+
+
+def extract_query_entities(text: str, nlp) -> Set[str]:
+    """Extract entities from a user query.
+
+    Deterministic (spaCy-based).
+    """
+    if not text.strip():
+        return set()
+
+    doc = nlp(text)
+    return {ent.text.strip() for ent in doc.ents if len(ent.text.strip()) >= 3}
+
+
+def expand_entities(
+    graph: nx.Graph,
+    entities: Iterable[str],
+    hops: int = 1,
+) -> Set[str]:
+    """Expand entities via graph traversal.
+
+    hops=1 → direct neighbors
+    hops=2 → neighbors of neighbors
+    """
+    expanded: Set[str] = set(entities)
+
+    for _ in range(hops):
+        neighbors: Set[str] = set()
+        for entity in expanded:
+            if entity in graph:
+                neighbors.update(graph.neighbors(entity))
+        expanded |= neighbors
+
+    return expanded
+
+
+def chunks_from_entities(
+    chunks: List[Chunk],
+    entities: Set[str],
+) -> List[Chunk]:
+    """Recall chunks mentioning any of the given entities.
+
+    THIS is the Graph-RAG recall step.
+    """
+    matched_chunk_ids: Set[str] = set()
+
+    for entity in entities:
+        matched_chunk_ids |= _ENTITY_TO_CHUNKS.get(entity, set())
+
+    return [chunk for chunk in chunks if chunk.chunk_id in matched_chunk_ids]

backend/app/retrieval/hybrid.py

@@ -1,31 +0,0 @@
-"""Hybrid retrieval (Vector + BM25)."""
-
-from typing import List
-
-from app.models.retrieval import ScoredChunk
-from app.retrieval.keyword_index import bm25_search
-from app.retrieval.normalize import normalize_scores
-from app.retrieval.vector_store import vector_search
-
-
-def hybrid_search(query: str, top_k: int = 5) -> List[ScoredChunk]:
-    """Run hybrid retrieval with score fusion."""
-    vector_chunks = normalize_scores(vector_search(query, top_k=10))
-    bm25_chunks = normalize_scores(bm25_search(query, top_k=10))
-
-    merged: dict[str, ScoredChunk] = {}
-
-    for chunk in vector_chunks + bm25_chunks:
-        cid = chunk.chunk.chunk_id
-        if cid not in merged:
-            merged[cid] = chunk
-        else:
-            merged[cid].score += chunk.score
-
-    ranked = sorted(
-        merged.values(),
-        key=lambda x: x.score,
-        reverse=True,
-    )
-
-    return ranked[:top_k]

backend/app/retrieval/normalize.py

@@ -1,25 +0,0 @@
-"""Score normalization utilities."""
-
-from typing import List
-
-from app.models.retrieval import ScoredChunk
-
-
-def normalize_scores(chunks: List[ScoredChunk]) -> List[ScoredChunk]:
-    """Min-max normalize scores to [0, 1]."""
-    if not chunks:
-        return []
-
-    scores = [c.score for c in chunks]
-    min_score = min(scores)
-    max_score = max(scores)
-
-    if min_score == max_score:
-        for c in chunks:
-            c.score = 1.0
-        return chunks
-
-    for c in chunks:
-        c.score = (c.score - min_score) / (max_score - min_score)
-
-    return chunks

backend/app/retrieval/retrieve.py

@@ -0,0 +1,61 @@
+"""Unified Hybrid + Graph retrieval.
+
+Pipeline:
+Vector → BM25 → merge → graph recall → rank
+"""
+
+from typing import Dict, List
+
+from app.ingestion.entities import NLP
+from app.models.retrieval import ScoredChunk
+from app.retrieval.chunk_registry import get_chunks
+from app.retrieval.graph_utils import (
+    build_graph,
+    chunks_from_entities,
+    expand_entities,
+    extract_query_entities,
+)
+from app.retrieval.keyword_index import bm25_search
+from app.retrieval.vector_store import vector_search
+
+
+def hybrid_graph_search(query: str, top_k: int) -> List[ScoredChunk]:
+    """Hybrid + Graph-RAG retrieval.
+
+    Important:
+    - top_k controls FINAL context size
+    - retrieval breadth is independent
+    """
+    # 1. Broad seed retrieval
+    seed_k = max(top_k * 4, 8)
+
+    vector_results = vector_search(query, top_k=seed_k)
+    bm25_results = bm25_search(query, top_k=seed_k)
+
+    combined: Dict[str, ScoredChunk] = {sc.chunk.chunk_id: sc for sc in vector_results}
+
+    for sc in bm25_results:
+        combined.setdefault(sc.chunk.chunk_id, sc)
+
+    # 2. Graph recall expansion
+    all_chunks = get_chunks()
+    graph = build_graph(all_chunks)
+
+    query_entities = extract_query_entities(query, NLP)
+
+    if query_entities:
+        expanded_entities = expand_entities(graph, query_entities, hops=1)
+        graph_chunks = chunks_from_entities(all_chunks, expanded_entities)
+
+        for chunk in graph_chunks:
+            if chunk.chunk_id not in combined:
+                combined[chunk.chunk_id] = ScoredChunk(
+                    chunk=chunk,
+                    score=0.25,  # low but non-zero recall score
+                )
+
+    # 3. Rank and return
+    results = list(combined.values())
+    results.sort(key=lambda x: x.score, reverse=True)
+
+    return results[:top_k]

requirements.txt

@@ -19,6 +19,7 @@ chromadb==0.5.0
 # Text Processing & NLP
 pymupdf==1.24.7
 spacy==3.7.4
+https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1-py3-none-any.whl
 sentence-transformers==2.6.1
 rank-bm25==0.2.2
 whoosh==2.7.4