Add repo key

Author: faysou

src/cocoindex_code/cli.py

@@ -104,6 +104,7 @@ def __call__(
         query: str,
         languages: list[str] | None = None,
         paths: list[str] | None = None,
+        repo_keys: list[str] | None = None,
         limit: int = 5,
         offset: int = 0,
         on_waiting: Callable[[], None] | None = None,
@@ -194,6 +195,7 @@ def search_response_json_payload(response: SearchResponse) -> dict[str, object]:
         "results": [
             {
                 "file_path": r.file_path,
+                "repo_key": r.repo_key,
                 "language": r.language,
                 "content": r.content,
                 "start_line": r.start_line,
@@ -317,6 +319,7 @@ def handle_bridge_jsonrpc_request(
             query=_required_str(params, "query"),
             languages=_optional_str_list(params, "languages"),
             paths=_optional_str_list(params, "paths"),
+            repo_keys=_optional_str_list(params, "repo_keys"),
             limit=_positive_int_param(params, "limit", 10),
             offset=_non_negative_int_param(params, "offset", 0),
         )
@@ -405,6 +408,7 @@ def _search_with_wait_spinner(
     query: str,
     languages: list[str] | None = None,
     paths: list[str] | None = None,
+    repo_keys: list[str] | None = None,
     limit: int = 10,
     offset: int = 0,
 ) -> SearchResponse:
@@ -430,6 +434,7 @@ def _on_waiting() -> None:
             query=query,
             languages=languages,
             paths=paths,
+            repo_keys=repo_keys,
             limit=limit,
             offset=offset,
             on_waiting=_on_waiting,
@@ -722,6 +727,7 @@ def search(
     query: list[str] = _typer.Argument(..., help="Search query"),
     lang: list[str] = _typer.Option([], "--lang", help="Filter by language"),
     path: str | None = _typer.Option(None, "--path", help="Filter by file path glob"),
+    repo_key: list[str] = _typer.Option([], "--repo-key", help="Filter by indexed repo key"),
     offset: int = _typer.Option(0, "--offset", help="Number of results to skip"),
     limit: int = _typer.Option(10, "--limit", help="Maximum results to return"),
     refresh: bool = _typer.Option(False, "--refresh", help="Refresh index before searching"),
@@ -748,6 +754,7 @@ def search(
         query=query_str,
         languages=lang or None,
         paths=paths,
+        repo_keys=repo_key or None,
         limit=limit,
         offset=offset,
     )

src/cocoindex_code/client.py

@@ -278,6 +278,7 @@ def search(
     query: str,
     languages: list[str] | None = None,
     paths: list[str] | None = None,
+    repo_keys: list[str] | None = None,
     limit: int = 5,
     offset: int = 0,
     on_waiting: Callable[[], None] | None = None,
@@ -298,6 +299,7 @@ def search(
                     query=query,
                     languages=languages,
                     paths=paths,
+                    repo_keys=repo_keys,
                     limit=limit,
                     offset=offset,
                 )

src/cocoindex_code/daemon.py

@@ -275,6 +275,7 @@ async def _search_with_wait(
             query=req.query,
             languages=req.languages,
             paths=req.paths,
+            repo_keys=req.repo_keys,
             limit=req.limit,
             offset=req.offset,
         )
@@ -488,6 +489,7 @@ async def _dispatch(
                 query=req.query,
                 languages=req.languages,
                 paths=req.paths,
+                repo_keys=req.repo_keys,
                 limit=req.limit,
                 offset=req.offset,
             )

src/cocoindex_code/indexer.py

@@ -33,6 +33,22 @@
 splitter = RecursiveSplitter()
 
 
+def repo_key_for_path(file_path: PurePath, project_root: Path) -> str:
+    """Return the relative Git repo root for fast scoped search."""
+    directory = file_path.parent
+    while True:
+        if (project_root / directory / ".git").exists():
+            repo_key = directory.as_posix()
+            return repo_key if repo_key != "." else "."
+
+        if directory in (PurePath("."), PurePath("")):
+            break
+        directory = directory.parent
+
+    parts = file_path.parts
+    return parts[0] if len(parts) > 1 else "."
+
+
 def _normalize_gitignore_lines(lines: Iterable[str], directory: PurePath) -> list[str]:
     """Normalize .gitignore lines to root-relative gitignore patterns."""
     if directory in (PurePath("."), PurePath("")):
@@ -151,8 +167,9 @@ async def process_file(
     if not content.strip():
         return
 
-    suffix = file.file_path.path.suffix
     project_root = coco.use_context(CODEBASE_DIR)
+    suffix = file.file_path.path.suffix
+    repo_key = repo_key_for_path(file.file_path.path, project_root)
     ps = load_project_settings(project_root)
     ext_lang_map = {f".{lo.ext}": lo.lang for lo in ps.language_overrides}
     language = (
@@ -183,6 +200,7 @@ async def process(chunk: Chunk) -> None:
             row=CodeChunk(
                 id=await id_gen.next_id(chunk.text),
                 file_path=file.file_path.path.as_posix(),
+                repo_key=repo_key,
                 language=language,
                 content=chunk.text,
                 start_line=chunk.start.line,
@@ -209,7 +227,7 @@ async def indexer_main() -> None:
             primary_key=["id"],
         ),
         virtual_table_def=Vec0TableDef(
-            partition_key_columns=["language"],
+            partition_key_columns=["repo_key", "language"],
             auxiliary_columns=["file_path", "content", "start_line", "end_line"],
         ),
     )

src/cocoindex_code/project.py

@@ -179,6 +179,7 @@ async def search(
         query: str,
         languages: list[str] | None = None,
         paths: list[str] | None = None,
+        repo_keys: list[str] | None = None,
         limit: int = 5,
         offset: int = 0,
     ) -> list[SearchResult]:
@@ -192,10 +193,12 @@ async def search(
             offset=offset,
             languages=languages,
             paths=paths,
+            repo_keys=repo_keys,
         )
         return [
             SearchResult(
                 file_path=r.file_path,
+                repo_key=r.repo_key,
                 language=r.language,
                 content=r.content,
                 start_line=r.start_line,

src/cocoindex_code/protocol.py

@@ -22,6 +22,7 @@ class SearchRequest(_msgspec.Struct, tag="search"):
     query: str
     languages: list[str] | None = None
     paths: list[str] | None = None
+    repo_keys: list[str] | None = None
     limit: int = 5
     offset: int = 0
 
@@ -111,6 +112,7 @@ class SearchResult(_msgspec.Struct):
     start_line: int
     end_line: int
     score: float
+    repo_key: str | None = None
 
 
 class SearchResponse(_msgspec.Struct, tag="search"):

src/cocoindex_code/query.py

@@ -21,26 +21,29 @@ def _knn_query(
     embedding_bytes: bytes,
     k: int,
     language: str | None = None,
+    repo_key: str | None = None,
+    has_repo_key: bool = False,
 ) -> list[tuple[Any, ...]]:
     """Run a vec0 KNN query, optionally constrained to a language partition."""
+    conditions = ["embedding MATCH ?", "k = ?"]
+    params: list[Any] = [embedding_bytes, k]
+    if repo_key is not None:
+        conditions.append("repo_key = ?")
+        params.append(repo_key)
     if language is not None:
-        return conn.execute(
-            """
-            SELECT file_path, language, content, start_line, end_line, distance
-            FROM code_chunks_vec
-            WHERE embedding MATCH ? AND k = ? AND language = ?
-            ORDER BY distance
-            """,
-            (embedding_bytes, k, language),
-        ).fetchall()
+        conditions.append("language = ?")
+        params.append(language)
+
+    repo_key_select = "repo_key" if has_repo_key else "NULL"
     return conn.execute(
-        """
-        SELECT file_path, language, content, start_line, end_line, distance
+        f"""
+        SELECT file_path, {repo_key_select} as repo_key,
+               language, content, start_line, end_line, distance
         FROM code_chunks_vec
-        WHERE embedding MATCH ? AND k = ?
+        WHERE {" AND ".join(conditions)}
         ORDER BY distance
         """,
-        (embedding_bytes, k),
+        params,
     ).fetchall()
 
 
@@ -51,27 +54,42 @@ def _full_scan_query(
     offset: int,
     languages: list[str] | None = None,
     paths: list[str] | None = None,
+    repo_keys: list[str] | None = None,
 ) -> list[tuple[Any, ...]]:
     """Full scan with SQL-level distance computation and filtering."""
     conditions: list[str] = []
     params: list[Any] = [embedding_bytes]
 
+    has_repo_key = _table_has_column(conn, "code_chunks_vec", "repo_key")
+
     if languages:
         placeholders = ",".join("?" for _ in languages)
         conditions.append(f"language IN ({placeholders})")
         params.extend(languages)
 
+    if repo_keys:
+        if has_repo_key:
+            placeholders = ",".join("?" for _ in repo_keys)
+            conditions.append(f"repo_key IN ({placeholders})")
+            params.extend(repo_keys)
+        else:
+            repo_key_paths = [
+                f"{repo_key.rstrip('/')}/*" for repo_key in repo_keys if repo_key != "."
+            ]
+            paths = [*(paths or []), *repo_key_paths] or paths
+
     if paths:
         path_clauses = " OR ".join("file_path GLOB ?" for _ in paths)
         conditions.append(f"({path_clauses})")
         params.extend(paths)
 
+    repo_key_select = "repo_key" if has_repo_key else "NULL as repo_key"
     where = f"WHERE {' AND '.join(conditions)}" if conditions else ""
     params.extend([limit, offset])
 
     return conn.execute(
         f"""
-        SELECT file_path, language, content, start_line, end_line,
+        SELECT file_path, {repo_key_select}, language, content, start_line, end_line,
                vec_distance_L2(embedding, ?) as distance
         FROM code_chunks_vec
         {where}
@@ -82,6 +100,22 @@ def _full_scan_query(
     ).fetchall()
 
 
+def _table_has_column(conn: sqlite3.Connection, table_name: str, column_name: str) -> bool:
+    return any(row[1] == column_name for row in conn.execute(f"PRAGMA table_info({table_name})"))
+
+
+def _repo_key_candidates(repo_keys: list[str] | None) -> list[str | None]:
+    if repo_keys:
+        return list(repo_keys)
+    return [None]
+
+
+def _language_candidates(languages: list[str] | None) -> list[str | None]:
+    if languages:
+        return list(languages)
+    return [None]
+
+
 async def query_codebase(
     query: str,
     target_sqlite_db_path: Path,
@@ -90,13 +124,16 @@ async def query_codebase(
     offset: int = 0,
     languages: list[str] | None = None,
     paths: list[str] | None = None,
+    repo_keys: list[str] | None = None,
 ) -> list[QueryResult]:
     """
     Perform vector similarity search using vec0 KNN index.
 
     Uses sqlite-vec's vec0 virtual table for indexed nearest-neighbor search.
     Language filtering uses vec0 partition keys for exact index-level filtering.
     Path filtering triggers a full scan with distance computation.
+    Repo-key filtering uses the vec0 partition key when available, and
+    falls back to equivalent path filters for older indexes.
     """
     if not target_sqlite_db_path.exists():
         raise RuntimeError(
@@ -114,34 +151,46 @@ async def query_codebase(
     embedding_bytes = query_embedding.astype("float32").tobytes()
 
     with db.readonly() as conn:
+        has_repo_key = _table_has_column(conn, "code_chunks_vec", "repo_key")
         if paths:
-            rows = _full_scan_query(conn, embedding_bytes, limit, offset, languages, paths)
-        elif not languages or len(languages) == 1:
+            rows = _full_scan_query(
+                conn, embedding_bytes, limit, offset, languages, paths, repo_keys
+            )
+        elif repo_keys and not has_repo_key:
+            rows = _full_scan_query(
+                conn, embedding_bytes, limit, offset, languages, None, repo_keys
+            )
+        elif (not languages or len(languages) == 1) and (not repo_keys or len(repo_keys) == 1):
             lang = languages[0] if languages else None
-            rows = _knn_query(conn, embedding_bytes, limit + offset, lang)
+            repo_key = repo_keys[0] if repo_keys else None
+            rows = _knn_query(conn, embedding_bytes, limit + offset, lang, repo_key, has_repo_key)
         else:
             fetch_k = limit + offset
             rows = heapq.nsmallest(
                 fetch_k,
                 (
                     row
-                    for lang in languages
-                    for row in _knn_query(conn, embedding_bytes, fetch_k, lang)
+                    for repo_key in _repo_key_candidates(repo_keys)
+                    for lang in _language_candidates(languages)
+                    for row in _knn_query(
+                        conn, embedding_bytes, fetch_k, lang, repo_key, has_repo_key
+                    )
                 ),
-                key=lambda r: r[5],
+                key=lambda r: r[6],
             )
 
-    if not paths:
+    if not paths and not (repo_keys and not has_repo_key):
         rows = rows[offset:]
 
     return [
         QueryResult(
             file_path=file_path,
+            repo_key=repo_key,
             language=language,
             content=content,
             start_line=start_line,
             end_line=end_line,
             score=_l2_to_score(distance),
         )
-        for file_path, language, content, start_line, end_line, distance in rows
+        for file_path, repo_key, language, content, start_line, end_line, distance in rows
     ]

src/cocoindex_code/schema.py

@@ -10,6 +10,7 @@ class CodeChunk:
 
     id: int
     file_path: str
+    repo_key: str
     language: str
     content: str
     start_line: int
@@ -22,6 +23,7 @@ class QueryResult:
     """Result from a vector similarity query."""
 
     file_path: str
+    repo_key: str | None
     language: str
     content: str
     start_line: int