odin-engine 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

retrieval/cache.py

@@ -1,159 +1,159 @@
-"""
-Caching layer for GraphAccessor to prevent network hammering during PPR.
-
-PRODUCTION FIX: PPR's push algorithm repeatedly calls accessor.iter_out(u) for
-the same nodes, causing excessive network traffic. CachedGraphAccessor wraps any
-accessor and caches neighbor lookups.
-"""
-
-from __future__ import annotations
-from typing import Iterable, Optional, List, Tuple
-from collections import OrderedDict
-
-from .adapters import GraphAccessor, NodeId, RelId
-
-
-class CachedGraphAccessor:
-    """
-    Wraps a GraphAccessor with LRU caching for neighbor queries.
-    
-    Critical for production: Prevents "network hammer" issue where PPR
-    makes repeated calls to iter_out() for the same nodes, each hitting
-    the database/network.
-    
-    Usage:
-        base_accessor = ArangoCommunityAccessor(db, community_id="insurance")
-        cached_accessor = CachedGraphAccessor(base_accessor, cache_size=5000)
-        
-        # Now PPR won't hammer the network
-        orchestrator = RetrievalOrchestrator(accessor=cached_accessor, ...)
-    """
-    
-    def __init__(self, base: GraphAccessor, cache_size: int = 5000):
-        """
-        Args:
-            base: The underlying GraphAccessor to wrap
-            cache_size: Maximum number of nodes to cache (default: 5000)
-        """
-        self.base = base
-        self.cache_size = cache_size
-        
-        # LRU caches for outbound and inbound neighbors
-        self._out_cache: OrderedDict[NodeId, List[Tuple[NodeId, RelId, float]]] = OrderedDict()
-        self._in_cache: OrderedDict[NodeId, List[Tuple[NodeId, RelId, float]]] = OrderedDict()
-        
-        # Stats for monitoring
-        self._hits = 0
-        self._misses = 0
-    
-    def iter_out(self, node: NodeId) -> Iterable[Tuple[NodeId, RelId, float]]:
-        """Get outbound neighbors with caching."""
-        if node in self._out_cache:
-            # Cache hit - move to end (LRU)
-            self._out_cache.move_to_end(node)
-            self._hits += 1
-            return iter(self._out_cache[node])
-        
-        # Cache miss - fetch from base accessor
-        self._misses += 1
-        neighbors = list(self.base.iter_out(node))
-        
-        # Store in cache with LRU eviction
-        if len(self._out_cache) >= self.cache_size:
-            self._out_cache.popitem(last=False)  # Remove oldest
-        self._out_cache[node] = neighbors
-        
-        return iter(neighbors)
-    
-    def iter_in(self, node: NodeId) -> Iterable[Tuple[NodeId, RelId, float]]:
-        """Get inbound neighbors with caching."""
-        if node in self._in_cache:
-            # Cache hit - move to end (LRU)
-            self._in_cache.move_to_end(node)
-            self._hits += 1
-            return iter(self._in_cache[node])
-        
-        # Cache miss - fetch from base accessor
-        self._misses += 1
-        neighbors = list(self.base.iter_in(node))
-        
-        # Store in cache with LRU eviction
-        if len(self._in_cache) >= self.cache_size:
-            self._in_cache.popitem(last=False)  # Remove oldest
-        self._in_cache[node] = neighbors
-        
-        return iter(neighbors)
-    
-    def nodes(self, community_id: Optional[str] = None) -> Iterable[NodeId]:
-        """Pass through to base accessor (no caching)."""
-        return self.base.nodes(community_id)
-    
-    def get_node(self, node_id: NodeId, fields: Optional[List[str]] = None) -> dict:
-        """Pass through to base accessor (node lookups are typically one-off)."""
-        return self.base.get_node(node_id, fields)
-    
-    def degree(self, node: NodeId) -> int:
-        """Pass through to base accessor."""
-        return self.base.degree(node)
-    
-    def community_seed_norm(self, community_id: str, seeds: List[str]) -> List[str]:
-        """Pass through to base accessor."""
-        return self.base.community_seed_norm(community_id, seeds)
-    
-    def clear_cache(self):
-        """Clear all caches. Useful for memory management or testing."""
-        self._out_cache.clear()
-        self._in_cache.clear()
-        self._hits = 0
-        self._misses = 0
-    
-    def cache_stats(self) -> dict:
-        """
-        Return cache statistics for monitoring.
-        
-        Returns:
-            Dict with hit rate, sizes, and utilization metrics
-        """
-        total = self._hits + self._misses
-        hit_rate = self._hits / total if total > 0 else 0
-        
-        return {
-            "hits": self._hits,
-            "misses": self._misses,
-            "total_requests": total,
-            "hit_rate": hit_rate,
-            "out_cache_size": len(self._out_cache),
-            "in_cache_size": len(self._in_cache),
-            "max_cache_size": self.cache_size,
-            "out_cache_utilization": len(self._out_cache) / self.cache_size,
-            "in_cache_utilization": len(self._in_cache) / self.cache_size,
-        }
-    
-    def warm_cache(self, nodes: List[NodeId], direction: str = "out"):
-        """
-        Pre-populate cache for a list of nodes.
-        Useful for batch operations where you know which nodes will be accessed.
-        
-        Args:
-            nodes: List of node IDs to pre-fetch
-            direction: "out" or "in" for outbound/inbound neighbors
-        """
-        if direction == "out":
-            for node in nodes:
-                if node not in self._out_cache:
-                    neighbors = list(self.base.iter_out(node))
-                    if len(self._out_cache) >= self.cache_size:
-                        self._out_cache.popitem(last=False)
-                    self._out_cache[node] = neighbors
-        elif direction == "in":
-            for node in nodes:
-                if node not in self._in_cache:
-                    neighbors = list(self.base.iter_in(node))
-                    if len(self._in_cache) >= self.cache_size:
-                        self._in_cache.popitem(last=False)
-                    self._in_cache[node] = neighbors
-    
-    # Delegate all other methods to base accessor
-    def __getattr__(self, name):
-        """Delegate unknown methods to the base accessor."""
-        return getattr(self.base, name)
+"""
+Caching layer for GraphAccessor to prevent network hammering during PPR.
+
+PRODUCTION FIX: PPR's push algorithm repeatedly calls accessor.iter_out(u) for
+the same nodes, causing excessive network traffic. CachedGraphAccessor wraps any
+accessor and caches neighbor lookups.
+"""
+
+from __future__ import annotations
+from typing import Iterable, Optional, List, Tuple
+from collections import OrderedDict
+
+from .adapters import GraphAccessor, NodeId, RelId
+
+
+class CachedGraphAccessor:
+    """
+    Wraps a GraphAccessor with LRU caching for neighbor queries.
+    
+    Critical for production: Prevents "network hammer" issue where PPR
+    makes repeated calls to iter_out() for the same nodes, each hitting
+    the database/network.
+    
+    Usage:
+        base_accessor = ArangoCommunityAccessor(db, community_id="insurance")
+        cached_accessor = CachedGraphAccessor(base_accessor, cache_size=5000)
+        
+        # Now PPR won't hammer the network
+        orchestrator = RetrievalOrchestrator(accessor=cached_accessor, ...)
+    """
+    
+    def __init__(self, base: GraphAccessor, cache_size: int = 5000):
+        """
+        Args:
+            base: The underlying GraphAccessor to wrap
+            cache_size: Maximum number of nodes to cache (default: 5000)
+        """
+        self.base = base
+        self.cache_size = cache_size
+        
+        # LRU caches for outbound and inbound neighbors
+        self._out_cache: OrderedDict[NodeId, List[Tuple[NodeId, RelId, float]]] = OrderedDict()
+        self._in_cache: OrderedDict[NodeId, List[Tuple[NodeId, RelId, float]]] = OrderedDict()
+        
+        # Stats for monitoring
+        self._hits = 0
+        self._misses = 0
+    
+    def iter_out(self, node: NodeId) -> Iterable[Tuple[NodeId, RelId, float]]:
+        """Get outbound neighbors with caching."""
+        if node in self._out_cache:
+            # Cache hit - move to end (LRU)
+            self._out_cache.move_to_end(node)
+            self._hits += 1
+            return iter(self._out_cache[node])
+        
+        # Cache miss - fetch from base accessor
+        self._misses += 1
+        neighbors = list(self.base.iter_out(node))
+        
+        # Store in cache with LRU eviction
+        if len(self._out_cache) >= self.cache_size:
+            self._out_cache.popitem(last=False)  # Remove oldest
+        self._out_cache[node] = neighbors
+        
+        return iter(neighbors)
+    
+    def iter_in(self, node: NodeId) -> Iterable[Tuple[NodeId, RelId, float]]:
+        """Get inbound neighbors with caching."""
+        if node in self._in_cache:
+            # Cache hit - move to end (LRU)
+            self._in_cache.move_to_end(node)
+            self._hits += 1
+            return iter(self._in_cache[node])
+        
+        # Cache miss - fetch from base accessor
+        self._misses += 1
+        neighbors = list(self.base.iter_in(node))
+        
+        # Store in cache with LRU eviction
+        if len(self._in_cache) >= self.cache_size:
+            self._in_cache.popitem(last=False)  # Remove oldest
+        self._in_cache[node] = neighbors
+        
+        return iter(neighbors)
+    
+    def nodes(self, community_id: Optional[str] = None) -> Iterable[NodeId]:
+        """Pass through to base accessor (no caching)."""
+        return self.base.nodes(community_id)
+    
+    def get_node(self, node_id: NodeId, fields: Optional[List[str]] = None) -> dict:
+        """Pass through to base accessor (node lookups are typically one-off)."""
+        return self.base.get_node(node_id, fields)
+    
+    def degree(self, node: NodeId) -> int:
+        """Pass through to base accessor."""
+        return self.base.degree(node)
+    
+    def community_seed_norm(self, community_id: str, seeds: List[str]) -> List[str]:
+        """Pass through to base accessor."""
+        return self.base.community_seed_norm(community_id, seeds)
+    
+    def clear_cache(self):
+        """Clear all caches. Useful for memory management or testing."""
+        self._out_cache.clear()
+        self._in_cache.clear()
+        self._hits = 0
+        self._misses = 0
+    
+    def cache_stats(self) -> dict:
+        """
+        Return cache statistics for monitoring.
+        
+        Returns:
+            Dict with hit rate, sizes, and utilization metrics
+        """
+        total = self._hits + self._misses
+        hit_rate = self._hits / total if total > 0 else 0
+        
+        return {
+            "hits": self._hits,
+            "misses": self._misses,
+            "total_requests": total,
+            "hit_rate": hit_rate,
+            "out_cache_size": len(self._out_cache),
+            "in_cache_size": len(self._in_cache),
+            "max_cache_size": self.cache_size,
+            "out_cache_utilization": len(self._out_cache) / self.cache_size,
+            "in_cache_utilization": len(self._in_cache) / self.cache_size,
+        }
+    
+    def warm_cache(self, nodes: List[NodeId], direction: str = "out"):
+        """
+        Pre-populate cache for a list of nodes.
+        Useful for batch operations where you know which nodes will be accessed.
+        
+        Args:
+            nodes: List of node IDs to pre-fetch
+            direction: "out" or "in" for outbound/inbound neighbors
+        """
+        if direction == "out":
+            for node in nodes:
+                if node not in self._out_cache:
+                    neighbors = list(self.base.iter_out(node))
+                    if len(self._out_cache) >= self.cache_size:
+                        self._out_cache.popitem(last=False)
+                    self._out_cache[node] = neighbors
+        elif direction == "in":
+            for node in nodes:
+                if node not in self._in_cache:
+                    neighbors = list(self.base.iter_in(node))
+                    if len(self._in_cache) >= self.cache_size:
+                        self._in_cache.popitem(last=False)
+                    self._in_cache[node] = neighbors
+    
+    # Delegate all other methods to base accessor
+    def __getattr__(self, name):
+        """Delegate unknown methods to the base accessor."""
+        return getattr(self.base, name)

