fittok 0.3.0__py3-none-any.whl

fittok/__init__.py

@@ -0,0 +1,27 @@
+"""Fittok — relevant-code retrieval within a token budget.
+
+Usable three ways, each standalone:
+  - As a library:   ``from fittok import optimize``
+  - As a CLI:       ``fittok query <path> "<question>"``
+  - As an MCP server: ``python -m fittok`` (for AI clients)
+"""
+
+__version__ = "0.3.0"
+
+__all__ = ["optimize", "index", "__version__"]
+
+
+def optimize(codebase_path: str, query: str, token_budget: int = 0) -> dict:
+    """Return the most relevant source code for *query*, within *token_budget*.
+
+    token_budget=0 → adaptive sizing. No MCP client required. Returns a dict with
+    ``optimized_context``, ``graph_stats``, ``slurp_stats`` and ``savings``.
+    """
+    from .server import optimize_context_tool
+    return optimize_context_tool(codebase_path, query, token_budget)
+
+
+def index(codebase_path: str) -> dict:
+    """Pre-build + cache the graph and embeddings for a codebase."""
+    from .indexer import index_codebase
+    return index_codebase(codebase_path)

fittok/__main__.py

@@ -0,0 +1,5 @@
+"""Allow running as `python -m fittok`."""
+
+from .server import main
+
+main()

fittok/cache.py

