Add pricing tracking for Anthropic model and refactor pricing functions

Author: recursix

src/agentlab/llm/response_api.py

@@ -9,6 +9,8 @@
 from anthropic import Anthropic
 from openai import OpenAI
 
+from agentlab.llm import tracking
+
 from .base_api import BaseModelArgs
 
 type ContentItem = Dict[str, Any]
@@ -269,6 +271,20 @@ def __init__(
             max_tokens=max_tokens,
             extra_kwargs=extra_kwargs,
         )
+
+        # Get pricing information
+
+        try:
+            pricing = tracking.get_pricing_anthropic()
+            self.input_cost = float(pricing[model_name]["prompt"])
+            self.output_cost = float(pricing[model_name]["completion"])
+        except KeyError:
+            logging.warning(
+                f"Model {model_name} not found in the pricing information, prices are set to 0. Maybe try upgrading langchain_community."
+            )
+            self.input_cost = 0.0
+            self.output_cost = 0.0
+
         self.client = Anthropic(api_key=api_key)
 
     def _call_api(self, messages: list[dict | MessageBuilder]) -> dict:
@@ -286,6 +302,17 @@ def _call_api(self, messages: list[dict | MessageBuilder]) -> dict:
                 max_tokens=self.max_tokens,
                 **self.extra_kwargs,
             )
+            input_tokens = response.usage.input_tokens
+            output_tokens = response.usage.output_tokens
+            cost = input_tokens * self.input_cost + output_tokens * self.output_cost
+
+            print(f"response.usage: {response.usage}")
+
+            if hasattr(tracking.TRACKER, "instance") and isinstance(
+                tracking.TRACKER.instance, tracking.LLMTracker
+            ):
+                tracking.TRACKER.instance(input_tokens, output_tokens, cost)
+
             return response
         except Exception as e:
             logging.error(f"Failed to get a response from the API: {e}")

src/agentlab/llm/tracking.py

@@ -1,10 +1,11 @@
 import os
+import re
 import threading
 from contextlib import contextmanager
 from functools import cache
 
 import requests
-from langchain_community.callbacks.openai_info import MODEL_COST_PER_1K_TOKENS
+from langchain_community.callbacks import bedrock_anthropic_callback, openai_info
 
 TRACKER = threading.local()
 
@@ -85,7 +86,7 @@ def get_pricing_openrouter():
 
 
 def get_pricing_openai():
-    cost_dict = MODEL_COST_PER_1K_TOKENS
+    cost_dict = openai_info.MODEL_COST_PER_1K_TOKENS
     cost_dict = {k: v / 1000 for k, v in cost_dict.items()}
     res = {}
     for k in cost_dict:
@@ -99,3 +100,25 @@ def get_pricing_openai():
                 "completion": cost_dict[completion_key],
             }
     return res
+
+
+def _remove_version_suffix(model_name):
+    no_version = re.sub(r"-v\d+(?:[.:]\d+)?$", "", model_name)
+    return re.sub(r"anthropic.", "", no_version)
+
+
+def get_pricing_anthropic():
+    input_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_INPUT_TOKENS
+    output_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_OUTPUT_TOKENS
+
+    res = {}
+    for k, v in input_cost_dict.items():
+        k = _remove_version_suffix(k)
+        res[k] = {"prompt": v / 1000}
+
+    for k, v in output_cost_dict.items():
+        k = _remove_version_suffix(k)
+        if k not in res:
+            res[k] = {}
+        res[k]["completion"] = v / 1000
+    return res