Feat: Delete embedding when delete document

src/document/document_service.py

@@ -1,7 +1,8 @@
 """Service for managing documents."""
 
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy import select, delete
+from sqlalchemy import select, delete, text
+from src.db.postgres.connection import _pgvector_engine
 from src.db.postgres.models import Document
 from src.storage.az_blob.az_blob import blob_storage
 from src.middlewares.logging import get_logger
@@ -77,6 +78,20 @@ class DocumentService:
         # Delete from blob storage
         await blob_storage.delete_file(document.blob_name)
 
+        async with _pgvector_engine.begin() as conn:
+            await conn.execute(
+                text("""
+                    DELETE FROM langchain_pg_embedding
+                    WHERE cmetadata->>'user_id' = :user_id
+                      AND cmetadata->>'source_type' = 'document'
+                      AND cmetadata->'data'->>'document_id' = :doc_id
+                      AND collection_id = (
+                        SELECT uuid FROM langchain_pg_collection WHERE name = 'document_embeddings'
+                      )
+                """),
+                {"user_id": document.user_id, "doc_id": document_id},
+            )
+
         # Delete from database
         await db.execute(
             delete(Document).where(Document.id == document_id)

src/pipeline/db_pipeline/db_pipeline_service.py

@@ -112,23 +112,23 @@ class DbPipelineService:
                 location=credentials.get("location", "US"),
             )
 
-        if db_type == "snowflake":
-            from snowflake.sqlalchemy import URL as SnowflakeURL
-
-            url = SnowflakeURL(
-                account=credentials["account"],
-                user=credentials["username"],
-                password=credentials["password"],
-                database=credentials["database"],
-                schema=(
-                    credentials.get("db_schema")
-                    or credentials.get("schema")
-                    or "PUBLIC"
-                ),
-                warehouse=credentials["warehouse"],
-                role=credentials.get("role") or "",
-            )
-            return create_engine(url)
+        # if db_type == "snowflake":
+        #     from snowflake.sqlalchemy import URL as SnowflakeURL
+
+        #     url = SnowflakeURL(
+        #         account=credentials["account"],
+        #         user=credentials["username"],
+        #         password=credentials["password"],
+        #         database=credentials["database"],
+        #         schema=(
+        #             credentials.get("db_schema")
+        #             or credentials.get("schema")
+        #             or "PUBLIC"
+        #         ),
+        #         warehouse=credentials["warehouse"],
+        #         role=credentials.get("role") or "",
+        #     )
+        #     return create_engine(url)
 
         raise NotImplementedError(f"Unsupported db_type: {db_type}")