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

kodit/infrastructure/sqlalchemy/embedding_repository.py

@@ -14,59 +14,79 @@ def create_embedding_repository(
     session_factory: Callable[[], AsyncSession],
 ) -> "SqlAlchemyEmbeddingRepository":
     """Create an embedding repository."""
-    uow = SqlAlchemyUnitOfWork(session_factory=session_factory)
-    return SqlAlchemyEmbeddingRepository(uow)
+    return SqlAlchemyEmbeddingRepository(session_factory=session_factory)
 
 
 class SqlAlchemyEmbeddingRepository:
     """SQLAlchemy implementation of embedding repository."""
 
-    def __init__(self, uow: SqlAlchemyUnitOfWork) -> None:
+    def __init__(self, session_factory: Callable[[], AsyncSession]) -> None:
         """Initialize the SQLAlchemy embedding repository."""
-        self.uow = uow
+        self.session_factory = session_factory
 
     async def create_embedding(self, embedding: Embedding) -> None:
         """Create a new embedding record in the database."""
-        async with self.uow:
-            self.uow.session.add(embedding)
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            session.add(embedding)
 
     async def get_embedding_by_snippet_id_and_type(
         self, snippet_id: int, embedding_type: EmbeddingType
     ) -> Embedding | None:
         """Get an embedding by its snippet ID and type."""
-        async with self.uow:
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
             query = select(Embedding).where(
                 Embedding.snippet_id == snippet_id,
                 Embedding.type == embedding_type,
             )
-            result = await self.uow.session.execute(query)
+            result = await session.execute(query)
             return result.scalar_one_or_none()
 
     async def list_embeddings_by_type(
         self, embedding_type: EmbeddingType
     ) -> list[Embedding]:
         """List all embeddings of a given type."""
-        async with self.uow:
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
             query = select(Embedding).where(Embedding.type == embedding_type)
-            result = await self.uow.session.execute(query)
+            result = await session.execute(query)
             return list(result.scalars())
 
-    async def delete_embeddings_by_snippet_id(self, snippet_id: int) -> None:
+    async def delete_embeddings_by_snippet_id(self, snippet_id: str) -> None:
         """Delete all embeddings for a snippet."""
-        async with self.uow:
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
             query = select(Embedding).where(Embedding.snippet_id == snippet_id)
-            result = await self.uow.session.execute(query)
+            result = await session.execute(query)
             embeddings = result.scalars().all()
             for embedding in embeddings:
-                await self.uow.session.delete(embedding)
+                await session.delete(embedding)
+
+    async def list_embeddings_by_snippet_ids_and_type(
+        self, snippet_ids: list[str], embedding_type: EmbeddingType
+    ) -> list[Embedding]:
+        """Get all embeddings for the given snippet IDs."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            query = select(Embedding).where(
+                Embedding.snippet_id.in_(snippet_ids),
+                Embedding.type == embedding_type,
+            )
+            result = await session.execute(query)
+            return list(result.scalars())
+
+    async def get_embeddings_by_snippet_ids(
+        self, snippet_ids: list[str]
+    ) -> list[Embedding]:
+        """Get all embeddings for the given snippet IDs."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            query = select(Embedding).where(Embedding.snippet_id.in_(snippet_ids))
+            result = await session.execute(query)
+            return list(result.scalars())
 
     async def list_semantic_results(
         self,
         embedding_type: EmbeddingType,
         embedding: list[float],
         top_k: int = 10,
-        snippet_ids: list[int] | None = None,
-    ) -> list[tuple[int, float]]:
+        snippet_ids: list[str] | None = None,
+    ) -> list[tuple[str, float]]:
         """List semantic results using cosine similarity.
 
         This implementation fetches all embeddings of the given type and computes
@@ -97,8 +117,8 @@ class SqlAlchemyEmbeddingRepository:
         return self._get_top_k_results(similarities, embeddings, top_k)
 
     async def _list_embedding_values(
-        self, embedding_type: EmbeddingType, snippet_ids: list[int] | None = None
-    ) -> list[tuple[int, list[float]]]:
+        self, embedding_type: EmbeddingType, snippet_ids: list[str] | None = None
+    ) -> list[tuple[str, list[float]]]:
         """List all embeddings of a given type from the database.
 
         Args:
@@ -109,7 +129,7 @@ class SqlAlchemyEmbeddingRepository:
             List of (snippet_id, embedding) tuples
 
         """
-        async with self.uow:
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
             query = select(Embedding.snippet_id, Embedding.embedding).where(
                 Embedding.type == embedding_type
             )
@@ -118,11 +138,11 @@ class SqlAlchemyEmbeddingRepository:
             if snippet_ids is not None:
                 query = query.where(Embedding.snippet_id.in_(snippet_ids))
 
-            rows = await self.uow.session.execute(query)
+            rows = await session.execute(query)
             return [tuple(row) for row in rows.all()]  # Convert Row objects to tuples
 
     def _prepare_vectors(
-        self, embeddings: list[tuple[int, list[float]]], query_embedding: list[float]
+        self, embeddings: list[tuple[str, list[float]]], query_embedding: list[float]
     ) -> tuple[np.ndarray, np.ndarray]:
         """Convert embeddings to numpy arrays.
 
@@ -191,9 +211,9 @@ class SqlAlchemyEmbeddingRepository:
     def _get_top_k_results(
         self,
         similarities: np.ndarray,
-        embeddings: list[tuple[int, list[float]]],
+        embeddings: list[tuple[str, list[float]]],
         top_k: int,
-    ) -> list[tuple[int, float]]:
+    ) -> list[tuple[str, float]]:
         """Get top-k results by similarity score.
 
         Args:

kodit/infrastructure/sqlalchemy/enrichment_v2_repository.py

@@ -0,0 +1,118 @@
+"""EnrichmentV2 repository."""
+
+from collections.abc import Callable, Sequence
+
+import structlog
+from sqlalchemy import delete, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.enrichments.enrichment import EnrichmentV2
+from kodit.infrastructure.mappers.enrichment_mapper import EnrichmentMapper
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+from kodit.infrastructure.sqlalchemy.unit_of_work import SqlAlchemyUnitOfWork
+
+
+class EnrichmentV2Repository:
+    """Repository for managing enrichments and their associations."""
+
+    def __init__(
+        self,
+        session_factory: Callable[[], AsyncSession],
+    ) -> None:
+        """Initialize the repository."""
+        self.session_factory = session_factory
+        self.mapper = EnrichmentMapper()
+        self.log = structlog.get_logger(__name__)
+
+    async def enrichments_for_entity_type(
+        self,
+        entity_type: str,
+        entity_ids: list[str],
+    ) -> list[EnrichmentV2]:
+        """Get all enrichments for multiple entities of the same type."""
+        if not entity_ids:
+            return []
+
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = (
+                select(
+                    db_entities.EnrichmentV2,
+                    db_entities.EnrichmentAssociation.entity_id,
+                )
+                .join(db_entities.EnrichmentAssociation)
+                .where(
+                    db_entities.EnrichmentAssociation.entity_type == entity_type,
+                    db_entities.EnrichmentAssociation.entity_id.in_(entity_ids),
+                )
+            )
+
+            result = await session.execute(stmt)
+            rows = result.all()
+
+            return [
+                self.mapper.to_domain(db_enrichment, entity_type, entity_id)
+                for db_enrichment, entity_id in rows
+            ]
+
+    async def bulk_save_enrichments(
+        self,
+        enrichments: Sequence[EnrichmentV2],
+    ) -> None:
+        """Bulk save enrichments with their associations."""
+        if not enrichments:
+            return
+
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            enrichment_records = []
+            for enrichment in enrichments:
+                db_enrichment = db_entities.EnrichmentV2(
+                    type=enrichment.type,
+                    subtype=enrichment.subtype,
+                    content=enrichment.content,
+                )
+                session.add(db_enrichment)
+                enrichment_records.append((enrichment, db_enrichment))
+
+            await session.flush()
+
+            for enrichment, db_enrichment in enrichment_records:
+                db_association = db_entities.EnrichmentAssociation(
+                    enrichment_id=db_enrichment.id,
+                    entity_type=enrichment.entity_type_key(),
+                    entity_id=enrichment.entity_id,
+                )
+                session.add(db_association)
+
+    async def bulk_delete_enrichments(
+        self,
+        entity_type: str,
+        entity_ids: list[str],
+    ) -> None:
+        """Bulk delete enrichments for multiple entities of the same type."""
+        if not entity_ids:
+            return
+
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            stmt = select(db_entities.EnrichmentAssociation.enrichment_id).where(
+                db_entities.EnrichmentAssociation.entity_type == entity_type,
+                db_entities.EnrichmentAssociation.entity_id.in_(entity_ids),
+            )
+            result = await session.execute(stmt)
+            enrichment_ids = result.scalars().all()
+
+            if enrichment_ids:
+                await session.execute(
+                    delete(db_entities.EnrichmentV2).where(
+                        db_entities.EnrichmentV2.id.in_(enrichment_ids)
+                    )
+                )
+
+    async def delete_enrichment(self, enrichment_id: int) -> bool:
+        """Delete a specific enrichment by ID."""
+        async with SqlAlchemyUnitOfWork(self.session_factory) as session:
+            result = await session.execute(
+                delete(db_entities.EnrichmentV2).where(
+                    db_entities.EnrichmentV2.id == enrichment_id
+                )
+            )
+            return result.rowcount > 0