cognitive-cache 0.1.1__py3-none-any.whl

cognitive_cache/__init__.py

@@ -0,0 +1,12 @@
+from cognitive_cache.api import RepoIndex, select_context, select_context_from_repo
+from cognitive_cache.models import Source, Task, ScoredSource, SelectionResult
+
+__all__ = [
+    "RepoIndex",
+    "select_context",
+    "select_context_from_repo",
+    "Source",
+    "Task",
+    "ScoredSource",
+    "SelectionResult",
+]

cognitive_cache/api.py

@@ -0,0 +1,265 @@
+"""Public API for Cognitive Cache: RepoIndex and select_context.
+
+This module is the single entry point for library consumers. The CLI
+and MCP server are thin wrappers around RepoIndex and select_context;
+everything else (indexing, scoring, selection) is internal machinery.
+"""
+
+import os
+import re
+import subprocess
+from dataclasses import dataclass
+
+from cognitive_cache.models import Source, Task, SelectionResult
+from cognitive_cache.indexer.repo_indexer import (
+    index_repo,
+    SOURCE_EXTENSIONS,
+    SKIP_DIRS,
+)
+from cognitive_cache.indexer.git_analyzer import GitAnalyzer
+from cognitive_cache.indexer.graph_builder import build_dependency_graph
+from cognitive_cache.signals.embedding_sim import EmbeddingSimilaritySignal
+from cognitive_cache.core.value_function import ValueFunction
+from cognitive_cache.core.selector import GreedySelector
+from cognitive_cache.core.orderer import order_context
+
+
+def _extract_task_symbols(
+    title: str, body: str, sources: list[Source]
+) -> frozenset[str]:
+    """Extract symbols from task text using exact and selective substring matching.
+
+    Strategy:
+    1. Exact match: symbol name appears verbatim in issue text (strongest signal).
+    2. Substring match: issue words (6+ chars) appear in symbol names.
+       Short common words like 'error', 'class', 'test' cause too many false
+       matches, so we only use longer words.
+    """
+    all_symbols = set()
+    for s in sources:
+        all_symbols.update(s.symbols)
+
+    text = f"{title} {body}".lower()
+    task_words_long = set(re.findall(r"\b[a-z_][a-z0-9_]{5,}\b", text))
+    stop_words = {
+        "return",
+        "import",
+        "should",
+        "string",
+        "number",
+        "before",
+        "after",
+        "called",
+        "values",
+        "object",
+        "update",
+        "create",
+        "delete",
+        "method",
+        "function",
+        "default",
+        "option",
+        "options",
+        "config",
+        "module",
+        "result",
+        "response",
+        "request",
+        "handler",
+        "callback",
+        "parameter",
+    }
+    task_words_long -= stop_words
+
+    matches = set()
+    for sym in all_symbols:
+        sym_lower = sym.lower()
+        if sym_lower in text and len(sym_lower) >= 4:
+            matches.add(sym)
+            continue
+        for word in task_words_long:
+            if word in sym_lower:
+                matches.add(sym)
+                break
+
+    return frozenset(matches)
+
+
+def _find_entry_points(task_symbols: frozenset[str], sources: list[Source]) -> set[str]:
+    entry_points = set()
+    for s in sources:
+        if s.symbols & task_symbols:
+            entry_points.add(s.path)
+    return entry_points
+
+
+def _get_head_commit(repo_path: str) -> str:
+    try:
+        result = subprocess.run(
+            ["git", "rev-parse", "HEAD"],
+            cwd=repo_path,
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode != 0:
+            return ""
+        return result.stdout.strip()
+    except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
+        return ""
+
+
+def _collect_mtimes(repo_path: str, sources: list[Source]) -> dict[str, float]:
+    mtimes: dict[str, float] = {}
+    for s in sources:
+        full_path = os.path.join(repo_path, s.path)
+        try:
+            mtimes[s.path] = os.path.getmtime(full_path)
+        except OSError:
+            pass
+    return mtimes
+
+
+@dataclass
+class RepoIndex:
+    """An indexed snapshot of a repository, ready for context selection.
+
+    Build one with ``RepoIndex.build(repo_path)`` and reuse it across
+    multiple ``select_context`` calls. Call ``refresh()`` to cheaply detect
+    whether the repo has changed since the last index.
+    """
+
+    repo_path: str
+    sources: list[Source]
+    recency_data: dict[str, float]
+    graph: object  # DependencyGraph
+    embedding_signal: EmbeddingSimilaritySignal
+    file_mtimes: dict[str, float]
+    head_commit: str
+
+    @classmethod
+    def build(cls, repo_path: str) -> "RepoIndex":
+        """Build a full index from scratch.
+
+        Raises FileNotFoundError if repo_path does not exist.
+        """
+        if not os.path.exists(repo_path):
+            raise FileNotFoundError(f"Repository path does not exist: {repo_path}")
+
+        sources = index_repo(repo_path)
+        graph = build_dependency_graph(sources)
+
+        git_analyzer = GitAnalyzer(repo_path)
+        recency_data = git_analyzer.recency_scores()
+
+        embedding_signal = EmbeddingSimilaritySignal()
+        if sources:
+            embedding_signal.fit(sources)
+
+        return cls(
+            repo_path=repo_path,
+            sources=sources,
+            recency_data=recency_data,
+            graph=graph,
+            embedding_signal=embedding_signal,
+            file_mtimes=_collect_mtimes(repo_path, sources),
+            head_commit=_get_head_commit(repo_path),
+        )
+
+    def refresh(self) -> "RepoIndex":
+        """Check whether the repo has changed and rebuild if needed.
+
+        Returns ``self`` if nothing changed. Returns a new RepoIndex if
+        files were modified or HEAD moved. Only re-fetches git recency
+        data when HEAD has actually moved, since that's the expensive part.
+        """
+        current_head = _get_head_commit(self.repo_path)
+        head_changed = current_head != self.head_commit
+
+        # Walk the source tree and compare mtimes against our snapshot.
+        current_mtimes: dict[str, float] = {}
+        for root, dirs, files in os.walk(self.repo_path):
+            dirs[:] = [d for d in dirs if d not in SKIP_DIRS]
+            for filename in files:
+                ext = os.path.splitext(filename)[1]
+                if ext not in SOURCE_EXTENSIONS:
+                    continue
+                full_path = os.path.join(root, filename)
+                rel_path = os.path.relpath(full_path, self.repo_path).replace("\\", "/")
+                try:
+                    current_mtimes[rel_path] = os.path.getmtime(full_path)
+                except OSError:
+                    pass
+
+        files_changed = current_mtimes != self.file_mtimes
+
+        if not files_changed and not head_changed:
+            return self
+
+        # Something changed, so rebuild the index.
+        sources = index_repo(self.repo_path)
+        graph = build_dependency_graph(sources)
+        embedding_signal = EmbeddingSimilaritySignal()
+        if sources:
+            embedding_signal.fit(sources)
+
+        recency_data = self.recency_data
+        if head_changed:
+            git_analyzer = GitAnalyzer(self.repo_path)
+            recency_data = git_analyzer.recency_scores()
+
+        return RepoIndex(
+            repo_path=self.repo_path,
+            sources=sources,
+            recency_data=recency_data,
+            graph=graph,
+            embedding_signal=embedding_signal,
+            file_mtimes=_collect_mtimes(self.repo_path, sources),
+            head_commit=current_head,
+        )
+
+
+def select_context(
+    index: RepoIndex,
+    task: str | Task,
+    budget: int = 12_000,
+) -> SelectionResult:
+    """Select the most valuable context for a task from an indexed repo.
+
+    If *task* is a plain string it gets wrapped into a Task with symbols
+    extracted automatically. If it's already a Task object, its symbols
+    are used directly.
+    """
+    if not index.sources:
+        return SelectionResult(selected=[], total_tokens=0, budget=budget)
+
+    if isinstance(task, str):
+        task_symbols = _extract_task_symbols(task, "", index.sources)
+        task = Task(title=task, body="", symbols=task_symbols)
+
+    entry_points = _find_entry_points(task.symbols, index.sources)
+    vf = ValueFunction(
+        graph=index.graph,
+        recency_data=index.recency_data,
+        embedding_signal=index.embedding_signal,
+        entry_points=entry_points,
+    )
+    selector = GreedySelector(value_function=vf)
+    result = selector.select(index.sources, task, budget)
+    result.selected = order_context(result.selected)
+    return result
+
+
+def select_context_from_repo(
+    repo_path: str,
+    task: str | Task,
+    budget: int = 12_000,
+) -> SelectionResult:
+    """Convenience wrapper that builds an index and selects context in one call.
+
+    Useful for one-shot usage where you don't need to reuse the index.
+    For repeated queries against the same repo, build a RepoIndex once
+    and call select_context directly.
+    """
+    index = RepoIndex.build(repo_path)
+    return select_context(index, task, budget)

cognitive_cache/baselines/base.py

@@ -0,0 +1,11 @@
+"""Base interface for context selection strategies (baselines + our algorithm)."""
+
+from abc import ABC, abstractmethod
+
+from cognitive_cache.models import Source, Task, SelectionResult
+
+
+class BaselineStrategy(ABC):
+    @abstractmethod
+    def select(self, sources: list[Source], task: Task, budget: int) -> SelectionResult:
+        ...

cognitive_cache/baselines/embedding_select.py

@@ -0,0 +1,37 @@
+"""Baseline 3: Embedding Similarity (simulates RAG). Top-k by TF-IDF cosine similarity."""
+
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+
+from cognitive_cache.models import Source, Task, ScoredSource, SelectionResult
+from cognitive_cache.baselines.base import BaselineStrategy
+
+
+class EmbeddingStrategy(BaselineStrategy):
+    def select(self, sources: list[Source], task: Task, budget: int) -> SelectionResult:
+        if not sources:
+            return SelectionResult(selected=[], total_tokens=0, budget=budget)
+
+        corpus = [task.full_text] + [s.content for s in sources]
+        try:
+            vectorizer = TfidfVectorizer(max_features=5000, token_pattern=r"(?u)\b\w+\b")
+            matrix = vectorizer.fit_transform(corpus)
+        except ValueError:
+            return SelectionResult(selected=[], total_tokens=0, budget=budget)
+
+        task_vec = matrix[0:1]
+        source_vecs = matrix[1:]
+        similarities = cosine_similarity(task_vec, source_vecs)[0]
+
+        ranked = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)
+
+        selected = []
+        total = 0
+        for idx, sim in ranked:
+            s = sources[idx]
+            if total + s.token_count > budget:
+                continue
+            selected.append(ScoredSource(source=s, score=float(sim)))
+            total += s.token_count
+
+        return SelectionResult(selected=selected, total_tokens=total, budget=budget)

