Add Layer 2: Inference Orchestration (dual pipeline, adaptive routing, comparison)

graphrag/layers/orchestration_layer.py

@@ -0,0 +1,335 @@
+"""
+Layer 2: Inference Orchestration — Dual Pipeline Manager
+========================================================
+Routes queries through Baseline RAG and GraphRAG pipelines,
+collects metrics, and provides adaptive routing.
+"""
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Tuple
+
+from .graph_layer import GraphLayer, cosine_similarity
+from .llm_layer import LLMLayer, LLMResponse, TokenTracker
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PipelineResult:
+    """Result from a single pipeline execution."""
+    answer: str = ""
+    contexts: List[str] = field(default_factory=list)
+    total_tokens: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    latency_ms: float = 0.0
+    cost_usd: float = 0.0
+    pipeline_type: str = ""
+    entities_found: List[Dict] = field(default_factory=list)
+    relations_traversed: List[str] = field(default_factory=list)
+    hops_used: int = 0
+    complexity_score: float = 0.0
+    query_type: str = ""
+    token_breakdown: Dict = field(default_factory=dict)
+
+
+@dataclass
+class ComparisonResult:
+    """Side-by-side comparison of both pipelines."""
+    query: str = ""
+    baseline: PipelineResult = field(default_factory=PipelineResult)
+    graphrag: PipelineResult = field(default_factory=PipelineResult)
+    token_savings_pct: float = 0.0
+    latency_diff_ms: float = 0.0
+    cost_diff_usd: float = 0.0
+    recommended_pipeline: str = ""
+    routing_reason: str = ""
+
+
+class EmbeddingManager:
+    """Manages embedding generation (OpenAI or local)."""
+
+    def __init__(self, provider="openai", model="text-embedding-3-small",
+                 api_key="", dimension=1536):
+        self.provider = provider
+        self.model = model
+        self._api_key = api_key
+        self.dimension = dimension
+        self._client = None
+        self._local_model = None
+
+    def initialize(self):
+        if self.provider == "openai":
+            try:
+                from openai import OpenAI
+                import os
+                key = self._api_key or os.getenv("OPENAI_API_KEY", "")
+                if key:
+                    self._client = OpenAI(api_key=key)
+                    logger.info(f"OpenAI embeddings: {self.model}")
+                else:
+                    self._init_local()
+            except ImportError:
+                self._init_local()
+        else:
+            self._init_local()
+
+    def _init_local(self):
+        try:
+            from sentence_transformers import SentenceTransformer
+            self._local_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+            self.dimension = 384
+            self.provider = "local"
+            logger.info("Local embeddings: all-MiniLM-L6-v2")
+        except ImportError:
+            logger.warning("No embedding model available — zero vectors")
+
+    def embed(self, texts: List[str]) -> List[List[float]]:
+        if not texts: return []
+        if self.provider == "openai" and self._client:
+            try:
+                resp = self._client.embeddings.create(input=texts, model=self.model)
+                return [item.embedding for item in resp.data]
+            except Exception as e:
+                logger.error(f"Embedding error: {e}")
+                return [[0.0] * self.dimension for _ in texts]
+        elif self._local_model:
+            return [emb.tolist() for emb in self._local_model.encode(texts)]
+        return [[0.0] * self.dimension for _ in texts]
+
+    def embed_single(self, text: str) -> List[float]:
+        r = self.embed([text])
+        return r[0] if r else [0.0] * self.dimension
+
+
+class InferenceOrchestrator:
+    """
+    Layer 2: Manages both pipelines and routes queries.
+    """
+
+    def __init__(self, graph_layer=None, llm_layer=None, embedder=None, config=None):
+        self.graph = graph_layer or GraphLayer()
+        self.llm = llm_layer or LLMLayer()
+        self.embedder = embedder or EmbeddingManager()
+        self.config = config or {}
+        self.baseline_tracker = TokenTracker()
+        self.graphrag_tracker = TokenTracker()
+        self.comparison_history: List[ComparisonResult] = []
+
+    def initialize(self):
+        self.llm.initialize()
+        self.embedder.initialize()
+        logger.info("Inference Orchestrator initialized.")
+
+    # ── Pipeline A: Baseline RAG ────────────────────────────
+
+    def run_baseline_rag(self, query, passages=None, top_k=5):
+        """
+        Pipeline A: Query → Embed → Vector Search → Top-K Chunks → LLM → Answer
+        """
+        start = time.perf_counter()
+        result = PipelineResult(pipeline_type="baseline")
+        ti = to = cost = 0.0
+
+        if passages:
+            query_emb = self.embedder.embed_single(query)
+            passage_embs = self.embedder.embed(passages)
+            scored = sorted(
+                [(cosine_similarity(query_emb, emb), p) for p, emb in zip(passages, passage_embs)],
+                reverse=True
+            )
+            result.contexts = [p for _, p in scored[:top_k]]
+        elif self.graph.is_connected:
+            query_emb = self.embedder.embed_single(query)
+            chunks = self.graph.vector_search_chunks(query_emb, top_k)
+            result.contexts = [c.get("text", "") for c in chunks]
+        else:
+            result.contexts = ["[No context available — connect TigerGraph or provide passages]"]
+
+        ctx_text = "\n\n".join(result.contexts[:top_k])
+        resp = self.llm.generate_answer(query, ctx_text)
+        result.answer = resp.content
+        ti += resp.input_tokens; to += resp.output_tokens; cost += resp.cost_usd
+        result.input_tokens = int(ti); result.output_tokens = int(to)
+        result.total_tokens = int(ti + to); result.cost_usd = cost
+        result.latency_ms = (time.perf_counter() - start) * 1000
+        self.baseline_tracker.record(resp, "baseline")
+        return result
+
+    # ── Pipeline B: GraphRAG ────────────────────────────────
+
+    def run_graphrag(self, query, passages=None, seed_entities=5, hops=2, max_ctx=10):
+        """
+        Pipeline B: Query → Keywords → Entity Search → Graph Traverse → Structured Context → LLM
+        Novelties: Dual-level keywords, schema-bounded extraction, graph reasoning
+        """
+        start = time.perf_counter()
+        result = PipelineResult(pipeline_type="graphrag")
+        ti = to = cost = 0.0
+
+        # Step 1: Extract dual-level keywords (LightRAG-inspired)
+        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")
+
+        try:
+            kws = json.loads(kw_resp.content)
+        except json.JSONDecodeError:
+            kws = {"high_level": [], "low_level": [query]}
+
+        low_level = kws.get("low_level", [])
+
+        if self.graph.is_connected:
+            # Step 2: Find seed entities via vector search
+            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)
+            seed_ids = [e.get("entity_id", "") for e in ents]
+            result.entities_found = [
+                {"name": e.get("name",""), "entity_type": e.get("entity_type",""),
+                 "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
+                )
+
+                # 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 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]]))
+        if result.relations_traversed:
+            ctx_parts.append("### Relationships:\n" + "\n".join(
+                [f"- {r}" for r in result.relations_traversed[:15]]))
+        if result.contexts:
+            ctx_parts.append("### Retrieved Passages:\n" + "\n\n".join(
+                [f"[Passage {i+1}]: {c}" for i, c in enumerate(result.contexts[:max_ctx])]))
+
+        structured = "\n\n".join(ctx_parts)
+        sys_prompt = (
+            "You are a knowledgeable assistant with access to a knowledge graph. "
+            "Use the structured context including entities, relationships, and passages "
+            "to answer accurately. Follow relationship chains for multi-hop reasoning. Be concise."
+        )
+        gen_resp = self.llm.generate_answer(query, structured, sys_prompt)
+        ti += gen_resp.input_tokens; to += gen_resp.output_tokens; cost += gen_resp.cost_usd
+        self.graphrag_tracker.record(gen_resp, "graphrag_gen")
+
+        result.answer = gen_resp.content
+        result.input_tokens = int(ti); result.output_tokens = int(to)
+        result.total_tokens = int(ti + to); result.cost_usd = cost
+        result.latency_ms = (time.perf_counter() - start) * 1000
+        return result
+
+    # ── Adaptive Query Router (Novelty) ─────────────────────
+
+    def analyze_complexity(self, query):
+        """Analyze query complexity for adaptive routing."""
+        resp = self.llm.analyze_query_complexity(query)
+        try:
+            a = json.loads(resp.content)
+            return float(a.get("complexity_score", 0.5)), a.get("query_type", "unknown"), a.get("reasoning", "")
+        except (json.JSONDecodeError, ValueError):
+            return 0.5, "unknown", "Analysis failed"
+
+    def run_comparison(self, query, passages=None, top_k=5, hops=2):
+        """Run both pipelines and compare."""
+        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)
+        if b.total_tokens > 0:
+            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.graphrag.complexity_score = score
+        comp.graphrag.query_type = qtype
+        if score >= threshold:
+            comp.recommended_pipeline = "graphrag"
+            comp.routing_reason = f"Complex query (score={score:.2f}, type={qtype}): {reasoning}"
+        else:
+            comp.recommended_pipeline = "baseline"
+            comp.routing_reason = f"Simple query (score={score:.2f}, type={qtype}): {reasoning}"
+        return comp
+
+    def explain_graphrag_reasoning(self, query, graphrag_result):
+        """Generate reasoning path explanation (novelty)."""
+        resp = self.llm.generate_graph_explanation(
+            query, graphrag_result.entities_found,
+            graphrag_result.relations_traversed, graphrag_result.answer)
+        return resp.content
+
+    def get_aggregate_metrics(self):
+        if not self.comparison_history: return {"message": "No comparisons"}
+        n = len(self.comparison_history)
+        return {
+            "total_queries": 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,
+                "total_cost": sum(c.baseline.cost_usd for c in self.comparison_history),
+                "avg_latency": sum(c.baseline.latency_ms for c in self.comparison_history) / n,
+            },
+            "graphrag": {
+                "total_tokens": sum(c.graphrag.total_tokens for c in self.comparison_history),
+                "avg_tokens": sum(c.graphrag.total_tokens for c in self.comparison_history) / n,
+                "total_cost": sum(c.graphrag.cost_usd for c in self.comparison_history),
+                "avg_latency": sum(c.graphrag.latency_ms for c in self.comparison_history) / n,
+            },
+        }