kodit 0.1.14__py3-none-any.whl → 0.1.16__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.14'
-__version_tuple__ = version_tuple = (0, 1, 14)
+__version__ = version = '0.1.16'
+__version_tuple__ = version_tuple = (0, 1, 16)

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,17 +10,17 @@ 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.enrichment.enrichment_factory import enrichment_factory
 from kodit.indexing.indexing_repository import IndexRepository
-from kodit.indexing.indexing_service import IndexService
+from kodit.indexing.indexing_service import IndexService, SearchRequest
 from kodit.log import configure_logging, configure_telemetry, log_event
-from kodit.search.search_repository import SearchRepository
-from kodit.search.search_service import SearchRequest, SearchService
 from kodit.source.source_repository import SourceRepository
 from kodit.source.source_service import SourceService
 
@@ -68,10 +68,16 @@ 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),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
     )
 
     if not sources:
@@ -128,11 +134,20 @@ async def code(
 
     This works best if your query is code.
     """
-    repository = SearchRepository(session)
-    service = SearchService(
-        repository,
-        app_context.get_data_dir(),
-        embedding_service=embedding_factory(app_context.get_default_openai_client()),
+    source_repository = SourceRepository(session)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
+    repository = IndexRepository(session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
     )
 
     snippets = await service.search(SearchRequest(code_query=query, top_k=top_k))
@@ -144,6 +159,7 @@ async def code(
     for snippet in snippets:
         click.echo("-" * 80)
         click.echo(f"{snippet.uri}")
+        click.echo(f"Original scores: {snippet.original_scores}")
         click.echo(snippet.content)
         click.echo("-" * 80)
         click.echo()
@@ -161,11 +177,20 @@ async def keyword(
     top_k: int,
 ) -> None:
     """Search for snippets using keyword search."""
-    repository = SearchRepository(session)
-    service = SearchService(
-        repository,
-        app_context.get_data_dir(),
-        embedding_service=embedding_factory(app_context.get_default_openai_client()),
+    source_repository = SourceRepository(session)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
+    repository = IndexRepository(session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
     )
 
     snippets = await service.search(SearchRequest(keywords=keywords, top_k=top_k))
@@ -177,6 +202,53 @@ async def keyword(
     for snippet in snippets:
         click.echo("-" * 80)
         click.echo(f"{snippet.uri}")
+        click.echo(f"Original scores: {snippet.original_scores}")
+        click.echo(snippet.content)
+        click.echo("-" * 80)
+        click.echo()
+
+
+@search.command()
+@click.argument("query")
+@click.option("--top-k", default=10, help="Number of snippets to retrieve")
+@with_app_context
+@with_session
+async def text(
+    session: AsyncSession,
+    app_context: AppContext,
+    query: str,
+    top_k: int,
+) -> None:
+    """Search for snippets using semantic text search.
+
+    This works best if your query is text.
+    """
+    source_repository = SourceRepository(session)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
+    repository = IndexRepository(session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
+    )
+
+    snippets = await service.search(SearchRequest(text_query=query, top_k=top_k))
+
+    if len(snippets) == 0:
+        click.echo("No snippets found")
+        return
+
+    for snippet in snippets:
+        click.echo("-" * 80)
+        click.echo(f"{snippet.uri}")
+        click.echo(f"Original scores: {snippet.original_scores}")
         click.echo(snippet.content)
         click.echo("-" * 80)
         click.echo()
@@ -186,28 +258,44 @@ async def keyword(
 @click.option("--top-k", default=10, help="Number of snippets to retrieve")
 @click.option("--keywords", required=True, help="Comma separated list of keywords")
 @click.option("--code", required=True, help="Semantic code search query")
+@click.option("--text", required=True, help="Semantic text search query")
 @with_app_context
 @with_session
-async def hybrid(
+async def hybrid(  # noqa: PLR0913
     session: AsyncSession,
     app_context: AppContext,
     top_k: int,
     keywords: str,
     code: str,
+    text: str,
 ) -> None:
     """Search for snippets using hybrid search."""
-    repository = SearchRepository(session)
-    service = SearchService(
-        repository,
-        app_context.get_data_dir(),
-        embedding_service=embedding_factory(app_context.get_default_openai_client()),
+    source_repository = SourceRepository(session)
+    source_service = SourceService(app_context.get_clone_dir(), source_repository)
+    repository = IndexRepository(session)
+    service = IndexService(
+        repository=repository,
+        source_service=source_service,
+        keyword_search_provider=keyword_search_factory(app_context, session),
+        code_search_service=embedding_factory(
+            task_name="code", app_context=app_context, session=session
+        ),
+        text_search_service=embedding_factory(
+            task_name="text", app_context=app_context, session=session
+        ),
+        enrichment_service=enrichment_factory(app_context),
     )
 
     # Parse keywords into a list of strings
     keywords_list = [k.strip().lower() for k in keywords.split(",")]
 
     snippets = await service.search(
-        SearchRequest(keywords=keywords_list, code_query=code, top_k=top_k)
+        SearchRequest(
+            text_query=text,
+            keywords=keywords_list,
+            code_query=code,
+            top_k=top_k,
+        )
     )
 
     if len(snippets) == 0:
@@ -217,6 +305,7 @@ async def hybrid(
     for snippet in snippets:
         click.echo("-" * 80)
         click.echo(f"{snippet.uri}")
+        click.echo(f"Original scores: {snippet.original_scores}")
         click.echo(snippet.content)
         click.echo("-" * 80)
         click.echo()

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(
+    task_name: str, 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(task_name, 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."""