Rename chimera_router_v2.1.py to chimera_router_v2_1.py (fix Python import syntax)

nexus_os_v2/chimera_router_v2_1.py

@@ -0,0 +1,644 @@
+"""
+ChimeraRouter v2.1 — Production Telemetry-Integrated Inference Orchestrator
+
+Integrates:
+  - Sulphur prompt enhancement
+  - QWAVE budget allocation
+  - Multi-source retrieval (Pinecone + Milvus + ERNIE)
+  - TWAVE token-level thermodynamic tracking via Ollama telemetry
+  - CK-PLUG confidence gain coupling
+  - EPR + Spilled Energy unified detector
+  - Cloud API adapters (DeepSeek, Qwen, Kimi, GLM, OpenAI, Claude)
+  - Model fallback controller (reflection → grounding → switch → cloud)
+  - Stochastic resonance optimal temperature
+
+Architecture:
+  1. Sulphur enhancement
+  2. QWAVE budget → model selection
+  3. Multi-source retrieval
+  4. Stochastic resonance T_eff optimization
+  5. TWAVE tracker + unified detector initialization
+  6. Generation (Ollama or Cloud API)
+  7. Post-hoc telemetry + per_token_debug
+  8. Fallback controller if risk too high
+  9. Result assembly
+"""
+import os
+import time
+import json
+from typing import List, Dict, Optional, Any, Tuple
+from dataclasses import dataclass, field
+from enum import Enum
+
+from .model_registry import REGISTRY, SULPHUR, get, by_tier, by_cap, Tier, Capability, ModelProfile
+from .sulphur_enhancer import SulphurEnhancer, MockSulphurEnhancer, EnhancedPrompt
+from .twave_tracker import TWAVETracker, TokenState, StochasticResonance, GenerationTrajectory
+from .ckplug_retriever import CKPLUGCoupling, get_preset_epsilon
+from .pinecone_client import PineconeRetriever, MockPineconeRetriever
+from .milvus_client import MilvusRetriever, MockMilvusRetriever
+from .ernie_adapter import ERNIEAdapter, MockERNIEAdapter
+from .ollama_telemetry import OllamaStreamingClient, OllamaTelemetryExtractor, estimate_entropy_from_response
+from .per_token_debug import PerTokenDebug, GenerationTelemetry
+from .unified_detector import UnifiedThermodynamicDetector, FusionMode, Action
+from .epr_detector import EPRDetector, SequenceEPR
+from .spilled_energy import SpilledEnergyDetector, CombinedThermodynamicDetector
+from .cloud_api_adapters import CloudAPIManager, CloudResponse
+
+
+class RoutingDecision(Enum):
+    LOCAL_OLLAMA = "local_ollama"
+    CLOUD_API = "cloud_api"
+    FALLBACK = "fallback"
+    REFLECTION = "reflection"
+
+
+@dataclass
+class QWAVEBudget:
+    max_tokens: int = 4096
+    target_latency_ms: float = 2000.0
+    vram_budget_gb: float = 8.0
+    cloud_budget_cents: float = 5.0
+    allow_cloud: bool = True
+    allow_uncensored: bool = True
+    require_vision: bool = False
+    require_safety: bool = False
+    require_tools: bool = False
+
+
+@dataclass
+class RouterResult:
+    selected_model: str
+    model_profile: ModelProfile
+    tier: str
+    enhanced_prompt: str
+    response: str
+    token_states: List[TokenState] = field(default_factory=list)
+    per_token_debug: List[PerTokenDebug] = field(default_factory=list)
+    generation_telemetry: Optional[GenerationTelemetry] = None
+    reflection_count: int = 0
+    grounding_score: float = 0.0
+    hallucination_risk: float = 0.0
+    latency_ms: float = 0.0
+    tokens_generated: int = 0
+    detector_verdict: Optional[Any] = None
+    fallback_history: List[str] = field(default_factory=list)
+    debug: Dict[str, Any] = field(default_factory=dict)
+
+
+class ModelFallbackController:
+    """
+    Handles generation failures and high hallucination risk through escalation:
+      Level 1: Increase retrieval grounding (re-ground prompt)
+      Level 2: Backtrack and regenerate (reflection)
+      Level 3: Switch to more capable model (larger params or cloud)
+      Level 4: Fallback to cloud API (guaranteed generation)
+    """
+    
+    MAX_RETRIES = 3
+    
+    def __init__(self, router: "ChimeraRouterV2"):
+        self.router = router
+        self.retry_count = 0
+    
+    def escalate(
+        self,
+        prompt: str,
+        enhanced: EnhancedPrompt,
+        evidence: Dict[str, Any],
+        budget: QWAVEBudget,
+        previous_model_id: str,
+        previous_risk: float,
+    ) -> Tuple[str, str, ModelProfile, List[str]]:
+        """
+        Escalate to next level. Returns (new_model_id, new_prompt, profile, history).
+        """
+        history = []
+        
+        # Level 1: Increase grounding if retrieval available
+        if evidence.get("aggregated") and self.retry_count == 0:
+            top_evidence = "\n".join([
+                f"[HIGH-PRIORITY EVIDENCE] {e.get('text', '')[:400]}"
+                for e in evidence.get("aggregated", [])[:5]
+            ])
+            new_prompt = f"""CRITICAL: Use ONLY the following verified evidence to answer. 
+Do not rely on parametric knowledge if it conflicts with the evidence.
+
+{top_evidence}
+
+---
+
+{enhanced.enhanced}"""
+            history.append("re-grounded with more evidence")
+            self.retry_count += 1
+            return previous_model_id, new_prompt, get(previous_model_id), history
+        
+        # Level 2: Switch to next larger model in same tier
+        current_profile = get(previous_model_id)
+        if current_profile and current_profile.tier != Tier.CLOUD_API:
+            current_tier_models = by_tier(current_profile.tier)
+            larger_models = [m for m in current_tier_models 
+                           if m.params_b > current_profile.params_b 
+                           and m.size_gb <= budget.vram_budget_gb]
+            if larger_models and self.retry_count < 2:
+                larger = max(larger_models, key=lambda m: m.params_b)
+                for k, v in REGISTRY.items():
+                    if v == larger:
+                        history.append(f"switched to larger model {k} ({larger.params_b:.1f}B)")
+                        self.retry_count += 1
+                        return k, enhanced.enhanced, larger, history
+        
+        # Level 3: Upgrade to next tier
+        tier_upgrade = {
+            Tier.LOCAL_8GB: Tier.LOCAL_16GB,
+            Tier.LOCAL_16GB: Tier.LOCAL_24GB,
+            Tier.LOCAL_24GB: Tier.LOCAL_48GB,
+            Tier.LOCAL_48GB: Tier.CLOUD_API,
+        }
+        
+        if current_profile and current_profile.tier in tier_upgrade:
+            next_tier = tier_upgrade[current_profile.tier]
+            next_tier_models = by_tier(next_tier)
+            available = [m for m in next_tier_models 
+                        if m.tier == Tier.CLOUD_API or m.size_gb <= budget.vram_budget_gb]
+            
+            if available and self.retry_count < self.MAX_RETRIES:
+                best = max(available, key=lambda m: m.params_b)
+                for k, v in REGISTRY.items():
+                    if v == best:
+                        history.append(f"upgraded to {next_tier.value} with {k}")
+                        self.retry_count += 1
+                        return k, enhanced.enhanced, best, history
+        
+        # Level 4: Cloud fallback
+        if budget.allow_cloud:
+            cloud_models = by_tier(Tier.CLOUD_API)
+            if cloud_models:
+                best_cloud = max(cloud_models, key=lambda m: m.params_b)
+                for k, v in REGISTRY.items():
+                    if v == best_cloud:
+                        history.append(f"cloud fallback to {k}")
+                        self.retry_count += 1
+                        return k, enhanced.enhanced, best_cloud, history
+        
+        # Exhausted all options
+        history.append("all fallback options exhausted")
+        return previous_model_id, enhanced.enhanced, current_profile, history
+
+
+class ChimeraRouterV2:
+    """
+    Production router with telemetry integration and fallback controller.
+    """
+
+    def __init__(
+        self,
+        sulphur: Optional[Any] = None,
+        pinecone: Optional[Any] = None,
+        milvus: Optional[Any] = None,
+        ernie: Optional[Any] = None,
+        twave: Optional[TWAVETracker] = None,
+        ollama_host: str = "http://localhost:11434",
+        default_budget: Optional[QWAVEBudget] = None,
+        use_telemetry: bool = True,
+        detector_fusion: FusionMode = FusionMode.WEIGHTED,
+    ):
+        self.sulphur = sulphur or MockSulphurEnhancer()
+        self.pinecone = pinecone or MockPineconeRetriever()
+        self.milvus = milvus or MockMilvusRetriever()
+        self.ernie = ernie or MockERNIEAdapter()
+        self.twave = twave
+        self.ollama_host = ollama_host
+        self.default_budget = default_budget or QWAVEBudget()
+        self.use_telemetry = use_telemetry
+        self.detector_fusion = detector_fusion
+        
+        # Subsystems
+        self.cloud_manager = CloudAPIManager()
+        self._ollama_client: Optional[OllamaStreamingClient] = None
+        self._fallback: Optional[ModelFallbackController] = None
+
+    @property
+    def ollama_client(self) -> OllamaStreamingClient:
+        if self._ollama_client is None:
+            telemetry = None
+            if self.use_telemetry:
+                telemetry = OllamaTelemetryExtractor(
+                    ollama_host=self.ollama_host,
+                    embedding_model="functiongemma:latest",
+                    telemetry_interval=5,
+                )
+            self._ollama_client = OllamaStreamingClient(
+                ollama_host=self.ollama_host,
+                telemetry_extractor=telemetry,
+            )
+        return self._ollama_client
+
+    def _enhance(self, prompt: str) -> EnhancedPrompt:
+        return self.sulphur.enhance(prompt)
+
+    def _select_model(self, enhanced: EnhancedPrompt, budget: QWAVEBudget) -> Tuple[str, ModelProfile]:
+        required_caps = []
+        for tag in enhanced.intent_tags:
+            cap_map = {
+                "coding": Capability.CODING,
+                "reasoning": Capability.REASONING,
+                "vision": Capability.VISION,
+                "creative": Capability.INSTRUCT,
+                "factual": Capability.REASONING,
+                "safety": Capability.SAFETY,
+                "fast": Capability.FAST,
+                "long_context": Capability.LONG_CONTEXT,
+            }
+            if tag.lower() in cap_map:
+                required_caps.append(cap_map[tag.lower()])
+
+        if budget.require_safety:
+            exclude = [Capability.ABLITERATED, Capability.UNCHAINED]
+        elif budget.allow_uncensored:
+            exclude = []
+        else:
+            exclude = [Capability.ABLITERATED]
+
+        optimal_T_ratio = StochasticResonance.recommend_temperature(
+            enhanced.complexity_score,
+            T_c=1.0,
+        )
+
+        if enhanced.complexity_score > 0.8 and budget.allow_cloud:
+            preferred_tiers = [Tier.CLOUD_API, Tier.LOCAL_48GB, Tier.LOCAL_24GB]
+        elif enhanced.complexity_score > 0.6:
+            preferred_tiers = [Tier.LOCAL_24GB, Tier.LOCAL_16GB, Tier.CLOUD_API]
+        elif enhanced.complexity_score > 0.4:
+            preferred_tiers = [Tier.LOCAL_16GB, Tier.LOCAL_8GB]
+        else:
+            preferred_tiers = [Tier.LOCAL_8GB]
+
+        candidates = []
+        for tier in preferred_tiers:
+            tier_models = by_tier(tier)
+            for m in tier_models:
+                if required_caps and not all(c in m.capabilities for c in required_caps):
+                    continue
+                if any(c in m.capabilities for c in exclude):
+                    continue
+                if tier != Tier.CLOUD_API and m.size_gb > budget.vram_budget_gb:
+                    continue
+                candidates.append(m)
+
+        if not candidates:
+            all_models = list(REGISTRY.values())
+            candidates = [m for m in all_models if m.tier != Tier.CLOUD_API and m.size_gb <= budget.vram_budget_gb]
+            if not candidates and budget.allow_cloud:
+                candidates = by_tier(Tier.CLOUD_API)
+
+        if not candidates:
+            raise RuntimeError("No models available for this request.")
+
+        def score_model(m: ModelProfile) -> float:
+            s = 0.0
+            for cap in required_caps:
+                if cap in m.capabilities:
+                    s += 10.0
+            if enhanced.complexity_score > 0.7:
+                s += m.params_b * 2.0
+            else:
+                s += (10.0 - m.params_b) * 0.5
+            if Capability.FAST in m.capabilities and "fast" in enhanced.intent_tags:
+                s += 5.0
+            s -= m.size_gb * 0.1
+            temp_diff = abs(m.default_temp - optimal_T_ratio)
+            s -= temp_diff * 2.0
+            return s
+
+        candidates.sort(key=score_model, reverse=True)
+        best = candidates[0]
+
+        for k, v in REGISTRY.items():
+            if v == best:
+                return k, best
+
+        raise RuntimeError("Model selected but not found in registry.")
+
+    def _retrieve(self, query: str) -> Dict[str, Any]:
+        results = {"pinecone": [], "milvus": [], "ernie": [], "aggregated": []}
+
+        try:
+            results["pinecone"] = self.pinecone.get_evidence_for_ckplug(query)
+        except Exception as e:
+            results["pinecone_error"] = str(e)
+
+        try:
+            results["milvus"] = self.milvus.get_evidence("nexus_docs", query)
+        except Exception as e:
+            results["milvus_error"] = str(e)
+
+        try:
+            if self.ernie.is_available():
+                results["ernie"] = self.ernie.get_evidence(query)
+        except Exception as e:
+            results["ernie_error"] = str(e)
+
+        all_evidence = []
+        for src in [results["pinecone"], results["milvus"], results["ernie"]]:
+            for item in src:
+                all_evidence.append({
+                    "text": item.get("text", ""),
+                    "relevance": item.get("relevance", 0.0),
+                    "source": item.get("type", item.get("collection", "unknown")),
+                })
+
+        all_evidence.sort(key=lambda x: x["relevance"], reverse=True)
+        results["aggregated"] = all_evidence[:10]
+        results["top_score"] = all_evidence[0]["relevance"] if all_evidence else 0.0
+
+        return results
+
+    def _generate_with_telemetry(
+        self,
+        model_tag: str,
+        prompt: str,
+        profile: ModelProfile,
+        budget: QWAVEBudget,
+    ) -> Tuple[str, List[PerTokenDebug], GenerationTelemetry, Any]:
+        twave = TWAVETracker(
+            T_c=profile.T_c,
+            mu_0=profile.mu_base,
+            kappa=profile.kappa,
+        )
+        
+        temperature = profile.default_temp
+        
+        client = self.ollama_client
+        
+        response, tokens, trajectory = client.generate(
+            model_tag=model_tag,
+            prompt=prompt,
+            system="You are a helpful assistant. Use retrieved evidence when answering.",
+            temperature=temperature,
+            max_tokens=budget.max_tokens,
+        )
+        
+        debugs = client.telemetry.to_per_token_debug(
+            trajectory=trajectory,
+            twave=twave,
+            model_id=profile.name,
+            tier=profile.tier.value,
+        )
+        
+        gen_telemetry = GenerationTelemetry(
+            request_id=f"req_{int(time.time())}",
+            prompt=prompt,
+            tokens=debugs,
+            total_tokens=len(debugs),
+            selected_model=profile.name,
+            model_family=profile.family,
+            model_params_b=profile.params_b,
+            model_quantization=profile.quantization,
+        )
+        gen_telemetry.compute_aggregates()
+        
+        # Run unified detector for post-hoc analysis
+        detector = UnifiedThermodynamicDetector(
+            fusion_mode=self.detector_fusion,
+            enable_epr=True,
+            enable_spilled=True,
+            enable_ckplug=False,  # Post-hoc, no RAG context per token
+            enable_twave=True,
+        )
+        
+        token_verdicts = []
+        for i, debug in enumerate(debugs):
+            verdict = detector.evaluate_token(
+                position=i,
+                token_str=debug.token_str,
+                topk_probs=None,  # Not available post-hoc
+                token_id=debug.token_id or 0,
+                full_logits=None,
+                sampled_token_id=debug.token_id or 0,
+                probs_distribution=None,
+                log_prob_policy=0.0,
+                log_prob_ref=0.0,
+                visual_attention=debug.attention_mass_to_image or 1.0,
+                prev_psi=debugs[i-1].twave_psi if i > 0 else 0.0,
+            )
+            token_verdicts.append(verdict)
+        
+        sequence_verdict = detector.evaluate_sequence(token_verdicts)
+        
+        return response, debugs, gen_telemetry, sequence_verdict
+
+    def _generate_non_streaming(
+        self,
+        model_tag: str,
+        prompt: str,
+        profile: ModelProfile,
+        budget: QWAVEBudget,
+    ) -> str:
+        client = OllamaStreamingClient(ollama_host=self.ollama_host)
+        return client.generate_non_streaming(
+            model_tag=model_tag,
+            prompt=prompt,
+            system="You are a helpful assistant. Use retrieved evidence when answering.",
+            temperature=profile.default_temp,
+            max_tokens=budget.max_tokens,
+        )
+
+    def _generate_cloud(
+        self,
+        model_id: str,
+        prompt: str,
+        budget: QWAVEBudget,
+        profile: ModelProfile,
+    ) -> str:
+        """Generate via cloud API using the appropriate adapter."""
+        if not self.cloud_manager.is_available(profile.family):
+            return f"[CLOUD: {profile.name}] {prompt[:200]}... (no API key configured for {profile.family})"
+        
+        try:
+            response = self.cloud_manager.generate(
+                model_family=profile.family,
+                prompt=prompt,
+                max_tokens=budget.max_tokens,
+                temperature=profile.default_temp,
+                system="You are a helpful assistant. Use retrieved evidence when answering.",
+            )
+            return response.text
+        except RuntimeError as e:
+            return f"[CLOUD ERROR: {profile.name}] {str(e)}"
+
+    def route(
+        self,
+        prompt: str,
+        budget: Optional[QWAVEBudget] = None,
+        custom_model: Optional[str] = None,
+        use_telemetry: Optional[bool] = None,
+        max_retries: int = 3,
+    ) -> RouterResult:
+        """
+        Main routing entry point with full telemetry and fallback.
+        
+        Pipeline:
+          1. Enhance prompt
+          2. Select model
+          3. Retrieve evidence
+          4. Generate (with telemetry if available)
+          5. Evaluate risk
+          6. Fallback if needed
+          7. Assemble result
+        """
+        budget = budget or self.default_budget
+        use_telemetry = use_telemetry if use_telemetry is not None else self.use_telemetry
+        
+        # Step 1: Enhance
+        enhanced = self._enhance(prompt)
+        
+        # Step 2-3: Select model
+        if custom_model:
+            model_id = custom_model
+            profile = get(model_id)
+            if not profile:
+                raise ValueError(f"Unknown model: {custom_model}")
+        else:
+            model_id, profile = self._select_model(enhanced, budget)
+        
+        # Step 4: Retrieve
+        evidence = self._retrieve(enhanced.enhanced)
+        top_evidence = "\n".join([
+            f"[{e.get('source', 'unknown')}] {e.get('text', '')[:300]}"
+            for e in evidence.get("aggregated", [])[:3]
+        ])
+        
+        final_prompt = f"""Retrieved evidence:
+{top_evidence}
+
+---
+
+{enhanced.enhanced}"""
+        
+        # Step 5: Generate with fallback loop
+        response = ""
+        debugs = []
+        gen_telemetry = None
+        sequence_verdict = None
+        tokens_est = 0
+        fallback_history = []
+        
+        self._fallback = ModelFallbackController(self)
+        
+        for attempt in range(max_retries + 1):
+            t0 = time.time()
+            
+            try:
+                if profile.tier == Tier.CLOUD_API:
+                    response = self._generate_cloud(model_id, final_prompt, budget, profile)
+                    tokens_est = len(response.split())
+                elif use_telemetry and attempt == 0:
+                    # Try telemetry generation on first attempt
+                    response, debugs, gen_telemetry, sequence_verdict = self._generate_with_telemetry(
+                        profile.ollama_tag or model_id,
+                        final_prompt,
+                        profile,
+                        budget,
+                    )
+                    tokens_est = len(debugs)
+                else:
+                    # Fallback: non-streaming
+                    response = self._generate_non_streaming(
+                        profile.ollama_tag or model_id,
+                        final_prompt,
+                        profile,
+                        budget,
+                    )
+                    tokens_est = len(response.split())
+                
+                latency_ms = (time.time() - t0) * 1000
+                
+                # Evaluate risk
+                risk = 0.0
+                if sequence_verdict:
+                    risk = sequence_verdict.avg_fused_score
+                elif gen_telemetry:
+                    risk = gen_telemetry.hallucination_risk_score
+                else:
+                    risk = estimate_entropy_from_response(response)
+                
+                # Check if we need to escalate
+                if risk > 0.6 and attempt < max_retries:
+                    new_model, new_prompt, new_profile, history = self._fallback.escalate(
+                        final_prompt, enhanced, evidence, budget, model_id, risk
+                    )
+                    if new_model != model_id or new_prompt != final_prompt:
+                        model_id = new_model
+                        profile = new_profile
+                        final_prompt = new_prompt
+                        fallback_history.extend(history)
+                        continue
+                
+                # Success or exhausted retries
+                break
+                
+            except RuntimeError as e:
+                latency_ms = (time.time() - t0) * 1000
+                if attempt < max_retries:
+                    # Try fallback
+                    new_model, new_prompt, new_profile, history = self._fallback.escalate(
+                        final_prompt, enhanced, evidence, budget, model_id, 0.0
+                    )
+                    model_id = new_model
+                    profile = new_profile
+                    final_prompt = new_prompt
+                    fallback_history.extend([f"error: {str(e)}"] + history)
+                else:
+                    response = f"[ERROR] Generation failed after {max_retries} attempts: {str(e)}"
+                    break
+        
+        # Step 7: Assemble result
+        return RouterResult(
+            selected_model=model_id,
+            model_profile=profile,
+            tier=profile.tier.value,
+            enhanced_prompt=enhanced.enhanced,
+            response=response,
+            per_token_debug=debugs,
+            generation_telemetry=gen_telemetry,
+            reflection_count=gen_telemetry.reflection_count if gen_telemetry else 0,
+            grounding_score=evidence.get("top_score", 0.0),
+            hallucination_risk=sequence_verdict.avg_fused_score if sequence_verdict else estimate_entropy_from_response(response),
+            latency_ms=latency_ms,
+            tokens_generated=tokens_est,
+            detector_verdict=sequence_verdict,
+            fallback_history=fallback_history,
+            debug={
+                "enhancement": enhanced,
+                "evidence_summary": evidence,
+                "budget": budget,
+                "ckplug_epsilon": get_preset_epsilon(profile.family),
+                "optimal_temp_ratio": StochasticResonance.recommend_temperature(enhanced.complexity_score),
+                "fallback_attempts": len(fallback_history),
+            },
+        )
+
+    def quick_route(self, prompt: str, budget: Optional[QWAVEBudget] = None) -> str:
+        return self.route(prompt, budget=budget, use_telemetry=False).response
+
+    def get_available_models(self, budget: Optional[QWAVEBudget] = None) -> List[Dict[str, Any]]:
+        budget = budget or self.default_budget
+        available = []
+        for name, profile in REGISTRY.items():
+            fits = True
+            if profile.tier != Tier.CLOUD_API and profile.size_gb > budget.vram_budget_gb:
+                fits = False
+            if profile.tier == Tier.CLOUD_API and not budget.allow_cloud:
+                fits = False
+            available.append({
+                "id": name,
+                "name": profile.name,
+                "tier": profile.tier.value,
+                "size_gb": profile.size_gb,
+                "params_b": profile.params_b,
+                "capabilities": [c.value for c in profile.capabilities],
+                "fits_budget": fits,
+                "T_c": profile.T_c,
+                "mu_base": profile.mu_base,
+            })
+        return available