cognitive_cache/baselines/full_stuff.py

@@ -0,0 +1,19 @@
+"""Baseline 2: Full Stuff. Cram as many files as fit, in filesystem order."""
+
+from cognitive_cache.models import Source, Task, ScoredSource, SelectionResult
+from cognitive_cache.baselines.base import BaselineStrategy
+
+
+class FullStuffStrategy(BaselineStrategy):
+    def select(self, sources: list[Source], task: Task, budget: int) -> SelectionResult:
+        sorted_sources = sorted(sources, key=lambda s: s.path)
+
+        selected = []
+        total = 0
+        for s in sorted_sources:
+            if total + s.token_count > budget:
+                continue
+            selected.append(ScoredSource(source=s, score=0.0))
+            total += s.token_count
+
+        return SelectionResult(selected=selected, total_tokens=total, budget=budget)

cognitive_cache/baselines/grep_select.py

@@ -0,0 +1,32 @@
+"""Baseline 4: Grep Strategy (simulates Claude Code / Copilot symbol search)."""
+
+from cognitive_cache.models import Source, Task, ScoredSource, SelectionResult
+from cognitive_cache.baselines.base import BaselineStrategy
+
+
+class GrepStrategy(BaselineStrategy):
+    def select(self, sources: list[Source], task: Task, budget: int) -> SelectionResult:
+        if not task.symbols:
+            keywords = set(task.full_text.lower().split())
+        else:
+            keywords = {s.lower() for s in task.symbols}
+
+        scored = []
+        for source in sources:
+            content_lower = source.content.lower()
+            matches = sum(1 for kw in keywords if kw in content_lower)
+            if matches > 0:
+                scored.append((source, matches))
+
+        scored.sort(key=lambda x: x[1], reverse=True)
+
+        selected = []
+        total = 0
+        for source, matches in scored:
+            if total + source.token_count > budget:
+                continue
+            score = matches / max(len(keywords), 1)
+            selected.append(ScoredSource(source=source, score=score))
+            total += source.token_count
+
+        return SelectionResult(selected=selected, total_tokens=total, budget=budget)

