vector ranker

Author: carsontung666

README.md

@@ -172,6 +172,17 @@ all      500 queries   no filter, includes ~48% path-signal-less queries
 Output goes to `bench/runs/<timestamp>__<tier>/`: `report.md`, `summary.json`,
 `per_query.jsonl`.
 
+Block mode can optionally rerank only the cross-block merge candidates before
+the file/directory split:
+
+```bash
+python bench/run_swebench_filetree.py --tier medium --strategy block --ranker vector
+```
+
+Available rankers are `none`, `bm25`, and `vector`. The vector ranker uses
+LiteLLM embeddings (`--embedding-provider`, `--embedding-model`) and leaves
+the default `ranker=none` unchanged.
+
 #### Latest Run (Claude Sonnet 4.6, `--strategy block --ranker none`, top-k=10)
 
 The cutoff for each query is its gold-file count: one-gold queries use top-1,

bench/README.md

@@ -20,6 +20,13 @@ Outputs:
 - `report.md`: markdown report
 - `bench.sqlite`: temporary benchmark database
 
+Ranker options:
+
+- `--ranker none`: preserve traversal and block-local LLM order.
+- `--ranker bm25`: lexical path ordering for cross-block merge candidates.
+- `--ranker vector`: embedding path ordering for cross-block merge candidates;
+  configure with `--embedding-provider` and `--embedding-model`.
+
 ### Latest Full Run
 
 Claude Sonnet 4.6, `tier=all`, `strategy=block`, `ranker=none`, `top_k=10`.

bench/run_swebench_filetree.py

@@ -31,20 +31,25 @@
 from contextdb.api.condb import ConDB
 from contextdb.retriever.algorithm.beam_retriever import BeamRetriever
 from contextdb.retriever.algorithm.block_retriever import BlockRetriever
-from contextdb.retriever.algorithm.ranker import BM25PathRanker
+from contextdb.retriever.algorithm.ranker import make_ranker
 
 DEFAULT_MODEL = "claude-sonnet-4-6"
 DEFAULT_DATA_DIR = Path("data/swebench_pathonly")
 
 
 def make_filesystem_retriever(db: ConDB, args, node_count: int):
-    ranker = BM25PathRanker() if args.ranker == "bm25" else None
     strategy = args.strategy
     if strategy == "auto":
         strategy = "beam" if node_count <= 50 else "block"
     if strategy == "beam":
         return BeamRetriever(db.storage, db._llm, mode="filesystem")
     if strategy == "block":
