cocoindex-code 0.1.0__tar.gz

cocoindex_code-0.1.0/.gitignore

@@ -0,0 +1,45 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+.env
+.venv
+env/
+venv/
+ENV/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Testing
+.tox/
+.coverage
+.coverage.*
+htmlcov/
+.pytest_cache/
+.mypy_cache/
+
+# CocoIndex
+.cocoindex_code/

cocoindex_code-0.1.0/PKG-INFO

@@ -0,0 +1,203 @@
+Metadata-Version: 2.4
+Name: cocoindex-code
+Version: 0.1.0
+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
+Project-URL: Issues, https://github.com/cocoindex-io/cocoindex-code/issues
+License-Expression: MIT
+Keywords: cocoindex,codebase,indexing,mcp,vector-search
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+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==1.0.0a10
+Requires-Dist: mcp>=1.0.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: sentence-transformers>=2.2.0
+Requires-Dist: sqlite-vec>=0.1.0
+Provides-Extra: dev
+Requires-Dist: mypy>=1.0.0; extra == 'dev'
+Requires-Dist: prek>=0.1.0; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
+Requires-Dist: pytest>=7.0.0; extra == 'dev'
+Requires-Dist: ruff>=0.1.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# CocoIndex Code
+
+An MCP (Model Context Protocol) server for indexing and querying codebases using [CocoIndex](https://cocoindex.io).
+
+## Features
+
+- **Semantic Code Search**: Find relevant code using natural language queries
+- **Incremental Indexing**: Only re-indexes changed files for fast updates
+- **Multi-Language Support**: Python, JavaScript/TypeScript, Rust, Go
+- **Vector Embeddings**: Uses sentence-transformers for semantic similarity
+- **SQLite Storage**: Portable, no external database required
+
+## Installation
+
+```bash
+pip install cocoindex-code
+```
+
+Or with uv:
+
+```bash
+uv pip install cocoindex-code
+```
+
+## Usage with Claude Code
+
+Add to your Claude Code MCP configuration (`.claude/mcp_config.json`):
+
+```json
+{
+  "mcpServers": {
+    "cocoindex-code": {
+      "command": "cocoindex-code",
+      "env": {
+        "COCOINDEX_CODE_ROOT_PATH": "/path/to/your/codebase"
+      }
+    }
+  }
+}
+```
+
+Or without explicit path (auto-discovers from current directory):
+
+```json
+{
+  "mcpServers": {
+    "cocoindex-code": {
+      "command": "cocoindex-code"
+    }
+  }
+}
+```
+
+## Configuration
+
+Environment variables:
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `COCOINDEX_CODE_ROOT_PATH` | Root path of the codebase | Auto-discovered (see below) |
+| `COCOINDEX_CODE_EMBEDDING_MODEL` | Embedding model to use | `sentence-transformers/all-MiniLM-L6-v2` |
+
+### Root Path Discovery
+
+If `COCOINDEX_CODE_ROOT_PATH` is not set, the codebase root is discovered by:
+
+1. Finding the nearest parent directory containing `.cocoindex_code/`
+2. Finding the nearest parent directory containing `.git/`
+3. Falling back to the current working directory
+
+## MCP Tools
+
+### `query`
+
+Search the codebase using semantic similarity.
+
+```
+query(
+    query: str,               # Natural language query or code snippet
+    limit: int = 10,          # Maximum results (1-100)
+    offset: int = 0,          # Pagination offset
+    refresh_index: bool = True  # Refresh index before querying
+)
+```
+
+The `refresh_index` parameter controls whether the index is refreshed before searching:
+
+- `True` (default): Refreshes the index to include any recent changes
+- `False`: Skip refresh for faster consecutive queries
+
+Returns matching code chunks with:
+
+- File path
+- Language
+- Code content
+- Line numbers (start/end)
+- Similarity score
+
+## Index Storage
+
+The index is stored in `.cocoindex_code/` under your codebase root:
+
+```
+your-project/
+├── .cocoindex_code/
+│   ├── target_sqlite.db  # Vector index (SQLite + sqlite-vec)
+│   └── cocoindex.db/     # CocoIndex state
+├── src/
+│   └── ...
+```
+
+Add `.cocoindex_code/` to your `.gitignore`.
+
+## Supported File Types
+
+- **Python**: `.py`, `.pyi`
+- **JavaScript**: `.js`, `.jsx`, `.mjs`, `.cjs`
+- **TypeScript**: `.ts`, `.tsx`
+- **Rust**: `.rs`
+- **Go**: `.go`
+
+Common generated directories are automatically excluded:
+
+- `__pycache__/`
+- `node_modules/`
+- `target/`
+- `dist/`
+- `build/`
+- `.git/`
+
+## Development
+
+### Local Testing with Claude Code
+
+To test locally without installing the package, use the Claude Code CLI:
+
+```bash
+claude mcp add cocoindex-code \
+  -- uv run --project /path/to/cocoindex-code cocoindex-code
+```
+
+Or add to `.mcp.json` in your project root:
+
+```json
+{
+  "mcpServers": {
+    "cocoindex-code": {
+      "command": "uv",
+      "args": ["run", "--project", "/path/to/cocoindex-code", "cocoindex-code"]
+    }
+  }
+}
+```
+
+### Running Tests
+
+```bash
+# Install dev dependencies
+uv sync --group dev
+
+# Run tests
+uv run pytest tests/ -v
+
+# Run pre-commit hooks
+uv run pre-commit run --all-files
+```
+
+## License
+
+MIT

cocoindex_code-0.1.0/README.md

@@ -0,0 +1,170 @@
+# CocoIndex Code
+
+An MCP (Model Context Protocol) server for indexing and querying codebases using [CocoIndex](https://cocoindex.io).
+
+## Features
+
+- **Semantic Code Search**: Find relevant code using natural language queries
+- **Incremental Indexing**: Only re-indexes changed files for fast updates
+- **Multi-Language Support**: Python, JavaScript/TypeScript, Rust, Go
+- **Vector Embeddings**: Uses sentence-transformers for semantic similarity
+- **SQLite Storage**: Portable, no external database required
+
+## Installation
+
+```bash
+pip install cocoindex-code
+```
+
+Or with uv:
+
+```bash
+uv pip install cocoindex-code
+```
+
+## Usage with Claude Code
+
+Add to your Claude Code MCP configuration (`.claude/mcp_config.json`):
+
+```json
+{
+  "mcpServers": {
+    "cocoindex-code": {
+      "command": "cocoindex-code",
+      "env": {
+        "COCOINDEX_CODE_ROOT_PATH": "/path/to/your/codebase"
+      }
+    }
+  }
+}
+```
+
+Or without explicit path (auto-discovers from current directory):
+
+```json
+{
+  "mcpServers": {
+    "cocoindex-code": {
+      "command": "cocoindex-code"
+    }
+  }
+}
+```
+
+## Configuration
+
+Environment variables:
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `COCOINDEX_CODE_ROOT_PATH` | Root path of the codebase | Auto-discovered (see below) |
+| `COCOINDEX_CODE_EMBEDDING_MODEL` | Embedding model to use | `sentence-transformers/all-MiniLM-L6-v2` |
+
+### Root Path Discovery
+
+If `COCOINDEX_CODE_ROOT_PATH` is not set, the codebase root is discovered by:
+
+1. Finding the nearest parent directory containing `.cocoindex_code/`
+2. Finding the nearest parent directory containing `.git/`
+3. Falling back to the current working directory
+
+## MCP Tools
+
+### `query`
+
+Search the codebase using semantic similarity.
+
+```
+query(
+    query: str,               # Natural language query or code snippet
+    limit: int = 10,          # Maximum results (1-100)
+    offset: int = 0,          # Pagination offset
+    refresh_index: bool = True  # Refresh index before querying
+)
+```
+
+The `refresh_index` parameter controls whether the index is refreshed before searching:
+
+- `True` (default): Refreshes the index to include any recent changes
+- `False`: Skip refresh for faster consecutive queries
+
+Returns matching code chunks with:
+
+- File path
+- Language
+- Code content
+- Line numbers (start/end)
+- Similarity score
+
+## Index Storage
+
+The index is stored in `.cocoindex_code/` under your codebase root:
+
+```
+your-project/
+├── .cocoindex_code/
+│   ├── target_sqlite.db  # Vector index (SQLite + sqlite-vec)
+│   └── cocoindex.db/     # CocoIndex state
+├── src/
+│   └── ...
+```
+
+Add `.cocoindex_code/` to your `.gitignore`.
+
+## Supported File Types
+
+- **Python**: `.py`, `.pyi`
+- **JavaScript**: `.js`, `.jsx`, `.mjs`, `.cjs`
+- **TypeScript**: `.ts`, `.tsx`
+- **Rust**: `.rs`
+- **Go**: `.go`
+
+Common generated directories are automatically excluded:
+
+- `__pycache__/`
+- `node_modules/`
+- `target/`
+- `dist/`
+- `build/`
+- `.git/`
+
+## Development
+
+### Local Testing with Claude Code
+
+To test locally without installing the package, use the Claude Code CLI:
+
+```bash
+claude mcp add cocoindex-code \
+  -- uv run --project /path/to/cocoindex-code cocoindex-code
+```
+
+Or add to `.mcp.json` in your project root:
+
+```json
+{
+  "mcpServers": {
+    "cocoindex-code": {
+      "command": "uv",
+      "args": ["run", "--project", "/path/to/cocoindex-code", "cocoindex-code"]
+    }
+  }
+}
+```
+
+### Running Tests
+
+```bash
+# Install dev dependencies
+uv sync --group dev
+
+# Run tests
+uv run pytest tests/ -v
+
+# Run pre-commit hooks
+uv run pre-commit run --all-files
+```
+
+## License
+
+MIT

cocoindex_code-0.1.0/pyproject.toml

@@ -0,0 +1,85 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "cocoindex-code"
+version = "0.1.0"
+description = "MCP server for indexing and querying codebases using CocoIndex"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.11"
+keywords = ["mcp", "codebase", "indexing", "vector-search", "cocoindex"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+
+dependencies = [
+    "mcp>=1.0.0",
+    "cocoindex==1.0.0a10",
+    "sentence-transformers>=2.2.0",
+    "sqlite-vec>=0.1.0",
+    "pydantic>=2.0.0",
+    "numpy>=1.24.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-asyncio>=0.21.0",
+    "pytest-cov>=4.0.0",
+    "ruff>=0.1.0",
+    "mypy>=1.0.0",
+    "prek>=0.1.0",
+]
+
+[project.scripts]
+cocoindex-code = "cocoindex_code:main"
+
+[project.urls]
+Homepage = "https://github.com/cocoindex-io/cocoindex-code"
+Repository = "https://github.com/cocoindex-io/cocoindex-code"
+Issues = "https://github.com/cocoindex-io/cocoindex-code/issues"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/cocoindex_code"]
+
+[tool.hatch.build.targets.sdist]
+include = ["/src", "/README.md", "/LICENSE"]
+
+[dependency-groups]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-asyncio>=0.21.0",
+    "pytest-cov>=4.0.0",
+    "ruff>=0.1.0",
+    "mypy>=1.0.0",
+    "prek>=0.1.0",
+]
+
+[tool.uv]
+prerelease = "explicit"
+
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "N", "W", "UP"]
+
+[tool.mypy]
+python_version = "3.11"
+strict = true
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_functions = ["test_*"]
+addopts = "-v --tb=short"

cocoindex_code-0.1.0/src/cocoindex_code/__init__.py

@@ -0,0 +1,7 @@
+"""CocoIndex Code - MCP server for indexing and querying codebases."""
+
+from .config import Config
+from .server import main, mcp
+
+__version__ = "0.1.0"
+__all__ = ["Config", "main", "mcp"]

cocoindex_code-0.1.0/src/cocoindex_code/__main__.py

@@ -0,0 +1,6 @@
+"""Entry point for `python -m cocoindex_code`."""
+
+from .server import main
+
+if __name__ == "__main__":
+    main()

cocoindex_code-0.1.0/src/cocoindex_code/config.py

@@ -0,0 +1,87 @@
+"""Configuration for CocoIndex Code MCP server."""
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+
+def _find_root_with_marker(start_dir: Path, marker: str) -> Path | None:
+    """Find the nearest parent directory containing the given marker directory."""
+    current = start_dir.resolve()
+    while current != current.parent:
+        if (current / marker).is_dir():
+            return current
+        current = current.parent
+    # Check root directory too
+    if (current / marker).is_dir():
+        return current
+    return None
+
+
+def _discover_codebase_root() -> Path:
+    """
+    Discover the codebase root directory.
+
+    Discovery order:
+    1. Find nearest parent with `.cocoindex_code` directory
+    2. Find nearest parent with `.git` directory
+    3. Fall back to current working directory
+    """
+    cwd = Path.cwd()
+
+    # First, look for existing .cocoindex_code directory
+    root = _find_root_with_marker(cwd, ".cocoindex_code")
+    if root is not None:
+        return root
+
+    # Then, look for .git directory
+    root = _find_root_with_marker(cwd, ".git")
+    if root is not None:
+        return root
+
+    # Fall back to current working directory
+    return cwd
+
+
+@dataclass
+class Config:
+    """Configuration loaded from environment variables."""
+
+    codebase_root_path: Path
+    embedding_model: str
+    index_dir: Path
+
+    @classmethod
+    def from_env(cls) -> "Config":
+        """Load configuration from environment variables."""
+        # Get root path from env or discover it
+        root_path_str = os.environ.get("COCOINDEX_CODE_ROOT_PATH")
+        if root_path_str:
+            root = Path(root_path_str).resolve()
+        else:
+            root = _discover_codebase_root()
+
+        # Get embedding model
+        embedding_model = os.environ.get(
+            "COCOINDEX_CODE_EMBEDDING_MODEL",
+            "sentence-transformers/all-MiniLM-L6-v2",
+        )
+
+        # Index directory is always under the root
+        index_dir = root / ".cocoindex_code"
+
+        return cls(
+            codebase_root_path=root,
+            embedding_model=embedding_model,
+            index_dir=index_dir,
+        )
+
+    @property
+    def target_sqlite_db_path(self) -> Path:
+        """Path to the vector index SQLite database."""
+        return self.index_dir / "target_sqlite.db"
+
+    @property
+    def cocoindex_db_path(self) -> Path:
+        """Path to the CocoIndex state database."""
+        return self.index_dir / "cocoindex.db"

cocoindex_code-0.1.0/src/cocoindex_code/indexer.py

@@ -0,0 +1,164 @@
+"""CocoIndex app for indexing codebases."""
+
+import asyncio
+
+import cocoindex.asyncio as coco_aio
+from cocoindex.connectors import localfs, sqlite
+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 .shared import SQLITE_DB, CodeChunk, config, embedder
+
+# File patterns for supported languages
+INCLUDED_PATTERNS = [
+    "*.py",  # Python
+    "*.pyi",  # Python stubs
+    "*.js",  # JavaScript
+    "*.jsx",  # JavaScript React
+    "*.ts",  # TypeScript
+    "*.tsx",  # TypeScript React
+    "*.mjs",  # JavaScript ES modules
+    "*.cjs",  # JavaScript CommonJS
+    "*.rs",  # Rust
+    "*.go",  # Go
+]
+
+EXCLUDED_PATTERNS = [
+    ".*/**",  # Hidden directories
+    "**/__pycache__/**",  # Python cache
+    "**/node_modules/**",  # Node.js dependencies
+    "**/target/**",  # Rust/Maven build output
+    "**/dist/**",  # Distribution directories
+    "**/build/**",  # Build directories
+    "**/vendor/**",  # Go vendor directory
+    "**/.git/**",  # Git directory
+    "**/.cocoindex_code/**",  # Our own index directory
+    "*.min.js",  # Minified JavaScript
+    "*.min.css",  # Minified CSS
+    "*.lock",  # Lock files
+    "**/package-lock.json",  # NPM lock
+    "**/yarn.lock",  # Yarn lock
+    "**/Cargo.lock",  # Cargo lock
+    "**/go.sum",  # Go sum
+    "**/*.pyc",  # Python bytecode
+    "**/*.pyo",  # Python optimized bytecode
+    "**/*.so",  # Shared objects
+    "**/*.dylib",  # macOS dynamic libraries
+    "**/*.dll",  # Windows dynamic libraries
+]
+
+# Chunking configuration
+CHUNK_SIZE = 1000
+MIN_CHUNK_SIZE = 300
+CHUNK_OVERLAP = 200
+
+# Chunking splitter (stateless, can be module-level)
+splitter = RecursiveSplitter()
+
+
+@coco_aio.function
+async def process_chunk(
+    file_path: str,
+    chunk: Chunk,
+    language: str,
+    id_gen: IdGenerator,
+    table: sqlite.TableTarget,
+) -> None:
+    """Process a single chunk: embed and store."""
+    id, chunk_embedding = await asyncio.gather(
+        id_gen.next_id(chunk.text),
+        embedder.embed(chunk.text),
+    )
+    table.declare_row(
+        row=CodeChunk(  # type: ignore[arg-type]
+            id=id,
+            file_path=file_path,
+            language=language,
+            content=chunk.text,
+            start_line=chunk.start.line,
+            end_line=chunk.end.line,
+            embedding=chunk_embedding,
+        )
+    )
+
+
+@coco_aio.function(memo=True)
+async def process_file(
+    file: localfs.File,
+    table: sqlite.TableTarget,
+) -> None:
+    """Process a single file: chunk, embed, and store."""
+    # Read file content
+    try:
+        content = file.read_text()
+    except UnicodeDecodeError:
+        # Skip binary files
+        return
+
+    if not content.strip():
+        return
+
+    # Get relative path and detect language
+    language = detect_code_language(filename=file.file_path.path.name) or "text"
+
+    # Split into chunks
+    chunks = splitter.split(
+        content,
+        chunk_size=CHUNK_SIZE,
+        min_chunk_size=MIN_CHUNK_SIZE,
+        chunk_overlap=CHUNK_OVERLAP,
+        language=language,
+    )
+
+    id_gen = IdGenerator()
+    await asyncio.gather(
+        *(
+            process_chunk(str(file.file_path.path), chunk, language, id_gen, table)
+            for chunk in chunks
+        )
+    )
+
+
+@coco_aio.function
+async def app_main() -> None:
+    """Main indexing function - walks files and processes each."""
+    db = coco_aio.use_context(SQLITE_DB)
+
+    # Declare the table target for storing embeddings
+    table = await coco_aio.mount_run(
+        coco_aio.component_subpath("setup", "table"),
+        db.declare_table_target,
+        table_name="code_chunks",
+        table_schema=await sqlite.TableSchema.from_class(
+            CodeChunk,
+            primary_key=["id"],
+        ),
+    ).result()
+
+    # Walk source directory
+    files = localfs.walk_dir(
+        config.codebase_root_path,
+        recursive=True,
+        path_matcher=PatternFilePathMatcher(
+            included_patterns=INCLUDED_PATTERNS,
+            excluded_patterns=EXCLUDED_PATTERNS,
+        ),
+    )
+
+    # Process each file
+    for f in files:
+        coco_aio.mount(
+            coco_aio.component_subpath("process", str(f.file_path.path)),
+            process_file,
+            f,
+            table,
+        )
+
+
+# Create the app
+app = coco_aio.App(
+    coco_aio.AppConfig(name="CocoIndexCode"),
+    app_main,
+)

cocoindex_code-0.1.0/src/cocoindex_code/query.py

@@ -0,0 +1,64 @@
+"""Query implementation for codebase search."""
+
+import cocoindex as coco
+
+from .schema import QueryResult
+from .shared import SQLITE_DB, config, 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
+    db = coco.default_env().get_context(SQLITE_DB)
+
+    # Generate query embedding
+    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()
+        ]