@@ -0,0 +1,223 @@
+"""3-level caching layer for fittok.
+
+Levels:
+  1. Graph cache — keyed by (root_path, file_mtimes_hash)
+  2. Query cache — keyed by (graph_path, query_hash, token_budget)
+  3. Compression cache — keyed by (context_hash, question_hash, target_tokens)
+
+Uses diskcache for persistent storage. Cache dir: ~/.cache/fittok/
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import os
+from pathlib import Path
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+CACHE_DIR = os.environ.get(
+    "FITTOK_CACHE_DIR",
+    os.path.expanduser("~/.cache/fittok"),
+)
+MAX_CACHE_SIZE = int(os.environ.get("FITTOK_CACHE_MAX_MB", "500")) * 1024 * 1024
+
+# Stats tracking
+_stats: dict[str, int] = {"graph_hits": 0, "graph_misses": 0,
+                          "query_hits": 0, "query_misses": 0,
+                          "compression_hits": 0, "compression_misses": 0}
+
+_cache = None
+
+
+def _get_cache():
+    """Lazily initialize diskcache."""
+    global _cache
+    if _cache is not None:
+        return _cache
+    try:
+        import diskcache
+        os.makedirs(CACHE_DIR, exist_ok=True)
+        _cache = diskcache.Cache(CACHE_DIR, size_limit=MAX_CACHE_SIZE)
+        return _cache
+    except ImportError:
+        logger.warning("diskcache not installed — caching disabled. Install with: pip install diskcache")
+        return None
+
+
+# ── Hashing helpers ───────────────────────────────────────────────────────────
+
+def _hash_str(text: str) -> str:
+    return hashlib.sha256(text.encode()).hexdigest()[:16]
+
+
+def _hash_mtimes(root_path: str) -> str:
+    """Hash of all code file mtimes under root_path."""
+    from .graphify import _EXT_TO_LANG
+    mtimes: list[str] = []
+    root = Path(root_path).resolve()
+    if not root.is_dir():
+        return _hash_str(root_path)
+    skip_dirs = {
+        "node_modules", ".git", "__pycache__", ".venv", "venv",
+        "dist", "build", ".tox", ".mypy_cache", ".pytest_cache",
+        ".next", ".nuxt", "target", "vendor", ".gradle",
+    }
+    for dirpath, dirnames, filenames in os.walk(root):
+        dirnames[:] = [d for d in dirnames if d not in skip_dirs]
+        for f in sorted(filenames):
+            p = Path(dirpath) / f
+            if p.suffix in _EXT_TO_LANG:
+                try:
+                    mtimes.append(f"{p.relative_to(root)}:{p.stat().st_mtime}")
+                except OSError:
+                    pass
+    return _hash_str("|".join(mtimes))
+
+
+# ── Graph cache ───────────────────────────────────────────────────────────────
+
+def get_cached_graph(root_path: str) -> Optional[Any]:
+    """Look up a cached graph by root path + file mtimes."""
+    cache = _get_cache()
+    if cache is None:
+        return None
+    key = f"graph:{root_path}:{_hash_mtimes(root_path)}"
+    result = cache.get(key)
+    if result is not None:
+        _stats["graph_hits"] += 1
+        logger.debug("Graph cache HIT: %s", root_path)
+        from .models import KnowledgeGraph
+        return KnowledgeGraph.model_validate(result)
+    _stats["graph_misses"] += 1
+    return None
+
+
+def set_cached_graph(root_path: str, graph: Any) -> None:
+    """Store a graph in cache."""
+    cache = _get_cache()
+    if cache is None:
+        return
+    key = f"graph:{root_path}:{_hash_mtimes(root_path)}"
+    cache.set(key, graph.model_dump(), expire=3600 * 24)  # 24h TTL
+
+
+# ── Query cache ──────────────────────────────────────────────────────────────
+
+def get_cached_query(graph_path: str, query: str, token_budget: int) -> Optional[dict]:
+    """Look up a cached query result."""
+    cache = _get_cache()
+    if cache is None:
+        return None
+    key = f"query:{graph_path}:{_hash_str(query)}:{token_budget}"
+    result = cache.get(key)
+    if result is not None:
+        _stats["query_hits"] += 1
+        logger.debug("Query cache HIT: %s", query[:40])
+        return result  # type: ignore[return-value]
+    _stats["query_misses"] += 1
+    return None
+
+
+def set_cached_query(graph_path: str, query: str, token_budget: int, result: dict) -> None:
+    """Store a query result in cache."""
+    cache = _get_cache()
+    if cache is None:
+        return
+    key = f"query:{graph_path}:{_hash_str(query)}:{token_budget}"
+    cache.set(key, result, expire=3600)  # 1h TTL
+
+
+# ── Compression cache ────────────────────────────────────────────────────────
+
+def get_cached_compression(context: str, question: str, target_tokens: int) -> Optional[dict]:
+    """Look up a cached compression result."""
+    cache = _get_cache()
+    if cache is None:
+        return None
+    key = f"compress:{_hash_str(context)}:{_hash_str(question)}:{target_tokens}"
+    result = cache.get(key)
+    if result is not None:
+        _stats["compression_hits"] += 1
+        logger.debug("Compression cache HIT")
+        return result  # type: ignore[return-value]
+    _stats["compression_misses"] += 1
+    return None
+
+
+def set_cached_compression(context: str, question: str, target_tokens: int, result: dict) -> None:
+    """Store a compression result in cache."""
+    cache = _get_cache()
+    if cache is None:
+        return
+    key = f"compress:{_hash_str(context)}:{_hash_str(question)}:{target_tokens}"
+    cache.set(key, result, expire=3600)
+
+
+# ── Embedding cache (persistent, content-keyed) ──────────────────────────────
+# Content-keyed means this is incremental by nature: an unchanged function keeps
+# its embedding across restarts and edits; only new/changed code is re-embedded.
+
+def get_cached_embedding(content_hash: str):
+    """Return a persisted embedding vector for a content hash, or None."""
+    cache = _get_cache()
+    if cache is None:
+        return None
+    return cache.get(f"emb:{content_hash}")
+
+
+def set_cached_embedding(content_hash: str, vector) -> None:
+    """Persist an embedding vector (30-day TTL)."""
+    cache = _get_cache()
+    if cache is None:
+        return
+    cache.set(f"emb:{content_hash}", vector, expire=3600 * 24 * 30)
+
+
+# ── Cache management ─────────────────────────────────────────────────────────
+
+def clear_cache(scope: str = "all") -> dict:
+    """Clear cache by scope: all, graph, query, compression."""
+    cache = _get_cache()
+    if cache is None:
+        return {"error": "Cache not available (diskcache not installed)"}
+
+    removed = 0
+    for key in list(cache.iterkeys()):
+        if scope == "all":
+            cache.delete(key)
+            removed += 1
+        elif scope == "graph" and str(key).startswith("graph:"):
+            cache.delete(key)
+            removed += 1
+        elif scope == "query" and str(key).startswith("query:"):
+            cache.delete(key)
+            removed += 1
+        elif scope == "compression" and str(key).startswith("compress:"):
+            cache.delete(key)
+            removed += 1
+
+    return {"cleared": removed, "scope": scope}
+
+
+def cache_stats() -> dict:
+    """Return cache hit/miss stats and size."""
+    cache = _get_cache()
+    size = 0
+    count = 0
+    if cache is not None:
+        try:
+            count = len(cache)  # type: ignore[arg-type]
+            size = cache.volume()
+        except Exception:
+            pass
+
+    return {
+        "stats": dict(_stats),
+        "entries": count,
+        "size_bytes": size,
+        "cache_dir": CACHE_DIR,
+        "available": cache is not None,
+    }

fittok/cli.py

@@ -0,0 +1,81 @@
+"""Command-line interface — lets the package do its job WITHOUT an MCP client.
+
+Subcommands:
+    fittok serve              Run the MCP server (stdio) — for AI clients.
+    fittok index <path>       Pre-build graph + embeddings for a repo.
+    fittok query <path> "<q>" Print the most relevant code for a query.
+
+`query` is the standalone equivalent of the MCP `optimize_context` tool: it
+returns the same readable, budget-bounded context, straight to your terminal.
+With no subcommand, defaults to `serve` (so existing MCP registrations that
+launch the bare `fittok` command keep working).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        prog="fittok",
+        description="Retrieve the most relevant source code for a query, within a token budget.",
+    )
+    sub = parser.add_subparsers(dest="cmd")
+
+    sub.add_parser("serve", help="Run the MCP server over stdio (for AI clients).")
+
+    p_index = sub.add_parser("index", help="Pre-build the knowledge graph + embeddings for a repo.")
+    p_index.add_argument("path", help="Path to the codebase to index.")
+
+    p_query = sub.add_parser("query", help="Print the most relevant code for a query (no MCP needed).")
+    p_query.add_argument("path", help="Path to the codebase.")
+    p_query.add_argument("query", help="Natural-language question about the codebase.")
+    p_query.add_argument("--budget", type=int, default=0,
+                         help="Token budget (0 = adaptive, the default).")
+    p_query.add_argument("--json", action="store_true",
+                         help="Emit the full result dict as JSON.")
+
+    args = parser.parse_args(argv)
+
+    # Default (or `serve`) → run the MCP server.
+    if args.cmd in (None, "serve"):
+        from .server import main as serve_main
+        serve_main()
+        return
+
+    # For human-facing commands, keep stderr clean (INFO logs are server-mode noise).
+    import logging
+    logging.disable(logging.INFO)
+
+    if args.cmd == "index":
+        from .indexer import index_codebase
+        r = index_codebase(args.path)
+        print(f"Indexed {r['nodes']} nodes / {r['edges']} edges, {r['embedded']} embeddings "
+              f"(parse {r['parse_s']}s, embed {r['embed_s']}s). Cached.")
+        return
+
+    if args.cmd == "query":
+        from .server import optimize_context_tool
+        res = optimize_context_tool(args.path, args.query, args.budget)
+        if "error" in res:
+            print(f"Error: {res['error']}", file=sys.stderr)
+            sys.exit(1)
+        if args.json:
+            print(json.dumps(res, indent=2))
+            return
+        # Human mode: stats to stderr, the actual context to stdout (pipe-friendly).
+        s, sv = res["slurp_stats"], res["savings"]
+        print(
+            f"# {s['selected_nodes']} nodes · {s['tokens_sent']} tokens · "
+            f"confidence {s['confidence_label']} ({s['confidence']}) · {sv['summary']}",
+            file=sys.stderr,
+        )
+        print(res["optimized_context"])
+        return
+
+
+if __name__ == "__main__":
+    main()

