kodit 0.2.8__py3-none-any.whl → 0.2.9__py3-none-any.whl

kodit/infrastructure/bm25/vectorchord_bm25_repository.py

@@ -8,10 +8,10 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from kodit.domain.services.bm25_service import BM25Repository
 from kodit.domain.value_objects import (
-    BM25DeleteRequest,
-    BM25IndexRequest,
-    BM25SearchRequest,
-    BM25SearchResult,
+    DeleteRequest,
+    IndexRequest,
+    SearchRequest,
+    SearchResult,
 )
 
 TABLE_NAME = "vectorchord_bm25_documents"
@@ -80,6 +80,17 @@ SEARCH_QUERY = f"""
     ORDER BY bm25_score
     LIMIT :limit
 """  # noqa: S608
+SEARCH_QUERY_WITH_FILTER = f"""
+    SELECT
+        snippet_id,
+        embedding <&>
+            to_bm25query('{INDEX_NAME}', tokenize(:query_text, '{TOKENIZER_NAME}'))
+    AS bm25_score
+    FROM {TABLE_NAME}
+    WHERE snippet_id = ANY(:snippet_ids)
+    ORDER BY bm25_score
+    LIMIT :limit
+"""  # noqa: S608
 DELETE_QUERY = f"""
 DELETE FROM {TABLE_NAME}
 WHERE snippet_id IN :snippet_ids
@@ -146,7 +157,7 @@ class VectorChordBM25Repository(BM25Repository):
         """Commit the session."""
         await self.__session.commit()
 
-    async def index_documents(self, request: BM25IndexRequest) -> None:
+    async def index_documents(self, request: IndexRequest) -> None:
         """Index documents for BM25 search."""
         # Filter out any documents that don't have a snippet_id or text
         valid_documents = [
@@ -172,27 +183,36 @@ class VectorChordBM25Repository(BM25Repository):
         await self._execute(text(UPDATE_QUERY))
         await self._commit()
 
-    async def search(self, request: BM25SearchRequest) -> list[BM25SearchResult]:
+    async def search(self, request: SearchRequest) -> list[SearchResult]:
         """Search documents using BM25."""
         if not request.query or request.query == "":
             return []
 
-        sql = text(SEARCH_QUERY).bindparams(
-            query_text=request.query, limit=request.top_k
-        )
+        if request.snippet_ids is not None:
+            sql = text(SEARCH_QUERY_WITH_FILTER).bindparams(
+                query_text=request.query,
+                limit=request.top_k,
+                snippet_ids=request.snippet_ids,
+            )
+        else:
+            sql = text(SEARCH_QUERY).bindparams(
+                query_text=request.query,
+                limit=request.top_k,
+            )
+
         try:
             result = await self._execute(sql)
             rows = result.mappings().all()
 
             return [
-                BM25SearchResult(snippet_id=row["snippet_id"], score=row["bm25_score"])
+                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: BM25DeleteRequest) -> None:
+    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)),

kodit/infrastructure/cloning/metadata.py

@@ -38,6 +38,7 @@ class BaseFileMetadataExtractor:
                 uri=path.as_uri(),
                 sha256=sha,
                 size_bytes=len(content),
+                extension=path.suffix.removeprefix(".").lower(),
             )
 
     async def _get_timestamps(

kodit/infrastructure/embedding/embedding_providers/hash_embedding_provider.py

@@ -26,40 +26,29 @@ class HashEmbeddingProvider(EmbeddingProvider):
         self.embedding_size = embedding_size
         self.log = structlog.get_logger(__name__)
 
-    def embed(
+    async def embed(
         self, data: list[EmbeddingRequest]
     ) -> AsyncGenerator[list[EmbeddingResponse], None]:
         """Embed a list of strings using a simple hash-based approach."""
         if not data:
-
-            async def empty_generator() -> AsyncGenerator[
-                list[EmbeddingResponse], None
-            ]:
-                if False:
-                    yield []
-
-            return empty_generator()
+            yield []
 
         # Process in batches
         batch_size = 10
-
-        async def _embed_batches() -> AsyncGenerator[list[EmbeddingResponse], None]:
-            for i in range(0, len(data), batch_size):
-                batch = data[i : i + batch_size]
-                responses = []
-
-                for request in batch:
-                    # Generate a deterministic embedding based on the text
-                    embedding = self._generate_embedding(request.text)
-                    responses.append(
-                        EmbeddingResponse(
-                            snippet_id=request.snippet_id, embedding=embedding
-                        )
+        for i in range(0, len(data), batch_size):
+            batch = data[i : i + batch_size]
+            responses = []
+
+            for request in batch:
+                # Generate a deterministic embedding based on the text
+                embedding = self._generate_embedding(request.text)
+                responses.append(
+                    EmbeddingResponse(
+                        snippet_id=request.snippet_id, embedding=embedding
                     )
+                )
 
-                yield responses
-
-        return _embed_batches()
+            yield responses
 
     def _generate_embedding(self, text: str) -> list[float]:
         """Generate a deterministic embedding for the given text."""

kodit/infrastructure/embedding/local_vector_search_repository.py

@@ -3,7 +3,6 @@
 from collections.abc import AsyncGenerator
 
 import structlog
-import tiktoken
 
 from kodit.domain.entities import Embedding, EmbeddingType
 from kodit.domain.services.embedding_service import (
@@ -12,10 +11,10 @@ from kodit.domain.services.embedding_service import (
 )
 from kodit.domain.value_objects import (
     EmbeddingRequest,
+    IndexRequest,
     IndexResult,
-    VectorIndexRequest,
-    VectorSearchQueryRequest,
-    VectorSearchResult,
+    SearchRequest,
+    SearchResult,
 )
 from kodit.infrastructure.sqlalchemy.embedding_repository import (
     SqlAlchemyEmbeddingRepository,
@@ -27,35 +26,29 @@ class LocalVectorSearchRepository(VectorSearchRepository):
 
     def __init__(
         self,
-        embedding_repository: SqlAlchemyEmbeddingRepository,
         embedding_provider: EmbeddingProvider,
-        embedding_type: EmbeddingType = EmbeddingType.CODE,
+        embedding_repository: SqlAlchemyEmbeddingRepository,
+        embedding_type: EmbeddingType,
     ) -> None:
         """Initialize the local vector search repository.
 
         Args:
