kodit 0.5.4__py3-none-any.whl → 0.5.6__py3-none-any.whl

kodit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.5.4'
-__version_tuple__ = version_tuple = (0, 5, 4)
+__version__ = version = '0.5.6'
+__version_tuple__ = version_tuple = (0, 5, 6)
 
 __commit_id__ = commit_id = None

kodit/app.py

@@ -63,6 +63,8 @@ async def app_lifespan(_: FastAPI) -> AsyncIterator[AppLifespanState]:
                 )
             )
         )
+    except StopAsyncIteration:
+        pass
     except Exception as e:
         raise ValueError("Embedding service is not accessible") from e
     try:

kodit/application/factories/server_factory.py

@@ -23,10 +23,13 @@ from kodit.domain.enrichments.architecture.physical.formatter import (
 )
 from kodit.domain.enrichments.enricher import Enricher
 from kodit.domain.protocols import (
+    EnrichmentAssociationRepository,
+    EnrichmentV2Repository,
     FusionService,
     GitAdapter,
     GitBranchRepository,
     GitCommitRepository,
+    GitFileRepository,
     GitRepoRepository,
     GitTagRepository,
     SnippetRepositoryV2,
@@ -47,6 +50,9 @@ from kodit.infrastructure.bm25.vectorchord_bm25_repository import (
     VectorChordBM25Repository,
 )
 from kodit.infrastructure.cloning.git.git_python_adaptor import GitPythonAdapter
+from kodit.infrastructure.database_schema.database_schema_detector import (
+    DatabaseSchemaDetector,
+)
 from kodit.infrastructure.embedding.embedding_factory import (
     embedding_domain_service_factory,
 )
@@ -64,8 +70,11 @@ from kodit.infrastructure.sqlalchemy.embedding_repository import (
     SqlAlchemyEmbeddingRepository,
     create_embedding_repository,
 )
+from kodit.infrastructure.sqlalchemy.enrichment_association_repository import (
+    create_enrichment_association_repository,
+)
 from kodit.infrastructure.sqlalchemy.enrichment_v2_repository import (
-    EnrichmentV2Repository,
+    create_enrichment_v2_repository,
 )
 from kodit.infrastructure.sqlalchemy.git_branch_repository import (
     create_git_branch_repository,
@@ -73,13 +82,13 @@ from kodit.infrastructure.sqlalchemy.git_branch_repository import (
 from kodit.infrastructure.sqlalchemy.git_commit_repository import (
     create_git_commit_repository,
 )
+from kodit.infrastructure.sqlalchemy.git_file_repository import (
+    create_git_file_repository,
+)
 from kodit.infrastructure.sqlalchemy.git_repository import create_git_repo_repository
 from kodit.infrastructure.sqlalchemy.git_tag_repository import (
     create_git_tag_repository,
 )
-from kodit.infrastructure.sqlalchemy.snippet_v2_repository import (
-    create_snippet_v2_repository,
-)
 from kodit.infrastructure.sqlalchemy.task_status_repository import (
     create_task_status_repository,
 )
@@ -121,10 +130,14 @@ class ServerFactory:
             None
         )
         self._git_commit_repository: GitCommitRepository | None = None
+        self._git_file_repository: GitFileRepository | None = None
         self._git_branch_repository: GitBranchRepository | None = None
         self._git_tag_repository: GitTagRepository | None = None
         self._architecture_service: PhysicalArchitectureService | None = None
         self._enrichment_v2_repository: EnrichmentV2Repository | None = None
+        self._enrichment_association_repository: (
+            EnrichmentAssociationRepository | None
+        ) = None
         self._architecture_formatter: PhysicalArchitectureFormatter | None = None
         self._trackable_resolution_service: TrackableResolutionService | None = None
         self._enrichment_query_service: EnrichmentQueryService | None = None
@@ -146,11 +159,23 @@ class ServerFactory:
     def enrichment_v2_repository(self) -> EnrichmentV2Repository:
         """Create a EnrichmentV2Repository instance."""
         if not self._enrichment_v2_repository:
-            self._enrichment_v2_repository = EnrichmentV2Repository(
+            self._enrichment_v2_repository = create_enrichment_v2_repository(
                 session_factory=self.session_factory
             )
         return self._enrichment_v2_repository
 
+    def enrichment_association_repository(
+        self,
+    ) -> EnrichmentAssociationRepository:
+        """Create a EnrichmentAssociationRepository instance."""
+        if not self._enrichment_association_repository:
+            self._enrichment_association_repository = (
+                create_enrichment_association_repository(
+                    session_factory=self.session_factory
+                )
+            )
+        return self._enrichment_association_repository
+
     def queue_service(self) -> QueueService:
         """Create a QueueService instance."""
         if not self._queue_service:
@@ -217,27 +242,27 @@ class ServerFactory:
     def commit_indexing_application_service(self) -> CommitIndexingApplicationService:
         """Create a CommitIndexingApplicationService instance."""
         if not self._commit_indexing_application_service:
-            self._commit_indexing_application_service = (
-                CommitIndexingApplicationService(
-                    snippet_v2_repository=self.snippet_v2_repository(),
-                    repo_repository=self.repo_repository(),
-                    git_commit_repository=self.git_commit_repository(),
-                    git_branch_repository=self.git_branch_repository(),
-                    git_tag_repository=self.git_tag_repository(),
-                    operation=self.operation(),
-                    scanner=self.scanner(),
-                    cloner=self.cloner(),
-                    snippet_repository=self.snippet_v2_repository(),
-                    slicer=self.slicer(),
-                    queue=self.queue_service(),
-                    bm25_service=self.bm25_service(),
-                    code_search_service=self.code_search_service(),
-                    text_search_service=self.text_search_service(),
-                    embedding_repository=self.embedding_repository(),
-                    architecture_service=self.architecture_service(),
-                    enrichment_v2_repository=self.enrichment_v2_repository(),
-                    enricher_service=self.enricher(),
-                )
+            self._commit_indexing_application_service = CommitIndexingApplicationService(  # noqa: E501
+                repo_repository=self.repo_repository(),
+                git_commit_repository=self.git_commit_repository(),
+                git_file_repository=self.git_file_repository(),
+                git_branch_repository=self.git_branch_repository(),
+                git_tag_repository=self.git_tag_repository(),
+                operation=self.operation(),
+                scanner=self.scanner(),
+                cloner=self.cloner(),
+                slicer=self.slicer(),
+                queue=self.queue_service(),
+                bm25_service=self.bm25_service(),
+                code_search_service=self.code_search_service(),
+                text_search_service=self.text_search_service(),
+                embedding_repository=self.embedding_repository(),
+                architecture_service=self.architecture_service(),
+                database_schema_detector=DatabaseSchemaDetector(),
+                enrichment_v2_repository=self.enrichment_v2_repository(),
+                enricher_service=self.enricher(),
+                enrichment_association_repository=self.enrichment_association_repository(),
+                enrichment_query_service=self.enrichment_query_service(),
             )
 
         return self._commit_indexing_application_service
@@ -279,13 +304,13 @@ class ServerFactory:
             )
         return self._cloner
 
-    def snippet_v2_repository(self) -> SnippetRepositoryV2:
-        """Create a SnippetRepositoryV2 instance."""
-        if not self._snippet_v2_repository:
-            self._snippet_v2_repository = create_snippet_v2_repository(
+    def git_file_repository(self) -> GitFileRepository:
+        """Create a GitFileRepository instance."""
+        if not self._git_file_repository:
+            self._git_file_repository = create_git_file_repository(
                 session_factory=self.session_factory
             )
-        return self._snippet_v2_repository
+        return self._git_file_repository
 
     def enricher(self) -> Enricher:
         """Create a EnricherDomainService instance."""
@@ -324,8 +349,8 @@ class ServerFactory:
                 code_search_service=self.code_search_service(),
                 text_search_service=self.text_search_service(),
                 progress_tracker=self.operation(),
-                snippet_repository=self.snippet_v2_repository(),
                 fusion_service=self.fusion_service(),
+                enrichment_query_service=self.enrichment_query_service(),
             )
         return self._code_search_application_service
 
@@ -369,5 +394,6 @@ class ServerFactory:
             self._enrichment_query_service = EnrichmentQueryService(
                 trackable_resolution=self.trackable_resolution_service(),
                 enrichment_repo=self.enrichment_v2_repository(),
+                enrichment_association_repository=self.enrichment_association_repository(),
             )
         return self._enrichment_query_service

kodit/application/services/code_search_application_service.py

@@ -1,15 +1,17 @@
 """Service for searching the indexes."""
 
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import structlog
 
 from kodit.application.services.reporting import ProgressTracker
 from kodit.domain.entities.git import SnippetV2
-from kodit.domain.protocols import FusionService, SnippetRepositoryV2
+from kodit.domain.protocols import FusionService
 from kodit.domain.services.bm25_service import BM25DomainService
 from kodit.domain.services.embedding_service import EmbeddingDomainService
 from kodit.domain.value_objects import (
+    Enrichment,
     FusionRequest,
     MultiSearchRequest,
     SearchRequest,
@@ -17,6 +19,11 @@ from kodit.domain.value_objects import (
 )
 from kodit.log import log_event
 
+if TYPE_CHECKING:
+    from kodit.application.services.enrichment_query_service import (
+        EnrichmentQueryService,
+    )
+
 
 @dataclass
 class MultiSearchResult:
@@ -53,16 +60,16 @@ class CodeSearchApplicationService:
         code_search_service: EmbeddingDomainService,
         text_search_service: EmbeddingDomainService,
         progress_tracker: ProgressTracker,
-        snippet_repository: SnippetRepositoryV2,
         fusion_service: FusionService,
+        enrichment_query_service: "EnrichmentQueryService",
     ) -> None:
         """Initialize the code search application service."""
         self.bm25_service = bm25_service
         self.code_search_service = code_search_service
         self.text_search_service = text_search_service
         self.progress_tracker = progress_tracker
-        self.snippet_repository = snippet_repository
         self.fusion_service = fusion_service
+        self.enrichment_query_service = enrichment_query_service
         self.log = structlog.get_logger(__name__)
 
     async def search(self, request: MultiSearchRequest) -> list[MultiSearchResult]:
@@ -108,17 +115,32 @@ class CodeSearchApplicationService:
 
         # Semantic text search
         if request.text_query:
-            query_results = await self.text_search_service.search(
+            # These contain a pointer to the enrichment ID that represents the summary
+            summary_results = await self.text_search_service.search(
                 SearchRequest(
                     query=request.text_query,
                     top_k=request.top_k,
                     snippet_ids=filtered_snippet_ids,
                 )
             )
-            fusion_list.append(
-                [FusionRequest(id=x.snippet_id, score=x.score) for x in query_results]
+
+            summary_to_snippet_map = (
+                await self.enrichment_query_service.summary_to_snippet_map(
+                    summary_ids=[int(x.snippet_id) for x in summary_results]
+                )
             )
 
+            # Build fusion list in the correct order
+            fusion_items = [
+                FusionRequest(
+                    id=str(summary_to_snippet_map[int(result.snippet_id)]),
+                    score=result.score,
+                )
+                for result in summary_results
+                if int(result.snippet_id) in summary_to_snippet_map
+            ]
+            fusion_list.append(fusion_items)
+
         if len(fusion_list) == 0:
             return []
 
@@ -131,14 +153,69 @@ class CodeSearchApplicationService:
         # Keep only top_k results
         final_results = final_results[: request.top_k]
 
-        # Get snippet details
-        ids = [x.id for x in final_results]
-        search_results = await self.snippet_repository.get_by_ids(ids)
-        search_results.sort(key=lambda x: ids.index(x.id))
+        # Get enrichment details
+        enrichment_ids = [int(x.id) for x in final_results]
+
+        self.log.info(
+            "found enrichments",
+            len_enrichments=len(enrichment_ids),
+        )
+        final_enrichments = await self.enrichment_query_service.get_enrichments_by_ids(
+            enrichment_ids
+        )
+
+        # Get enrichments pointing to these enrichments
+        extra_enrichments = (
+            await self.enrichment_query_service.get_enrichments_pointing_to_enrichments(
+                enrichment_ids
+            )
+        )
+
+        self.log.info(
+            "final enrichments",
+            len_final_enrichments=len(final_enrichments),
+        )
+
+        # Convert enrichments to SnippetV2 domain objects
+        # Map enrichment ID to snippet for correct ordering
+        enrichment_id_to_snippet: dict[int | None, SnippetV2] = {}
+        for enrichment in final_enrichments:
+            # Get extra enrichments for this enrichment (only if ID is not None)
+            enrichment_extras = (
+                extra_enrichments[enrichment.id] if enrichment.id is not None else []
+            )
+            enrichment_id_to_snippet[enrichment.id] = SnippetV2(
+                sha=str(enrichment.id),  # The snippet SHA
+                content=enrichment.content,  # The code content
+                extension="",  # Not available in enrichment
+                derives_from=[],  # Not available in enrichment
+                created_at=enrichment.created_at,
+                updated_at=enrichment.updated_at,
+                enrichments=[
+                    Enrichment(
+                        type=enrichment.subtype or enrichment.type,
+                        content=enrichment.content,
+                    )
+                    for enrichment in enrichment_extras
+                ],
+            )
+
+        # Sort by the original fusion ranking order
+        snippets = [
+            enrichment_id_to_snippet[eid]
+            for eid in enrichment_ids
+            if eid in enrichment_id_to_snippet
+        ]
+
         return [
             MultiSearchResult(
                 snippet=snippet,
-                original_scores=[x.score for x in final_results if x.id == snippet.id],
+                original_scores=[
+                    x.score
+                    for x in final_results
+                    if int(x.id) in enrichment_id_to_snippet
+                    and enrichment_id_to_snippet[int(x.id)].sha == snippet.sha
+                ],
             )
-            for snippet in search_results
+            for snippet in snippets
         ]