cognitive_cache/baselines/llm_triage.py

@@ -0,0 +1,42 @@
+"""Baseline 5: LLM Triage (let the brain pick its own context).
+
+Two-pass strategy — requires an LLM adapter. Only used during benchmark runs.
+"""
+
+from cognitive_cache.models import Source, Task, ScoredSource, SelectionResult
+from cognitive_cache.baselines.base import BaselineStrategy
+
+
+class LLMTriageStrategy(BaselineStrategy):
+    def __init__(self, llm_adapter=None):
+        self._llm = llm_adapter
+
+    def select(self, sources: list[Source], task: Task, budget: int) -> SelectionResult:
+        if self._llm is None:
+            raise RuntimeError("LLMTriageStrategy requires an LLM adapter")
+
+        file_listing = "\n".join(s.path for s in sources)
+        triage_prompt = (
+            f"You are helping fix a bug. Here is the issue:\n\n"
+            f"Title: {task.title}\n"
+            f"Body: {task.body}\n\n"
+            f"Here are all the files in the repository:\n\n"
+            f"{file_listing}\n\n"
+            f"Which files are most likely relevant to this issue? "
+            f"List just the file paths, one per line. Pick at most 20 files."
+        )
+
+        response = self._llm.complete(triage_prompt, max_tokens=500)
+
+        source_by_path = {s.path: s for s in sources}
+        selected = []
+        total = 0
+        for line in response.strip().split("\n"):
+            path = line.strip().strip("- ").strip("`")
+            if path in source_by_path:
+                source = source_by_path[path]
+                if total + source.token_count <= budget:
+                    selected.append(ScoredSource(source=source, score=1.0))
+                    total += source.token_count
+
+        return SelectionResult(selected=selected, total_tokens=total, budget=budget)

