From e0c4fb811fbe703f91b30da014d277b23bbbed05 Mon Sep 17 00:00:00 2001
From: Lenny Liu <hengyu2@illinois.edu>
Date: Sun, 24 Mar 2024 23:42:31 +0800
Subject: [PATCH 1/9] add set_text function

---
 ufo/ui_control/executor.py | 35 ++++++++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/ufo/ui_control/executor.py b/ufo/ui_control/executor.py
index 0ecad975..ab2b0577 100644
--- a/ufo/ui_control/executor.py
+++ b/ufo/ui_control/executor.py
@@ -90,7 +90,26 @@ def __summary(self, args_dict):
         return args_dict.get("text")
     
 
-
+    def set_text(self, method_name:str, args:dict):
+            try: 
+                result = self.atomic_execution(self.control, method_name, args)
+                if configs["INPUT_TEXT_ENTER"]:
+                    self.atomic_execution(self.control, "type_keys", args = {"keys": "{ENTER}"})
+                # the following "if" is tentative, could be eliminated as it will cause some wierd behavior. 
+                # But is a good way to check if set_text worked. 
+                if args["text"] not in self.control.window_text():
+                    raise Exception(f"Failed to use set_text: {args['text']}")
+                return result
+            except:
+                print_with_color(f"{self.control} doesn't have a method named {method_name}, trying another 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)
+    
     def __set_edit_text(self, args_dict:dict):
         """
         Set the edit text of the control element.
@@ -100,17 +119,23 @@ def __set_edit_text(self, args_dict:dict):
         if configs["INPUT_TEXT_API"] == "type_keys":
             method_name = "type_keys"
             args = {"keys": args_dict["text"], "pause": 0.1, "with_spaces": True}
+        elif configs["INPUT_TEXT_API"] == "set_text":
+            method_name = "set_text"
+            args = {"text": args_dict["text"]}   
         else:
-            args = {"text": args_dict["text"]}
+            raise ValueError(f"INPUT_TEXT_API {configs['INPUT_TEXT_API']} not supported")
 
         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 == "type_keys":
+                result = self.atomic_execution(self.control, method_name, args)
+            else:
+                result = self.set_text(method_name, args)
+            if configs["INPUT_TEXT_ENTER"] and method_name in ["type_keys", "set_edit_text", "set_text"]:
                 self.atomic_execution(self.control, "type_keys", args = {"keys": "{ENTER}"})
             return result
         except Exception as e:
             return f"An error occurred: {e}"
-    
+ 
 
 
     def __texts(self, args_dict:dict):

From 8ad4f0f21e7bb27d6fef5c2d5049f019674a7fb6 Mon Sep 17 00:00:00 2001
From: Lenny Liu <hengyu2@illinois.edu>
Date: Sun, 24 Mar 2024 23:51:06 +0800
Subject: [PATCH 2/9] add comment

---
 ufo/ui_control/executor.py | 45 ++++++++++++++++++++++----------------
 1 file changed, 26 insertions(+), 19 deletions(-)

diff --git a/ufo/ui_control/executor.py b/ufo/ui_control/executor.py
index ab2b0577..f617b03e 100644
--- a/ufo/ui_control/executor.py
+++ b/ufo/ui_control/executor.py
@@ -91,25 +91,32 @@ def __summary(self, args_dict):
     
 
     def set_text(self, method_name:str, args:dict):
-            try: 
-                result = self.atomic_execution(self.control, method_name, args)
-                if configs["INPUT_TEXT_ENTER"]:
-                    self.atomic_execution(self.control, "type_keys", args = {"keys": "{ENTER}"})
-                # the following "if" is tentative, could be eliminated as it will cause some wierd behavior. 
-                # But is a good way to check if set_text worked. 
-                if args["text"] not in self.control.window_text():
-                    raise Exception(f"Failed to use set_text: {args['text']}")
-                return result
-            except:
-                print_with_color(f"{self.control} doesn't have a method named {method_name}, trying another 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)
-    
+        """
+        Use set_text method to take the input action, 
+        will automatically switch to type_keys if set_text fails, 
+        and make sure the text is correctly set by clearing the exist values.
+        :param method_names: set edit text method.
+        :param args: The arguments of the set edit text method.
+        :return: The result of the set edit text action."""
+        try: 
+            result = self.atomic_execution(self.control, method_name, args)
+            if configs["INPUT_TEXT_ENTER"]:
+                self.atomic_execution(self.control, "type_keys", args = {"keys": "{ENTER}"})
+            # the following "if" is tentative, could be eliminated as it will cause some wierd behavior. 
+            # But is a good way to check if set_text worked. 
+            if args["text"] not in self.control.window_text():
+                raise Exception(f"Failed to use set_text: {args['text']}")
+            return result
+        except:
+            print_with_color(f"{self.control} doesn't have a method named {method_name}, trying another 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)
+
     def __set_edit_text(self, args_dict:dict):
         """
         Set the edit text of the control element.

From 95d730b581c23cba0547181a2c7974bb3d5c2124 Mon Sep 17 00:00:00 2001
From: Lenny Liu <hengyu2@illinois.edu>
Date: Mon, 25 Mar 2024 01:30:54 +0800
Subject: [PATCH 3/9] modified according to comments, merged 2 functions

---
 ufo/ui_control/executor.py | 58 +++++++++++++-------------------------
 1 file changed, 19 insertions(+), 39 deletions(-)

diff --git a/ufo/ui_control/executor.py b/ufo/ui_control/executor.py
index f617b03e..304097c9 100644
--- a/ufo/ui_control/executor.py
+++ b/ufo/ui_control/executor.py
@@ -90,58 +90,38 @@ def __summary(self, args_dict):
         return args_dict.get("text")
     
 
-    def set_text(self, method_name:str, args:dict):
-        """
-        Use set_text method to take the input action, 
-        will automatically switch to type_keys if set_text fails, 
-        and make sure the text is correctly set by clearing the exist values.
-        :param method_names: set edit text method.
-        :param args: The arguments of the set edit text method.
-        :return: The result of the set edit text action."""
-        try: 
-            result = self.atomic_execution(self.control, method_name, args)
-            if configs["INPUT_TEXT_ENTER"]:
-                self.atomic_execution(self.control, "type_keys", args = {"keys": "{ENTER}"})
-            # the following "if" is tentative, could be eliminated as it will cause some wierd behavior. 
-            # But is a good way to check if set_text worked. 
-            if args["text"] not in self.control.window_text():
-                raise Exception(f"Failed to use set_text: {args['text']}")
-            return result
-        except:
-            print_with_color(f"{self.control} doesn't have a method named {method_name}, trying another 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)
-
     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":
-            method_name = "type_keys"
-            args = {"keys": args_dict["text"], "pause": 0.1, "with_spaces": True}
-        elif configs["INPUT_TEXT_API"] == "set_text":
+        
+        if configs["INPUT_TEXT_API"] == "set_text":
             method_name = "set_text"
             args = {"text": args_dict["text"]}   
         else:
-            raise ValueError(f"INPUT_TEXT_API {configs['INPUT_TEXT_API']} not supported")
-
+            method_name = "type_keys"
+            args = {"keys": args_dict["text"], "pause": 0.1, "with_spaces": True}
         try:
-            if method_name == "type_keys":
-                result = self.atomic_execution(self.control, method_name, args)
-            else:
-                result = self.set_text(method_name, args)
-            if configs["INPUT_TEXT_ENTER"] and method_name in ["type_keys", "set_edit_text", "set_text"]:
+            result = self.atomic_execution(self.control, method_name, args)
+            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}"
  
 
 

From ebad32a9887f0d32d795702f9c6897ab9dfa8847 Mon Sep 17 00:00:00 2001
From: Mac0q <gmacq0209@gmail.com>
Date: Tue, 26 Mar 2024 17:02:59 +0800
Subject: [PATCH 4/9] price

---
 ufo/llm/openai.py  |  4 +++-
 ufo/llm/util.py    | 48 ++++++++++++++++++++++++++++++++++++++++++++++
 ufo/module/flow.py | 22 +++++++++++++++------
 3 files changed, 67 insertions(+), 7 deletions(-)
 create mode 100644 ufo/llm/util.py

diff --git a/ufo/llm/openai.py b/ufo/llm/openai.py
index 92687ad1..b15d9be8 100644
--- a/ufo/llm/openai.py
+++ b/ufo/llm/openai.py
@@ -3,6 +3,8 @@
 import openai
 from openai import AzureOpenAI, OpenAI
 
+from ufo.llm.util import get_cost_estimator
+
 
 
 class OpenAIService:
@@ -61,7 +63,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 = get_cost_estimator(self.config_llm["API_TYPE"], model, prompt_tokens, completion_tokens)
 
             return response.choices[0].message.content, cost
 
diff --git a/ufo/llm/util.py b/ufo/llm/util.py
new file mode 100644
index 00000000..de451ce3
--- /dev/null
+++ b/ufo/llm/util.py
@@ -0,0 +1,48 @@
+# Source: https://openai.com/pricing
+# Prices in $ per 1000 tokens
+# Last updated: 2024-01-26
+OPENAI_PRICES = {
+    "gpt-4-0613": {"input": 0.03, "output": 0.06},
+    "gpt-3.5-turbo-0613": {"input": 0.0015, "output": 0.002},
+    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-1106-vision-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4": {"input": 0.03, "output": 0.06},
+    "gpt-4-32k": {"input": 0.06, "output": 0.12},
+    "gpt-3.5-turbo-0125": {"input": 0.0005, "output": 0.0015},
+    "gpt-3.5-turbo-1106": {"input": 0.001, "output": 0.002},
+    "gpt-3.5-turbo-instruct": {"input": 0.0015, "output": 0.002},
+    "gpt-3.5-turbo-16k-0613": {"input": 0.003, "output": 0.004},
+    "whisper-1": {"input": 0.006, "output": 0.006},
+    "tts-1": {"input": 0.015, "output": 0.015},
+    "tts-hd-1": {"input": 0.03, "output": 0.03},
+    "text-embedding-ada-002-v2": {"input": 0.0001, "output": 0.0001},
+    "text-davinci:003": {"input": 0.02, "output": 0.02},
+    "text-ada-001": {"input": 0.0004, "output": 0.0004},
+    }
+
+AZURE_PRICES = {
+    "gpt-35-turbo-20220309":{"input": 0.0015, "output": 0.002},
+    "gpt-35-turbo-20230613":{"input": 0.0015, "output": 0.002},
+    "gpt-35-turbo-16k-20230613":{"input": 0.003, "output": 0.004},
+    "gpt-35-turbo-1106":{"input": 0.001, "output": 0.002},
+    "gpt-4-20230321":{"input": 0.03, "output": 0.06},
+    "gpt-4-32k-20230321":{"input": 0.06, "output": 0.12},
+    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-visual-preview": {"input": 0.01, "output": 0.03},
+}
+
+
+def get_cost_estimator(type: str, model: str, prompt_tokens: int, completion_tokens: int) -> float:
+    if type == "openai":
+        cost = prompt_tokens * OPENAI_PRICES[model]["input"]/1000 + completion_tokens * OPENAI_PRICES[model]["output"]/1000
+    elif type == "azure_ad" or type == "aoai":
+        if model in AZURE_PRICES:
+            cost = prompt_tokens * AZURE_PRICES[model]["input"]/1000 + completion_tokens * AZURE_PRICES[model]["output"]/1000
+        else:
+            print(f"Model {model} not found in Azure prices")
+            return None
+    else:
+        return None 
+    return cost
\ No newline at end of file
diff --git a/ufo/module/flow.py b/ufo/module/flow.py
index 1e70704a..28179042 100644
--- a/ufo/module/flow.py
+++ b/ufo/module/flow.py
@@ -56,7 +56,7 @@ def __init__(self, task):
         self.plan = ""
         self.request = ""
         self.results = ""
-        self.cost = 0
+        self.cost = 0.0
         self.offline_doc_retriever = None
         self.online_doc_retriever = None
         self.experience_retriever = None
@@ -108,8 +108,10 @@ def process_application_selection(self):
             self.request_logger.info(log)
             self.status = "ERROR"
             return
-
-        self.cost += cost
+        if isinstance(cost, float) and isinstance(self.cost, float):
+            self.cost += cost
+        else:
+            self.cost = None
 
         try:
             response_json = json_parser(response_string)
@@ -257,7 +259,10 @@ def process_action_selection(self):
                 time.sleep(configs["SLEEP_TIME"])
                 return 
             
-            self.cost += cost
+            if isinstance(cost, float) and isinstance(self.cost, float):
+                self.cost += cost
+            else:
+                self.cost = None
 
             try:
                 response_json = json_parser(response_string)
@@ -414,8 +419,11 @@ def experience_saver(self):
         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
+        
+        if isinstance(total_cost, float) and isinstance(self.cost, float):
+            self.cost += total_cost
+        else:
+            self.cost = None
         print_with_color("The experience has been saved.", "cyan")
 
 
@@ -476,6 +484,8 @@ def get_cost(self):
         Get the cost of the session.
         return: The cost of the session.
         """
+        if not isinstance(self.cost, float):
+            return "Cost is not available. Please try to update the price."
         return self.cost
     
     def get_application_window(self):

From 4b1b66153f1dc7956a579653cd78717e9cedb71b Mon Sep 17 00:00:00 2001
From: Mac0q <gmacq0209@gmail.com>
Date: Tue, 26 Mar 2024 17:19:12 +0800
Subject: [PATCH 5/9] price

---
 ufo/module/flow.py | 2 +-
 ufo/ufo.py         | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/ufo/module/flow.py b/ufo/module/flow.py
index 9cca13e9..8c5091b3 100644
--- a/ufo/module/flow.py
+++ b/ufo/module/flow.py
@@ -484,7 +484,7 @@ def get_cost(self):
         return: The cost of the session.
         """
         if not isinstance(self.cost, float):
-            return "Cost is not available. Please try to update the price."
+            return None
         return self.cost
     
     def get_application_window(self):
diff --git a/ufo/ufo.py b/ufo/ufo.py
index db618564..6e24ef78 100644
--- a/ufo/ufo.py
+++ b/ufo/ufo.py
@@ -83,8 +83,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 total_cost:
+        formatted_cost = '${:.2f}'.format(total_cost)
+        print_with_color(f"Request total cost is {formatted_cost}", "yellow")
 
     return status
 

From c3c7986df86f7a08a46f766578b4fc08dc1912d7 Mon Sep 17 00:00:00 2001
From: Mac0q <gmacq0209@gmail.com>
Date: Tue, 26 Mar 2024 20:40:48 +0800
Subject: [PATCH 6/9] price config

---
 ufo/config/config.yaml.template | 35 +++++++++++++++++++++++-
 ufo/experience/summarizer.py    |  2 +-
 ufo/llm/openai.py               |  5 ++--
 ufo/llm/util.py                 | 48 ---------------------------------
 ufo/utils/__init__.py           | 15 ++++++++++-
 5 files changed, 52 insertions(+), 53 deletions(-)
 delete mode 100644 ufo/llm/util.py

diff --git a/ufo/config/config.yaml.template b/ufo/config/config.yaml.template
index 52734bee..b0af10c9 100644
--- a/ufo/config/config.yaml.template
+++ b/ufo/config/config.yaml.template
@@ -96,6 +96,39 @@ RAG_EXPERIENCE: True  # Whether to use the RAG from its self-experience.
 RAG_EXPERIENCE_RETRIEVED_TOPK: 5  # The topk for the offline retrieved documents
 
 
-
+# Source: https://openai.com/pricing
+# Prices in $ per 1000 tokens
+# Last updated: 2024-01-26
+OPENAI_PRICES = {
+    "gpt-4-0613": {"input": 0.03, "output": 0.06},
+    "gpt-3.5-turbo-0613": {"input": 0.0015, "output": 0.002},
+    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-1106-vision-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4": {"input": 0.03, "output": 0.06},
+    "gpt-4-32k": {"input": 0.06, "output": 0.12},
+    "gpt-3.5-turbo-0125": {"input": 0.0005, "output": 0.0015},
+    "gpt-3.5-turbo-1106": {"input": 0.001, "output": 0.002},
+    "gpt-3.5-turbo-instruct": {"input": 0.0015, "output": 0.002},
+    "gpt-3.5-turbo-16k-0613": {"input": 0.003, "output": 0.004},
+    "whisper-1": {"input": 0.006, "output": 0.006},
+    "tts-1": {"input": 0.015, "output": 0.015},
+    "tts-hd-1": {"input": 0.03, "output": 0.03},
+    "text-embedding-ada-002-v2": {"input": 0.0001, "output": 0.0001},
+    "text-davinci:003": {"input": 0.02, "output": 0.02},
+    "text-ada-001": {"input": 0.0004, "output": 0.0004},
+    }
+
+AZURE_PRICES = {
+    "gpt-35-turbo-20220309":{"input": 0.0015, "output": 0.002},
+    "gpt-35-turbo-20230613":{"input": 0.0015, "output": 0.002},
+    "gpt-35-turbo-16k-20230613":{"input": 0.003, "output": 0.004},
+    "gpt-35-turbo-1106":{"input": 0.001, "output": 0.002},
+    "gpt-4-20230321":{"input": 0.03, "output": 0.06},
+    "gpt-4-32k-20230321":{"input": 0.06, "output": 0.12},
+    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-visual-preview": {"input": 0.01, "output": 0.03},
+}
 
 
diff --git a/ufo/experience/summarizer.py b/ufo/experience/summarizer.py
index 180809ca..b0d1bc71 100644
--- a/ufo/experience/summarizer.py
+++ b/ufo/experience/summarizer.py
@@ -78,7 +78,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)
diff --git a/ufo/llm/openai.py b/ufo/llm/openai.py
index b15d9be8..0e3db248 100644
--- a/ufo/llm/openai.py
+++ b/ufo/llm/openai.py
@@ -3,7 +3,8 @@
 import openai
 from openai import AzureOpenAI, OpenAI
 
-from ufo.llm.util import get_cost_estimator
+from ..utils import get_cost_estimator
+
 
 
 
@@ -63,7 +64,7 @@ def chat_completion(
             prompt_tokens = usage.prompt_tokens
             completion_tokens = usage.completion_tokens
 
-            cost = get_cost_estimator(self.config_llm["API_TYPE"], model, prompt_tokens, completion_tokens)
+            cost = get_cost_estimator(self.config, prompt_tokens, completion_tokens)
 
             return response.choices[0].message.content, cost
 
diff --git a/ufo/llm/util.py b/ufo/llm/util.py
deleted file mode 100644
index de451ce3..00000000
--- a/ufo/llm/util.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# Source: https://openai.com/pricing
-# Prices in $ per 1000 tokens
-# Last updated: 2024-01-26
-OPENAI_PRICES = {
-    "gpt-4-0613": {"input": 0.03, "output": 0.06},
-    "gpt-3.5-turbo-0613": {"input": 0.0015, "output": 0.002},
-    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-1106-vision-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4": {"input": 0.03, "output": 0.06},
-    "gpt-4-32k": {"input": 0.06, "output": 0.12},
-    "gpt-3.5-turbo-0125": {"input": 0.0005, "output": 0.0015},
-    "gpt-3.5-turbo-1106": {"input": 0.001, "output": 0.002},
-    "gpt-3.5-turbo-instruct": {"input": 0.0015, "output": 0.002},
-    "gpt-3.5-turbo-16k-0613": {"input": 0.003, "output": 0.004},
-    "whisper-1": {"input": 0.006, "output": 0.006},
-    "tts-1": {"input": 0.015, "output": 0.015},
-    "tts-hd-1": {"input": 0.03, "output": 0.03},
-    "text-embedding-ada-002-v2": {"input": 0.0001, "output": 0.0001},
-    "text-davinci:003": {"input": 0.02, "output": 0.02},
-    "text-ada-001": {"input": 0.0004, "output": 0.0004},
-    }
-
-AZURE_PRICES = {
-    "gpt-35-turbo-20220309":{"input": 0.0015, "output": 0.002},
-    "gpt-35-turbo-20230613":{"input": 0.0015, "output": 0.002},
-    "gpt-35-turbo-16k-20230613":{"input": 0.003, "output": 0.004},
-    "gpt-35-turbo-1106":{"input": 0.001, "output": 0.002},
-    "gpt-4-20230321":{"input": 0.03, "output": 0.06},
-    "gpt-4-32k-20230321":{"input": 0.06, "output": 0.12},
-    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-visual-preview": {"input": 0.01, "output": 0.03},
-}
-
-
-def get_cost_estimator(type: str, model: str, prompt_tokens: int, completion_tokens: int) -> float:
-    if type == "openai":
-        cost = prompt_tokens * OPENAI_PRICES[model]["input"]/1000 + completion_tokens * OPENAI_PRICES[model]["output"]/1000
-    elif type == "azure_ad" or type == "aoai":
-        if model in AZURE_PRICES:
-            cost = prompt_tokens * AZURE_PRICES[model]["input"]/1000 + completion_tokens * AZURE_PRICES[model]["output"]/1000
-        else:
-            print(f"Model {model} not found in Azure prices")
-            return None
-    else:
-        return None 
-    return cost
\ No newline at end of file
diff --git a/ufo/utils/__init__.py b/ufo/utils/__init__.py
index d55f06ea..574208fb 100644
--- a/ufo/utils/__init__.py
+++ b/ufo/utils/__init__.py
@@ -178,7 +178,20 @@ def revise_line_breaks(args: dict):
 
 
     
-
+def get_cost_estimator(configs, prompt_tokens, completion_tokens) -> float:
+        model_type = configs["API_TYPE"]
+        model = configs["API_MODEL"]
+        if model_type == "openai":
+            cost = prompt_tokens * configs["OPENAI_PRICES"][model]["input"]/1000 + completion_tokens * configs["OPENAI_PRICES"][model]["output"]/1000
+        elif model_type == "azure_ad" or type == "aoai":
+            if model in configs["OPENAI_PRICES"]:
+                cost = prompt_tokens * configs["AZURE_PRICES"][model]["input"]/1000 + completion_tokens * configs["AZURE_PRICES"][model]["output"]/1000
+            else:
+                print(f"Model {model} not found in Azure prices")
+                return None
+        else:
+            return None 
+        return cost
 
 
 

From 36d7fac8ac25a513c6510738047f3816a14f1b25 Mon Sep 17 00:00:00 2001
From: Mac0q <gmacq0209@gmail.com>
Date: Wed, 27 Mar 2024 09:30:13 +0800
Subject: [PATCH 7/9] price

---
 ufo/utils/__init__.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/ufo/utils/__init__.py b/ufo/utils/__init__.py
index 574208fb..b6646297 100644
--- a/ufo/utils/__init__.py
+++ b/ufo/utils/__init__.py
@@ -182,7 +182,11 @@ def get_cost_estimator(configs, prompt_tokens, completion_tokens) -> float:
         model_type = configs["API_TYPE"]
         model = configs["API_MODEL"]
         if model_type == "openai":
-            cost = prompt_tokens * configs["OPENAI_PRICES"][model]["input"]/1000 + completion_tokens * configs["OPENAI_PRICES"][model]["output"]/1000
+            if model in configs["OPENAI_PRICES"]:
+                cost = prompt_tokens * configs["OPENAI_PRICES"][model]["input"]/1000 + completion_tokens * configs["OPENAI_PRICES"][model]["output"]/1000
+            else:
+                print(f"Model {model} not found in OpenAI prices")
+                return None
         elif model_type == "azure_ad" or type == "aoai":
             if model in configs["OPENAI_PRICES"]:
                 cost = prompt_tokens * configs["AZURE_PRICES"][model]["input"]/1000 + completion_tokens * configs["AZURE_PRICES"][model]["output"]/1000

From 4b4d344e0f377befc3a101752de0f070b0bdddd7 Mon Sep 17 00:00:00 2001
From: Mac0q <gmacq0209@gmail.com>
Date: Wed, 27 Mar 2024 13:49:56 +0800
Subject: [PATCH 8/9] price

---
 ufo/config/config.py            |  4 ++
 ufo/config/config.yaml.template | 40 +-------------------
 ufo/config/config_prices.yaml   | 36 ++++++++++++++++++
 ufo/llm/openai.py               | 66 ++++++++++++++++++++-------------
 ufo/module/flow.py              | 24 ++++++------
 ufo/ufo.py                      |  2 +-
 ufo/utils/__init__.py           | 27 +-------------
 7 files changed, 96 insertions(+), 103 deletions(-)
 create mode 100644 ufo/config/config_prices.yaml

diff --git a/ufo/config/config.py b/ufo/config/config.py
index 576a9cdd..e45fe09c 100644
--- a/ufo/config/config.py
+++ b/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")
diff --git a/ufo/config/config.yaml.template b/ufo/config/config.yaml.template
index b0af10c9..d7c21ef3 100644
--- a/ufo/config/config.yaml.template
+++ b/ufo/config/config.yaml.template
@@ -93,42 +93,4 @@ RAG_ONLINE_RETRIEVED_TOPK: 1 # The topk for the online retrieved documents
 
 ## RAG Configuration for experience
 RAG_EXPERIENCE: True  # Whether to use the RAG from its self-experience.
-RAG_EXPERIENCE_RETRIEVED_TOPK: 5  # The topk for the offline retrieved documents
-
-
-# Source: https://openai.com/pricing
-# Prices in $ per 1000 tokens
-# Last updated: 2024-01-26
-OPENAI_PRICES = {
-    "gpt-4-0613": {"input": 0.03, "output": 0.06},
-    "gpt-3.5-turbo-0613": {"input": 0.0015, "output": 0.002},
-    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-1106-vision-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4": {"input": 0.03, "output": 0.06},
-    "gpt-4-32k": {"input": 0.06, "output": 0.12},
-    "gpt-3.5-turbo-0125": {"input": 0.0005, "output": 0.0015},
-    "gpt-3.5-turbo-1106": {"input": 0.001, "output": 0.002},
-    "gpt-3.5-turbo-instruct": {"input": 0.0015, "output": 0.002},
-    "gpt-3.5-turbo-16k-0613": {"input": 0.003, "output": 0.004},
-    "whisper-1": {"input": 0.006, "output": 0.006},
-    "tts-1": {"input": 0.015, "output": 0.015},
-    "tts-hd-1": {"input": 0.03, "output": 0.03},
-    "text-embedding-ada-002-v2": {"input": 0.0001, "output": 0.0001},
-    "text-davinci:003": {"input": 0.02, "output": 0.02},
-    "text-ada-001": {"input": 0.0004, "output": 0.0004},
-    }
-
-AZURE_PRICES = {
-    "gpt-35-turbo-20220309":{"input": 0.0015, "output": 0.002},
-    "gpt-35-turbo-20230613":{"input": 0.0015, "output": 0.002},
-    "gpt-35-turbo-16k-20230613":{"input": 0.003, "output": 0.004},
-    "gpt-35-turbo-1106":{"input": 0.001, "output": 0.002},
-    "gpt-4-20230321":{"input": 0.03, "output": 0.06},
-    "gpt-4-32k-20230321":{"input": 0.06, "output": 0.12},
-    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-visual-preview": {"input": 0.01, "output": 0.03},
-}
-
-
+RAG_EXPERIENCE_RETRIEVED_TOPK: 5  # The topk for the offline retrieved documents
\ No newline at end of file
diff --git a/ufo/config/config_prices.yaml b/ufo/config/config_prices.yaml
new file mode 100644
index 00000000..6223e451
--- /dev/null
+++ b/ufo/config/config_prices.yaml
@@ -0,0 +1,36 @@
+# Source: https://openai.com/pricing
+# Prices in $ per 1000 tokens
+# Last updated: 2024-03-27
+OPENAI_PRICES = {
+    "gpt-4-0613": {"input": 0.03, "output": 0.06},
+    "gpt-3.5-turbo-0613": {"input": 0.0015, "output": 0.002},
+    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-1106-vision-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4": {"input": 0.03, "output": 0.06},
+    "gpt-4-32k": {"input": 0.06, "output": 0.12},
+    "gpt-3.5-turbo-0125": {"input": 0.0005, "output": 0.0015},
+    "gpt-3.5-turbo-1106": {"input": 0.001, "output": 0.002},
+    "gpt-3.5-turbo-instruct": {"input": 0.0015, "output": 0.002},
+    "gpt-3.5-turbo-16k-0613": {"input": 0.003, "output": 0.004},
+    "whisper-1": {"input": 0.006, "output": 0.006},
+    "tts-1": {"input": 0.015, "output": 0.015},
+    "tts-hd-1": {"input": 0.03, "output": 0.03},
+    "text-embedding-ada-002-v2": {"input": 0.0001, "output": 0.0001},
+    "text-davinci:003": {"input": 0.02, "output": 0.02},
+    "text-ada-001": {"input": 0.0004, "output": 0.0004},
+    }
+
+AZURE_PRICES = {
+    "gpt-35-turbo-20220309":{"input": 0.0015, "output": 0.002},
+    "gpt-35-turbo-20230613":{"input": 0.0015, "output": 0.002},
+    "gpt-35-turbo-16k-20230613":{"input": 0.003, "output": 0.004},
+    "gpt-35-turbo-1106":{"input": 0.001, "output": 0.002},
+    "gpt-4-20230321":{"input": 0.03, "output": 0.06},
+    "gpt-4-32k-20230321":{"input": 0.06, "output": 0.12},
+    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-visual-preview": {"input": 0.01, "output": 0.03},
+}
+
+
diff --git a/ufo/llm/openai.py b/ufo/llm/openai.py
index 0e3db248..04ce2b1b 100644
--- a/ufo/llm/openai.py
+++ b/ufo/llm/openai.py
@@ -1,37 +1,38 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 import datetime
 from typing import Any, Optional
 import openai
 from openai import AzureOpenAI, OpenAI
 
-from ..utils import get_cost_estimator
-
-
-
 
 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.client: OpenAI = (
-            OpenAI(
-                base_url=self.config_llm["API_BASE"],
-                api_key=self.config_llm["API_KEY"],
-                max_retries=max_retry,
-                timeout=self.config["TIMEOUT"],
-            )
-            if api_type == "openai"
-            else AzureOpenAI(
-                max_retries=max_retry,
+        self.api_type = self.config_llm["API_TYPE"].lower()
+        self.max_retry = self.config["MAX_RETRY"]
+        assert self.api_type in ["openai", "aoai", "azure_ad"], "Invalid API type"
+        if self.api_type == "openai":
+            self.client = OpenAI(
+                    base_url=self.config_llm["API_BASE"],
+                    api_key=self.config_llm["API_KEY"],
+                    max_retries=self.max_retry,
+                    timeout=self.config["TIMEOUT"],
+                )
+            self.prices = self.config["OPENAI_PRICES"]
+        else:
+            self.client = AzureOpenAI(
+                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":
+            self.prices = self.config["AZURE_PRICES"]
+
+        if self.api_type == "azure_ad":
             self.auto_refresh_token()
 
     def chat_completion(
@@ -64,7 +65,7 @@ def chat_completion(
             prompt_tokens = usage.prompt_tokens
             completion_tokens = usage.completion_tokens
 
-            cost = get_cost_estimator(self.config, prompt_tokens, completion_tokens)
+            cost = self.get_cost_estimator(model, self.prices, prompt_tokens, completion_tokens)
 
             return response.choices[0].message.content, cost
 
@@ -90,9 +91,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',
@@ -199,7 +197,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',
@@ -266,3 +263,22 @@ def stop():
 
         return stop
 
+    def get_cost_estimator(self, 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 model in prices:
+            cost = prompt_tokens * prices[model]["input"]/1000 + completion_tokens * prices[model]["output"]/1000
+        else:
+            print(f"Model {model} not found in prices")
+            return None
+        return cost
diff --git a/ufo/module/flow.py b/ufo/module/flow.py
index 8c5091b3..fc671d3e 100644
--- a/ufo/module/flow.py
+++ b/ufo/module/flow.py
@@ -109,10 +109,7 @@ def process_application_selection(self):
             self.request_logger.info(log)
             self.status = "ERROR"
             return
-        if isinstance(cost, float) and isinstance(self.cost, float):
-            self.cost += cost
-        else:
-            self.cost = None
+        self.update_cost(cost=cost)
 
         try:
             response_json = json_parser(response_string)
@@ -263,10 +260,7 @@ def process_action_selection(self):
             time.sleep(configs["SLEEP_TIME"])
             return 
             
-        if isinstance(cost, float) and isinstance(self.cost, float):
-            self.cost += cost
-        else:
-            self.cost = None
+        self.update_cost(cost=cost)
 
         try:
             response_json = json_parser(response_string)
@@ -419,10 +413,7 @@ def experience_saver(self):
         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"))
         
-        if isinstance(total_cost, float) and isinstance(self.cost, float):
-            self.cost += total_cost
-        else:
-            self.cost = None
+        self.update_cost(cost=total_cost)
         print_with_color("The experience has been saved.", "cyan")
 
 
@@ -532,6 +523,15 @@ def error_logger(self, response_str, error):
         log = json.dumps({"step": self.step, "status": "ERROR", "response": response_str, "error": error})
         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, log_filename):
         """
diff --git a/ufo/ufo.py b/ufo/ufo.py
index 6e24ef78..9897a4ce 100644
--- a/ufo/ufo.py
+++ b/ufo/ufo.py
@@ -83,7 +83,7 @@ def main():
 
     # Print the total cost
     total_cost = session.get_cost()
-    if total_cost:
+    if isinstance(total_cost, float):
         formatted_cost = '${:.2f}'.format(total_cost)
         print_with_color(f"Request total cost is {formatted_cost}", "yellow")
 
diff --git a/ufo/utils/__init__.py b/ufo/utils/__init__.py
index b6646297..f86a0a59 100644
--- a/ufo/utils/__init__.py
+++ b/ufo/utils/__init__.py
@@ -173,29 +173,4 @@ def revise_line_breaks(args: dict):
         if isinstance(args[key], str):
             args[key] = args[key].replace('\\n', '\n')
 
-    return args
-
-
-
-    
-def get_cost_estimator(configs, prompt_tokens, completion_tokens) -> float:
-        model_type = configs["API_TYPE"]
-        model = configs["API_MODEL"]
-        if model_type == "openai":
-            if model in configs["OPENAI_PRICES"]:
-                cost = prompt_tokens * configs["OPENAI_PRICES"][model]["input"]/1000 + completion_tokens * configs["OPENAI_PRICES"][model]["output"]/1000
-            else:
-                print(f"Model {model} not found in OpenAI prices")
-                return None
-        elif model_type == "azure_ad" or type == "aoai":
-            if model in configs["OPENAI_PRICES"]:
-                cost = prompt_tokens * configs["AZURE_PRICES"][model]["input"]/1000 + completion_tokens * configs["AZURE_PRICES"][model]["output"]/1000
-            else:
-                print(f"Model {model} not found in Azure prices")
-                return None
-        else:
-            return None 
-        return cost
-
-
-
+    return args
\ No newline at end of file

From d06671289f6ec4d9b888daae44cf8969ad37367d Mon Sep 17 00:00:00 2001
From: Mac0q <gmacq0209@gmail.com>
Date: Thu, 28 Mar 2024 15:20:49 +0800
Subject: [PATCH 9/9] prices

---
 ufo/config/config.py          |  1 +
 ufo/config/config_prices.yaml | 57 +++++++++++++++++------------------
 ufo/llm/openai.py             | 37 ++++++++++++-----------
 3 files changed, 48 insertions(+), 47 deletions(-)

diff --git a/ufo/config/config.py b/ufo/config/config.py
index e45fe09c..d7f2a552 100644
--- a/ufo/config/config.py
+++ b/ufo/config/config.py
@@ -61,6 +61,7 @@ def update_api_base(configs, agent):
 def optimize_configs(configs):
     update_api_base(configs,'APP_AGENT')
     update_api_base(configs,'ACTION_AGENT')
+    update_api_base(configs,'BACKUP_AGENT')
     
     return configs
 
diff --git a/ufo/config/config_prices.yaml b/ufo/config/config_prices.yaml
index 6223e451..efe93a6a 100644
--- a/ufo/config/config_prices.yaml
+++ b/ufo/config/config_prices.yaml
@@ -1,36 +1,33 @@
 # Source: https://openai.com/pricing
 # Prices in $ per 1000 tokens
 # Last updated: 2024-03-27
-OPENAI_PRICES = {
-    "gpt-4-0613": {"input": 0.03, "output": 0.06},
-    "gpt-3.5-turbo-0613": {"input": 0.0015, "output": 0.002},
-    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-1106-vision-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4": {"input": 0.03, "output": 0.06},
-    "gpt-4-32k": {"input": 0.06, "output": 0.12},
-    "gpt-3.5-turbo-0125": {"input": 0.0005, "output": 0.0015},
-    "gpt-3.5-turbo-1106": {"input": 0.001, "output": 0.002},
-    "gpt-3.5-turbo-instruct": {"input": 0.0015, "output": 0.002},
-    "gpt-3.5-turbo-16k-0613": {"input": 0.003, "output": 0.004},
-    "whisper-1": {"input": 0.006, "output": 0.006},
-    "tts-1": {"input": 0.015, "output": 0.015},
-    "tts-hd-1": {"input": 0.03, "output": 0.03},
-    "text-embedding-ada-002-v2": {"input": 0.0001, "output": 0.0001},
-    "text-davinci:003": {"input": 0.02, "output": 0.02},
-    "text-ada-001": {"input": 0.0004, "output": 0.0004},
-    }
-
-AZURE_PRICES = {
-    "gpt-35-turbo-20220309":{"input": 0.0015, "output": 0.002},
-    "gpt-35-turbo-20230613":{"input": 0.0015, "output": 0.002},
-    "gpt-35-turbo-16k-20230613":{"input": 0.003, "output": 0.004},
-    "gpt-35-turbo-1106":{"input": 0.001, "output": 0.002},
-    "gpt-4-20230321":{"input": 0.03, "output": 0.06},
-    "gpt-4-32k-20230321":{"input": 0.06, "output": 0.12},
-    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
-    "gpt-4-visual-preview": {"input": 0.01, "output": 0.03},
+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}
 }
 
 
diff --git a/ufo/llm/openai.py b/ufo/llm/openai.py
index 04ce2b1b..46155f8c 100644
--- a/ufo/llm/openai.py
+++ b/ufo/llm/openai.py
@@ -13,25 +13,24 @@ def __init__(self, config, agent_type: str):
         self.config = config
         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"
-        if self.api_type == "openai":
-            self.client = OpenAI(
-                    base_url=self.config_llm["API_BASE"],
-                    api_key=self.config_llm["API_KEY"],
-                    max_retries=self.max_retry,
-                    timeout=self.config["TIMEOUT"],
-                )
-            self.prices = self.config["OPENAI_PRICES"]
-        else:
-            self.client = AzureOpenAI(
+        self.client: OpenAI = (
+            OpenAI(
+                base_url=self.config_llm["API_BASE"],
+                api_key=self.config_llm["API_KEY"],
+                max_retries=self.max_retry,
+                timeout=self.config["TIMEOUT"],
+            )
+            if self.api_type == "openai"
+            else AzureOpenAI(
                 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 self.api_type == 'aoai' else self.get_openai_token()),
             )
-            self.prices = self.config["AZURE_PRICES"]
-
+        )
         if self.api_type == "azure_ad":
             self.auto_refresh_token()
 
@@ -65,7 +64,7 @@ def chat_completion(
             prompt_tokens = usage.prompt_tokens
             completion_tokens = usage.completion_tokens
 
-            cost = self.get_cost_estimator(model, self.prices, prompt_tokens, completion_tokens)
+            cost = self.get_cost_estimator(self.api_type, model, self.prices, prompt_tokens, completion_tokens)
 
             return response.choices[0].message.content, cost
 
@@ -263,7 +262,7 @@ def stop():
 
         return stop
 
-    def get_cost_estimator(self, model, prices, prompt_tokens, completion_tokens) -> float:
+    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.
 
@@ -276,9 +275,13 @@ def get_cost_estimator(self, model, prices, prompt_tokens, completion_tokens) ->
         Returns:
             float: The estimated cost for using the model.
         """
-        if model in prices:
-            cost = prompt_tokens * prices[model]["input"]/1000 + completion_tokens * prices[model]["output"]/1000
+        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"Model {model} not found in prices")
+            print(f"{name} not found in prices")
             return None
         return cost