3
3
4
4
5
5
import os
6
+ from types import SimpleNamespace
6
7
7
8
from haystack_integrations .components .retrievers .qdrant import QdrantEmbeddingRetriever
8
9
from haystack_integrations .document_stores .qdrant import QdrantDocumentStore
@@ -26,7 +27,7 @@ class OpeaQDrantRetriever(OpeaComponent):
26
27
def __init__ (self , name : str , description : str , config : dict = None ):
27
28
super ().__init__ (name , ServiceType .RETRIEVER .name .lower (), description , config )
28
29
29
- self .retriever = self ._initialize_client ()
30
+ self .db_store , self . retriever = self ._initialize_client ()
30
31
health_status = self .check_health ()
31
32
if not health_status :
32
33
logger .error ("OpeaQDrantRetriever health check failed." )
@@ -43,7 +44,7 @@ def _initialize_client(self) -> QdrantEmbeddingRetriever:
43
44
44
45
retriever = QdrantEmbeddingRetriever (document_store = qdrant_store )
45
46
46
- return retriever
47
+ return qdrant_store , retriever
47
48
48
49
def check_health (self ) -> bool :
49
50
"""Checks the health of the retriever service.
@@ -55,7 +56,7 @@ def check_health(self) -> bool:
55
56
logger .info ("[ check health ] start to check health of QDrant" )
56
57
try :
57
58
# Check the status of the QDrant service
58
- _ = self .retriever .client
59
+ _ = self .db_store .client
59
60
logger .info ("[ check health ] Successfully connected to QDrant!" )
60
61
return True
61
62
except Exception as e :
@@ -75,6 +76,14 @@ async def invoke(self, input: EmbedDoc) -> list:
75
76
76
77
search_res = self .retriever .run (query_embedding = input .embedding )["documents" ]
77
78
79
+ # format result to align with the standard output in opea_retrievers_microservice.py
80
+ final_res = []
81
+ for res in search_res :
82
+ dict_res = res .meta
83
+ res_obj = SimpleNamespace (** dict_res )
84
+ final_res .append (res_obj )
85
+
78
86
if logflag :
79
- logger .info (f"[ similarity search ] search result: { search_res } " )
80
- return search_res
87
+ logger .info (f"[ similarity search ] search result: { final_res } " )
88
+
89
+ return final_res
0 commit comments