Added cost tracking.

toolformers/camel.py

@@ -12,6 +12,19 @@ from camel.agents import ChatAgent
 from camel.toolkits.function_tool import FunctionTool
 from camel.configs.openai_config import ChatGPTConfig
 
+from utils import register_cost
+
+COSTS = {
+    'gpt-4o': {
+        'prompt_tokens': 2.5e-6,
+        'completion_tokens': 10e-6
+    },
+    'gpt-4o-mini': {
+        'prompt_tokens': 0.15e-6,
+        'completion_tokens': 0.6e-6
+    }
+}
+
 class CamelConversation(Conversation):
     def __init__(self, toolformer, agent, category=None):
         self.toolformer = toolformer
@@ -32,6 +45,14 @@ class CamelConversation(Conversation):
         
         response = self.agent.step(formatted_message)
 
+        if response.info.get('usage', None) is not None:
+            usage_data = response.info['usage']
+
+            total = 0
+            for cost_name in ['prompt_tokens', 'completion_tokens']:
+                total += COSTS[str(self.toolformer.model_type)][cost_name] * usage_data[cost_name]
+            register_cost(self.category, total)
+
         reply = response.msg.content
 
         if print_output:

toolformers/gemini.py

@@ -10,8 +10,17 @@ from toolformers.base import Conversation, Tool, Toolformer
 import google.generativeai as genai
 from google.generativeai.generative_models import ChatSession
 
+from utils import register_cost
+
 genai.configure(api_key=os.environ['GOOGLE_API_KEY'])
 
+COSTS = {
+    'gemini-1.5-pro': {
+        'prompt_tokens': 1.25e-6,
+        'completion_tokens': 5e-6
+    }
+}
+
 class GeminiConversation(Conversation):
     def __init__(self, model_name, chat_agent : ChatSession, category=None):
         self.model_name = model_name
@@ -60,6 +69,12 @@ class GeminiConversation(Conversation):
             'completion_tokens': response.usage_metadata.candidates_token_count
         }
 
+        total_cost = 0
+        for cost_name in ['prompt_tokens', 'completion_tokens']:
+            total_cost += COSTS[self.model_name][cost_name] * usage_info[cost_name]
+
+        register_cost(self.category, total_cost)
+
         #send_usage_to_db(
         #    usage_info,
         #    time_start,

toolformers/langchain_agent.py

@@ -15,10 +15,23 @@ from langchain_core.tools import tool as function_to_tool
 
 from toolformers.base import StringParameter, Toolformer, Conversation
 
+from utils import register_cost
 
 
+COSTS = {
+    'claude-3-5-sonnet-latest': {
+        'input_tokens': 3e-6,
+        'output_tokens': 15e-6
+    },
+    'claude-3-5-haiku-latest': {
+        'input_tokens': 1e-6,
+        'output_tokens': 5e-6
+    }
+}
+
 class LangChainConversation(Conversation):
-    def __init__(self, agent, messages, category=None):
+    def __init__(self, model_name, agent, messages, category=None):
+        self.model_name = model_name
         self.agent = agent
         self.messages = messages
         self.category = category
@@ -26,6 +39,9 @@ class LangChainConversation(Conversation):
     def chat(self, message, role='user', print_output=True) -> str:
         self.messages.append(HumanMessage(content=message))
         final_message = ''
+
+        aggregate = None
+
         for chunk in self.agent.stream({"messages": self.messages}, stream_mode="values"):
             print(chunk)
             print("----")
@@ -38,8 +54,18 @@ class LangChainConversation(Conversation):
                         for content_chunk in content:
                             if isinstance(content_chunk, str):
                                 final_message += content_chunk
+
+            aggregate = chunk if aggregate is None else (aggregate + chunk)
             #final_message += chunk['agent']['messages'].content
 
+        total_cost = 0
+        for message in aggregate['messages']:
+            if isinstance(message, AIMessage):
+                for cost_name in ['input_tokens', 'output_tokens']:
+                    total_cost += COSTS[self.model_name][cost_name] * message.usage_metadata[cost_name]
+
+        register_cost(self.category, total_cost)
+
         self.messages.append(AIMessage(content=final_message))
         #print(final_message)
 
@@ -55,7 +81,7 @@ class LangChainAnthropicToolformer(Toolformer):
         model = ChatAnthropic(model_name=self.model_name, api_key=self.api_key)
         agent_executor = create_react_agent(model, tools)
         
-        return LangChainConversation(agent_executor, [SystemMessage(prompt)], category)
+        return LangChainConversation(self.model_name, agent_executor, [SystemMessage(prompt)], category)
 
 
 #weather_tool = AgoraTool("WeatherForecastAPI", "A simple tool that returns the weather", [StringParameter(

toolformers/sambanova/core.py

@@ -5,6 +5,14 @@ from typing import List
 from toolformers.base import Conversation, Toolformer, Tool
 from toolformers.sambanova.function_calling import FunctionCallingLlm
 
+
+COSTS = {
+    'llama3-405b': {
+        'prompt_tokens': 5e-6,
+        'completion_tokens': 10e-6
+    }
+}
+
 class SambanovaConversation(Conversation):
     def __init__(self, model_name, function_calling_llm : FunctionCallingLlm, category=None):
         self.model_name = model_name
@@ -26,8 +34,11 @@ class SambanovaConversation(Conversation):
         print('Usage data:', usage_data)
         if print_output:
             print(response)
-        
-        #send_usage_to_db(usage_data, start_time, end_time, agent_id, self.category, self.model_name)
+
+        cost = 0
+
+        for cost_name in ['prompt_tokens', 'completion_tokens']:
+            cost += COSTS[self.model_name][cost_name] * usage_data[cost_name]
         
         return response
 

utils.py

@@ -1,6 +1,8 @@
 import base64
 import hashlib
 
+import contextlib
+
 from functools import wraps
 from typing import Callable, Any
 from inspect import signature, Parameter, Signature
@@ -77,4 +79,20 @@ def add_params_and_annotations(name: str, description: str, params: dict, return
         print(docstring)
 
         return wrapper
-    return decorator
+    return decorator
+
+_cost_tracker = {}
+
+@contextlib.contextmanager
+def use_cost_tracker():
+    _cost_tracker.clear()
+    _cost_tracker['conversation'] = 0
+    _cost_tracker['negotiation'] = 0
+    _cost_tracker['programming'] = 0
+    yield
+
+def register_cost(category, cost):
+    _cost_tracker[category] += cost
+
+def get_costs():
+    return _cost_tracker