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

kodit/infrastructure/enrichment/legacy_enrichment_models.py

@@ -0,0 +1,42 @@
+"""Legacy enrichment models for backward compatibility."""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+
+
+@dataclass
+class EnrichmentRequest:
+    """Legacy enrichment request model."""
+
+    snippet_id: int
+    text: str
+
+
+@dataclass
+class EnrichmentResponse:
+    """Legacy enrichment response model."""
+
+    snippet_id: int
+    text: str
+
+
+class EnrichmentService(ABC):
+    """Legacy enrichment service interface."""
+
+    @abstractmethod
+    def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Enrich a list of requests."""
+
+
+class NullEnrichmentService(EnrichmentService):
+    """Null enrichment service for testing."""
+
+    async def enrich(
+        self, data: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Return empty responses for all requests."""
+        for request in data:
+            yield EnrichmentResponse(snippet_id=request.snippet_id, text="")

kodit/{enrichment/enrichment_provider → infrastructure/enrichment}/local_enrichment_provider.py

@@ -1,4 +1,4 @@
-"""Local embedding service."""
+"""Local enrichment provider implementation."""
 
 import os
 from collections.abc import AsyncGenerator
@@ -6,29 +6,33 @@ from collections.abc import AsyncGenerator
 import structlog
 import tiktoken
 
-from kodit.embedding.embedding_provider.embedding_provider import (
-    EmbeddingRequest,
-)
-from kodit.enrichment.enrichment_provider.enrichment_provider import (
-    ENRICHMENT_SYSTEM_PROMPT,
-    EnrichmentProvider,
-    EnrichmentRequest,
-    EnrichmentResponse,
-)
+from kodit.domain.services.enrichment_service import EnrichmentProvider
+from kodit.domain.value_objects import EnrichmentRequest, EnrichmentResponse
+
+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_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."""
+    """Local enrichment provider implementation."""
 
     def __init__(
         self,
         model_name: str = DEFAULT_ENRICHMENT_MODEL,
         context_window: int = DEFAULT_CONTEXT_WINDOW_SIZE,
     ) -> None:
-        """Initialize the local enrichment provider."""
+        """Initialize the local enrichment provider.
+
+        Args:
+            model_name: The model name to use for enrichment.
+            context_window: The context window size for the model.
+
+        """
         self.log = structlog.get_logger(__name__)
         self.model_name = model_name
         self.context_window = context_window
@@ -37,14 +41,22 @@ class LocalEnrichmentProvider(EnrichmentProvider):
         self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
 
     async def enrich(
-        self, data: list[EnrichmentRequest]
+        self, requests: list[EnrichmentRequest]
     ) -> AsyncGenerator[EnrichmentResponse, None]:
-        """Enrich a list of strings."""
+        """Enrich a list of requests using local model.
+
+        Args:
+            requests: List of enrichment requests.
+
+        Yields:
+            Enrichment responses as they are processed.
+
+        """
         # Remove empty snippets
-        data = [snippet for snippet in data if snippet.text]
+        requests = [req for req in requests if req.text]
 
-        if not data or len(data) == 0:
-            self.log.warning("Data is empty, skipping enrichment")
+        if not requests:
+            self.log.warning("No valid requests for enrichment")
             return
 
         from transformers.models.auto.modeling_auto import (
@@ -66,25 +78,25 @@ class LocalEnrichmentProvider(EnrichmentProvider):
             )
 
         # Prepare prompts
-        prompts: list[EmbeddingRequest] = [
-            EmbeddingRequest(
-                id=snippet.snippet_id,
-                text=self.tokenizer.apply_chat_template(
+        prompts = [
+            {
+                "id": req.snippet_id,
+                "text": self.tokenizer.apply_chat_template(
                     [
                         {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
-                        {"role": "user", "content": snippet.text},
+                        {"role": "user", "content": req.text},
                     ],
                     tokenize=False,
                     add_generation_prompt=True,
                     enable_thinking=False,
                 ),
-            )
-            for snippet in data
+            }
+            for req in requests
         ]
 
         for prompt in prompts:
             model_inputs = self.tokenizer(
-                prompt.text,
+                prompt["text"],
                 return_tensors="pt",
                 padding=True,
                 truncation=True,
@@ -98,6 +110,6 @@ class LocalEnrichmentProvider(EnrichmentProvider):
                 "\n"
             )
             yield EnrichmentResponse(
-                snippet_id=prompt.id,
+                snippet_id=prompt["id"],
                 text=content,
             )

kodit/infrastructure/enrichment/null_enrichment_provider.py

@@ -0,0 +1,25 @@
+"""Null enrichment provider for testing."""
+
+from collections.abc import AsyncGenerator
+
+from kodit.domain.services.enrichment_service import EnrichmentProvider
+from kodit.domain.value_objects import EnrichmentRequest, EnrichmentResponse
+
+
+class NullEnrichmentProvider(EnrichmentProvider):
+    """Null enrichment provider that returns empty responses."""
+
+    async def enrich(
+        self, requests: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Return empty responses for all requests.
+
+        Args:
+            requests: List of enrichment requests.
+
+        Yields:
+            Empty enrichment responses.
+
+        """
+        for request in requests:
+            yield EnrichmentResponse(snippet_id=request.snippet_id, text="")

