backup format

pengmingqiang · Mar 16, 2024 · 7c61c0d · 7c61c0d
1 parent d420cb4
commit 7c61c0d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 16 deletions.
diff --git a/ufo/llm/llm_call.py b/ufo/llm/llm_call.py
@@ -8,15 +8,15 @@
 configs = load_config()
 
 
-def get_completion(messages, is_visual=True, agent: str='APP', backup_engine: bool=True):
+def get_completion(messages, is_visual=True, agent: str='APP', use_backup_engine: bool=True):
     """
     Get completion for the given messages.
 
     Args:
         messages (list): List of messages to be used for completion.
         is_visual (bool, optional): Flag indicating whether the completion is visual or not. 
         agent (str, optional): Type of agent. Possible values are 'APP', 'ACTION' or 'BACKUP'.
-        backup_engine (bool, optional): Flag indicating whether to use the backup engine or not.
+        use_backup_engine (bool, optional): Flag indicating whether to use the backup engine or not.
 
     Returns:
         tuple: A tuple containing the completion response (str) and the cost (float).
@@ -32,16 +32,16 @@ def get_completion(messages, is_visual=True, agent: str='APP', backup_engine: bo
         raise ValueError(f'Agent {agent} not supported')
 
     api_type =  configs[agent_type]['VISUAL' if is_visual else 'NON_VISUAL']['API_TYPE']
-    if api_type.lower() in ['openai', 'aoai', 'azure_ad']:
-        from .openai import OpenAIService
-        try:
+    try:
+        if api_type.lower() in ['openai', 'aoai', 'azure_ad']:
+            from .openai import OpenAIService
             response, cost = OpenAIService(configs, is_visual=is_visual, agent_type=agent_type).chat_completion(messages)
-        except Exception as e:
-            if backup_engine:
-                print_with_color(f"OpenAI API request failed: {e}, try to use the backup engine", "red")
-                return get_completion(messages, is_visual=is_visual, agent='backup', backup_engine=False)
-            else:
-                raise e
-        return response, cost
-    else:
-        raise ValueError(f'API_TYPE {api_type} not supported')
+            return response, cost
+        else:
+            raise ValueError(f'API_TYPE {api_type} not supported')
+    except Exception as e:
+        if use_backup_engine:
+            print_with_color(f"The API request of {agent_type} failed: {e}, try to use the backup engine", "red")
+            return get_completion(messages, is_visual=is_visual, agent='backup', use_backup_engine=False)
+        else:
+            raise e
diff --git a/ufo/module/flow.py b/ufo/module/flow.py
@@ -105,7 +105,7 @@ def process_application_selection(self):
         self.request_logger.debug(json.dumps({"step": self.step, "prompt": app_selection_prompt_message, "status": ""}))
 
         try:
-            response_string, cost = llm_call.get_completion(app_selection_prompt_message, configs["APP_AGENT_VISUAL_MODE"], "APP", backup_engine=True)
+            response_string, cost = llm_call.get_completion(app_selection_prompt_message, configs["APP_AGENT_VISUAL_MODE"], "APP", use_backup_engine=True)
 
         except Exception as e:
             log = json.dumps({"step": self.step, "status": str(e), "prompt": app_selection_prompt_message})
@@ -248,7 +248,7 @@ def process_action_selection(self):
             self.request_logger.debug(json.dumps({"step": self.step, "prompt": action_selection_prompt_message, "status": ""}))
 
             try:
-                response_string, cost = llm_call.get_completion(action_selection_prompt_message, configs["ACTION_AGENT_VISUAL_MODE"], "ACTION", backup_engine=True)
+                response_string, cost = llm_call.get_completion(action_selection_prompt_message, configs["ACTION_AGENT_VISUAL_MODE"], "ACTION", use_backup_engine=True)
             except Exception as e:
                 log = json.dumps({"step": self.step, "status": str(e), "prompt": action_selection_prompt_message})
                 print_with_color("Error occurs when calling LLM: {e}".format(e=str(e)), "red")