feat: strict Graph-RAG reasoning engine with JSON path output + validation script

backend/rag.py

@@ -16,11 +16,56 @@ FALLBACK_MODELS = [
     "mistralai/Mistral-Small-24B-Instruct-2501",
 ]
 
-SYSTEM_PROMPT = (
-    "You are a helpful GraphRAG assistant. Use only the provided context when possible. "
-    "If context is insufficient, say what is missing. Keep answers concise. "
-    "When answering from context, cite supporting sources inline as [Chunk X]."
-)
+SYSTEM_PROMPT = """\
+You are a Graph-RAG reasoning engine.
+
+Your job is NOT to answer directly from text.
+Your job is to construct answers by traversing entity relationships.
+
+STRICT RULES:
+
+1. ENTITY EXTRACTION
+   - Identify all entities in the question.
+   - Do not answer yet.
+
+2. RELATIONSHIP TRAVERSAL
+   - Use retrieved context to find explicit relationships between entities.
+   - Build a step-by-step path using only these relationships.
+
+3. MULTI-HOP ENFORCEMENT
+   - If the question requires multiple steps, you MUST show intermediate entities.
+   - Do NOT skip steps even if the answer is obvious.
+
+4. NO SHORTCUTS
+   - Do NOT answer from a single chunk if it contains the full answer.
+   - You MUST validate at least one intermediate relationship (bridge).
+
+5. FAITHFULNESS
+   - Use ONLY retrieved context. Do NOT use prior knowledge.
+   - If the path cannot be constructed -> return INSUFFICIENT_CONTEXT.
+
+6. NEGATIVE GUARD
+   - If the question asks for unknown/private/unsupported info -> return INSUFFICIENT_CONTEXT.
+
+7. OUTPUT FORMAT (MANDATORY) - Return JSON ONLY:
+{
+  "answer": "...",
+  "reasoning_type": "direct | multi-hop | insufficient",
+  "path": ["Entity1 -> Entity2", "Entity2 -> Entity3"],
+  "used_chunks": ["0", "1"],
+  "justification": "Explain briefly how the path leads to the answer using retrieved data only."
+}
+
+8. FAILURE CONDITIONS - Return:
+{
+  "answer": "INSUFFICIENT_CONTEXT",
+  "reasoning_type": "insufficient",
+  "path": [],
+  "used_chunks": [],
+  "justification": "No valid relationship path found in retrieved context."
+}
+IF: No relationship path exists | Only one-hop shortcut found | Information is missing
+"""
 
 STOPWORDS = {
     "a", "an", "and", "are", "as", "at", "be", "by", "for", "from", "has", "he", "in", "is", "it",
@@ -309,15 +354,19 @@ class RAGEngine:
 
         if answer_mode == "strict_grounded" and confidence_label == "Low":
             return {
-                "answer": "Insufficient context from document.",
+                "answer": "INSUFFICIENT_CONTEXT",
                 "chunks": chunks_payload,
                 "model_used": None,
                 "top_score": round(top_score, 4),
                 "confidence_label": confidence_label,
+                "reasoning_type": "insufficient",
+                "path": [],
+                "used_chunks": [],
+                "justification": "Low retrieval confidence; no valid path found.",
             }
 
         context_block = "\n\n".join(
-            f"Chunk {c.index} (score={c.score:.4f}):\n{c.text}" for c in retrieved
+            f"[chunk_id: {c.index}]\n{c.text}" for c in retrieved
         )
 
         hf_token = os.environ.get("HF_TOKEN")
@@ -328,6 +377,10 @@ class RAGEngine:
                 "model_used": None,
                 "top_score": round(top_score, 4),
                 "confidence_label": confidence_label,
+                "reasoning_type": "insufficient",
+                "path": [],
+                "used_chunks": [],
+                "justification": "",
             }
 
         messages = self._build_messages(query, context_block, history)
@@ -341,46 +394,77 @@ class RAGEngine:
         yield {"type": "chunks", "data": chunks_payload}
 
         if answer_mode == "strict_grounded" and confidence_label == "Low":
-            yield {"type": "token", "data": "Insufficient context from document."}
-            yield {"type": "done", "model_used": None}
+            yield {"type": "token", "data": "INSUFFICIENT_CONTEXT"}
+            yield {
+                "type": "done", "model_used": None,
+                "reasoning_type": "insufficient", "path": [],
+                "used_chunks": [], "justification": "Low retrieval confidence; no valid path found.",
+            }
             return
 
         hf_token = os.environ.get("HF_TOKEN")
         if not hf_token:
             yield {"type": "token", "data": "HF_TOKEN not set."}
-            yield {"type": "done", "model_used": None}
+            yield {"type": "done", "model_used": None, "reasoning_type": "insufficient", "path": [], "used_chunks": [], "justification": ""}
             return
 
         context_block = "\n\n".join(
-            f"Chunk {c.index} (score={c.score:.4f}):\n{c.text}" for c in retrieved
+            f"[chunk_id: {c.index}]\n{c.text}" for c in retrieved
         )
         messages = self._build_messages(query, context_block, history)
