kodit 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

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/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/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/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/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/enrichment/enrichment_provider/local_enrichment_provider.py

@@ -3,61 +3,86 @@
 import os
 
 import structlog
-from transformers.models.auto.modeling_auto import AutoModelForCausalLM
-from transformers.models.auto.tokenization_auto import AutoTokenizer
+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 = "Qwen/Qwen3-0.6B") -> None:
+    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)
+            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",
             )
 
-        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,
+        # 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,
             )
-            model_inputs = self.tokenizer([text], return_tensors="pt").to(
-                self.model.device
-            )
+            for snippet in data
+        ]
 
-            # 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"
+        # 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
             )
-
-            results.append(content)
-
+            # 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/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.0.dist-info → kodit-0.2.1.dist-info}/METADATA

@@ -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.dist-info → kodit-0.2.1.dist-info}/RECORD

@@ -1,9 +1,9 @@
 kodit/.gitignore,sha256=ztkjgRwL9Uud1OEi36hGQeDGk3OLK1NfDEO8YqGYy8o,11
 kodit/__init__.py,sha256=aEKHYninUq1yh6jaNfvJBYg-6fenpN132nJt1UU6Jxs,59
-kodit/_version.py,sha256=iB5DfB5V6YB5Wo4JmvS-txT42QtmGaWcWp3udRT7zCI,511
+kodit/_version.py,sha256=UoNvMtd4wCG76RwoSpNCUtaFyTwakGcZolfjXzNVSMY,511
 kodit/app.py,sha256=Mr5BFHOHx5zppwjC4XPWVvHjwgl1yrKbUjTWXKubJQM,891
 kodit/cli.py,sha256=i7eEt0FdIQGEfXKFte-8fBcZZGE8BPXBp40aGwJDQGI,11323
-kodit/config.py,sha256=2W2u5J8j-Mbt-C4xzOuK-PeuDCx0S_rnCXPhBwvfLT4,4353
+kodit/config.py,sha256=KAxs6qdIvgRZEGTdDIl30QC0g56U0mdk-W8VlaVzghs,4472
 kodit/database.py,sha256=WB1KpVxUYPgiJGU0gJa2hqytYB8wJEJ5z3WayhWzNMU,2403
 kodit/log.py,sha256=HU1OmuxO4FcVw61k4WW7Y4WM7BrDaeplw1PcBHhuIZY,5434
 kodit/mcp.py,sha256=QruyPskWB0_x59pkfj5BBeXuR13GMny5TAZEa2j4U9s,5752
@@ -11,7 +11,7 @@ kodit/middleware.py,sha256=I6FOkqG9-8RH5kR1-0ZoQWfE4qLCB8lZYv8H_OCH29o,2714
 kodit/bm25/__init__.py,sha256=j8zyriNWhbwE5Lbybzg1hQAhANlU9mKHWw4beeUR6og,19
 kodit/bm25/keyword_search_factory.py,sha256=rp-wx3DJsc2KlELK1V337EyeYvmwnMQwUqOo1WVPSmg,631
 kodit/bm25/keyword_search_service.py,sha256=aBbWQKgQmi2re3EIHdXFS00n7Wj3b2D0pZsLZ4qmHfE,754
-kodit/bm25/local_bm25.py,sha256=AAbFhbQDqyL3d7jsPL7W4HsLxdoYctaDsREUXOLy6jM,3260
+kodit/bm25/local_bm25.py,sha256=cDx_hT9hXyEWz8LlFGyQs5-0ufK2Y4gArMfZv7-D9zQ,3621
 kodit/bm25/vectorchord_bm25.py,sha256=_nGrkUReYLLV-L8RIuIVLwjuhSYZl9T532n5OVf0kWs,6393
 kodit/embedding/__init__.py,sha256=h9NXzDA1r-K23nvBajBV-RJzHJN0p3UJ7UQsmdnOoRw,24
 kodit/embedding/embedding_factory.py,sha256=UGnFRyyQXazSUOwyW4Hg7Vq2-kfAoDj9lD4CTLu8x04,1630
@@ -21,16 +21,16 @@ kodit/embedding/local_vector_search_service.py,sha256=hkF0qlfzjyGt400qIX9Mr6B7b7
 kodit/embedding/vector_search_service.py,sha256=pQJ129QjGrAWOXzqkywmgtDRpy8_gtzYgkivyqF9Vrs,1009
 kodit/embedding/vectorchord_vector_search_service.py,sha256=63Xf7_nAz3xWOwrmZibw8Q-xoRdCrPDDpdSA_WE7mrc,5131
 kodit/embedding/embedding_provider/__init__.py,sha256=h9NXzDA1r-K23nvBajBV-RJzHJN0p3UJ7UQsmdnOoRw,24
