Copy nexus_os_v2/ollama_telemetry.py from dataset for module imports

nexus_os_v2/ollama_telemetry.py

@@ -0,0 +1,458 @@
+"""
+Ollama Telemetry Extractor for NEXUS OS v2
+Extracts per-token thermodynamic order parameters from Ollama generation.
+
+Since Ollama does not expose raw logits, we use a dual-tier approach:
+  Tier 1 (fast): Token-level surface metrics from stream timing + text patterns
+  Tier 2 (deep): Embedding-space trajectory divergence as entropy proxy
+
+The embedding divergence correlates with generation coherence:
+  - Smooth trajectory -> low entropy -> coherent (condensate phase)
+  - Sharp jumps -> high entropy -> bifurcation risk (near T_c)
+
+Uses Ollama's /api/embeddings endpoint for vector representations.
+"""
+import json
+import math
+import time
+import urllib.request
+import urllib.error
+from typing import List, Dict, Optional, Any, Tuple
+from dataclasses import dataclass, field
+
+# Try numpy, fallback to pure Python
+import sys
+HAS_NUMPY = True
+try:
+    import numpy as np
+except ImportError:
+    HAS_NUMPY = False
+
+
+@dataclass
+class TokenTelemetry:
+    """Surface-level metrics for a single generated token."""
+    position: int
+    token_text: str
+    timestamp_ms: float              # Relative to generation start
+    time_since_prev_ms: float        # Inter-token latency
+    char_length: int
+    is_punctuation: bool
+    is_whitespace: bool
+    repetition_count: int            # How many times this token appeared recently
+    embedding: Optional[List[float]] = None  # Ollama embedding vector
+
+
+@dataclass
+class EmbeddingTrajectory:
+    """Sequence of embeddings with computed divergences."""
+    tokens: List[TokenTelemetry]
+    divergences: List[float]         # L2 distance between consecutive embeddings
+    cosine_similarities: List[float]
+    trajectory_curvature: List[float]  # Angle change in embedding space
+    cumulative_drift: List[float]    # Running sum of divergences
+    
+    def entropy_proxy(self, position: int) -> float:
+        """Compute entropy proxy from embedding divergence at position."""
+        if position < 1 or position >= len(self.divergences):
+            return 0.0
+        div = self.divergences[position]
+        cum = self.cumulative_drift[position] if position < len(self.cumulative_drift) else div
+        if cum < 1e-6:
+            return 0.0
+        return min(1.0, div / (cum / max(1, position + 1)))
+
+
+class OllamaTelemetryExtractor:
+    """
+    Extract per-token telemetry from Ollama generation streams.
+    
+    Dual-tier architecture:
+      Tier 1: Fast surface metrics from stream (timing, repetition)
+      Tier 2: Deep embedding-space analysis via /api/embeddings
+    """
+    
+    def __init__(
+        self,
+        ollama_host: str = "http://localhost:11434",
+        embedding_model: str = "functiongemma:latest",
+        telemetry_interval: int = 5,
+        history_window: int = 20,
+    ):
+        self.ollama_host = ollama_host.rstrip("/")
+        self.embedding_model = embedding_model
+        self.telemetry_interval = telemetry_interval
+        self.history_window = history_window
+        self._token_buffer: List[TokenTelemetry] = []
+        self._text_buffer: str = ""
+        self._start_time: Optional[float] = None
+        self._embedding_cache: Dict[str, List[float]] = {}
+        
+    def _get_embedding(self, text: str) -> Optional[List[float]]:
+        if text in self._embedding_cache:
+            return self._embedding_cache[text]
+        
+        payload = json.dumps({
+            "model": self.embedding_model,
+            "prompt": text,
+        }).encode("utf-8")
+        
+        req = urllib.request.Request(
+            f"{self.ollama_host}/api/embeddings",
+            data=payload,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        
+        try:
+            with urllib.request.urlopen(req, timeout=30) as resp:
+                data = json.loads(resp.read().decode("utf-8"))
+                embedding = data.get("embedding")
+                if embedding:
+                    self._embedding_cache[text] = embedding
+                return embedding
+        except Exception:
+            return None
+    
+    def _compute_repetition(self, token_text: str) -> int:
+        recent = [t.token_text for t in self._token_buffer[-self.history_window:]]
+        return recent.count(token_text)
+    
+    def on_token(self, token_text: str, position: int) -> TokenTelemetry:
+        now = time.time()
+        if self._start_time is None:
+            self._start_time = now
+        
+        elapsed_ms = (now - self._start_time) * 1000
+        prev_time = self._token_buffer[-1].timestamp_ms if self._token_buffer else elapsed_ms
+        time_since_prev = elapsed_ms - prev_time
+        
+        self._text_buffer += token_text
+        
+        embedding = None
+        if position % self.telemetry_interval == 0 and position > 0:
+            embedding = self._get_embedding(self._text_buffer)
+        
+        telemetry = TokenTelemetry(
+            position=position,
+            token_text=token_text,
+            timestamp_ms=elapsed_ms,
+            time_since_prev_ms=time_since_prev,
+            char_length=len(token_text),
+            is_punctuation=token_text.strip() in ".,;:!?",
+            is_whitespace=token_text.strip() == "",
+            repetition_count=self._compute_repetition(token_text),
+            embedding=embedding,
+        )
+        
+        self._token_buffer.append(telemetry)
+        return telemetry
+    
+    def build_embedding_trajectory(self) -> EmbeddingTrajectory:
+        """Compute embedding-space trajectory metrics."""
+        embeddings = [t.embedding for t in self._token_buffer if t.embedding is not None]
+        
+        if len(embeddings) < 2:
+            return EmbeddingTrajectory(
+                tokens=self._token_buffer,
+                divergences=[],
+                cosine_similarities=[],
+                trajectory_curvature=[],
+                cumulative_drift=[],
+            )
+        
+        if HAS_NUMPY:
+            emb_array = np.array(embeddings)
+            divergences = []
+            for i in range(1, len(emb_array)):
+                div = float(np.linalg.norm(emb_array[i] - emb_array[i-1]))
+                divergences.append(div)
+            cos_sims = []
+            for i in range(1, len(emb_array)):
+                dot = np.dot(emb_array[i], emb_array[i-1])
+                norm = np.linalg.norm(emb_array[i]) * np.linalg.norm(emb_array[i-1])
+                cos_sims.append(float(dot / norm) if norm > 0 else 0.0)
+            curvature = []
+            for i in range(2, len(emb_array)):
+                v1 = emb_array[i-1] - emb_array[i-2]
+                v2 = emb_array[i] - emb_array[i-1]
+                cross = np.linalg.norm(np.cross(v1, v2))
+                dot = np.dot(v1, v2)
+                angle = math.atan2(cross, dot) if len(v1) == 3 else 0.0
+                curvature.append(float(angle))
+        else:
+            # Pure Python fallback
+            divergences = []
+            for i in range(1, len(embeddings)):
+                div = math.sqrt(sum((a - b) ** 2 for a, b in zip(embeddings[i], embeddings[i-1])))
+                divergences.append(div)
+            cos_sims = []
+            for i in range(1, len(embeddings)):
+                dot = sum(a * b for a, b in zip(embeddings[i], embeddings[i-1]))
+                norm1 = math.sqrt(sum(a * a for a in embeddings[i]))
+                norm2 = math.sqrt(sum(a * a for a in embeddings[i-1]))
+                cos_sims.append(dot / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0)
+            curvature = []
+        
+        cumulative = []
+        total = 0.0
+        for div in divergences:
+            total += div
+            cumulative.append(total)
+        
+        return EmbeddingTrajectory(
+            tokens=self._token_buffer,
+            divergences=divergences,
+            cosine_similarities=cos_sims,
+            trajectory_curvature=curvature,
+            cumulative_drift=cumulative,
+        )
+    
+    def compute_surface_entropy(self, token: TokenTelemetry) -> float:
+        """
+        Fast heuristic entropy from surface metrics.
+        Maps timing + repetition patterns to entropy proxy.
+        """
+        entropy = 0.0
+        
+        if token.time_since_prev_ms > 200:
+            entropy += 0.2
+        if token.time_since_prev_ms > 500:
+            entropy += 0.3
+        
+        if token.repetition_count >= 3:
+            entropy -= 0.3
+        
+        if token.is_punctuation and token.time_since_prev_ms > 100:
+            entropy += 0.1
+        
+        return max(0.0, min(1.0, entropy))
+    
+    def to_per_token_debug(self, trajectory, twave, model_id: str, tier: str):
+        """Convert telemetry to PerTokenDebug schema."""
+        from .per_token_debug import PerTokenDebug
+        debugs = []
+        
+        for i, token in enumerate(self._token_buffer):
+            surface_H = self.compute_surface_entropy(token)
+            
+            embedding_H = 0.0
+            if i > 0 and i <= len(trajectory.divergences):
+                embedding_H = trajectory.entropy_proxy(i)
+            
+            H = 0.3 * surface_H + 0.7 * embedding_H if embedding_H > 0 else surface_H
+            H = max(0.0, min(1.0, H))
+            
+            coherence = twave.compute_coherence(H * twave.H_max)
+            T_eff = twave.T_c * coherence
+            
+            CG = 0.0
+            mu_ret = twave.compute_chemical_potential(CG)
+            
+            psi = twave.compute_order_parameter(coherence, mu_ret)
+            f_density = twave.compute_free_energy_density(psi, coherence, mu_ret)
+            
+            prev_psi = debugs[-1].twave_psi if debugs else 0.0
+            k_local = abs(psi - prev_psi) if prev_psi else 0.0
+            E_exc = twave.compute_bogoliubov_energy(psi, k_local, mu_ret)
+            
+            debug = PerTokenDebug(
+                position=i,
+                token_id=0,
+                token_str=token.token_text,
+                entropy=H * twave.H_max,
+                entropy_normalized=H,
+                twave_T_eff=T_eff,
+                twave_coherence=coherence,
+                twave_psi=psi,
+                twave_f_density=f_density,
+                twave_mu_ret=mu_ret,
+                twave_E_exc=E_exc,
+                twave_k_local=k_local,
+                generation_time_ms=token.time_since_prev_ms,
+                model_id=model_id,
+                tier=tier,
+            )
+            debugs.append(debug)
+        
+        for i in range(len(debugs)):
+            if i >= 2:
+                ent_history = [d.entropy for d in debugs[max(0, i-2):i+1]]
+                debugs[i].twave_C_V = twave.compute_specific_heat(ent_history)
+            debugs[i].jarzynski_W_i = 0.0
+        
+        return debugs
+    
+    def reset(self):
+        self._token_buffer = []
+        self._text_buffer = ""
+        self._start_time = None
+
+
+class OllamaStreamingClient:
+    """
+    Production Ollama client with streaming + telemetry extraction.
+    """
+    
+    def __init__(
+        self,
+        ollama_host: str = "http://localhost:11434",
+        telemetry_extractor: Optional[OllamaTelemetryExtractor] = None,
+    ):
+        self.ollama_host = ollama_host.rstrip("/")
+        self.telemetry = telemetry_extractor or OllamaTelemetryExtractor(ollama_host)
+    
+    def generate(
+        self,
+        model_tag: str,
+        prompt: str,
+        system: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: int = 2048,
+        top_p: float = 0.95,
+        stream_callback=None,
+    ) -> Tuple[str, List[TokenTelemetry], EmbeddingTrajectory]:
+        """
+        Generate text via Ollama with full telemetry extraction.
+        
+        Returns: (full_text, token_telemetry_list, embedding_trajectory)
+        """
+        self.telemetry.reset()
+        
+        messages = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+        
+        payload = json.dumps({
+            "model": model_tag,
+            "messages": messages,
+            "stream": True,
+            "options": {
+                "temperature": temperature,
+                "num_predict": max_tokens,
+                "top_p": top_p,
+            },
+        }).encode("utf-8")
+        
+        req = urllib.request.Request(
+            f"{self.ollama_host}/api/chat",
+            data=payload,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        
+        full_text = ""
+        position = 0
+        
+        try:
+            with urllib.request.urlopen(req, timeout=300) as resp:
+                for line in resp:
+                    if not line:
+                        continue
+                    try:
+                        data = json.loads(line.decode("utf-8"))
+                    except json.JSONDecodeError:
+                        continue
+                    
+                    if data.get("done", False):
+                        break
+                    
+                    token_text = data.get("message", {}).get("content", "")
+                    if not token_text:
+                        continue
+                    
+                    full_text += token_text
+                    
+                    telemetry = self.telemetry.on_token(token_text, position)
+                    position += 1
+                    
+                    if stream_callback:
+                        stream_callback(token_text, telemetry)
+                    
+        except urllib.error.URLError as e:
+            raise RuntimeError(f"Ollama connection failed: {e}")
+        
+        trajectory = self.telemetry.build_embedding_trajectory()
+        
+        return full_text, self.telemetry._token_buffer, trajectory
+    
+    def generate_non_streaming(
+        self,
+        model_tag: str,
+        prompt: str,
+        system: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: int = 2048,
+        top_p: float = 0.95,
+    ) -> str:
+        """Simple non-streaming generation (faster, no telemetry)."""
+        messages = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+        
+        payload = json.dumps({
+            "model": model_tag,
+            "messages": messages,
+            "stream": False,
+            "options": {
+                "temperature": temperature,
+                "num_predict": max_tokens,
+                "top_p": top_p,
+            },
+        }).encode("utf-8")
+        
+        req = urllib.request.Request(
+            f"{self.ollama_host}/api/chat",
+            data=payload,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        
+        try:
+            with urllib.request.urlopen(req, timeout=300) as resp:
+                data = json.loads(resp.read().decode("utf-8"))
+                return data.get("message", {}).get("content", "")
+        except urllib.error.URLError as e:
+            raise RuntimeError(f"Ollama connection failed: {e}")
+
+
+def estimate_entropy_from_response(response: str, chunk_size: int = 10) -> float:
+    """
+    Post-hoc entropy estimation from completed response.
+    Uses lexical diversity as a proxy for generation entropy.
+    
+    Returns: Estimated normalized entropy [0, 1]
+    """
+    if not response:
+        return 0.0
+    
+    words = response.split()
+    if not words:
+        return 0.0
+    
+    unique_words = len(set(w.lower() for w in words))
+    lexical_diversity = unique_words / len(words)
+    
+    from collections import Counter
+    word_counts = Counter(w.lower() for w in words)
+    max_repeat = max(word_counts.values()) if word_counts else 1
+    repetition_penalty = min(1.0, max_repeat / max(1, len(words) * 0.1))
+    
+    sentences = response.split(".")
+    sentence_lengths = [len(s.split()) for s in sentences if s.strip()]
+    if len(sentence_lengths) > 1:
+        mean_len = sum(sentence_lengths) / len(sentence_lengths)
+        variance = sum((x - mean_len) ** 2 for x in sentence_lengths) / len(sentence_lengths)
+        length_variance = variance / max(1, mean_len ** 2)
+    else:
+        length_variance = 0.0
+    
+    entropy = (
+        0.5 * lexical_diversity +
+        0.3 * (1.0 - repetition_penalty) +
+        0.2 * min(1.0, length_variance)
+    )
+    
+    return min(1.0, max(0.0, entropy))