-        yield from self._stream_llm(hf_token, messages)
 
-    def _stream_llm(self, token: str, messages: list[dict]):
+        # Buffer full LLM output so we can parse the JSON before emitting clean answer tokens.
+        full_text, model_used = self._buffer_llm(hf_token, messages)
+        graph = self._parse_graph_response(full_text)
+        answer_text = graph.get("answer", full_text)
+
+        # Emit answer text word-by-word so the frontend stream still assembles naturally.
+        for token_chunk in re.split(r'(\s+)', answer_text):
+            if token_chunk:
+                yield {"type": "token", "data": token_chunk}
+
+        yield {
+            "type": "done",
+            "model_used": model_used,
+            "reasoning_type": graph.get("reasoning_type", "direct"),
+            "path": graph.get("path", []),
+            "used_chunks": graph.get("used_chunks", []),
+            "justification": graph.get("justification", ""),
+        }
+
+    def _buffer_llm(self, token: str, messages: list[dict]) -> tuple[str, str | None]:
+        """Non-streaming call used by stream_answer to enable JSON parsing before token emission."""
         client = InferenceClient(api_key=token)
         candidates = list(dict.fromkeys([self.llm_model] + FALLBACK_MODELS))
-
         for model in candidates:
             try:
-                for chunk in client.chat_completion(
-                    model=model,
-                    messages=messages,
-                    max_tokens=512,
-                    temperature=0.2,
-                    stream=True,
-                ):
-                    delta = chunk.choices[0].delta.content
-                    if delta:
-                        yield {"type": "token", "data": delta}
-                yield {"type": "done", "model_used": model}
-                return
+                resp = client.chat_completion(model=model, messages=messages, max_tokens=600, temperature=0.2)
+                return resp.choices[0].message.content, model
             except Exception as e:
-                log.warning("Streaming LLM %s failed: %s", model, e)
+                log.warning("Buffered LLM %s failed: %s", model, e)
                 continue
-
-        yield {"type": "token", "data": "All candidate models failed."}
-        yield {"type": "done", "model_used": None}
+        return "All candidate models failed.", None
+
+    def _parse_graph_response(self, text: str) -> dict:
+        """Extract and parse the mandatory JSON object from the LLM response."""
+        import json
+        # Strip optional markdown fences
+        cleaned = re.sub(r'^```(?:json)?\s*|\s*```$', '', text.strip(), flags=re.MULTILINE)
+        # Grab outermost JSON object
+        match = re.search(r'\{.*\}', cleaned, re.DOTALL)
+        if match:
+            try:
+                return json.loads(match.group())
+            except json.JSONDecodeError:
+                pass
+        # Fallback: treat raw text as answer, mark as direct
+        return {
+            "answer": text.strip(),
+            "reasoning_type": "direct",
+            "path": [],
+            "used_chunks": [],
+            "justification": "JSON parse failed; raw answer returned.",
+        }
 
     def _build_messages(self, query: str, context: str, history: list[dict] | None) -> list[dict]:
         messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
@@ -391,11 +475,13 @@ class RAGEngine:
                 "role": "user",
                 "content": (
                     f"Question:\n{query}\n\n"
+                    "Retrieved Context (use the chunk_id values in your used_chunks field):\n"
+                    f"{context}\n\n"
                     "Instructions:\n"
-                    "- Use only the provided context.\n"
-                    "- Cite supporting evidence inline as [Chunk X].\n"
-                    "- If context is insufficient, say context is insufficient.\n\n"
-                    f"Context:\n{context}"
+                    "- Step 1: Extract entities from the question.\n"
+                    "- Step 2: Traverse relationships across chunks to build a path.\n"
+                    "- Step 3: Return ONLY a valid JSON object matching the mandatory format.\n"
+                    "- Do NOT include any text outside the JSON object."
                 ),
             }
         )
@@ -414,13 +500,19 @@ class RAGEngine:
 
         for model in candidates:
             try:
-                resp = client.chat_completion(model=model, messages=messages, max_tokens=350, temperature=0.2)
+                resp = client.chat_completion(model=model, messages=messages, max_tokens=600, temperature=0.2)
+                raw = resp.choices[0].message.content
+                graph = self._parse_graph_response(raw)
                 return {
-                    "answer": resp.choices[0].message.content,
+                    "answer": graph.get("answer", raw),
                     "chunks": chunks_payload,
                     "model_used": model,
                     "top_score": round(top_score, 4),
                     "confidence_label": confidence_label,
+                    "reasoning_type": graph.get("reasoning_type", "direct"),
+                    "path": graph.get("path", []),
+                    "used_chunks": graph.get("used_chunks", []),
+                    "justification": graph.get("justification", ""),
                 }
             except Exception as e:
                 log.warning("LLM %s failed: %s", model, e)
