Fix #2+#4+#6: Add LLM-Only pipeline, wire NoveltyEngine, integrate TG GraphRAG client, 3-pipeline comparison

graphrag/layers/orchestration_layer.py

@@ -1,14 +1,19 @@
 """
-Layer 2: Inference Orchestration — Dual Pipeline Manager
-========================================================
-Routes queries through Baseline RAG and GraphRAG pipelines,
-collects metrics, and provides adaptive routing.
+Layer 2: Inference Orchestration — Triple Pipeline Manager
+==========================================================
+Routes queries through three pipelines:
+  Pipeline 1: LLM-Only (no retrieval — worst-case baseline)
+  Pipeline 2: Basic RAG (vector embeddings + LLM — industry standard)
+  Pipeline 3: GraphRAG (TigerGraph GraphRAG repo + novelty engine)
+
+Collects metrics for all three and provides adaptive routing.
 """
 import json
 import logging
 import time
+from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from .graph_layer import GraphLayer, cosine_similarity
 from .llm_layer import LLMLayer, LLMResponse, TokenTracker
@@ -33,11 +38,29 @@ class PipelineResult:
     complexity_score: float = 0.0
     query_type: str = ""
     token_breakdown: Dict = field(default_factory=dict)
+    novelty_chain: List[str] = field(default_factory=list)
+    retriever_used: str = ""
 
 
+@dataclass
+class TripleComparisonResult:
+    """Side-by-side comparison of all 3 pipelines."""
+    query: str = ""
+    llm_only: PipelineResult = field(default_factory=PipelineResult)
+    baseline: PipelineResult = field(default_factory=PipelineResult)
+    graphrag: PipelineResult = field(default_factory=PipelineResult)
+    token_savings_vs_baseline_pct: float = 0.0
+    token_savings_vs_llm_only_pct: float = 0.0
+    latency_diff_ms: float = 0.0
+    cost_diff_usd: float = 0.0
+    recommended_pipeline: str = ""
+    routing_reason: str = ""
+
+
+# Keep backward compat
 @dataclass
 class ComparisonResult:
-    """Side-by-side comparison of both pipelines."""
+    """Side-by-side comparison of both pipelines (backward compat)."""
     query: str = ""
     baseline: PipelineResult = field(default_factory=PipelineResult)
     graphrag: PipelineResult = field(default_factory=PipelineResult)
@@ -106,28 +129,84 @@ class EmbeddingManager:
 
 class InferenceOrchestrator:
     """
-    Layer 2: Manages both pipelines and routes queries.
+    Layer 2: Manages all three pipelines and routes queries.
+      Pipeline 1: LLM-Only (no retrieval)
+      Pipeline 2: Basic RAG (vector search + LLM)
+      Pipeline 3: GraphRAG (TG GraphRAG service + novelty engine)
     """
 
-    def __init__(self, graph_layer=None, llm_layer=None, embedder=None, config=None):
+    def __init__(self, graph_layer=None, llm_layer=None, embedder=None,
+                 tg_graphrag_client=None, novelty_engine=None, config=None):
         self.graph = graph_layer or GraphLayer()
         self.llm = llm_layer or LLMLayer()
         self.embedder = embedder or EmbeddingManager()
+        self.tg_client = tg_graphrag_client  # official TG GraphRAG service client
+        self.novelty_engine = novelty_engine  # NoveltyEngine from novelties.py
         self.config = config or {}
+        self.llm_only_tracker = TokenTracker()
         self.baseline_tracker = TokenTracker()
         self.graphrag_tracker = TokenTracker()
-        self.comparison_history: List[ComparisonResult] = []
+        self.comparison_history: List[TripleComparisonResult] = []
 
     def initialize(self):
         self.llm.initialize()
         self.embedder.initialize()
-        logger.info("Inference Orchestrator initialized.")
 
