diff --git a/README.md b/README.md index da1b895..0071657 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ A small local code retrieval model and a tool for giving a larger coding model useful source context. -Micro-scout indexes a repository, combines lexical and neural search, and returns **verified file paths, line ranges, and bounded source snippets**. Its MCP server keeps the model in memory between requests. +Micro-scout indexes a repository, supports neural, lexical, and hybrid search, and returns **verified file paths, line ranges, and bounded source snippets**. Its MCP server keeps the model in memory between requests. ## Version 0.1 diff --git a/docs/TRAINING.md b/docs/TRAINING.md index d5a270d..f89823d 100644 --- a/docs/TRAINING.md +++ b/docs/TRAINING.md @@ -6,7 +6,7 @@ Version 0.1 fine-tunes `sentence-transformers/all-MiniLM-L6-v2` as a **shared bi For each batch, the query's paired function is the labeled positive. Other functions in the batch are treated as negatives. The loss averages query-to-code and code-to-query cross-entropy. A temperature scales cosine scores. The model learns ranking without generating text. -At inference, the repository's vectors are already in the index. Each request encodes only its query, compares vectors, and combines rankings with BM25. This makes persistent inference practical without Ollama or vLLM. +At inference, the repository's vectors are already in the index. Each request encodes only its query and compares vectors. Optional hybrid mode combines rankings with BM25. This makes persistent inference practical without Ollama or vLLM. Dense mode is the default because it performed best on complete validation; the fusion weights were not tuned. ## Reproduce the local run @@ -73,7 +73,7 @@ OPENBLAS_NUM_THREADS=1 uv run --no-sync micro-scout benchmark \ --output runs/minilm-v1/latency.json ``` -This measures the warm harness, including search, context assembly, and checking returned files. Startup is reported separately. Five repeated development queries are used; these are not representative production traffic. Index construction is also reported separately. +This measures the warm harness, including search, context assembly, and checking returned files. `startup_ms` covers model/index initialization after the core imports, not the entire interpreter startup. Five repeated development queries are used; these are not representative production traffic. Index construction is also reported separately. ## What these results cannot establish diff --git a/docs/USAGE.md b/docs/USAGE.md index 8ffc762..467621d 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -4,6 +4,8 @@ `scout_search` accepts a natural-language query, `top_k` from 1 to 50, and `max_chars` from 100 to 100,000. It returns ranked snippets and up to three graph neighbors within the shared source-character budget. Overlapping source lines are not repeated. Scores are not probabilities. +With a model, the default `mode` is `dense`, selected from the complete validation results. `lexical` uses BM25, which is useful for exact identifiers, and `hybrid` combines the two rankings. Without a model the CLI and MCP server default to `lexical`. The CLI exposes the same choice as `--mode`. + Markdown is excluded from search by default; set `include_docs=true` to include it. An optional `language` filter narrows candidates before rank fusion, for example to `python` or `typescript`. CLI equivalents are `--include-docs` and `--language python`. Each snippet includes an opaque symbol ID, repository-relative path, line range, source text, file SHA-256, and verification/truncation flags. The response also identifies the index snapshot and model fingerprint. Truncation occurs at line boundaries. @@ -40,6 +42,8 @@ The server uses the official [MCP Python SDK v1 maintenance line](https://py.sdk Indexing honors Git's ignored-file rules when available and skips hidden directories, common dependency/build directories, symlinks, unsupported extensions, and files above 1 MB. Python symbols use AST boundaries. Other supported text/code formats use 60-line chunks. Incomplete Python falls back to chunks. +Python functions longer than 48 lines also get 32-line search fragments with an 8-line overlap. The full symbol remains available through each fragment's `parent_id`. This gives the retriever access to code beyond the beginning of a long function. Neural input is still capped at 256 tokens per candidate; very long individual lines can exceed that representation budget. BM25 searches the complete candidate text. + An index is one SQLite file, atomically replaced after a successful build. A model fingerprint prevents combining incompatible query and code embeddings. Reindexing with unchanged weights reuses embeddings for unchanged normalized code. Returned files are checked against their indexed content hashes. Modified or missing files are excluded with warnings; reads of stale IDs fail. New files require reindexing. This is a snapshot workflow, not a filesystem watcher. Restart the server after rebuilding an index so it loads the new snapshot. diff --git a/src/micro_scout/cli.py b/src/micro_scout/cli.py index 92b8f4f..b149676 100644 --- a/src/micro_scout/cli.py +++ b/src/micro_scout/cli.py @@ -32,6 +32,7 @@ def main() -> None: if name == "benchmark": command.add_argument("--iterations", type=int, default=100) command.add_argument("--output", type=Path) + command.add_argument("--mode", choices=["lexical", "dense", "hybrid"], default=None) for command in (index_parser, *(sub.choices[n] for n in ("search", "serve", "benchmark"))): command.add_argument("--model", help="Local weights directory or Hugging Face model ID") command.add_argument("--device", default="cpu") @@ -62,7 +63,7 @@ def main() -> None: args.query, top_k=args.top_k, max_chars=args.max_chars, - mode=args.mode or ("hybrid" if encoder else "lexical"), + mode=args.mode or ("dense" if encoder else "lexical"), expand=not args.no_expand, language=args.language, include_docs=args.include_docs, @@ -80,7 +81,7 @@ def main() -> None: "remove documentation strings from Python functions", "combine lexical and neural search rankings", ] - mode = "hybrid" if encoder else "lexical" + mode = args.mode or ("dense" if encoder else "lexical") for query in queries: scout.search(query, mode=mode) times = [] diff --git a/src/micro_scout/encoder.py b/src/micro_scout/encoder.py index 39b215e..f12a5ae 100644 --- a/src/micro_scout/encoder.py +++ b/src/micro_scout/encoder.py @@ -61,6 +61,17 @@ class Encoder: self.dimension = self.model.config.hidden_size self.parameter_count = sum(p.numel() for p in self.model.parameters()) signature = hashlib.sha256(json.dumps(self.config, sort_keys=True).encode()) + # Weights alone do not identify an encoder: tokenization and model + # configuration can change the vectors without changing the weights. + tokenizer_config = json.loads(self.tokenizer.backend_tokenizer.to_str()) + for runtime_option in ("padding", "truncation"): + tokenizer_config.pop(runtime_option, None) + model_config = self.model.config.to_dict() + for provenance in ("_name_or_path", "transformers_version"): + model_config.pop(provenance, None) + signature.update(json.dumps(tokenizer_config, sort_keys=True).encode()) + signature.update(json.dumps(self.tokenizer.special_tokens_map, sort_keys=True).encode()) + signature.update(json.dumps(model_config, sort_keys=True).encode()) if local: weights = sorted(Path(model).glob("*.safetensors")) if not weights: diff --git a/src/micro_scout/evaluate.py b/src/micro_scout/evaluate.py index 6093e25..945f887 100644 --- a/src/micro_scout/evaluate.py +++ b/src/micro_scout/evaluate.py @@ -65,6 +65,8 @@ def evaluate( "candidates_per_query": len(rows), "repositories": len({r["repo"] for r in rows}), "dataset_sha256": dataset_hash, + "evaluator_source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "fusion": {"dense_weight": 0.5, "per_method_limit": 100, "k": 60}, "evaluation_ids_sha256": hashlib.sha256( json.dumps([r["id"] for r in rows]).encode() ).hexdigest(), @@ -115,6 +117,9 @@ def evaluate( report["trained_hybrid_vs_bm25_mrr"] = paired_mrr_interval( ranks["trained_hybrid"], ranks["bm25"], [r["repo"] for r in rows] ) + report["trained_dense_vs_bm25_mrr"] = paired_mrr_interval( + ranks["trained_dense"], ranks["bm25"], [r["repo"] for r in rows] + ) report["elapsed_seconds"] = time.perf_counter() - started atomic_json(output / f"{split}-metrics.json", report) write_jsonl( diff --git a/src/micro_scout/scout.py b/src/micro_scout/scout.py index 3cd4a5c..9174d87 100644 --- a/src/micro_scout/scout.py +++ b/src/micro_scout/scout.py @@ -88,6 +88,7 @@ class Scout: "path": symbol.path, "name": symbol.name, "kind": symbol.kind, + "parent_id": symbol.parent, "language": symbol.language, "start_line": symbol.start_line, "end_line": end, @@ -104,7 +105,7 @@ class Scout: *, top_k: int = 6, max_chars: int = 12_000, - mode: str = "hybrid", + mode: str = "dense", expand: bool = True, language: str | None = None, include_docs: bool = False, diff --git a/src/micro_scout/server.py b/src/micro_scout/server.py index 62e09ab..1cb251d 100644 --- a/src/micro_scout/server.py +++ b/src/micro_scout/server.py @@ -30,18 +30,21 @@ def create_server(scout: Scout): max_chars: int = 12_000, language: str | None = None, include_docs: bool = False, + mode: str | None = None, ) -> dict[str, Any]: """Find relevant code. English queries are the evaluated language. Returns verified paths, line ranges, bounded source text, and approximate graph neighbors. max_chars counts source characters, not model tokens. Markdown is excluded unless include_docs is true. Optional language filter examples: python, typescript, rust. + mode: dense (default with a model), lexical, or hybrid. Dense performed best + on held-out English descriptions. Try lexical for an exact identifier. No code is executed. No source files are changed. """ return scout.search( query, top_k=top_k, max_chars=max_chars, - mode="hybrid" if scout.encoder else "lexical", + mode=mode or ("dense" if scout.encoder else "lexical"), language=language, include_docs=include_docs, ) diff --git a/src/micro_scout/symbols.py b/src/micro_scout/symbols.py index e2e90c8..9513c15 100644 --- a/src/micro_scout/symbols.py +++ b/src/micro_scout/symbols.py @@ -167,6 +167,16 @@ def parse_source(relative: str, text: str, *, chunk_lines: int = 60) -> list[Sym walk(child, prefix, parent) walk(tree) + # Long functions need searchable windows beyond the encoder's first + # 256 tokens. Keep the full symbol so callers can read its context. + for symbol in list(symbols): + if symbol.kind != "function" or symbol.end_line - symbol.start_line + 1 <= 48: + continue + for start in range(symbol.start_line, symbol.end_line + 1, 24): + end = min(start + 31, symbol.end_line) + add(f"{symbol.name}::<{start}-{end}>", "fragment", start, end, symbol.id) + if end == symbol.end_line: + break covered = set() for symbol in symbols: covered.update(range(symbol.start_line, symbol.end_line + 1)) diff --git a/tests/test_encoder.py b/tests/test_encoder.py index 7b4c6b0..906787a 100644 --- a/tests/test_encoder.py +++ b/tests/test_encoder.py @@ -119,3 +119,18 @@ def test_training_updates_weights_and_can_resume(tiny_model, tmp_path): after.save(output / "last") with pytest.raises(ValueError, match="weights and optimizer state"): train(data, output, config, "cpu", output / "last") + + +def test_fingerprint_detects_tokenizer_change(tiny_model): + original = Encoder(str(tiny_model), max_length=32, query_length=16) + path = tiny_model / "tokenizer.json" + config = json.loads(path.read_text()) + vocab = config["model"]["vocab"] + vocab["read"], vocab["write"] = vocab["write"], vocab["read"] + path.write_text(json.dumps(config)) + changed = Encoder(str(tiny_model), max_length=32, query_length=16) + assert changed.fingerprint != original.fingerprint + assert ( + changed.tokenize(["read file"])["input_ids"].tolist() + != original.tokenize(["read file"])["input_ids"].tolist() + ) diff --git a/tests/test_retrieval.py b/tests/test_retrieval.py index 42575a3..e7074e0 100644 --- a/tests/test_retrieval.py +++ b/tests/test_retrieval.py @@ -263,3 +263,20 @@ def test_source_that_grows_after_indexing_is_bounded(repository, tmp_path): (repository / "numbers.py").write_text("x" * 1_000_001) with pytest.raises(ValueError, match="file-size limit"): scout.read(symbol.id) + + +def test_long_function_tail_has_a_searchable_fragment(repository, tmp_path): + lines = ["def long_function():"] + [f" value_{i} = {i}" for i in range(90)] + lines.append(" return unique_tail_marker") + (repository / "long.py").write_text("\n".join(lines) + "\n") + path = tmp_path / "index.sqlite" + build_index(repository, path) + scout = Scout(Index(path)) + result = scout.search("unique_tail_marker", mode="lexical", top_k=1) + hit = result["results"][0] + assert hit["kind"] == "fragment" + assert "unique_tail_marker" in hit["content"] + assert hit["start_line"] > 48 + parent = scout.read(hit["parent_id"]) + assert parent["name"] == "long_function" + assert parent["start_line"] == 1