Return filename on search request

weni-ai · May 2, 2024 · 94d1f5a · 94d1f5a
1 parent a67af8c
commit 94d1f5a
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/app/handlers/content_bases.py b/app/handlers/content_bases.py
@@ -33,7 +33,7 @@ class ContentBaseSearchRequest(BaseModel):
 
 
 class ContentBaseSearchResponse(BaseModel):
-    response: List[str]
+    response: List[dict]
 
 
 class ContentBaseDeleteRequest(BaseModel):

diff --git a/app/indexer/content_bases.py b/app/indexer/content_bases.py
@@ -47,7 +47,20 @@ def index_batch(self):
 
     def search(self, search, filter=None, threshold=0.1) -> list[Product]:
         matched_responses = self.storage.search(search, filter, threshold)
-        return set([doc.metadata.get("full_page") for doc in matched_responses])
+
+        seen = set()
+        return_list = []
+
+        for doc in matched_responses:
+            full_page = doc.metadata.get("full_page")
+            if full_page not in seen:
+                seen.add(full_page)
+                return_list.append({
+                    "full_page": full_page,
+                    "filename": doc.metadata.get("filename")
+                })
+
+        return return_list
 
     def _search_docs_by_content_base_uuid(self, content_base_uuid: UUID, file_uuid: str = None):
         search_filter = {

diff --git a/app/store/elasticsearch_vector_store.py b/app/store/elasticsearch_vector_store.py
@@ -124,6 +124,7 @@ def search_delete(self, search_filter: dict, scroll_id: str = None) -> tuple[str
 
     def search(self, search: str, filter=None, threshold=0.1) -> list[Document]:
         docs = self.vectorstore.similarity_search_with_score(query=search, k=5, filter=filter)
+        print(("[---] Response: ", docs[0][0].metadata.get('filename')))
         return [doc[0] for doc in docs if doc[1] > threshold]
 
     def delete(self, ids: list[str] = []) -> bool: