flowllm 0.1.0__py3-none-any.whl

llmflow/vector_store/base_vector_store.py

@@ -0,0 +1,136 @@
+import fcntl
+import json
+from abc import ABC
+from pathlib import Path
+from typing import List, Iterable
+
+from loguru import logger
+from pydantic import BaseModel, Field
+from tqdm import tqdm
+
+from llmflow.embedding_model.base_embedding_model import BaseEmbeddingModel
+from llmflow.schema.vector_node import VectorNode
+
+
+class BaseVectorStore(BaseModel, ABC):
+    embedding_model: BaseEmbeddingModel | None = Field(default=None)
+    batch_size: int = Field(default=1024)
+
+    @staticmethod
+    def _load_from_path(workspace_id: str, path: str | Path, callback_fn=None, **kwargs) -> Iterable[VectorNode]:
+        workspace_path = Path(path) / f"{workspace_id}.jsonl"
+        if not workspace_path.exists():
+            logger.warning(f"workspace_path={workspace_path} is not exists!")
+            return
+
+        with workspace_path.open() as f:
+            fcntl.flock(f, fcntl.LOCK_SH)
+            try:
+                for line in tqdm(f, desc="load from path"):
+                    if line.strip():
+                        node_dict = json.loads(line.strip())
+                        if callback_fn:
+                            node = callback_fn(node_dict)
+                        else:
+                            node = VectorNode(**node_dict, **kwargs)
+                        node.workspace_id = workspace_id
+                        yield node
+
+            finally:
+                fcntl.flock(f, fcntl.LOCK_UN)
+
+    @staticmethod
+    def _dump_to_path(nodes: Iterable[VectorNode], workspace_id: str, path: str | Path = "", callback_fn=None,
+                      ensure_ascii: bool = False, **kwargs):
+        dump_path: Path = Path(path)
+        dump_path.mkdir(parents=True, exist_ok=True)
+        dump_file = dump_path / f"{workspace_id}.jsonl"
+
+        count = 0
+        with dump_file.open("w") as f:
+            fcntl.flock(f, fcntl.LOCK_EX)
+            try:
+                for node in tqdm(nodes, desc="dump to path"):
+                    node.workspace_id = workspace_id
+                    if callback_fn:
+                        node_dict = callback_fn(node)
+                    else:
+                        node_dict = node.model_dump()
+                    assert isinstance(node_dict, dict)
+                    f.write(json.dumps(node_dict, ensure_ascii=ensure_ascii, **kwargs))
+                    f.write("\n")
+                    count += 1
+
+                return {"size": count}
+            finally:
+                fcntl.flock(f, fcntl.LOCK_UN)
+
+    def exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        raise NotImplementedError
+
+    def delete_workspace(self, workspace_id: str, **kwargs):
+        raise NotImplementedError
+
+    def create_workspace(self, workspace_id: str, **kwargs):
+        raise NotImplementedError
+
+    def _iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
+        raise NotImplementedError
+
+    def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs):
+        if not self.exist_workspace(workspace_id=workspace_id, **kwargs):
+            logger.warning(f"workspace_id={workspace_id} is not exist!")
+            return {}
+
+        return self._dump_to_path(nodes=self._iter_workspace_nodes(workspace_id=workspace_id, **kwargs),
+                                  workspace_id=workspace_id,
+                                  path=path,
+                                  callback_fn=callback_fn,
+                                  **kwargs)
+
+    def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None, callback_fn=None,
+                       **kwargs):
+        if self.exist_workspace(workspace_id, **kwargs):
+            self.delete_workspace(workspace_id=workspace_id, **kwargs)
+            logger.info(f"delete workspace_id={workspace_id}")
+
+        self.create_workspace(workspace_id=workspace_id, **kwargs)
+
+        all_nodes: List[VectorNode] = []
+        if nodes:
+            all_nodes.extend(nodes)
+        for node in self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs):
+            all_nodes.append(node)
+        self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs)
+        return {"size": len(all_nodes)}
+
+    def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs):
+        if not self.exist_workspace(workspace_id=src_workspace_id, **kwargs):
+            logger.warning(f"src_workspace_id={src_workspace_id} is not exist!")
+            return {}
+
+        if not self.exist_workspace(dest_workspace_id, **kwargs):
+            self.create_workspace(workspace_id=dest_workspace_id, **kwargs)
+
+        nodes = []
+        node_size = 0
+        for node in self._iter_workspace_nodes(workspace_id=src_workspace_id, **kwargs):
+            nodes.append(node)
+            node_size += 1
+            if len(nodes) >= self.batch_size:
+                self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs)
+                nodes.clear()
+
+        if nodes:
+            self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs)
+        return {"size": node_size}
+
+    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+        raise NotImplementedError
+
+    def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+        raise NotImplementedError
+
+    def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs):
+        raise NotImplementedError
+

