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

kodit/embedding/embedding_provider/embedding_provider.py

@@ -1,92 +0,0 @@
-"""Embedding provider."""
-
-from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator
-from dataclasses import dataclass
-
-import structlog
-import tiktoken
-
-OPENAI_MAX_EMBEDDING_SIZE = 8192
-
-Vector = list[float]
-
-
-@dataclass
-class EmbeddingRequest:
-    """Embedding request."""
-
-    id: int
-    text: str
-
-
-@dataclass
-class EmbeddingResponse:
-    """Embedding response."""
-
-    id: int
-    embedding: Vector
-
-
-class EmbeddingProvider(ABC):
-    """Embedding provider."""
-
-    @abstractmethod
-    def embed(
-        self, data: list[EmbeddingRequest]
-    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
-        """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[EmbeddingRequest],
-    max_context_window: int = OPENAI_MAX_EMBEDDING_SIZE,
-) -> list[list[EmbeddingRequest]]:
-    """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.text.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.text, disallowed_special=()))
-
-            if item_tokens > max_context_window:
-                # Optimise truncation by operating on tokens directly instead of
-                # removing one character at a time and repeatedly re-encoding.
-                tokens = encoding.encode(next_item.text, disallowed_special=())
-                if len(tokens) > max_context_window:
-                    # Keep only the first *max_context_window* tokens.
-                    tokens = tokens[:max_context_window]
-                    # Convert back to text. This requires only one decode call and
-                    # guarantees that the resulting string fits the token budget.
-                    next_item.text = encoding.decode(tokens)
-                    item_tokens = max_context_window  # We know the exact size now
-
-                    data_to_process[0] = next_item
-
-                    log.warning(
-                        "Truncated snippet because it was too long to embed",
-                        snippet=next_item.text[:100] + "...",
-                    )
-
-            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,86 +0,0 @@
-"""Hash embedding provider, for use in tests only."""
-
-import asyncio
-import hashlib
-import math
-from collections.abc import AsyncGenerator, Generator, Sequence
-
-from kodit.embedding.embedding_provider.embedding_provider import (
-    EmbeddingProvider,
-    EmbeddingRequest,
-    EmbeddingResponse,
-    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[EmbeddingRequest]
-    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
-        """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:
-            yield []
-
-        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([i.text for i in data], self.batch_size)
-        ]
-
-        for task in tasks:
-            result = await task
-            yield [
-                EmbeddingResponse(
-                    id=item.id,
-                    embedding=embedding,
-                )
-                for item, embedding in zip(data, result, strict=True)
-            ]
-
-    @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,96 +0,0 @@
-"""Local embedding service."""
-
-from __future__ import annotations
-
-import os
-from time import time
-from typing import TYPE_CHECKING
-
-import structlog
-
-from kodit.embedding.embedding_provider.embedding_provider import (
-    EmbeddingProvider,
-    EmbeddingRequest,
-    EmbeddingResponse,
-    split_sub_batches,
-)
-
-if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator
-
-    from sentence_transformers import SentenceTransformer
-    from tiktoken import Encoding
-
-
-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.encoding_name = "text-embedding-3-small"
-        self.embedding_model = None
-        self.encoding = None
-
-    def _encoding(self) -> 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."""
-        model = self._model()
-
-        batched_data = split_sub_batches(self._encoding(), data)
-
-        for batch in batched_data:
-            embeddings = model.encode(
-                [i.text for i in batch], show_progress_bar=False, batch_size=4
-            )
-            yield [
-                EmbeddingResponse(
-                    id=item.id,
-                    embedding=[float(x) for x in embedding],
-                )
-                for item, embedding in zip(batch, embeddings, strict=True)
-            ]

kodit/embedding/embedding_provider/openai_embedding_provider.py

@@ -1,73 +0,0 @@
-"""OpenAI embedding service."""
-
-import asyncio
-from collections.abc import AsyncGenerator
-
-import structlog
-import tiktoken
-from openai import AsyncOpenAI
-
-from kodit.embedding.embedding_provider.embedding_provider import (
-    EmbeddingProvider,
-    EmbeddingRequest,
-    EmbeddingResponse,
-    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[EmbeddingRequest]
-    ) -> AsyncGenerator[list[EmbeddingResponse], None]:
-        """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)
-
-        async def process_batch(
-            data: list[EmbeddingRequest],
-        ) -> list[EmbeddingResponse]:
-            async with sem:
-                try:
-                    response = await self.openai_client.embeddings.create(
-                        model=self.model_name,
-                        input=[i.text for i in data],
-                    )
-                    return [
-                        EmbeddingResponse(
-                            id=item.id,
-                            embedding=embedding.embedding,
-                        )
-                        for item, embedding in zip(data, response.data, strict=True)
-                    ]
-                except Exception as e:
-                    self.log.exception("Error embedding batch", error=str(e))
-                    return []
-
-        # Create tasks for all batches
-        tasks = [process_batch(batch) for batch in batched_data]
-
-        # Process all batches and yield results as they complete
-        for task in asyncio.as_completed(tasks):
-            result = await task
-            yield result

kodit/embedding/local_vector_search_service.py

@@ -1,87 +0,0 @@
-"""Local vector search."""
-
-from collections.abc import AsyncGenerator
-
-import structlog
-import tiktoken
-
-from kodit.embedding.embedding_models import Embedding, EmbeddingType
-from kodit.embedding.embedding_provider.embedding_provider import (
-    EmbeddingProvider,
-    EmbeddingRequest,
-)
-from kodit.embedding.embedding_repository import EmbeddingRepository
-from kodit.embedding.vector_search_service import (
-    IndexResult,
-    VectorSearchRequest,
-    VectorSearchResponse,
-    VectorSearchService,
-)
-
-
-class LocalVectorSearchService(VectorSearchService):
-    """Local vector search."""
-
-    def __init__(
-        self,
-        embedding_repository: EmbeddingRepository,
-        embedding_provider: EmbeddingProvider,
-        embedding_type: EmbeddingType = EmbeddingType.CODE,
-    ) -> 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")
-        self.embedding_type = embedding_type
-
-    async def index(
-        self, data: list[VectorSearchRequest]
-    ) -> AsyncGenerator[list[IndexResult], None]:
-        """Embed a list of documents."""
-        if not data or len(data) == 0:
-            return
-
-        requests = [EmbeddingRequest(id=doc.snippet_id, text=doc.text) for doc in data]
-
-        async for batch in self.embedding_provider.embed(requests):
-            for result in batch:
-                await self.embedding_repository.create_embedding(
-                    Embedding(
-                        snippet_id=result.id,
-                        embedding=result.embedding,
-                        type=self.embedding_type,
-                    )
-                )
-                yield [IndexResult(snippet_id=result.id)]
-
-    async def retrieve(self, query: str, top_k: int = 10) -> list[VectorSearchResponse]:
-        """Query the embedding model."""
-        # Build a single-item request and collect its embedding.
-        req = EmbeddingRequest(id=0, text=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, top_k
-        )
-        return [
-            VectorSearchResponse(snippet_id, 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/vector_search_service.py

@@ -1,55 +0,0 @@
-"""Embedding service."""
-
-from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator
-from typing import NamedTuple
-
-from kodit.embedding.embedding_models import EmbeddingType
-
-
-class VectorSearchResponse(NamedTuple):
-    """Embedding result."""
-
-    snippet_id: int
-    score: float
-
-
-class VectorSearchRequest(NamedTuple):
-    """Input for embedding."""
-
-    snippet_id: int
-    text: str
-
-
-class IndexResult(NamedTuple):
-    """Result of indexing."""
-
-    snippet_id: int
-
-
-class VectorSearchService(ABC):
-    """Semantic search service interface."""
-
-    @abstractmethod
-    def index(
-        self, data: list[VectorSearchRequest]
-    ) -> AsyncGenerator[list[IndexResult], 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."""
-
-    @abstractmethod
-    async def has_embedding(
-        self, snippet_id: int, embedding_type: EmbeddingType
-    ) -> bool:
-        """Check if a snippet has an embedding."""

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,36 +0,0 @@
-"""Enrichment provider."""
-
-from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator
-from dataclasses import dataclass
-
-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.
-"""
-
-
-@dataclass
-class EnrichmentRequest:
-    """Enrichment request."""
-
-    snippet_id: int
-    text: str
-
-
-@dataclass
-class EnrichmentResponse:
-    """Enrichment response."""
-
-    snippet_id: int
-    text: str
-
-
-class EnrichmentProvider(ABC):
-    """Enrichment provider."""
-
-    @abstractmethod
-    def enrich(
-        self, data: list[EnrichmentRequest]
-    ) -> AsyncGenerator[EnrichmentResponse, None]:
-        """Enrich a list of strings."""

kodit/enrichment/enrichment_provider/openai_enrichment_provider.py

@@ -1,79 +0,0 @@
-"""OpenAI embedding service."""
-
-import asyncio
-from collections.abc import AsyncGenerator
-
-import structlog
-import tiktoken
-from openai import AsyncOpenAI
-
-from kodit.enrichment.enrichment_provider.enrichment_provider import (
-    ENRICHMENT_SYSTEM_PROMPT,
-    EnrichmentProvider,
-    EnrichmentRequest,
-    EnrichmentResponse,
-)
-
-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[EnrichmentRequest]
-    ) -> AsyncGenerator[EnrichmentResponse, None]:
-        """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)
-
-        async def process_data(data: EnrichmentRequest) -> EnrichmentResponse:
-            async with sem:
-                if not data.text:
-                    return EnrichmentResponse(
-                        snippet_id=data.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": data.text},
-                        ],
-                    )
-                    return EnrichmentResponse(
-                        snippet_id=data.snippet_id,
-                        text=response.choices[0].message.content or "",
-                    )
-                except Exception as e:
-                    self.log.exception("Error enriching data", error=str(e))
-                    return EnrichmentResponse(
-                        snippet_id=data.snippet_id,
-                        text="",
-                    )
-
-        # Create tasks for all data
-        tasks = [process_data(snippet) for snippet in data]
-
-        # Process all data and yield results as they complete
-        for task in asyncio.as_completed(tasks):
-            yield await task