cocoindex_code-0.1.0/src/cocoindex_code/schema.py

@@ -0,0 +1,29 @@
+"""Data models for CocoIndex Code."""
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class CodeChunk:
+    """Represents an indexed code chunk stored in SQLite."""
+
+    id: int
+    file_path: str
+    language: str
+    content: str
+    start_line: int
+    end_line: int
+    embedding: Any  # NDArray - type hint relaxed for compatibility
+
+
+@dataclass
+class QueryResult:
+    """Result from a vector similarity query."""
+
+    file_path: str
+    language: str
+    content: str
+    start_line: int
+    end_line: int
+    score: float

cocoindex_code-0.1.0/src/cocoindex_code/server.py

@@ -0,0 +1,150 @@
+"""MCP server for codebase indexing and querying."""
+
+import asyncio
+
+from mcp.server.fastmcp import FastMCP
+from pydantic import BaseModel, Field
+
+from .indexer import app as indexer_app
+from .query import query_codebase
+
+# Initialize MCP server
+mcp = FastMCP(
+    "cocoindex-code",
+    instructions="""
+This server provides semantic code search for the codebase.
+This allows you to quickly and cheaply search for code related to a concept or functionality
+across the entire codebase.
+
+Use the `query` tool when you need to:
+- Find code related to a concept or functionality
+- Search for implementations of specific features
+- Discover how something is done in the codebase
+- Find similar code patterns
+
+The `query` tool has a `refresh_index` parameter (default: True) that refreshes
+the index before searching. Set it to False for consecutive queries to avoid
+redundant refreshes.
+
+The search uses vector embeddings for semantic similarity, so you can describe
+what you're looking for in natural language rather than exact text matches.
+""".strip(),
+)
+
+# Lock to prevent concurrent index updates
+_index_lock = asyncio.Lock()
+
+
+async def _refresh_index() -> None:
+    """Refresh the index. Uses lock to prevent concurrent updates."""
+    async with _index_lock:
+        await indexer_app.update(report_to_stdout=False)
+
+
+# === Pydantic Models for Tool Inputs/Outputs ===
+
+
+class CodeChunkResult(BaseModel):
+    """A single code chunk result."""
+
+    file_path: str = Field(description="Relative path to the file")
+    language: str = Field(description="Programming language")
+    content: str = Field(description="The code content")
+    start_line: int = Field(description="Starting line number (1-indexed)")
+    end_line: int = Field(description="Ending line number (1-indexed)")
+    score: float = Field(description="Similarity score (0-1, higher is better)")
+
+
+class QueryResultModel(BaseModel):
+    """Result from query tool."""
+
+    success: bool
+    results: list[CodeChunkResult] = Field(default_factory=list)
+    total_returned: int = Field(default=0)
+    offset: int = Field(default=0)
+    message: str | None = None
+
+
+# === MCP Tools ===
+
+
+@mcp.tool(
+    name="query",
+    description=(
+        "Search the codebase using semantic similarity. "
+        "Returns relevant code chunks with file locations and similarity scores. "
+        "Use natural language queries or code snippets to find related code."
+    ),
+)
+async def query(
+    query: str = Field(description="Natural language query or code snippet to search for"),
+    limit: int = Field(
+        default=10,
+        ge=1,
+        le=100,
+        description="Maximum number of results to return (1-100)",
+    ),
+    offset: int = Field(
+        default=0,
+        ge=0,
+        description="Number of results to skip for pagination",
+    ),
+    refresh_index: bool = Field(
+        default=True,
+        description=(
+            "Whether to refresh the index before querying. "
+            "Set to False for consecutive queries to skip redundant refreshes."
+        ),
+    ),
+) -> QueryResultModel:
+    """Query the codebase index."""
+    try:
+        # Refresh index if requested
+        if refresh_index:
+            await _refresh_index()
+
+        results = await query_codebase(query=query, limit=limit, offset=offset)
+
+        return QueryResultModel(
+            success=True,
+            results=[
+                CodeChunkResult(
+                    file_path=r.file_path,
+                    language=r.language,
+                    content=r.content,
+                    start_line=r.start_line,
+                    end_line=r.end_line,
+                    score=r.score,
+                )
+                for r in results
+            ],
+            total_returned=len(results),
+            offset=offset,
+        )
+    except RuntimeError as e:
+        # Index doesn't exist
+        return QueryResultModel(
+            success=False,
+            message=str(e),
+        )
+    except Exception as e:
+        return QueryResultModel(
+            success=False,
+            message=f"Query failed: {e!s}",
+        )
+
+
+async def _async_main() -> None:
+    """Async entry point for the MCP server."""
+    # Refresh index in background so startup isn't blocked
+    asyncio.create_task(_refresh_index())
+    await mcp.run_stdio_async()
+
+
+def main() -> None:
+    """Entry point for the MCP server."""
+    asyncio.run(_async_main())
+
+
+if __name__ == "__main__":
+    main()

