chore: Adjust MongoDB Atlas Vector Store serialization

This commit is contained in:
joaoguilhermeS 2024-06-22 01:34:17 -03:00 • committed by Gabriel Luiz Freitas Almeida
commit 5ba076cd71

View file

@ -83,31 +83,21 @@ class MongoVectorStoreComponent(LCVectorStoreComponent):
return vector_store return vector_store
def search_documents(self) -> List[Data]: def search_documents(self) -> List[Data]:
from typing import Optional, Union
from bson import ObjectId from bson import ObjectId
from langchain.schema import Document from langchain.schema import Document
import json
vector_store = self._build_mongodb_atlas() vector_store = self._build_mongodb_atlas()
class JSONEncoder(json.JSONEncoder):
def default(self, o):
if isinstance(o, Union[ObjectId, Document]):
return str(o)
return super(JSONEncoder, self).default(o)
if self.search_input and isinstance(self.search_input, str): if self.search_input and isinstance(self.search_input, str):
docs = vector_store.similarity_search( docs = vector_store.similarity_search(
query=self.search_input, query=self.search_input,
k=self.number_of_results, k=self.number_of_results,
) )
for index in range(len(docs)): for doc in docs:
docs[index].metadata = { doc.metadata = {
key: str(docs[index].metadata[key]) key: str(value) if isinstance(value, ObjectId) else value for key, value in doc.metadata.items()
if isinstance(docs[index].metadata[key], ObjectId)
else docs[index].metadata[key]
for key in docs[index].metadata
} }
data = docs_to_data(docs) data = docs_to_data(docs)
self.status = data self.status = data
return data return data