kodit/enrichment/enrichment_service.py

@@ -1,45 +0,0 @@
-"""Enrichment service."""
-
-from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator
-
-from kodit.enrichment.enrichment_provider.enrichment_provider import (
-    EnrichmentProvider,
-    EnrichmentRequest,
-    EnrichmentResponse,
-)
-
-
-class EnrichmentService(ABC):
-    """Enrichment service."""
-
-    @abstractmethod
-    def enrich(
-        self, data: list[EnrichmentRequest]
-    ) -> AsyncGenerator[EnrichmentResponse, None]:
-        """Enrich a list of strings."""
-
-
-class NullEnrichmentService(EnrichmentService):
-    """Null enrichment service."""
-
-    async def enrich(
-        self, data: list[EnrichmentRequest]
-    ) -> AsyncGenerator[EnrichmentResponse, None]:
-        """Enrich a list of strings."""
-        for request in data:
-            yield EnrichmentResponse(snippet_id=request.snippet_id, text="")
-
-
-class LLMEnrichmentService(EnrichmentService):
-    """Enrichment service using an LLM."""
-
-    def __init__(self, enrichment_provider: EnrichmentProvider) -> None:
-        """Initialize the enrichment service."""
-        self.enrichment_provider = enrichment_provider
-
-    def enrich(
-        self, data: list[EnrichmentRequest]
-    ) -> AsyncGenerator[EnrichmentResponse, None]:
-        """Enrich a list of snippets."""
-        return self.enrichment_provider.enrich(data)

kodit/indexing/__init__.py

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