graphrag/layers/orchestration_layer.py

"""
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, Optional, 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)
    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 (backward compat)."""
    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):
        import os
        if os.getenv("EMBEDDING_PROVIDER", "local") == "openai" and self.provider == "openai":
            try:
                from openai import OpenAI
                key = self._api_key or os.getenv("OPENAI_API_KEY", "")
                if key:
                    base_url = os.getenv("OPENAI_BASE_URL", "")
                    self._client = OpenAI(api_key=key, base_url=base_url) if base_url else 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 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,
                 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[TripleComparisonResult] = []

    def initialize(self):
        self.llm.initialize()
        self.embedder.initialize()

        # 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 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")
        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 3: GraphRAG (TG GraphRAG + Novelties) ──────

    def run_graphrag(self, query, passages=None, seed_entities=5, hops=2,
                     max_ctx=10, retriever="hybrid", community_level=1):
        """
        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", retriever_used=retriever)
        ti = to = cost = 0.0

        # 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")

        try:
            kws = json.loads(kw_resp.content)
        except json.JSONDecodeError:
            kws = {"high_level": [], "low_level": [query]}

        low_level = kws.get("low_level", [])

        # 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)
            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
            ]
            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
            tg_used = True

        # 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:
            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]]))
        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_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 (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)
        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
        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_triple_comparison(query, passages)
        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,
            "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,
                "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,
            },
        }