store api breaking changes

LazyAGI · Oct 14, 2024 · e143940 · e143940
1 parent 8e0d791
commit e143940
Show file tree

Hide file tree

Showing 6 changed files with 393 additions and 86 deletions.
diff --git a/lazyllm/tools/rag/doc_impl.py b/lazyllm/tools/rag/doc_impl.py
@@ -52,9 +52,9 @@ def _lazy_init(self) -> None:
 
         self.store = self._get_store()
         self.index = DefaultIndex(self.embed, self.store)
-        if not self.store.has_nodes(LAZY_ROOT_NAME):
+        if not self.store.has_group(LAZY_ROOT_NAME):
             root_nodes = self._reader.load_data(self._list_files())
-            self.store.add_nodes(root_nodes)
+            self.store.update_nodes(root_nodes)
             LOG.debug(f"building {LAZY_ROOT_NAME} nodes: {root_nodes}")
 
     def _get_store(self) -> BaseStore:
@@ -63,7 +63,6 @@ def _get_store(self) -> BaseStore:
             store = MapStore(node_groups=self.node_groups.keys())
         elif rag_store_type == "chroma":
             store = ChromadbStore(node_groups=self.node_groups.keys(), embed=self.embed)
-            store.try_load_store()
         else:
             raise NotImplementedError(
                 f"Not implemented store type for {rag_store_type}"
@@ -165,13 +164,13 @@ def _add_files(self, input_files: List[str]):
         self._lazy_init()
         root_nodes = self._reader.load_data(input_files)
         temp_store = self._get_store()
-        temp_store.add_nodes(root_nodes)
+        temp_store.update_nodes(root_nodes)
         active_groups = self.store.active_groups()
         LOG.info(f"add_files: Trying to merge store with {active_groups}")
         for group in active_groups:
             # Duplicate group will be discarded automatically
             nodes = self._get_nodes(group, temp_store)
-            self.store.add_nodes(nodes)
+            self.store.update_nodes(nodes)
             LOG.debug(f"Merge {group} with {nodes}")
 
     def _delete_files(self, input_files: List[str]) -> None:
@@ -202,7 +201,7 @@ def gather_children(node: DocNode):
             LOG.debug(f"Removed nodes from group {group} for node IDs: {node_uids}")
 
     def _dynamic_create_nodes(self, group_name: str, store: BaseStore) -> None:
-        if store.has_nodes(group_name):
+        if store.has_group(group_name):
             return
         node_group = self.node_groups.get(group_name)
         if node_group is None:
@@ -212,7 +211,7 @@ def _dynamic_create_nodes(self, group_name: str, store: BaseStore) -> None:
         transform = AdaptiveTransform(t) if isinstance(t, list) or t.pattern else make_transform(t)
         parent_nodes = self._get_nodes(node_group["parent"], store)
         nodes = transform.batch_forward(parent_nodes, group_name)
-        store.add_nodes(nodes)
+        store.update_nodes(nodes)
         LOG.debug(f"building {group_name} nodes: {nodes}")
 
     def _get_nodes(self, group_name: str, store: Optional[BaseStore] = None) -> List[DocNode]:

diff --git a/lazyllm/tools/rag/index.py b/lazyllm/tools/rag/index.py
@@ -96,7 +96,7 @@ def query(
             assert len(query) > 0, "Query should not be empty."
             query_embedding = {k: self.embed[k](query) for k in (embed_keys or self.embed.keys())}
             nodes = self._parallel_do_embedding(nodes)
-            self.store.try_save_nodes(nodes)
+            self.store.update_nodes(nodes)
             similarities = similarity_func(query_embedding, nodes, topk=topk, **kwargs)
         elif mode == "text":
             similarities = similarity_func(query, nodes, topk=topk, **kwargs)