""" SwiftContext — Production-Grade Repository Explorer ==================================================== A deterministic, zero-LLM replacement for FastContext. New vs FastContext ────────────────────────────────────────────────────────────────── Feature FastContext SwiftContext ─────────────────────────────────────────────────────────────────── Search ranking (LLM confidence) Okapi BM25 + 4-signal Persistent disk index No Yes (.swiftcontext/) Incremental re-index No Yes (MD5 per file) Symbol table No Yes (kind/sig/docstr) Call graph No Yes (AST call walk) Reverse call graph No Yes (O(k) callers) Import resolver Partial Full AST resolution trace(symbol) API Not supported Yes [NEW] explain(symbol) API Not supported Yes [NEW] GPU for queries Required (4B LLM) Not needed LLM tokens / query ~2 000 0 Line number accuracy ~70 % (hallucin.) 100 % (reads file) Output Plain file:Lnn JSON+relevance+reason +docstring+deps+snippet ─────────────────────────────────────────────────────────────────── Latency (after first run with cached index) pinpoint_cite : ~2 ms (FastContext: ~1-2 s LLM call) targeted_search : ~10 ms (FastContext: ~1-2 s LLM call) broad_scan : ~30 ms (FastContext: ~3-5 s LLM call) trace() : ~5 ms (FastContext: not supported) explain() : ~1 ms (FastContext: not supported) Usage ───── from inference import SwiftContextPipeline sc = SwiftContextPipeline("./model/final") # Citation search result = sc.explore("Find the BM25Index class", repo_path=".") # Call chain chain = sc.trace("_build", repo_path=".") # Documentation doc = sc.explain("BM25Index", repo_path=".") """ from __future__ import annotations import ast import hashlib import json import math import os import re import time from collections import defaultdict from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Optional # Optional: sentence-transformers for semantic code search (~22 MB model) # Bridges vocabulary gaps BM25 cannot handle ("authentication" → login()) # Install: pip install sentence-transformers try: from sentence_transformers import SentenceTransformer import numpy as _np _HAS_SEMANTIC = True except ImportError: _HAS_SEMANTIC = False import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer # ───────────────────────────────────────────────────────────────────────────── # Constants # ───────────────────────────────────────────────────────────────────────────── _INDEX_VERSION = 2 _INDEX_DIR = ".swiftcontext" _INDEX_FILE = "index.json" BM25_K1 = 1.5 # term-frequency saturation BM25_B = 0.75 # document-length normalisation STRATEGY_LABELS = {0: "broad_scan", 1: "targeted_search", 2: "pinpoint_cite"} CONFIDENCE_FALLBACK = 0.40 _SKIP_DIRS = { ".git", "__pycache__", "node_modules", ".venv", "venv", "dist", "build", ".next", "target", ".mypy_cache", ".pytest_cache", ".tox", _INDEX_DIR, } _CODE_EXTS = { ".py", ".js", ".ts", ".jsx", ".tsx", ".java", ".cs", ".go", ".rs", ".cpp", ".c", ".h", ".rb", ".php", ".swift", ".kt", ".scala", ".r", ".lua", ".ex", ".exs", } _LANG_MAP = { ".py": "Python", ".js": "JavaScript", ".ts": "TypeScript", ".jsx": "JSX", ".tsx": "TSX", ".java": "Java", ".cs": "C#", ".go": "Go", ".rs": "Rust", ".cpp": "C++", ".c": "C", ".h": "C/C++ Header", ".rb": "Ruby", ".php": "PHP", ".swift": "Swift", ".kt": "Kotlin", ".scala": "Scala", ".r": "R", ".lua": "Lua", ".ex": "Elixir", ".exs": "Elixir", } _FC_BASELINE_TURNS = {"broad_scan": 3, "targeted_search": 2, "pinpoint_cite": 2} _FC_AVG_TOKENS = 2_000 # ───────────────────────────────────────────────────────────────────────────── # Data structures # ───────────────────────────────────────────────────────────────────────────── @dataclass class SymbolInfo: """Rich metadata for one code symbol — extracted entirely by AST.""" name: str kind: str # class | function | async_function | method | async_method | symbol file: str start_line: int end_line: int signature: str # e.g. "def foo(x: int) -> str:" docstring: str # first docstring paragraph or "" parent: Optional[str] # enclosing class name for methods language: str @dataclass class Citation: """Single code citation with full metadata. FastContext returns plain text.""" file: str start_line: int end_line: int snippet: str # actual source lines (FastContext: absent) relevance: float # 0.0-1.0 BM25+multi-signal (FastContext: absent) reason: str # why this is relevant (FastContext: absent) symbol: Optional[SymbolInfo] # (FastContext: absent) docstring: str # extracted docstring (FastContext: absent) deps: list[str] = field(default_factory=list) # import deps @dataclass class ExploreResult: """Result of explore() — code citation search.""" citations: list[Citation] confidence: float strategy_used: str turns_used: int tokens_used: int # always 0; FastContext avg ~2 000 / query tokens_saved_pct: float latency_ms: float index_chunks: int index_symbols: int @dataclass class TraceResult: """Call-chain result from trace(). Not available in FastContext.""" symbol: str definition: Optional[Citation] callers: list[Citation] # functions that call this symbol callees: list[Citation] # functions called by this symbol latency_ms: float @dataclass class ExplainResult: """Documentation result from explain(). Not available in FastContext.""" symbol: str kind: str signature: str docstring: str file: str start_line: int end_line: int language: str deps: list[str] latency_ms: float @dataclass class CodeBehavior: """Structured behavior extracted from AST — no LLM required.""" reads: list[str] # self.x attributes that are read writes: list[str] # self.x attributes that are written calls: list[str] # function / method names called raises: list[str] # exception types raised returns: str # return-type annotation or "" @dataclass class SummarizeResult: """Natural-language behavior summary. Not available in FastContext.""" symbol: str kind: str summary: str # "BM25Index is a Python class that ranks…" behavior: CodeBehavior file: str start_line: int end_line: int latency_ms: float @dataclass class ContextResult: """ Multi-file context window. FastContext built this through 2-3 LLM turns of repo browsing. SwiftContext builds it deterministically in <50 ms from the AST index. Pass to_llm_context() into any downstream LLM (GPT-4, Claude …) for deep reasoning over real code with zero hallucination of file contents. """ query: str primary: list[Citation] # direct query matches caller_context: list[Citation] # who calls the primary matches callee_context: list[Citation] # what the primary matches call summaries: dict[str, str] # symbol → one-line summary total_tokens_est: int # estimated tokens for the window latency_ms: float def to_llm_context(self) -> str: """ Format as a ready-to-use context string for any downstream LLM. Replaces what FastContext produced through 2-3 expensive LLM turns. """ parts = [f"# Repository Context: {self.query}\n"] if self.primary: parts.append("## Primary Matches\n") for c in self.primary: parts.append(f"### {c.file}:L{c.start_line}-{c.end_line}") if c.docstring: parts.append(f"*{c.docstring}*\n") parts.append(f"```\n{c.snippet}\n```\n") if c.symbol and c.symbol.name in self.summaries: parts.append(f"*Summary: {self.summaries[c.symbol.name]}*\n") if self.caller_context: parts.append("## Functions That Call The Above\n") for c in self.caller_context[:3]: parts.append(f"### {c.file}:L{c.start_line}-{c.end_line}") parts.append(f"```\n{c.snippet[:300]}\n```\n") if self.callee_context: parts.append("## Functions Called By The Above\n") for c in self.callee_context[:3]: parts.append(f"### {c.file}:L{c.start_line}-{c.end_line}") parts.append(f"```\n{c.snippet[:300]}\n```\n") return "\n".join(parts) # ───────────────────────────────────────────────────────────────────────────── # Python AST extractor # ───────────────────────────────────────────────────────────────────────────── def _safe_docstring(node: ast.AST) -> str: try: ds = ast.get_docstring(node) or "" return ds.splitlines()[0] if ds else "" except Exception: return "" def _signature(node: ast.AST, lines: list[str]) -> str: try: start = node.lineno - 1 # type: ignore[attr-defined] parts = [] for i in range(start, min(start + 6, len(lines))): parts.append(lines[i].rstrip()) if lines[i].rstrip().endswith(":"): break return " ".join(p.strip() for p in parts) except Exception: return "" class PythonExtractor: """ Extracts symbols and imports from Python via the built-in `ast` module. Falls back to regex on SyntaxError (handles Python 2 / stub files). """ def extract( self, filepath: Path, rel: str ) -> tuple[list[dict], list[SymbolInfo]]: chunks: list[dict] = [] symbols: list[SymbolInfo] = [] try: src = filepath.read_text(encoding="utf-8", errors="ignore") lines = src.splitlines() try: tree = ast.parse(src) except SyntaxError: return self._regex(src, lines, rel) # Map node id -> enclosing class name class_of: dict[int, str] = {} for node in ast.walk(tree): if isinstance(node, ast.ClassDef): for child in ast.walk(node): if child is not node: class_of[id(child)] = node.name for node in ast.walk(tree): if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): continue start = node.lineno end = getattr(node, "end_lineno", min(node.lineno + 30, len(lines))) parent = class_of.get(id(node)) if isinstance(node, ast.ClassDef): kind = "class" elif isinstance(node, ast.AsyncFunctionDef): kind = "async_method" if parent else "async_function" else: kind = "method" if parent else "function" calls: list[str] = [] if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): for child in ast.walk(node): if isinstance(child, ast.Call): if isinstance(child.func, ast.Name): calls.append(child.func.id) elif isinstance(child.func, ast.Attribute): calls.append(child.func.attr) sym = SymbolInfo( name=node.name, kind=kind, file=rel, start_line=start, end_line=end, signature=_signature(node, lines), docstring=_safe_docstring(node), parent=parent, language="Python", ) symbols.append(sym) chunks.append({ "file": rel, "start_line": start, "end_line": end, "text": "\n".join(lines[start - 1: end]), "symbols": [node.name], "kind": kind, "calls": list(dict.fromkeys(calls)), "language": "Python", }) except Exception: pass return chunks, symbols def _regex( self, src: str, lines: list[str], rel: str ) -> tuple[list[dict], list[SymbolInfo]]: chunks, symbols = [], [] for m in re.finditer(r"^(async\s+def|def|class)\s+(\w+)", src, re.MULTILINE): ln = src[: m.start()].count("\n") + 1 name = m.group(2) kw = m.group(1).strip() end = min(ln + 30, len(lines)) kind = "async_function" if "async" in kw else ("class" if "class" in kw else "function") chunks.append({ "file": rel, "start_line": ln, "end_line": end, "text": "\n".join(lines[ln - 1: end]), "symbols": [name], "kind": kind, "calls": [], "language": "Python", }) symbols.append(SymbolInfo( name=name, kind=kind, file=rel, start_line=ln, end_line=end, signature="", docstring="", parent=None, language="Python", )) return chunks, symbols def extract_imports(self, filepath: Path) -> list[str]: imports: list[str] = [] try: src = filepath.read_text(encoding="utf-8", errors="ignore") tree = ast.parse(src) for node in ast.walk(tree): if isinstance(node, ast.Import): imports.extend(a.name for a in node.names) elif isinstance(node, ast.ImportFrom) and node.module: imports.append(node.module) except Exception: pass return imports # ───────────────────────────────────────────────────────────────────────────── # Generic extractor (JS / TS / Java / Go / Rust / C# …) # ───────────────────────────────────────────────────────────────────────────── class GenericExtractor: _PAT = re.compile( r"^(?:(?:export\s+)?(?:async\s+)?(?:function|class|def|fn|func|" r"pub(?:\s+(?:async\s+)?fn)?|" r"(?:private|public|protected|static)" r"(?:\s+(?:async\s+)?(?:function|class|void|int|str|bool|string))?)" r"\s+(\w+))", re.MULTILINE, ) def extract( self, filepath: Path, rel: str ) -> tuple[list[dict], list[SymbolInfo]]: chunks, symbols = [], [] try: src = filepath.read_text(encoding="utf-8", errors="ignore") lines = src.splitlines() lang = _LANG_MAP.get(filepath.suffix.lower(), "Unknown") for m in self._PAT.finditer(src): if not m.group(1): continue ln = src[: m.start()].count("\n") + 1 end = min(ln + 40, len(lines)) name = m.group(1) chunks.append({ "file": rel, "start_line": ln, "end_line": end, "text": "\n".join(lines[ln - 1: end]), "symbols": [name], "kind": "symbol", "calls": [], "language": lang, }) symbols.append(SymbolInfo( name=name, kind="symbol", file=rel, start_line=ln, end_line=end, signature=lines[ln - 1].strip() if lines else "", docstring="", parent=None, language=lang, )) if not chunks and 0 < len(lines) <= 200: chunks.append({ "file": rel, "start_line": 1, "end_line": len(lines), "text": src, "symbols": [], "kind": "file", "calls": [], "language": lang, }) except Exception: pass return chunks, symbols # ───────────────────────────────────────────────────────────────────────────── # Import resolver # ───────────────────────────────────────────────────────────────────────────── class ImportResolver: """Maps Python module names to relative file paths inside the repo.""" def __init__(self, repo_path: Path) -> None: self._map: dict[str, str] = {} for root, dirs, files in os.walk(repo_path): dirs[:] = [d for d in dirs if d not in _SKIP_DIRS] for fname in files: if not fname.endswith(".py"): continue fpath = Path(root) / fname try: rel = str(fpath.relative_to(repo_path)) module = rel.replace(os.sep, ".").removesuffix(".py") self._map[module] = rel self._map[module.split(".")[-1]] = rel except ValueError: pass def resolve_many(self, modules: list[str]) -> list[str]: out = [] for m in modules: r = self._map.get(m) if r and r not in out: out.append(r) return out # ───────────────────────────────────────────────────────────────────────────── # Code behavior summarizer (AST-driven, no LLM) # ───────────────────────────────────────────────────────────────────────────── class CodeSummarizer: """ Analyzes Python AST to extract structured behavior and generate a natural-language summary — entirely without an LLM. Answers the FastContext question "What does this function DO?" using: - self.x reads/writes (state access patterns) - function calls (collaborators) - exceptions raised (error contracts) - return type annotation Non-Python symbols get a signature-only summary. """ def analyze(self, chunk: dict, sym_info: Optional[SymbolInfo]) -> CodeBehavior: reads: list[str] = [] writes: list[str] = [] calls: list[str] = [] raises_: list[str] = [] ret_ann: str = "" if sym_info and sym_info.language != "Python": return CodeBehavior(reads=[], writes=[], calls=chunk.get("calls", [])[:6], raises=[], returns="") try: tree = ast.parse(chunk.get("text", "")) for node in ast.walk(tree): # Reads: self.attr if (isinstance(node, ast.Attribute) and isinstance(node.ctx, ast.Load) and isinstance(node.value, ast.Name) and node.value.id == "self"): attr = f"self.{node.attr}" if attr not in reads: reads.append(attr) # Writes: self.attr = ... if isinstance(node, (ast.Assign, ast.AugAssign)): targets = (node.targets if isinstance(node, ast.Assign) else [node.target]) for t in targets: if (isinstance(t, ast.Attribute) and isinstance(t.value, ast.Name) and t.value.id == "self"): attr = f"self.{t.attr}" if attr not in writes: writes.append(attr) # Calls if isinstance(node, ast.Call): if isinstance(node.func, ast.Name): name = node.func.id elif isinstance(node.func, ast.Attribute): name = f"{node.func.attr}()" else: name = None if name and name not in calls: calls.append(name) # Raises if isinstance(node, ast.Raise) and node.exc: if (isinstance(node.exc, ast.Call) and isinstance(node.exc.func, ast.Name)): raises_.append(node.exc.func.id) elif isinstance(node.exc, ast.Name): raises_.append(node.exc.id) # Return annotation from signature if sym_info and sym_info.signature: m = re.search(r"->\s*(.+?):", sym_info.signature) if m: ret_ann = m.group(1).strip() except Exception: pass return CodeBehavior( reads = reads[:6], writes = writes[:4], calls = calls[:8], raises = list(dict.fromkeys(raises_))[:4], returns = ret_ann, ) def summarize( self, sym_info: Optional[SymbolInfo], behavior: CodeBehavior, chunk: dict, ) -> str: """Produce a one-paragraph natural-language summary from AST data.""" parts: list[str] = [] if sym_info: kind_label = sym_info.kind.replace("_", " ") parts.append(f"`{sym_info.name}` is a {sym_info.language} {kind_label}") if sym_info.parent: parts.append(f" on class `{sym_info.parent}`") if sym_info.docstring: parts.append(f" that {sym_info.docstring.rstrip('.').lower()}") else: parts.append(f"Code block in `{chunk['file']}`") details: list[str] = [] if behavior.reads: details.append(f"reads {', '.join(behavior.reads[:3])}") if behavior.writes: details.append(f"writes {', '.join(behavior.writes[:2])}") top_calls = [c for c in behavior.calls[:4] if not c.startswith("self.")] if top_calls: details.append(f"calls {', '.join(top_calls)}") if behavior.raises: details.append(f"raises {', '.join(behavior.raises)}") if behavior.returns: details.append(f"returns {behavior.returns}") if details: parts.append(". It " + ", ".join(details)) if sym_info: parts.append(f". Defined at {sym_info.file}:L{sym_info.start_line}.") return "".join(parts) # ───────────────────────────────────────────────────────────────────────────── # Semantic Index (sentence-transformers, optional) # ───────────────────────────────────────────────────────────────────────────── class SemanticIndex: """ Embedding-based semantic search using sentence-transformers. Bridges the vocabulary gap that defeats BM25: Query: "user authentication flow" BM25 finds: files containing those exact words SemanticIndex finds: login(), verify_token(), check_session() even without keyword overlap Model: all-MiniLM-L6-v2 (22 MB, 22M params, <5 ms on CPU per query). Embeddings are cached to .swiftcontext/embeddings.npy — instant on reload. Gracefully disabled if sentence-transformers is not installed. Install: pip install sentence-transformers """ MODEL_NAME = "all-MiniLM-L6-v2" _EMBS_FILE = "embeddings.npy" def __init__(self) -> None: self._model: Optional[object] = None self._embeds: Optional[object] = None # np.ndarray (N, D) self._built = False if not _HAS_SEMANTIC: return try: self._model = SentenceTransformer(self.MODEL_NAME, device="cpu") except Exception: pass @property def available(self) -> bool: return self._model is not None and self._built def build(self, chunks: list[dict], cache_dir: Optional[Path] = None) -> None: """Encode all chunks. Saves to cache_dir/embeddings.npy if provided.""" if not self._model: return try: texts = [ (c.get("text", "") + " " + " ".join(c.get("symbols", [])))[:512] for c in chunks ] self._embeds = self._model.encode( # type: ignore[union-attr] texts, batch_size=64, show_progress_bar=False, normalize_embeddings=True, ) self._built = True if cache_dir is not None and self._embeds is not None: try: _np.save(str(cache_dir / self._EMBS_FILE), self._embeds) except Exception: pass except Exception: pass def load(self, cache_dir: Path) -> bool: """Load pre-built embeddings from disk. Returns True on success.""" if not self._model or not _HAS_SEMANTIC: return False try: emb_path = cache_dir / self._EMBS_FILE if not emb_path.exists(): return False self._embeds = _np.load(str(emb_path)) self._built = True return True except Exception: return False def search(self, query: str, top_k: int = 10) -> list[tuple[int, float]]: """Return (chunk_index, cosine_similarity) pairs sorted descending.""" if not self.available or self._embeds is None: return [] try: q_emb = self._model.encode( # type: ignore[union-attr] [query], normalize_embeddings=True, show_progress_bar=False ) sims = (_np.dot(self._embeds, q_emb.T)).flatten() top = _np.argsort(-sims)[:top_k] return [(int(i), float(sims[i])) for i in top if sims[i] > 0.20] except Exception: return [] # ───────────────────────────────────────────────────────────────────────────── # BM25 Index (Okapi BM25) # ───────────────────────────────────────────────────────────────────────────── class BM25Index: """ Okapi BM25 — industry-standard IR ranking. Advantages over TF-IDF used in the previous prototype: - Term saturation: diminishing returns for repeated terms - Document-length normalisation: no bias toward long files - Smooth IDF: handles very common vs. rare tokens correctly Parameters k1=1.5, b=0.75 are established defaults; no tuning needed. """ def __init__( self, chunks: list[dict], k1: float = BM25_K1, b: float = BM25_B ) -> None: self.chunks = chunks self.k1, self.b = k1, b self.tf: list[dict[str, int]] = [] self.dl: list[int] = [] self.idf: dict[str, float] = {} self.avgdl: float = 1.0 self._build() @staticmethod def tokenize(text: str) -> list[str]: raw = re.findall(r"[a-zA-Z_][a-zA-Z0-9_]*", text) out: list[str] = [] for t in raw: parts = re.sub(r"([A-Z])", r" \1", t).lower().split() out.extend(parts) out.append(t.lower()) return [t for t in out if len(t) > 1] def _build(self) -> None: N = len(self.chunks) if N == 0: return df: dict[str, int] = defaultdict(int) for chunk in self.chunks: text = chunk["text"] + " " + " ".join(chunk.get("symbols", [])) tokens = self.tokenize(text) tf: dict[str, int] = defaultdict(int) for t in tokens: tf[t] += 1 self.tf.append(dict(tf)) self.dl.append(len(tokens)) for t in tf: df[t] += 1 self.avgdl = sum(self.dl) / N self.idf = { t: math.log((N - cnt + 0.5) / (cnt + 0.5) + 1.0) for t, cnt in df.items() } def search(self, query: str, top_k: int = 10) -> list[tuple[int, float]]: """(chunk_index, bm25_score) sorted by descending score.""" q_tokens = set(self.tokenize(query)) scores: list[tuple[int, float]] = [] for i, tf in enumerate(self.tf): score = 0.0 dl_ratio = self.dl[i] / self.avgdl for t in q_tokens: if t not in tf: continue freq = tf[t] score += self.idf.get(t, 0.0) * ( freq * (self.k1 + 1) / (freq + self.k1 * (1 - self.b + self.b * dl_ratio)) ) if score > 0: scores.append((i, round(score, 6))) scores.sort(key=lambda x: -x[1]) return scores[:top_k] def score_one(self, query: str, idx: int) -> float: if idx >= len(self.tf): return 0.0 q_tokens = set(self.tokenize(query)) tf = self.tf[idx] dl_ratio = self.dl[idx] / self.avgdl score = 0.0 for t in q_tokens: if t not in tf: continue freq = tf[t] score += self.idf.get(t, 0.0) * ( freq * (self.k1 + 1) / (freq + self.k1 * (1 - self.b + self.b * dl_ratio)) ) return round(score, 6) # ───────────────────────────────────────────────────────────────────────────── # Disk cache # ───────────────────────────────────────────────────────────────────────────── class DiskCache: """ Persistent on-disk index stored at {repo}/.swiftcontext/index.json. Uses per-file MD5 hashes so only changed files are re-indexed. The .swiftcontext directory is auto-gitignored on creation. """ def __init__(self, repo_path: Path) -> None: self._dir = repo_path / _INDEX_DIR self._file = self._dir / _INDEX_FILE def load(self) -> Optional[dict]: if not self._file.exists(): return None try: data = json.loads(self._file.read_text(encoding="utf-8")) return data if data.get("version") == _INDEX_VERSION else None except Exception: return None def save(self, data: dict) -> None: try: self._dir.mkdir(parents=True, exist_ok=True) gi = self._dir / ".gitignore" if not gi.exists(): gi.write_text("*\n") self._file.write_text( json.dumps(data, indent=2, default=str), encoding="utf-8" ) except Exception: pass @staticmethod def hash_file(path: Path) -> str: try: return hashlib.md5(path.read_bytes()).hexdigest() except Exception: return "" # ───────────────────────────────────────────────────────────────────────────── # Repository index # ───────────────────────────────────────────────────────────────────────────── class RepoIndex: """ Full repository index: BM25, symbol table, call graph, dep graph. On first run for a repo: walks all code files, builds everything, saves to disk. On subsequent runs: loads from disk in <100 ms. When files change: detects via MD5 and rebuilds only the affected state. """ def __init__(self, repo_path: str | Path, verbose: bool = False) -> None: self.repo_path = Path(repo_path).resolve() self.chunks: list[dict] = [] self.symbols: list[SymbolInfo] = [] self.sym_map: dict[str, list[SymbolInfo]] = defaultdict(list) self.call_graph: dict[str, list[str]] = defaultdict(list) self.rev_call_graph: dict[str, list[str]] = defaultdict(list) self.dep_graph: dict[str, list[str]] = defaultdict(list) self.bm25: Optional[BM25Index] = None self.semantic: SemanticIndex = SemanticIndex() self._verbose = verbose self._build() def _rel(self, path: Path) -> str: try: return str(path.relative_to(self.repo_path)) except ValueError: return str(path) def _scan_files(self) -> dict[str, str]: files: dict[str, str] = {} for root, dirs, fnames in os.walk(self.repo_path): dirs[:] = [d for d in dirs if d not in _SKIP_DIRS] for fname in fnames: fpath = Path(root) / fname if fpath.suffix.lower() in _CODE_EXTS: files[self._rel(fpath)] = DiskCache.hash_file(fpath) return files def _build(self) -> None: cache = DiskCache(self.repo_path) current = self._scan_files() cached = cache.load() if cached and cached.get("file_hashes") == current: self._from_cache(cached) if self._verbose: print(f" [Index] cache hit — {len(self.chunks)} chunks, " f"{len(self.symbols)} symbols") return py_ext = PythonExtractor() gen_ext = GenericExtractor() raw_imports: dict[str, list[str]] = {} for root, dirs, fnames in os.walk(self.repo_path): dirs[:] = [d for d in dirs if d not in _SKIP_DIRS] for fname in fnames: fpath = Path(root) / fname ext = fpath.suffix.lower() if ext not in _CODE_EXTS: continue rel = self._rel(fpath) if ext == ".py": chunks, syms = py_ext.extract(fpath, rel) raw_imports[rel] = py_ext.extract_imports(fpath) else: chunks, syms = gen_ext.extract(fpath, rel) self.chunks.extend(chunks) self.symbols.extend(syms) for s in syms: self.sym_map[s.name.lower()].append(s) for chunk in chunks: caller = (chunk.get("symbols") or [None])[0] if caller and chunk.get("calls"): ca = caller.lower() for callee in chunk["calls"]: ce = callee.lower() if ce not in self.call_graph[ca]: self.call_graph[ca].append(ce) if ca not in self.rev_call_graph[ce]: self.rev_call_graph[ce].append(ca) resolver = ImportResolver(self.repo_path) for file, mods in raw_imports.items(): self.dep_graph[file] = resolver.resolve_many(mods)[:5] self.bm25 = BM25Index(self.chunks) self.semantic.build(self.chunks, cache_dir=cache._dir) cache.save({ "version": _INDEX_VERSION, "file_hashes": current, "chunks": self.chunks, "symbols": [asdict(s) for s in self.symbols], "call_graph": dict(self.call_graph), "rev_call_graph": dict(self.rev_call_graph), "dep_graph": dict(self.dep_graph), }) if self._verbose: n_files = len({c["file"] for c in self.chunks}) print(f" [Index] built — {len(self.chunks)} chunks, " f"{len(self.symbols)} symbols, {n_files} files") def _from_cache(self, cached: dict) -> None: self.chunks = cached.get("chunks", []) fields = set(SymbolInfo.__dataclass_fields__) for s in cached.get("symbols", []): sym = SymbolInfo(**{k: v for k, v in s.items() if k in fields}) self.symbols.append(sym) self.sym_map[sym.name.lower()].append(sym) self.call_graph = defaultdict(list, cached.get("call_graph", {})) self.rev_call_graph = defaultdict(list, cached.get("rev_call_graph", {})) self.dep_graph = defaultdict(list, cached.get("dep_graph", {})) self.bm25 = BM25Index(self.chunks) if not self.semantic.load(DiskCache(self.repo_path)._dir): self.semantic.build(self.chunks) # ───────────────────────────────────────────────────────────────────────────── # Multi-signal ranker # ───────────────────────────────────────────────────────────────────────────── class MultiSignalRanker: """ Combines four independent relevance signals: Signal 1 BM25 score — normalised to [0, 1] Signal 2 Exact symbol match — +0.40 when query token matches symbol name Signal 3 Path relevance — +0.15 when query mentions dir / filename Signal 4 Kind bonus — +0.20 for definitions in pinpoint queries Final score is capped at 1.0. """ def rank( self, query: str, candidates: list[tuple[int, float]], chunks: list[dict], sym_map: dict[str, list[SymbolInfo]], strategy: str, top_k: int, ) -> list[tuple[int, float]]: if not candidates: return [] max_bm25 = max((s for _, s in candidates), default=1.0) or 1.0 q_lower = query.lower() q_tokens = set(re.findall(r"[a-z][a-z0-9_]{1,}", q_lower)) out: list[tuple[int, float]] = [] for idx, raw in candidates: if idx >= len(chunks): continue chunk = chunks[idx] score = raw / max_bm25 # signal 1 for sym_name in chunk.get("symbols", []): sl = sym_name.lower() if sl in q_lower or any(t in sl for t in q_tokens if len(t) > 3): score += 0.40 # signal 2 break for part in Path(chunk["file"]).parts: if part.lower().removesuffix(".py") in q_lower: score += 0.15 # signal 3 break if strategy == "pinpoint_cite" and chunk.get("kind") in { "class", "function", "async_function", "method", "async_method" }: score += 0.20 # signal 4 out.append((idx, round(min(score, 1.0), 6))) out.sort(key=lambda x: -x[1]) return out[:top_k] # ───────────────────────────────────────────────────────────────────────────── # Strategy router (DistilBERT + heuristic layer) # ───────────────────────────────────────────────────────────────────────────── class SwiftContextRouter: """ 66M DistilBERT router with a heuristic pre-classification layer. Classification pipeline (first match wins): 1. Broad pattern → broad_scan "how does X work" 2. Targeted pattern → targeted_search "all callers of X" 3. Pinpoint pattern → pinpoint_cite "find class X" 4. DistilBERT model → any strategy 5. Low-confidence → broad_scan (safe fallback) The heuristic layer corrects the most common misrouting cases for queries phrased differently from the synthetic training data. """ _BROAD = re.compile( r"\b(?:how\s+does|explain\s+(?:the|how)|overview\s+of|" r"architecture\s+of|walk\s+me\s+through|understand\s+(?:the|how)|" r"what\s+(?:is|are)\s+the\s+(?:main|overall|whole|general|full))\b", re.IGNORECASE, ) _TARGETED = re.compile( r"\b(?:all\s+(?:usages?|calls?|references?|callers?|occurrences?)|" r"who\s+calls?|where\s+(?:is\s+)?(?:it\s+)?(?:called|used|imported)|" r"every\s+(?:place|location|file)\s+(?:that\s+)?(?:uses?|calls?)|" r"usages?\s+of|references?\s+to)\b", re.IGNORECASE, ) _PINPOINT_ACTION = re.compile( r"\b(?:find|locate|show\s+me|where\s+is|jump\s+to|go\s+to|" r"definition\s+of|source\s+of|implementation\s+of)\b", re.IGNORECASE, ) _IDENTIFIER = re.compile( r"`[^`]+`|[A-Z][a-zA-Z0-9]{2,}|[a-z][a-z0-9]*(?:_[a-z0-9]+){1,}" ) def __init__(self, model_path: str) -> None: device = "cuda" if torch.cuda.is_available() else "cpu" self.tokenizer = AutoTokenizer.from_pretrained(model_path) self.model = AutoModelForSequenceClassification.from_pretrained(model_path) self.model.eval() self.model.to(device) self.device = device def predict(self, query: str) -> tuple[str, float]: if self._BROAD.search(query): return "broad_scan", 0.92 if self._TARGETED.search(query): return "targeted_search", 0.88 if self._PINPOINT_ACTION.search(query) and self._IDENTIFIER.search(query): return "pinpoint_cite", 0.85 inputs = self.tokenizer( query, return_tensors="pt", truncation=True, padding="max_length", max_length=128, ).to(self.device) with torch.no_grad(): probs = torch.softmax(self.model(**inputs).logits, dim=-1)[0] idx = int(probs.argmax()) conf = float(probs[idx]) if conf < CONFIDENCE_FALLBACK: return "broad_scan", conf return STRATEGY_LABELS[idx], conf # ───────────────────────────────────────────────────────────────────────────── # Pipeline # ───────────────────────────────────────────────────────────────────────────── class SwiftContextPipeline: """ SwiftContext production pipeline — three API methods. explore(query, repo_path) — BM25 + multi-signal citation search trace(symbol, repo_path) — call chain: callers + callees [NEW vs FC] explain(symbol, repo_path) — signature, docstring, deps [NEW vs FC] Index is disk-persisted and incrementally updated. All three methods consume 0 LLM tokens. """ def __init__(self, router_path: str) -> None: self.router = SwiftContextRouter(router_path) self._ranker = MultiSignalRanker() self._summarizer = CodeSummarizer() self._cache: dict[str, RepoIndex] = {} # ── internal helpers ───────────────────────────────────────────────────── def _index(self, repo_path: str, verbose: bool = False) -> RepoIndex: key = str(Path(repo_path).resolve()) if key not in self._cache: self._cache[key] = RepoIndex(repo_path, verbose=verbose) return self._cache[key] def _make_citation( self, chunk: dict, query: str, strategy: str, score: float, index: RepoIndex, ) -> Citation: sym_info: Optional[SymbolInfo] = None for s in index.symbols: if s.file == chunk["file"] and s.start_line == chunk["start_line"]: sym_info = s break label = f"`{chunk['symbols'][0]}`" if chunk.get("symbols") else "this block" q50 = query[:50].rstrip() if strategy == "pinpoint_cite": reason = f"Direct definition of {label} — exact AST symbol match" elif strategy == "targeted_search": reason = f"{label} is directly relevant to '{q50}'" else: reason = f"{label} is broadly relevant to the query scope" snippet = chunk["text"] if len(snippet) > 400: snippet = snippet[:400] + "..." return Citation( file = chunk["file"], start_line = chunk["start_line"], end_line = chunk["end_line"], snippet = snippet, relevance = round(min(score, 1.0), 4), reason = reason, symbol = sym_info, docstring = sym_info.docstring if sym_info else "", deps = index.dep_graph.get(chunk["file"], [])[:3], ) def _pinpoint_hits( self, query: str, idx: RepoIndex, top_k: int ) -> list[tuple[int, float]]: """Exact symbol name lookup via symbol table — O(1) per candidate.""" backtick = re.findall(r"`([^`]+)`", query) camel = re.findall(r"\b([A-Z][a-zA-Z0-9]{1,})\b", query) snake = re.findall(r"\b([a-z][a-z0-9]*(?:_[a-z0-9]+){1,})\b", query) cands = list(dict.fromkeys(c.lower() for c in backtick + camel + snake)) hits: list[tuple[int, float]] = [] seen: set[tuple] = set() for cand in cands: for sym_info in idx.sym_map.get(cand, []): key = (sym_info.file, sym_info.start_line) if key in seen: continue seen.add(key) for ci, chunk in enumerate(idx.chunks): if (chunk["file"] == sym_info.file and chunk["start_line"] == sym_info.start_line): hits.append((ci, 2.0)) # boosted above any BM25 score break return hits[:top_k] # ── explore() ──────────────────────────────────────────────────────────── def explore( self, query: str, repo_path: str, top_k: int = 5, verbose: bool = False, ) -> ExploreResult: """ Find relevant code citations in repo_path for the given query. Uses BM25 retrieval + 4-signal ranking + exact symbol lookup. Persistent disk index means subsequent calls for the same repo are instant (no rebuild). Zero LLM tokens consumed. Args: query : natural-language or code-specific query repo_path : root of the repository to explore top_k : max citations to return (default 5) verbose : print routing + index details Returns: ExploreResult with structured citations, strategy, and metrics. """ t0 = time.perf_counter() strategy, conf = self.router.predict(query) if verbose: print(f" [Router] strategy={strategy} confidence={conf:.3f}") idx = self._index(repo_path, verbose=verbose) turns = 1 k = top_k * (3 if strategy == "broad_scan" else 4) bm25_hits = idx.bm25.search(query, k) if idx.bm25 else [] # Semantic blending: merge embedding results for conceptual queries. # Covers vocabulary gaps BM25 cannot handle ("authentication" → login()). if idx.semantic.available: sem_hits = idx.semantic.search(query, top_k * 2) existing = {i for i, _ in bm25_hits} # Scale semantic scores into BM25 range before merging max_bm25 = max((s for _, s in bm25_hits), default=1.0) or 1.0 for si, sscore in sem_hits: if si not in existing: bm25_hits.append((si, sscore * max_bm25 * 0.6)) if strategy == "pinpoint_cite": exact = self._pinpoint_hits(query, idx, top_k) existing = {i for i, _ in bm25_hits} for ci, boost in exact: if ci not in existing: bm25_hits.insert(0, (ci, boost)) elif strategy == "broad_scan": turns = 2 ranked = self._ranker.rank( query, bm25_hits, idx.chunks, idx.sym_map, strategy, top_k ) citations: list[Citation] = [] seen: set[tuple] = set() for ci, score in ranked: if ci >= len(idx.chunks): continue chunk = idx.chunks[ci] key = (chunk["file"], chunk["start_line"]) if key in seen: continue seen.add(key) citations.append(self._make_citation(chunk, query, strategy, score, idx)) citations.sort(key=lambda c: -c.relevance) fc_turns = _FC_BASELINE_TURNS[strategy] saved_pct = round( min(max(0, fc_turns - turns) * 800 / _FC_AVG_TOKENS, 1.0) * 100, 1 ) return ExploreResult( citations = citations, confidence = round(conf, 4), strategy_used = strategy, turns_used = turns, tokens_used = 0, tokens_saved_pct = saved_pct, latency_ms = round((time.perf_counter() - t0) * 1000, 1), index_chunks = len(idx.chunks), index_symbols = len(idx.symbols), ) # ── trace() ────────────────────────────────────────────────────────────── def trace( self, symbol: str, repo_path: str, verbose: bool = False, ) -> TraceResult: """ Call-chain analysis for `symbol`. NOT available in FastContext. Walks the AST-derived call graph to find: - definition : exact file + line where the symbol is defined - callers : all functions that call this symbol - callees : all functions called by this symbol Uses the reverse call graph for O(k) caller lookup instead of O(n*k). Args: symbol : exact symbol name (case-insensitive) repo_path : root of the repository Returns: TraceResult with definition, callers, and callees as Citations. """ t0 = time.perf_counter() idx = self._index(repo_path) sym_lo = symbol.lower() # Definition definition: Optional[Citation] = None for sym_info in idx.sym_map.get(sym_lo, [])[:1]: for chunk in idx.chunks: if (chunk["file"] == sym_info.file and chunk["start_line"] == sym_info.start_line): definition = self._make_citation( chunk, symbol, "pinpoint_cite", 1.0, idx ) break # Callees — functions called BY this symbol callees: list[Citation] = [] seen_ce: set[str] = set() for callee_name in idx.call_graph.get(sym_lo, [])[:15]: if callee_name in seen_ce: continue seen_ce.add(callee_name) for sym_info in idx.sym_map.get(callee_name, [])[:1]: for chunk in idx.chunks: if (chunk["file"] == sym_info.file and chunk["start_line"] == sym_info.start_line): callees.append(self._make_citation( chunk, callee_name, "pinpoint_cite", 0.80, idx )) break # Callers — functions that call this symbol (O(k) via reverse graph) callers: list[Citation] = [] seen_ca: set[str] = set() for caller_name in idx.rev_call_graph.get(sym_lo, [])[:15]: if caller_name in seen_ca: continue seen_ca.add(caller_name) for sym_info in idx.sym_map.get(caller_name, [])[:1]: for chunk in idx.chunks: if (chunk["file"] == sym_info.file and chunk["start_line"] == sym_info.start_line): callers.append(self._make_citation( chunk, caller_name, "pinpoint_cite", 0.70, idx )) break if verbose: print( f" [Trace] '{symbol}': " f"{len(callers)} caller(s), {len(callees)} callee(s)" ) return TraceResult( symbol = symbol, definition = definition, callers = callers, callees = callees, latency_ms = round((time.perf_counter() - t0) * 1000, 1), ) # ── explain() ──────────────────────────────────────────────────────────── def explain( self, symbol: str, repo_path: str, ) -> Optional[ExplainResult]: """ Extract documentation for `symbol`. NOT available in FastContext. Returns the symbol's signature, docstring, language, and the files it directly imports — all from the AST index, no LLM required. Args: symbol : exact symbol name (case-insensitive) repo_path : root of the repository Returns: ExplainResult, or None if the symbol is not found in the index. """ t0 = time.perf_counter() idx = self._index(repo_path) sym_lo = symbol.lower() found = idx.sym_map.get(sym_lo, []) if not found: return None s = found[0] return ExplainResult( symbol = s.name, kind = s.kind, signature = s.signature, docstring = s.docstring, file = s.file, start_line = s.start_line, end_line = s.end_line, language = s.language, deps = idx.dep_graph.get(s.file, [])[:5], latency_ms = round((time.perf_counter() - t0) * 1000, 1), ) # ── summarize() ────────────────────────────────────────────────────────── def summarize( self, symbol: str, repo_path: str, ) -> Optional[SummarizeResult]: """ Generate a natural-language behavior summary for `symbol`. NOT available in FastContext. Analyzes AST to answer "What does this symbol DO?" without any LLM: - What state does it read / write (self.x) - What functions / methods does it call - What exceptions does it raise - What does it return Works for all Python symbols. Non-Python symbols return a signature-only summary. Args: symbol : exact symbol name (case-insensitive) repo_path : root of the repository Returns: SummarizeResult or None if symbol not found. """ t0 = time.perf_counter() idx = self._index(repo_path) sym_lo = symbol.lower() found = idx.sym_map.get(sym_lo, []) if not found: return None sym_info = found[0] chunk: Optional[dict] = None for c in idx.chunks: if c["file"] == sym_info.file and c["start_line"] == sym_info.start_line: chunk = c break if chunk is None: return None behavior = self._summarizer.analyze(chunk, sym_info) summary = self._summarizer.summarize(sym_info, behavior, chunk) return SummarizeResult( symbol = sym_info.name, kind = sym_info.kind, summary = summary, behavior = behavior, file = sym_info.file, start_line = sym_info.start_line, end_line = sym_info.end_line, latency_ms = round((time.perf_counter() - t0) * 1000, 1), ) # ── context() ──────────────────────────────────────────────────────────── def context( self, query: str, repo_path: str, top_k: int = 3, verbose: bool = False, ) -> ContextResult: """ Build a multi-file context window for `query`. NOT available in FastContext. FastContext had to call a 4B LLM 2-3 times to browse the repo and build this context. SwiftContext does it deterministically in <50 ms. The returned ContextResult.to_llm_context() produces a ready-to-use context string you can pass to ANY downstream LLM (GPT-4, Claude …) for deep reasoning over real code — zero hallucination of file contents. Args: query : conceptual question, e.g. "why does auth fail on expiry?" repo_path : root of the repository top_k : max primary citations (default 3) verbose : print context-building details Returns: ContextResult with primary citations, caller/callee context, per-symbol summaries, and to_llm_context() formatter. """ t0 = time.perf_counter() idx = self._index(repo_path, verbose=verbose) # Primary: best matches for the query primary = self.explore(query, repo_path, top_k=top_k).citations # Expand: callers + callees of primary matches (cross-file context) caller_context: list[Citation] = [] callee_context: list[Citation] = [] seen_keys: set[tuple] = {(c.file, c.start_line) for c in primary} for cit in primary[:2]: if not cit.symbol: continue tr = self.trace(cit.symbol.name, repo_path) for c in tr.callers[:2]: k = (c.file, c.start_line) if k not in seen_keys: caller_context.append(c) seen_keys.add(k) for c in tr.callees[:3]: k = (c.file, c.start_line) if k not in seen_keys: callee_context.append(c) seen_keys.add(k) # Natural-language summaries for primary symbols summaries: dict[str, str] = {} for cit in primary: if cit.symbol: sr = self.summarize(cit.symbol.name, repo_path) if sr: summaries[cit.symbol.name] = sr.summary # Estimate LLM token cost (~4 chars / token) total_chars = sum( len(c.snippet) for c in primary + caller_context + callee_context ) token_est = total_chars // 4 if verbose: print(f" [Context] {len(primary)} primary, " f"{len(caller_context)} caller, " f"{len(callee_context)} callee, ~{token_est} tokens") return ContextResult( query = query, primary = primary, caller_context = caller_context, callee_context = callee_context, summaries = summaries, total_tokens_est = token_est, latency_ms = round((time.perf_counter() - t0) * 1000, 1), ) # ───────────────────────────────────────────────────────────────────────────── # Demo # ───────────────────────────────────────────────────────────────────────────── def demo(repo_path: str = ".") -> None: """Live demo — SwiftContext explores the SwiftContext codebase itself.""" ROUTER = "./model/final" W = 70 print("=" * W) print("SwiftContext — Production Demo (zero LLM tokens, no FastContext)") print(f"Router : {ROUTER}") print(f"Repo : {Path(repo_path).resolve()}") print("=" * W) sc = SwiftContextPipeline(router_path=ROUTER) # ── explore() ──────────────────────────────────────────────────────────── print(f"\n{'─'*W}") print(" explore() — BM25 + 4-signal ranked code citation search") print(f"{'─'*W}") for q in [ "Find the BM25Index class", "Where is the SwiftContextRouter predict method?", "How does the whole pipeline indexing work?", ]: r = sc.explore(q, repo_path, top_k=3, verbose=True) print(f" query : {q!r}") print(f" strategy : {r.strategy_used} conf={r.confidence} " f"latency={r.latency_ms} ms tokens={r.tokens_used} " f"(FC avg ~{_FC_AVG_TOKENS}) saved={r.tokens_saved_pct}%") for c in r.citations[:2]: print(f" [{c.relevance:.2f}] {c.file}:L{c.start_line}-{c.end_line} {c.reason}") if c.docstring: print(f" doc: {c.docstring[:80]}") print() # ── trace() ────────────────────────────────────────────────────────────── print(f"{'─'*W}") print(" trace() — call-chain analysis [NEW — not in FastContext]") print(f"{'─'*W}") for sym in ["explore", "_build", "search"]: tr = sc.trace(sym, repo_path, verbose=True) cname = lambda c: c.symbol.name if c.symbol else "?" print(f" {tr.symbol!r} ({tr.latency_ms} ms)") if tr.definition: print(f" defined : {tr.definition.file}:L{tr.definition.start_line}") print(f" callers : {[cname(c) for c in tr.callers[:5]]}") print(f" callees : {[cname(c) for c in tr.callees[:5]]}") print() # ── explain() ──────────────────────────────────────────────────────────── print(f"{'─'*W}") print(" explain() — symbol documentation [NEW — not in FastContext]") print(f"{'─'*W}") for sym in ["BM25Index", "SwiftContextRouter", "RepoIndex", "MultiSignalRanker"]: ex = sc.explain(sym, repo_path) if ex: print(f" {ex.symbol} ({ex.kind}, {ex.language})") print(f" sig : {ex.signature}") print(f" docstring : {ex.docstring[:90] or "(none)"}") print(f" location : {ex.file}:L{ex.start_line}-{ex.end_line}") print(f" deps : {ex.deps}") print(f" latency : {ex.latency_ms} ms") print() # ── summarize() ────────────────────────────────────────────────────────── print(f"{'─'*W}") print(" summarize() — AST behavior analysis [NEW — not in FastContext]") print(f"{'─'*W}") for sym in ["search", "_build", "rank"]: sr = sc.summarize(sym, repo_path) if sr: print(f" {sr.symbol} ({sr.kind}) {sr.latency_ms} ms") print(f" summary : {sr.summary[:130]}") if sr.behavior.reads: print(f" reads : {sr.behavior.reads[:3]}") if sr.behavior.calls: print(f" calls : {[c for c in sr.behavior.calls if not c.startswith('self.')][:4]}") if sr.behavior.raises: print(f" raises : {sr.behavior.raises}") print() # ── context() ──────────────────────────────────────────────────────────── print(f"{'─'*W}") print(" context() — LLM-ready context window [replaces FC's LLM browsing]") print(f"{'─'*W}") ctx = sc.context( "How does the BM25 search ranking work end to end?", repo_path, top_k=2, verbose=True, ) print(f" query : {ctx.query!r}") print(f" primary : {len(ctx.primary)} citations") print(f" caller context : {len(ctx.caller_context)} citations") print(f" callee context : {len(ctx.callee_context)} citations") print(f" summaries : {list(ctx.summaries.keys())}") print(f" ~{ctx.total_tokens_est} LLM tokens | latency {ctx.latency_ms} ms") print(f" (FastContext built equivalent context in 2-3 LLM turns = ~{_FC_AVG_TOKENS * 3} tokens)") print() print(" to_llm_context() preview (first 600 chars):") llm_ctx = ctx.to_llm_context() print(" " + llm_ctx[:600].replace("\n", "\n ")) print() if __name__ == "__main__": demo()