Add usage summary for agents (#1269)

* update

* update

* Update notebook/oai_client_cost.ipynb

Co-authored-by: Chi Wang <[email protected]>

* update doc and test

---------

Co-authored-by: Qingyun Wu <[email protected]>
Co-authored-by: Chi Wang <[email protected]>

Author: 3 people

autogen/agent_utils.py

@@ -0,0 +1,51 @@
+from typing import List, Dict, Tuple
+from autogen import Agent
+
+
+def gather_usage_summary(agents: List[Agent]) -> Tuple[Dict[str, any], Dict[str, any]]:
+    """Gather usage summary from all agents.
+
+    Args:
+        agents: (list): List of agents.
+
+    Returns:
+        tuple: (total_usage_summary, actual_usage_summary)
+
+    Example return:
+        total_usage_summary = {
+            'total_cost': 0.0006090000000000001,
+            'gpt-35-turbo':
+                {
+                    'cost': 0.0006090000000000001,
+                    'prompt_tokens': 242,
+                    'completion_tokens': 123,
+                    'total_tokens': 365
+                }
+        }
+        `actual_usage_summary` follows the same format.
+        If none of the agents incurred any cost (not having a client), then the total_usage_summary and actual_usage_summary will be {'total_cost': 0}.
+    """
+
+    def aggregate_summary(usage_summary: Dict[str, any], agent_summary: Dict[str, any]) -> None:
+        if agent_summary is None:
+            return
+        usage_summary["total_cost"] += agent_summary.get("total_cost", 0)
+        for model, data in agent_summary.items():
+            if model != "total_cost":
+                if model not in usage_summary:
+                    usage_summary[model] = data.copy()
+                else:
+                    usage_summary[model]["cost"] += data.get("cost", 0)
+                    usage_summary[model]["prompt_tokens"] += data.get("prompt_tokens", 0)
+                    usage_summary[model]["completion_tokens"] += data.get("completion_tokens", 0)
+                    usage_summary[model]["total_tokens"] += data.get("total_tokens", 0)
+
+    total_usage_summary = {"total_cost": 0}
+    actual_usage_summary = {"total_cost": 0}
+
+    for agent in agents:
+        if agent.client:
+            aggregate_summary(total_usage_summary, agent.client.total_usage_summary)
+            aggregate_summary(actual_usage_summary, agent.client.actual_usage_summary)
+
+    return total_usage_summary, actual_usage_summary

autogen/agentchat/conversable_agent.py

@@ -699,6 +699,8 @@ def reset(self):
         self.clear_history()
         self.reset_consecutive_auto_reply_counter()
         self.stop_reply_at_receive()
+        if self.client is not None:
+            self.client.clear_usage_summary()
         for reply_func_tuple in self._reply_func_list:
             if reply_func_tuple["reset_config"] is not None:
                 reply_func_tuple["reset_config"](reply_func_tuple["config"])
@@ -1890,3 +1892,25 @@ def process_last_message(self, messages):
         messages = messages.copy()
         messages[-1]["content"] = processed_user_text
         return messages
+
+    def print_usage_summary(self, mode: Union[str, List[str]] = ["actual", "total"]) -> None:
+        """Print the usage summary."""
+        if self.client is None:
+            print(f"No cost incurred from agent '{self.name}'.")
+        else:
+            print(f"Agent '{self.name}':")
+            self.client.print_usage_summary(mode)
+
+    def get_actual_usage(self) -> Union[None, Dict[str, int]]:
+        """Get the actual usage summary."""
+        if self.client is None:
+            return None
+        else:
+            return self.client.actual_usage_summary
+
+    def get_total_usage(self) -> Union[None, Dict[str, int]]:
+        """Get the total usage summary."""
+        if self.client is None:
+            return None
+        else:
+            return self.client.total_usage_summary

autogen/oai/openai_utils.py

@@ -42,7 +42,7 @@
     "gpt-4-0613": (0.03, 0.06),
     "gpt-4-32k-0613": (0.06, 0.12),
     # 11-06
-    "gpt-3.5-turbo": (0.001, 0.002),
+    "gpt-3.5-turbo": (0.0015, 0.002),  # default is still 0613
     "gpt-3.5-turbo-1106": (0.001, 0.002),
     "gpt-35-turbo-1106": (0.001, 0.002),
     "gpt-4-1106-preview": (0.01, 0.03),