kodit 0.4.3__py3-none-any.whl → 0.5.1__py3-none-any.whl

kodit/infrastructure/embedding/vectorchord_vector_search_repository.py

@@ -1,10 +1,10 @@
 """VectorChord vector search repository implementation."""
 
-from collections.abc import AsyncGenerator
-from typing import Any, Literal
+from collections.abc import AsyncGenerator, Callable
+from typing import Literal
 
 import structlog
-from sqlalchemy import Result, TextClause, text
+from sqlalchemy import text
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.domain.services.embedding_service import (
@@ -19,6 +19,7 @@ from kodit.domain.value_objects import (
     SearchResult,
 )
 from kodit.infrastructure.sqlalchemy.entities import EmbeddingType
+from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
 
 # SQL Queries
 CREATE_VCHORD_EXTENSION = """
@@ -72,6 +73,10 @@ CHECK_VCHORD_EMBEDDING_EXISTS = """
 SELECT EXISTS(SELECT 1 FROM {TABLE_NAME} WHERE snippet_id = :snippet_id)
 """
 
+CHECK_VCHORD_EMBEDDING_EXISTS_MULTIPLE = """
+SELECT snippet_id FROM {TABLE_NAME} WHERE snippet_id = ANY(:snippet_ids)
+"""
+
 TaskName = Literal["code", "text"]
 
 
@@ -80,8 +85,8 @@ class VectorChordVectorSearchRepository(VectorSearchRepository):
 
     def __init__(
         self,
+        session_factory: Callable[[], AsyncSession],
         task_name: TaskName,
-        session: AsyncSession,
         embedding_provider: EmbeddingProvider,
     ) -> None:
         """Initialize the VectorChord vector search repository.
@@ -93,7 +98,7 @@ class VectorChordVectorSearchRepository(VectorSearchRepository):
 
         """
         self.embedding_provider = embedding_provider
-        self._session = session
+        self.session_factory = session_factory
         self._initialized = False
         self.table_name = f"vectorchord_{task_name}_embeddings"
         self.index_name = f"{self.table_name}_idx"
@@ -111,12 +116,12 @@ class VectorChordVectorSearchRepository(VectorSearchRepository):
 
     async def _create_extensions(self) -> None:
         """Create the necessary extensions."""
-        await self._session.execute(text(CREATE_VCHORD_EXTENSION))
-        await self._commit()
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            await session.execute(text(CREATE_VCHORD_EXTENSION))
 
     async def _create_tables(self) -> None:
         """Create the necessary tables."""
-        req = EmbeddingRequest(snippet_id=0, text="dimension")
+        req = EmbeddingRequest(snippet_id="0", text="dimension")
         vector_dim: list[float] | None = None
         async for batch in self.embedding_provider.embed([req]):
             if batch:
@@ -125,79 +130,85 @@ class VectorChordVectorSearchRepository(VectorSearchRepository):
         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} (
-                    id SERIAL PRIMARY KEY,
-                    snippet_id INT NOT NULL UNIQUE,
-                    embedding VECTOR({len(vector_dim)}) NOT NULL
-                );"""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            await session.execute(
+                text(
+                    f"""CREATE TABLE IF NOT EXISTS {self.table_name} (
+                        id SERIAL PRIMARY KEY,
+                        snippet_id VARCHAR(255) NOT NULL UNIQUE,
+                        embedding VECTOR({len(vector_dim)}) NOT NULL
+                    );"""
+                )
             )
-        )
-        await self._session.execute(
-            text(
-                CREATE_VCHORD_INDEX.format(
-                    TABLE_NAME=self.table_name, INDEX_NAME=self.index_name
+            await session.execute(
+                text(
+                    CREATE_VCHORD_INDEX.format(
+                        TABLE_NAME=self.table_name, INDEX_NAME=self.index_name
+                    )
                 )
             )
-        )
-        result = await self._session.execute(
-            text(CHECK_VCHORD_EMBEDDING_DIMENSION.format(TABLE_NAME=self.table_name))
-        )
-        vector_dim_from_db = result.scalar_one()
-        if vector_dim_from_db != len(vector_dim):
-            msg = (
-                f"Embedding vector dimension does not match database, "
-                f"please delete your index: {vector_dim_from_db} != {len(vector_dim)}"
+            result = await session.execute(
+                text(
+                    CHECK_VCHORD_EMBEDDING_DIMENSION.format(TABLE_NAME=self.table_name)
+                )
             )
-            raise ValueError(msg)
-        await self._commit()
-
-    async def _execute(
-        self, query: TextClause, param_list: list[Any] | dict[str, Any] | None = None
-    ) -> Result:
-        """Execute a query."""
-        if not self._initialized:
-            await self._initialize()
-        return await self._session.execute(query, param_list)
-
-    async def _commit(self) -> None:
-        """Commit the session."""
-        await self._session.commit()
+            vector_dim_from_db = result.scalar_one()
+            if vector_dim_from_db != len(vector_dim):
+                msg = (
+                    f"Embedding vector dimension does not match database, please "
+                    f"delete your index: {vector_dim_from_db} != {len(vector_dim)}"
+                )
+                raise ValueError(msg)
 
     async def index_documents(
         self, request: IndexRequest
     ) -> AsyncGenerator[list[IndexResult], None]:
         """Index documents for vector search."""
+        if not self._initialized:
+            await self._initialize()
+
         if not request.documents:
             yield []
 
+        # Search for existing embeddings
+        existing_ids = await self._get_existing_ids(
+            [doc.snippet_id for doc in request.documents]
+        )
+        new_documents = [
+            doc for doc in request.documents if doc.snippet_id not in existing_ids
+        ]
+        if not new_documents:
+            self.log.info("No new documents to index")
+            return
+
         # Convert to embedding requests
-        requests = [
+        embedding_requests = [
             EmbeddingRequest(snippet_id=doc.snippet_id, text=doc.text)
-            for doc in request.documents
+            for doc in new_documents
         ]
 
-        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]
+        async for batch in self.embedding_provider.embed(embedding_requests):
+            async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+                await session.execute(
+                    text(INSERT_QUERY.format(TABLE_NAME=self.table_name)),
+                    [
+                        {
+                            "snippet_id": result.snippet_id,
+                            "embedding": str(result.embedding),
+                        }
+                        for result in batch
+                    ],
+                )
+                yield [IndexResult(snippet_id=result.snippet_id) for result in batch]
 
     async def search(self, request: SearchRequest) -> list[SearchResult]:
         """Search documents using vector similarity."""
+        if not self._initialized:
+            await self._initialize()
         if not request.query or not request.query.strip():
             return []
 
-        req = EmbeddingRequest(snippet_id=0, text=request.query)
+        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:
@@ -207,39 +218,55 @@ class VectorChordVectorSearchRepository(VectorSearchRepository):
         if not embedding_vec:
             return []
 
-        # Use filtered query if snippet_ids are provided
-        if request.snippet_ids is not None:
-            result = await self._execute(
-                text(SEARCH_QUERY_WITH_FILTER.format(TABLE_NAME=self.table_name)),
-                {
-                    "query": str(embedding_vec),
-                    "top_k": request.top_k,
-                    "snippet_ids": request.snippet_ids,
-                },
-            )
-        else:
-            result = await self._execute(
-                text(SEARCH_QUERY.format(TABLE_NAME=self.table_name)),
-                {"query": str(embedding_vec), "top_k": request.top_k},
-            )
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Use filtered query if snippet_ids are provided
+            if request.snippet_ids is not None:
+                result = await session.execute(
+                    text(SEARCH_QUERY_WITH_FILTER.format(TABLE_NAME=self.table_name)),
+                    {
+                        "query": str(embedding_vec),
+                        "top_k": request.top_k,
+                        "snippet_ids": request.snippet_ids,
+                    },
+                )
+            else:
+                result = await session.execute(
+                    text(SEARCH_QUERY.format(TABLE_NAME=self.table_name)),
+                    {"query": str(embedding_vec), "top_k": request.top_k},
+                )
 
-        rows = result.mappings().all()
+            rows = result.mappings().all()
 
-        return [
-            SearchResult(snippet_id=row["snippet_id"], score=row["score"])
-            for row in rows
-        ]
+            return [
+                SearchResult(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."""
+        if not self._initialized:
+            await self._initialize()
         # 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())
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            result = await session.execute(
+                text(CHECK_VCHORD_EMBEDDING_EXISTS.format(TABLE_NAME=self.table_name)),
+                {"snippet_id": snippet_id},
+            )
+            return bool(result.scalar())
+
+    async def _get_existing_ids(self, snippet_ids: list[str]) -> set[str]:
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            result = await session.execute(
+                text(
+                    CHECK_VCHORD_EMBEDDING_EXISTS_MULTIPLE.format(
+                        TABLE_NAME=self.table_name
+                    )
+                ),
+                {"snippet_ids": snippet_ids},
+            )
+            return {row[0] for row in result.fetchall()}

kodit/infrastructure/enricher/__init__.py

@@ -0,0 +1 @@
+"""Generic enricher infrastructure implementations."""

kodit/infrastructure/enricher/enricher_factory.py

@@ -0,0 +1,53 @@
+"""Enricher factory for creating generic enricher domain services."""
+
+from kodit.config import AppContext, Endpoint
+from kodit.domain.enrichments.enricher import Enricher
+from kodit.infrastructure.enricher.litellm_enricher import LiteLLMEnricher
+from kodit.infrastructure.enricher.local_enricher import LocalEnricher
+from kodit.infrastructure.enricher.null_enricher import NullEnricher
+from kodit.log import log_event
+
+
+def _get_endpoint_configuration(app_context: AppContext) -> Endpoint | None:
+    """Get the endpoint configuration for the enricher service.
+
+    Args:
+        app_context: The application context.
+
+    Returns:
+        The endpoint configuration or None.
+
+    """
+    return app_context.enrichment_endpoint or None
+
+
+def enricher_domain_service_factory(
+    app_context: AppContext,
+    *,
+    use_null_enricher: bool = False,
+) -> Enricher:
+    """Create an enricher domain service.
+
+    Args:
+        app_context: The application context.
+        use_null_enricher: Whether to use the null enricher instead.
+
+    Returns:
+        An enricher domain service instance.
+
+    """
+    enricher: Enricher
+
+    if use_null_enricher:
+        log_event("kodit.enricher", {"provider": "null"})
+        enricher = NullEnricher()
+    else:
+        endpoint = _get_endpoint_configuration(app_context)
+        if endpoint:
+            log_event("kodit.enricher", {"provider": "litellm"})
+            enricher = LiteLLMEnricher(endpoint=endpoint)
+        else:
+            log_event("kodit.enricher", {"provider": "local"})
+            enricher = LocalEnricher()
+
+    return enricher

kodit/infrastructure/{enrichment/litellm_enrichment_provider.py → enricher/litellm_enricher.py}

@@ -1,4 +1,4 @@
-"""LiteLLM enrichment provider implementation."""
+"""LiteLLM enricher implementation."""
 
 import asyncio
 from collections.abc import AsyncGenerator
@@ -10,27 +10,22 @@ import structlog
 from litellm import acompletion
 
 from kodit.config import Endpoint
-from kodit.domain.services.enrichment_service import EnrichmentProvider
-from kodit.domain.value_objects import EnrichmentRequest, EnrichmentResponse
-from kodit.infrastructure.enrichment.utils import clean_thinking_tags
+from kodit.domain.enrichments.enricher import Enricher
+from kodit.domain.enrichments.request import EnrichmentRequest
+from kodit.domain.enrichments.response import EnrichmentResponse
+from kodit.infrastructure.enricher.utils import clean_thinking_tags
 
-ENRICHMENT_SYSTEM_PROMPT = """
-You are a professional software developer. You will be given a snippet of code.
-Please provide a concise explanation of the code.
-"""
-
-# Default tuned conservatively for broad provider compatibility
 DEFAULT_NUM_PARALLEL_TASKS = 20
 
 
-class LiteLLMEnrichmentProvider(EnrichmentProvider):
-    """LiteLLM enrichment provider that supports 100+ providers."""
+class LiteLLMEnricher(Enricher):
+    """LiteLLM enricher that supports 100+ providers."""
 
     def __init__(
         self,
         endpoint: Endpoint,
     ) -> None:
-        """Initialize the LiteLLM enrichment provider.
+        """Initialize the LiteLLM enricher.
 
         Args:
             endpoint: The endpoint configuration containing all settings.
@@ -44,23 +39,20 @@ class LiteLLMEnrichmentProvider(EnrichmentProvider):
         self.num_parallel_tasks = (
             endpoint.num_parallel_tasks or DEFAULT_NUM_PARALLEL_TASKS
         )
-        self.timeout = endpoint.timeout or 30.0
+        self.timeout = endpoint.timeout
         self.extra_params = endpoint.extra_params or {}
 
-        # Configure LiteLLM with custom HTTPX client for Unix socket support if needed
         self._setup_litellm_client()
 
     def _setup_litellm_client(self) -> None:
         """Set up LiteLLM with custom HTTPX client for Unix socket support."""
         if self.socket_path:
-            # Create HTTPX client with Unix socket transport
             transport = httpx.AsyncHTTPTransport(uds=self.socket_path)
             unix_client = httpx.AsyncClient(
                 transport=transport,
-                base_url="http://localhost",  # Base URL for Unix socket
+                base_url="http://localhost",
                 timeout=self.timeout,
             )
-            # Set as LiteLLM's async client session
             litellm.aclient_session = unix_client
 
     async def _call_chat_completion(self, messages: list[dict[str, str]]) -> Any:
@@ -79,20 +71,17 @@ class LiteLLMEnrichmentProvider(EnrichmentProvider):
             "timeout": self.timeout,
         }
 
-        # Add API key if provided
         if self.api_key:
             kwargs["api_key"] = self.api_key
 
-        # Add base_url if provided
         if self.base_url:
             kwargs["api_base"] = self.base_url
 
-        # Add extra parameters
         kwargs.update(self.extra_params)
 
         try:
-            # Use litellm's async completion function
             response = await acompletion(**kwargs)
+            self.log.debug("enrichment request", request=kwargs, response=response)
             return (
                 response.model_dump() if hasattr(response, "model_dump") else response
             )
@@ -108,62 +97,53 @@ class LiteLLMEnrichmentProvider(EnrichmentProvider):
         """Enrich a list of requests using LiteLLM.
 
         Args:
-            requests: List of enrichment requests.
+            requests: List of generic enrichment requests.
 
         Yields:
-            Enrichment responses as they are processed.
+            Generic enrichment responses as they are processed.
 
         """
         if not requests:
             self.log.warning("No requests for enrichment")
             return
 
-        # Process requests in parallel with a semaphore to limit concurrent requests
         sem = asyncio.Semaphore(self.num_parallel_tasks)
 
-        async def process_request(request: EnrichmentRequest) -> EnrichmentResponse:
+        async def process_request(
+            request: EnrichmentRequest,
+        ) -> EnrichmentResponse:
             async with sem:
                 if not request.text:
                     return EnrichmentResponse(
-                        snippet_id=request.snippet_id,
-                        text="",
-                    )
-                try:
-                    messages = [
-                        {
-                            "role": "system",
-                            "content": ENRICHMENT_SYSTEM_PROMPT,
-                        },
-                        {"role": "user", "content": request.text},
-                    ]
-                    response = await self._call_chat_completion(messages)
-                    content = (
-                        response.get("choices", [{}])[0]
-                        .get("message", {})
-                        .get("content", "")
-                    )
-                    # Remove thinking tags from the response
-                    cleaned_content = clean_thinking_tags(content or "")
-                    return EnrichmentResponse(
-                        snippet_id=request.snippet_id,
-                        text=cleaned_content,
-                    )
-                except Exception as e:
-                    self.log.exception("Error enriching request", error=str(e))
-                    return EnrichmentResponse(
-                        snippet_id=request.snippet_id,
+                        id=request.id,
                         text="",
                     )
+                messages = [
+                    {
+                        "role": "system",
+                        "content": request.system_prompt,
+                    },
+                    {"role": "user", "content": request.text},
+                ]
+                response = await self._call_chat_completion(messages)
+                content = (
+                    response.get("choices", [{}])[0]
+                    .get("message", {})
+                    .get("content", "")
+                )
+                cleaned_content = clean_thinking_tags(content or "")
+                return EnrichmentResponse(
+                    id=request.id,
+                    text=cleaned_content,
+                )
 
-        # Create tasks for all requests
         tasks = [process_request(request) for request in requests]
 
-        # Process all requests and yield results as they complete
         for task in asyncio.as_completed(tasks):
             yield await task
 
     async def close(self) -> None:
-        """Close the provider and cleanup HTTPX client if using Unix sockets."""
+        """Close the enricher and cleanup HTTPX client if using Unix sockets."""
         if (
             self.socket_path
             and hasattr(litellm, "aclient_session")

kodit/infrastructure/{enrichment/local_enrichment_provider.py → enricher/local_enricher.py}

@@ -1,4 +1,4 @@
-"""Local enrichment provider implementation."""
+"""Local enricher implementation."""
 
 import asyncio
 import os
@@ -8,28 +8,24 @@ from typing import Any
 import structlog
 import tiktoken
 
-from kodit.domain.services.enrichment_service import EnrichmentProvider
-from kodit.domain.value_objects import EnrichmentRequest, EnrichmentResponse
-from kodit.infrastructure.enrichment.utils import clean_thinking_tags
+from kodit.domain.enrichments.enricher import Enricher
+from kodit.domain.enrichments.request import EnrichmentRequest
+from kodit.domain.enrichments.response import EnrichmentResponse
+from kodit.infrastructure.enricher.utils import clean_thinking_tags
 
-ENRICHMENT_SYSTEM_PROMPT = """
-You are a professional software developer. You will be given a snippet of code.
-Please provide a concise explanation of the code.
-"""
+DEFAULT_ENRICHER_MODEL = "Qwen/Qwen3-0.6B"
+DEFAULT_CONTEXT_WINDOW_SIZE = 2048
 
-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 enrichment provider implementation."""
+class LocalEnricher(Enricher):
+    """Local enricher implementation using local models."""
 
     def __init__(
         self,
-        model_name: str = DEFAULT_ENRICHMENT_MODEL,
+        model_name: str = DEFAULT_ENRICHER_MODEL,
         context_window: int = DEFAULT_CONTEXT_WINDOW_SIZE,
     ) -> None:
-        """Initialize the local enrichment provider.
+        """Initialize the local enricher.
 
         Args:
             model_name: The model name to use for enrichment.
@@ -49,13 +45,13 @@ class LocalEnrichmentProvider(EnrichmentProvider):
         """Enrich a list of requests using local model.
 
         Args:
-            requests: List of enrichment requests.
+            requests: List of generic enrichment requests.
 
         Yields:
-            Enrichment responses as they are processed.
+            Generic enrichment responses as they are processed.
 
         """
-        # Remove empty snippets
+        # Remove empty requests
         requests = [req for req in requests if req.text]
 
         if not requests:
@@ -73,7 +69,7 @@ class LocalEnrichmentProvider(EnrichmentProvider):
                     self.model_name, padding_side="left"
                 )
             if self.model is None:
-                os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
+                os.environ["TOKENIZERS_PARALLELISM"] = "false"
                 self.model = AutoModelForCausalLM.from_pretrained(
                     self.model_name,
                     torch_dtype="auto",
@@ -83,13 +79,13 @@ class LocalEnrichmentProvider(EnrichmentProvider):
 
         await asyncio.to_thread(_init_model)
 
-        # Prepare prompts
+        # Prepare prompts with custom system prompts
         prompts = [
             {
-                "id": req.snippet_id,
+                "id": req.id,
                 "text": self.tokenizer.apply_chat_template(  # type: ignore[attr-defined]
                     [
-                        {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
+                        {"role": "system", "content": req.system_prompt},
                         {"role": "user", "content": req.text},
                     ],
                     tokenize=False,
@@ -121,9 +117,8 @@ class LocalEnrichmentProvider(EnrichmentProvider):
                 )
 
             content = await asyncio.to_thread(process_prompt, prompt)
-            # Remove thinking tags from the response
             cleaned_content = clean_thinking_tags(content)
             yield EnrichmentResponse(
-                snippet_id=prompt["id"],
+                id=prompt["id"],
                 text=cleaned_content,
             )

kodit/infrastructure/enricher/null_enricher.py

@@ -0,0 +1,36 @@
+"""Null enricher implementation."""
+
+from collections.abc import AsyncGenerator
+
+import structlog
+
+from kodit.domain.enrichments.enricher import Enricher
+from kodit.domain.enrichments.request import EnrichmentRequest
+from kodit.domain.enrichments.response import EnrichmentResponse
+
+
+class NullEnricher(Enricher):
+    """Null enricher that returns empty responses."""
+
+    def __init__(self) -> None:
+        """Initialize the null enricher."""
+        self.log = structlog.get_logger(__name__)
+
+    async def enrich(
+        self, requests: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Return empty responses for all requests.
+
+        Args:
+            requests: List of generic enrichment requests.
+
+        Yields:
+            Empty generic enrichment responses.
+
+        """
+        self.log.info("NullEnricher: returning empty responses", count=len(requests))
+        for request in requests:
+            yield EnrichmentResponse(
+                id=request.id,
+                text="",
+            )

kodit/infrastructure/indexing/fusion_service.py

@@ -2,7 +2,7 @@
 
 from collections import defaultdict
 
-from kodit.domain.services.index_query_service import FusionService
+from kodit.domain.protocols import FusionService
 from kodit.domain.value_objects import FusionRequest, FusionResult