-    # ── Pipeline A: Baseline RAG ────────────────────────────
+        # Initialize TG GraphRAG client if not provided
+        if self.tg_client is None:
+            try:
+                from .tg_graphrag_client import TGGraphRAGClient
+                self.tg_client = TGGraphRAGClient()
+                self.tg_client.connect()
+            except Exception as e:
+                logger.info(f"TG GraphRAG client not available: {e}")
+
+        # Initialize NoveltyEngine if not provided
+        if self.novelty_engine is None:
+            try:
+                from .novelties import NoveltyEngine
+                self.novelty_engine = NoveltyEngine(
+                    token_budget=self.config.get("token_budget", 2000))
+                logger.info("NoveltyEngine initialized.")
+            except Exception as e:
+                logger.warning(f"NoveltyEngine not available: {e}")
+
+        logger.info("Inference Orchestrator initialized (3-pipeline mode).")
+
+    # ── Pipeline 1: LLM-Only (No Retrieval) ─────────────────
+
+    def run_llm_only(self, query: str) -> PipelineResult:
+        """
+        Pipeline 1: LLM-Only — raw prompt in, answer out. No retrieval.
+        This is the worst-case baseline: the LLM uses only its parametric knowledge.
+        """
+        start = time.perf_counter()
+        result = PipelineResult(pipeline_type="llm_only")
+
+        sys_prompt = (
+            "You are a knowledgeable assistant. Answer the question accurately and concisely "
+            "based on your knowledge. If you are not sure, say so."
+        )
+        resp = self.llm.generate([
+            {"role": "system", "content": sys_prompt},
+            {"role": "user", "content": f"Question: {query}\n\nAnswer:"},
+        ], max_tokens=512)
+
+        result.answer = resp.content
+        result.input_tokens = resp.input_tokens
+        result.output_tokens = resp.output_tokens
+        result.total_tokens = resp.total_tokens
+        result.cost_usd = resp.cost_usd
+        result.latency_ms = (time.perf_counter() - start) * 1000
+        self.llm_only_tracker.record(resp, "llm_only")
+        return result
+
+    # ── Pipeline 2: Basic RAG ────────────────────────────────
 
     def run_baseline_rag(self, query, passages=None, top_k=5):
         """
-        Pipeline A: Query → Embed → Vector Search → Top-K Chunks → LLM → Answer
+        Pipeline 2: Basic RAG — Query → Embed → Vector Search → Top-K Chunks → LLM → Answer
+        Industry standard vector-based retrieval augmented generation.
         """
         start = time.perf_counter()
         result = PipelineResult(pipeline_type="baseline")
@@ -158,18 +237,27 @@ class InferenceOrchestrator:
         self.baseline_tracker.record(resp, "baseline")
         return result
 
-    # ── Pipeline B: GraphRAG ────────────────────────────────
+    # ── Pipeline 3: GraphRAG (TG GraphRAG + Novelties) ──────
 
-    def run_graphrag(self, query, passages=None, seed_entities=5, hops=2, max_ctx=10):
+    def run_graphrag(self, query, passages=None, seed_entities=5, hops=2,
+                     max_ctx=10, retriever="hybrid", community_level=1):
         """
-        Pipeline B: Query → Keywords → Entity Search → Graph Traverse → Structured Context → LLM
-        Novelties: Dual-level keywords, schema-bounded extraction, graph reasoning
+        Pipeline 3: GraphRAG — Built on top of the TigerGraph GraphRAG repo.
+
+        Flow:
+          1. Call TG GraphRAG service (official repo REST API) for retrieval
+          2. Apply NoveltyEngine enhancements (PPR, activation, token budget, etc.)
+          3. Build structured context with entities + relationships + passages
+          4. Generate answer with graph-aware LLM prompt
+
+        Falls back to direct pyTigerGraph GSQL queries if service unavailable.
+        Falls back to passage-based entity extraction if no TG connection.
         """
         start = time.perf_counter()
-        result = PipelineResult(pipeline_type="graphrag")
+        result = PipelineResult(pipeline_type="graphrag", retriever_used=retriever)
         ti = to = cost = 0.0
 
