kodit 0.1.14__py3-none-any.whl → 0.1.15__py3-none-any.whl

kodit/embedding/embedding.py

@@ -1,203 +0,0 @@
-"""Embedding service."""
-
-import asyncio
-import os
-from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator
-from typing import NamedTuple
-
-import structlog
-import tiktoken
-from openai import AsyncOpenAI
-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 EmbeddingInput(NamedTuple):
-    """Input for embedding."""
-
-    id: int
-    text: str
-
-
-class EmbeddingOutput(NamedTuple):
-    """Output for embedding."""
-
-    id: int
-    embedding: list[float]
-
-
-class Embedder(ABC):
-    """Embedder interface."""
-
-    @abstractmethod
-    def embed(
-        self, data: list[EmbeddingInput]
-    ) -> AsyncGenerator[EmbeddingOutput, 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
-    def query(self, data: list[str]) -> AsyncGenerator[list[float], None]:
-        """Query the embedding model."""
-
-
-def embedding_factory(openai_client: AsyncOpenAI | None = None) -> Embedder:
-    """Create an embedding service."""
-    if openai_client is not None:
-        return OpenAIEmbedder(openai_client)
-    return LocalEmbedder(model_name=TINY)
-
-
-class LocalEmbedder(Embedder):
-    """Local embedder."""
-
-    def __init__(self, model_name: str) -> None:
-        """Initialize the local embedder."""
-        self.log = structlog.get_logger(__name__)
-        self.log.info("Creating local embedder", model_name=model_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
-            self.embedding_model = SentenceTransformer(
-                self.model_name,
-                trust_remote_code=True,
-                device="cpu",  # Force CPU so we don't have to install accelerate, etc.
-            )
-        return self.embedding_model
-
-    async def embed(
-        self, data: list[EmbeddingInput]
-    ) -> AsyncGenerator[EmbeddingOutput, None]:
-        """Embed a list of documents."""
-        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
-            )
-            for i, x in zip(batch, embeddings, strict=False):
-                yield EmbeddingOutput(i.id, [float(y) for y in x])
-
-    async def query(self, data: list[str]) -> AsyncGenerator[list[float], None]:
-        """Query the embedding model."""
-        model = self._model()
-        embeddings = model.encode(data, show_progress_bar=False, batch_size=4)
-        for embedding in embeddings:
-            yield [float(x) for x in embedding]
-
-
-OPENAI_MAX_EMBEDDING_SIZE = 8192
-OPENAI_NUM_PARALLEL_TASKS = 10
-
-
-def _split_sub_batches(
-    encoding: tiktoken.Encoding, data: list[EmbeddingInput]
-) -> list[list[EmbeddingInput]]:
-    """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))
-
-            if item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
-                log.warning("Skipping too long snippet", snippet=data_to_process.pop(0))
-                continue
-
-            if current_tokens + item_tokens > OPENAI_MAX_EMBEDDING_SIZE:
-                break
-
-            next_batch.append(data_to_process.pop(0))
-            current_tokens += item_tokens
-
-        if next_batch:
-            result.append(next_batch)
-
-    return result
-
-
-class OpenAIEmbedder(Embedder):
-    """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.log.info("Creating OpenAI embedder", model_name=model_name)
-        self.openai_client = openai_client
-        self.encoding = tiktoken.encoding_for_model(model_name)
-        self.log = structlog.get_logger(__name__)
-
-    async def embed(
-        self,
-        data: list[EmbeddingInput],
-    ) -> AsyncGenerator[EmbeddingOutput, 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(batch: list[EmbeddingInput]) -> list[EmbeddingOutput]:
-            async with sem:
-                try:
-                    response = await self.openai_client.embeddings.create(
-                        model="text-embedding-3-small",
-                        input=[i.text for i in batch],
-                    )
-                    return [
-                        EmbeddingOutput(i.id, x.embedding)
-                        for i, x in zip(batch, response.data, strict=False)
-                    ]
-                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):
-            embeddings = await task
-            for e in embeddings:
-                yield e
-
-    async def query(self, data: list[str]) -> AsyncGenerator[list[float], None]:
-        """Query the embedding model."""
-        async for e in self.embed(
-            [EmbeddingInput(i, text) for i, text in enumerate(data)]
-        ):
-            yield e.embedding

kodit-0.1.14.dist-info/RECORD

@@ -1,44 +0,0 @@
-kodit/.gitignore,sha256=ztkjgRwL9Uud1OEi36hGQeDGk3OLK1NfDEO8YqGYy8o,11
-kodit/__init__.py,sha256=aEKHYninUq1yh6jaNfvJBYg-6fenpN132nJt1UU6Jxs,59
-kodit/_version.py,sha256=O_r2EWoixTKREu-RyeL8e93UHfqprj1LCIlwiWXfHcg,513
-kodit/app.py,sha256=Mr5BFHOHx5zppwjC4XPWVvHjwgl1yrKbUjTWXKubJQM,891
-kodit/cli.py,sha256=VLoXFS1xJnQ0TLy3_cO8-B9tCb4NJHiYPfzZtHxpgRY,7784
-kodit/config.py,sha256=TDcLt6fiJn9cI1PoO5AqBqsL_Bxmm9JV5GqRxhj1tLw,4202
-kodit/database.py,sha256=kekSdyEATdb47jxzQemkSOXMNOwnUwmVVTpn9hYaDK8,2356
-kodit/log.py,sha256=HU1OmuxO4FcVw61k4WW7Y4WM7BrDaeplw1PcBHhuIZY,5434
-kodit/mcp.py,sha256=ot5CIH240mSXK3sJcxTf4lBfthq0tcMS8XBGTaHY-n8,5088
-kodit/middleware.py,sha256=I6FOkqG9-8RH5kR1-0ZoQWfE4qLCB8lZYv8H_OCH29o,2714
-kodit/bm25/__init__.py,sha256=j8zyriNWhbwE5Lbybzg1hQAhANlU9mKHWw4beeUR6og,19
-kodit/bm25/bm25.py,sha256=JtgJfsHz-2SHx96zxWjkPFSH7fXkahFMp01cDwl4YBg,2298
-kodit/embedding/__init__.py,sha256=h9NXzDA1r-K23nvBajBV-RJzHJN0p3UJ7UQsmdnOoRw,24
-kodit/embedding/embedding.py,sha256=EMJpHK8ICZk_FjiO9Aqr2IO20qkGOmj_PfA1hyfI7Vk,6745
-kodit/embedding/embedding_models.py,sha256=rN90vSs86dYiqoawcp8E9jtwY31JoJXYfaDlsJK7uqc,656
-kodit/indexing/__init__.py,sha256=cPyi2Iej3G1JFWlWr7X80_UrsMaTu5W5rBwgif1B3xo,75
-kodit/indexing/indexing_models.py,sha256=6NX9HVcj6Pu9ePwHC7n-PWSyAgukpJq0nCNmUIigtbo,1282
-kodit/indexing/indexing_repository.py,sha256=7bkAiBwtr3qlkdhNIalwMwbxezVz_RQGOhLVWPKHwNk,5506
-kodit/indexing/indexing_service.py,sha256=VGfKgbkYEAYP_gIubvhMxo3yThT20ndS5xdg2LxwRgA,6685
-kodit/migrations/README,sha256=ISVtAOvqvKk_5ThM5ioJE-lMkvf9IbknFUFVU_vPma4,58
-kodit/migrations/__init__.py,sha256=lP5MuwlyWRMO6UcDWnQcQ3G-GYHcFb6rl9gYPHJ1sjo,40
-kodit/migrations/env.py,sha256=w1M7OZh-ZeR2dPHS0ByXAUxQjfZQ8xIzMseWuzLDTWw,2469
-kodit/migrations/script.py.mako,sha256=zWziKtiwYKEWuwPV_HBNHwa9LCT45_bi01-uSNFaOOE,703
-kodit/migrations/versions/7c3bbc2ab32b_add_embeddings_table.py,sha256=-61qol9PfQKILCDQRA5jEaats9aGZs9Wdtp-j-38SF4,1644
-kodit/migrations/versions/85155663351e_initial.py,sha256=Cg7zlF871o9ShV5rQMQ1v7hRV7fI59veDY9cjtTrs-8,3306
-kodit/migrations/versions/__init__.py,sha256=9-lHzptItTzq_fomdIRBegQNm4Znx6pVjwD4MiqRIdo,36
-kodit/search/__init__.py,sha256=4QbdjbrlhNKMovmuKHxJnUeZT7KNjTTFU0GdnuwUHdQ,36
-kodit/search/search_repository.py,sha256=r1fkV6-cy9BKsy5J4WTHaY_FcjMaT1PV5qqqq0gvjZw,5833
-kodit/search/search_service.py,sha256=KePkqCAc3CUcrpNsbDc5DqbF6W2m0TG6TDa9-VSJZS0,4227
-kodit/snippets/__init__.py,sha256=-2coNoCRjTixU9KcP6alpmt7zqf37tCRWH3D7FPJ8dg,48
-kodit/snippets/method_snippets.py,sha256=EVHhSNWahAC5nSXv9fWVFJY2yq25goHdCSCuENC07F8,4145
-kodit/snippets/snippets.py,sha256=mwN0bM1Msu8ZeEsUHyQ7tx3Hj3vZsm8G7Wu4eWSkLY8,1539
-kodit/snippets/languages/__init__.py,sha256=Bj5KKZSls2MQ8ZY1S_nHg447MgGZW-2WZM-oq6vjwwA,1187
-kodit/snippets/languages/csharp.scm,sha256=gbBN4RiV1FBuTJF6orSnDFi8H9JwTw-d4piLJYsWUsc,222
-kodit/snippets/languages/python.scm,sha256=ee85R9PBzwye3IMTE7-iVoKWd_ViU3EJISTyrFGrVeo,429
-kodit/source/__init__.py,sha256=1NTZyPdjThVQpZO1Mp1ColVsS7sqYanOVLqnoqV9Ipo,83
-kodit/source/source_models.py,sha256=xb42CaNDO1CUB8SIW-xXMrB6Ji8cFw-yeJ550xBEg9Q,2398
-kodit/source/source_repository.py,sha256=0EksMpoLzdkfe8S4eeCm4Sf7TuxsOzOzaF4BBsMYo-4,3163
-kodit/source/source_service.py,sha256=qBV9FCFQbJppeFrVo4uMgvC_mzWRIKldymp5yqLx9pw,9255
-kodit-0.1.14.dist-info/METADATA,sha256=acFpcf0ODyUSnA1hg4BPlLexpOEh-0yuaqsaWUNopOs,2349
-kodit-0.1.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-kodit-0.1.14.dist-info/entry_points.txt,sha256=hoTn-1aKyTItjnY91fnO-rV5uaWQLQ-Vi7V5et2IbHY,40
-kodit-0.1.14.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-kodit-0.1.14.dist-info/RECORD,,