fittok/diff.py

@@ -0,0 +1,110 @@
+"""Graph diffing — compare two knowledge graphs and report structural differences."""
+
+from __future__ import annotations
+
+import hashlib
+from typing import Any
+
+from .models import KnowledgeGraph
+
+
+def _content_hash(content: str) -> str:
+    return hashlib.sha256(content.encode()).hexdigest()[:16]
+
+
+def diff_graphs(graph_a: KnowledgeGraph, graph_b: KnowledgeGraph) -> dict:
+    """Compare two knowledge graphs and return structural differences.
+
+    Args:
+        graph_a: First (earlier) graph.
+        graph_b: Second (later) graph.
+
+    Returns:
+        Dict with nodes_added, nodes_removed, nodes_modified, edges_added,
+        edges_removed, and a human-readable summary.
+    """
+    # Index nodes by ID
+    nodes_a = {n.id: n for n in graph_a.nodes}
+    nodes_b = {n.id: n for n in graph_b.nodes}
+
+    ids_a = set(nodes_a.keys())
+    ids_b = set(nodes_b.keys())
+
+    nodes_added = [
+        {"id": n.id, "name": n.name, "type": n.type.value, "file": n.file}
+        for nid in sorted(ids_b - ids_a)
+        if (n := nodes_b[nid])
+    ]
+
+    nodes_removed = [
+        {"id": n.id, "name": n.name, "type": n.type.value, "file": n.file}
+        for nid in sorted(ids_a - ids_b)
+        if (n := nodes_a[nid])
+    ]
+
+    # Modified: same ID but different content hash
+    nodes_modified: list[dict] = []
+    for nid in sorted(ids_a & ids_b):
+        a, b = nodes_a[nid], nodes_b[nid]
+        if _content_hash(a.content) != _content_hash(b.content):
+            nodes_modified.append({
+                "id": nid,
+                "name": b.name,
+                "type": b.type.value,
+                "file": b.file,
+                "line_start": b.line_start,
+                "line_end": b.line_end,
+            })
+
+    # Index edges by (source, target, type)
+    def _edge_key(e: Any) -> str:
+        return f"{e.source}|{e.target}|{e.type.value}"
+
+    edges_a = {_edge_key(e): e for e in graph_a.edges}
+    edges_b = {_edge_key(e): e for e in graph_b.edges}
+
+    edge_keys_a = set(edges_a.keys())
+    edge_keys_b = set(edges_b.keys())
+
+    edges_added = [
+        {"source": e.source, "target": e.target, "type": e.type.value}
+        for k in sorted(edge_keys_b - edge_keys_a)
+        if (e := edges_b[k])
+    ]
+
+    edges_removed = [
+        {"source": e.source, "target": e.target, "type": e.type.value}
+        for k in sorted(edge_keys_a - edge_keys_b)
+        if (e := edges_a[k])
+    ]
+
+    # Build summary
+    parts: list[str] = []
+    if nodes_added:
+        files_added = {n["file"] for n in nodes_added}
+        parts.append(f"{len(nodes_added)} nodes added across {len(files_added)} file(s)")
+    if nodes_removed:
+        files_removed = {n["file"] for n in nodes_removed}
+        parts.append(f"{len(nodes_removed)} nodes removed across {len(files_removed)} file(s)")
+    if nodes_modified:
+        files_modified = {n["file"] for n in nodes_modified}
+        parts.append(f"{len(nodes_modified)} nodes modified across {len(files_modified)} file(s)")
+    if edges_added:
+        parts.append(f"{len(edges_added)} edges added")
+    if edges_removed:
+        parts.append(f"{len(edges_removed)} edges removed")
+
+    summary = "; ".join(parts) if parts else "No changes detected"
+
+    return {
+        "nodes_added": nodes_added,
+        "nodes_removed": nodes_removed,
+        "nodes_modified": nodes_modified,
+        "edges_added": edges_added,
+        "edges_removed": edges_removed,
+        "stats": {
+            "graph_a": {"nodes": len(graph_a.nodes), "edges": len(graph_a.edges)},
+            "graph_b": {"nodes": len(graph_b.nodes), "edges": len(graph_b.edges)},
+        },
+        "summary": summary,
+    }

