kodit 0.4.3__py3-none-any.whl → 0.5.1__py3-none-any.whl

kodit/infrastructure/bm25/vectorchord_bm25_repository.py

@@ -1,9 +1,9 @@
 """VectorChord BM25 repository implementation."""
 
-from typing import Any
+from collections.abc import Callable
 
 import structlog
-from sqlalchemy import Result, TextClause, bindparam, text
+from sqlalchemy import bindparam, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.domain.services.bm25_service import BM25Repository
@@ -13,6 +13,7 @@ from kodit.domain.value_objects import (
     SearchRequest,
     SearchResult,
 )
+from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
 
 TABLE_NAME = "vectorchord_bm25_documents"
 INDEX_NAME = f"{TABLE_NAME}_idx"
@@ -29,13 +30,17 @@ SET search_path TO
 CREATE_BM25_TABLE = f"""
 CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
     id SERIAL PRIMARY KEY,
-    snippet_id BIGINT NOT NULL,
+    snippet_id VARCHAR(255) NOT NULL,
     passage TEXT NOT NULL,
     embedding bm25vector,
     UNIQUE(snippet_id)
 )
 """
-
+CHECK_EXISTING_IDS = f"""
+    SELECT snippet_id
+    FROM {TABLE_NAME}
+    WHERE snippet_id = ANY(:snippet_ids)
+"""  # noqa: S608
 CREATE_BM25_INDEX = f"""
 CREATE INDEX IF NOT EXISTS {INDEX_NAME}
 ON {TABLE_NAME}
@@ -103,14 +108,14 @@ WHERE snippet_id IN :snippet_ids
 class VectorChordBM25Repository(BM25Repository):
     """VectorChord BM25 repository implementation."""
 
-    def __init__(self, session: AsyncSession) -> None:
+    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
         """Initialize the VectorChord BM25 repository.
 
         Args:
             session: The SQLAlchemy async session to use for database operations
 
         """
-        self.__session = session
+        self.session_factory = session_factory
         self._initialized = False
         self.log = structlog.get_logger(__name__)
 
@@ -127,41 +132,39 @@ class VectorChordBM25Repository(BM25Repository):
 
     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 with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            await session.execute(text(CREATE_VCHORD_EXTENSION))
+            await session.execute(text(CREATE_PG_TOKENIZER))
+            await session.execute(text(CREATE_VCHORD_BM25))
+            await session.execute(text(SET_SEARCH_PATH))
 
     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 with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            # Check if tokenizer exists in the catalog
+            result = await session.execute(text(TOKENIZER_NAME_CHECK_QUERY))
+            if result.scalar_one_or_none() is None:
+                # Tokenizer doesn't exist, create it
+                await session.execute(text(LOAD_TOKENIZER))
 
     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 with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            await session.execute(text(CREATE_BM25_TABLE))
+            await session.execute(text(CREATE_BM25_INDEX))
+
+    async def _get_existing_ids(self, snippet_ids: list[str]) -> set[int]:
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            result = await session.execute(
+                text(CHECK_EXISTING_IDS), {"snippet_ids": snippet_ids}
+            )
+            return {row[0] for row in result.fetchall()}
 
     async def index_documents(self, request: IndexRequest) -> None:
         """Index documents for BM25 search."""
+        if not self._initialized:
+            await self._initialize()
+
         # Filter out any documents that don't have a snippet_id or text
         valid_documents = [
             doc
@@ -173,21 +176,35 @@ class VectorChordBM25Repository(BM25Repository):
             self.log.warning("Corpus is empty, skipping bm25 index")
             return
 
-        # Execute inserts
-        await self._execute(
-            text(INSERT_QUERY),
-            [
-                {"snippet_id": doc.snippet_id, "passage": doc.text}
-                for doc in valid_documents
-            ],
+        # Filter out documents that have already been indexed
+        existing_ids = await self._get_existing_ids(
+            [doc.snippet_id for doc in valid_documents]
         )
+        valid_documents = [
+            doc for doc in valid_documents if doc.snippet_id not in existing_ids
+        ]
 
-        # Tokenize the new documents with schema qualification
-        await self._execute(text(UPDATE_QUERY))
-        await self._commit()
+        if not valid_documents:
+            self.log.info("No new documents to index")
+            return
+
+        # Execute inserts
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            await session.execute(
+                text(INSERT_QUERY),
+                [
+                    {"snippet_id": doc.snippet_id, "passage": doc.text}
+                    for doc in valid_documents
+                ],
+            )
+
+            # Tokenize the new documents with schema qualification
+            await session.execute(text(UPDATE_QUERY))
 
     async def search(self, request: SearchRequest) -> list[SearchResult]:
         """Search documents using BM25."""
+        if not self._initialized:
+            await self._initialize()
         if not request.query or request.query.strip() == "":
             return []
 
@@ -203,22 +220,21 @@ class VectorChordBM25Repository(BM25Repository):
                 limit=request.top_k,
             )
 
-        try:
-            result = await self._execute(sql)
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            result = await session.execute(sql)
             rows = result.mappings().all()
 
             return [
                 SearchResult(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
 
     async def delete_documents(self, request: DeleteRequest) -> None:
         """Delete documents from the index."""
-        await self._execute(
-            text(DELETE_QUERY).bindparams(bindparam("snippet_ids", expanding=True)),
-            {"snippet_ids": request.snippet_ids},
-        )
-        await self._commit()
+        if not self._initialized:
+            await self._initialize()
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            await session.execute(
+                text(DELETE_QUERY).bindparams(bindparam("snippet_ids", expanding=True)),
+                {"snippet_ids": request.snippet_ids},
+            )