new backend

OpenAIChatAtomicFlow.py

@@ -1,40 +1,42 @@
 from copy import deepcopy
-
 import hydra
 
 import time
 
 from typing import Dict, Optional, Any
 
-from langchain import PromptTemplate
-from langchain.schema import HumanMessage, AIMessage, SystemMessage
-
 from flows.base_flows import AtomicFlow
 from flows.datasets import GenericDemonstrationsDataset
 
 from flows.utils import logging
 from flows.messages.flow_message import UpdateMessage_ChatMessage
 
+from flows.prompt_template import JinjaPrompt
+
+from backends.llm_lite import LiteLLMBackend
+
 log = logging.get_logger(__name__)
 
 
 class OpenAIChatAtomicFlow(AtomicFlow):
-    REQUIRED_KEYS_CONFIG = ["model_name", "generation_parameters"]
+    REQUIRED_KEYS_CONFIG = ["backend"]
 
     SUPPORTS_CACHING: bool = True
 
-    system_message_prompt_template: PromptTemplate
-    human_message_prompt_template: PromptTemplate
+    system_message_prompt_template: JinjaPrompt
+    human_message_prompt_template: JinjaPrompt
 
-    init_human_message_prompt_template: Optional[PromptTemplate] = None
+    backend: LiteLLMBackend
+    init_human_message_prompt_template: Optional[JinjaPrompt] = None
     demonstrations: GenericDemonstrationsDataset = None
     demonstrations_k: Optional[int] = None
-    demonstrations_response_prompt_template: PromptTemplate = None
+    demonstrations_response_prompt_template: str = None
 
     def __init__(self,
                  system_message_prompt_template,
                  human_message_prompt_template,
                  init_human_message_prompt_template,
+                 backend,
                  demonstrations_response_prompt_template=None,
                  demonstrations=None,
                  **kwargs):
@@ -45,7 +47,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         self.demonstrations_response_prompt_template = demonstrations_response_prompt_template
         self.demonstrations = demonstrations
         self.demonstrations_k = self.flow_config.get("demonstrations_k", None)
-
+        self.backend = backend
         assert self.flow_config["name"] not in [
             "system",
             "user",
@@ -59,20 +61,29 @@ class OpenAIChatAtomicFlow(AtomicFlow):
     @classmethod
     def _set_up_prompts(cls, config):
         kwargs = {}
-
+        
         kwargs["system_message_prompt_template"] = \
             hydra.utils.instantiate(config['system_message_prompt_template'], _convert_="partial")
         kwargs["init_human_message_prompt_template"] = \
             hydra.utils.instantiate(config['init_human_message_prompt_template'], _convert_="partial")
         kwargs["human_message_prompt_template"] = \
             hydra.utils.instantiate(config['human_message_prompt_template'], _convert_="partial")
-
+        
         if "demonstrations_response_prompt_template" in config:
             kwargs["demonstrations_response_prompt_template"] = \
                 hydra.utils.instantiate(config['demonstrations_response_prompt_template'], _convert_="partial")
             kwargs["demonstrations"] = GenericDemonstrationsDataset(**config['demonstrations'])
 
         return kwargs
+    
+    @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):
@@ -82,6 +93,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
 
         # ~~~ Set up prompts ~~~
         kwargs.update(cls._set_up_prompts(flow_config))
+        kwargs.update(cls._set_up_backend(flow_config))
 
         # ~~~ Instantiate flow ~~~
         return cls(**kwargs)
@@ -106,7 +118,6 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         template_kwargs = {}
         for input_variable in prompt_template.input_variables:
             template_kwargs[input_variable] = input_data[input_variable]
-
         msg_content = prompt_template.format(**template_kwargs)
         return msg_content
 
@@ -140,19 +151,16 @@ class OpenAIChatAtomicFlow(AtomicFlow):
                                        role: str,
                                        content: str) -> None:
 
