hotfix: rm useless test file & add vector_stores for multiple stores (#…

…512)
modelscope · Jul 3, 2024 · f2d878f · f2d878f
1 parent bddd086
commit f2d878f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/modelscope_agent/rag/knowledge.py b/modelscope_agent/rag/knowledge.py
@@ -74,6 +74,9 @@ def __init__(self,
                  vector_store: Union[BasePydanticVectorStore,
                                      Type[BasePydanticVectorStore],
                                      None] = None,
+                 vector_stores: Dict[
+                     str, Union[BasePydanticVectorStore,
+                                Type[BasePydanticVectorStore]]] = {},
                  image_store: Union[BasePydanticVectorStore,
                                     Type[BasePydanticVectorStore],
                                     None] = None,
@@ -90,6 +93,10 @@ def __init__(self,
         vector_store = self.get_storage(vector_store)
         image_store = self.get_storage(image_store)
         graph_store = self.get_storage(graph_store)
+        v_stores = dict()
+        if vector_stores:
+            for key, value in vector_stores.items():
+                v_stores[key] = self.get_storage(value)
 
         documents.extend(self.read(files))
 
@@ -108,6 +115,7 @@ def __init__(self,
             docstore=docstore,
             index_store=index_store,
             vector_store=vector_store,
+            vector_stores=v_stores,
             image_store=image_store,
             graph_store=graph_store,
             **kwargs)
@@ -219,6 +227,8 @@ def get_index(self,
                   docstore: Optional[BaseDocumentStore] = None,
                   index_store: Optional[BaseIndexStore] = None,
                   vector_store: Optional[BasePydanticVectorStore] = None,
+                  vector_stores: Optional[Dict[
+                      str, BasePydanticVectorStore]] = None,
                   image_store: Optional[BasePydanticVectorStore] = None,
                   graph_store: Optional[GraphStore] = None,
                   **kwargs) -> BaseIndex:
@@ -235,6 +245,7 @@ def get_index(self,
                         docstore=docstore,
                         index_store=index_store,
                         vector_store=vector_store,
+                        vector_stores=vector_stores,
                         image_store=image_store,
                         graph_store=graph_store,
                         persist_dir=self.cache_dir)
@@ -271,6 +282,7 @@ def get_root_retriever(
             docstore: Optional[BaseDocumentStore] = None,
             index_store: Optional[BaseIndexStore] = None,
             vector_store: Optional[BasePydanticVectorStore] = None,
+            vector_stores: Optional[Dict[str, BasePydanticVectorStore]] = None,
             image_store: Optional[BasePydanticVectorStore] = None,
             graph_store: Optional[GraphStore] = None,
             **kwargs) -> BaseRetriever:
@@ -280,6 +292,7 @@ def get_root_retriever(
             docstore=docstore,
             index_store=index_store,
             vector_store=vector_store,
+            vector_stores=vector_stores,
             image_store=image_store,
             graph_store=graph_store,
             **kwargs)

diff --git a/tests/samples/code.py b/tests/samples/code.py