fix: bug when deleting sources

zaldivards · Aug 23, 2024 · af069e6 · af069e6
1 parent c6bdb7a
commit af069e6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/api/contextqa/services/sources.py b/api/contextqa/services/sources.py
@@ -84,17 +84,20 @@ def remove_sources(session: Session, sources: list[str], client: StoreClient) ->
         Specific store client
     """
     store_settings = get_or_set(kind="store")
+    db_index = store_settings.store_params.get("collection", "index")
+    complete_index_name = (
+        f"{store_settings.store_params['home']}/{db_index}" if store_settings.store == "chroma" else db_index
+    )
     sources_to_remove = (
         session.query(Source.id)
         .join(Index)
         .join(VectorStore)
         .filter(
             Source.name.in_(sources),
-            Index.name == store_settings.store_params.get("collection", "index"),
+            Index.name == complete_index_name,
             VectorStore.name == store_settings.store,
         )
     )
-
     removed_sources = session.query(Source).filter(Source.id.in_(sources_to_remove)).delete(synchronize_session=False)
 
     chunks_to_remove = client.get_ids(sources)

diff --git a/api/contextqa/utils/clients.py b/api/contextqa/utils/clients.py
@@ -39,7 +39,8 @@ def get_ids(self, sources: list[str]) -> list[str]:
         chunks_to_remove = []
         for source in sources:
             if source.endswith(".pdf"):
-                source = f"{get_or_set('extra').media_dir}/{source}"
+                # pdf sources have the template prefix: {}
+                source = f"{{}}/{source}"
             chunks = self.get(source)
             chunks_to_remove.extend(chunks)