Add MilvusStore and re-implement MapStore and ChromadbStore, support multi index for one store

lazyllm/tools/rag/chroma_store.py

@@ -10,23 +10,28 @@
 from .default_index import DefaultIndex
 from .map_store import MapStore
 import pickle
+import base64
 
 # ---------------------------------------------------------------------------- #
 
 class ChromadbStore(StoreBase):
-    def __init__(self, dir: str, embed: Dict[str, Callable], embed_dim: Dict[str, int], **kwargs) -> None:
+    def __init__(self, dir: str, node_groups: List[str], embed: Dict[str, Callable],
+                 embed_dim: Dict[str, int], **kwargs) -> None:
         self._db_client = chromadb.PersistentClient(path=dir)
         LOG.success(f"Initialzed chromadb in path: {dir}")
-        self._collections: Dict[str, Collection] = {}
+        self._collections: Dict[str, Collection] = {
+            group: self._db_client.get_or_create_collection(group)
+            for group in node_groups
+        }
+
+        self._map_store = MapStore(node_groups=node_groups, embed=embed)
+        self._load_store(embed_dim)
 
         self._name2index = {
             'default': DefaultIndex(embed, self._map_store),
             'file_node_map': _FileNodeIndex(),
         }
 
-        self._map_store = MapStore(embed=embed)
-        self._load_store(embed_dim)
-
     @override
     def update_nodes(self, nodes: List[DocNode]) -> None:
         self._map_store.update_nodes(nodes)
@@ -52,11 +57,6 @@ def is_group_active(self, name: str) -> bool:
     def all_groups(self) -> List[str]:
         return self._map_store.all_groups()
 
-    @override
-    def add_group(self, name: str, embed_keys: Optional[List[str]] = None) -> None:
-        self._collections[name] = self._db_client.get_or_create_collection(name)
-        self._map_store.add_group(name, embed_keys)
-
     @override
     def query(self, *args, **kwargs) -> List[DocNode]:
         return self.get_index('default').query(*args, **kwargs)
