codeembed 0.1.0__py3-none-any.whl

codeembed/llm/openai_adapter.py

@@ -0,0 +1,96 @@
+from typing import List, Optional, Type, TypeVar, cast
+
+from openai import OpenAI
+from openai._types import omit
+from openai.types.chat import ChatCompletionMessageParam
+from pydantic import BaseModel
+
+from codeembed.llm.base import LLMServiceBase
+from codeembed.llm.models import ChatMessage, LLMResponse, StructuredLLMResponse
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class OpenAILLMService(LLMServiceBase):
+    def __init__(self, client: OpenAI):
+        self._client = client
+
+    def generate_structured_output(
+        self,
+        messages: List[ChatMessage],
+        llm_model: str,
+        output_format: Type[T],
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+    ) -> StructuredLLMResponse[T]:
+
+        openai_messages = cast(List[ChatCompletionMessageParam], messages)
+
+        # newer OpenAI models requires max_completion_tokens instead of max_tokens...
+        _max_tokens = max_tokens if max_tokens is not None else omit
+        max_completion_tokens = omit
+        if self._is_reasoning_model(llm_model):
+            max_completion_tokens = _max_tokens
+            _max_tokens = omit
+
+        completion = self._client.beta.chat.completions.parse(
+            messages=openai_messages,
+            model=llm_model,
+            response_format=output_format,
+            max_tokens=_max_tokens,
+            max_completion_tokens=max_completion_tokens,
+            temperature=temperature if temperature is not None else omit,
+        )
+
+        parsed = completion.choices[0].message.parsed
+
+        if parsed is None:
+            raise ValueError("LLM did not return structured output")
+
+        return StructuredLLMResponse(
+            input_tokens=completion.usage.prompt_tokens if completion.usage else 0,
+            output_tokens=completion.usage.completion_tokens if completion.usage else 0,
+            data=parsed,
+            llm_model=llm_model,
+        )
+
+    def generate_response(
+        self,
+        messages: List[ChatMessage],
+        llm_model: str,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+    ) -> LLMResponse:
+
+        openai_messages = cast(List[ChatCompletionMessageParam], messages)
+
+        # newer OpenAI models requires max_completion_tokens instead of max_tokens...
+        _max_tokens = max_tokens if max_tokens is not None else omit
+        max_completion_tokens = omit
+        if self._is_reasoning_model(llm_model):
+            max_completion_tokens = _max_tokens
+            _max_tokens = omit
+
+        completion = self._client.chat.completions.create(
+            messages=openai_messages,
+            model=llm_model,
+            max_tokens=_max_tokens,
+            max_completion_tokens=max_completion_tokens,
+            temperature=temperature if temperature is not None else omit,
+            response_format={"type": "text"},
+        )
+
+        response = completion.choices[0].message.content
+
+        if response is None:
+            raise ValueError("LLM did not return a response")
+
+        return LLMResponse(
+            input_tokens=completion.usage.prompt_tokens if completion.usage else 0,
+            output_tokens=completion.usage.completion_tokens if completion.usage else 0,
+            response=response,
+            llm_model=llm_model,
+        )
+
+    def _is_reasoning_model(self, llm_model: str) -> bool:
+        return not llm_model.startswith("gpt-4") and not llm_model.startswith("gpt-3")

codeembed/mcp_server.py

@@ -0,0 +1,45 @@
+import asyncio
+from contextlib import asynccontextmanager, suppress
+
+from mcp.server.fastmcp import FastMCP
+
+from codeembed.bootstrap.services import embed_loop, get_search_service
+
+
+@asynccontextmanager
+async def lifespan(server):
+    task = asyncio.create_task(embed_loop())
+    try:
+        yield
+    finally:
+        task.cancel()
+        with suppress(asyncio.CancelledError):
+            await task
+
+
+mcp = FastMCP("Codebase Embedder", lifespan=lifespan, json_response=True)
+
+
+@mcp.tool()
+def search(query: str, top_n: int = 10) -> str:
+    """Searches the embedded codebases using semantic similarity.
+
+    Use this tool as the FIRST step for any question about the codebase — how something
+    works, where something is defined, what calls what, etc. Prefer this over grep or
+    file reads for exploratory questions; it returns ranked, summarized results instantly.
+
+    Args:
+        query: A natural-language description of what you're looking for.
+               Examples: "how are deltas computed", "LLM service abstraction",
+               "error handling in the embedding pipeline".
+        top_n: Number of results to return (default 10).
+    """
+    search_service = get_search_service()
+    return search_service.search(query, top_n)
+
+
+# NOTE: We could, e.g., add a tool to add codebases, but that sounds very risky.
+
+
+if __name__ == "__main__":
+    mcp.run(transport="stdio")

