kodit 0.2.0__tar.gz → 0.2.1__tar.gz

{kodit-0.2.0 → kodit-0.2.1}/.github/workflows/docker.yaml

@@ -26,6 +26,9 @@ jobs:
           load: true
           tags: ${{ env.TEST_TAG }}
 
+      - name: Docker smoke test
+        run: ./tests/docker-smoke.sh
+
   push_to_registry:
     name: Push Docker image to registry
     runs-on: ubuntu-latest

{kodit-0.2.0 → kodit-0.2.1}/.gitignore

@@ -128,7 +128,7 @@ celerybeat.pid
 *.sage.py
 
 # Environments
-.env
+.env*
 .venv
 env/
 venv/

{kodit-0.2.0 → kodit-0.2.1}/Dockerfile

@@ -14,7 +14,8 @@ apt-get update -qy
 apt-get install -qyy \
     -o APT::Install-Recommends=false \
     -o APT::Install-Suggests=false \
-    git 
+    git \
+    build-essential
 EOT
 
 # Security-conscious organizations should package/review uv themselves.
@@ -63,6 +64,14 @@ RUN --mount=type=cache,target=/root/.cache \
 FROM python:3.13.4-slim-bookworm
 SHELL ["sh", "-exc"]
 
+RUN <<EOT
+apt-get update -qy
+apt-get install -qyy \
+    -o APT::Install-Recommends=false \
+    -o APT::Install-Suggests=false \
+    git 
+EOT
+
 ENV PATH=/app/bin:$PATH
 
 # Don't run your app as root.

