merge

ufo/automator/ui_control/controller.py

@@ -91,27 +91,39 @@ def __summary(self, args_dict):
         return args_dict.get("text")
 
 
-
     def __set_edit_text(self, args_dict:dict):
         """
         Set the edit text of the control element.
         :param args: The arguments of the set edit text method.
         :return: The result of the set edit text action.
         """
-        if configs["INPUT_TEXT_API"] == "type_keys":
+
+        if configs["INPUT_TEXT_API"] == "set_text":
+            method_name = "set_text"
+            args = {"text": args_dict["text"]}   
+        else:
             method_name = "type_keys"
             args = {"keys": args_dict["text"], "pause": 0.1, "with_spaces": True}
-        else:
-            args = {"text": args_dict["text"]}
-
         try:
             result = self.atomic_execution(self.control, method_name, args)
-            if configs["INPUT_TEXT_ENTER"] and method_name in ["type_keys", "set_edit_text"]:
+            if method_name == "set_text" and args["text"] not in self.control.window_text():
+                raise Exception(f"Failed to use set_text: {args['text']}")
+            if configs["INPUT_TEXT_ENTER"] and method_name in ["type_keys", "set_text"]:
                 self.atomic_execution(self.control, "type_keys", args = {"keys": "{ENTER}"})
             return result
         except Exception as e:
-            return f"An error occurred: {e}"
-
+            if method_name == "set_text":
+                print_with_color(f"{self.control} doesn't have a method named {method_name}, trying default input method", "yellow")
+                method_name = "type_keys"
+                clear_text_keys = "^a{BACKSPACE}"
+                text_to_type = args["text"]
+                keys_to_send = clear_text_keys + text_to_type
+                method_name = "type_keys"
+                args = {"keys": keys_to_send, "pause": 0.1, "with_spaces": True}
+                return self.atomic_execution(self.control, method_name, args)
+            else:
+                return f"An error occurred: {e}"
+
 
 
     def __texts(self, args_dict:dict) -> str:

ufo/config/config.py

@@ -27,9 +27,13 @@ def load_config(config_path="ufo/config/"):
             configs.update(yaml_data)
         with open(path + "config_dev.yaml", "r") as file:
             yaml_dev_data = yaml.safe_load(file)
+        with open(path + "config_prices.yaml", "r") as file:
+            yaml_prices_data = yaml.safe_load(file)
         # Update configs with YAML data
         if yaml_data:
             configs.update(yaml_dev_data)
+        if yaml_prices_data:
+            configs.update(yaml_prices_data)
     except FileNotFoundError:
         print_with_color(
             f"Warning: Config file not found at {config_path}. Using only environment variables.", "yellow")

ufo/config/config_prices.yaml

@@ -0,0 +1,33 @@
+# Source: https://openai.com/pricing
+# Prices in $ per 1000 tokens
+# Last updated: 2024-03-27
+PRICES: { 
+    "openai/gpt-4-0613": {"input": 0.03, "output": 0.06},
+    "openai/gpt-3.5-turbo-0613": {"input": 0.0015, "output": 0.002},
+    "openai/gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "openai/gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "openai/gpt-4-1106-vision-preview": {"input": 0.01, "output": 0.03},
+    "openai/gpt-4": {"input": 0.03, "output": 0.06},
+    "openai/gpt-4-32k": {"input": 0.06, "output": 0.12},
+    "openai/gpt-3.5-turbo-0125": {"input": 0.0005, "output": 0.0015},
+    "openai/gpt-3.5-turbo-1106": {"input": 0.001, "output": 0.002},
+    "openai/gpt-3.5-turbo-instruct": {"input": 0.0015, "output": 0.002},
+    "openai/gpt-3.5-turbo-16k-0613": {"input": 0.003, "output": 0.004},
+    "openai/whisper-1": {"input": 0.006, "output": 0.006},
+    "openai/tts-1": {"input": 0.015, "output": 0.015},
+    "openai/tts-hd-1": {"input": 0.03, "output": 0.03},
+    "openai/text-embedding-ada-002-v2": {"input": 0.0001, "output": 0.0001},
+    "openai/text-davinci:003": {"input": 0.02, "output": 0.02},
+    "openai/text-ada-001": {"input": 0.0004, "output": 0.0004},
+    "azure/gpt-35-turbo-20220309":{"input": 0.0015, "output": 0.002},
+    "azure/gpt-35-turbo-20230613":{"input": 0.0015, "output": 0.002},
+    "azure/gpt-35-turbo-16k-20230613":{"input": 0.003, "output": 0.004},
+    "azure/gpt-35-turbo-1106":{"input": 0.001, "output": 0.002},
+    "azure/gpt-4-20230321":{"input": 0.03, "output": 0.06},
+    "azure/gpt-4-32k-20230321":{"input": 0.06, "output": 0.12},
+    "azure/gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "azure/gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "azure/gpt-4-visual-preview": {"input": 0.01, "output": 0.03}
+}
+
+

