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
def search_documents(self) -> List[Data]:
from typing import Optional, Union
from bson import ObjectId
from langchain.schema import Document
import json
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):
docs = vector_store.similarity_search(
query=self.search_input,
k=self.number_of_results,
)
for index in range(len(docs)):
docs[index].metadata = {
key: str(docs[index].metadata[key])
if isinstance(docs[index].metadata[key], ObjectId)
else docs[index].metadata[key]
for key in docs[index].metadata
for doc in docs:
doc.metadata = {
key: str(value) if isinstance(value, ObjectId) else value for key, value in doc.metadata.items()
}
data = docs_to_data(docs)
self.status = data
return data