-kodit/embedding/embedding_provider/embedding_provider.py,sha256=Tf3bwUsUMzAgoyLFM5qBtOLqPp1qr03TzrwGczkDvy0,1835
+kodit/embedding/embedding_provider/embedding_provider.py,sha256=IC7fZaZ_ze-DxpxKfK44pRDwHWUQhVIqVKKQ3alO5Qc,1882
 kodit/embedding/embedding_provider/hash_embedding_provider.py,sha256=nAhlhh8j8PqqCCbhVl26Y8ntFBm2vJBCtB4X04g5Wwg,2638
-kodit/embedding/embedding_provider/local_embedding_provider.py,sha256=4ER-UPq506Y0TWU6qcs0nUqw6bSKQkSrdog-DhNQWM8,1906
+kodit/embedding/embedding_provider/local_embedding_provider.py,sha256=WP8lw6XG7v1_5Mw4_rhIOETooYRsxhkwmFaXCqCouQU,1977
 kodit/embedding/embedding_provider/openai_embedding_provider.py,sha256=V_jdUXiaGdslplwxMlfgFc4_hAVS2eaJXMTs2C7RiLI,2666
 kodit/enrichment/__init__.py,sha256=vBEolHpKaHUhfINX0dSGyAPlvgpLNAer9YzFtdvCB24,18
 kodit/enrichment/enrichment_factory.py,sha256=vKjkUTdhj74IW2S4GENDWdWMJx6BwUSZjJGDC0i7DSk,787
 kodit/enrichment/enrichment_service.py,sha256=87Sd3gGbEMJYb_wVrHG8L1yGIZmQNR7foUS4_y94azI,977
 kodit/enrichment/enrichment_provider/__init__.py,sha256=klf8iuLVWX4iRz-DZQauFFNAoJC5CByczh48TBZPW-o,27
 kodit/enrichment/enrichment_provider/enrichment_provider.py,sha256=E0H5rq3OENM0yYbA8K_3nSnj5lUHCpoIOqpWLo-2MVU,413
-kodit/enrichment/enrichment_provider/local_enrichment_provider.py,sha256=bR6HR1gH7wtZdMLOwaKdASjvllRo1FlNW9GyZC11zAM,2164
+kodit/enrichment/enrichment_provider/local_enrichment_provider.py,sha256=kh_9X9m-WEziRi5TV6QflKXyXQTos9kzpmpGil7pywM,3196
 kodit/enrichment/enrichment_provider/openai_enrichment_provider.py,sha256=gYuFTAeIVdQNlCUvNSPgRoiRwCvRD0C8419h8ubyABA,2725
 kodit/indexing/__init__.py,sha256=cPyi2Iej3G1JFWlWr7X80_UrsMaTu5W5rBwgif1B3xo,75
 kodit/indexing/fusion.py,sha256=TZb4fPAedXdEUXzwzOofW98QIOymdbclBOP1KOijuEk,1674
@@ -57,11 +57,11 @@ kodit/snippets/languages/typescript.scm,sha256=U-ujbbv4tylbUBj9wuhL-e5cW6hmgPCNs
 kodit/source/__init__.py,sha256=1NTZyPdjThVQpZO1Mp1ColVsS7sqYanOVLqnoqV9Ipo,83
 kodit/source/source_models.py,sha256=kcC59XPSDDMth2mOYK3FakqTN0jxKFaTDch0ejyD9Sw,2446
 kodit/source/source_repository.py,sha256=0EksMpoLzdkfe8S4eeCm4Sf7TuxsOzOzaF4BBsMYo-4,3163
-kodit/source/source_service.py,sha256=u_GaH07ewakThQJRfT8O_yZ54A52qLtJuM1bF3xUT2A,9633
+kodit/source/source_service.py,sha256=dyXWf_t2qGvD9YAY3C5Zcxc63BbyrheL-jgTzXV7gqo,9156
 kodit/util/__init__.py,sha256=bPu6CtqDWCRGU7VgW2_aiQrCBi8G89FS6k1PjvDajJ0,37
 kodit/util/spinner.py,sha256=R9bzrHtBiIH6IfLbmsIVHL53s8vg-tqW4lwGGALu4dw,1932
-kodit-0.2.0.dist-info/METADATA,sha256=0CdegivoI9rcZLpmwzGTFfW_bui1D1tjNtz7ajXFOJk,5735
-kodit-0.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-kodit-0.2.0.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
-kodit-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-kodit-0.2.0.dist-info/RECORD,,
+kodit-0.2.1.dist-info/METADATA,sha256=Frs10m-Bc3DzgvxXSA27u0dFXDdV0fR_I4zzyhhPYfY,5768
+kodit-0.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+kodit-0.2.1.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
+kodit-0.2.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+kodit-0.2.1.dist-info/RECORD,,