cocoindex-code 0.1.7__tar.gz → 0.1.9__tar.gz

{cocoindex_code-0.1.7 → cocoindex_code-0.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex-code
-Version: 0.1.7
+Version: 0.1.9
 Summary: MCP server for indexing and querying codebases using CocoIndex
 Project-URL: Homepage, https://github.com/cocoindex-io/cocoindex-code
 Project-URL: Repository, https://github.com/cocoindex-io/cocoindex-code
@@ -17,7 +17,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.11
-Requires-Dist: cocoindex[litellm]==1.0.0a22
+Requires-Dist: cocoindex[litellm]==1.0.0a24
 Requires-Dist: einops>=0.8.2
 Requires-Dist: mcp>=1.0.0
 Requires-Dist: numpy>=1.24.0
@@ -40,6 +40,10 @@ Description-Content-Type: text/markdown
 
 <h1 align="center">light weight MCP for code that just works </h1>
 
+![effect](https://github.com/user-attachments/assets/cb3a4cae-0e1f-49c4-890b-7bb93317ab60)
+
+
+
 
 A super light-weight, effective embedded MCP **(AST-based)** that understand and searches your codebase that just works! Using [CocoIndex](https://github.com/cocoindex-io/cocoindex) - an Rust-based ultra performant data transformation engine. No blackbox. Works for Claude, Codex, Cursor - any coding agent.
 
@@ -116,6 +120,20 @@ Or use opencode.json:
 
 Optionally, you can run `cocoindex-code index` to create or update the index. Without running it, the MCP server will automatically build and keep the index up-to-date in the background.
 
+## When Is the MCP Triggered?
+
+Once configured, your coding agent (Claude Code, Codex, Cursor, etc.) automatically decides when semantic code search is helpful — especially for finding code by description, exploring unfamiliar codebases, fuzzy/conceptual matches, or locating implementations without knowing exact names.
+
+You can also nudge the agent explicitly, e.g. *"Use the cocoindex-code MCP to find how user sessions are managed."* For persistent instructions, add guidance to your project's `AGENTS.md` or `CLAUDE.md`:
+
+```
+Use the cocoindex-code MCP server for semantic code search when:
+- Searching for code by meaning or description rather than exact text
+- Exploring unfamiliar parts of the codebase
+- Looking for implementations without knowing exact names
+- Finding similar code patterns or related functionality
+```
+
 ## Features
 - **Semantic Code Search**: Find relevant code using natural language queries when grep doesn't work well, and save tokens immediately.
 - **Ultra Performant to code changes**:⚡ Built on top of ultra performant [Rust indexing engine](https://github.com/cocoindex-io/cocoindex/edit/main/README.md). Only re-indexes changed files for fast updates.
@@ -131,6 +149,7 @@ Optionally, you can run `cocoindex-code index` to create or update the index. Wi
 | `COCOINDEX_CODE_ROOT_PATH` | Root path of the codebase | Auto-discovered (see below) |
 | `COCOINDEX_CODE_EMBEDDING_MODEL` | Embedding model (see below) | `sbert/sentence-transformers/all-MiniLM-L6-v2` |
 | `COCOINDEX_CODE_BATCH_SIZE` | Max batch size for local embedding model | `16` |
+| `COCOINDEX_CODE_EXTRA_EXTENSIONS` | Additional file extensions to index (comma-separated, e.g. `"inc:php,yaml,toml"` — use `ext:lang` to override language detection) | _(none)_ |
 
 
 ### Root Path Discovery

{cocoindex_code-0.1.7 → cocoindex_code-0.1.9}/README.md

@@ -5,6 +5,10 @@
 
 <h1 align="center">light weight MCP for code that just works </h1>
 
+![effect](https://github.com/user-attachments/assets/cb3a4cae-0e1f-49c4-890b-7bb93317ab60)
+
+
+
 
 A super light-weight, effective embedded MCP **(AST-based)** that understand and searches your codebase that just works! Using [CocoIndex](https://github.com/cocoindex-io/cocoindex) - an Rust-based ultra performant data transformation engine. No blackbox. Works for Claude, Codex, Cursor - any coding agent.
 
@@ -81,6 +85,20 @@ Or use opencode.json:
 
 Optionally, you can run `cocoindex-code index` to create or update the index. Without running it, the MCP server will automatically build and keep the index up-to-date in the background.
 
+## When Is the MCP Triggered?
+
+Once configured, your coding agent (Claude Code, Codex, Cursor, etc.) automatically decides when semantic code search is helpful — especially for finding code by description, exploring unfamiliar codebases, fuzzy/conceptual matches, or locating implementations without knowing exact names.
+
+You can also nudge the agent explicitly, e.g. *"Use the cocoindex-code MCP to find how user sessions are managed."* For persistent instructions, add guidance to your project's `AGENTS.md` or `CLAUDE.md`:
+
+```
+Use the cocoindex-code MCP server for semantic code search when:
+- Searching for code by meaning or description rather than exact text
+- Exploring unfamiliar parts of the codebase
+- Looking for implementations without knowing exact names
+- Finding similar code patterns or related functionality
+```
+
 ## Features
 - **Semantic Code Search**: Find relevant code using natural language queries when grep doesn't work well, and save tokens immediately.
 - **Ultra Performant to code changes**:⚡ Built on top of ultra performant [Rust indexing engine](https://github.com/cocoindex-io/cocoindex/edit/main/README.md). Only re-indexes changed files for fast updates.
@@ -96,6 +114,7 @@ Optionally, you can run `cocoindex-code index` to create or update the index. Wi
 | `COCOINDEX_CODE_ROOT_PATH` | Root path of the codebase | Auto-discovered (see below) |
 | `COCOINDEX_CODE_EMBEDDING_MODEL` | Embedding model (see below) | `sbert/sentence-transformers/all-MiniLM-L6-v2` |
 | `COCOINDEX_CODE_BATCH_SIZE` | Max batch size for local embedding model | `16` |
+| `COCOINDEX_CODE_EXTRA_EXTENSIONS` | Additional file extensions to index (comma-separated, e.g. `"inc:php,yaml,toml"` — use `ext:lang` to override language detection) | _(none)_ |
 
 
 ### Root Path Discovery

{cocoindex_code-0.1.7 → cocoindex_code-0.1.9}/pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
 
 dependencies = [
     "mcp>=1.0.0",
-    "cocoindex[litellm]==1.0.0a22",
+    "cocoindex[litellm]==1.0.0a24",
     "sentence-transformers>=2.2.0",
     "sqlite-vec>=0.1.0",
     "pydantic>=2.0.0",
@@ -82,10 +82,6 @@ python_version = "3.11"
 strict = true
 ignore_missing_imports = true
 
-[[tool.mypy.overrides]]
-module = "cocoindex_code.embedder"
-warn_unused_ignores = false
-
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 python_files = ["test_*.py"]

{cocoindex_code-0.1.7 → cocoindex_code-0.1.9}/src/cocoindex_code/config.py

@@ -6,7 +6,6 @@ import os
 from dataclasses import dataclass
 from pathlib import Path
 
-_SBERT_PREFIX = "sbert/"
 _DEFAULT_MODEL = "sbert/sentence-transformers/all-MiniLM-L6-v2"
 
 
@@ -65,7 +64,7 @@ class Config:
     index_dir: Path
     device: str
     trust_remote_code: bool
-    batch_size: int
+    extra_extensions: dict[str, str | None]
 
     @classmethod
     def from_env(cls) -> Config:
@@ -100,18 +99,18 @@ class Config:
             "yes",
         )
 
-        # Batch size for local embedding model
-        _raw_batch_size = os.environ.get("COCOINDEX_CODE_BATCH_SIZE", "16")
-        try:
-            batch_size = int(_raw_batch_size)
-        except ValueError:
-            raise ValueError(
-                f"COCOINDEX_CODE_BATCH_SIZE must be a positive integer, got: {_raw_batch_size!r}"
-            ) from None
-        if batch_size <= 0:
-            raise ValueError(
-                f"COCOINDEX_CODE_BATCH_SIZE must be a positive integer, got: {batch_size}"
-            )
+        # Extra file extensions (format: "inc:php,yaml,toml" — optional lang after colon)
+        raw_extra = os.environ.get("COCOINDEX_CODE_EXTRA_EXTENSIONS", "")
+        extra_extensions: dict[str, str | None] = {}
+        for token in raw_extra.split(","):
+            token = token.strip()
+            if not token:
+                continue
+            if ":" in token:
+                ext, lang = token.split(":", 1)
+                extra_extensions[f".{ext.strip()}"] = lang.strip() or None
+            else:
+                extra_extensions[f".{token}"] = None
 
         return cls(
             codebase_root_path=root,
@@ -119,7 +118,7 @@ class Config:
             index_dir=index_dir,
             device=device,
             trust_remote_code=trust_remote_code,
-            batch_size=batch_size,
+            extra_extensions=extra_extensions,
         )
 
     @property

{cocoindex_code-0.1.7 → cocoindex_code-0.1.9}/src/cocoindex_code/indexer.py

@@ -2,16 +2,17 @@
 
 import cocoindex as coco
 from cocoindex.connectors import localfs, sqlite
+from cocoindex.connectors.sqlite import Vec0TableDef
 from cocoindex.ops.text import RecursiveSplitter, detect_code_language
 from cocoindex.resources.chunk import Chunk
 from cocoindex.resources.file import PatternFilePathMatcher
 from cocoindex.resources.id import IdGenerator
 
 from .config import config
-from .shared import SQLITE_DB, CodeChunk, embedder
+from .shared import CODEBASE_DIR, SQLITE_DB, CodeChunk, embedder
 
 # File patterns for supported languages
-INCLUDED_PATTERNS = [
+DEFAULT_INCLUDED_PATTERNS = [
     "**/*.py",  # Python
     "**/*.pyi",  # Python stubs
     "**/*.js",  # JavaScript
@@ -43,6 +44,13 @@ INCLUDED_PATTERNS = [
     "**/*.php",  # PHP
 ]
 
+INCLUDED_PATTERNS = DEFAULT_INCLUDED_PATTERNS + [f"**/*{ext}" for ext in config.extra_extensions]
+
+# Language overrides from extra_extensions (e.g. ".inc" -> "php")
+LANGUAGE_OVERRIDES: dict[str, str] = {
+    ext: lang for ext, lang in config.extra_extensions.items() if lang is not None
+}
+
 EXCLUDED_PATTERNS = [
     "**/.*",  # Hidden directories
     "**/__pycache__",  # Python cache
@@ -56,9 +64,9 @@ EXCLUDED_PATTERNS = [
 ]
 
 # Chunking configuration
-CHUNK_SIZE = 4000
-MIN_CHUNK_SIZE = 500
-CHUNK_OVERLAP = 400
+CHUNK_SIZE = 2000
+MIN_CHUNK_SIZE = 300
+CHUNK_OVERLAP = 200
 
 # Chunking splitter (stateless, can be module-level)
 splitter = RecursiveSplitter()
@@ -66,7 +74,7 @@ splitter = RecursiveSplitter()
 
 @coco.fn(memo=True)
 async def process_file(
-    file: localfs.AsyncFile,
+    file: localfs.File,
     table: sqlite.TableTarget[CodeChunk],
 ) -> None:
     """Process a single file: chunk, embed, and store."""
@@ -81,7 +89,12 @@ async def process_file(
         return
 
     # Get relative path and detect language
-    language = detect_code_language(filename=file.file_path.path.name) or "text"
+    suffix = file.file_path.path.suffix
+    language = (
+        LANGUAGE_OVERRIDES.get(suffix)
+        or detect_code_language(filename=file.file_path.path.name)
+        or "text"
+    )
 
     # Split into chunks
     chunks = splitter.split(
@@ -119,16 +132,20 @@ async def app_main() -> None:
 
     # Declare the table target for storing embeddings
     table = await db.mount_table_target(
-        table_name="code_chunks",
+        table_name="code_chunks_vec",
         table_schema=await sqlite.TableSchema.from_class(
             CodeChunk,
             primary_key=["id"],
         ),
+        virtual_table_def=Vec0TableDef(
+            partition_key_columns=["language"],
+            auxiliary_columns=["file_path", "content", "start_line", "end_line"],
+        ),
     )
 
     # Walk source directory
     files = localfs.walk_dir(
-        config.codebase_root_path,
+        coco.use_context(CODEBASE_DIR),
         recursive=True,
         path_matcher=PatternFilePathMatcher(
             included_patterns=INCLUDED_PATTERNS,

cocoindex_code-0.1.9/src/cocoindex_code/query.py

@@ -0,0 +1,148 @@
+"""Query implementation for codebase search."""
+
+import heapq
+import sqlite3
+from typing import Any
+
+import cocoindex as coco
+
+from .config import config
+from .schema import QueryResult
+from .shared import SQLITE_DB, embedder, query_prompt_name
+
+
+def _l2_to_score(distance: float) -> float:
+    """Convert L2 distance to cosine similarity (exact for unit vectors)."""
+    return 1.0 - distance * distance / 2.0
+
+
+def _knn_query(
+    conn: sqlite3.Connection,
+    embedding_bytes: bytes,
+    k: int,
+    language: str | None = None,
+) -> list[tuple[Any, ...]]:
+    """Run a vec0 KNN query, optionally constrained to a language partition."""
+    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()
+    return conn.execute(
+        """
+        SELECT file_path, language, content, start_line, end_line, distance
+        FROM code_chunks_vec
+        WHERE embedding MATCH ? AND k = ?
+        ORDER BY distance
+        """,
+        (embedding_bytes, k),
+    ).fetchall()
+
+
+def _full_scan_query(
+    conn: sqlite3.Connection,
+    embedding_bytes: bytes,
+    limit: int,
+    offset: int,
+    languages: list[str] | None = None,
+    paths: list[str] | None = None,
+) -> list[tuple[Any, ...]]:
+    """Full scan with SQL-level distance computation and filtering."""
+    conditions: list[str] = []
+    params: list[Any] = [embedding_bytes]
+
+    if languages:
+        placeholders = ",".join("?" for _ in languages)
+        conditions.append(f"language IN ({placeholders})")
+        params.extend(languages)
+
+    if paths:
+        path_clauses = " OR ".join("file_path GLOB ?" for _ in paths)
+        conditions.append(f"({path_clauses})")
+        params.extend(paths)
+
+    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,
+               vec_distance_L2(embedding, ?) as distance
+        FROM code_chunks_vec
+        {where}
+        ORDER BY distance
+        LIMIT ? OFFSET ?
+        """,
+        params,
+    ).fetchall()
+
+
+async def query_codebase(
+    query: str,
+    limit: int = 10,
+    offset: int = 0,
+    languages: list[str] | None = None,
+    paths: 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.
+    """
+    if not config.target_sqlite_db_path.exists():
+        raise RuntimeError(
+            f"Index database not found at {config.target_sqlite_db_path}. "
+            "Please run a query with refresh_index=True first."
+        )
+
+    coco_env = await coco.default_env()
+    db = coco_env.get_context(SQLITE_DB)
+
+    # Generate query embedding.
+    query_embedding = await embedder.embed(query, True, query_prompt_name)
+
+    embedding_bytes = query_embedding.astype("float32").tobytes()
+
+    with db.value.readonly() as conn:
+        if paths:
+            # Path filter → full scan (vec0 can't filter on auxiliary columns).
+            # LIMIT/OFFSET handled in SQL.
+            rows = _full_scan_query(conn, embedding_bytes, limit, offset, languages, paths)
+        elif not languages or len(languages) == 1:
+            # Single language or no filter: one KNN query.
+            lang = languages[0] if languages else None
+            rows = _knn_query(conn, embedding_bytes, limit + offset, lang)
+        else:
+            # Multiple languages: separate KNN per partition, merge by distance.
+            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)
+                ),
+                key=lambda r: r[5],  # distance column
+            )
+
+    if not paths:
+        rows = rows[offset:]
+
+    return [
+        QueryResult(
+            file_path=file_path,
+            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
+    ]

{cocoindex_code-0.1.7 → cocoindex_code-0.1.9}/src/cocoindex_code/server.py

@@ -2,14 +2,15 @@
 
 import argparse
 import asyncio
-import sqlite3
 
+import cocoindex as coco
 from mcp.server.fastmcp import FastMCP
 from pydantic import BaseModel, Field
 
 from .config import config
 from .indexer import app as indexer_app
 from .query import query_codebase
+from .shared import SQLITE_DB
 
 # Initialize MCP server
 mcp = FastMCP(
@@ -76,6 +77,8 @@ class SearchResultModel(BaseModel):
         " or code snippets."
         " Returns matching code chunks with file paths,"
         " line numbers, and relevance scores."
+        " Start with a small limit (e.g., 5);"
+        " if most results look relevant, use offset to paginate for more."
     ),
 )
 async def search(
@@ -89,7 +92,7 @@ async def search(
         )
     ),
     limit: int = Field(
-        default=10,
+        default=5,
         ge=1,
         le=100,
         description="Maximum number of results to return (1-100)",
@@ -107,6 +110,17 @@ async def search(
             " when the codebase hasn't changed."
         ),
     ),
+    languages: list[str] | None = Field(
+        default=None,
+        description=("Filter by programming language(s). Example: ['python', 'typescript']"),
+    ),
+    paths: list[str] | None = Field(
+        default=None,
+        description=(
+            "Filter by file path pattern(s) using GLOB wildcards (* and ?)."
+            " Example: ['src/utils/*', '*.py']"
+        ),
+    ),
 ) -> SearchResultModel:
     """Query the codebase index."""
     try:
@@ -114,7 +128,13 @@ async def search(
         if refresh_index:
             await _refresh_index()
 
-        results = await query_codebase(query=query, limit=limit, offset=offset)
+        results = await query_codebase(
+            query=query,
+            limit=limit,
+            offset=offset,
+            languages=languages,
+            paths=paths,
+        )
 
         return SearchResultModel(
             success=True,
@@ -155,35 +175,36 @@ async def _async_serve() -> None:
 async def _async_index() -> None:
     """Async entry point for the index command."""
     await indexer_app.update(report_to_stdout=True)
-    _print_index_stats()
+    await _print_index_stats()
 
 
-def _print_index_stats() -> None:
+async def _print_index_stats() -> None:
     """Print index statistics from the database."""
     db_path = config.target_sqlite_db_path
     if not db_path.exists():
         print("No index database found.")
         return
 
-    conn = sqlite3.connect(str(db_path))
-    try:
-        total_chunks = conn.execute("SELECT COUNT(*) FROM code_chunks").fetchone()[0]
-        total_files = conn.execute("SELECT COUNT(DISTINCT file_path) FROM code_chunks").fetchone()[
-            0
-        ]
+    coco_env = await coco.default_env()
+    db = coco_env.get_context(SQLITE_DB)
+
+    with db.value.readonly() as conn:
+        total_chunks = conn.execute("SELECT COUNT(*) FROM code_chunks_vec").fetchone()[0]
+        total_files = conn.execute(
+            "SELECT COUNT(DISTINCT file_path) FROM code_chunks_vec"
+        ).fetchone()[0]
         langs = conn.execute(
-            "SELECT language, COUNT(*) as cnt FROM code_chunks GROUP BY language ORDER BY cnt DESC"
+            "SELECT language, COUNT(*) as cnt FROM code_chunks_vec"
+            " GROUP BY language ORDER BY cnt DESC"
         ).fetchall()
 
-        print("\nIndex stats:")
-        print(f"  Chunks: {total_chunks}")
-        print(f"  Files:  {total_files}")
-        if langs:
-            print("  Languages:")
-            for lang, count in langs:
-                print(f"    {lang}: {count} chunks")
-    finally:
-        conn.close()
+    print("\nIndex stats:")
+    print(f"  Chunks: {total_chunks}")
+    print(f"  Files:  {total_files}")
+    if langs:
+        print("  Languages:")
+        for lang, count in langs:
+            print(f"    {lang}: {count} chunks")
 
 
 def main() -> None:

{cocoindex_code-0.1.7 → cocoindex_code-0.1.9}/src/cocoindex_code/shared.py

@@ -9,12 +9,12 @@ from typing import TYPE_CHECKING, Annotated
 
 import cocoindex as coco
 from cocoindex.connectors import sqlite
+from cocoindex.connectors.localfs import FilePath, register_base_dir
 from numpy.typing import NDArray
 
 if TYPE_CHECKING:
     from cocoindex.ops.litellm import LiteLLMEmbedder
-
-    from .embedder import LocalEmbedder
+    from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder
 
 from .config import config
 
@@ -23,22 +23,21 @@ logger = logging.getLogger(__name__)
 SBERT_PREFIX = "sbert/"
 
 # Initialize embedder at module level based on model prefix
-embedder: LocalEmbedder | LiteLLMEmbedder
+embedder: SentenceTransformerEmbedder | LiteLLMEmbedder
 if config.embedding_model.startswith(SBERT_PREFIX):
-    from .embedder import LocalEmbedder
+    from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder
 
     _model_name = config.embedding_model[len(SBERT_PREFIX) :]
     # Models that define a "query" prompt for asymmetric retrieval.
     _QUERY_PROMPT_MODELS = {"nomic-ai/nomic-embed-code", "nomic-ai/CodeRankEmbed"}
-    _query_prompt_name: str | None = "query" if _model_name in _QUERY_PROMPT_MODELS else None
+    query_prompt_name: str | None = "query" if _model_name in _QUERY_PROMPT_MODELS else None
     # Models whose custom remote code is known-compatible with transformers 5.x.
     _KNOWN_REMOTE_CODE_MODELS = {"nomic-ai/CodeRankEmbed"}
     _trust = config.trust_remote_code or _model_name in _KNOWN_REMOTE_CODE_MODELS
-    embedder = LocalEmbedder(
+    embedder = SentenceTransformerEmbedder(
         _model_name,
         device=config.device,
         trust_remote_code=_trust,
-        query_prompt_name=_query_prompt_name,
     )
     logger.info(
         "Embedding model: %s | device: %s | trust_remote_code: %s",
@@ -50,10 +49,13 @@ else:
     from cocoindex.ops.litellm import LiteLLMEmbedder
 
     embedder = LiteLLMEmbedder(config.embedding_model)
+    query_prompt_name = None
     logger.info("Embedding model (LiteLLM): %s", config.embedding_model)
 
 # Context key for SQLite database (connection managed in lifespan)
 SQLITE_DB = coco.ContextKey[sqlite.SqliteDatabase]("sqlite_db")
+# Context key for codebase root directory (provided in lifespan)
+CODEBASE_DIR = coco.ContextKey[FilePath]("codebase_dir")
 
 
 @coco.lifespan
@@ -65,6 +67,9 @@ def coco_lifespan(builder: coco.EnvironmentBuilder) -> Iterator[None]:
     # Set CocoIndex state database path
     builder.settings.db_path = config.cocoindex_db_path
 
+    # Provide codebase root directory to environment
+    builder.provide(CODEBASE_DIR, register_base_dir("codebase", config.codebase_root_path))
+
     # Connect to SQLite with vector extension
     conn = sqlite.connect(str(config.target_sqlite_db_path), load_vec="auto")
     builder.provide(SQLITE_DB, sqlite.register_db("index_db", conn))

cocoindex_code-0.1.7/src/cocoindex_code/embedder.py

@@ -1,117 +0,0 @@
-"""Local SentenceTransformer embedder with device and trust_remote_code support."""
-
-from __future__ import annotations
-
-import threading
-from typing import TYPE_CHECKING, Any
-
-import cocoindex as coco
-import numpy as np
-from cocoindex.resources import schema as _schema
-from numpy.typing import NDArray
-
-from .config import config as _config
-
-if TYPE_CHECKING:
-    from sentence_transformers import SentenceTransformer
-
-
-class LocalEmbedder(_schema.VectorSchemaProvider):
-    """SentenceTransformer embedder with explicit device and trust_remote_code support.
-
-    Drop-in replacement for cocoindex's SentenceTransformerEmbedder that supports:
-    - Explicit device selection (e.g. "cuda", "cpu")
-    - trust_remote_code for models with custom pooling (e.g. Jina models)
-    """
-
-    def __init__(
-        self,
-        model_name_or_path: str,
-        *,
-        device: str = "cpu",
-        trust_remote_code: bool = False,
-        normalize_embeddings: bool = True,
-        query_prompt_name: str | None = None,
-    ) -> None:
-        self._model_name_or_path = model_name_or_path
-        self._device = device
-        self._trust_remote_code = trust_remote_code
-        self._normalize_embeddings = normalize_embeddings
-        self._query_prompt_name = query_prompt_name
-        self._model: SentenceTransformer | None = None
-        self._lock = threading.Lock()
-
-    def __getstate__(self) -> dict[str, Any]:
-        return {
-            "model_name_or_path": self._model_name_or_path,
-            "device": self._device,
-            "trust_remote_code": self._trust_remote_code,
-            "normalize_embeddings": self._normalize_embeddings,
-            "query_prompt_name": self._query_prompt_name,
-        }
-
-    def __setstate__(self, state: dict[str, Any]) -> None:
-        self._model_name_or_path = state["model_name_or_path"]
-        self._device = state["device"]
-        self._trust_remote_code = state["trust_remote_code"]
-        self._normalize_embeddings = state["normalize_embeddings"]
-        self._query_prompt_name = state.get("query_prompt_name")
-        self._model = None
-        self._lock = threading.Lock()
-
-    def _get_model(self) -> SentenceTransformer:
-        """Lazy-load the model with thread-safe double-checked locking."""
-        if self._model is None:
-            with self._lock:
-                if self._model is None:
-                    from sentence_transformers import SentenceTransformer
-
-                    self._model = SentenceTransformer(
-                        self._model_name_or_path,
-                        device=self._device,
-                        trust_remote_code=self._trust_remote_code,
-                    )
-        return self._model
-
-    @coco.fn.as_async(batching=True, runner=coco.GPU, memo=True, max_batch_size=_config.batch_size)
-    def embed(self, texts: list[str]) -> list[NDArray[np.float32]]:
-        """Embed a batch of texts into float32 vectors."""
-        model = self._get_model()
-        embeddings: NDArray[np.float32] = model.encode(
-            texts,
-            convert_to_numpy=True,
-            normalize_embeddings=self._normalize_embeddings,
-        )  # type: ignore[assignment]
-        return list(embeddings)
-
-    @coco.fn.as_async(batching=True, runner=coco.GPU, memo=True, max_batch_size=_config.batch_size)
-    def embed_query(self, texts: list[str]) -> list[NDArray[np.float32]]:
-        """Embed query texts, applying query_prompt_name if configured."""
-        model = self._get_model()
-        embeddings: NDArray[np.float32] = model.encode(
-            texts,
-            prompt_name=self._query_prompt_name,
-            convert_to_numpy=True,
-            normalize_embeddings=self._normalize_embeddings,
-        )  # type: ignore[assignment]
-        return list(embeddings)
-
-    @coco.fn.as_async(runner=coco.GPU, memo=True)
-    def __coco_vector_schema__(self) -> _schema.VectorSchema:
-        """Return the vector schema (dimension + dtype) for this model."""
-        model = self._get_model()
-        dim = model.get_sentence_embedding_dimension()
-        if dim is None:
-            raise RuntimeError(
-                f"Embedding dimension is unknown for model {self._model_name_or_path}."
-            )
-        return _schema.VectorSchema(dtype=np.dtype(np.float32), size=dim)
-
-    def __coco_memo_key__(self) -> object:
-        return (
-            self._model_name_or_path,
-            self._device,
-            self._trust_remote_code,
-            self._normalize_embeddings,
-            self._query_prompt_name,
-        )

cocoindex_code-0.1.7/src/cocoindex_code/query.py

@@ -1,71 +0,0 @@
-"""Query implementation for codebase search."""
-
-import cocoindex as coco
-
-from .config import config
-from .schema import QueryResult
-from .shared import SQLITE_DB, embedder
-
-
-async def query_codebase(
-    query: str,
-    limit: int = 10,
-    offset: int = 0,
-) -> list[QueryResult]:
-    """
-    Perform vector similarity search.
-
-    Uses sqlite-vec's vec_distance_cosine for similarity scoring.
-    """
-    if not config.target_sqlite_db_path.exists():
-        raise RuntimeError(
-            f"Index database not found at {config.target_sqlite_db_path}. "
-            "Please run a query with refresh_index=True first."
-        )
-
-    # Get the database connection from CocoIndex environment
-    coco_env = await coco.default_env()
-    db = coco_env.get_context(SQLITE_DB)
-
-    # Generate query embedding — use embed_query if available (supports asymmetric
-    # prompting for models like nomic-embed-code that use different prefixes for
-    # queries vs indexed documents).
-    if hasattr(embedder, "embed_query"):
-        query_embedding = await embedder.embed_query(query)
-    else:
-        query_embedding = await embedder.embed(query)
-
-    # Convert to bytes for sqlite-vec (float32)
-    embedding_bytes = query_embedding.astype("float32").tobytes()
-
-    # Query using sqlite-vec with readonly transaction
-    # vec_distance_cosine returns distance (lower is better),
-    # so we convert to similarity score (1 - distance)
-    with db.value.readonly() as conn:
-        cursor = conn.execute(
-            """
-            SELECT
-                file_path,
-                language,
-                content,
-                start_line,
-                end_line,
-                (1.0 - vec_distance_cosine(embedding, ?)) as score
-            FROM code_chunks
-            ORDER BY vec_distance_cosine(embedding, ?) ASC
-            LIMIT ? OFFSET ?
-            """,
-            (embedding_bytes, embedding_bytes, limit, offset),
-        )
-
-        return [
-            QueryResult(
-                file_path=row[0],
-                language=row[1],
-                content=row[2],
-                start_line=row[3],
-                end_line=row[4],
-                score=row[5],
-            )
-            for row in cursor.fetchall()
-        ]