sharp-context 0.1.0__py3-none-any.whl

sharp_context/prefetch.py

@@ -0,0 +1,297 @@
+"""
+Predictive Context Pre-fetcher
+===============================
+
+When an agent reads a code symbol, predict what it will need next
+and pre-load it into the context cache — BEFORE it asks.
+
+This is **CPU cache prefetching applied to LLM context windows**.
+
+The Problem:
+  An agent debugging function `process_payment()` will inevitably need:
+    1. Callers of `process_payment()` — who triggers this?
+    2. Callees from `process_payment()` — what does it depend on?
+    3. Test file for `process_payment()` — how is it tested?
+    4. Type definitions used — what are the data structures?
+
+  Without pre-fetching, the agent makes 4 sequential tool calls,
+  each adding latency and token cost. With pre-fetching, these are
+  already in the context cache when the agent asks.
+
+Heuristics:
+  1. **Static call graph**: Extract function/method calls from source
+  2. **Import graph**: Follow import statements to related modules
+  3. **Naming conventions**: foo.py → test_foo.py, foo_test.py
+  4. **Co-access patterns**: Track which files are accessed together
+     across sessions (associative learning)
+
+References:
+  - CPU prefetch: Smith, J. "Sequential Program Prefetching" (1978)
+  - Agentic Plan Caching (arXiv 2025) — reusing structured plans
+  - Proximity (arXiv 2026) — LSH-bucketed pre-warming for caches
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from collections import Counter, defaultdict
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Set, Tuple
+
+
+@dataclass
+class PrefetchResult:
+    """A predicted context fragment that might be needed next."""
+
+    path: str
+    """File path or symbol identifier."""
+
+    reason: str
+    """Why this was predicted (e.g., 'callee', 'test_file', 'co_access')."""
+
+    confidence: float
+    """Prediction confidence [0, 1]."""
+
+    content: Optional[str] = None
+    """Pre-loaded content (if available)."""
+
+
+# ── Static Analysis Patterns ───────────────────────────────────────────
+
+# Python function/method calls
+_PY_CALL_RE = re.compile(
+    r"(?:self\.)?([a-zA-Z_]\w*)\s*\(", re.MULTILINE
+)
+
+# Python imports
+_PY_IMPORT_RE = re.compile(
+    r"(?:from\s+([\w.]+)\s+import|import\s+([\w.]+))", re.MULTILINE
+)
+
+# Python class inheritance
+_PY_CLASS_RE = re.compile(
+    r"class\s+\w+\s*\(\s*([\w.,\s]+)\s*\)\s*:", re.MULTILINE
+)
+
+# Rust use/mod statements
+_RS_USE_RE = re.compile(
+    r"(?:use\s+([\w:]+)|mod\s+(\w+))", re.MULTILINE
+)
+
+# TypeScript/JavaScript imports
+_TS_IMPORT_RE = re.compile(
+    r"(?:import\s+.*?\s+from\s+['\"]([^'\"]+)['\"]|"
+    r"require\s*\(\s*['\"]([^'\"]+)['\"]\s*\))",
+    re.MULTILINE,
+)
+
+
+def extract_callees(source: str, language: str = "python") -> List[str]:
+    """
+    Extract function/method names called from a source code fragment.
+
+    Returns a list of callee names (not fully qualified — just the
+    function name as it appears in source).
+    """
+    if language == "python":
+        return _PY_CALL_RE.findall(source)
+    return []
+
+
+def extract_imports(source: str, language: str = "python") -> List[str]:
+    """
+    Extract import targets from a source code fragment.
+
+    Returns module/path strings that could be resolved to files.
+    """
+    if language == "python":
+        results = []
+        for match in _PY_IMPORT_RE.finditer(source):
+            mod = match.group(1) or match.group(2)
+            if mod:
+                results.append(mod)
+        return results
+    elif language in ("typescript", "javascript"):
+        results = []
+        for match in _TS_IMPORT_RE.finditer(source):
+            path = match.group(1) or match.group(2)
+            if path:
+                results.append(path)
+        return results
+    elif language == "rust":
+        results = []
+        for match in _RS_USE_RE.finditer(source):
+            mod = match.group(1) or match.group(2)
+            if mod:
+                results.append(mod)
+        return results
+    return []
+
+
+def infer_test_files(file_path: str) -> List[str]:
+    """
+    Infer likely test file paths from a source file path.
+
+    Heuristics:
+      foo.py       → test_foo.py, foo_test.py, tests/test_foo.py
+      utils/bar.py → tests/test_bar.py, utils/test_bar.py
+      src/baz.rs   → tests/baz.rs, src/baz_test.rs
+    """
+    path = Path(file_path)
+    stem = path.stem
+    suffix = path.suffix
+    parent = path.parent
+
+    candidates = [
+        str(parent / f"test_{stem}{suffix}"),
+        str(parent / f"{stem}_test{suffix}"),
+        str(parent / "tests" / f"test_{stem}{suffix}"),
+        str(parent.parent / "tests" / f"test_{stem}{suffix}"),
+    ]
+
+    # Rust-specific
+    if suffix == ".rs":
+        candidates.append(str(parent / "tests" / f"{stem}{suffix}"))
+
+    return candidates
+
+
+def module_to_file_candidates(
+    module_path: str,
+    base_dir: str = "",
+    language: str = "python",
+) -> List[str]:
+    """
+    Convert a module path (e.g., 'utils.helpers') to candidate file paths.
+
+    Python: utils.helpers → utils/helpers.py, utils/helpers/__init__.py
+    """
+    if language == "python":
+        parts = module_path.split(".")
+        candidates = [
+            os.path.join(base_dir, *parts) + ".py",
+            os.path.join(base_dir, *parts, "__init__.py"),
+        ]
+        return candidates
+    return []
+
+
+class PrefetchEngine:
+    """
+    Predictive pre-fetcher that learns co-access patterns across
+    sessions and combines them with static analysis for predictions.
+
+    Two prediction strategies:
+      1. **Static**: Parse imports, calls, and naming conventions
+      2. **Learned**: Track which files are accessedtogether and
+         predict based on historical co-access frequency
+
+    The learned component uses a simple co-occurrence counter:
+      When file A and file B are accessed within K turns of each other,
+      increment co_access[A][B] and co_access[B][A].
+
+    Confidence is computed as:
+      - Static predictions: fixed confidence (0.7 for imports, 0.5 for tests)
+      - Learned predictions: normalized co-access count
+    """
+
+    def __init__(self, co_access_window: int = 5):
+        self.co_access_window = co_access_window
+
+        # co_access[file_a][file_b] = count of times accessed together
+        self._co_access: Dict[str, Counter] = defaultdict(Counter)
+
+        # Recent access history for learning
+        self._recent_accesses: List[Tuple[str, int]] = []  # (path, turn)
+
+    def record_access(self, file_path: str, turn: int) -> None:
+        """
+        Record that a file was accessed at a given turn.
+
+        Updates co-access counts with all files accessed within
+        the co-access window.
+        """
+        # Update co-access with recent files
+        for prev_path, prev_turn in self._recent_accesses:
+            if turn - prev_turn <= self.co_access_window and prev_path != file_path:
+                self._co_access[file_path][prev_path] += 1
+                self._co_access[prev_path][file_path] += 1
+
+        self._recent_accesses.append((file_path, turn))
+
+        # Prune old accesses (keep last 100)
+        if len(self._recent_accesses) > 100:
+            self._recent_accesses = self._recent_accesses[-100:]
+
+    def predict(
+        self,
+        file_path: str,
+        source_content: str,
+        language: str = "python",
+        max_results: int = 10,
+    ) -> List[PrefetchResult]:
+        """
+        Predict what context fragments will be needed next, given
+        that the agent just accessed `file_path` with `source_content`.
+
+        Combines static analysis and learned co-access patterns.
+        Results are sorted by confidence (highest first).
+        """
+        predictions: List[PrefetchResult] = []
+        seen_paths: Set[str] = set()
+
+        # 1. Import graph (confidence: 0.7)
+        imports = extract_imports(source_content, language)
+        base_dir = str(Path(file_path).parent)
+        for imp in imports:
+            candidates = module_to_file_candidates(imp, base_dir, language)
+            for candidate in candidates:
+                if candidate not in seen_paths:
+                    seen_paths.add(candidate)
+                    predictions.append(PrefetchResult(
+                        path=candidate,
+                        reason="import",
+                        confidence=0.70,
+                    ))
+
+        # 2. Test files (confidence: 0.5)
+        test_candidates = infer_test_files(file_path)
+        for tc in test_candidates:
+            if tc not in seen_paths:
+                seen_paths.add(tc)
+                predictions.append(PrefetchResult(
+                    path=tc,
+                    reason="test_file",
+                    confidence=0.50,
+                ))
+
+        # 3. Learned co-access patterns (confidence: normalized count)
+        if file_path in self._co_access:
+            co_counts = self._co_access[file_path]
+            if co_counts:
+                max_count = max(co_counts.values())
+                for co_path, count in co_counts.most_common(max_results):
+                    if co_path not in seen_paths:
+                        seen_paths.add(co_path)
+                        confidence = min(count / max(max_count, 1), 1.0) * 0.80
+                        predictions.append(PrefetchResult(
+                            path=co_path,
+                            reason="co_access",
+                            confidence=round(confidence, 2),
+                        ))
+
+        # Sort by confidence and limit results
+        predictions.sort(key=lambda p: p.confidence, reverse=True)
+        return predictions[:max_results]
+
+    def stats(self) -> dict:
+        total_pairs = sum(
+            len(targets) for targets in self._co_access.values()
+        )
+        return {
+            "tracked_files": len(self._co_access),
+            "co_access_pairs": total_pairs // 2,  # Undirected
+            "recent_accesses": len(self._recent_accesses),
+        }