retrieval/confidence.py

@@ -1,88 +1,88 @@
-from __future__ import annotations
-from typing import Protocol, Tuple, Optional, Callable, List
-from collections import OrderedDict
-from math import isfinite
-import torch
-
-NodeId = str
-RelId = str
-
-
-class EdgeConfidenceProvider(Protocol):
-    def confidence(self, u: NodeId, rel: RelId, v: NodeId) -> float: ...
-    def confidence_batch(self, edges: List[Tuple[NodeId, RelId, NodeId]]) -> List[float]: ...
-
-
-class ConstantConfidence:
-    def __init__(self, value: float = 0.8):
-        self.value = max(1e-6, min(1.0, value))
-
-    def confidence(self, u, rel, v) -> float:
-        return self.value
-    
-    def confidence_batch(self, edges: List[Tuple[NodeId, RelId, NodeId]]) -> List[float]:
-        return [self.value] * len(edges)
-
-
-class NPLLConfidence:
-    """
-    Wraps an NPLL model for retrieval-time scoring with LRU caching.
-    Uses the model's scoring module for batched probabilities when available.
-    
-    PRODUCTION FIX: Uses bounded LRU cache to prevent memory leaks in long-running processes.
-    """
-
-    def __init__(self, npll_model, cache_size: int = 10000):
-        """
-        Args:
-            npll_model: Trained NPLL model for scoring triples
-            cache_size: Maximum number of cached confidence scores (default: 10K)
-        """
-        self.model = npll_model
-        self.cache_size = cache_size
-        self._cache: OrderedDict[Tuple[str, str, str], float] = OrderedDict()
-
-    def confidence(self, u: NodeId, rel: RelId, v: NodeId) -> float:
-        return self.confidence_batch([(u, rel, v)])[0]
-
-    def confidence_batch(self, edges: List[Tuple[NodeId, RelId, NodeId]]) -> List[float]:
-        todo = [(u, r, v) for (u, r, v) in edges if (u, r, v) not in self._cache]
-        if todo:
-            heads, rels, tails = zip(*todo)
-            self.model.eval()
-            with torch.no_grad():
-                scores = self.model.scoring_module.forward_with_names(list(heads), list(rels), list(tails))
-                # Don't apply per-group temperature scaling (requires group_ids we don't have)
-                probs = self.model.probability_transform(scores, apply_temperature=False)
-            for (u, r, v), p in zip(todo, probs.tolist()):
-                confidence = max(1e-6, float(p)) if isfinite(p) else 1e-6
-                
-                # LRU eviction: remove oldest if at capacity
-                if len(self._cache) >= self.cache_size:
-                    self._cache.popitem(last=False)  # Remove oldest (FIFO)
-                
-                self._cache[(u, r, v)] = confidence
-        
-        # Move accessed items to end (LRU behavior)
-        result = []
-        for edge in edges:
-            conf = self._cache[edge]
-            # Move to end to mark as recently used
-            self._cache.move_to_end(edge)
-            result.append(conf)
-        
-        return result
-
-    def clear_cache(self):
-        """Clear the confidence cache. Useful for testing or memory management."""
-        self._cache.clear()
-
-    def cache_stats(self) -> dict:
-        """Return cache statistics for monitoring."""
-        return {
-            "size": len(self._cache),
-            "max_size": self.cache_size,
-            "utilization": len(self._cache) / self.cache_size if self.cache_size > 0 else 0,
-        }
-
-
+from __future__ import annotations
+from typing import Protocol, Tuple, Optional, Callable, List
+from collections import OrderedDict
+from math import isfinite
+import torch
+
+NodeId = str
+RelId = str
+
+
+class EdgeConfidenceProvider(Protocol):
+    def confidence(self, u: NodeId, rel: RelId, v: NodeId) -> float: ...
+    def confidence_batch(self, edges: List[Tuple[NodeId, RelId, NodeId]]) -> List[float]: ...
+
+
+class ConstantConfidence:
+    def __init__(self, value: float = 0.8):
+        self.value = max(1e-6, min(1.0, value))
+
+    def confidence(self, u, rel, v) -> float:
+        return self.value
+    
+    def confidence_batch(self, edges: List[Tuple[NodeId, RelId, NodeId]]) -> List[float]:
+        return [self.value] * len(edges)
+
+
+class NPLLConfidence:
+    """
+    Wraps an NPLL model for retrieval-time scoring with LRU caching.
+    Uses the model's scoring module for batched probabilities when available.
+    
+    PRODUCTION FIX: Uses bounded LRU cache to prevent memory leaks in long-running processes.
+    """
+
+    def __init__(self, npll_model, cache_size: int = 10000):
+        """
+        Args:
+            npll_model: Trained NPLL model for scoring triples
+            cache_size: Maximum number of cached confidence scores (default: 10K)
+        """
+        self.model = npll_model
+        self.cache_size = cache_size
+        self._cache: OrderedDict[Tuple[str, str, str], float] = OrderedDict()
+
+    def confidence(self, u: NodeId, rel: RelId, v: NodeId) -> float:
+        return self.confidence_batch([(u, rel, v)])[0]
+
+    def confidence_batch(self, edges: List[Tuple[NodeId, RelId, NodeId]]) -> List[float]:
+        todo = [(u, r, v) for (u, r, v) in edges if (u, r, v) not in self._cache]
+        if todo:
+            heads, rels, tails = zip(*todo)
+            self.model.eval()
+            with torch.no_grad():
+                scores = self.model.scoring_module.forward_with_names(list(heads), list(rels), list(tails))
+                # Don't apply per-group temperature scaling (requires group_ids we don't have)
+                probs = self.model.probability_transform(scores, apply_temperature=False)
+            for (u, r, v), p in zip(todo, probs.tolist()):
+                confidence = max(1e-6, float(p)) if isfinite(p) else 1e-6
+                
+                # LRU eviction: remove oldest if at capacity
+                if len(self._cache) >= self.cache_size:
+                    self._cache.popitem(last=False)  # Remove oldest (FIFO)
+                
+                self._cache[(u, r, v)] = confidence
+        
+        # Move accessed items to end (LRU behavior)
+        result = []
+        for edge in edges:
+            conf = self._cache[edge]
+            # Move to end to mark as recently used
+            self._cache.move_to_end(edge)
+            result.append(conf)
+        
+        return result
+
+    def clear_cache(self):
+        """Clear the confidence cache. Useful for testing or memory management."""
+        self._cache.clear()
+
+    def cache_stats(self) -> dict:
+        """Return cache statistics for monitoring."""
+        return {
+            "size": len(self._cache),
+            "max_size": self.cache_size,
+            "utilization": len(self._cache) / self.cache_size if self.cache_size > 0 else 0,
+        }
+
+

