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

kodit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.13'
-__version_tuple__ = version_tuple = (0, 1, 13)
+__version__ = version = '0.1.15'
+__version_tuple__ = version_tuple = (0, 1, 15)

kodit/bm25/keyword_search_factory.py

@@ -0,0 +1,17 @@
+"""Factory for creating keyword search providers."""
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.bm25.keyword_search_service import KeywordSearchProvider
+from kodit.bm25.local_bm25 import BM25Service
+from kodit.bm25.vectorchord_bm25 import VectorChordBM25
+from kodit.config import AppContext
+
+
+def keyword_search_factory(
+    app_context: AppContext, session: AsyncSession
+) -> KeywordSearchProvider:
+    """Create a keyword search provider."""
+    if app_context.default_search.provider == "vectorchord":
+        return VectorChordBM25(session=session)
+    return BM25Service(data_dir=app_context.get_data_dir())

kodit/bm25/keyword_search_service.py

@@ -0,0 +1,34 @@
+"""Keyword search service."""
+
+from abc import ABC, abstractmethod
+from typing import NamedTuple
+
+
+class BM25Document(NamedTuple):
+    """BM25 document."""
+
+    snippet_id: int
+    text: str
+
+
+class BM25Result(NamedTuple):
+    """BM25 result."""
+
+    snippet_id: int
+    score: float
+
+
+class KeywordSearchProvider(ABC):
+    """Interface for keyword search providers."""
+
+    @abstractmethod
+    async def index(self, corpus: list[BM25Document]) -> None:
+        """Index a new corpus."""
+
+    @abstractmethod
+    async def retrieve(self, query: str, top_k: int = 2) -> list[BM25Result]:
+        """Retrieve from the index."""
+
+    @abstractmethod
+    async def delete(self, snippet_ids: list[int]) -> None:
+        """Delete documents from the index."""

kodit/bm25/{bm25.py → local_bm25.py}

@@ -1,23 +1,36 @@
-"""BM25 service."""
+"""Locally hosted BM25 service primarily for use with SQLite."""
 
+import json
 from pathlib import Path
 
+import aiofiles
 import bm25s
 import Stemmer
 import structlog
 from bm25s.tokenization import Tokenized
 
+from kodit.bm25.keyword_search_service import (
+    BM25Document,
+    BM25Result,
+    KeywordSearchProvider,
+)
 
-class BM25Service:
-    """Service for BM25."""
+SNIPPET_IDS_FILE = "snippet_ids.jsonl"
+
+
+class BM25Service(KeywordSearchProvider):
+    """LocalBM25 service."""
 
     def __init__(self, data_dir: Path) -> None:
         """Initialize the BM25 service."""
         self.log = structlog.get_logger(__name__)
         self.index_path = data_dir / "bm25s_index"
+        self.snippet_ids: list[int] = []
         try:
             self.log.debug("Loading BM25 index")
             self.retriever = bm25s.BM25.load(self.index_path, mmap=True)
+            with Path(self.index_path / SNIPPET_IDS_FILE).open() as f:
+                self.snippet_ids = json.load(f)
         except FileNotFoundError:
             self.log.debug("BM25 index not found, creating new index")
             self.retriever = bm25s.BM25()
@@ -33,28 +46,34 @@ class BM25Service:
             show_progress=True,
         )
 
-    def index(self, corpus: list[str]) -> None:
+    async def index(self, corpus: list[BM25Document]) -> None:
         """Index a new corpus."""
         self.log.debug("Indexing corpus")
-        vocab = self._tokenize(corpus)
+        vocab = self._tokenize([doc.text for doc in corpus])
         self.retriever = bm25s.BM25()
         self.retriever.index(vocab, show_progress=False)
         self.retriever.save(self.index_path)
+        self.snippet_ids = self.snippet_ids + [doc.snippet_id for doc in corpus]
+        async with aiofiles.open(self.index_path / SNIPPET_IDS_FILE, "w") as f:
+            await f.write(json.dumps(self.snippet_ids))
 
-    def retrieve(
-        self, doc_ids: list[int], query: str, top_k: int = 2
-    ) -> list[tuple[int, float]]:
+    async def retrieve(self, query: str, top_k: int = 2) -> list[BM25Result]:
         """Retrieve from the index."""
         if top_k == 0:
             self.log.warning("Top k is 0, returning empty list")
             return []
-        if len(doc_ids) == 0:
-            self.log.warning("No documents to retrieve from, returning empty list")
+
+        # Get the number of documents in the index
+        num_docs = self.retriever.scores["num_docs"]
+        if num_docs == 0:
             return []
 
-        top_k = min(top_k, len(self.retriever.scores))
+        # Adjust top_k to not exceed corpus size
+        top_k = min(top_k, num_docs)
         self.log.debug(
-            "Retrieving from index", query=query, top_k=top_k, num_docs=len(doc_ids)
+            "Retrieving from index",
+            query=query,
+            top_k=top_k,
         )
 
         query_tokens = self._tokenize([query])
