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

kodit/embedding/embedding_provider/embedding_provider.py

@@ -1,64 +0,0 @@
-"""Embedding provider."""
-
-from abc import ABC, abstractmethod
-
-import structlog
-import tiktoken
-
-OPENAI_MAX_EMBEDDING_SIZE = 8192
-
-Vector = list[float]
-
-
-class EmbeddingProvider(ABC):
-    """Embedding provider."""
-
-    @abstractmethod
-    async def embed(self, data: list[str]) -> list[Vector]:
-        """Embed a list of strings.
-
-        The embedding provider is responsible for embedding a list of strings into a
-        list of vectors. The embedding provider is responsible for splitting the list of
-        strings into smaller sub-batches and embedding them in parallel.
-        """
-
-
-def split_sub_batches(
-    encoding: tiktoken.Encoding,
-    data: list[str],
-    max_context_window: int = OPENAI_MAX_EMBEDDING_SIZE,
-) -> list[list[str]]:
-    """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.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, disallowed_special=()))
-
-            if item_tokens > max_context_window:
-                # Loop around trying to truncate the snippet until it fits in the max
-                # embedding size
-                while item_tokens > max_context_window:
-                    next_item = next_item[:-1]
-                    item_tokens = len(encoding.encode(next_item, disallowed_special=()))
-
-                data_to_process[0] = next_item
-
-                log.warning("Truncated snippet", snippet=next_item)
-
-            if current_tokens + item_tokens > max_context_window:
-                break
-
-            next_batch.append(data_to_process.pop(0))
-            current_tokens += item_tokens
-
-        if next_batch:
-            result.append(next_batch)
-
-    return result

kodit/embedding/embedding_provider/hash_embedding_provider.py

@@ -1,77 +0,0 @@
-"""Hash embedding provider, for use in tests only."""
-
-import asyncio
-import hashlib
-import math
-from collections.abc import Generator, Sequence
-
-from kodit.embedding.embedding_provider.embedding_provider import (
-    EmbeddingProvider,
-    Vector,
-)
-
-
-class HashEmbeddingProvider(EmbeddingProvider):
-    """A minimal test-time embedding provider.
-
-    • Zero third-party dependencies (uses only std-lib)
-    • Distinguishes strings by hashing with SHA-256
-    • Maps the digest to a fixed-size float vector, then ℓ₂-normalises
-    • Splits work into small asynchronous chunks for speed in event loops
-    """
-
-    def __init__(self, dim: int = 16, batch_size: int = 64) -> None:
-        """Initialize the hash embedding provider."""
-        if dim <= 0:
-            msg = f"dim must be > 0, got {dim}"
-            raise ValueError(msg)
-        if batch_size <= 0:
-            msg = f"batch_size must be > 0, got {batch_size}"
-            raise ValueError(msg)
-        self.dim = dim
-        self.batch_size = batch_size
-
-    async def embed(self, data: list[str]) -> list[Vector]:
-        """Embed every string in *data*, preserving order.
-
-        Work is sliced into *batch_size* chunks and scheduled concurrently
-        (still CPU-bound, but enough to cooperate with an asyncio loop).
-        """
-        if not data:
-            return []
-
-        async def _embed_chunk(chunk: Sequence[str]) -> list[Vector]:
-            return [self._string_to_vector(text) for text in chunk]
-
-        tasks = [
-            asyncio.create_task(_embed_chunk(chunk))
-            for chunk in self._chunked(data, self.batch_size)
-        ]
-
-        vectors: list[Vector] = []
-        for task in tasks:
-            vectors.extend(await task)
-        return vectors
-
-    @staticmethod
-    def _chunked(seq: Sequence[str], size: int) -> Generator[Sequence[str], None, None]:
-        """Yield successive *size*-sized slices from *seq*."""
-        for i in range(0, len(seq), size):
-            yield seq[i : i + size]
-
-    def _string_to_vector(self, text: str) -> Vector:
-        """Deterministically convert *text* to a normalised float vector."""
-        digest = hashlib.sha256(text.encode("utf-8")).digest()
-
-        # Build the vector from 4-byte windows of the digest.
-        vec = [
-            int.from_bytes(
-                digest[(i * 4) % len(digest) : (i * 4) % len(digest) + 4], "big"
-            )
-            / 0xFFFFFFFF
-            for i in range(self.dim)
-        ]
-
-        # ℓ₂-normalise so magnitudes are comparable.
-        norm = math.sqrt(sum(x * x for x in vec)) or 1.0
-        return [x / norm for x in vec]

kodit/embedding/embedding_provider/local_embedding_provider.py

@@ -1,64 +0,0 @@
-"""Local embedding service."""
-
-from __future__ import annotations
-
-import os
-from typing import TYPE_CHECKING
-
-import structlog
-import tiktoken
-from tqdm import tqdm
-
-from kodit.embedding.embedding_provider.embedding_provider import (
-    EmbeddingProvider,
-    Vector,
-    split_sub_batches,
-)
-
-if TYPE_CHECKING:
-    from sentence_transformers import SentenceTransformer
-
-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 embedder."""
-
-    def __init__(self, model_name: str) -> None:
-        """Initialize the local embedder."""
-        self.log = structlog.get_logger(__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."""
-        if self.embedding_model is None:
-            os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid warnings
-            from sentence_transformers import SentenceTransformer
-
-            self.embedding_model = SentenceTransformer(
-                self.model_name,
-                trust_remote_code=True,
-            )
-        return self.embedding_model
-
-    async def embed(self, data: list[str]) -> list[Vector]:
-        """Embed a list of strings."""
-        model = self._model()
-
-        batched_data = split_sub_batches(self.encoding, data)
-
-        results: list[Vector] = []
-        for batch in tqdm(batched_data, total=len(batched_data), leave=False):
-            embeddings = model.encode(batch, show_progress_bar=False, batch_size=4)
-            results.extend([[float(x) for x in embedding] for embedding in embeddings])
-        return results

kodit/embedding/embedding_provider/openai_embedding_provider.py

@@ -1,77 +0,0 @@
-"""OpenAI embedding service."""
-
-import asyncio
-
-import structlog
-import tiktoken
-from openai import AsyncOpenAI
-
-from kodit.embedding.embedding_provider.embedding_provider import (
-    EmbeddingProvider,
-    Vector,
-    split_sub_batches,
-)
-
-OPENAI_NUM_PARALLEL_TASKS = 10
-
-
-class OpenAIEmbeddingProvider(EmbeddingProvider):
-    """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.openai_client = openai_client
-        self.model_name = model_name
-        self.encoding = tiktoken.encoding_for_model(
-            "text-embedding-3-small"
-        )  # Sensible default
-
-    async def embed(self, data: list[str]) -> list[Vector]:
-        """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)
-
-        # Create a list of tuples with a temporary id for each batch
-        # We need to do this so that we can return the results in the same order as the
-        # input data
-        input_data = [(i, batch) for i, batch in enumerate(batched_data)]
-
-        async def process_batch(
-            data: tuple[int, list[str]],
-        ) -> tuple[int, list[Vector]]:
-            batch_id, batch = data
-            async with sem:
-                try:
-                    response = await self.openai_client.embeddings.create(
-                        model=self.model_name,
-                        input=batch,
-                    )
-                    return batch_id, [
-                        [float(x) for x in embedding.embedding]
-                        for embedding in response.data
-                    ]
-                except Exception as e:
-                    self.log.exception("Error embedding batch", error=str(e))
-                    return batch_id, []
-
-        # Create tasks for all batches
-        tasks = [process_batch(batch) for batch in input_data]
-
-        # Process all batches and yield results as they complete
-        results: list[tuple[int, list[Vector]]] = []
-        for task in asyncio.as_completed(tasks):
-            result = await task
-            results.append(result)
-
-        # Output in the same order as the input data
-        ordered_results = [result for _, result in sorted(results, key=lambda x: x[0])]
-        return [item for sublist in ordered_results for item in sublist]

kodit/embedding/local_vector_search_service.py

@@ -1,54 +0,0 @@
-"""Local vector search."""
-
-import structlog
-import tiktoken
-
-from kodit.embedding.embedding_models import Embedding, EmbeddingType
-from kodit.embedding.embedding_provider.embedding_provider import EmbeddingProvider
-from kodit.embedding.embedding_repository import EmbeddingRepository
-from kodit.embedding.vector_search_service import (
-    VectorSearchRequest,
-    VectorSearchResponse,
-    VectorSearchService,
-)
-
-
-class LocalVectorSearchService(VectorSearchService):
-    """Local vector search."""
-
-    def __init__(
-        self,
-        embedding_repository: EmbeddingRepository,
-        embedding_provider: EmbeddingProvider,
-    ) -> None:
-        """Initialize the local embedder."""
-        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")
-
-    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([i.text for i in data])
-        for i, x in zip(data, embeddings, strict=False):
-            await self.embedding_repository.create_embedding(
-                Embedding(
-                    snippet_id=i.snippet_id,
-                    embedding=[float(y) for y in x],
-                    type=EmbeddingType.CODE,
-                )
-            )
-
-    async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
-        """Query the embedding model."""
-        embedding = (await self.embedding_provider.embed([query]))[0]
-        results = await self.embedding_repository.list_semantic_results(
-            EmbeddingType.CODE, [float(x) for x in embedding], top_k
-        )
-        return [
-            VectorSearchResponse(snippet_id, score) for snippet_id, score in results
-        ]

kodit/embedding/vector_search_service.py

@@ -1,38 +0,0 @@
-"""Embedding service."""
-
-from abc import ABC, abstractmethod
-from typing import NamedTuple
-
-
-class VectorSearchResponse(NamedTuple):
-    """Embedding result."""
-
-    snippet_id: int
-    score: float
-
-
-class VectorSearchRequest(NamedTuple):
-    """Input for embedding."""
-
-    snippet_id: int
-    text: str
-
-
-class VectorSearchService(ABC):
-    """Semantic search service interface."""
-
-    @abstractmethod
-    async def index(self, data: list[VectorSearchRequest]) -> 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
-    async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
-        """Query the embedding model."""

kodit/enrichment/__init__.py

@@ -1 +0,0 @@
-"""Enrichment."""

kodit/enrichment/enrichment_provider/__init__.py

@@ -1 +0,0 @@
-"""Enrichment provider."""

kodit/enrichment/enrichment_provider/enrichment_provider.py

@@ -1,16 +0,0 @@
-"""Enrichment provider."""
-
-from abc import ABC, abstractmethod
-
-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.
-"""
-
-
-class EnrichmentProvider(ABC):
-    """Enrichment provider."""
-
-    @abstractmethod
-    async def enrich(self, data: list[str]) -> list[str]:
-        """Enrich a list of strings."""

kodit/enrichment/enrichment_provider/local_enrichment_provider.py

@@ -1,92 +0,0 @@
-"""Local embedding service."""
-
-import os
-
-import structlog
-import tiktoken
-from tqdm import tqdm
-
-from kodit.embedding.embedding_provider.embedding_provider import split_sub_batches
-from kodit.enrichment.enrichment_provider.enrichment_provider import (
-    ENRICHMENT_SYSTEM_PROMPT,
-    EnrichmentProvider,
-)
-
-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."""
-
-    def __init__(
-        self,
-        model_name: str = DEFAULT_ENRICHMENT_MODEL,
-        context_window: int = DEFAULT_CONTEXT_WINDOW_SIZE,
-    ) -> None:
-        """Initialize the local enrichment provider."""
-        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, data: list[str]) -> list[str]:
-        """Enrich a list of strings."""
-        if not data or len(data) == 0:
-            self.log.warning("Data is empty, skipping 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 = [
-            self.tokenizer.apply_chat_template(
-                [
-                    {"role": "system", "content": ENRICHMENT_SYSTEM_PROMPT},
-                    {"role": "user", "content": snippet},
-                ],
-                tokenize=False,
-                add_generation_prompt=True,
-                enable_thinking=False,
-            )
-            for snippet in data
-        ]
-
-        # Batch prompts using split_sub_batches
-        batched_prompts = split_sub_batches(
-            self.encoding, prompts, max_context_window=self.context_window
-        )
-        results = []
-        for batch in tqdm(batched_prompts, leave=False, total=len(batched_prompts)):
-            model_inputs = self.tokenizer(
-                batch, return_tensors="pt", padding=True, truncation=True
-            ).to(self.model.device)
-            generated_ids = self.model.generate(
-                **model_inputs, max_new_tokens=self.context_window
-            )
-            # For each prompt in the batch, decode only the generated part
-            for i, input_ids in enumerate(model_inputs["input_ids"]):
-                output_ids = generated_ids[i][len(input_ids) :].tolist()
-                content = self.tokenizer.decode(
-                    output_ids, skip_special_tokens=True
-                ).strip("\n")
-                results.append(content)
-        return results

kodit/enrichment/enrichment_provider/openai_enrichment_provider.py

@@ -1,81 +0,0 @@
-"""OpenAI embedding service."""
-
-import asyncio
-
-import structlog
-import tiktoken
-from openai import AsyncOpenAI
-from tqdm import tqdm
-
-from kodit.enrichment.enrichment_provider.enrichment_provider import (
-    ENRICHMENT_SYSTEM_PROMPT,
-    EnrichmentProvider,
-)
-
-OPENAI_NUM_PARALLEL_TASKS = 10
-
-
-class OpenAIEnrichmentProvider(EnrichmentProvider):
-    """OpenAI enrichment provider."""
-
-    def __init__(
-        self,
-        openai_client: AsyncOpenAI,
-        model_name: str = "gpt-4o-mini",
-    ) -> None:
-        """Initialize the OpenAI enrichment provider."""
-        self.log = structlog.get_logger(__name__)
-        self.openai_client = openai_client
-        self.model_name = model_name
-        self.encoding = tiktoken.encoding_for_model("gpt-4o-mini")  # Approximation
-
-    async def enrich(self, data: list[str]) -> list[str]:
-        """Enrich a list of documents."""
-        if not data or len(data) == 0:
-            self.log.warning("Data is empty, skipping enrichment")
-            return []
-
-        # Process batches in parallel with a semaphore to limit concurrent requests
-        sem = asyncio.Semaphore(OPENAI_NUM_PARALLEL_TASKS)
-
-        # Create a list of tuples with a temporary id for each snippet
-        # We need to do this so that we can return the results in the same order as the
-        # input data
-        input_data = [(i, snippet) for i, snippet in enumerate(data)]
-
-        async def process_data(data: tuple[int, str]) -> tuple[int, str]:
-            snippet_id, snippet = data
-            if not snippet:
-                return snippet_id, ""
-            async with sem:
-                try:
-                    response = await self.openai_client.chat.completions.create(
-                        model=self.model_name,
-                        messages=[
-                            {
-                                "role": "system",
-                                "content": ENRICHMENT_SYSTEM_PROMPT,
-                            },
-                            {"role": "user", "content": snippet},
-                        ],
-                    )
-                    return snippet_id, response.choices[0].message.content or ""
-                except Exception as e:
-                    self.log.exception("Error enriching data", error=str(e))
-                    return snippet_id, ""
-
-        # Create tasks for all data
-        tasks = [process_data(snippet) for snippet in input_data]
-
-        # Process all data and yield results as they complete
-        results: list[tuple[int, str]] = []
-        for task in tqdm(
-            asyncio.as_completed(tasks),
-            total=len(tasks),
-            leave=False,
-        ):
-            result = await task
-            results.append(result)
-
-        # Output in the same order as the input data
-        return [result for _, result in sorted(results, key=lambda x: x[0])]

kodit/enrichment/enrichment_service.py

@@ -1,33 +0,0 @@
-"""Enrichment service."""
-
-from abc import ABC, abstractmethod
-
-from kodit.enrichment.enrichment_provider.enrichment_provider import EnrichmentProvider
-
-
-class EnrichmentService(ABC):
-    """Enrichment service."""
-
-    @abstractmethod
-    async def enrich(self, data: list[str]) -> list[str]:
-        """Enrich a list of strings."""
-
-
-class NullEnrichmentService(EnrichmentService):
-    """Null enrichment service."""
-
-    async def enrich(self, data: list[str]) -> list[str]:
-        """Enrich a list of strings."""
-        return [""] * len(data)
-
-
-class LLMEnrichmentService(EnrichmentService):
-    """Enrichment service using an LLM."""
-
-    def __init__(self, enrichment_provider: EnrichmentProvider) -> None:
-        """Initialize the enrichment service."""
-        self.enrichment_provider = enrichment_provider
-
-    async def enrich(self, data: list[str]) -> list[str]:
-        """Enrich a list of strings."""
-        return await self.enrichment_provider.enrich(data)

kodit/indexing/__init__.py

@@ -1 +0,0 @@
-"""Indexing package for managing code indexes and search functionality."""

kodit/indexing/fusion.py

@@ -1,67 +0,0 @@
-"""Fusion functions for combining search results."""
-
-from collections import defaultdict
-from dataclasses import dataclass
-
-
-@dataclass
-class FusionResult:
-    """Result of a fusion operation."""
-
-    id: int
-    score: float
-    original_scores: list[float]
-
-
-@dataclass
-class FusionRequest:
-    """Result of a RRF operation."""
-
-    id: int
-    score: float
-
-
-def reciprocal_rank_fusion(
-    rankings: list[list[FusionRequest]], k: float = 60
-) -> list[FusionResult]:
-    """RRF prioritises results that are present in all 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:
-        Dictionary of ids and their 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
-    ]