@@ -432,4 +524,8 @@ class RAGEngine:
             "model_used": None,
             "top_score": round(top_score, 4),
             "confidence_label": confidence_label,
+            "reasoning_type": "insufficient",
+            "path": [],
+            "used_chunks": [],
+            "justification": "All LLM candidates failed.",
         }

validate_graphrag.py

@@ -0,0 +1,76 @@
+"""Standalone validation: Graph-RAG engine (no LLM token required)."""
+import sys, json
+sys.path.insert(0, ".")
+from backend.rag import RAGEngine
+
+engine = RAGEngine(chunk_size=95, chunk_overlap=0)
+
+doc = (
+    "Alpha Dynamics acquired Beta Labs in 2023 to expand its diagnostics portfolio.\n\n"
+    "Beta Labs later formed a strategic alliance with Orion Health for hospital analytics.\n\n"
+    "Orion Health announced a joint research program with Nova BioSystems focused on predictive care."
+)
+engine.ingest(doc)
+
+# ── 1. Graph index ────────────────────────────────────────────────────────────
+print("=== GRAPH INDEX ===")
+nodes = len(engine._entity_to_chunks)
+edges = sum(len(v) for v in engine._entity_graph.values()) // 2
+print(f"  nodes  : {nodes}")
+print(f"  edges  : {edges}")
+print(f"  chunks : {len(engine.chunks)}")
+top = sorted(engine._entity_to_chunks.items(), key=lambda x: len(x[1]), reverse=True)[:5]
+print("  top entities:", [(e, sorted(c)) for e, c in top])
+assert nodes > 0, "Graph must have nodes"
+assert edges > 0, "Graph must have edges"
+
+# ── 2. Retrieval – 2-hop ──────────────────────────────────────────────────────
+print("\n=== RETRIEVAL: 2-hop ===")
+r = engine.retrieve("How is Orion Health connected to Alpha Dynamics?")
+for c in r:
+    print(f"  chunk {c.index}  score={c.score:.4f}  {c.text[:70]}")
+assert any(c.index == 0 for c in r), "Must include Alpha Dynamics chunk"
+assert any(c.index == 1 for c in r), "Must include bridge chunk (Beta Labs)"
+
+# ── 3. Retrieval – 3-hop ──────────────────────────────────────────────────────
+print("\n=== RETRIEVAL: 3-hop ===")
+r2 = engine.retrieve("What links Nova BioSystems to Alpha Dynamics?")
+for c in r2:
+    print(f"  chunk {c.index}  score={c.score:.4f}  {c.text[:70]}")
+assert len(r2) >= 2, "Needs at least 2 bridge chunks for 3-hop"
+
+# ── 4. JSON parser – clean JSON ───────────────────────────────────────────────
+print("\n=== JSON PARSER ===")
+sample = {
+    "answer": "Alpha Dynamics -> Beta Labs -> Orion Health",
+    "reasoning_type": "multi-hop",
+    "path": ["Alpha Dynamics -> Beta Labs", "Beta Labs -> Orion Health"],
+    "used_chunks": ["0", "1"],
+    "justification": "Alpha acquired Beta, Beta allied with Orion.",
+}
+p = engine._parse_graph_response(json.dumps(sample))
+assert p["reasoning_type"] == "multi-hop"
+assert len(p["path"]) == 2
+print(f"  reasoning_type : {p['reasoning_type']}")
+print(f"  path           : {p['path']}")
+print(f"  used_chunks    : {p['used_chunks']}")
+
+# ── 5. JSON parser – markdown fenced ─────────────────────────────────────────
+fenced = "```json\n" + json.dumps(sample) + "\n```"
+p2 = engine._parse_graph_response(fenced)
+assert p2["reasoning_type"] == "multi-hop", "Must strip markdown fences"
+print(f"  fenced input parsed OK : {p2['reasoning_type']}")
+
+# ── 6. JSON parser – fallback on non-JSON ────────────────────────────────────
+p3 = engine._parse_graph_response("I cannot answer that question.")
+assert p3["reasoning_type"] == "direct"
+print(f"  fallback answer: {p3['answer'][:50]}")
+
+# ── 7. Negative: completely unknown entity (not in graph at all) ──────────────
+print("\n=== RETRIEVAL: negative (unknown entity) ===")
+r3 = engine.retrieve("Who is the founder of ZetaCorp Robotics?")
+_, conf = engine._confidence_from_retrieved(r3)
+print(f"  confidence_label: {conf}  (expected Low — ZetaCorp is not in graph)")
+assert conf == "Low", "Unknown-entity queries should produce Low confidence"
+
+print("\nALL ASSERTIONS PASSED - Graph-RAG engine fully validated.")