@@ -62,10 +81,17 @@ class BM25Service:
         self.log.debug("Query tokens", query_tokens=query_tokens)
 
         results, scores = self.retriever.retrieve(
-            query_tokens=query_tokens, corpus=doc_ids, k=top_k
+            query_tokens=query_tokens,
+            corpus=self.snippet_ids,
+            k=top_k,
         )
         self.log.debug("Raw results", results=results, scores=scores)
         return [
-            (int(result), float(score))
+            BM25Result(snippet_id=int(result), score=float(score))
             for result, score in zip(results[0], scores[0], strict=False)
+            if score > 0.0
         ]
+
+    async def delete(self, snippet_ids: list[int]) -> None:  # noqa: ARG002
+        """Delete documents from the index."""
+        self.log.warning("Deletion not supported for local BM25 index")

kodit/bm25/vectorchord_bm25.py

@@ -0,0 +1,193 @@
+"""VectorChord repository for document operations."""
+
+from typing import Any
+
+from sqlalchemy import Result, TextClause, bindparam, text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.bm25.keyword_search_service import (
+    BM25Document,
+    BM25Result,
+    KeywordSearchProvider,
+)
+
+TABLE_NAME = "vectorchord_bm25_documents"
+INDEX_NAME = f"{TABLE_NAME}_idx"
+TOKENIZER_NAME = "bert"
+
+# SQL statements
+CREATE_VCHORD_EXTENSION = "CREATE EXTENSION IF NOT EXISTS vchord CASCADE;"
+CREATE_PG_TOKENIZER = "CREATE EXTENSION IF NOT EXISTS pg_tokenizer CASCADE;"
+CREATE_VCHORD_BM25 = "CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE;"
+SET_SEARCH_PATH = """
+SET search_path TO
+    "$user", public, bm25_catalog, pg_catalog, information_schema, tokenizer_catalog;
+"""
+CREATE_BM25_TABLE = f"""
+CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
+    id SERIAL PRIMARY KEY,
+    snippet_id BIGINT NOT NULL,
+    passage TEXT NOT NULL,
+    embedding bm25vector,
+    UNIQUE(snippet_id)
+)
+"""
+
+CREATE_BM25_INDEX = f"""
+CREATE INDEX IF NOT EXISTS {INDEX_NAME}
+ON {TABLE_NAME}
+USING bm25 (embedding bm25_ops)
+"""
+TOKENIZER_NAME_CHECK_QUERY = (
+    f"SELECT 1 FROM tokenizer_catalog.tokenizer WHERE name = '{TOKENIZER_NAME}'"  # noqa: S608
+)
+LOAD_TOKENIZER = """
+SELECT create_tokenizer('bert', $$
+model = "llmlingua2"
+pre_tokenizer = "unicode_segmentation"  # Unicode Standard Annex #29
+[[character_filters]]
+to_lowercase = {}                       # convert all characters to lowercase
+[[character_filters]]
+unicode_normalization = "nfkd"          # Unicode Normalization Form KD
+[[token_filters]]
+skip_non_alphanumeric = {}              # remove non-alphanumeric tokens
+[[token_filters]]
+stopwords = "nltk_english"              # remove stopwords using the nltk dictionary
+[[token_filters]]
+stemmer = "english_porter2"             # stem tokens using the English Porter2 stemmer
+$$)
+"""
+INSERT_QUERY = f"""
+INSERT INTO {TABLE_NAME} (snippet_id, passage)
+VALUES (:snippet_id, :passage)
+ON CONFLICT (snippet_id) DO UPDATE
+SET passage = EXCLUDED.passage
+"""  # noqa: S608
+UPDATE_QUERY = f"""
+UPDATE {TABLE_NAME}
+SET embedding = tokenize(passage, '{TOKENIZER_NAME}')
+"""  # noqa: S608
+SEARCH_QUERY = f"""
+    SELECT
+        snippet_id,
+        embedding <&>
+            to_bm25query('{INDEX_NAME}', tokenize(:query_text, '{TOKENIZER_NAME}'))
+    AS bm25_score
+    FROM {TABLE_NAME}
+    ORDER BY bm25_score
+    LIMIT :limit
+"""  # noqa: S608
+DELETE_QUERY = f"""
+DELETE FROM {TABLE_NAME}
+WHERE snippet_id IN :snippet_ids
+"""  # noqa: S608
+
+
+class VectorChordBM25(KeywordSearchProvider):
+    """BM25 using VectorChord."""
+
+    def __init__(
+        self,
+        session: AsyncSession,
+    ) -> None:
+        """Initialize the VectorChord BM25."""
+        self.__session = session
+        self._initialized = False
+
+    async def _initialize(self) -> None:
+        """Initialize the VectorChord environment."""
+        try:
+            await self._create_extensions()
+            await self._create_tokenizer_if_not_exists()
+            await self._create_tables()
+            self._initialized = True
+        except Exception as e:
+            msg = f"Failed to initialize VectorChord repository: {e}"
+            raise RuntimeError(msg) from e
+
+    async def _create_extensions(self) -> None:
+        """Create the necessary extensions."""
+        await self.__session.execute(text(CREATE_VCHORD_EXTENSION))
+        await self.__session.execute(text(CREATE_PG_TOKENIZER))
+        await self.__session.execute(text(CREATE_VCHORD_BM25))
+        await self.__session.execute(text(SET_SEARCH_PATH))
+        await self._commit()
+
+    async def _create_tokenizer_if_not_exists(self) -> None:
+        """Create the tokenizer if it doesn't exist."""
+        # Check if tokenizer exists in the catalog
+        result = await self.__session.execute(text(TOKENIZER_NAME_CHECK_QUERY))
+        if result.scalar_one_or_none() is None:
+            # Tokenizer doesn't exist, create it
+            await self.__session.execute(text(LOAD_TOKENIZER))
+            await self._commit()
+
+    async def _create_tables(self) -> None:
+        """Create the necessary tables in the correct order."""
+        await self.__session.execute(text(CREATE_BM25_TABLE))
+        await self.__session.execute(text(CREATE_BM25_INDEX))
+        await self._commit()
+
+    async def _execute(
+        self, query: TextClause, param_list: list[Any] | dict[str, Any] | None = None
+    ) -> Result:
+        """Execute a query."""
+        if not self._initialized:
+            await self._initialize()
+        return await self.__session.execute(query, param_list)
+
+    async def _commit(self) -> None:
+        """Commit the session."""
+        await self.__session.commit()
+
+    async def index(self, corpus: list[BM25Document]) -> None:
+        """Index a new corpus."""
+        # Filter out any documents that don't have a snippet_id or text
+        corpus = [
+            doc
+            for doc in corpus
+            if doc.snippet_id is not None and doc.text is not None and doc.text != ""
+        ]
+
+        if not corpus:
+            return
+
+        # Execute inserts
+        await self._execute(
+            text(INSERT_QUERY),
+            [{"snippet_id": doc.snippet_id, "passage": doc.text} for doc in corpus],
+        )
+
+        # Tokenize the new documents with schema qualification
+        await self._execute(text(UPDATE_QUERY))
+        await self._commit()
+
+    async def delete(self, snippet_ids: list[int]) -> None:
+        """Delete documents from the index."""
+        await self._execute(
+            text(DELETE_QUERY).bindparams(bindparam("snippet_ids", expanding=True)),
+            {"snippet_ids": snippet_ids},
+        )
+        await self._commit()
+
+    async def retrieve(
+        self,
+        query: str,
+        top_k: int = 10,
+    ) -> list[BM25Result]:
+        """Search documents using BM25 similarity."""
+        if not query or query == "":
+            return []
+
+        sql = text(SEARCH_QUERY).bindparams(query_text=query, limit=top_k)
+        try:
+            result = await self._execute(sql)
+            rows = result.mappings().all()
+
+            return [
+                BM25Result(snippet_id=row["snippet_id"], score=row["bm25_score"])
+                for row in rows
+            ]
+        except Exception as e:
+            msg = f"Error during BM25 search: {e}"
+            raise RuntimeError(msg) from e