ufo/experience/summarizer.py

@@ -80,7 +80,7 @@ def get_summary_list(self, logs: list) -> Tuple[list, float]:
         return: The summary list and the total cost.
         """
         summaries = []
-        total_cost = 0
+        total_cost = 0.0
         for log_partition in logs:
             prompt = self.build_prompt(log_partition)
             summary, cost = self.get_summary(prompt)

ufo/llm/openai.py

@@ -1,34 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 import datetime
 from typing import Any, Optional
 import openai
 from openai import AzureOpenAI, OpenAI
 
 
-
 class OpenAIService:
     def __init__(self, config, agent_type: str):
         self.config_llm = config[agent_type]
         self.config = config
-        api_type = self.config_llm["API_TYPE"].lower()
-        max_retry = self.config["MAX_RETRY"]
-        assert api_type in ["openai", "aoai", "azure_ad"], "Invalid API type"
+        self.api_type = self.config_llm["API_TYPE"].lower()
+        self.max_retry = self.config["MAX_RETRY"]
+        self.prices = self.config["PRICES"]
+        assert self.api_type in ["openai", "aoai", "azure_ad"], "Invalid API type"
         self.client: OpenAI = (
             OpenAI(
                 base_url=self.config_llm["API_BASE"],
                 api_key=self.config_llm["API_KEY"],
-                max_retries=max_retry,
+                max_retries=self.max_retry,
                 timeout=self.config["TIMEOUT"],
             )
-            if api_type == "openai"
+            if self.api_type == "openai"
             else AzureOpenAI(
-                max_retries=max_retry,
+                max_retries=self.max_retry,
                 timeout=self.config["TIMEOUT"],
                 api_version=self.config_llm["API_VERSION"],
                 azure_endpoint=self.config_llm["API_BASE"],
-                api_key=(self.config_llm["API_KEY"] if api_type == 'aoai' else self.get_openai_token()),
+                api_key=(self.config_llm["API_KEY"] if self.api_type == 'aoai' else self.get_openai_token()),
             )
         )
-        if api_type == "azure_ad":
+        if self.api_type == "azure_ad":
             self.auto_refresh_token()
 
     def chat_completion(
@@ -64,7 +67,7 @@ def chat_completion(
             prompt_tokens = usage.prompt_tokens
             completion_tokens = usage.completion_tokens
 
-            cost = prompt_tokens / 1000 * 0.01 + completion_tokens / 1000 * 0.03
+            cost = self.get_cost_estimator(self.api_type, model, self.prices, prompt_tokens, completion_tokens)
 
             return [response.choices[i].message.content for i in range(n)], cost
 
@@ -91,9 +94,6 @@ def chat_completion(
             # Handle API error, e.g. retry or log
             raise Exception(f"OpenAI API returned an API Error: {e}")
 
-
-
-
     def get_openai_token(
         self,
         token_cache_file: str = 'apim-token-cache.bin',
@@ -200,7 +200,6 @@ def save_cache():
                 raise Exception(
                     "Authentication failed for acquiring AAD token for your organization")
 
-
     def auto_refresh_token(
         self,
         token_cache_file: str = 'apim-token-cache.bin',
@@ -267,3 +266,26 @@ def stop():
 
         return stop
 
+    def get_cost_estimator(self, api_type, model, prices, prompt_tokens, completion_tokens) -> float:
+        """
+        Calculates the cost estimate for using a specific model based on the number of prompt tokens and completion tokens.
+
+        Args:
+            model (str): The name of the model.
+            prices (dict): A dictionary containing the prices for different models.
+            prompt_tokens (int): The number of prompt tokens used.
+            completion_tokens (int): The number of completion tokens used.
+
+        Returns:
+            float: The estimated cost for using the model.
+        """
+        if api_type.lower() == "openai":
+            name = str(api_type+'/'+model)
+        else:
+            name = str('azure/'+model)
+        if name in prices:
+            cost = prompt_tokens * prices[name]["input"]/1000 + completion_tokens * prices[name]["output"]/1000
+        else:
+            print(f"{name} not found in prices")
+            return None
+        return cost

ufo/module/flow.py

@@ -53,7 +53,7 @@ def __init__(self, task):
         self.plan = ""
         self.request = ""
 
-        self._cost = 0
+        self._cost = 0.0
         self.control_reannotate = None
 
         welcome_text = """
