kodit 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

kodit/infrastructure/api/v1/schemas/snippet.py

@@ -0,0 +1,58 @@
+"""Snippet JSON-API schemas."""
+
+from datetime import datetime
+
+from pydantic import BaseModel
+
+
+class SnippetContentSchema(BaseModel):
+    """Snippet content schema following JSON-API spec."""
+
+    value: str
+    language: str
+
+
+class GitFileSchema(BaseModel):
+    """Git file schema following JSON-API spec."""
+
+    blob_sha: str
+    path: str
+    mime_type: str
+    size: int
+
+
+class EnrichmentSchema(BaseModel):
+    """Enrichment schema following JSON-API spec."""
+
+    type: str
+    content: str
+
+
+class SnippetAttributes(BaseModel):
+    """Snippet attributes following JSON-API spec."""
+
+    created_at: datetime | None = None
+    updated_at: datetime | None = None
+    derives_from: list[GitFileSchema]
+    content: SnippetContentSchema
+    enrichments: list[EnrichmentSchema]
+
+
+class SnippetData(BaseModel):
+    """Snippet data following JSON-API spec."""
+
+    type: str = "snippet"
+    id: str
+    attributes: SnippetAttributes
+
+
+class SnippetResponse(BaseModel):
+    """Single snippet response following JSON-API spec."""
+
+    data: SnippetData
+
+
+class SnippetListResponse(BaseModel):
+    """Snippet list response following JSON-API spec."""
+
+    data: list[SnippetData]

kodit/infrastructure/api/v1/schemas/tag.py

@@ -0,0 +1,31 @@
+"""Tag JSON-API schemas."""
+
+from pydantic import BaseModel
+
+
+class TagAttributes(BaseModel):
+    """Tag attributes following JSON-API spec."""
+
+    name: str
+    target_commit_sha: str
+    is_version_tag: bool
+
+
+class TagData(BaseModel):
+    """Tag data following JSON-API spec."""
+
+    type: str = "tag"
+    id: str  # The tag name
+    attributes: TagAttributes
+
+
+class TagResponse(BaseModel):
+    """Single tag response following JSON-API spec."""
+
+    data: TagData
+
+
+class TagListResponse(BaseModel):
+    """Tag list response following JSON-API spec."""
+
+    data: list[TagData]

kodit/infrastructure/api/v1/schemas/task_status.py

@@ -0,0 +1,41 @@
+"""JSON:API schemas for task status operations."""
+
+from datetime import datetime
+
+from pydantic import BaseModel, Field
+
+
+class TaskStatusAttributes(BaseModel):
+    """Task status attributes for JSON:API responses."""
+
+    step: str = Field(..., description="Name of the task/operation")
+    state: str = Field(..., description="Current state of the task")
+    progress: float = Field(
+        default=0.0, ge=0.0, le=100.0, description="Progress percentage (0-100)"
+    )
+    total: int = Field(default=0, description="Total number of items to process")
+    current: int = Field(default=0, description="Current number of items processed")
+    created_at: datetime | None = Field(default=None, description="Task start time")
+    updated_at: datetime | None = Field(default=None, description="Last update time")
+    error: str = Field(default="", description="Error message")
+    message: str = Field(default="", description="Message")
+
+
+class TaskStatusData(BaseModel):
+    """Task status data for JSON:API responses."""
+
+    type: str = "task_status"
+    id: str
+    attributes: TaskStatusAttributes
+
+
+class TaskStatusResponse(BaseModel):
+    """JSON:API response for single task status."""
+
+    data: TaskStatusData
+
+
+class TaskStatusListResponse(BaseModel):
+    """JSON:API response for task status list."""
+
+    data: list[TaskStatusData]

kodit/infrastructure/bm25/local_bm25_repository.py

@@ -37,7 +37,7 @@ class LocalBM25Repository(BM25Repository):
         """
         self.log = structlog.get_logger(__name__)
         self.index_path = data_dir / "bm25s_index"
-        self.snippet_ids: list[int] = []
+        self.snippet_ids: list[str] = []
         self.stemmer = Stemmer.Stemmer("english")
         self.__retriever: bm25s.BM25 | None = None
 
@@ -76,11 +76,23 @@ class LocalBM25Repository(BM25Repository):
             self.log.warning("Corpus is empty, skipping bm25 index")
             return
 
-        vocab = self._tokenize([doc.text for doc in request.documents])
+        if not self.snippet_ids and (self.index_path / SNIPPET_IDS_FILE).exists():
+            async with aiofiles.open(self.index_path / SNIPPET_IDS_FILE) as f:
+                self.snippet_ids = json.loads(await f.read())
+
+        # Filter out documents that have already been indexed
+        new_documents = [
+            doc for doc in request.documents if doc.snippet_id not in self.snippet_ids
+        ]
+        if not new_documents:
+            self.log.info("No new documents to index")
+            return
+
+        vocab = self._tokenize([doc.text for doc in new_documents])
         self._retriever().index(vocab, show_progress=False)
         self._retriever().save(self.index_path)
         # Replace snippet_ids instead of appending, since the BM25 index is rebuilt
-        self.snippet_ids = [doc.snippet_id for doc in request.documents]
+        self.snippet_ids = [doc.snippet_id for doc in new_documents]
         async with aiofiles.open(self.index_path / SNIPPET_IDS_FILE, "w") as f:
             await f.write(json.dumps(self.snippet_ids))
 
@@ -121,7 +133,7 @@ class LocalBM25Repository(BM25Repository):
         # Filter results by snippet_ids if provided
         filtered_results = []
         for result, score in zip(results[0], scores[0], strict=True):
-            snippet_id = int(result)
+            snippet_id = result
             if score > 0.0 and (
                 request.snippet_ids is None or snippet_id in request.snippet_ids
             ):

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},
+            )