cognitive_cache/baselines/random_select.py

@@ -0,0 +1,26 @@
+"""Baseline 1: Random Selection. Lower bound."""
+
+import random as random_module
+
+from cognitive_cache.models import Source, Task, ScoredSource, SelectionResult
+from cognitive_cache.baselines.base import BaselineStrategy
+
+
+class RandomStrategy(BaselineStrategy):
+    def __init__(self, seed: int | None = None):
+        self._seed = seed
+
+    def select(self, sources: list[Source], task: Task, budget: int) -> SelectionResult:
+        rng = random_module.Random(self._seed)
+        shuffled = list(sources)
+        rng.shuffle(shuffled)
+
+        selected = []
+        total = 0
+        for s in shuffled:
+            if total + s.token_count > budget:
+                continue
+            selected.append(ScoredSource(source=s, score=0.0))
+            total += s.token_count
+
+        return SelectionResult(selected=selected, total_tokens=total, budget=budget)

cognitive_cache/cli.py

@@ -0,0 +1,141 @@
+"""CLI entry point for cognitive-cache.
+
+Usage:
+    cognitive-cache select --repo . --task "fix the login bug" --budget 12000
+"""
+
+import argparse
+import json
+import sys
+
+from cognitive_cache.api import RepoIndex, select_context
+
+
+def _format_human_readable(result) -> str:
+    lines = []
+    for ss in result.selected:
+        path = ss.source.path
+        score = ss.score
+        signals = ss.signal_scores
+
+        sig_parts = []
+        for key in [
+            "symbol_overlap",
+            "graph_distance",
+            "change_recency",
+            "embedding_sim",
+            "file_role_prior",
+            "redundancy",
+        ]:
+            val = signals.get(key, 0.0)
+            short = {
+                "symbol_overlap": "sym",
+                "graph_distance": "graph",
+                "change_recency": "recency",
+                "embedding_sim": "embed",
+                "file_role_prior": "role",
+                "redundancy": "redund",
+            }[key]
+            sig_parts.append(f"{short}={val:.1f}")
+
+        sig_str = " ".join(sig_parts)
+        lines.append(f"{path:<45} {score:.3f}  [{sig_str}]")
+
+    lines.append("")
+    lines.append(
+        f"{len(result.selected)} files selected | "
+        f"{result.total_tokens:,} / {result.budget:,} tokens used"
+    )
+    return "\n".join(lines)
+
+
+def _format_json(result) -> str:
+    data = {
+        "files": [
+            {
+                "path": ss.source.path,
+                "score": round(ss.score, 4),
+                "signals": {k: round(v, 4) for k, v in ss.signal_scores.items()},
+                "token_count": ss.source.token_count,
+            }
+            for ss in result.selected
+        ],
+        "total_tokens": result.total_tokens,
+        "budget": result.budget,
+        "budget_remaining": result.budget_remaining,
+    }
+    return json.dumps(data, indent=2)
+
+
+def _write_context_file(result, output_path: str):
+    with open(output_path, "w") as f:
+        for ss in result.selected:
+            f.write(f"# --- {ss.source.path} ---\n")
+            f.write(ss.source.content)
+            f.write("\n\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="cognitive-cache",
+        description="Optimal context selection for LLMs",
+    )
+    subparsers = parser.add_subparsers(dest="command")
+
+    select_parser = subparsers.add_parser(
+        "select", help="Select context files for a task"
+    )
+    select_parser.add_argument(
+        "--repo", required=True, help="Path to the repository root"
+    )
+    select_parser.add_argument("--task", required=True, help="Task description")
+    select_parser.add_argument(
+        "--budget", type=int, default=12000, help="Token budget (default: 12000)"
+    )
+    select_parser.add_argument(
+        "--json", action="store_true", dest="json_output", help="Output as JSON"
+    )
+    select_parser.add_argument("--output", help="Write full context to this file")
+
+    args = parser.parse_args()
+
+    if args.command != "select":
+        parser.print_help()
+        sys.exit(1)
+
+    try:
+        index = RepoIndex.build(args.repo)
+    except FileNotFoundError as e:
+        print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    if not index.sources:
+        print("No source files found.", file=sys.stderr)
+        sys.exit(1)
+
+    result = select_context(index, args.task, budget=args.budget)
+
+    if not result.selected:
+        print("No files selected.", file=sys.stderr)
+        sys.exit(0)
+
+    # Warn if no symbols matched (vague task description)
+    if all(
+        ss.signal_scores.get("symbol_overlap", 0.0) == 0.0 for ss in result.selected
+    ):
+        print(
+            "Warning: No symbol matches found, results may be less precise.",
+            file=sys.stderr,
+        )
+
+    if args.json_output:
+        print(_format_json(result))
+    else:
+        print(_format_human_readable(result))
+
+    if args.output:
+        _write_context_file(result, args.output)
+
+
+if __name__ == "__main__":
+    main()