kodit 0.1.10__py3-none-any.whl → 0.1.12__py3-none-any.whl

kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.10'
-__version_tuple__ = version_tuple = (0, 1, 10)
+__version__ = version = '0.1.12'
+__version_tuple__ = version_tuple = (0, 1, 12)

kodit/bm25/bm25.py

@@ -52,7 +52,7 @@ class BM25Service:
             self.log.warning("No documents to retrieve from, returning empty list")
             return []
 
-        top_k = min(top_k, len(doc_ids))
+        top_k = min(top_k, len(self.retriever.scores))
         self.log.debug(
             "Retrieving from index", query=query, top_k=top_k, num_docs=len(doc_ids)
         )

kodit/cli.py

@@ -12,35 +12,21 @@ from pytable_formatter import Cell, Table
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.config import (
-    DEFAULT_BASE_DIR,
-    DEFAULT_DB_URL,
-    DEFAULT_DISABLE_TELEMETRY,
-    DEFAULT_EMBEDDING_MODEL_NAME,
-    DEFAULT_LOG_FORMAT,
-    DEFAULT_LOG_LEVEL,
     AppContext,
     with_app_context,
     with_session,
 )
-from kodit.indexing.repository import IndexRepository
-from kodit.indexing.service import IndexService
+from kodit.embedding.embedding import embedding_factory
+from kodit.indexing.indexing_repository import IndexRepository
+from kodit.indexing.indexing_service import IndexService
 from kodit.log import configure_logging, configure_telemetry, log_event
-from kodit.retreival.repository import RetrievalRepository
-from kodit.retreival.service import RetrievalRequest, RetrievalService
-from kodit.sources.repository import SourceRepository
-from kodit.sources.service import SourceService
+from kodit.search.search_repository import SearchRepository
+from kodit.search.search_service import SearchRequest, SearchService
+from kodit.source.source_repository import SourceRepository
+from kodit.source.source_service import SourceService
 
 
 @click.group(context_settings={"max_content_width": 100})
-@click.option("--log-level", help=f"Log level [default: {DEFAULT_LOG_LEVEL}]")
-@click.option("--log-format", help=f"Log format [default: {DEFAULT_LOG_FORMAT}]")
-@click.option(
-    "--disable-telemetry",
-    is_flag=True,
-    help=f"Disable telemetry [default: {DEFAULT_DISABLE_TELEMETRY}]",
-)
-@click.option("--db-url", help=f"Database URL [default: {DEFAULT_DB_URL}]")
-@click.option("--data-dir", help=f"Data directory [default: {DEFAULT_BASE_DIR}]")
 @click.option(
     "--env-file",
     help="Path to a .env file [default: .env]",
@@ -52,13 +38,8 @@ from kodit.sources.service import SourceService
     ),
 )
 @click.pass_context