llmflow/vector_store/chroma_vector_store.py

@@ -0,0 +1,188 @@
+from typing import List, Iterable
+
+import chromadb
+from chromadb import Collection
+from chromadb.config import Settings
+from loguru import logger
+from pydantic import Field, PrivateAttr, model_validator
+
+from llmflow.embedding_model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel
+from llmflow.schema.vector_node import VectorNode
+from llmflow.vector_store import VECTOR_STORE_REGISTRY
+from llmflow.vector_store.base_vector_store import BaseVectorStore
+
+
+@VECTOR_STORE_REGISTRY.register("chroma")
+class ChromaVectorStore(BaseVectorStore):
+    store_dir: str = Field(default="./chroma_vector_store")
+    collections: dict = Field(default_factory=dict)
+    _client: chromadb.Client = PrivateAttr()
+
+    @model_validator(mode="after")
+    def init_client(self):
+        self._client = chromadb.Client(Settings(persist_directory=self.store_dir))
+        return self
+
+    def _get_collection(self, workspace_id: str) -> Collection:
+        if workspace_id not in self.collections:
+            self.collections[workspace_id] = self._client.get_or_create_collection(workspace_id)
+        return self.collections[workspace_id]
+
+    def exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        return workspace_id in [c.name for c in self._client.list_collections()]
+
+    def delete_workspace(self, workspace_id: str, **kwargs):
+        self._client.delete_collection(workspace_id)
+        if workspace_id in self.collections:
+            del self.collections[workspace_id]
+
+    def create_workspace(self, workspace_id: str, **kwargs):
+        self.collections[workspace_id] = self._client.get_or_create_collection(workspace_id)
+
+    def _iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
+        collection: Collection = self._get_collection(workspace_id)
+        results = collection.get()
+        for i in range(len(results["ids"])):
+            node = VectorNode(workspace_id=workspace_id,
+                              unique_id=results["ids"][i],
+                              content=results["documents"][i],
+                              metadata=results["metadatas"][i])
+            yield node
+
+    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+        if not self.exist_workspace(workspace_id=workspace_id):
+            logger.warning(f"workspace_id={workspace_id} is not exists!")
+            return []
+
+        collection: Collection = self._get_collection(workspace_id)
+        query_vector = self.embedding_model.get_embeddings(query)
+        results = collection.query(query_embeddings=[query_vector], n_results=top_k)
+        nodes = []
+        for i in range(len(results["ids"][0])):
+            node = VectorNode(workspace_id=workspace_id,
+                              unique_id=results["ids"][0][i],
+                              content=results["documents"][0][i],
+                              metadata=results["metadatas"][0][i])
+            nodes.append(node)
+        return nodes
+
+    def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+        if not self.exist_workspace(workspace_id=workspace_id):
+            self.create_workspace(workspace_id=workspace_id)
+
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        embedded_nodes = [node for node in nodes if node.vector]
+        not_embedded_nodes = [node for node in nodes if not node.vector]
+        now_embedded_nodes = self.embedding_model.get_node_embeddings(not_embedded_nodes)
+        all_nodes = embedded_nodes + now_embedded_nodes
+
+        collection: Collection = self._get_collection(workspace_id)
+        collection.add(ids=[n.unique_id for n in all_nodes],
+                       embeddings=[n.vector for n in all_nodes],
+                       documents=[n.content for n in all_nodes],
+                       metadatas=[n.metadata for n in all_nodes])
+
+    def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs):
+        if not self.exist_workspace(workspace_id=workspace_id):
+            logger.warning(f"workspace_id={workspace_id} is not exists!")
+            return
+
+        if isinstance(node_ids, str):
+            node_ids = [node_ids]
+
+        collection: Collection = self._get_collection(workspace_id)
+        collection.delete(ids=node_ids)
+
+
+def main():
+    from dotenv import load_dotenv
+    load_dotenv()
+
+    embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
+    workspace_id = "chroma_test_index"
+
+    chroma_store = ChromaVectorStore(
+        embedding_model=embedding_model,
+        store_dir="./chroma_test_db"
+    )
+
+    if chroma_store.exist_workspace(workspace_id):
+        chroma_store.delete_workspace(workspace_id)
+    chroma_store.create_workspace(workspace_id)
+
+    sample_nodes = [
+        VectorNode(
+            unique_id="node1",
+            workspace_id=workspace_id,
+            content="Artificial intelligence is a technology that simulates human intelligence.",
+            metadata={
+                "node_type": "n1",
+                "category": "tech"
+            }
+        ),
+        VectorNode(
+            unique_id="node2",
+            workspace_id=workspace_id,
+            content="AI is the future of mankind.",
+            metadata={
+                "node_type": "n1",
+                "category": "tech"
+            }
+        ),
+        VectorNode(
+            unique_id="node3",
+            workspace_id=workspace_id,
+            content="I want to eat fish!",
+            metadata={
+                "node_type": "n2",
+                "category": "food"
+            }
+        ),
+        VectorNode(
+            unique_id="node4",
+            workspace_id=workspace_id,
+            content="The bigger the storm, the more expensive the fish.",
+            metadata={
+                "node_type": "n1",
+                "category": "food"
+            }
+        ),
+    ]
+
+    chroma_store.insert(sample_nodes, workspace_id=workspace_id)
+
+    logger.info("=" * 20)
+    results = chroma_store.search("What is AI?", top_k=5, workspace_id=workspace_id)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
+    node2_update = VectorNode(
+        unique_id="node2",
+        workspace_id=workspace_id,
+        content="AI is the future of humanity and technology.",
+        metadata={
+            "node_type": "n1",
+            "category": "tech",
+            "updated": True
+        }
+    )
+    chroma_store.delete(node2_update.unique_id, workspace_id=workspace_id)
+    chroma_store.insert(node2_update, workspace_id=workspace_id)
+
+    logger.info("Updated Result:")
+    results = chroma_store.search("fish?", top_k=10, workspace_id=workspace_id)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
+    chroma_store.dump_workspace(workspace_id=workspace_id)
+
+    chroma_store.delete_workspace(workspace_id=workspace_id)
+
+
+if __name__ == "__main__":
+    main()
+    # launch with: python -m llmflow.storage.chroma_vector_store