-            embedding_repository: The SQLAlchemy embedding repository
             embedding_provider: The embedding provider for generating embeddings
+            embedding_repository: The embedding repository for persistence
             embedding_type: The type of embedding to use
 
         """
-        self.log = structlog.get_logger(__name__)
-        self.embedding_repository = embedding_repository
         self.embedding_provider = embedding_provider
-        self.encoding = tiktoken.encoding_for_model("text-embedding-3-small")
+        self.embedding_repository = embedding_repository
         self.embedding_type = embedding_type
+        self.log = structlog.get_logger(__name__)
 
-    def index_documents(
-        self, request: VectorIndexRequest
+    async def index_documents(
+        self, request: IndexRequest
     ) -> AsyncGenerator[list[IndexResult], None]:
         """Index documents for vector search."""
-        if not request.documents:
-
-            async def empty_generator() -> AsyncGenerator[list[IndexResult], None]:
-                if False:
-                    yield []
-
-            return empty_generator()
+        if not request.documents or len(request.documents) == 0:
+            yield []
 
         # Convert to embedding requests
         requests = [
@@ -63,25 +56,20 @@ class LocalVectorSearchRepository(VectorSearchRepository):
             for doc in request.documents
         ]
 
-        async def _index_batches() -> AsyncGenerator[list[IndexResult], None]:
-            async for batch in self.embedding_provider.embed(requests):
-                results = []
-                for result in batch:
-                    await self.embedding_repository.create_embedding(
-                        Embedding(
-                            snippet_id=result.snippet_id,
-                            embedding=result.embedding,
-                            type=self.embedding_type,
-                        )
+        async for batch in self.embedding_provider.embed(requests):
+            results = []
+            for result in batch:
+                await self.embedding_repository.create_embedding(
+                    Embedding(
+                        snippet_id=result.snippet_id,
+                        embedding=result.embedding,
+                        type=self.embedding_type,
                     )
-                    results.append(IndexResult(snippet_id=result.snippet_id))
-                yield results
-
-        return _index_batches()
+                )
+                results.append(IndexResult(snippet_id=result.snippet_id))
+            yield results
 
-    async def search(
-        self, request: VectorSearchQueryRequest
-    ) -> list[VectorSearchResult]:
+    async def search(self, request: SearchRequest) -> list[SearchResult]:
         """Search documents using vector similarity."""
         # Build a single-item request and collect its embedding
         req = EmbeddingRequest(snippet_id=0, text=request.query)
@@ -95,10 +83,10 @@ class LocalVectorSearchRepository(VectorSearchRepository):
             return []
 
         results = await self.embedding_repository.list_semantic_results(
-            self.embedding_type, embedding_vec, request.top_k
+            self.embedding_type, embedding_vec, request.top_k, request.snippet_ids
         )
         return [
-            VectorSearchResult(snippet_id=snippet_id, score=score)
+            SearchResult(snippet_id=snippet_id, score=score)
             for snippet_id, score in results
         ]
 

kodit/infrastructure/embedding/vectorchord_vector_search_repository.py

@@ -14,10 +14,10 @@ from kodit.domain.services.embedding_service import (
 )
 from kodit.domain.value_objects import (
     EmbeddingRequest,
+    IndexRequest,
     IndexResult,
-    VectorIndexRequest,
-    VectorSearchQueryRequest,
-    VectorSearchResult,
+    SearchRequest,
+    SearchResult,
 )
 
 # SQL Queries
@@ -59,6 +59,15 @@ ORDER BY score ASC
 LIMIT :top_k;
 """
 