@@ -133,13 +133,14 @@ def _build_nodes_from_chroma(self, results: Dict[str, List], embed_dim: Dict[str
             chroma_metadata = results['metadatas'][i]
 
             parent = chroma_metadata['parent']
-            fields = pickle.loads(chroma_metadata['fields']) if parent else None
+            fields = pickle.loads(base64.b64decode(chroma_metadata['fields'].encode('utf-8')))\
+                if parent else None
 
             node = DocNode(
                 uid=uid,
                 text=results["documents"][i],
                 group=chroma_metadata["group"],
-                embedding=pickle.loads(chroma_metadata['embedding']),
+                embedding=pickle.loads(base64.b64decode(chroma_metadata['embedding'].encode('utf-8'))),
                 parent=parent,
                 fields=fields,
             )
@@ -168,11 +169,11 @@ def _make_chroma_metadata(self, node: DocNode) -> Dict[str, Any]:
         metadata = {
             "group": node.group,
             "parent": node.parent.uid if node.parent else "",
-            "embedding": pickle.dumps(node.embedding),
+            "embedding": base64.b64encode(pickle.dumps(node.embedding)).decode('utf-8'),
         }
 
         if node.parent:
-            metadata["fields"] = pickle.dumps(node.fields)
+            metadata["fields"] = base64.b64encode(pickle.dumps(node.fields)).decode('utf-8')
 
         return metadata
 

lazyllm/tools/rag/doc_field_desc.py

@@ -0,0 +1,8 @@
+from typing import Optional
+
+class DocFieldDesc:
+    DTYPE_VARCHAR = 0
+
+    def __init__(self, data_type: int = DTYPE_VARCHAR, max_length: Optional[int] = 65535):
+        self.data_type = data_type
+        self.max_length = max_length

lazyllm/tools/rag/doc_impl.py

@@ -13,7 +13,7 @@
 from .doc_node import DocNode
 from .data_loaders import DirectoryReader
 from .utils import DocListManager
-from .doc_field_info import DocFieldInfo
+from .doc_field_desc import DocFieldDesc
 import threading
 import time
 
@@ -38,7 +38,7 @@ class DocImpl:
 
     def __init__(self, embed: Dict[str, Callable], dlm: Optional[DocListManager] = None,
                  doc_files: Optional[str] = None, kb_group_name: Optional[str] = None,
-                 fields_info: Dict[str, DocFieldInfo] = None, store_conf: Optional[Dict] = None):
+                 fields_desc: Dict[str, DocFieldDesc] = None, store_conf: Optional[Dict] = None):
         super().__init__()
         assert (dlm is None) ^ (doc_files is None), 'Only one of dataset_path or doc_files should be provided'
         self._local_file_reader: Dict[str, Callable] = {}
@@ -48,8 +48,9 @@ def __init__(self, embed: Dict[str, Callable], dlm: Optional[DocListManager] = N
         self.node_groups: Dict[str, Dict] = {LAZY_ROOT_NAME: {}}
         self.embed = {k: embed_wrapper(e) for k, e in embed.items()}
         self._embed_dim = None
-        self._fields_info = fields_info
+        self._fields_desc = fields_desc
         self.store = store_conf  # NOTE: will be initialized in _lazy_init()
+        self._activated_embeddings = set()
 
     @once_wrapper(reset_on_pickle=True)
     def _lazy_init(self) -> None:
@@ -71,7 +72,6 @@ def _lazy_init(self) -> None:
             raise ValueError(f'store type [{type(self.store)}] is not a dict.')
 
         if not self.store.is_group_active(LAZY_ROOT_NAME):
-            self.store.add_group(name=LAZY_ROOT_NAME, fields_info=self._fields_info, embed_keys={})
             ids, pathes = self._list_files()
             root_nodes = self._reader.load_data(pathes)
             self.store.update_nodes(root_nodes)
@@ -94,11 +94,14 @@ def _create_store(self, store_conf: Optional[Dict]) -> StoreBase:
             raise ValueError('`kwargs` in store conf is not a dict.')
 
         if store_type == "map":
-            store = MapStore(embed=self.embed, **kwargs)
+            store = MapStore(node_groups=self.node_groups, embed=self.embed, **kwargs)
         elif store_type == "chroma":
-            store = ChromadbStore(embed_dim=self.embed_dim, embed=self.embed, **kwargs)
+            store = ChromadbStore(node_groups=self.node_groups, embed=self.embed,
+                                  embed_dim=self.embed_dim, **kwargs)
         elif store_type == "milvus":
-            store = MilvusStore(embed=self.embed, fields_info=self._fields_info, **kwargs)
+            store = MilvusStore(node_groups=self.node_groups, embed=self.embed,
+                                embed_keys=self._activated_embeddings,
+                                fields_desc=self._fields_desc, **kwargs)
         else:
             raise NotImplementedError(
                 f"Not implemented store type for {store_type}"
@@ -216,7 +219,7 @@ def _add_files(self, input_files: List[str]):
         if len(input_files) == 0:
             return
         root_nodes = self._reader.load_data(input_files)
-        temp_store = self._create_store("map")
+        temp_store = self._create_store({"type": "map"})
         temp_store.update_nodes(root_nodes)
         all_groups = self.store.all_groups()
         LOG.info(f"add_files: Trying to merge store with {all_groups}")
@@ -277,9 +280,6 @@ def retrieve(self, query: str, group_name: str, similarity: str, similarity_cut_
                  index: str, topk: int, similarity_kws: dict, embed_keys: Optional[List[str]] = None) -> List[DocNode]:
         self._lazy_init()
 
-        if not self.store.is_group_active(group_name):
-            self.store.add_group(group_name, embed_keys=embed_keys)
-
         if type is None or type == 'default':
             return self.store.query(query=query, group_name=group_name, similarity_name=similarity,
                                     similarity_cut_off=similarity_cut_off, topk=topk,

lazyllm/tools/rag/doc_node.py

@@ -30,13 +30,12 @@ def __init__(self, uid: Optional[str] = None, text: Optional[str] = None, group:
         self.parent: Optional["DocNode"] = parent
         self.children: Dict[str, List["DocNode"]] = defaultdict(list)
         self.is_saved: bool = False
-        self._docpath = None
         self._lock = threading.Lock()
         self._embedding_state = set()
 
         if fields and parent:
             raise ValueError('only ROOT node can set fields.')
-        self._fields = fields
+        self._fields = fields if fields else {}
 
     @property
     def root_node(self) -> Optional["DocNode"]:
@@ -76,12 +75,12 @@ def excluded_llm_metadata_keys(self, excluded_llm_metadata_keys: List) -> None:
 
     @property
     def docpath(self) -> str:
-        return self.root_node._docpath or ''
+        return self.root_node._fields.get('lazyllm_doc_path', '')
 
     @docpath.setter
     def docpath(self, path):
         assert not self.parent, 'Only root node can set docpath'
-        self._docpath = str(path)
+        self._fields['lazyllm_doc_path'] = str(path)
 
     def get_children_str(self) -> str:
         return str(

lazyllm/tools/rag/document.py

@@ -10,7 +10,8 @@
 from .doc_node import DocNode
 from .store_base import LAZY_ROOT_NAME, EMBED_DEFAULT_KEY
 from .utils import DocListManager
-from .doc_field_info import DocFieldInfo
+from .doc_field_desc import DocFieldDesc
+from .web import DocWebModule
 import copy
 import functools
 
@@ -22,9 +23,9 @@ def __call__(self, cls, *args, **kw):
 class Document(ModuleBase):
     class _Impl(ModuleBase):
         def __init__(self, dataset_path: str, embed: Optional[Union[Callable, Dict[str, Callable]]] = None,
-                     manager: bool = False, server: bool = False, name: Optional[str] = None,
+                     manager: Union[bool, str] = False, server: bool = False, name: Optional[str] = None,
                      launcher: Optional[Launcher] = None, store_conf: Optional[Dict] = None,
-                     fields_info: Optional[Dict[str, DocFieldInfo]] = None):
+                     fields_desc: Optional[Dict[str, DocFieldDesc]] = None):
             super().__init__()
             if not os.path.exists(dataset_path):
                 defatult_path = os.path.join(lazyllm.config["data_path"], dataset_path)
@@ -39,19 +40,21 @@ def __init__(self, dataset_path: str, embed: Optional[Union[Callable, Dict[str,
                     self._submodules.append(embed)
             self._dlm = DocListManager(dataset_path, name).init_tables()
             self._kbs = CallableDict({DocListManager.DEDAULT_GROUP_NAME:
-                                      DocImpl(embed=self._embed, dlm=self._dlm, store_conf=store_conf)})
+                                      DocImpl(embed=self._embed, dlm=self._dlm, fields_desc=fields_desc,
+                                              store_conf=store_conf)})
             if manager: self._manager = ServerModule(DocManager(self._dlm))
             if manager == 'ui': self._docweb = DocWebModule(doc_server=self._manager)
             if server: self._kbs = ServerModule(self._kbs)
-            self._fields_info = fields_info
+            self._fields_desc = fields_desc
 
-        def add_kb_group(self, name, store_conf: Optional[Dict] = None):
+        def add_kb_group(self, name, fields_desc: Optional[Dict[str, DocFieldDesc]] = None,
+                         store_conf: Optional[Dict] = None):
             if isinstance(self._kbs, ServerModule):
                 self._kbs._impl._m[name] = DocImpl(dlm=self._dlm, embed=self._embed, kb_group_name=name,
-                                                   store_conf=store_conf)
+                                                   fields_desc=fields_desc, store_conf=store_conf)
             else:
                 self._kbs[name] = DocImpl(dlm=self._dlm, embed=self._embed, kb_group_name=name,
-                                          store_conf=store_conf)
+                                          fields_desc=fields_desc, store_conf=store_conf)
             self._dlm.add_kb_group(name)
 
         def get_doc_by_kb_group(self, name):
@@ -66,14 +69,14 @@ def __call__(self, *args, **kw):
             return self._kbs(*args, **kw)
 
     def __init__(self, dataset_path: str, embed: Optional[Union[Callable, Dict[str, Callable]]] = None,
-                 create_ui: bool = False, manager: bool = False, server: bool = False,
+                 create_ui: bool = False, manager: Union[bool, str] = False, server: bool = False,
                  name: Optional[str] = None, launcher: Optional[Launcher] = None,
-                 fields_info: Dict[str, DocFieldInfo] = None, store_conf: Optional[Dict] = None):
+                 fields_desc: Dict[str, DocFieldDesc] = None, store_conf: Optional[Dict] = None):
         super().__init__()
         if create_ui:
             lazyllm.LOG.warning('`create_ui` for Document is deprecated, use `manager` instead')
         self._impls = Document._Impl(dataset_path, embed, create_ui or manager, server, name,
-                                     launcher, store_conf, fields_info)
+                                     launcher, store_conf, fields_desc)
         self._curr_group = DocListManager.DEDAULT_GROUP_NAME
 
     def create_kb_group(self, name: str, store_conf: Optional[Dict] = None) -> "Document":

lazyllm/tools/rag/map_store.py

@@ -16,9 +16,11 @@ def _remove_from_indices(name2index: Dict[str, IndexBase], uids: List[str],
         index.remove(uids, group_name)
 
 class MapStore(StoreBase):
-    def __init__(self, embed: Dict[str, Callable], **kwargs):
+    def __init__(self, node_groups: List[str], embed: Dict[str, Callable], **kwargs):
         # Dict[group_name, Dict[uuid, DocNode]]
-        self._group2docs: Dict[str, Dict[str, DocNode]] = {}
+        self._group2docs: Dict[str, Dict[str, DocNode]] = {
+            group: {} for group in node_groups
+        }
 
         self._name2index = {
             'default': DefaultIndex(embed, self),
@@ -72,10 +74,6 @@ def is_group_active(self, name: str) -> bool:
     def all_groups(self) -> List[str]:
         return self._group2docs.keys()
 
-    @override
-    def add_group(self, name: str, embed_keys: Optional[List[str]] = None) -> None:
-        self._group2docs.setdefault(name, {})
-
     @override
     def query(self, *args, **kwargs) -> List[DocNode]:
         return self.get_index('default').query(*args, **kwargs)