kodit 0.2.3__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/infrastructure/enrichment/local_enrichment_provider.py

@@ -0,0 +1,115 @@
+"""Local enrichment provider implementation."""
+
+import os
+from collections.abc import AsyncGenerator
+
+import structlog
+import tiktoken
+
+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 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.
+
+        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
+        self.model = None
+        self.tokenizer = None
+        self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
+
+    async def enrich(
+        self, requests: list[EnrichmentRequest]
+    ) -> AsyncGenerator[EnrichmentResponse, None]:
+        """Enrich a list of requests using local model.
+
+        Args:
+            requests: List of enrichment requests.
+
+        Yields:
+            Enrichment responses as they are processed.
+
+        """
+        # Remove empty snippets
+        requests = [req for req in requests if req.text]
+
+        if not requests:
+            self.log.warning("No valid requests for enrichment")
+            return
+
+        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, 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",
+            )
+
+        # Prepare prompts
+        prompts = [
+            {
+                "id": req.snippet_id,
+                "text": self.tokenizer.apply_chat_template(
+                    [
+                        {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
+                        {"role": "user", "content": req.text},
+                    ],
+                    tokenize=False,
+                    add_generation_prompt=True,
+                    enable_thinking=False,
+                ),
+            }
+            for req in requests
+        ]
+
+        for prompt in prompts:
+            model_inputs = self.tokenizer(
+                prompt["text"],
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+            ).to(self.model.device)
+            generated_ids = self.model.generate(
+                **model_inputs, max_new_tokens=self.context_window
+            )
+            input_ids = model_inputs["input_ids"][0]
+            output_ids = generated_ids[0][len(input_ids) :].tolist()
+            content = self.tokenizer.decode(output_ids, skip_special_tokens=True).strip(
+                "\n"
+            )
+            yield EnrichmentResponse(
+                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
+        ]