cocoindex_code-0.1.0/src/cocoindex_code/shared.py

@@ -0,0 +1,52 @@
+"""Shared resources for CocoIndex Code."""
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Annotated
+
+import cocoindex as coco
+from cocoindex.connectors import sqlite
+from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder
+from numpy.typing import NDArray
+
+from .config import Config
+
+# Load configuration at module level
+config = Config.from_env()
+
+# Initialize embedder at module level
+embedder = SentenceTransformerEmbedder(config.embedding_model)
+
+# Context key for SQLite database (connection managed in lifespan)
+SQLITE_DB = coco.ContextKey[sqlite.SqliteDatabase]("sqlite_db")
+
+
+@coco.lifespan
+def coco_lifespan(builder: coco.EnvironmentBuilder) -> Iterator[None]:
+    """Set up database connection."""
+    # Ensure index directory exists
+    config.index_dir.mkdir(parents=True, exist_ok=True)
+
+    # Set CocoIndex state database path
+    builder.settings.db_path = config.cocoindex_db_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))
+
+    yield
+
+    conn.close()
+
+
+@dataclass
+class CodeChunk:
+    """Schema for storing code chunks in SQLite."""
+
+    id: int
+    file_path: str
+    language: str
+    content: str
+    start_line: int
+    end_line: int
+    embedding: Annotated[NDArray, embedder]  # type: ignore[type-arg]