First commit.

ChromaDBFlow.py

@@ -0,0 +1,68 @@
+import os
+from typing import Dict, List, Any
+
+import uuid
+
+from langchain.embeddings import OpenAIEmbeddings
+
+from chromadb import Client as ChromaClient
+
+from flows.base_flows import AtomicFlow
+
+
+class ChromaDBFlow(AtomicFlow):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.client = ChromaClient()
+        self.collection = self.client.get_or_create_collection(name=self.flow_config["name"])
+
+    def get_input_keys(self) -> List[str]:
+        return self.flow_config["input_keys"]
+
+    def get_output_keys(self) -> List[str]:
+        return self.flow_config["output_keys"]
+
+    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+
+        api_information = self._get_from_state("api_information")
+
+        if api_information.backend_used == "openai":
+            embeddings = OpenAIEmbeddings(openai_api_key=api_information.api_key)
+        else:
+            # ToDo: Add support for Azure
+            embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("OPENAI_API_KEY"))
+        response = {}
+
+        operation = input_data["operation"]
+        if operation not in ["write", "read"]:
+            raise ValueError(f"Operation '{operation}' not supported")
+
+        content = input_data["content"]
+        if operation == "read":
+            if not isinstance(content, str):
+                raise ValueError(f"content(query) must be a string during read, got {type(content)}: {content}")
+            if content == "":
+                response["retrieved"] = [[""]]
+                return response
+            query = content
+            query_result = self.collection.query(
+                query_embeddings=embeddings.embed_query(query),
+                n_results=self.flow_config["n_results"]
+            )
+
+            response["retrieved"] = [doc for doc in query_result["documents"]]
+
+        elif operation == "write":
+            if content != "":
+                if not isinstance(content, list):
+                    content = [content]
+                documents = content
+                self.collection.add(
+                    ids=[str(uuid.uuid4()) for _ in range(len(documents))],
+                    embeddings=embeddings.embed_documents(documents),
+                    documents=documents
+                )
+            response["retrieved"] = ""
+
+        return response

ChromaDBFlow.yaml

@@ -0,0 +1,10 @@
+name: chroma_db
+description: ChromaDB is a document store that uses vector embeddings to store and retrieve documents
+
+input_keys:
+  - operation
+  - content
+output_keys:
+  - retrieved
+
+n_results: 5 # number of results to retrieve when query

README.md

@@ -0,0 +1,25 @@
+---
+license: mit
+---
+## Description
+ToDo
+
+< Flow description >
+
+## Configuration parameters
+
+< Name 1 > (< Type 1 >): < Description 1 >. Required parameter.
+
+< Name 2 > (< Type 2 >): < Description 2 >. Default value is: < value 2 >
+
+## Input interface
+
+< Name 1 > (< Type 1 >): < Description 1 >. 
+
+(Note that the interface might depend on the state of the Flow.)
+
+## Output interface
+
+< Name 1 > (< Type 1 >): < Description 1 >. 
+
+(Note that the interface might depend on the state of the Flow.)

VectorStoreFlow.py

@@ -0,0 +1,84 @@
+from copy import deepcopy
+from typing import Dict, List, Any, Optional
+
+import faiss
+
+from langchain.docstore import InMemoryDocstore
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.schema import Document
+from langchain.vectorstores import Chroma, FAISS
+from langchain.vectorstores.base import VectorStoreRetriever
+
+from flows.base_flows import AtomicFlow
+
+
+class VectorStoreFlow(AtomicFlow):
+    REQUIRED_KEYS_CONFIG = ["type", "api_keys"]
+
+    vector_db: VectorStoreRetriever
+
+    def __init__(self, vector_db, **kwargs):
+        super().__init__(**kwargs)
+        self.vector_db = vector_db
+
+    @classmethod
+    def _set_up_retriever(cls, config: Dict[str, Any]) -> Dict[str, Any]:
+        embeddings = OpenAIEmbeddings(openai_api_key=config["api_keys"]["openai"])
+        kwargs = {}
+
+        vs_type = config["type"]
+
+        if vs_type == "chroma":
+            vectorstore = Chroma(config["name"], embedding_function=embeddings)
+        elif vs_type == "faiss":
+            index = faiss.IndexFlatL2(config.get("embedding_size", 1536))
+            vectorstore = FAISS(
+                embedding_function=embeddings.embed_query,
+                index=index,
+                docstore=InMemoryDocstore({}),
+                index_to_docstore_id={}
+            )
+        else:
+            raise NotImplementedError(f"Vector store '{vs_type}' not implemented")
+
+        kwargs["vector_db"] = vectorstore.as_retriever(**config.get("retriever_config", {}))
+
+        return kwargs
+
+    @classmethod
+    def instantiate_from_config(cls, config: Dict[str, Any]):
+        flow_config = deepcopy(config)
+
+        kwargs = {"flow_config": flow_config}
+
+        kwargs.update(cls._set_up_retriever(flow_config))
+
+        return cls(**kwargs)
+
+    @staticmethod
+    def package_documents(documents: List[str]) -> List[Document]:
+        # TODO(yeeef): support metadata
+        return [Document(page_content=doc, metadata={"": ""}) for doc in documents]
+
+    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        response = {}
+
+        operation = input_data["operation"]
+        assert operation in ["write", "read"], f"Operation '{operation}' not supported"
+
+        content = input_data["content"]
+        if operation == "read":
+            assert isinstance(content, str), f"Content must be a string, got {type(content)}"
+            query = content
+            retrieved_documents = self.vector_db.get_relevant_documents(query)
+            response["retrieved"] = [doc.page_content for doc in retrieved_documents]
+        elif operation == "write":
+            if isinstance(content, str):
+                content = [content]
+            assert isinstance(content, list), f"Content must be a list of strings, got {type(content)}"
+            documents = content
+            documents = self.package_documents(documents)
+            self.vector_db.add_documents(documents)
+            response["retrieved"] = ""
+
+        return response

VectorStoreFlow.yaml

@@ -0,0 +1,15 @@
+name: "VectorStoreFlow"
+description: "VectorStoreFlow"
+
+input_keys:
+  - "operation" # read or write
+  - "content"
+
+output_keys:
+  - "retrieved"
+
+type: "chroma"
+api_keys:
+  openai: "YOUR_OPENAI_API_KEY"
+
+

__init__.py

@@ -0,0 +1,2 @@
+from .VectorStoreFlow import VectorStoreFlow
+from .ChromaDBFlow import ChromaDBFlow

pip_requirements.py

@@ -0,0 +1 @@
+# ToDo