kodit/infrastructure/enrichment/openai_enrichment_provider.py

@@ -0,0 +1,89 @@
+"""OpenAI enrichment provider implementation."""
+
+import asyncio
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import structlog
+
+from kodit.domain.services.enrichment_service import EnrichmentProvider
+from kodit.domain.value_objects import EnrichmentRequest, EnrichmentResponse
+
+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 to approximately fit within OpenAI's rate limit of 500 / RPM
+OPENAI_NUM_PARALLEL_TASKS = 40
+
+
+class OpenAIEnrichmentProvider(EnrichmentProvider):
+    """OpenAI enrichment provider implementation."""
+
+    def __init__(self, openai_client: Any, model_name: str = "gpt-4o-mini") -> None:
+        """Initialize the OpenAI enrichment provider.
+
+        Args:
+            openai_client: The OpenAI client instance.
+            model_name: The model name to use for enrichment.
+
+        """
+        self.log = structlog.get_logger(__name__)
+        self.openai_client = openai_client
+        self.model_name = model_name
+
+    async def enrich(
+        self, requests: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Enrich a list of requests using OpenAI API.
+
+        Args:
+            requests: List of enrichment requests.
+
+        Yields:
+            Enrichment responses as they are processed.
+
+        """
+        if not requests:
+            self.log.warning("No requests for enrichment")
+            return
+
+        # Process batches in parallel with a semaphore to limit concurrent requests
+        sem = asyncio.Semaphore(OPENAI_NUM_PARALLEL_TASKS)
+
+        async def process_request(request: EnrichmentRequest) -> EnrichmentResponse:
+            async with sem:
+                if not request.text:
+                    return EnrichmentResponse(
+                        snippet_id=request.snippet_id,
+                        text="",
+                    )
+                try:
+                    response = await self.openai_client.chat.completions.create(
+                        model=self.model_name,
+                        messages=[
+                            {
+                                "role": "system",
+                                "content": ENRICHMENT_SYSTEM_PROMPT,
+                            },
+                            {"role": "user", "content": request.text},
+                        ],
+                    )
+                    return EnrichmentResponse(
+                        snippet_id=request.snippet_id,
+                        text=response.choices[0].message.content or "",
+                    )
+                except Exception as e:
+                    self.log.exception("Error enriching request", error=str(e))
+                    return EnrichmentResponse(
+                        snippet_id=request.snippet_id,
+                        text="",
+                    )
+
+        # 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

kodit/infrastructure/git/__init__.py

@@ -0,0 +1 @@
+"""Git infrastructure module."""

kodit/{source/git.py → infrastructure/git/git_utils.py}

@@ -1,4 +1,4 @@
-"""Git utilities."""
+"""Git utilities for infrastructure operations."""
 
 import tempfile
 
@@ -6,7 +6,15 @@ import git
 
 
 def is_valid_clone_target(target: str) -> bool:
-    """Return True if the target is clonable."""
+    """Return True if the target is clonable.
+
+    Args:
+        target: The git repository URL or path to validate.
+
+    Returns:
+        True if the target can be cloned, False otherwise.
+
+    """
     with tempfile.TemporaryDirectory() as temp_dir:
         try:
             git.Repo.clone_from(target, temp_dir)

kodit/infrastructure/ignore/__init__.py

@@ -0,0 +1 @@
+"""Ignore infrastructure module."""

kodit/{source/ignore.py → infrastructure/ignore/ignore_pattern_provider.py}

@@ -1,18 +1,27 @@
-"""Ignore patterns."""
+"""Infrastructure implementation of ignore pattern provider."""
 
 from pathlib import Path
 
 import git
 import pathspec
 
-from kodit.source.git import is_valid_clone_target
+from kodit.domain.services.ignore_service import IgnorePatternProvider
+from kodit.infrastructure.git.git_utils import is_valid_clone_target
 
 
-class IgnorePatterns:
-    """Ignore patterns."""
+class GitIgnorePatternProvider(IgnorePatternProvider):
+    """Ignore pattern provider for git repositories."""
 
     def __init__(self, base_dir: Path) -> None:
-        """Initialize the ignore patterns."""
+        """Initialize the ignore pattern provider.
+
+        Args:
+            base_dir: The base directory to check for ignore patterns.
+
+        Raises:
+            ValueError: If the base directory is not a directory.
+
+        """
         if not base_dir.is_dir():
             msg = f"Base directory is not a directory: {base_dir}"
             raise ValueError(msg)
@@ -25,7 +34,15 @@ class IgnorePatterns:
             self.git_repo = git.Repo(base_dir)
 
     def should_ignore(self, path: Path) -> bool:
-        """Check if a path should be ignored."""
+        """Check if a path should be ignored.
+
+        Args:
+            path: The path to check.
+
+        Returns:
+            True if the path should be ignored, False otherwise.
+
+        """
         if path.is_dir():
             return False
 

kodit/infrastructure/indexing/__init__.py

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

kodit/infrastructure/indexing/fusion_service.py

@@ -0,0 +1,55 @@
+"""Infrastructure implementation of the fusion service."""
+
+from collections import defaultdict
+
+from kodit.domain.services.indexing_service import FusionService
+from kodit.domain.value_objects import FusionRequest, FusionResult
+
+
+class ReciprocalRankFusionService(FusionService):
+    """Infrastructure implementation of reciprocal rank fusion."""
+
+    def reciprocal_rank_fusion(
+        self, rankings: list[list[FusionRequest]], k: float = 60
+    ) -> list[FusionResult]:
+        """Perform reciprocal rank fusion on search results.
+
+        Args:
+            rankings: List of rankers, each containing a list of document ids.
+                Top of the list is considered to be the best result.
+            k: Parameter for RRF.
+
+        Returns:
+            List of fused results with scores.
+
+        """
+        scores = {}
+        for ranker in rankings:
+            for rank in ranker:
+                scores[rank.id] = float(0)
+
+        for ranker in rankings:
+            for i, rank in enumerate(ranker):
+                scores[rank.id] += 1.0 / (k + i)
+
+        # Create a list of tuples of ids and their scores
+        results = [(rank, scores[rank]) for rank in scores]
+
+        # Sort results by score
+        results.sort(key=lambda x: x[1], reverse=True)
+
+        # Create a map of original scores to ids
+        original_scores_to_ids = defaultdict(list)
+        for ranker in rankings:
+            for rank in ranker:
+                original_scores_to_ids[rank.id].append(rank.score)
+
+        # Rebuild a list of final results with their original scores
+        return [
+            FusionResult(
+                id=result[0],
+                score=result[1],
+                original_scores=original_scores_to_ids[result[0]],
+            )
+            for result in results
+        ]