retrieval/eval.py

@@ -1,49 +1,49 @@
-from __future__ import annotations
-from typing import List, Tuple, Set
-from sklearn.isotonic import IsotonicRegression
-import numpy as np
-import numpy as np
-
-
-def recall_at_k(predicted: List[Tuple[str, float]], relevant: Set[str], k: int = 10) -> float:
-    if not predicted or not relevant:
-        return 0.0
-    top = [n for n, _ in predicted[:k]]
-    hits = sum(1 for n in top if n in relevant)
-    denom = min(k, len(relevant)) or 1
-    return hits / denom
-
-
-def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
-    bins = np.linspace(0.0, 1.0, n_bins + 1)
-    ece = 0.0
-    N = len(probs)
-    for i in range(n_bins):
-        lo, hi = bins[i], bins[i + 1]
-        mask = (probs > lo) & (probs <= hi)
-        if not np.any(mask):
-            continue
-        acc = labels[mask].mean()
-        conf = probs[mask].mean()
-        ece += np.abs(conf - acc) * (mask.sum() / max(N, 1))
-    return float(ece)
-
-
-class SimpleLLMCalibrator:
-    """
-    Wraps isotonic regression to calibrate LLM self-reported confidences.
-    """
-    def __init__(self):
-        self.iso = IsotonicRegression(out_of_bounds='clip')
-        self.is_fitted = False
-
-    def fit(self, raw_conf: np.ndarray, labels: np.ndarray):
-        self.iso.fit(raw_conf, labels)
-        self.is_fitted = True
-
-    def transform(self, raw_conf: np.ndarray) -> np.ndarray:
-        if not self.is_fitted:
-            return raw_conf
-        return self.iso.transform(raw_conf)
-
-
+from __future__ import annotations
+from typing import List, Tuple, Set
+from sklearn.isotonic import IsotonicRegression
+import numpy as np
+import numpy as np
+
+
+def recall_at_k(predicted: List[Tuple[str, float]], relevant: Set[str], k: int = 10) -> float:
+    if not predicted or not relevant:
+        return 0.0
+    top = [n for n, _ in predicted[:k]]
+    hits = sum(1 for n in top if n in relevant)
+    denom = min(k, len(relevant)) or 1
+    return hits / denom
+
+
+def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
+    bins = np.linspace(0.0, 1.0, n_bins + 1)
+    ece = 0.0
+    N = len(probs)
+    for i in range(n_bins):
+        lo, hi = bins[i], bins[i + 1]
+        mask = (probs > lo) & (probs <= hi)
+        if not np.any(mask):
+            continue
+        acc = labels[mask].mean()
+        conf = probs[mask].mean()
+        ece += np.abs(conf - acc) * (mask.sum() / max(N, 1))
+    return float(ece)
+
+
+class SimpleLLMCalibrator:
+    """
+    Wraps isotonic regression to calibrate LLM self-reported confidences.
+    """
+    def __init__(self):
+        self.iso = IsotonicRegression(out_of_bounds='clip')
+        self.is_fitted = False
+
+    def fit(self, raw_conf: np.ndarray, labels: np.ndarray):
+        self.iso.fit(raw_conf, labels)
+        self.is_fitted = True
+
+    def transform(self, raw_conf: np.ndarray) -> np.ndarray:
+        if not self.is_fitted:
+            return raw_conf
+        return self.iso.transform(raw_conf)
+
+