kodit 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

kodit/infrastructure/embedding/embedding_providers/hash_embedding_provider.py

@@ -0,0 +1,79 @@
+"""Hash-based embedding provider for testing purposes."""
+
+import hashlib
+from collections.abc import AsyncGenerator
+
+import structlog
+
+from kodit.domain.services.embedding_service import EmbeddingProvider
+from kodit.domain.value_objects import EmbeddingRequest, EmbeddingResponse
+
+# Constants for different embedding sizes
+TINY = 64
+CODE = 1536
+
+
+class HashEmbeddingProvider(EmbeddingProvider):
+    """Hash-based embedding that generates deterministic embeddings for testing."""
+
+    def __init__(self, embedding_size: int = CODE) -> None:
+        """Initialize the hash embedding provider.
+
+        Args:
+            embedding_size: The size of the embedding vectors to generate
+
+        """
+        self.embedding_size = embedding_size
+        self.log = structlog.get_logger(__name__)
+
+    def embed(
+        self, data: list[EmbeddingRequest]
+    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
+        """Embed a list of strings using a simple hash-based approach."""
+        if not data:
+
+            async def empty_generator() -> AsyncGenerator[
+                list[EmbeddingResponse], None
+            ]:
+                if False:
+                    yield []
+
+            return empty_generator()
+
+        # Process in batches
+        batch_size = 10
+
+        async def _embed_batches() -> AsyncGenerator[list[EmbeddingResponse], None]:
+            for i in range(0, len(data), batch_size):
+                batch = data[i : i + batch_size]
+                responses = []
+
+                for request in batch:
+                    # Generate a deterministic embedding based on the text
+                    embedding = self._generate_embedding(request.text)
+                    responses.append(
+                        EmbeddingResponse(
+                            snippet_id=request.snippet_id, embedding=embedding
+                        )
+                    )
+
+                yield responses
+
+        return _embed_batches()
+
+    def _generate_embedding(self, text: str) -> list[float]:
+        """Generate a deterministic embedding for the given text."""
+        # Use SHA-256 hash of the text as a seed
+        hash_obj = hashlib.sha256(text.encode("utf-8"))
+        hash_bytes = hash_obj.digest()
+
+        # Convert hash bytes to a list of floats
+        embedding = []
+        for i in range(self.embedding_size):
+            # Use different bytes for each dimension
+            byte_index = i % len(hash_bytes)
+            # Convert byte to float between -1 and 1
+            value = (hash_bytes[byte_index] - 128) / 128.0
+            embedding.append(value)
+
+        return embedding

kodit/infrastructure/embedding/embedding_providers/local_embedding_provider.py

@@ -0,0 +1,129 @@
+"""Local embedding provider implementation."""
+
+import os
+from collections.abc import AsyncGenerator
+from time import time
+from typing import TYPE_CHECKING
+
+import structlog
+
+from kodit.domain.services.embedding_service import EmbeddingProvider
+from kodit.domain.value_objects import EmbeddingRequest, EmbeddingResponse
+
+from .batching import split_sub_batches
+
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+    from tiktoken import Encoding
+
+# Constants for different embedding models
+TINY = "tiny"
+CODE = "code"
+TEST = "test"
+
+COMMON_EMBEDDING_MODELS = {
+    TINY: "ibm-granite/granite-embedding-30m-english",
+    CODE: "flax-sentence-embeddings/st-codesearch-distilroberta-base",
+    TEST: "minishlab/potion-base-4M",
+}
+
+
+class LocalEmbeddingProvider(EmbeddingProvider):
+    """Local embedding provider that uses sentence-transformers."""
+
+    def __init__(self, model_name: str = CODE) -> None:
+        """Initialize the local embedding provider.
+
+        Args:
+            model_name: The model name to use for embeddings. Can be a preset
+                       ('tiny', 'code', 'test') or a full model name.
+
+        """
+        self.log = structlog.get_logger(__name__)
+        self.model_name = COMMON_EMBEDDING_MODELS.get(model_name, model_name)
+        self.encoding_name = "text-embedding-3-small"
+        self.embedding_model: SentenceTransformer | None = None
+        self.encoding: Encoding | None = None
+
+    def _encoding(self) -> "Encoding":
+        """Get the tiktoken encoding."""
+        if self.encoding is None:
+            from tiktoken import encoding_for_model
+
+            start_time = time()
+            self.encoding = encoding_for_model(self.encoding_name)
+            self.log.debug(
+                "Encoding loaded",
+                model_name=self.encoding_name,
+                duration=time() - start_time,
+            )
+        return self.encoding
+
+    def _model(self) -> "SentenceTransformer":
+        """Get the embedding model."""
+        if self.embedding_model is None:
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+            from sentence_transformers import SentenceTransformer
+
+            start_time = time()
+            self.embedding_model = SentenceTransformer(
+                self.model_name,
+                trust_remote_code=True,
+            )
+            self.log.debug(
+                "Model loaded",
+                model_name=self.model_name,
+                duration=time() - start_time,
+            )
+        return self.embedding_model
+
+    async def embed(
+        self, data: list[EmbeddingRequest]
+    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
+        """Embed a list of strings using the local model."""
+        if not data:
+            yield []
+
+        model = self._model()
+        encoding = self._encoding()
+
+        # Split into sub-batches based on token limits
+        batched_data = self._split_sub_batches(encoding, data)
+
+        for batch in batched_data:
+            try:
+                # Encode the texts using the model
+                embeddings = model.encode(
+                    [item.text for item in batch],
+                    show_progress_bar=False,
+                    batch_size=4,
+                )
+
+                # Convert to our response format
+                responses = [
+                    EmbeddingResponse(
+                        snippet_id=item.snippet_id,
+                        embedding=[float(x) for x in embedding],
+                    )
+                    for item, embedding in zip(batch, embeddings, strict=True)
+                ]
+
+                yield responses
+
+            except Exception as e:
+                self.log.exception("Error generating embeddings", error=str(e))
+                # Return zero embeddings on error
+                responses = [
+                    EmbeddingResponse(
+                        snippet_id=item.snippet_id,
+                        embedding=[0.0] * 1536,  # Default embedding size
+                    )
+                    for item in batch
+                ]
+                yield responses
+
+    def _split_sub_batches(
+        self, encoding: "Encoding", data: list[EmbeddingRequest]
+    ) -> list[list[EmbeddingRequest]]:
+        """Proxy to the shared batching utility (kept for backward-compat)."""
+        return split_sub_batches(encoding, data)

kodit/infrastructure/embedding/embedding_providers/openai_embedding_provider.py

@@ -0,0 +1,113 @@
+"""OpenAI embedding provider implementation."""
+
+import asyncio
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import structlog
+import tiktoken
+from tiktoken import Encoding
+
+from kodit.domain.services.embedding_service import EmbeddingProvider
+from kodit.domain.value_objects import EmbeddingRequest, EmbeddingResponse
+
+from .batching import split_sub_batches
+
+# Constants
+MAX_TOKENS = 8192  # Conservative token limit for the embedding model
+BATCH_SIZE = (
+    10  # Maximum number of items per API call (keeps existing test expectations)
+)
+OPENAI_NUM_PARALLEL_TASKS = 25  # Semaphore limit for concurrent OpenAI requests
+
+
+class OpenAIEmbeddingProvider(EmbeddingProvider):
+    """OpenAI embedding provider that uses OpenAI's embedding API."""
+
+    def __init__(
+        self, openai_client: Any, model_name: str = "text-embedding-3-small"
+    ) -> None:
+        """Initialize the OpenAI embedding provider.
+
+        Args:
+            openai_client: The OpenAI client instance
+            model_name: The model name to use for embeddings
+
+        """
+        self.openai_client = openai_client
+        self.model_name = model_name
+        self.log = structlog.get_logger(__name__)
+
+        # Lazily initialised token encoding
+        self._encoding: Encoding | None = None
+
+    # ---------------------------------------------------------------------
+    # Helper utilities
+    # ---------------------------------------------------------------------
+
+    def _get_encoding(self) -> "Encoding":
+        """Return (and cache) the tiktoken encoding for the chosen model."""
+        if self._encoding is None:
+            self._encoding = tiktoken.encoding_for_model(self.model_name)
+        return self._encoding
+
+    def _split_sub_batches(
+        self, encoding: "Encoding", data: list[EmbeddingRequest]
+    ) -> list[list[EmbeddingRequest]]:
+        """Proxy to the shared batching utility (kept for backward-compat)."""
+        return split_sub_batches(
+            encoding,
+            data,
+            max_tokens=MAX_TOKENS,
+            batch_size=BATCH_SIZE,
+        )
+
+    async def embed(
+        self, data: list[EmbeddingRequest]
+    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
+        """Embed a list of strings using OpenAI's API."""
+        if not data:
+            yield []
+
+        encoding = self._get_encoding()
+
+        # First, split by token limits (and max batch size)
+        batched_data = self._split_sub_batches(encoding, data)
+
+        # -----------------------------------------------------------------
+        # Process batches concurrently (but bounded by a semaphore)
+        # -----------------------------------------------------------------
+
+        sem = asyncio.Semaphore(OPENAI_NUM_PARALLEL_TASKS)
+
+        async def _process_batch(
+            batch: list[EmbeddingRequest],
+        ) -> list[EmbeddingResponse]:
+            async with sem:
+                try:
+                    response = await self.openai_client.embeddings.create(
+                        model=self.model_name,
+                        input=[item.text for item in batch],
+                    )
+
+                    return [
+                        EmbeddingResponse(
+                            snippet_id=item.snippet_id,
+                            embedding=embedding.embedding,
+                        )
+                        for item, embedding in zip(batch, response.data, strict=True)
+                    ]
+                except Exception as e:
+                    self.log.exception("Error embedding batch", error=str(e))
+                    # Fall back to zero embeddings so pipeline can continue
+                    return [
+                        EmbeddingResponse(
+                            snippet_id=item.snippet_id,
+                            embedding=[0.0] * 1536,  # Default OpenAI dim
+                        )
+                        for item in batch
+                    ]
+
+        tasks = [_process_batch(batch) for batch in batched_data]
+        for task in asyncio.as_completed(tasks):
+            yield await task

kodit/infrastructure/embedding/local_vector_search_repository.py

@@ -0,0 +1,114 @@
+"""Local vector search repository implementation."""
+
+from collections.abc import AsyncGenerator
+
+import structlog
+import tiktoken
+
+from kodit.domain.entities import Embedding, EmbeddingType
+from kodit.domain.services.embedding_service import (
+    EmbeddingProvider,
+    VectorSearchRepository,
+)
+from kodit.domain.value_objects import (
+    EmbeddingRequest,
+    IndexResult,
+    VectorIndexRequest,
+    VectorSearchQueryRequest,
+    VectorSearchResult,
+)
+from kodit.infrastructure.sqlalchemy.embedding_repository import (
+    SqlAlchemyEmbeddingRepository,
+)
+
+
+class LocalVectorSearchRepository(VectorSearchRepository):
+    """Local vector search repository implementation."""
+
+    def __init__(
+        self,
+        embedding_repository: SqlAlchemyEmbeddingRepository,
+        embedding_provider: EmbeddingProvider,
+        embedding_type: EmbeddingType = EmbeddingType.CODE,
+    ) -> None:
+        """Initialize the local vector search repository.
+
+        Args:
+            embedding_repository: The SQLAlchemy embedding repository
+            embedding_provider: The embedding provider for generating embeddings
+            embedding_type: The type of embedding to use
+
+        """
+        self.log = structlog.get_logger(__name__)
+        self.embedding_repository = embedding_repository
+        self.embedding_provider = embedding_provider
+        self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
+        self.embedding_type = embedding_type
+
+    def index_documents(
+        self, request: VectorIndexRequest
+    ) -> AsyncGenerator[list[IndexResult], None]:
+        """Index documents for vector search."""
+        if not request.documents:
+
+            async def empty_generator() -> AsyncGenerator[list[IndexResult], None]:
+                if False:
+                    yield []
+
+            return empty_generator()
+
+        # Convert to embedding requests
+        requests = [
+            EmbeddingRequest(snippet_id=doc.snippet_id, text=doc.text)
+            for doc in request.documents
+        ]
+
+        async def _index_batches() -> AsyncGenerator[list[IndexResult], None]:
+            async for batch in self.embedding_provider.embed(requests):
+                results = []
+                for result in batch:
+                    await self.embedding_repository.create_embedding(
+                        Embedding(
+                            snippet_id=result.snippet_id,
+                            embedding=result.embedding,
+                            type=self.embedding_type,
+                        )
+                    )
+                    results.append(IndexResult(snippet_id=result.snippet_id))
+                yield results
+
+        return _index_batches()
+
+    async def search(
+        self, request: VectorSearchQueryRequest
+    ) -> list[VectorSearchResult]:
+        """Search documents using vector similarity."""
+        # Build a single-item request and collect its embedding
+        req = EmbeddingRequest(snippet_id=0, text=request.query)
+        embedding_vec: list[float] | None = None
+        async for batch in self.embedding_provider.embed([req]):
+            if batch:
+                embedding_vec = [float(v) for v in batch[0].embedding]
+                break
+
+        if not embedding_vec:
+            return []
+
+        results = await self.embedding_repository.list_semantic_results(
+            self.embedding_type, embedding_vec, request.top_k
+        )
+        return [
+            VectorSearchResult(snippet_id=snippet_id, score=score)
+            for snippet_id, score in results
+        ]
+
+    async def has_embedding(
+        self, snippet_id: int, embedding_type: EmbeddingType
+    ) -> bool:
+        """Check if a snippet has an embedding."""
+        return (
+            await self.embedding_repository.get_embedding_by_snippet_id_and_type(
+                snippet_id, embedding_type
+            )
+            is not None
+        )

kodit/{embedding/vectorchord_vector_search_service.py → infrastructure/embedding/vectorchord_vector_search_repository.py}

@@ -1,16 +1,23 @@
-"""Vectorchord vector search."""
+"""VectorChord vector search repository implementation."""
 
+from collections.abc import AsyncGenerator
 from typing import Any, Literal
 
 import structlog
 from sqlalchemy import Result, TextClause, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.embedding.embedding_provider.embedding_provider import EmbeddingProvider
-from kodit.embedding.vector_search_service import (
-    VectorSearchRequest,
-    VectorSearchResponse,
-    VectorSearchService,
+from kodit.domain.entities import EmbeddingType
+from kodit.domain.services.embedding_service import (
+    EmbeddingProvider,
+    VectorSearchRepository,
+)
+from kodit.domain.value_objects import (
+    EmbeddingRequest,
+    IndexResult,
+    VectorIndexRequest,
+    VectorSearchQueryRequest,
+    VectorSearchResult,
 )
 
 # SQL Queries
@@ -52,11 +59,15 @@ ORDER BY score ASC
 LIMIT :top_k;
 """
 
+CHECK_VCHORD_EMBEDDING_EXISTS = """
+SELECT EXISTS(SELECT 1 FROM {TABLE_NAME} WHERE snippet_id = :snippet_id)
+"""
+
 TaskName = Literal["code", "text"]
 
 
-class VectorChordVectorSearchService(VectorSearchService):
-    """VectorChord vector search."""
+class VectorChordVectorSearchRepository(VectorSearchRepository):
+    """VectorChord vector search repository implementation."""
 
     def __init__(
         self,
@@ -64,7 +75,14 @@ class VectorChordVectorSearchService(VectorSearchService):
         session: AsyncSession,
         embedding_provider: EmbeddingProvider,
     ) -> None:
-        """Initialize the VectorChord BM25."""
+        """Initialize the VectorChord vector search repository.
+
+        Args:
+            task_name: The task name (code or text)
+            session: The SQLAlchemy async session
+            embedding_provider: The embedding provider for generating embeddings
+
+        """
         self.embedding_provider = embedding_provider
         self._session = session
         self._initialized = False
@@ -89,7 +107,15 @@ class VectorChordVectorSearchService(VectorSearchService):
 
     async def _create_tables(self) -> None:
         """Create the necessary tables."""
-        vector_dim = (await self.embedding_provider.embed(["dimension"]))[0]
+        req = EmbeddingRequest(snippet_id=0, text="dimension")
+        vector_dim: list[float] | None = None
+        async for batch in self.embedding_provider.embed([req]):
+            if batch:
+                vector_dim = batch[0].embedding
+                break
+        if vector_dim is None:
+            msg = "Failed to obtain embedding dimension from provider"
+            raise RuntimeError(msg)
         await self._session.execute(
             text(
                 f"""CREATE TABLE IF NOT EXISTS {self.table_name} (
@@ -130,35 +156,75 @@ class VectorChordVectorSearchService(VectorSearchService):
         """Commit the session."""
         await self._session.commit()
 
-    async def index(self, data: list[VectorSearchRequest]) -> None:
-        """Embed a list of documents."""
-        if not data or len(data) == 0:
-            self.log.warning("Embedding data is empty, skipping embedding")
-            return
-
-        embeddings = await self.embedding_provider.embed([doc.text for doc in data])
-        # Execute inserts
-        await self._execute(
-            text(INSERT_QUERY.format(TABLE_NAME=self.table_name)),
-            [
-                {"snippet_id": doc.snippet_id, "embedding": str(embedding)}
-                for doc, embedding in zip(data, embeddings, strict=True)
-            ],
-        )
-        await self._commit()
+    def index_documents(
+        self, request: VectorIndexRequest
+    ) -> AsyncGenerator[list[IndexResult], None]:
+        """Index documents for vector search."""
+        if not request.documents:
+
+            async def empty_generator() -> AsyncGenerator[list[IndexResult], None]:
+                if False:
+                    yield []
 
-    async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
-        """Query the embedding model."""
-        embedding = await self.embedding_provider.embed([query])
-        if len(embedding) == 0 or len(embedding[0]) == 0:
+            return empty_generator()
+
+        # Convert to embedding requests
+        requests = [
+            EmbeddingRequest(snippet_id=doc.snippet_id, text=doc.text)
+            for doc in request.documents
+        ]
+
+        async def _index_batches() -> AsyncGenerator[list[IndexResult], None]:
+            async for batch in self.embedding_provider.embed(requests):
+                await self._execute(
+                    text(INSERT_QUERY.format(TABLE_NAME=self.table_name)),
+                    [
+                        {
+                            "snippet_id": result.snippet_id,
+                            "embedding": str(result.embedding),
+                        }
+                        for result in batch
+                    ],
+                )
+                await self._commit()
+                yield [IndexResult(snippet_id=result.snippet_id) for result in batch]
+
+        return _index_batches()
+
+    async def search(
+        self, request: VectorSearchQueryRequest
+    ) -> list[VectorSearchResult]:
+        """Search documents using vector similarity."""
+        req = EmbeddingRequest(snippet_id=0, text=request.query)
+        embedding_vec: list[float] | None = None
+        async for batch in self.embedding_provider.embed([req]):
+            if batch:
+                embedding_vec = batch[0].embedding
+                break
+
+        if not embedding_vec:
             return []
         result = await self._execute(
             text(SEARCH_QUERY.format(TABLE_NAME=self.table_name)),
-            {"query": str(embedding[0]), "top_k": top_k},
+            {"query": str(embedding_vec), "top_k": request.top_k},
         )
         rows = result.mappings().all()
 
         return [
-            VectorSearchResponse(snippet_id=row["snippet_id"], score=row["score"])
+            VectorSearchResult(snippet_id=row["snippet_id"], score=row["score"])
             for row in rows
         ]
+
+    async def has_embedding(
+        self, snippet_id: int, embedding_type: EmbeddingType
+    ) -> bool:
+        """Check if a snippet has an embedding."""
+        # For VectorChord, we check if the snippet exists in the table
+        # Note: embedding_type is ignored since VectorChord uses separate
+        # tables per task
+        # ruff: noqa: ARG002
+        result = await self._execute(
+            text(CHECK_VCHORD_EMBEDDING_EXISTS.format(TABLE_NAME=self.table_name)),
+            {"snippet_id": snippet_id},
+        )
+        return bool(result.scalar())

kodit/infrastructure/enrichment/__init__.py

@@ -0,0 +1 @@
+"""Infrastructure enrichment module."""

kodit/{enrichment → infrastructure/enrichment}/enrichment_factory.py

@@ -1,28 +1,42 @@
-"""Embedding service."""
+"""Enrichment factory for creating enrichment domain services."""
 
 from kodit.config import AppContext, Endpoint
-from kodit.enrichment.enrichment_provider.local_enrichment_provider import (
+from kodit.domain.services.enrichment_service import EnrichmentDomainService
+from kodit.infrastructure.enrichment.local_enrichment_provider import (
     LocalEnrichmentProvider,
 )
-from kodit.enrichment.enrichment_provider.openai_enrichment_provider import (
+from kodit.infrastructure.enrichment.openai_enrichment_provider import (
     OpenAIEnrichmentProvider,
 )
-from kodit.enrichment.enrichment_service import (
-    EnrichmentService,
-    LLMEnrichmentService,
-)
 from kodit.log import log_event
 
 
 def _get_endpoint_configuration(app_context: AppContext) -> Endpoint | None:
-    """Get the endpoint configuration for the enrichment service."""
+    """Get the endpoint configuration for the enrichment service.
+
+    Args:
+        app_context: The application context.
+
+    Returns:
+        The endpoint configuration or None.
+
+    """
     return app_context.enrichment_endpoint or app_context.default_endpoint or None
 
 
-def enrichment_factory(app_context: AppContext) -> EnrichmentService:
-    """Create an enrichment service."""
+def create_enrichment_domain_service(
+    app_context: AppContext,
+) -> EnrichmentDomainService:
+    """Create an enrichment domain service.
+
+    Args:
+        app_context: The application context.
+
+    Returns:
+        An enrichment domain service instance.
+
+    """
     endpoint = _get_endpoint_configuration(app_context)
-    endpoint = app_context.enrichment_endpoint or app_context.default_endpoint or None
 
     if endpoint and endpoint.type == "openai":
         log_event("kodit.enrichment", {"provider": "openai"})
@@ -32,6 +46,8 @@ def enrichment_factory(app_context: AppContext) -> EnrichmentService:
             openai_client=AsyncOpenAI(
                 api_key=endpoint.api_key or "default",
                 base_url=endpoint.base_url or "https://api.openai.com/v1",
+                timeout=60,
+                max_retries=2,
             ),
             model_name=endpoint.model or "gpt-4o-mini",
         )
@@ -39,4 +55,4 @@ def enrichment_factory(app_context: AppContext) -> EnrichmentService:
         log_event("kodit.enrichment", {"provider": "local"})
         enrichment_provider = LocalEnrichmentProvider()
 
-    return LLMEnrichmentService(enrichment_provider=enrichment_provider)
+    return EnrichmentDomainService(enrichment_provider=enrichment_provider)