@@ -99,8 +99,8 @@ def process_application_selection(self):
             self.request_logger.info(log)
             self._status = "ERROR"
             return
-
-        self._cost += cost
+        
+        self.update_cost(cost=cost)
 
         try:
             response_json = self.HostAgent.response_to_dict(response_string)
@@ -242,8 +242,8 @@ def process_action_selection(self):
             self._status = "ERROR"
             time.sleep(configs["SLEEP_TIME"])
             return 
-
-        self._cost += cost
+            
+        self.update_cost(cost=cost)
 
         try:
             response_json = self.AppAgent.response_to_dict(response_string)
@@ -384,8 +384,8 @@ def experience_saver(self) -> None:
         utils.create_folder(experience_path)
         summarizer.create_or_update_yaml(summaries, os.path.join(experience_path, "experience.yaml"))
         summarizer.create_or_update_vector_db(summaries, os.path.join(experience_path, "experience_db"))
-
-        self._cost += total_cost
+        
+        self.update_cost(cost=total_cost)
         utils.print_with_color("The experience has been saved.", "cyan")
 
     def set_new_round(self) -> None:
@@ -451,14 +451,14 @@ def get_results(self) -> list:
         Get the results of the session.
         return: The results of the session.
         """
-
-        if len(self.action_history) > 0:
-            result = self.action_history[-1].get("Results")
-        else:
-            result
-        return result
+        return self.results
 
-
+    def get_cost(self):
+        """
+        Get the cost of the session.
+        return: The cost of the session.
+        """
+        return self.cost
 
     def get_application_window(self) -> object:
         """
@@ -511,6 +511,16 @@ def error_logger(self, response_str: str, error: str) -> None:
         self.logger.info(log)
 
 
+    def update_cost(self, cost):
+        """
+        Update the cost of the session.
+        """
+        if isinstance(cost, float) and isinstance(self.cost, float):
+            self._cost += cost
+        else:
+            self._cost = None
+
+
     @staticmethod
     def initialize_logger(log_path: str, log_filename: str) -> logging.Logger:
         """

ufo/ufo.py

@@ -82,8 +82,9 @@ def main():
 
     # Print the total cost
     total_cost = session.get_cost()
-    formatted_cost = '${:.2f}'.format(total_cost)
-    print_with_color(f"Request total cost is {formatted_cost}", "yellow")
+    if isinstance(total_cost, float):
+        formatted_cost = '${:.2f}'.format(total_cost)
+        print_with_color(f"Request total cost is {formatted_cost}", "yellow")
 
     return status