llmflow/vector_store/es_vector_store.py

@@ -0,0 +1,227 @@
+import os
+from typing import List, Tuple, Iterable
+
+from elasticsearch import Elasticsearch
+from elasticsearch.helpers import bulk
+from loguru import logger
+from pydantic import Field, PrivateAttr, model_validator
+
+from llmflow.embedding_model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel
+from llmflow.schema.vector_node import VectorNode
+from llmflow.vector_store import VECTOR_STORE_REGISTRY
+from llmflow.vector_store.base_vector_store import BaseVectorStore
+
+
+@VECTOR_STORE_REGISTRY.register("elasticsearch")
+class EsVectorStore(BaseVectorStore):
+    hosts: str | List[str] = Field(default_factory=lambda: os.getenv("ES_HOSTS", "http://localhost:9200"))
+    basic_auth: str | Tuple[str, str] | None = Field(default=None)
+    retrieve_filters: List[dict] = []
+    _client: Elasticsearch = PrivateAttr()
+
+    @model_validator(mode="after")
+    def init_client(self):
+        if isinstance(self.hosts, str):
+            self.hosts = [self.hosts]
+        self._client = Elasticsearch(hosts=self.hosts, basic_auth=self.basic_auth)
+        return self
+
+    def exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        return self._client.indices.exists(index=workspace_id)
+
+    def delete_workspace(self, workspace_id: str, **kwargs):
+        return self._client.indices.delete(index=workspace_id, **kwargs)
+
+    def create_workspace(self, workspace_id: str, **kwargs):
+        body = {
+            "mappings": {
+                "properties": {
+                    "workspace_id": {"type": "keyword"},
+                    "content": {"type": "text"},
+                    "metadata": {"type": "object"},
+                    "vector": {
+                        "type": "dense_vector",
+                        "dims": self.embedding_model.dimensions
+                    }
+                }
+            }
+        }
+        return self._client.indices.create(index=workspace_id, body=body)
+
+    def _iter_workspace_nodes(self, workspace_id: str, max_size: int = 10000, **kwargs) -> Iterable[VectorNode]:
+        response = self._client.search(index=workspace_id, body={"query": {"match_all": {}}, "size": max_size})
+        for doc in response['hits']['hits']:
+            yield self.doc2node(doc, workspace_id)
+
+    def refresh(self, workspace_id: str):
+        self._client.indices.refresh(index=workspace_id)
+
+    @staticmethod
+    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
+        return node
+
+    def add_term_filter(self, key: str, value):
+        if key:
+            self.retrieve_filters.append({"term": {key: value}})
+        return self
+
+    def add_range_filter(self, key: str, gte=None, lte=None):
+        if key:
+            if gte is not None and lte is not None:
+                self.retrieve_filters.append({"range": {key: {"gte": gte, "lte": lte}}})
+            elif gte is not None:
+                self.retrieve_filters.append({"range": {key: {"gte": gte}}})
+            elif lte is not None:
+                self.retrieve_filters.append({"range": {key: {"lte": lte}}})
+        return self
+
+    def clear_filter(self):
+        self.retrieve_filters.clear()
+        return self
+
+    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+        if not self.exist_workspace(workspace_id=workspace_id):
+            logger.warning(f"workspace_id={workspace_id} is not exists!")
+            return []
+
+        query_vector = self.embedding_model.get_embeddings(query)
+        body = {
+            "query": {
+                "script_score": {
+                    "query": {"bool": {"must": self.retrieve_filters}},
+                    "script": {
+                        "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
+                        "params": {"query_vector": query_vector},
+                    }
+                }
+            },
+            "size": top_k
+        }
+        response = self._client.search(index=workspace_id, body=body, **kwargs)
+
+        nodes: List[VectorNode] = []
+        for doc in response['hits']['hits']:
+            nodes.append(self.doc2node(doc, workspace_id))
+
+        self.retrieve_filters.clear()
+        return nodes
+
+    def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, refresh: bool = False, **kwargs):
+        if not self.exist_workspace(workspace_id=workspace_id):
+            self.create_workspace(workspace_id=workspace_id)
+
+        if isinstance(nodes, VectorNode):
+            nodes = [nodes]
+
+        embedded_nodes = [node for node in nodes if node.vector]
+        not_embedded_nodes = [node for node in nodes if not node.vector]
+        now_embedded_nodes = self.embedding_model.get_node_embeddings(not_embedded_nodes)
+
+        docs = [
+            {
+                "_op_type": "index",
+                "_index": workspace_id,
+                "_id": node.unique_id,
+                "_source": {
+                    "workspace_id": workspace_id,
+                    "content": node.content,
+                    "metadata": node.metadata,
+                    "vector": node.vector
+                }
+            } for node in embedded_nodes + now_embedded_nodes]
+        status, error = bulk(self._client, docs, chunk_size=self.batch_size, **kwargs)
+        logger.info(f"insert docs.size={len(docs)} status={status} error={error}")
+
+        if refresh:
+            self.refresh(workspace_id=workspace_id)
+
+    def delete(self, node_ids: str | List[str], workspace_id: str, refresh: bool = False, **kwargs):
+        if not self.exist_workspace(workspace_id=workspace_id):
+            logger.warning(f"workspace_id={workspace_id} is not exists!")
+            return
+
+        if isinstance(node_ids, str):
+            node_ids = [node_ids]
+
+        actions = [
+            {
+                "_op_type": "delete",
+                "_index": workspace_id,
+                "_id": node_id
+            } for node_id in node_ids]
+        status, error = bulk(self._client, actions, chunk_size=self.batch_size, **kwargs)
+        logger.info(f"delete actions.size={len(actions)} status={status} error={error}")
+
+        if refresh:
+            self.refresh(workspace_id=workspace_id)
+
+
+def main():
+    from dotenv import load_dotenv
+    load_dotenv()
+
+    embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
+    workspace_id = "rag_nodes_index"
+    hosts = "http://11.160.132.46:8200"
+    es = EsVectorStore(hosts=hosts, embedding_model=embedding_model)
+    if es.exist_workspace(workspace_id=workspace_id):
+        es.delete_workspace(workspace_id=workspace_id)
+    es.create_workspace(workspace_id=workspace_id)
+
+    sample_nodes = [
+        VectorNode(
+            workspace_id=workspace_id,
+            content="Artificial intelligence is a technology that simulates human intelligence.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+        VectorNode(
+            workspace_id=workspace_id,
+            content="AI is the future of mankind.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+        VectorNode(
+            workspace_id=workspace_id,
+            content="I want to eat fish!",
+            metadata={
+                "node_type": "n2",
+            }
+        ),
+        VectorNode(
+            workspace_id=workspace_id,
+            content="The bigger the storm, the more expensive the fish.",
+            metadata={
+                "node_type": "n1",
+            }
+        ),
+    ]
+
+    es.insert(sample_nodes, workspace_id=workspace_id, refresh=True)
+
+    logger.info("=" * 20)
+    results = es.add_term_filter(key="metadata.node_type", value="n1") \
+        .search("What is AI?", top_k=5, workspace_id=workspace_id)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+
+    logger.info("=" * 20)
+    results = es.search("What is AI?", top_k=5, workspace_id=workspace_id)
+    for r in results:
+        logger.info(r.model_dump(exclude={"vector"}))
+    logger.info("=" * 20)
+    es.dump_workspace(workspace_id=workspace_id)
+    es.delete_workspace(workspace_id=workspace_id)
+
+
+if __name__ == "__main__":
+    main()
+    # launch with: python -m llmflow.storage.es_vector_store