-def cli(  # noqa: PLR0913
+def cli(
     ctx: click.Context,
-    log_level: str | None,
-    log_format: str | None,
-    disable_telemetry: bool | None,
-    db_url: str | None,
-    data_dir: str | None,
     env_file: Path | None,
 ) -> None:
     """kodit CLI - Code indexing for better AI code generation."""  # noqa: D403
@@ -67,17 +48,6 @@ def cli(  # noqa: PLR0913
     if env_file:
         config = AppContext(_env_file=env_file)  # type: ignore[reportCallIssue]
 
-    # Now override with CLI arguments, if set
-    if data_dir:
-        config.data_dir = Path(data_dir)
-    if db_url:
-        config.db_url = db_url
-    if log_level:
-        config.log_level = log_level
-    if log_format:
-        config.log_format = log_format
-    if disable_telemetry:
-        config.disable_telemetry = disable_telemetry
     configure_logging(config)
     configure_telemetry(config)
 
@@ -102,7 +72,7 @@ async def index(
         repository,
         source_service,
         app_context.get_data_dir(),
-        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+        embedding_service=embedding_factory(app_context.get_default_openai_client()),
     )
 
     if not sources:
@@ -159,14 +129,14 @@ async def code(
 
     This works best if your query is code.
     """
-    repository = RetrievalRepository(session)
-    service = RetrievalService(
+    repository = SearchRepository(session)
+    service = SearchService(
         repository,
         app_context.get_data_dir(),
-        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+        embedding_service=embedding_factory(app_context.get_default_openai_client()),
     )
 
-    snippets = await service.retrieve(RetrievalRequest(code_query=query, top_k=top_k))
+    snippets = await service.search(SearchRequest(code_query=query, top_k=top_k))
 
     if len(snippets) == 0:
         click.echo("No snippets found")
@@ -192,14 +162,14 @@ async def keyword(
     top_k: int,
 ) -> None:
     """Search for snippets using keyword search."""
-    repository = RetrievalRepository(session)
-    service = RetrievalService(
+    repository = SearchRepository(session)
+    service = SearchService(
         repository,
         app_context.get_data_dir(),
-        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+        embedding_service=embedding_factory(app_context.get_default_openai_client()),
     )
 
-    snippets = await service.retrieve(RetrievalRequest(keywords=keywords, top_k=top_k))
+    snippets = await service.search(SearchRequest(keywords=keywords, top_k=top_k))
 
     if len(snippets) == 0:
         click.echo("No snippets found")
@@ -227,18 +197,18 @@ async def hybrid(
     code: str,
 ) -> None:
     """Search for snippets using hybrid search."""
-    repository = RetrievalRepository(session)
-    service = RetrievalService(
+    repository = SearchRepository(session)
+    service = SearchService(
         repository,
         app_context.get_data_dir(),
-        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
+        embedding_service=embedding_factory(app_context.get_default_openai_client()),
     )
 
     # Parse keywords into a list of strings
     keywords_list = [k.strip().lower() for k in keywords.split(",")]
 
-    snippets = await service.retrieve(
-        RetrievalRequest(keywords=keywords_list, code_query=code, top_k=top_k)
+    snippets = await service.search(
+        SearchRequest(keywords=keywords_list, code_query=code, top_k=top_k)
     )
 
     if len(snippets) == 0:

kodit/config.py

@@ -4,10 +4,11 @@ import asyncio
 from collections.abc import Callable, Coroutine
 from functools import wraps
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import Any, Literal, TypeVar
 
 import click
-from pydantic import Field
+from openai import AsyncOpenAI
+from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from kodit.database import Database
@@ -22,16 +23,40 @@ DEFAULT_EMBEDDING_MODEL_NAME = TINY
 T = TypeVar("T")
 
 
+class Endpoint(BaseModel):
+    """Endpoint provides configuration for an AI service."""
+
+    type: Literal["openai"] = Field(default="openai")
+    api_key: str | None = None
+    base_url: str | None = None
+
+
 class AppContext(BaseSettings):
     """Global context for the kodit project. Provides a shared state for the app."""
 
-    model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        env_nested_delimiter="_",
+        nested_model_default_partial_update=True,
+        env_nested_max_split=1,
+    )
 
     data_dir: Path = Field(default=DEFAULT_BASE_DIR)
     db_url: str = Field(default=DEFAULT_DB_URL)
     log_level: str = Field(default=DEFAULT_LOG_LEVEL)
     log_format: str = Field(default=DEFAULT_LOG_FORMAT)
     disable_telemetry: bool = Field(default=DEFAULT_DISABLE_TELEMETRY)
+    default_endpoint: Endpoint | None = Field(
+        default=Endpoint(
+            type="openai",
+            base_url="https://api.openai.com/v1",
+        ),
+        description=(
+            "Default endpoint to use for all AI interactions "
+            "(can be overridden by task-specific configuration)."
+        ),
+    )
     _db: Database | None = None
 
     def model_post_init(self, _: Any) -> None:
@@ -58,6 +83,21 @@ class AppContext(BaseSettings):
             await self._db.run_migrations(self.db_url)
         return self._db
 
+    def get_default_openai_client(self) -> AsyncOpenAI | None:
+        """Get the default OpenAI client, if it is configured."""
+        endpoint = self.default_endpoint
+        if not (
+            endpoint
+            and endpoint.type == "openai"
+            and endpoint.api_key
+            and endpoint.base_url
+        ):
+            return None
+        return AsyncOpenAI(
+            api_key=endpoint.api_key,
+            base_url=endpoint.base_url,
+        )
+
 
 with_app_context = click.make_pass_decorator(AppContext)
 

kodit/embedding/embedding.py

@@ -1,9 +1,14 @@
 """Embedding service."""
 
+import asyncio
 import os
-from collections.abc import Generator
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from typing import NamedTuple
 
 import structlog
+import tiktoken
+from openai import AsyncOpenAI
 from sentence_transformers import SentenceTransformer
 
 TINY = "tiny"
@@ -17,14 +22,59 @@ COMMON_EMBEDDING_MODELS = {
 }
 
 
-class EmbeddingService:
-    """Service for embeddings."""
+class EmbeddingInput(NamedTuple):
+    """Input for embedding."""
+
+    id: int
+    text: str
+
+
+class EmbeddingOutput(NamedTuple):
+    """Output for embedding."""
+
+    id: int
+    embedding: list[float]
+
+
+class Embedder(ABC):
+    """Embedder interface."""
+
+    @abstractmethod
+    def embed(
+        self, data: list[EmbeddingInput]
+    ) -> AsyncGenerator[EmbeddingOutput, None]:
+        """Embed a list of documents.
+
+        The embedding service accepts a massive list of id,strings to embed. Behind the
+        scenes it batches up requests and parallelizes them for performance according to
+        the specifics of the embedding service.
+
+        The id reference is required because the parallelization may return results out
+        of order.
+        """
+
+    @abstractmethod
+    def query(self, data: list[str]) -> AsyncGenerator[list[float], None]:
+        """Query the embedding model."""
+
+
+def embedding_factory(openai_client: AsyncOpenAI | None = None) -> Embedder:
+    """Create an embedding service."""
+    if openai_client is not None:
+        return OpenAIEmbedder(openai_client)
+    return LocalEmbedder(model_name=TINY)
+
+
+class LocalEmbedder(Embedder):
+    """Local embedder."""
 
     def __init__(self, model_name: str) -> None:
-        """Initialize the embedding service."""
+        """Initialize the local embedder."""
         self.log = structlog.get_logger(__name__)
+        self.log.info("Creating local embedder", model_name=model_name)
         self.model_name = COMMON_EMBEDDING_MODELS.get(model_name, model_name)
         self.embedding_model = None
+        self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
 
     def _model(self) -> SentenceTransformer:
         """Get the embedding model."""
@@ -37,16 +87,117 @@ class EmbeddingService:
             )
         return self.embedding_model
 
-    def embed(self, snippets: list[str]) -> Generator[list[float], None, None]:
+    async def embed(
+        self, data: list[EmbeddingInput]
+    ) -> AsyncGenerator[EmbeddingOutput, None]:
         """Embed a list of documents."""
         model = self._model()
-        embeddings = model.encode(snippets, show_progress_bar=False, batch_size=4)
-        for embedding in embeddings:
-            yield [float(x) for x in embedding]
 
-    def query(self, query: list[str]) -> Generator[list[float], None, None]:
+        batched_data = _split_sub_batches(self.encoding, data)
+
+        for batch in batched_data:
+            embeddings = model.encode(
+                [i.text for i in batch], show_progress_bar=False, batch_size=4
+            )
+            for i, x in zip(batch, embeddings, strict=False):
+                yield EmbeddingOutput(i.id, [float(y) for y in x])
+
+    async def query(self, data: list[str]) -> AsyncGenerator[list[float], None]:
         """Query the embedding model."""
         model = self._model()
-        embeddings = model.encode(query, show_progress_bar=False, batch_size=4)
+        embeddings = model.encode(data, show_progress_bar=False, batch_size=4)
         for embedding in embeddings:
             yield [float(x) for x in embedding]
+
+
+OPENAI_MAX_EMBEDDING_SIZE = 8192
+OPENAI_NUM_PARALLEL_TASKS = 10
+
+
+def _split_sub_batches(
+    encoding: tiktoken.Encoding, data: list[EmbeddingInput]
+) -> list[list[EmbeddingInput]]:
+    """Split a list of strings into smaller sub-batches."""
+    log = structlog.get_logger(__name__)
+    result = []
+    data_to_process = [s for s in data if s.text.strip()]  # Filter out empty strings
+
+    while data_to_process:
+        next_batch = []
+        current_tokens = 0
+
+        while data_to_process:
+            next_item = data_to_process[0]
+            item_tokens = len(encoding.encode(next_item.text))
+
+            if item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+                log.warning("Skipping too long snippet", snippet=data_to_process.pop(0))
+                continue
+
+            if current_tokens + item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
+                break
+
+            next_batch.append(data_to_process.pop(0))
+            current_tokens += item_tokens
+
+        if next_batch:
+            result.append(next_batch)
+
+    return result
+
+
+class OpenAIEmbedder(Embedder):
+    """OpenAI embedder."""
+
+    def __init__(
+        self, openai_client: AsyncOpenAI, model_name: str = "text-embedding-3-small"
+    ) -> None:
+        """Initialize the OpenAI embedder."""
+        self.log = structlog.get_logger(__name__)
+        self.log.info("Creating OpenAI embedder", model_name=model_name)
+        self.openai_client = openai_client
+        self.encoding = tiktoken.encoding_for_model(model_name)
+        self.log = structlog.get_logger(__name__)
+
+    async def embed(
+        self,
+        data: list[EmbeddingInput],
+    ) -> AsyncGenerator[EmbeddingOutput, None]:
+        """Embed a list of documents."""
+        # First split the list into a list of list where each sublist has fewer than
+        # max tokens.
+        batched_data = _split_sub_batches(self.encoding, data)
+
+        # Process batches in parallel with a semaphore to limit concurrent requests
+        sem = asyncio.Semaphore(OPENAI_NUM_PARALLEL_TASKS)
+
+        async def process_batch(batch: list[EmbeddingInput]) -> list[EmbeddingOutput]:
+            async with sem:
+                try:
+                    response = await self.openai_client.embeddings.create(
+                        model="text-embedding-3-small",
+                        input=[i.text for i in batch],
+                    )
+                    return [
+                        EmbeddingOutput(i.id, x.embedding)
+                        for i, x in zip(batch, response.data, strict=False)
+                    ]
+                except Exception as e:
+                    self.log.exception("Error embedding batch", error=str(e))
+                    return []
+
+        # Create tasks for all batches
+        tasks = [process_batch(batch) for batch in batched_data]
+
+        # Process all batches and yield results as they complete
+        for task in asyncio.as_completed(tasks):
+            embeddings = await task
+            for e in embeddings:
+                yield e
+
+    async def query(self, data: list[str]) -> AsyncGenerator[list[float], None]:
+        """Query the embedding model."""
+        async for e in self.embed(
+            [EmbeddingInput(i, text) for i, text in enumerate(data)]
+        ):
+            yield e.embedding

kodit/indexing/{models.py → indexing_models.py}

@@ -31,8 +31,8 @@ class Snippet(Base, CommonMixin):
 
     __tablename__ = "snippets"
 
-    file_id: Mapped[int] = mapped_column(ForeignKey("files.id"))
-    index_id: Mapped[int] = mapped_column(ForeignKey("indexes.id"))
+    file_id: Mapped[int] = mapped_column(ForeignKey("files.id"), index=True)
+    index_id: Mapped[int] = mapped_column(ForeignKey("indexes.id"), index=True)
     content: Mapped[str] = mapped_column(UnicodeText, default="")
 
     def __init__(self, file_id: int, index_id: int, content: str) -> None:

kodit/indexing/{repository.py → indexing_repository.py}

@@ -11,9 +11,9 @@ from typing import TypeVar
 from sqlalchemy import delete, func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.embedding.models import Embedding
-from kodit.indexing.models import Index, Snippet
-from kodit.sources.models import File, Source
+from kodit.embedding.embedding_models import Embedding
+from kodit.indexing.indexing_models import Index, Snippet
+from kodit.source.source_models import File, Source
 
 T = TypeVar("T")
 
@@ -156,14 +156,14 @@ class IndexRepository:
         result = await self.session.execute(query)
         return list(result.scalars())
 
-    async def get_all_snippets(self) -> list[Snippet]:
+    async def get_all_snippets(self, index_id: int) -> list[Snippet]:
         """Get all snippets.
 
         Returns:
             A list of all snippets.
 
         """
-        query = select(Snippet).order_by(Snippet.id)
+        query = select(Snippet).where(Snippet.index_id == index_id).order_by(Snippet.id)
         result = await self.session.execute(query)
         return list(result.scalars())
 

kodit/indexing/{service.py → indexing_service.py}

@@ -14,12 +14,12 @@ import structlog
 from tqdm.asyncio import tqdm
 
 from kodit.bm25.bm25 import BM25Service
-from kodit.embedding.embedding import EmbeddingService
-from kodit.embedding.models import Embedding, EmbeddingType
-from kodit.indexing.models import Snippet
-from kodit.indexing.repository import IndexRepository
+from kodit.embedding.embedding import Embedder, EmbeddingInput
+from kodit.embedding.embedding_models import Embedding, EmbeddingType
+from kodit.indexing.indexing_models import Snippet
+from kodit.indexing.indexing_repository import IndexRepository
 from kodit.snippets.snippets import SnippetService
-from kodit.sources.service import SourceService
+from kodit.source.source_service import SourceService
 
 # List of MIME types that are blacklisted from being indexed
 MIME_BLACKLIST = ["unknown/unknown"]
@@ -52,7 +52,7 @@ class IndexService:
         repository: IndexRepository,
         source_service: SourceService,
         data_dir: Path,
-        embedding_model_name: str,
+        embedding_service: Embedder,
     ) -> None:
         """Initialize the index service.
 
@@ -66,7 +66,7 @@ class IndexService:
         self.snippet_service = SnippetService()
         self.log = structlog.get_logger(__name__)
         self.bm25 = BM25Service(data_dir)
-        self.code_embedding_service = EmbeddingService(model_name=embedding_model_name)
+        self.code_embedding_service = embedding_service
 
     async def create(self, source_id: int) -> IndexView:
         """Create a new index for a source.
@@ -132,7 +132,7 @@ class IndexService:
         # Create snippets for supported file types
         await self._create_snippets(index_id)
 
-        snippets = await self.repository.get_all_snippets()
+        snippets = await self.repository.get_all_snippets(index_id)
 
         self.log.info("Creating keyword index")
         self.bm25.index(
@@ -143,12 +143,17 @@ class IndexService:
         )
 
         self.log.info("Creating semantic code index")
-        for snippet in tqdm(snippets, total=len(snippets), leave=False):
-            embedding = next(self.code_embedding_service.embed([snippet.content]))
+        async for e in tqdm(
+            self.code_embedding_service.embed(
+                [EmbeddingInput(snippet.id, snippet.content) for snippet in snippets]
+            ),
+            total=len(snippets),
+            leave=False,
+        ):
             await self.repository.add_embedding(
                 Embedding(
-                    snippet_id=snippet.id,
-                    embedding=embedding,
+                    snippet_id=e.id,
+                    embedding=e.embedding,
                     type=EmbeddingType.CODE,
                 )
             )

kodit/log.py

@@ -93,6 +93,7 @@ def configure_logging(app_context: AppContext) -> None:
         "uvicorn.access",
         "bm25s",
         "sentence_transformers.SentenceTransformer",
+        "httpx",
     ]:
         if root_logger.getEffectiveLevel() == logging.DEBUG:
             logging.getLogger(_log).handlers.clear()