modified for new backend

.gitignore

@@ -0,0 +1 @@
+__pycache__/*

ChromaDBFlow.py

@@ -2,21 +2,44 @@ import os
 from typing import Dict, List, Any
 
 import uuid
-
+from copy import deepcopy
 from langchain.embeddings import OpenAIEmbeddings
 
 from chromadb import Client as ChromaClient
 
 from flows.base_flows import AtomicFlow
 
+import hydra
 
 class ChromaDBFlow(AtomicFlow):
 
-    def __init__(self, **kwargs):
+    def __init__(self, backend,**kwargs):
         super().__init__(**kwargs)
         self.client = ChromaClient()
         self.collection = self.client.get_or_create_collection(name=self.flow_config["name"])
+        self.backend = backend
+
+    @classmethod
+    def _set_up_backend(cls, config):
+        kwargs = {}
+
+        kwargs["backend"] = \
+            hydra.utils.instantiate(config['backend'], _convert_="partial")
+        
+        return kwargs
+    
+    @classmethod
+    def instantiate_from_config(cls, config):
+        flow_config = deepcopy(config)
+
+        kwargs = {"flow_config": flow_config}
+
+        # ~~~ Set up backend ~~~
+        kwargs.update(cls._set_up_backend(flow_config))
 
+        # ~~~ Instantiate flow ~~~
+        return cls(**kwargs)
+    
     def get_input_keys(self) -> List[str]:
         return self.flow_config["input_keys"]
 
@@ -25,7 +48,7 @@ class ChromaDBFlow(AtomicFlow):
 
     def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
 
-        api_information = self._get_from_state("api_information")
+        api_information = self.backend.get_key()
 
         if api_information.backend_used == "openai":
             embeddings = OpenAIEmbeddings(openai_api_key=api_information.api_key)

ChromaDBFlow.yaml

@@ -1,6 +1,10 @@
 name: chroma_db
 description: ChromaDB is a document store that uses vector embeddings to store and retrieve documents
 
+backend:
+  _target_: flows.backends.llm_lite.LiteLLMBackend
+  api_infos: ???
+
 input_keys:
   - operation
   - content

VectorStoreFlow.py

@@ -10,20 +10,33 @@ from langchain.vectorstores import Chroma, FAISS
 from langchain.vectorstores.base import VectorStoreRetriever
 
 from flows.base_flows import AtomicFlow
+import hydra
 
 
 class VectorStoreFlow(AtomicFlow):
-    REQUIRED_KEYS_CONFIG = ["type", "api_keys"]
+    REQUIRED_KEYS_CONFIG = ["type"]
 
     vector_db: VectorStoreRetriever
 
-    def __init__(self, vector_db, **kwargs):
+    def __init__(self, backend,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"])
+    def _set_up_backend(cls, config):
+        kwargs = {}
+
+        kwargs["backend"] = \
+            hydra.utils.instantiate(config['backend'], _convert_="partial")
+        
+        return kwargs
+    
+    
+    @classmethod
+    def _set_up_retriever(cls, api_information,config: Dict[str, Any]) -> Dict[str, Any]:
+       
+        embeddings = OpenAIEmbeddings(openai_api_key=api_information.api_key)
         kwargs = {}
 
         vs_type = config["type"]
@@ -50,9 +63,13 @@ class VectorStoreFlow(AtomicFlow):
         flow_config = deepcopy(config)
 
         kwargs = {"flow_config": flow_config}
-
-        kwargs.update(cls._set_up_retriever(flow_config))
-
+       
+        # ~~~ Set up backend ~~~
+        kwargs.update(cls._set_up_backend(flow_config))
+        api_information = kwargs["backend"].get_key()
+        
+        kwargs.update(cls._set_up_retriever(api_information,flow_config))
+        
         return cls(**kwargs)
 
     @staticmethod

VectorStoreFlow.yaml

@@ -1,6 +1,10 @@
 name: "VectorStoreFlow"
 description: "VectorStoreFlow"
 
+backend:
+  _target_: flows.backends.llm_lite.LiteLLMBackend
+  api_infos: ?
+  
 input_keys:
   - "operation" # read or write
   - "content"