update vector store

.gitignore

@@ -27,4 +27,5 @@ build/*
 cookbook/appworld/data/*
 cookbook/appworld/experiments/*
 cookbook/appworld/exp_result/*
-file_vector_store/*
+file_vector_store/*
+cookbook/appworld/file_vector_store/*

experiencemaker/vector_store/base_vector_store.py

@@ -28,7 +28,9 @@ class BaseVectorStore(BaseModel, ABC):
             try:
                 for line in tqdm(f, desc="load from path"):
                     if line.strip():
-                        yield VectorNode(**json.loads(line.strip(), **kwargs))
+                        node = VectorNode(**json.loads(line.strip(), **kwargs))
+                        node.workspace_id = workspace_id
+                        yield node
 
             finally:
                 fcntl.flock(f, fcntl.LOCK_UN)
@@ -45,6 +47,7 @@ class BaseVectorStore(BaseModel, ABC):
             fcntl.flock(f, fcntl.LOCK_EX)
             try:
                 for node in tqdm(nodes, desc="dump to path"):
+                    node.workspace_id = workspace_id
                     f.write(json.dumps(node.model_dump(), ensure_ascii=ensure_ascii, **kwargs))
                     f.write("\n")
                     count += 1

experiencemaker/vector_store/es_vector_store.py

@@ -51,14 +51,15 @@ class EsVectorStore(BaseVectorStore):
     def _iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
         response = self._client.search(index=workspace_id, body={"query": {"match_all": {}}})
         for doc in response['hits']['hits']:
-            yield self.doc2node(doc)
+            yield self.doc2node(doc, workspace_id)
 
     def refresh(self, workspace_id: str):
         self._client.indices.refresh(index=workspace_id)
 
     @staticmethod
-    def doc2node(doc) -> VectorNode:
+    def doc2node(doc, workspace_id: str) -> VectorNode:
         node = VectorNode(**doc["_source"])
+        node.workspace_id = workspace_id
         node.unique_id = doc["_id"]
         if "_score" in doc:
             node.metadata["_score"] = doc["_score"] - 1
@@ -105,7 +106,7 @@ class EsVectorStore(BaseVectorStore):
 
         nodes: List[VectorNode] = []
         for doc in response['hits']['hits']:
-            nodes.append(self.doc2node(doc))
+            nodes.append(self.doc2node(doc, workspace_id))
 
         self.retrieve_filters.clear()
         return nodes
@@ -124,10 +125,10 @@ class EsVectorStore(BaseVectorStore):
         docs = [
             {
                 "_op_type": "index",
-                "_index": node.workspace_id,
+                "_index": workspace_id,
                 "_id": node.unique_id,
                 "_source": {
-                    "workspace_id": node.workspace_id,
+                    "workspace_id": workspace_id,
                     "content": node.content,
                     "metadata": node.metadata,
                     "vector": node.vector

experiencemaker/vector_store/file_vector_store.py

@@ -84,6 +84,7 @@ class FileVectorStore(BaseVectorStore):
                            workspace_id=workspace_id,
                            path=self.store_path,
                            **kwargs)
+
         logger.info(f"update workspace_id={workspace_id} nodes.size={len(nodes)} all.size={len(all_node_dict)} "
                     f"update_cnt={update_cnt}")