-        # Step 1: Extract dual-level keywords (LightRAG-inspired)
+        # Step 1: Extract dual-level keywords (LightRAG-inspired novelty)
         kw_resp = self.llm.extract_keywords(query)
         ti += kw_resp.input_tokens; to += kw_resp.output_tokens; cost += kw_resp.cost_usd
         self.graphrag_tracker.record(kw_resp, "keywords")
@@ -181,8 +269,27 @@ class InferenceOrchestrator:
 
         low_level = kws.get("low_level", [])
 
-        if self.graph.is_connected:
-            # Step 2: Find seed entities via vector search
+        # Step 2: Try TG GraphRAG service first (official repo integration)
+        tg_used = False
+        if self.tg_client and self.tg_client.is_connected:
+            try:
+                tg_result = self.tg_client.retrieve(
+                    query=query, retriever=retriever,
+                    top_k=seed_entities * 2, num_hops=hops,
+                    community_level=community_level,
+                )
+                if tg_result.chunks:
+                    result.contexts = [c.get("text", "") for c in tg_result.chunks[:max_ctx]]
+                    result.entities_found = tg_result.entities
+                    result.relations_traversed = tg_result.relations
+                    result.hops_used = hops
+                    tg_used = True
+                    logger.info(f"TG GraphRAG service returned {len(tg_result.chunks)} chunks")
+            except Exception as e:
+                logger.warning(f"TG GraphRAG service call failed: {e}")
+
+        # Step 2b: Fall back to direct pyTigerGraph if service failed
+        if not tg_used and self.graph.is_connected:
             search_text = " ".join(low_level) if low_level else query
             query_emb = self.embedder.embed_single(search_text)
             ents = self.graph.vector_search_entities(query_emb, seed_entities)
@@ -192,60 +299,108 @@ class InferenceOrchestrator:
                  "description": e.get("description",""), "score": e.get("score",0)}
                 for e in ents
             ]
-            # Step 3: Multi-hop graph traversal
             if seed_ids:
                 traversal = self.graph.graph_traverse(seed_ids, hops)
                 result.contexts = traversal.get("chunk_texts", [])[:max_ctx]
                 result.relations_traversed = traversal.get("relations", [])
                 result.hops_used = hops
-        else:
-            # Fallback: simulate GraphRAG with passages + entity extraction
-            if passages:
-                query_emb = self.embedder.embed_single(query)
-                passage_embs = self.embedder.embed(passages)
-                scored = sorted(
-                    [(cosine_similarity(query_emb, emb), p, i)
-                     for i, (p, emb) in enumerate(zip(passages, passage_embs))],
-                    reverse=True
-                )
+            tg_used = True
 