kodit/cli.py

@@ -10,12 +10,13 @@ import uvicorn
 from pytable_formatter import Cell, Table
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from kodit.bm25.keyword_search_factory import keyword_search_factory
 from kodit.config import (
     AppContext,
     with_app_context,
     with_session,
 )
-from kodit.embedding.embedding import embedding_factory
+from kodit.embedding.embedding_factory import embedding_factory
 from kodit.indexing.indexing_repository import IndexRepository
 from kodit.indexing.indexing_service import IndexService
 from kodit.log import configure_logging, configure_telemetry, log_event
@@ -68,10 +69,12 @@ async def index(
     source_service = SourceService(app_context.get_clone_dir(), source_repository)
     repository = IndexRepository(session)
     service = IndexService(
-        repository,
-        source_service,
-        app_context.get_data_dir(),
-        embedding_service=embedding_factory(app_context.get_default_openai_client()),
+        repository=repository,
+        source_service=source_service,
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        vector_search_service=embedding_factory(
+            app_context=app_context, session=session
+        ),
     )
 
     if not sources:
@@ -131,8 +134,8 @@ async def code(
     repository = SearchRepository(session)
     service = SearchService(
         repository,
-        app_context.get_data_dir(),
-        embedding_service=embedding_factory(app_context.get_default_openai_client()),
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        embedding_service=embedding_factory(app_context=app_context, session=session),
     )
 
     snippets = await service.search(SearchRequest(code_query=query, top_k=top_k))
@@ -164,8 +167,8 @@ async def keyword(
     repository = SearchRepository(session)
     service = SearchService(
         repository,
-        app_context.get_data_dir(),
-        embedding_service=embedding_factory(app_context.get_default_openai_client()),
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        embedding_service=embedding_factory(app_context=app_context, session=session),
     )
 
     snippets = await service.search(SearchRequest(keywords=keywords, top_k=top_k))
@@ -199,8 +202,8 @@ async def hybrid(
     repository = SearchRepository(session)
     service = SearchService(
         repository,
-        app_context.get_data_dir(),
-        embedding_service=embedding_factory(app_context.get_default_openai_client()),
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        embedding_service=embedding_factory(app_context=app_context, session=session),
     )
 
     # Parse keywords into a list of strings

kodit/config.py

@@ -12,14 +12,12 @@ from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from kodit.database import Database
-from kodit.embedding.embedding import TINY
 
 DEFAULT_BASE_DIR = Path.home() / ".kodit"
 DEFAULT_DB_URL = f"sqlite+aiosqlite:///{DEFAULT_BASE_DIR}/kodit.db"
 DEFAULT_LOG_LEVEL = "INFO"
 DEFAULT_LOG_FORMAT = "pretty"
 DEFAULT_DISABLE_TELEMETRY = False
-DEFAULT_EMBEDDING_MODEL_NAME = TINY
 T = TypeVar("T")
 
 
@@ -31,6 +29,12 @@ class Endpoint(BaseModel):
     base_url: str | None = None
 
 
+class Search(BaseModel):
+    """Search provides configuration for a search engine."""
+
+    provider: Literal["sqlite", "vectorchord"] = Field(default="sqlite")
+
+
 class AppContext(BaseSettings):
     """Global context for the kodit project. Provides a shared state for the app."""
 
@@ -57,6 +61,9 @@ class AppContext(BaseSettings):
             "(can be overridden by task-specific configuration)."
         ),
     )
+    default_search: Search = Field(
+        default=Search(),
+    )
     _db: Database | None = None
 
     def model_post_init(self, _: Any) -> None:

kodit/database.py

@@ -27,10 +27,12 @@ class CommonMixin:
 
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
     created_at: Mapped[datetime] = mapped_column(
-        DateTime, default=lambda: datetime.now(UTC)
+        DateTime(timezone=True), default=lambda: datetime.now(UTC)
     )
     updated_at: Mapped[datetime] = mapped_column(
-        DateTime, default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)
+        DateTime(timezone=True),
+        default=lambda: datetime.now(UTC),
+        onupdate=lambda: datetime.now(UTC),
     )
 
 

kodit/embedding/embedding_factory.py

@@ -0,0 +1,44 @@
+"""Embedding service."""
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.config import AppContext
+from kodit.embedding.embedding_provider.local_embedding_provider import (
+    CODE,
+    LocalEmbeddingProvider,
+)
+from kodit.embedding.embedding_provider.openai_embedding_provider import (
+    OpenAIEmbeddingProvider,
+)
+from kodit.embedding.embedding_repository import EmbeddingRepository
+from kodit.embedding.local_vector_search_service import LocalVectorSearchService
+from kodit.embedding.vector_search_service import (
+    VectorSearchService,
+)
+from kodit.embedding.vectorchord_vector_search_service import (
+    VectorChordVectorSearchService,
+)
+
+
+def embedding_factory(
+    app_context: AppContext, session: AsyncSession
+) -> VectorSearchService:
+    """Create an embedding service."""
+    embedding_repository = EmbeddingRepository(session=session)
+    embedding_provider = None
+    openai_client = app_context.get_default_openai_client()
+    if openai_client is not None:
+        embedding_provider = OpenAIEmbeddingProvider(openai_client=openai_client)
+    else:
+        embedding_provider = LocalEmbeddingProvider(CODE)
+
+    if app_context.default_search.provider == "vectorchord":
+        return VectorChordVectorSearchService(session, embedding_provider)
+    if app_context.default_search.provider == "sqlite":
+        return LocalVectorSearchService(
+            embedding_repository=embedding_repository,
+            embedding_provider=embedding_provider,
+        )
+
+    msg = f"Invalid semantic search provider: {app_context.default_search.provider}"
+    raise ValueError(msg)

kodit/embedding/embedding_provider/__init__.py

@@ -0,0 +1 @@
+"""Embedding module."""

kodit/embedding/embedding_provider/embedding_provider.py

@@ -0,0 +1,53 @@
+"""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]) -> 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))
+
+            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

kodit/embedding/embedding_provider/hash_embedding_provider.py

@@ -0,0 +1,77 @@
+"""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]