{kodit-0.2.0 → kodit-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.2.0
+Version: 0.2.1
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/
@@ -18,6 +18,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Software Development :: Code Generators
 Requires-Python: >=3.12
+Requires-Dist: accelerate>=1.7.0
 Requires-Dist: aiofiles>=24.1.0
 Requires-Dist: aiosqlite>=0.20.0
 Requires-Dist: alembic>=1.15.2

{kodit-0.2.0 → kodit-0.2.1}/pyproject.toml

@@ -51,6 +51,7 @@ dependencies = [
     "tiktoken>=0.9.0",
     "asyncpg>=0.30.0",
     "transformers>=4.51.3",
+    "accelerate>=1.7.0",
 ]
 
 [dependency-groups]

{kodit-0.2.0 → kodit-0.2.1}/src/kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0'
-__version_tuple__ = version_tuple = (0, 2, 0)
+__version__ = version = '0.2.1'
+__version_tuple__ = version_tuple = (0, 2, 1)

{kodit-0.2.0 → kodit-0.2.1}/src/kodit/bm25/local_bm25.py

@@ -1,13 +1,14 @@
 """Locally hosted BM25 service primarily for use with SQLite."""
 
+from __future__ import annotations
+
 import json
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import aiofiles
-import bm25s
 import Stemmer
 import structlog
-from bm25s.tokenization import Tokenized
 
 from kodit.bm25.keyword_search_service import (
     BM25Document,
@@ -15,6 +16,11 @@ from kodit.bm25.keyword_search_service import (
     KeywordSearchProvider,
 )
 
+if TYPE_CHECKING:
+    import bm25s
+    from bm25s.tokenization import Tokenized
+
+
 SNIPPET_IDS_FILE = "snippet_ids.jsonl"
 
 
@@ -26,19 +32,28 @@ class BM25Service(KeywordSearchProvider):
         self.log = structlog.get_logger(__name__)
         self.index_path = data_dir / "bm25s_index"
         self.snippet_ids: list[int] = []
-        try:
-            self.log.debug("Loading BM25 index")
-            self.retriever = bm25s.BM25.load(self.index_path, mmap=True)
-            with Path(self.index_path / SNIPPET_IDS_FILE).open() as f:
-                self.snippet_ids = json.load(f)
-        except FileNotFoundError:
-            self.log.debug("BM25 index not found, creating new index")
-            self.retriever = bm25s.BM25()
-
         self.stemmer = Stemmer.Stemmer("english")
+        self.__retriever: bm25s.BM25 | None = None
+
+    def _retriever(self) -> bm25s.BM25:
+        """Get the BM25 retriever."""
+        if self.__retriever is None:
+            import bm25s
+
+            try:
+                self.log.debug("Loading BM25 index")
+                self.__retriever = bm25s.BM25.load(self.index_path, mmap=True)
+                with Path(self.index_path / SNIPPET_IDS_FILE).open() as f:
+                    self.snippet_ids = json.load(f)
+            except FileNotFoundError:
+                self.log.debug("BM25 index not found, creating new index")
+                self.__retriever = bm25s.BM25()
+        return self.__retriever
 
     def _tokenize(self, corpus: list[str]) -> list[list[str]] | Tokenized:
-        return bm25s.tokenize(
+        from bm25s import tokenize
+
+        return tokenize(
             corpus,
             stopwords="en",
             stemmer=self.stemmer,
@@ -50,9 +65,8 @@ class BM25Service(KeywordSearchProvider):
         """Index a new corpus."""
         self.log.debug("Indexing corpus")
         vocab = self._tokenize([doc.text for doc in corpus])
-        self.retriever = bm25s.BM25()
-        self.retriever.index(vocab, show_progress=False)
-        self.retriever.save(self.index_path)
+        self._retriever().index(vocab, show_progress=False)
+        self._retriever().save(self.index_path)
         self.snippet_ids = self.snippet_ids + [doc.snippet_id for doc in corpus]
         async with aiofiles.open(self.index_path / SNIPPET_IDS_FILE, "w") as f:
             await f.write(json.dumps(self.snippet_ids))
@@ -64,7 +78,7 @@ class BM25Service(KeywordSearchProvider):
             return []
 
         # Get the number of documents in the index
-        num_docs = self.retriever.scores["num_docs"]
+        num_docs = self._retriever().scores["num_docs"]
         if num_docs == 0:
             return []
 
@@ -80,7 +94,7 @@ class BM25Service(KeywordSearchProvider):
 
         self.log.debug("Query tokens", query_tokens=query_tokens)
 
-        results, scores = self.retriever.retrieve(
+        results, scores = self._retriever().retrieve(
             query_tokens=query_tokens,
             corpus=self.snippet_ids,
             k=top_k,

{kodit-0.2.0 → kodit-0.2.1}/src/kodit/config.py

@@ -1,16 +1,21 @@
 """Global configuration for the kodit project."""
 
+from __future__ import annotations
+
 import asyncio
-from collections.abc import Callable, Coroutine
 from functools import wraps
 from pathlib import Path
-from typing import Any, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
 import click
-from openai import AsyncOpenAI
 from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine
+
+    from openai import AsyncOpenAI
+
 from kodit.database import Database
 
 DEFAULT_BASE_DIR = Path.home() / ".kodit"
@@ -92,6 +97,8 @@ class AppContext(BaseSettings):
 
     def get_default_openai_client(self) -> AsyncOpenAI | None:
         """Get the default OpenAI client, if it is configured."""
+        from openai import AsyncOpenAI
+
         endpoint = self.default_endpoint
         if not (
             endpoint

{kodit-0.2.0 → kodit-0.2.1}/src/kodit/embedding/embedding_provider/embedding_provider.py

@@ -23,7 +23,11 @@ class EmbeddingProvider(ABC):
         """
 
 
-def split_sub_batches(encoding: tiktoken.Encoding, data: list[str]) -> list[list[str]]:
+def split_sub_batches(
+    encoding: tiktoken.Encoding,
+    data: list[str],
+    max_context_window: int = OPENAI_MAX_EMBEDDING_SIZE,
+) -> list[list[str]]:
     """Split a list of strings into smaller sub-batches."""
     log = structlog.get_logger(__name__)
     result = []
@@ -37,10 +41,10 @@ def split_sub_batches(encoding: tiktoken.Encoding, data: list[str]) -> list[list
             next_item = data_to_process[0]
             item_tokens = len(encoding.encode(next_item))
 
-            if item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+            if item_tokens > max_context_window:
                 # Loop around trying to truncate the snippet until it fits in the max
                 # embedding size
-                while item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+                while item_tokens > max_context_window:
                     next_item = next_item[:-1]
                     item_tokens = len(encoding.encode(next_item))
 
@@ -48,7 +52,7 @@ def split_sub_batches(encoding: tiktoken.Encoding, data: list[str]) -> list[list
 
                 log.warning("Truncated snippet", snippet=next_item)
 
-            if current_tokens + item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+            if current_tokens + item_tokens > max_context_window:
                 break
 
             next_batch.append(data_to_process.pop(0))

{kodit-0.2.0 → kodit-0.2.1}/src/kodit/embedding/embedding_provider/local_embedding_provider.py

@@ -1,10 +1,12 @@
 """Local embedding service."""
 
+from __future__ import annotations
+
 import os
+from typing import TYPE_CHECKING
 
 import structlog
 import tiktoken
-from sentence_transformers import SentenceTransformer
 from tqdm import tqdm
 
 from kodit.embedding.embedding_provider.embedding_provider import (
@@ -13,6 +15,9 @@ from kodit.embedding.embedding_provider.embedding_provider import (
     split_sub_batches,
 )
 
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
 TINY = "tiny"
 CODE = "code"
 TEST = "test"
@@ -38,10 +43,11 @@ class LocalEmbeddingProvider(EmbeddingProvider):
         """Get the embedding model."""
         if self.embedding_model is None:
             os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+            from sentence_transformers import SentenceTransformer
+
             self.embedding_model = SentenceTransformer(
                 self.model_name,
                 trust_remote_code=True,
-                device="cpu",  # Force CPU so we don't have to install accelerate, etc.
             )
         return self.embedding_model
 

kodit-0.2.1/src/kodit/enrichment/enrichment_provider/local_enrichment_provider.py

@@ -0,0 +1,88 @@
+"""Local embedding service."""
+
+import os
+
+import structlog
+import tiktoken
+from tqdm import tqdm
+
+from kodit.embedding.embedding_provider.embedding_provider import split_sub_batches
+from kodit.enrichment.enrichment_provider.enrichment_provider import (
+    ENRICHMENT_SYSTEM_PROMPT,
+    EnrichmentProvider,
+)
+
+DEFAULT_ENRICHMENT_MODEL = "Qwen/Qwen3-0.6B"
+DEFAULT_CONTEXT_WINDOW_SIZE = 2048  # Small so it works even on low-powered devices
+
+
+class LocalEnrichmentProvider(EnrichmentProvider):
+    """Local embedder."""
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_ENRICHMENT_MODEL,
+        context_window: int = DEFAULT_CONTEXT_WINDOW_SIZE,
+    ) -> None:
+        """Initialize the local enrichment provider."""
+        self.log = structlog.get_logger(__name__)
+        self.model_name = model_name
+        self.context_window = context_window
+        self.model = None
+        self.tokenizer = None
+        self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
+
+    async def enrich(self, data: list[str]) -> list[str]:
+        """Enrich a list of strings."""
+        from transformers.models.auto.modeling_auto import (
+            AutoModelForCausalLM,
+        )
+        from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+        if self.tokenizer is None:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_name, padding_side="left"
+            )
+        if self.model is None:
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype="auto",
+                trust_remote_code=True,
+                device_map="auto",
+            )
+
+        # Prepare prompts
+        prompts = [
+            self.tokenizer.apply_chat_template(
+                [
+                    {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
+                    {"role": "user", "content": snippet},
+                ],
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=False,
+            )
+            for snippet in data
+        ]
+
+        # Batch prompts using split_sub_batches
+        batched_prompts = split_sub_batches(
+            self.encoding, prompts, max_context_window=self.context_window
+        )
+        results = []
+        for batch in tqdm(batched_prompts, leave=False, total=len(batched_prompts)):
+            model_inputs = self.tokenizer(
+                batch, return_tensors="pt", padding=True, truncation=True
+            ).to(self.model.device)
+            generated_ids = self.model.generate(
+                **model_inputs, max_new_tokens=self.context_window
+            )
+            # For each prompt in the batch, decode only the generated part
+            for i, input_ids in enumerate(model_inputs["input_ids"]):
+                output_ids = generated_ids[i][len(input_ids) :].tolist()
+                content = self.tokenizer.decode(
+                    output_ids, skip_special_tokens=True
+                ).strip("\n")
+                results.append(content)
+        return results

{kodit-0.2.0 → kodit-0.2.1}/src/kodit/source/source_service.py

@@ -82,17 +82,7 @@ class SourceService:
         )
 
     async def create(self, uri_or_path_like: str) -> SourceView:
-        """Create a new source from a URI.
-
-        Args:
-            uri: The URI of the source to create. Can be a git-like URI or a local
-                directory.
-
-        Raises:
-            ValueError: If the source type is not supported or if the folder doesn't
-                exist.
-
-        """
+        """Create a new source from a URI or path."""
         if Path(uri_or_path_like).is_dir():
             return await self._create_folder_source(Path(uri_or_path_like))
         if isuri(uri_or_path_like):
@@ -103,18 +93,11 @@ class SourceService:
                 ".git"
             ):
                 return await self._create_git_source(uri_or_path_like)
-
-            # Try adding a .git suffix, sometimes people just pass the url
             if not uri_or_path_like.endswith(".git"):
-                uri_or_path_like = uri_or_path_like + ".git"
-                try:
-                    return await self._create_git_source(uri_or_path_like)
-                except git.GitCommandError:
-                    raise
-                except ValueError:
-                    pass
-
-        msg = f"Unsupported source type: {uri_or_path_like}"
+                uri_or_path_like = uri_or_path_like.strip("/") + ".git"
+                return await self._create_git_source(uri_or_path_like)
+
+        msg = f"Unsupported source: {uri_or_path_like}"
         raise ValueError(msg)
 
     async def _create_folder_source(self, directory: Path) -> SourceView:

kodit-0.2.1/tests/docker-smoke.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+set -e
+
+if [ -z "$TEST_TAG" ]; then
+    echo "TEST_TAG is not set"
+    exit 1
+fi
+
+# Get the directory of this script
+script_dir=$(dirname "$0")
+
+# Start the container, mount the smoke test and run it
+docker run -i -v $script_dir:/tests --entrypoint /bin/bash --env CI=true $TEST_TAG -c "/tests/smoke.sh"

{kodit-0.2.0 → kodit-0.2.1}/tests/kodit/bm25/local_bm25_test.py

@@ -34,14 +34,6 @@ def sample_documents():
     ]
 
 
-@pytest.mark.asyncio
-async def test_initialization(bm25_service):
-    """Test that the service initializes correctly."""
-    assert bm25_service.retriever is not None
-    assert isinstance(bm25_service.snippet_ids, list)
-    assert len(bm25_service.snippet_ids) == 0
-
-
 @pytest.mark.asyncio
 async def test_index_and_retrieve(bm25_service, sample_documents):
     """Test indexing and retrieving documents."""
@@ -134,6 +126,7 @@ async def test_persistence(bm25_service, sample_documents):
 
     # Create a new service instance with the same data directory
     new_service = BM25Service(bm25_service.index_path.parent)
+    new_service._retriever()
 
     # Verify the new instance loaded the correct data
     assert len(new_service.snippet_ids) == 3

{kodit-0.2.0 → kodit-0.2.1}/tests/kodit/source/source_service_test.py

@@ -32,7 +32,7 @@ async def test_create_source_nonexistent_path(service: SourceService) -> None:
     uri = nonexistent_path.as_uri()
 
     # Try to create a source with the nonexistent path
-    with pytest.raises(ValueError, match=f"Folder does not exist: {nonexistent_path}"):
+    with pytest.raises(ValueError):
         await service.create(uri)
 
 
@@ -41,7 +41,7 @@ async def test_create_source_invalid_path_and_uri(service: SourceService) -> Non
     """Test creating a source with an invalid path that is also not a valid URI."""
     # Try to create a source with an invalid path that is also not a valid URI
     invalid_path = "not/a/valid/path/or/uri"
-    with pytest.raises(ValueError, match=f"Unsupported source type: {invalid_path}"):
+    with pytest.raises(ValueError):
         await service.create(invalid_path)
 
 
@@ -65,7 +65,7 @@ async def test_create_source_already_added(
 async def test_create_source_unsupported_uri(service: SourceService) -> None:
     """Test creating a source with an unsupported URI."""
     # Try to create a source with an unsupported URI (e.g., http)
-    with pytest.raises(ValueError, match="Unsupported source type: http://example.com"):
+    with pytest.raises(ValueError):
         await service.create("http://example.com")
 
 
@@ -155,3 +155,11 @@ async def test_create_source_relative_path(
 
     # Should not raise an error
     await service.create(".")
+
+
+@pytest.mark.asyncio
+async def test_strip_trailing_slash_on_gh_url(service: SourceService) -> None:
+    """Test creating a source with a file URI."""
+
+    # Should work
+    await service.create("https://github.com/helixml/kodit/")

{kodit-0.2.0 → kodit-0.2.1}/uv.lock

@@ -6,6 +6,24 @@ resolution-markers = [
     "python_full_version < '3.13'",
 ]
 
+[[package]]
+name = "accelerate"
+version = "1.7.0"
+source = { registry = "https://pypi.org/simple/" }
+dependencies = [
+    { name = "huggingface-hub" },
+    { name = "numpy" },
+    { name = "packaging" },
+    { name = "psutil" },
+    { name = "pyyaml" },
+    { name = "safetensors" },
+    { name = "torch" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/97/33/47bbd507e3a851d33d19ce7b2141c5ea3689bfae91ba168044d7db24b0e9/accelerate-1.7.0.tar.gz", hash = "sha256:e8a2a5503d6237b9eee73cc8d36cf543f9c2d8dd2c6713450b322f5e6d53a610", size = 376026, upload-time = "2025-05-15T10:00:52.117Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f8/bb/be8146c196ad6e4dec78385d91e92591f8a433576c4e04c342a636fcd811/accelerate-1.7.0-py3-none-any.whl", hash = "sha256:cf57165cca28769c6cf2650812371c81b18e05743dfa3c748524b1bb4f2b272f", size = 362095, upload-time = "2025-05-15T10:00:49.914Z" },
+]
+
 [[package]]
 name = "aiofiles"
 version = "24.1.0"
@@ -840,6 +858,7 @@ wheels = [
 name = "kodit"
 source = { editable = "." }
 dependencies = [
+    { name = "accelerate" },
     { name = "aiofiles" },
     { name = "aiosqlite" },
     { name = "alembic" },
@@ -884,6 +903,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
+    { name = "accelerate", specifier = ">=1.7.0" },
     { name = "aiofiles", specifier = ">=24.1.0" },
     { name = "aiosqlite", specifier = ">=0.20.0" },
     { name = "alembic", specifier = ">=1.15.2" },

kodit-0.2.0/src/kodit/enrichment/enrichment_provider/local_enrichment_provider.py

@@ -1,63 +0,0 @@
-"""Local embedding service."""
-
-import os
-
-import structlog
-from transformers.models.auto.modeling_auto import AutoModelForCausalLM
-from transformers.models.auto.tokenization_auto import AutoTokenizer
-
-from kodit.enrichment.enrichment_provider.enrichment_provider import (
-    ENRICHMENT_SYSTEM_PROMPT,
-    EnrichmentProvider,
-)
-
-
-class LocalEnrichmentProvider(EnrichmentProvider):
-    """Local embedder."""
-
-    def __init__(self, model_name: str = "Qwen/Qwen3-0.6B") -> None:
-        """Initialize the local enrichment provider."""
-        self.log = structlog.get_logger(__name__)
-        self.model_name = model_name
-        self.model = None
-        self.tokenizer = None
-
-    async def enrich(self, data: list[str]) -> list[str]:
-        """Enrich a list of strings."""
-        if self.tokenizer is None:
-            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        if self.model is None:
-            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                torch_dtype="auto",
-                trust_remote_code=True,
-            )
-
-        results = []
-        for snippet in data:
-            # prepare the model input
-            messages = [
-                {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
-                {"role": "user", "content": snippet},
-            ]
-            text = self.tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True,
-                enable_thinking=False,
-            )
-            model_inputs = self.tokenizer([text], return_tensors="pt").to(
-                self.model.device
-            )
-
-            # conduct text completion
-            generated_ids = self.model.generate(**model_inputs, max_new_tokens=32768)
-            output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()
-            content = self.tokenizer.decode(output_ids, skip_special_tokens=True).strip(
-                "\n"
-            )
-
-            results.append(content)
-
-        return results