+        ranker = make_ranker(
+            args.ranker,
+            embedding_provider=args.embedding_provider,
+            embedding_model=args.embedding_model,
+            embedding_api_key=args.embedding_api_key,
+        )
         return BlockRetriever(
             db.storage,
             db._llm,
@@ -224,6 +229,8 @@ def run(args):
         "top_k": args.top_k,
         "strategy": args.strategy,
         "ranker": args.ranker,
+        "embedding_provider": args.embedding_provider if args.ranker == "vector" else None,
+        "embedding_model": args.embedding_model if args.ranker == "vector" else None,
         "limit": args.limit,
         "num_queries": len(queries),
         "num_snapshots": len(by_snap),
@@ -476,8 +483,11 @@ def main():
     p.add_argument("--provider", default="anthropic")
     p.add_argument("--top-k", type=int, default=10)
     p.add_argument("--strategy", choices=["auto", "beam", "block"], default="auto")
-    p.add_argument("--ranker", choices=["bm25", "none"], default="none",
+    p.add_argument("--ranker", choices=["bm25", "vector", "none"], default="none",
                    help="Optional path ordering for Block merge results")
+    p.add_argument("--embedding-provider", default="openai")
+    p.add_argument("--embedding-model", default="text-embedding-3-small")
+    p.add_argument("--embedding-api-key", default=None)
     p.add_argument("--max-parallel-blocks", type=int, default=None)
     p.add_argument("--max-turns", type=int, default=None)
     p.add_argument("--limit", type=int, default=0, help="0 = all")

contextdb/api/condb.py

@@ -14,6 +14,7 @@
     BlockRetriever,
     TreeFormatter,
 )
+from contextdb.retriever.algorithm.ranker import Ranker, make_ranker
 
 # ── Errors ──────────────────────────────────────────────────────────
 
@@ -107,6 +108,10 @@ def query(
         max_tokens_per_block: int = 16000,
         cache_enabled: bool = True,
         max_parallel_blocks: int = None,
+        ranker: str | Ranker | None = None,
+        embedding_provider: str = "openai",
+        embedding_model: str = None,
+        embedding_api_key: str = None,
         retriever: BaseRetriever = None,
     ) -> QueryResult:
         self._check_tree(tree_id)
@@ -118,6 +123,10 @@ def query(
                 max_tokens_per_block=max_tokens_per_block,
                 cache_enabled=cache_enabled,
                 max_parallel_blocks=max_parallel_blocks,
+                ranker=ranker,
+                embedding_provider=embedding_provider,
+                embedding_model=embedding_model,
+                embedding_api_key=embedding_api_key,
             )
 
         result = retriever.retrieve(tree_id, question,
@@ -136,6 +145,16 @@ def _make_retriever(self, tree_id, llm, strategy, **kwargs) -> BaseRetriever:
         if strategy == "auto":
             strategy = self._pick_strategy(tree_id)
         mode = self._tree_mode(tree_id)
+        ranker = (
+            make_ranker(
+                kwargs.get("ranker"),
+                embedding_provider=kwargs.get("embedding_provider", "openai"),
+                embedding_model=kwargs.get("embedding_model"),
+                embedding_api_key=kwargs.get("embedding_api_key"),
+            )
+            if strategy == "block"
+            else None
+        )
         return build_strategy_retriever(
             self.storage,
             llm,
@@ -146,7 +165,7 @@ def _make_retriever(self, tree_id, llm, strategy, **kwargs) -> BaseRetriever:
             cache_enabled=kwargs.get("cache_enabled", True),
             max_parallel_blocks=kwargs.get("max_parallel_blocks"),
             mode=mode,
-            ranker=kwargs.get("ranker"),
+            ranker=ranker,
         )
 
     def _tree_mode(self, tree_id: str) -> str:

contextdb/retriever/algorithm/block_retriever_filesystem.py

@@ -340,7 +340,11 @@ def _order_fs_node_id_groups_for_query(
             for node_id in node_ids
             if node_id in node_by_id
         ]
-        if not has_path_evidence(candidates, query):
+        should_rank = getattr(ranker, "should_rank", None)
+        if callable(should_rank):
+            if not should_rank(query, candidates, context={"mode": "filesystem", "tree_id": tree_id}):
+                return node_ids
+        elif not has_path_evidence(candidates, query):
             return node_ids
         scores = ranker.rank(
             query,

contextdb/retriever/algorithm/embeddings.py

@@ -0,0 +1,69 @@
+"""Embedding adapters used by retrieval rankers."""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class EmbeddingClient(Protocol):
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        """Return one embedding vector per input text."""
+        ...
+
+
+def cosine_similarity(a: list[float], b: list[float]) -> float:
+    if not a or not b:
+        return 0.0
+    n = min(len(a), len(b))
+    dot = sum(a[i] * b[i] for i in range(n))
+    norm_a = math.sqrt(sum(v * v for v in a[:n]))
+    norm_b = math.sqrt(sum(v * v for v in b[:n]))
+    if norm_a <= 0.0 or norm_b <= 0.0:
+        return 0.0
+    return dot / (norm_a * norm_b)
+
+
+def _resolve_embedding_model(provider: str, model: str | None) -> str:
+    model = model or "text-embedding-3-small"
+    if "/" in model:
+        return model
+    return f"{provider}/{model}"
+
+
+class LiteLLMEmbeddingClient:
+    """Thin adapter around litellm.embedding."""
+
+    def __init__(
+        self,
+        *,
+        provider: str = "openai",
+        model: str | None = None,
+        api_key: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        import litellm
+
+        self._litellm = litellm
+        self.model = _resolve_embedding_model(provider, model)
+        self.api_key = api_key
+        self.kwargs = kwargs
+        litellm.suppress_debug_info = True
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+        kwargs = {"model": self.model, "input": texts, **self.kwargs}
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        response = self._litellm.embedding(**kwargs)
+        rows = list(getattr(response, "data", None) or response["data"])
+        rows.sort(key=lambda row: _row_value(row, "index", 0))
+        return [list(_row_value(row, "embedding", [])) for row in rows]
+
+
+def _row_value(row: Any, key: str, default: Any) -> Any:
+    if isinstance(row, dict):
+        return row.get(key, default)
+    return getattr(row, key, default)