-        # Add the message to the previous messages list
-        if role == self.flow_config["system_name"]:
-            self.flow_state["previous_messages"].append(SystemMessage(content=content))
-        elif role == self.flow_config["user_name"]:
-            self.flow_state["previous_messages"].append(HumanMessage(content=content))
-        elif role == self.flow_config["assistant_name"]:
-            self.flow_state["previous_messages"].append(AIMessage(content=content))
+
+        acceptable_roles = [self.flow_config["system_name"],self.flow_config["user_name"],self.flow_config["assistant_name"]]
+        if role in acceptable_roles:
+            self.flow_state["previous_messages"].append({"role": role , "content": content})
+            
         else:
             raise Exception(f"Invalid role: `{role}`.\n"
                             f"Role should be one of: "
-                            f"`{self.flow_config['system_name']}`, "
-                            f"`{self.flow_config['user_name']}`, "
-                            f"`{self.flow_config['assistant_name']}`")
+                            f"`{acceptable_roles}`, ")
+      
 
         # Log the update to the flow messages list
         chat_message = UpdateMessage_ChatMessage(
@@ -174,49 +182,24 @@ class OpenAIChatAtomicFlow(AtomicFlow):
             return all_messages[:first_k] + all_messages[-last_k:]
         elif first_k:
             return all_messages[:first_k]
-
         return all_messages[-last_k:]
 
     def _call(self):
-        api_information = self._get_from_state("api_information")
-        api_key = api_information.api_key
-
-        if api_information.backend_used == 'azure':
-            from backends.azure_openai import SafeAzureChatOpenAI
-            endpoint = api_information.endpoint
-            backend = SafeAzureChatOpenAI(
-                openai_api_type='azure',
-                openai_api_key=api_key,
-                openai_api_base=endpoint,
-                openai_api_version='2023-05-15',
-                deployment_name=self.flow_config["model_name"],
-                **self.flow_config["generation_parameters"],
-            )
-        elif api_information.backend_used == 'openai':
-            from backends.openai import SafeChatOpenAI
-            backend = SafeChatOpenAI(
-                model_name=self.flow_config["model_name"],
-                openai_api_key=api_key,
-                openai_api_type="open_ai",
-                **self.flow_config["generation_parameters"],
-            )
-        else:
-            raise ValueError(f"Unsupported backend: {api_information.backend_used}")
-
+         
         messages = self._get_previous_messages()
-
         _success = False
         attempts = 1
         error = None
         response = None
         while attempts <= self.flow_config['n_api_retries']:
             try:
-                response = backend(messages).content
+                response = self.backend(messages=messages,mock_response=False) #set mock_response to True when debugging (fake API request)
+                response = [ answer["content"] for answer in response] # because n in the generation parameters can be > 1
                 _success = True
                 break
             except Exception as e:
                 log.error(
-                    f"Error {attempts} in calling backend: {e}. Key used: `{api_key}`. "
+                    f"Error {attempts} in calling backend: {e}. "
                     f"Retrying in {self.flow_config['wait_time_between_retries']} seconds..."
                 )
                 # log.error(
@@ -226,7 +209,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
                 attempts += 1
                 time.sleep(self.flow_config['wait_time_between_retries'])
                 error = e
-
+        
         if not _success:
             raise error
 
@@ -266,9 +249,12 @@ class OpenAIChatAtomicFlow(AtomicFlow):
 
         # ~~~ Call ~~~
         response = self._call()
-        self._state_update_add_chat_message(
-            role=self.flow_config["assistant_name"],
-            content=response
-        )
+        
+        #loop is in case there was more than one answer (n>1 in generation parameters)
+        for answer in response:
+            self._state_update_add_chat_message(
+                role=self.flow_config["assistant_name"],
+                content=answer
+            )
 
         return {"api_output": response}

OpenAIChatAtomicFlow.yaml

@@ -1,17 +1,6 @@
 # This is an abstract flow, therefore some required fields are not defined (and must be defined by the concrete flow)
 enable_cache: True
 
-model_name: "gpt-4"
-generation_parameters:
-  n: 1
-  max_tokens: 2000
-  temperature: 0.3
-
-  model_kwargs:
-    top_p: 0.2
-    frequency_penalty: 0
-    presence_penalty: 0
-
 n_api_retries: 6
 wait_time_between_retries: 20
 
@@ -19,26 +8,39 @@ system_name: system
 user_name: user
 assistant_name: assistant
 
+backend:
+  _target_: backends.llm_lite.LiteLLMBackend
+  api_infos: ???
+  model_name: "gpt-3.5-turbo"
+  n: 1
+  max_tokens: 2000
+  temperature: 0.3
+
+
+  top_p: 0.2
+  frequency_penalty: 0
+  presence_penalty: 0
+  stream: True
+
+
 system_message_prompt_template:
-  _target_: langchain.PromptTemplate
-  template_format: jinja2
+  _target_: flows.prompt_template.JinjaPrompt
+
 
 init_human_message_prompt_template:
-  _target_: langchain.PromptTemplate
-  template_format: jinja2
+  _target_: flows.prompt_template.JinjaPrompt
 
 human_message_prompt_template:
-  _target_: langchain.PromptTemplate
+  _target_: flows.prompt_template.JinjaPrompt
   template: "{{query}}"
   input_variables:
     - "query"
-  template_format: jinja2
 input_interface_initialized:
   - "query"
 
 query_message_prompt_template:
-  _target_: langchain.PromptTemplate
-  template_format: jinja2
+  _target_: flows.prompt_template.JinjaPrompt
+
 
 previous_messages:
   first_k: null  # Note that the first message is the system prompt

run.py

@@ -3,7 +3,8 @@ import os
 import hydra
 
 import flows
-from flows.flow_launchers import FlowLauncher, ApiInfo
+from flows.flow_launchers import FlowLauncher
+from backends.api_info import ApiInfo
 from flows.utils.general_helpers import read_yaml_file
 
 from flows import logging
@@ -23,25 +24,33 @@ flow_verse.sync_dependencies(dependencies)
 if __name__ == "__main__":
     # ~~~ Set the API information ~~~
     # OpenAI backend
-    # api_information = ApiInfo("openai", os.getenv("OPENAI_API_KEY"))
-    # Azure backend
-    api_information = ApiInfo("azure", os.getenv("AZURE_OPENAI_KEY"), os.getenv("AZURE_OPENAI_ENDPOINT"))
+
+    api_information = [ApiInfo(backend_used="openai",
+                              api_key = os.getenv("OPENAI_API_KEY"))]
+    
+    # # Azure backend
+    # api_information = ApiInfo(backend_used = "azure",
+    #                           api_base = os.getenv("AZURE_API_BASE"),
+    #                           api_key = os.getenv("AZURE_OPENAI_KEY"),
+    #                           api_version =  os.getenv("AZURE_API_VERSION") )
 
     root_dir = "."
     cfg_path = os.path.join(root_dir, "SimpleQA.yaml")
     cfg = read_yaml_file(cfg_path)
 
+    cfg["flow"]["backend"]["api_infos"] = api_information
+    # ~~~ Instantiate the Flow ~~~
     # ~~~ Instantiate the Flow ~~~
     flow_with_interfaces = {
         "flow": hydra.utils.instantiate(cfg['flow'], _recursive_=False, _convert_="partial"),
         "input_interface": (
             None
-            if getattr(cfg, "input_interface", None) is None
+            if cfg.get( "input_interface", None) is None
             else hydra.utils.instantiate(cfg['input_interface'], _recursive_=False)
         ),
         "output_interface": (
             None
-            if getattr(cfg, "output_interface", None) is None
+            if cfg.get( "output_interface", None) is None
             else hydra.utils.instantiate(cfg['output_interface'], _recursive_=False)
         ),
     }
@@ -58,7 +67,6 @@ if __name__ == "__main__":
         flow_with_interfaces=flow_with_interfaces,
         data=data,
         path_to_output_file=path_to_output_file,
-        api_information=api_information,
     )
 
     # ~~~ Print the output ~~~