+# Filtered search query with snippet_ids
+SEARCH_QUERY_WITH_FILTER = """
+SELECT snippet_id, embedding <=> :query as score
+FROM {TABLE_NAME}
+WHERE snippet_id = ANY(:snippet_ids)
+ORDER BY score ASC
+LIMIT :top_k;
+"""
+
 CHECK_VCHORD_EMBEDDING_EXISTS = """
 SELECT EXISTS(SELECT 1 FROM {TABLE_NAME} WHERE snippet_id = :snippet_id)
 """
@@ -156,17 +165,12 @@ class VectorChordVectorSearchRepository(VectorSearchRepository):
         """Commit the session."""
         await self._session.commit()
 
-    def index_documents(
-        self, request: VectorIndexRequest
+    async def index_documents(
+        self, request: IndexRequest
     ) -> AsyncGenerator[list[IndexResult], None]:
         """Index documents for vector search."""
         if not request.documents:
-
-            async def empty_generator() -> AsyncGenerator[list[IndexResult], None]:
-                if False:
-                    yield []
-
-            return empty_generator()
+            yield []
 
         # Convert to embedding requests
         requests = [
@@ -174,27 +178,25 @@ class VectorChordVectorSearchRepository(VectorSearchRepository):
             for doc in request.documents
         ]
 
-        async def _index_batches() -> AsyncGenerator[list[IndexResult], None]:
-            async for batch in self.embedding_provider.embed(requests):
-                await self._execute(
-                    text(INSERT_QUERY.format(TABLE_NAME=self.table_name)),
-                    [
-                        {
-                            "snippet_id": result.snippet_id,
-                            "embedding": str(result.embedding),
-                        }
-                        for result in batch
-                    ],
-                )
-                await self._commit()
-                yield [IndexResult(snippet_id=result.snippet_id) for result in batch]
-
-        return _index_batches()
+        async for batch in self.embedding_provider.embed(requests):
+            await self._execute(
+                text(INSERT_QUERY.format(TABLE_NAME=self.table_name)),
+                [
+                    {
+                        "snippet_id": result.snippet_id,
+                        "embedding": str(result.embedding),
+                    }
+                    for result in batch
+                ],
+            )
+            await self._commit()
+            yield [IndexResult(snippet_id=result.snippet_id) for result in batch]
 
-    async def search(
-        self, request: VectorSearchQueryRequest
-    ) -> list[VectorSearchResult]:
+    async def search(self, request: SearchRequest) -> list[SearchResult]:
         """Search documents using vector similarity."""
+        if not request.query or not request.query.strip():
+            return []
+
         req = EmbeddingRequest(snippet_id=0, text=request.query)
         embedding_vec: list[float] | None = None
         async for batch in self.embedding_provider.embed([req]):
@@ -204,14 +206,27 @@ class VectorChordVectorSearchRepository(VectorSearchRepository):
 
         if not embedding_vec:
             return []
-        result = await self._execute(
-            text(SEARCH_QUERY.format(TABLE_NAME=self.table_name)),
-            {"query": str(embedding_vec), "top_k": request.top_k},
-        )
+
+        # Use filtered query if snippet_ids are provided
+        if request.snippet_ids is not None:
+            result = await self._execute(
+                text(SEARCH_QUERY_WITH_FILTER.format(TABLE_NAME=self.table_name)),
+                {
+                    "query": str(embedding_vec),
+                    "top_k": request.top_k,
+                    "snippet_ids": request.snippet_ids,
+                },
+            )
+        else:
+            result = await self._execute(
+                text(SEARCH_QUERY.format(TABLE_NAME=self.table_name)),
+                {"query": str(embedding_vec), "top_k": request.top_k},
+            )
+
         rows = result.mappings().all()
 
         return [
-            VectorSearchResult(snippet_id=row["snippet_id"], score=row["score"])
+            SearchResult(snippet_id=row["snippet_id"], score=row["score"])
             for row in rows
         ]
 

kodit/infrastructure/enrichment/enrichment_factory.py

@@ -24,7 +24,7 @@ def _get_endpoint_configuration(app_context: AppContext) -> Endpoint | None:
     return app_context.enrichment_endpoint or app_context.default_endpoint or None
 
 
-def create_enrichment_domain_service(
+def enrichment_domain_service_factory(
     app_context: AppContext,
 ) -> EnrichmentDomainService:
     """Create an enrichment domain service.

kodit/infrastructure/indexing/indexing_factory.py

@@ -1,113 +1,30 @@
 """Factory for creating indexing services."""
 
-from typing import Any
-
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from kodit.application.services.indexing_application_service import (
-    IndexingApplicationService,
-)
-from kodit.application.services.snippet_application_service import (
-    SnippetApplicationService,
-)
-from kodit.domain.services.bm25_service import BM25DomainService
 from kodit.domain.services.indexing_service import IndexingDomainService
-from kodit.domain.services.source_service import SourceService
-from kodit.infrastructure.bm25.bm25_factory import bm25_repository_factory
-from kodit.infrastructure.embedding.embedding_factory import (
-    embedding_domain_service_factory,
-)
-from kodit.infrastructure.enrichment.enrichment_factory import (
-    create_enrichment_domain_service,
-)
 from kodit.infrastructure.indexing.fusion_service import ReciprocalRankFusionService
 from kodit.infrastructure.indexing.index_repository import SQLAlchemyIndexRepository
-from kodit.infrastructure.snippet_extraction.snippet_extraction_factory import (
-    create_snippet_extraction_domain_service,
-)
-from kodit.infrastructure.sqlalchemy.file_repository import (
-    SqlAlchemyFileRepository,
-)
-from kodit.infrastructure.sqlalchemy.snippet_repository import (
-    SqlAlchemySnippetRepository,
-)
-
-
-def create_snippet_application_service(
-    session: AsyncSession,
-) -> SnippetApplicationService:
-    """Create a snippet application service with all dependencies."""
-    # Create domain service
-    snippet_extraction_service = create_snippet_extraction_domain_service()
-
-    # Create repositories
-    snippet_repository = SqlAlchemySnippetRepository(session)
-    file_repository = SqlAlchemyFileRepository(session)
 
-    # Create application service
-    from kodit.application.services.snippet_application_service import (
-        SnippetApplicationService,
-    )
-
-    return SnippetApplicationService(
-        snippet_extraction_service=snippet_extraction_service,
-        snippet_repository=snippet_repository,
-        file_repository=file_repository,
-        session=session,
-    )
 
-
-def create_indexing_domain_service(session: AsyncSession) -> IndexingDomainService:
-    """Create an indexing domain service.
+def indexing_domain_service_factory(session: AsyncSession) -> IndexingDomainService:
+    """Create an indexing domain service with all dependencies.
 
     Args:
-        session: The database session.
+        session: SQLAlchemy session
 
     Returns:
-        An indexing domain service instance.
+        Configured indexing domain service
 
     """
+    # Create repositories
     index_repository = SQLAlchemyIndexRepository(session)
+
+    # Create fusion service
     fusion_service = ReciprocalRankFusionService()
 
+    # Create domain service
     return IndexingDomainService(
         index_repository=index_repository,
         fusion_service=fusion_service,
     )
-
-
-def create_indexing_application_service(
-    app_context: Any,
-    session: AsyncSession,
-    source_service: SourceService,
-    snippet_application_service: SnippetApplicationService,
-) -> IndexingApplicationService:
-    """Create an indexing application service.
-
-    Args:
-        app_context: The application context.
-        session: The database session.
-        source_service: The source service.
-        snippet_application_service: The snippet application service.
-
-    Returns:
-        An indexing application service instance.
-
-    """
-    # Create domain services
-    indexing_domain_service = create_indexing_domain_service(session)
-    bm25_service = BM25DomainService(bm25_repository_factory(app_context, session))
-    code_search_service = embedding_domain_service_factory("code", app_context, session)
-    text_search_service = embedding_domain_service_factory("text", app_context, session)
-    enrichment_service = create_enrichment_domain_service(app_context)
-
-    return IndexingApplicationService(
-        indexing_domain_service=indexing_domain_service,
-        source_service=source_service,
-        bm25_service=bm25_service,
-        code_search_service=code_search_service,
-        text_search_service=text_search_service,
-        enrichment_service=enrichment_service,
-        snippet_application_service=snippet_application_service,
-        session=session,
-    )

kodit/infrastructure/indexing/snippet_domain_service_factory.py

@@ -0,0 +1,37 @@
+"""Factory for creating snippet domain service."""
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.services.snippet_service import SnippetDomainService
+from kodit.infrastructure.snippet_extraction.snippet_extraction_factory import (
+    create_snippet_extraction_domain_service,
+)
+from kodit.infrastructure.sqlalchemy.file_repository import SqlAlchemyFileRepository
+from kodit.infrastructure.sqlalchemy.snippet_repository import (
+    SqlAlchemySnippetRepository,
+)
+
+
+def snippet_domain_service_factory(session: AsyncSession) -> SnippetDomainService:
+    """Create a snippet domain service with all dependencies.
+
+    Args:
+        session: The database session
+
+    Returns:
+        Configured snippet domain service
+
+    """
+    # Create domain service for snippet extraction
+    snippet_extraction_service = create_snippet_extraction_domain_service()
+
+    # Create repositories
+    snippet_repository = SqlAlchemySnippetRepository(session)
+    file_repository = SqlAlchemyFileRepository(session)
+
+    # Create and return the domain service
+    return SnippetDomainService(
+        snippet_extraction_service=snippet_extraction_service,
+        snippet_repository=snippet_repository,
+        file_repository=file_repository,
+    )

kodit/infrastructure/snippet_extraction/languages/java.scm

@@ -0,0 +1,12 @@
+(import_declaration
+  (scoped_identifier) @import.name
+)
+
+(method_declaration
+  name: (identifier) @function.name
+  body: (block) @function.body
+) @function.def
+
+(class_declaration
+  name: (identifier) @class.name
+) @class.def

kodit/infrastructure/snippet_extraction/snippet_extraction_factory.py

@@ -9,6 +9,7 @@ from kodit.domain.repositories import FileRepository, SnippetRepository
 from kodit.domain.services.snippet_extraction_service import (
     SnippetExtractionDomainService,
 )
+from kodit.domain.value_objects import LanguageMapping
 from kodit.infrastructure.snippet_extraction.language_detection_service import (
     FileSystemLanguageDetectionService,
 )
@@ -31,37 +32,8 @@ def create_snippet_extraction_domain_service() -> SnippetExtractionDomainService
         Configured snippet extraction domain service
 
     """
-    # Language mapping from the existing languages module
-    language_map = {
-        # JavaScript/TypeScript
-        "js": "javascript",
-        "jsx": "javascript",
-        "ts": "typescript",
-        "tsx": "typescript",
-        # Python
-        "py": "python",
-        # Rust
-        "rs": "rust",
-        # Go
-        "go": "go",
-        # C/C++
-        "cpp": "cpp",
-        "hpp": "cpp",
-        "c": "c",
-        "h": "c",
-        # C#
-        "cs": "csharp",
-        # Ruby
-        "rb": "ruby",
-        # Java
-        "java": "java",
-        # PHP
-        "php": "php",
-        # Swift
-        "swift": "swift",
-        # Kotlin
-        "kt": "kotlin",
-    }
+    # Use the unified language mapping from the domain layer
+    language_map = LanguageMapping.get_extension_to_language_map()
 
     # Create infrastructure services
     language_detector = FileSystemLanguageDetectionService(language_map)