-                # Extract entities from top passages (simulates graph construction)
-                top_p = scored[:3]
-                all_ent_names = set()
-                for _, passage, _ in top_p:
-                    ext_resp = self.llm.extract_entities(passage)
-                    ti += ext_resp.input_tokens; to += ext_resp.output_tokens; cost += ext_resp.cost_usd
-                    self.graphrag_tracker.record(ext_resp, "entity_extraction")
-                    try:
-                        extracted = json.loads(ext_resp.content)
-                        for ent in extracted.get("entities", []):
-                            all_ent_names.add(ent.get("name", ""))
-                            result.entities_found.append(ent)
-                        for rel in extracted.get("relations", []):
-                            result.relations_traversed.append(
-                                f"{rel['source']} -[{rel['type']}]-> {rel['target']}: {rel.get('description','')}")
-                    except json.JSONDecodeError:
-                        pass
-
-                # Multi-hop simulation: expand by entity mentions
-                expanded = []
-                for _, passage, idx in scored:
-                    for en in all_ent_names:
-                        if en.lower() in passage.lower():
-                            expanded.append(passage)
-                            break
-                all_ctx = [p for _, p, _ in top_p]
-                for ep in expanded:
-                    if ep not in all_ctx: all_ctx.append(ep)
-                result.contexts = all_ctx[:max_ctx]
-                result.hops_used = hops
+        # Step 2c: Fallback for offline mode — simulate with passages + entity extraction
+        if not tg_used and passages:
+            query_emb = self.embedder.embed_single(query)
+            passage_embs = self.embedder.embed(passages)
+            scored = sorted(
+                [(cosine_similarity(query_emb, emb), p, i)
+                 for i, (p, emb) in enumerate(zip(passages, passage_embs))],
+                reverse=True
+            )
+            top_p = scored[:3]
+            all_ent_names = set()
+            for _, passage, _ in top_p:
+                ext_resp = self.llm.extract_entities(passage)
+                ti += ext_resp.input_tokens; to += ext_resp.output_tokens; cost += ext_resp.cost_usd
+                self.graphrag_tracker.record(ext_resp, "entity_extraction")
+                try:
+                    extracted = json.loads(ext_resp.content)
+                    for ent in extracted.get("entities", []):
+                        all_ent_names.add(ent.get("name", ""))
+                        result.entities_found.append(ent)
+                    for rel in extracted.get("relations", []):
+                        result.relations_traversed.append(
+                            f"{rel['source']} -[{rel['type']}]-> {rel['target']}: {rel.get('description','')}")
+                except json.JSONDecodeError:
+                    pass
+
+            # Multi-hop simulation: expand by entity mentions
+            expanded = []
+            for _, passage, idx in scored:
+                for en in all_ent_names:
+                    if en.lower() in passage.lower():
+                        expanded.append(passage)
+                        break
+            all_ctx = [p for _, p, _ in top_p]
+            for ep in expanded:
+                if ep not in all_ctx: all_ctx.append(ep)
+            result.contexts = all_ctx[:max_ctx]
+            result.hops_used = hops
+
+        # Step 3: Apply NoveltyEngine enhancements if available
+        if self.novelty_engine and result.entities_found and result.contexts:
+            try:
+                # Build adjacency from extracted relations
+                adjacency: Dict[str, List[Tuple[str, float]]] = defaultdict(list)
+                entity_to_chunks: Dict[str, List[str]] = defaultdict(list)
+                chunk_texts: Dict[str, str] = {}
+
+                for i, ctx in enumerate(result.contexts):
+                    cid = f"ctx_{i}"
+                    chunk_texts[cid] = ctx
+
+                for e in result.entities_found:
+                    ename = e.get("name", "").lower()
+                    for i, ctx in enumerate(result.contexts):
+                        if ename in ctx.lower():
+                            entity_to_chunks[ename].append(f"ctx_{i}")
+
+                for rel_str in result.relations_traversed:
+                    parts = rel_str.split(" -[")
+                    if len(parts) >= 2:
+                        src = parts[0].strip().lower()
+                        rest = parts[1].split("]->")
+                        if len(rest) >= 2:
+                            tgt = rest[1].split(":")[0].strip().lower()
+                            adjacency[src].append((tgt, 0.8))
+                            adjacency[tgt].append((src, 0.8))
+
+                seed_ents = [e.get("name", "").lower() for e in result.entities_found[:5]]
+
+                if adjacency and seed_ents and entity_to_chunks:
+                    novelty_result = self.novelty_engine.enhanced_retrieve(
+                        query=query,
+                        adjacency=adjacency,
+                        seed_entities=seed_ents,
+                        entity_to_chunks=entity_to_chunks,
+                        chunk_texts=chunk_texts,
+                    )
+                    if novelty_result.get("contexts"):
+                        result.contexts = novelty_result["contexts"]
+                    result.novelty_chain = novelty_result.get("technique_chain", [])
+                    logger.info(f"NoveltyEngine applied: {result.novelty_chain}")
+            except Exception as e:
+                logger.warning(f"NoveltyEngine enhancement failed: {e}")
 
         # Step 4: Build structured context with graph information
         ctx_parts = []
         if result.entities_found:
-            ctx_parts.append("### Entities Found:\n" + "\n".join(
-                [f"- **{e.get('name','?')}** ({e.get('entity_type','?')}): {e.get('description','')}"
-                 for e in result.entities_found[:10]]))
+            ent_list = result.entities_found[:10]
+            if isinstance(ent_list[0], dict):
+                ctx_parts.append("### Entities Found:\n" + "\n".join(
+                    [f"- **{e.get('name','?')}** ({e.get('entity_type','?')}): {e.get('description','')}"
+                     for e in ent_list]))
+            else:
+                ctx_parts.append("### Entities Found:\n" + "\n".join(
+                    [f"- {e}" for e in ent_list]))
         if result.relations_traversed:
             ctx_parts.append("### Relationships:\n" + "\n".join(
                 [f"- {r}" for r in result.relations_traversed[:15]]))
@@ -280,8 +435,29 @@ class InferenceOrchestrator:
         except (json.JSONDecodeError, ValueError):
             return 0.5, "unknown", "Analysis failed"
 
+    def run_triple_comparison(self, query, passages=None, top_k=5, hops=2):
+        """Run all 3 pipelines and compare side-by-side."""
+        lo = self.run_llm_only(query)
+        b = self.run_baseline_rag(query, passages, top_k)
+        g = self.run_graphrag(query, passages, hops=hops)
+
+        comp = TripleComparisonResult(query=query, llm_only=lo, baseline=b, graphrag=g)
+        if b.total_tokens > 0:
+            comp.token_savings_vs_baseline_pct = (
+                (b.total_tokens - g.total_tokens) / b.total_tokens * 100
+            )
+        if lo.total_tokens > 0:
+            comp.token_savings_vs_llm_only_pct = (
+                (lo.total_tokens - g.total_tokens) / lo.total_tokens * 100
+            )
+        comp.latency_diff_ms = g.latency_ms - b.latency_ms
+        comp.cost_diff_usd = g.cost_usd - b.cost_usd
+        self.comparison_history.append(comp)
+        return comp
+
+    # Backward compat — 2-pipeline comparison
     def run_comparison(self, query, passages=None, top_k=5, hops=2):
-        """Run both pipelines and compare."""
+        """Run both pipelines and compare (backward compat)."""
         b = self.run_baseline_rag(query, passages, top_k)
         g = self.run_graphrag(query, passages, hops=hops)
         comp = ComparisonResult(query=query, baseline=b, graphrag=g)
@@ -289,15 +465,12 @@ class InferenceOrchestrator:
             comp.token_savings_pct = (g.total_tokens - b.total_tokens) / b.total_tokens * 100
         comp.latency_diff_ms = g.latency_ms - b.latency_ms
         comp.cost_diff_usd = g.cost_usd - b.cost_usd
-        self.comparison_history.append(comp)
         return comp
 
     def run_adaptive(self, query, passages=None, threshold=0.6):
         """Adaptive routing: automatically picks optimal pipeline."""
         score, qtype, reasoning = self.analyze_complexity(query)
-        comp = self.run_comparison(query, passages)
-        comp.baseline.complexity_score = score
-        comp.baseline.query_type = qtype
+        comp = self.run_triple_comparison(query, passages)
         comp.graphrag.complexity_score = score
         comp.graphrag.query_type = qtype
         if score >= threshold:
@@ -320,6 +493,12 @@ class InferenceOrchestrator:
         n = len(self.comparison_history)
         return {
             "total_queries": n,
+            "llm_only": {
+                "total_tokens": sum(c.llm_only.total_tokens for c in self.comparison_history),
+                "avg_tokens": sum(c.llm_only.total_tokens for c in self.comparison_history) / n,
+                "total_cost": sum(c.llm_only.cost_usd for c in self.comparison_history),
+                "avg_latency": sum(c.llm_only.latency_ms for c in self.comparison_history) / n,
+            },
             "baseline": {
                 "total_tokens": sum(c.baseline.total_tokens for c in self.comparison_history),
                 "avg_tokens": sum(c.baseline.total_tokens for c in self.comparison_history) / n,