Add support for removing embeddings (#13)

* WIP

* Add possibility to remove embeddings

api/delete.py

@@ -0,0 +1,14 @@
+from fastapi import APIRouter
+from models.delete import RequestPayload, ResponsePayload
+from service.vector_database import get_vector_service, VectorService
+
+router = APIRouter()
+
+
+@router.post("/delete", response_model=ResponsePayload)
+async def delete(payload: RequestPayload):
+    vector_service: VectorService = get_vector_service(
+        index_name=payload.index_name, credentials=payload.vector_database
+    )
+    await vector_service.delete(file_url=payload.file_url)
+    return {"success": True}

models/delete.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel
+from models.vector_database import VectorDatabase
+
+
+class RequestPayload(BaseModel):
+    index_name: str
+    file_url: str
+    vector_database: VectorDatabase
+
+
+class ResponsePayload(BaseModel):
+    success: bool

router.py

@@ -1,9 +1,10 @@
 from fastapi import APIRouter
 
-from api import ingest, query
+from api import ingest, query, delete
 
 router = APIRouter()
 api_prefix = "/api/v1"
 
 router.include_router(ingest.router, tags=["Ingest"], prefix=api_prefix)
 router.include_router(query.router, tags=["Query"], prefix=api_prefix)
+router.include_router(delete.router, tags=["Delete"], prefix=api_prefix)

service/vector_database.py

@@ -30,6 +30,10 @@ async def query():
     async def convert_to_rerank_format():
         pass
 
+    @abstractmethod
+    async def delete(self, file_url: str):
+        pass
+
     async def _generate_vectors(sefl, input: str):
         vectors = []
         embedding_object = embedding(
@@ -109,6 +113,9 @@ async def query(self, input: str, top_k: 4, include_metadata: bool = True):
         )
         return results["matches"]
 
+    async def delete(self, file_url: str) -> None:
+        self.index.delete(filter={"file_url": {"$eq": file_url}})
+
 
 class QdrantService(VectorService):
     def __init__(self, index_name: str, dimension: int, credentials: dict):
@@ -183,6 +190,20 @@ async def query(self, input: str, top_k: int) -> List:
         )
         return search_result
 
+    async def delete(self, file_url: str) -> None:
+        self.client.delete(
+            collection_name=self.index_name,
+            points_selector=rest.FilterSelector(
+                filter=rest.Filter(
+                    must=[
+                        rest.FieldCondition(
+                            key="file_url", match=rest.MatchValue(value=file_url)
+                        )
+                    ]
+                )
+            ),
+        )
+
 
 class WeaviateService(VectorService):
     def __init__(self, index_name: str, dimension: int, credentials: dict):
@@ -242,6 +263,12 @@ async def query(self, input: str, top_k: int = 4) -> List:
         )
         return result["data"]["Get"][self.index_name.capitalize()]
 
+    async def delete(self, file_url: str) -> None:
+        self.client.batch.delete_objects(
+            class_name=self.index_name,
+            where={"path": ["file_url"], "operator": "Equal", "valueText": file_url},
+        )
+
 
 class AstraService(VectorService):
     def __init__(self, index_name: str, dimension: int, credentials: dict):
@@ -290,6 +317,9 @@ async def query(self, input: str, top_k: int = 4) -> List:
         )
         return results
 
+    async def delete(self, file_url: str) -> None:
+        self.collection.delete_many(filter={"file_url": file_url})
+
 
 def get_vector_service(
     index_name: str, credentials: VectorDatabase, dimension: int = 1024