codeembed/setup_logger.py

@@ -0,0 +1,34 @@
+import logging
+import os
+from logging.handlers import TimedRotatingFileHandler
+
+_LOG_DIR = ".codeembed/logs"
+_LOG_FILE = os.path.join(_LOG_DIR, "codeembed.log")
+
+
+def setup_logger(level: int = logging.INFO) -> logging.Logger:
+    logger = logging.getLogger()
+
+    # Avoid duplicate handlers if called multiple times
+    if logger.handlers:
+        return logger
+
+    logger.setLevel(level)
+
+    formatter = logging.Formatter(
+        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    stream_handler = logging.StreamHandler()
+    stream_handler.setLevel(level)
+    stream_handler.setFormatter(formatter)
+    logger.addHandler(stream_handler)
+
+    os.makedirs(_LOG_DIR, exist_ok=True)
+    file_handler = TimedRotatingFileHandler(_LOG_FILE, when="midnight", backupCount=7, encoding="utf-8")
+    file_handler.setLevel(level)
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
+    return logger

codeembed/utils/__init__.py

@@ -0,0 +1,9 @@
+from codeembed.utils.checksum_utils import string_to_sha256
+from codeembed.utils.string_utils import truncate_string
+from codeembed.utils.time_utils import utc_now
+
+__all__ = [
+    "string_to_sha256",
+    "truncate_string",
+    "utc_now",
+]

codeembed/utils/checksum_utils.py

@@ -0,0 +1,5 @@
+import hashlib
+
+
+def string_to_sha256(s: str) -> str:
+    return hashlib.sha256(s.encode()).hexdigest()

codeembed/utils/string_utils.py

@@ -0,0 +1,5 @@
+def truncate_string(s: str, max_length: int) -> str:
+    """Truncate a string to a specified maximum length, adding ellipsis if truncated."""
+    if len(s) <= max_length:
+        return s
+    return s[: max_length - 3] + "..."

codeembed/utils/time_utils.py

@@ -0,0 +1,5 @@
+from datetime import datetime, timezone
+
+
+def utc_now() -> datetime:
+    return datetime.now(timezone.utc)

codeembed/vector_db/__init__.py

@@ -0,0 +1,9 @@
+from codeembed.vector_db.base import VectorDbBase
+from codeembed.vector_db.chromadb_adapter import ChromaDbAdapter
+from codeembed.vector_db.models import Chunk
+
+__all__ = [
+    "ChromaDbAdapter",
+    "Chunk",
+    "VectorDbBase",
+]

codeembed/vector_db/base.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Iterator, List, Optional
+from uuid import UUID
+
+from codeembed.vector_db.models import Chunk
+
+
+class VectorDbBase(ABC):
+    @abstractmethod
+    def add_chunks(self, chunks: List[Chunk]) -> None:
+        pass
+
+    @abstractmethod
+    def search(self, query: str, top_n: int) -> List[Chunk]:
+        """Vector search. Returns top_n most relevant results."""
+
+    @abstractmethod
+    def iter_chunks(self, where: Optional[Dict[str, str]] = None) -> Iterator[Chunk]:
+        """
+        Iterates all chunks stored in the vector database.
+
+        For simplicity exposes 'where' argument which is a ChromaDB specific filter.
+        """
+
+    @abstractmethod
+    def delete_chunks(self, chunk_ids: List[UUID]) -> None:
+        pass

codeembed/vector_db/chromadb_adapter.py

@@ -0,0 +1,130 @@
+from datetime import datetime
+from typing import Dict, Iterator, List, Optional, Type, TypeVar
+from uuid import UUID
+
+import chromadb
+from chromadb.api.types import Metadata, QueryResult
+
+from codeembed.vector_db.base import VectorDbBase
+from codeembed.vector_db.models import Chunk
+
+T = TypeVar("T")
+
+
+class ChromaDbAdapter(VectorDbBase):
+    def __init__(self, collection_name: str) -> None:
+        # TODO: Support adding EmbeddingServiceBase and replacing ChromaDB default embedder.
+        self._client = chromadb.PersistentClient(path="./.codeembed")
+        self._collection = self._client.get_or_create_collection(collection_name)
+
+    def add_chunks(self, chunks: List[Chunk]) -> None:
+        documents: List[str] = [chunk.content for chunk in chunks]
+        ids: List[str] = [str(chunk.id) for chunk in chunks]
+
+        metadatas: List[Metadata] = [
+            {
+                "modified_at": chunk.modified_at.isoformat(),
+                "file_path": chunk.file_path,
+                "line_start": chunk.line_start,
+                "line_end": chunk.line_end,
+                "raw_code": chunk.raw_code,  # Assume ChromaDB can handle None values.
+                "file_sha256_checksum": chunk.file_sha256_checksum,
+            }
+            for chunk in chunks
+        ]
+
+        self._collection.add(
+            documents=documents,
+            ids=ids,
+            metadatas=metadatas,
+        )
+
+    def search(self, query: str, top_n: int) -> List[Chunk]:
+        # TODO: Support filtering.
+        results: QueryResult = self._collection.query(
+            query_texts=[query],
+            n_results=top_n,
+        )
+
+        ids = results["ids"][0]
+        docs = results["documents"][0]  # type: ignore
+        metas = results["metadatas"][0]  # type: ignore
+
+        chunks_out: List[Chunk] = []
+
+        for i in range(len(ids)):
+            # Can be simplified by adding a "get_safe_val" or similar.
+            modified_at = self._get_safe_val(metas[i], "modified_at", str)
+            modified_at = datetime.fromisoformat(modified_at)
+            file_path = self._get_safe_val(metas[i], "file_path", str)
+            line_start = self._get_safe_val(metas[i], "line_start", int)
+            line_end = self._get_safe_val(metas[i], "line_end", int)
+            raw_code = self._get_safe_val(metas[i], "raw_code", str, allow_none=True)
+            file_sha256_checksum = self._get_safe_val(metas[i], "file_sha256_checksum", str)
+            chunks_out.append(
+                Chunk(
+                    id=UUID(ids[i]),
+                    content=docs[i],
+                    modified_at=modified_at,
+                    file_path=file_path,
+                    line_start=line_start,
+                    line_end=line_end,
+                    raw_code=raw_code,
+                    file_sha256_checksum=file_sha256_checksum,
+                )
+            )
+
+        return chunks_out
+
+    def iter_chunks(self, where: Optional[Dict[str, str]] = None) -> Iterator[Chunk]:
+        offset = 0
+        limit = 100
+
+        while True:
+            results = self._collection.get(
+                limit=limit,
+                offset=offset,
+                where=where,  # type: ignore
+            )
+
+            ids = results["ids"]
+
+            if not ids:
+                break
+
+            docs = results["documents"] or []
+            metas = results["metadatas"] or []
+
+            for i in range(len(ids)):
+                # Can be simplified by adding a "get_safe_val" or similar.
+                modified_at = self._get_safe_val(metas[i], "modified_at", str)
+                modified_at = datetime.fromisoformat(modified_at)
+                file_path = self._get_safe_val(metas[i], "file_path", str)
+                line_start = self._get_safe_val(metas[i], "line_start", int)
+                line_end = self._get_safe_val(metas[i], "line_end", int)
+                raw_code = self._get_safe_val(metas[i], "raw_code", str, allow_none=True)
+                file_sha256_checksum = self._get_safe_val(metas[i], "file_sha256_checksum", str)
+                yield Chunk(
+                    id=UUID(ids[i]),
+                    content=docs[i],
+                    modified_at=modified_at,
+                    file_path=file_path,
+                    line_start=line_start,
+                    line_end=line_end,
+                    raw_code=raw_code,
+                    file_sha256_checksum=file_sha256_checksum,
+                )
+
+            offset += limit
+
+    def delete_chunks(self, chunk_ids: List[UUID]) -> None:
+        # Maybe batch if list is very long? I hope ChromaDB does so internally.
+        self._collection.delete(ids=[str(chunk_id) for chunk_id in chunk_ids])
+
+    def _get_safe_val(self, meta: Metadata, key: str, expected_type: Type[T], allow_none: bool = False) -> T:
+        val = meta.get(key)
+        if val is None and allow_none:
+            return None  # type: ignore
+        if not isinstance(val, expected_type):
+            raise ValueError(f"Expected {expected_type}, got {type(val)} for ChromaDB metadata key '{key}'.")
+        return val

codeembed/vector_db/models.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional
+from uuid import UUID
+
+
+@dataclass
+class Chunk:
+    id: UUID
+    content: str
+    modified_at: datetime
+    file_path: str
+    line_start: int
+    line_end: int
+    raw_code: Optional[str]
+    file_sha256_checksum: str

codeembed-0.1.0.dist-info/METADATA

@@ -0,0 +1,292 @@
+Metadata-Version: 2.4
+Name: codeembed
+Version: 0.1.0
+Summary: Embeds your codebase and makes it available for quick LLM lookups via MCP.
+Project-URL: Homepage, https://github.com/robino16/codeembed
+Project-URL: Repository, https://github.com/robino16/codeembed
+Project-URL: Issues, https://github.com/robino16/codeembed/issues
+Author-email: robino16 <robinoms.dev@proton.me>
+License: MIT License
+        
+        Copyright (c) 2026 robino16/robinoms
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: codebase,embeddings,llm,mcp,rag,vector-search
+Classifier: Development Status :: 4 - Beta
+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: Programming Language :: Python :: 3.14
+Requires-Python: >=3.11
+Requires-Dist: chromadb<2,>=1.5
+Requires-Dist: mcp<2,>=1.26
+Requires-Dist: ollama<1,>=0.6
+Requires-Dist: pydantic<3,>=2.13
+Requires-Dist: python-dotenv<2,>=1.2.2
+Requires-Dist: tiktoken<1,>=0.12
+Provides-Extra: openai
+Requires-Dist: azure-identity<2,>=1.25; extra == 'openai'
+Requires-Dist: openai<3,>=2.33; extra == 'openai'
+Description-Content-Type: text/markdown
+
+# CodeEmbed
+
+Embeds your codebase into a local vector database and exposes it as an MCP tool, giving AI assistants like Claude Code fast semantic search over your code.
+
+Particularly useful for questions like:
+
+- How is X implemented in this repo?
+- Where is X defined or used?
+- Does this repo already have X?
+
+For other questions, the agent will fall back to normal lookups.
+CodeEmbed can improve lookup speed and accuracy, especially for finding existing implementations before writing new ones.
+Note that the biggest bottleneck in coding agents is LLM thinking and token generation — solid prompts and follow-up questions still matter.
+
+Uses [ChromaDB](https://github.com/chroma-core/chroma) for local vector storage and either [Ollama](https://github.com/ollama/ollama) or OpenAI (including OpenAI models via Azure AI Foundry) for LLM analysis.
+
+## Prerequisites
+
+- [Python](https://python.org) 3.11+
+- [uv](https://github.com/astral-sh/uv)
+- One of:
+  - [Ollama](https://ollama.com) running locally, **or**
+  - An OpenAI API key or Azure OpenAI endpoint
+
+## Installation
+
+**With Ollama:**
+
+```bash
+uv tool install codeembed
+```
+
+**With OpenAI / Azure OpenAI:**
+
+```bash
+uv tool install 'codeembed[openai]'
+```
+
+> **Supply chain safety:** To reduce the risk of newly-published malicious packages, consider adding `exclude-newer = "7 days"` to your global [`uv.toml`](https://docs.astral.sh/uv/reference/settings/#exclude-newer). This prevents `uv` from installing packages published in the last 7 days.
+
+### Manual installation (from source)
+
+If CodeEmbed is not published to PyPI, install it directly from source:
+
+```bash
+git clone https://github.com/robino16/codeembed
+cd codeembed
+
+# With Ollama
+uv tool install .
+
+# With OpenAI support
+uv tool install '.[openai]'
+```
+
+Then run `codeembed init` inside of your target repository.
+
+## Upgrading
+
+```bash
+uv tool upgrade codeembed
+```
+
+## Usage
+
+CodeEmbed is intended to be used within a single project — run all commands from your project root. Each project gets its own local vector database stored in `.codeembed/`.
+
+Supported file types: `.py`, `.md`, `.ts`, `.tsx`, `.js`, `.jsx`.
+
+**1. Initialize** (run once in your project root):
+
+```bash
+codeembed init
+```
+
+Creates a `codeembed.toml` config and configures your `.gitignore`. You'll be prompted to select a provider (Ollama or OpenAI) and a model. You'll also be offered the option to automatically configure Claude Code and/or GitHub Copilot.
+
+**2. Pre-populate the index:**
+
+```bash
+codeembed embed
+```
+
+Run this before starting the server to pre-populate the index. Searches will return empty results until the first file has been embedded.
+
+CodeEmbed respects your project's `.gitignore` and also excludes typical environment directories and files (`.env`, `venv`, `node_modules`, etc.) by default.
+
+**3. Start the MCP server:**
+
+```bash
+codeembed serve
+```
+
+Starts the MCP server.
+If the MCP server is added to Claude or GitHub Copilot, you do not need to do this.
+
+The `serve` command will embed your codebase in the background - by default it will scan for changes every 60 seconds.
+
+## Configuring OpenAI
+
+If you use the OpenAI provider, credentials are read from environment variables. The recommended approach is a `.env` file. `codeembed init` will ask for the path, and it will be stored in `codeembed.toml` so `codeembed serve` and `codeembed embed` loads the `.env` file automatically.
+
+### Standard OpenAI
+
+```env
+OPENAI_API_KEY=...
+```
+
+Optionally override the endpoint (for compatible APIs like vLLM, LM Studio, OpenRouter):
+
+```env
+OPENAI_API_KEY=...
+OPENAI_BASE_URL=...
+```
+
+### Azure OpenAI — API key
+
+```env
+AZURE_OPENAI_ENDPOINT=https://<your-resource>.openai.azure.com/openai/v1/
+AZURE_OPENAI_API_KEY=...
+```
+
+### Azure OpenAI — RBAC / Entra ID (keyless)
+
+Set only the endpoint; CodeEmbed will use `DefaultAzureCredential`, which automatically tries multiple credential sources in order — service principals (via env vars), workload identity, managed identity, VS Code Azure sign-in, `az login`, Azure PowerShell, and `azd auth login` — falling back to an interactive browser window if none are found automatically:
+
+```env
+AZURE_OPENAI_ENDPOINT=https://<your-resource>.openai.azure.com/openai/v1/
+```
+
+## Add to Claude Code or GitHub Copilot
+
+`codeembed init` will offer to configure these automatically. If you prefer to do it manually:
+
+**Claude Code** — add to `.mcp.json` in your project root:
+
+```json
+{
+  "mcpServers": {
+    "codeembed": {
+      "command": "codeembed",
+      "args": ["serve"]
+    }
+  }
+}
+```
+
+And add to `.claude/settings.local.json` to enable and pre-approve the tool:
+
+```json
+{
+  "enabledMcpjsonServers": ["codeembed"],
+  "permissions": {
+    "allow": ["mcp__codeembed__search"]
+  }
+}
+```
+
+**GitHub Copilot** — add to `.vscode/mcp.json`:
+
+```json
+{
+  "servers": {
+    "codeembed": {
+      "command": "codeembed",
+      "args": ["serve"]
+    }
+  }
+}
+```
+
+The MCP server exposes a single `search(query)` tool for semantic search over your codebase.
+
+## Contributing
+
+Clone this repo with:
+
+```bash
+git clone git@github.com:robino16/codeembed.git
+```
+
+```bash
+cd codeembed
+uv sync
+```
+
+Check for dependency conflicts with:
+
+```bash
+uv pip check
+```
+
+Check for package vulnerabilities with:
+
+```bash
+uv run pip-audit
+```
+
+(Optional) Add Ruff pre-commit with:
+
+```bash
+pre-commit install
+```
+
+Update init files:
+
+```bash
+uv run --no-sync scripts/generate_init_files.py
+```
+
+Run linter:
+
+```bash
+ruff check . --fix
+```
+
+Run formatter:
+
+```bash
+ruff format .
+```
+
+Run tests:
+
+```bash
+uv run --no-sync pytest
+```
+
+Build with:
+
+```bash
+uv build
+```
+
+Validate build with:
+
+```bash
+uv run twine check dist/*
+```
+
+> `--no-sync` is required for local dev commands when the MCP server is running, as uv holds a lock that blocks sync operations.