feat: extend VectorStore parameters

Signed-off-by: Tomáš Dvořák <toomas2d@gmail.com>

Author: Tomas2D

python/beeai_framework/adapters/langchain/backend/vector_store.py

@@ -30,9 +30,9 @@ def __init__(self, *, vector_store: LCVectorStore) -> None:
         super().__init__()
         self.vector_store: LCVectorStore = vector_store
 
-    async def add_documents(self, documents: list[Document]) -> list[str]:
+    async def add_documents(self, documents: list[Document], **kwargs: Any) -> list[str]:
         lc_documents = [document_to_lc_document(document) for document in documents]
-        return await self.vector_store.aadd_documents(lc_documents)
+        return await self.vector_store.aadd_documents(lc_documents, **kwargs)
 
     async def search(self, query: QueryLike, k: int = 4, **kwargs: Any) -> list[DocumentWithScore]:
         query_str = str(query)

python/beeai_framework/adapters/langchain/mappers/documents.py

@@ -13,8 +13,8 @@
 
 
 def lc_document_to_document(lc_document: Document) -> VectorStoreDocument:
-    return VectorStoreDocument(content=lc_document.page_content, metadata=lc_document.metadata)
+    return VectorStoreDocument(id=lc_document.id, content=lc_document.page_content, metadata=lc_document.metadata)
 
 
 def document_to_lc_document(document: VectorStoreDocument) -> Document:
-    return Document(page_content=document.content, metadata=document.metadata)
+    return Document(id=document.id, page_content=document.content, metadata=document.metadata)

python/beeai_framework/backend/types.py

@@ -247,14 +247,19 @@ class EmbeddingModelOutput(BaseModel):
 
 
 class Document(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    id: str | None = Field(default=None, coerce_numbers_to_str=True)
+    metadata: dict[str, str | int | float | bool] = Field(default_factory=dict)
     content: str
-    metadata: dict[str, str | int | float | bool]
 
     def __str__(self) -> str:
         return self.content
 
 
 class DocumentWithScore(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
     document: Document
     score: float
 

python/beeai_framework/tools/search/retrieval/vector_store.py

@@ -49,7 +49,7 @@ async def _run(
         self, input: VectorStoreSearchToolInput, options: ToolRunOptions | None, context: RunContext
     ) -> VectorStoreSearchToolOutput:
         # Search the vector store using the provided query and parameters
-        documents_with_score = await self.vector_store.search(query=input.query, k=input.k)
+        documents_with_score = await self.vector_store.search(**input.model_dump())
 
         # Convert DocumentWithScore objects to SearchToolResult format
         results = [