fittok/embeddings.py

@@ -0,0 +1,115 @@
+"""Semantic embeddings for relevance scoring.
+
+Optional and fail-safe: if sentence-transformers (or the model) is unavailable,
+every function degrades to returning ``None`` and the caller falls back to the
+lexical TF-IDF path. This is what lets natural-language queries
+(e.g. "real-time conversation with the AI") match code that uses different
+words (``WebSocket``, ``streamResponse``, ``transcript``) — something pure
+keyword matching cannot do.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_MODEL = os.environ.get(
+    "FITTOK_EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2"
+)
+
+_model = None
+_unavailable = False
+_EMB_CACHE: dict[str, object] = {}  # content-hash -> embedding vector (per process)
+
+
+def _get_model():
+    """Lazily load the embedding model; cache the result. Returns None if unavailable."""
+    global _model, _unavailable
+    if _model is not None:
+        return _model
+    if _unavailable:
+        return None
+    try:
+        from sentence_transformers import SentenceTransformer
+    except ImportError:
+        logger.info("sentence-transformers not installed; semantic scoring disabled")
+        _unavailable = True
+        return None
+    try:
+        _model = SentenceTransformer(_DEFAULT_MODEL)
+        logger.info("Loaded embedding model %s", _DEFAULT_MODEL)
+    except Exception:
+        logger.warning("Failed to load embedding model %s; semantic scoring disabled",
+                       _DEFAULT_MODEL, exc_info=True)
+        _unavailable = True
+        return None
+    return _model
+
+
+def is_available() -> bool:
+    """True if semantic scoring can run."""
+    return _get_model() is not None
+
+
+def semantic_scores(nodes: list, query: str) -> Optional[dict[str, float]]:
+    """Cosine similarity of each node to the query, in [0, 1]-ish (raw cosine).
+
+    Returns None if embeddings are unavailable (caller should fall back to TF-IDF).
+    Scores are raw cosine similarities — meaningful as an *absolute* confidence
+    signal, unlike min-max-normalized lexical scores.
+    """
+    if not nodes:
+        return {}
+    model = _get_model()
+    if model is None:
+        return None
+    # Represent each node by its name + content so both signal sources count.
+    node_texts = [f"{n.name}\n{n.content or ''}" for n in nodes]
+    try:
+        node_embs = _embed_cached(model, node_texts)        # reused across queries
+        query_vec = model.encode([query], normalize_embeddings=True,
+                                 show_progress_bar=False)[0]
+    except Exception:
+        logger.warning("Embedding encode failed; falling back to lexical scoring", exc_info=True)
+        return None
+    sims = node_embs @ query_vec  # normalized vectors → dot product == cosine
+    return {n.id: float(sims[i]) for i, n in enumerate(nodes)}
+
+
+def _embed_cached(model, texts: list[str]):
+    """Encode texts, reusing a per-process cache keyed by content.
+
+    Node content is stable across queries for a given graph, so without this the
+    server re-embedded thousands of nodes on every call (and N times for an
+    N-query batch) — seconds of wasted compute each time.
+    """
+    import hashlib
+    import numpy as np
+    from . import cache as _cache
+
+    keys = [hashlib.sha256(t.encode("utf-8", "ignore")).hexdigest() for t in texts]
+    out: list = [None] * len(texts)
+    missing_i, missing_t = [], []
+    for i, k in enumerate(keys):
+        # L1: in-process cache
+        vec = _EMB_CACHE.get(k)
+        if vec is None:
+            # L2: persistent disk cache (survives restarts; incremental by content)
+            vec = _cache.get_cached_embedding(k)
+            if vec is not None:
+                _EMB_CACHE[k] = vec
+        if vec is None:
+            missing_i.append(i)
+            missing_t.append(texts[i])
+        else:
+            out[i] = vec
+    if missing_t:
+        enc = model.encode(missing_t, normalize_embeddings=True, show_progress_bar=False)
+        for j, i in enumerate(missing_i):
+            _EMB_CACHE[keys[i]] = enc[j]
+            _cache.set_cached_embedding(keys[i], enc[j])
+            out[i] = enc[j]
+    return np.array(out)