From 65a8f78e625c9989a213f608d412677afcf26d67 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Wed, 11 Sep 2024 10:07:14 -0400 Subject: [PATCH 01/15] updating dependencies for github workflows --- requirements.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/requirements.txt b/requirements.txt index eae71c15..86890452 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,6 @@ +black[jupyter]==24.2.0 +blacken-docs +pre-commit pytest==7.3.2 pytest-xdist pytest-playwright From ab0f4ec2f06e5f0becbc8379a98611ffb8f8aab6 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Wed, 18 Sep 2024 12:12:06 -0400 Subject: [PATCH 02/15] openrouter tracker poc --- .../launch_command.py => launch_command.py | 11 +- src/agentlab/agents/generic_agent/__init__.py | 2 + .../agents/generic_agent/agent_configs.py | 8 +- .../agents/generic_agent/generic_agent.py | 3 +- src/agentlab/llm/chat_api.py | 3 +- src/agentlab/llm/tracking.py | 104 ++++++++++++++++++ 6 files changed, 120 insertions(+), 11 deletions(-) rename src/agentlab/experiments/launch_command.py => launch_command.py (76%) create mode 100644 src/agentlab/llm/tracking.py diff --git a/src/agentlab/experiments/launch_command.py b/launch_command.py similarity index 76% rename from src/agentlab/experiments/launch_command.py rename to launch_command.py index 01b48a7f..198bbe07 100644 --- a/src/agentlab/experiments/launch_command.py +++ b/launch_command.py @@ -7,21 +7,22 @@ import logging -from agentlab.agents.generic_agent import RANDOM_SEARCH_AGENT, AGENT_4o, AGENT_4o_MINI +from agentlab.agents.generic_agent import AGENT_CUSTOM, RANDOM_SEARCH_AGENT, AGENT_4o, AGENT_4o_MINI from agentlab.analyze.inspect_results import get_most_recent_folder from agentlab.experiments import study_generators from agentlab.experiments.exp_utils import RESULTS_DIR -from agentlab.experiments.launch_exp import make_study_dir, run_experiments, relaunch_study +from agentlab.experiments.launch_exp import make_study_dir, relaunch_study, run_experiments logging.getLogger().setLevel(logging.INFO) # choose your agent or provide a new agent -agent_args = AGENT_4o_MINI +agent_args = [AGENT_CUSTOM, AGENT_4o_MINI] # agent = AGENT_4o ## select the benchmark to run on benchmark = "miniwob" +benchmark = "miniwob_tiny_test" # benchmark = "workarena.l1" # benchmark = "workarena.l2" # benchmark = "workarena.l3" @@ -37,8 +38,8 @@ ## alternatively, relaunch an existing study -study_dir = get_most_recent_folder(RESULTS_DIR, contains=None) -exp_args_list, study_dir = relaunch_study(study_dir, relaunch_mode="incomplete_or_error") +# study_dir = get_most_recent_folder(RESULTS_DIR, contains=None) +# exp_args_list, study_dir = relaunch_study(study_dir, relaunch_mode="incomplete_or_error") ## Number of parallel jobs diff --git a/src/agentlab/agents/generic_agent/__init__.py b/src/agentlab/agents/generic_agent/__init__.py index d9839c4d..fec74910 100644 --- a/src/agentlab/agents/generic_agent/__init__.py +++ b/src/agentlab/agents/generic_agent/__init__.py @@ -2,6 +2,7 @@ AGENT_3_5, AGENT_8B, AGENT_70B, + AGENT_CUSTOM, RANDOM_SEARCH_AGENT, AGENT_4o, AGENT_4o_MINI, @@ -16,4 +17,5 @@ "AGENT_70B", "AGENT_8B", "RANDOM_SEARCH_AGENT", + "AGENT_CUSTOM", ] diff --git a/src/agentlab/agents/generic_agent/agent_configs.py b/src/agentlab/agents/generic_agent/agent_configs.py index a53046b2..da0a6799 100644 --- a/src/agentlab/agents/generic_agent/agent_configs.py +++ b/src/agentlab/agents/generic_agent/agent_configs.py @@ -1,9 +1,9 @@ -from .generic_agent_prompt import GenericPromptFlags from agentlab.agents import dynamic_prompting as dp -from .generic_agent import GenericAgentArgs -from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT from agentlab.experiments import args +from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT +from .generic_agent import GenericAgentArgs +from .generic_agent_prompt import GenericPromptFlags FLAGS_CUSTOM = GenericPromptFlags( obs=dp.ObsFlags( @@ -45,7 +45,7 @@ AGENT_CUSTOM = GenericAgentArgs( - chat_model_args=CHAT_MODEL_ARGS_DICT["openai/gpt-3.5-turbo-1106"], + chat_model_args=CHAT_MODEL_ARGS_DICT["openrouter/meta-llama/llama-3.1-8b-instruct"], flags=FLAGS_CUSTOM, ) diff --git a/src/agentlab/agents/generic_agent/generic_agent.py b/src/agentlab/agents/generic_agent/generic_agent.py index a53f1aeb..eb3b0d87 100644 --- a/src/agentlab/agents/generic_agent/generic_agent.py +++ b/src/agentlab/agents/generic_agent/generic_agent.py @@ -10,6 +10,7 @@ from agentlab.agents.utils import openai_monitored_agent from agentlab.llm.chat_api import BaseModelArgs from agentlab.llm.llm_utils import RetryError, retry_raise +from agentlab.llm.tracking import get_action_decorator from .generic_agent_prompt import GenericPromptFlags, MainPrompt @@ -65,7 +66,7 @@ def __init__( def obs_preprocessor(self, obs: dict) -> dict: return self._obs_preprocessor(obs) - @openai_monitored_agent + @get_action_decorator def get_action(self, obs): self.obs_history.append(obs) diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index 2f425da7..1246081a 100644 --- a/src/agentlab/llm/chat_api.py +++ b/src/agentlab/llm/chat_api.py @@ -12,6 +12,7 @@ HuggingFaceAPIChatModel, HuggingFaceURLChatModel, ) +from agentlab.llm.tracking import OpenRouterChatModel if TYPE_CHECKING: from langchain_core.language_models.chat_models import BaseChatModel @@ -86,7 +87,7 @@ class OpenRouterModelArgs(BaseModelArgs): model.""" def make_model(self): - return ChatOpenRouter( + return OpenRouterChatModel( model_name=self.model_name, temperature=self.temperature, max_tokens=self.max_new_tokens, diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py new file mode 100644 index 00000000..13b33455 --- /dev/null +++ b/src/agentlab/llm/tracking.py @@ -0,0 +1,104 @@ +import os +from contextlib import contextmanager +from typing import Any, List, Optional + +import requests +from langchain.schema import AIMessage, BaseMessage +from openai import OpenAI + +from agentlab.llm.langchain_utils import _convert_messages_to_dict + + +class LLMTracker: + def __init__(self): + self.input_tokens = 0 + self.output_tokens = 0 + self.cost = 0 + + def __call__(self, input_tokens: int, output_tokens: int, cost: float): + self.input_tokens += input_tokens + self.output_tokens += output_tokens + self.cost += cost + + @property + def stats(self): + return { + "input_tokens": self.input_tokens, + "output_tokens": self.output_tokens, + "cost": self.cost, + } + + +@contextmanager +def set_tracker(tracker: LLMTracker): + global current_tracker + previous_tracker = globals().get("current_tracker", None) + current_tracker = tracker + yield + current_tracker = previous_tracker + + +def get_action_decorator(get_action): + def wrapper(self, obs): + tracker = LLMTracker() + with set_tracker(tracker): + action, agent_info = get_action(self, obs) + agent_info.get("stats").update(tracker.stats) + return action, agent_info + + return wrapper + + +class OpenRouterChatModel: + def __init__( + self, + model_name, + openrouter_api_key=None, + openrouter_api_base="https://openrouter.ai/api/v1", + temperature=0.5, + max_tokens=100, + ): + self.model_name = model_name + self.openrouter_api_key = openrouter_api_key + self.openrouter_api_base = openrouter_api_base + self.temperature = temperature + self.max_tokens = max_tokens + + openrouter_api_key = openrouter_api_key or os.getenv("OPENROUTER_API_KEY") + + # query api to get model metadata + url = "https://openrouter.ai/api/v1/models" + headers = {"Authorization": f"Bearer {openrouter_api_key}"} + response = requests.get(url, headers=headers) + + if response.status_code != 200: + raise ValueError("Failed to get model metadata") + + model_metadata = response.json() + pricings = {model["id"]: model["pricing"] for model in model_metadata["data"]} + + self.input_cost = float(pricings[model_name]["prompt"]) + self.output_cost = float(pricings[model_name]["completion"]) + + self.client = OpenAI( + base_url=openrouter_api_base, + api_key=openrouter_api_key, + ) + + def __call__(self, messages: List[BaseMessage]) -> str: + messages_formated = _convert_messages_to_dict(messages) + completion = self.client.chat.completions.create( + model=self.model_name, messages=messages_formated + ) + input_tokens = completion.usage.prompt_tokens + output_tokens = completion.usage.completion_tokens + cost = input_tokens * self.input_cost + output_tokens * self.output_cost + + global current_tracker + if "current_tracker" in globals() and isinstance(current_tracker, LLMTracker): + current_tracker(input_tokens, output_tokens, cost) + + return AIMessage(content=completion.choices[0].message.content) + + def invoke(self, messages: List[BaseMessage]) -> AIMessage: + return self(messages) From ca50598a2d1efd3719f059745dc07c6dc815805f Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Wed, 18 Sep 2024 14:47:14 -0400 Subject: [PATCH 03/15] adding openai pricing request --- src/agentlab/llm/chat_api.py | 4 +- src/agentlab/llm/tracking.py | 159 +++++++++++++++++++++++++++++------ 2 files changed, 134 insertions(+), 29 deletions(-) diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index 1246081a..035838aa 100644 --- a/src/agentlab/llm/chat_api.py +++ b/src/agentlab/llm/chat_api.py @@ -12,7 +12,7 @@ HuggingFaceAPIChatModel, HuggingFaceURLChatModel, ) -from agentlab.llm.tracking import OpenRouterChatModel +from agentlab.llm.tracking import OpenAIChatModel, OpenRouterChatModel if TYPE_CHECKING: from langchain_core.language_models.chat_models import BaseChatModel @@ -100,7 +100,7 @@ class OpenAIModelArgs(BaseModelArgs): model.""" def make_model(self): - return ChatOpenAI( + return OpenAIChatModel( model_name=self.model_name, temperature=self.temperature, max_tokens=self.max_new_tokens, diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 13b33455..fcf92ae6 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -1,10 +1,12 @@ +import ast import os +from abc import ABC, abstractmethod from contextlib import contextmanager from typing import Any, List, Optional import requests from langchain.schema import AIMessage, BaseMessage -from openai import OpenAI +from openai import AzureOpenAI, OpenAI from agentlab.llm.langchain_utils import _convert_messages_to_dict @@ -49,46 +51,75 @@ def wrapper(self, obs): return wrapper -class OpenRouterChatModel: - def __init__( - self, - model_name, - openrouter_api_key=None, - openrouter_api_base="https://openrouter.ai/api/v1", - temperature=0.5, - max_tokens=100, - ): - self.model_name = model_name - self.openrouter_api_key = openrouter_api_key - self.openrouter_api_base = openrouter_api_base - self.temperature = temperature - self.max_tokens = max_tokens - - openrouter_api_key = openrouter_api_key or os.getenv("OPENROUTER_API_KEY") - +def get_pricing(api: str = "openrouter", api_key: str = None): + if api == "openrouter": + assert api_key, "OpenRouter API key is required" # query api to get model metadata url = "https://openrouter.ai/api/v1/models" - headers = {"Authorization": f"Bearer {openrouter_api_key}"} + headers = {"Authorization": f"Bearer {api_key}"} response = requests.get(url, headers=headers) if response.status_code != 200: raise ValueError("Failed to get model metadata") model_metadata = response.json() - pricings = {model["id"]: model["pricing"] for model in model_metadata["data"]} + return { + model["id"]: {k: float(v) for k, v in model["pricing"].items()} + for model in model_metadata["data"] + } + elif api == "openai": + url = "https://raw.githubusercontent.com/langchain-ai/langchain/master/libs/community/langchain_community/callbacks/openai_info.py" + response = requests.get(url) + + if response.status_code == 200: + content = response.text + tree = ast.parse(content) + cost_dict = None + for node in tree.body: + if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name): + if node.targets[0].id == "MODEL_COST_PER_1K_TOKENS": + cost_dict = ast.literal_eval(node.value) + break + if cost_dict: + cost_dict = {k: v / 1000 for k, v in cost_dict.items()} + res = {} + for k in cost_dict: + if k.endswith("-completion"): + continue + prompt_key = k + completion_key = k + "-completion" + if completion_key in cost_dict: + res[k] = { + "prompt": cost_dict[prompt_key], + "completion": cost_dict[completion_key], + } + return res + else: + raise ValueError("Cost dictionary not found.") + else: + raise ValueError(f"Failed to retrieve the file. Status code: {response.status_code}") + + +class ChatModel(ABC): + + @abstractmethod + def __init__(self, model_name, api_key=None, temperature=0.5, max_tokens=100): + self.model_name = model_name + self.temperature = temperature + self.max_tokens = max_tokens - self.input_cost = float(pricings[model_name]["prompt"]) - self.output_cost = float(pricings[model_name]["completion"]) + self.client = OpenAI() - self.client = OpenAI( - base_url=openrouter_api_base, - api_key=openrouter_api_key, - ) + self.input_cost = 0.0 + self.output_cost = 0.0 def __call__(self, messages: List[BaseMessage]) -> str: messages_formated = _convert_messages_to_dict(messages) completion = self.client.chat.completions.create( - model=self.model_name, messages=messages_formated + model=self.model_name, + messages=messages_formated, + temperature=self.temperature, + max_tokens=self.max_tokens, ) input_tokens = completion.usage.prompt_tokens output_tokens = completion.usage.completion_tokens @@ -102,3 +133,77 @@ def __call__(self, messages: List[BaseMessage]) -> str: def invoke(self, messages: List[BaseMessage]) -> AIMessage: return self(messages) + + +class OpenRouterChatModel(ChatModel): + def __init__( + self, + model_name, + api_key=None, + temperature=0.5, + max_tokens=100, + ): + self.model_name = model_name + self.temperature = temperature + self.max_tokens = max_tokens + + api_key = api_key or os.getenv("OPENROUTER_API_KEY") + + pricings = get_pricing(api="openrouter", api_key=api_key) + + self.input_cost = pricings[model_name]["prompt"] + self.output_cost = pricings[model_name]["completion"] + + self.client = OpenAI( + base_url="https://openrouter.ai/api/v1", + api_key=api_key, + ) + + +class OpenAIChatModel(ChatModel): + def __init__( + self, + model_name, + api_key=None, + temperature=0.5, + max_tokens=100, + ): + self.model_name = model_name + self.temperature = temperature + self.max_tokens = max_tokens + + api_key = api_key or os.getenv("OPENAI_API_KEY") + + pricings = get_pricing(api="openai") + + self.input_cost = float(pricings[model_name]["prompt"]) + self.output_cost = float(pricings[model_name]["completion"]) + + self.client = OpenAI( + api_key=api_key, + ) + + +class AzureChatModel(ChatModel): + def __init__( + self, + model_name, + api_key=None, + endpoint=None, + temperature=0.5, + max_tokens=100, + ): + self.model_name = model_name + self.temperature = temperature + self.max_tokens = max_tokens + + api_key = api_key or os.getenv("OPENAI_API_KEY") + + pricings = get_pricing(api="openai") + + self.input_cost = float(pricings[model_name]["prompt"]) + self.output_cost = float(pricings[model_name]["completion"]) + + self.client = AzureOpenAI( + api_key=api_key, azure_endpoint=endpoint, api_version="2024-02-01" + ) From 0cde22b6b76f2c82ebfffe24bfb17127f0d334b3 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Thu, 19 Sep 2024 10:56:53 -0400 Subject: [PATCH 04/15] switching back to langchain community for openai pricing --- src/agentlab/llm/tracking.py | 45 ++++++++++++------------------------ 1 file changed, 15 insertions(+), 30 deletions(-) diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index fcf92ae6..32777904 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -6,6 +6,7 @@ import requests from langchain.schema import AIMessage, BaseMessage +from langchain_community.callbacks.openai_info import MODEL_COST_PER_1K_TOKENS from openai import AzureOpenAI, OpenAI from agentlab.llm.langchain_utils import _convert_messages_to_dict @@ -68,36 +69,20 @@ def get_pricing(api: str = "openrouter", api_key: str = None): for model in model_metadata["data"] } elif api == "openai": - url = "https://raw.githubusercontent.com/langchain-ai/langchain/master/libs/community/langchain_community/callbacks/openai_info.py" - response = requests.get(url) - - if response.status_code == 200: - content = response.text - tree = ast.parse(content) - cost_dict = None - for node in tree.body: - if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name): - if node.targets[0].id == "MODEL_COST_PER_1K_TOKENS": - cost_dict = ast.literal_eval(node.value) - break - if cost_dict: - cost_dict = {k: v / 1000 for k, v in cost_dict.items()} - res = {} - for k in cost_dict: - if k.endswith("-completion"): - continue - prompt_key = k - completion_key = k + "-completion" - if completion_key in cost_dict: - res[k] = { - "prompt": cost_dict[prompt_key], - "completion": cost_dict[completion_key], - } - return res - else: - raise ValueError("Cost dictionary not found.") - else: - raise ValueError(f"Failed to retrieve the file. Status code: {response.status_code}") + cost_dict = MODEL_COST_PER_1K_TOKENS + cost_dict = {k: v / 1000 for k, v in cost_dict.items()} + res = {} + for k in cost_dict: + if k.endswith("-completion"): + continue + prompt_key = k + completion_key = k + "-completion" + if completion_key in cost_dict: + res[k] = { + "prompt": cost_dict[prompt_key], + "completion": cost_dict[completion_key], + } + return res class ChatModel(ABC): From 060e1e74b90bef6fcd18a6dc32a4d1307abf5430 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Thu, 19 Sep 2024 11:02:53 -0400 Subject: [PATCH 05/15] renaming launch_command.py to main.py --- launch_command.py => main.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) rename launch_command.py => main.py (93%) diff --git a/launch_command.py b/main.py similarity index 93% rename from launch_command.py rename to main.py index 198bbe07..bfb696c2 100644 --- a/launch_command.py +++ b/main.py @@ -16,7 +16,7 @@ logging.getLogger().setLevel(logging.INFO) # choose your agent or provide a new agent -agent_args = [AGENT_CUSTOM, AGENT_4o_MINI] +agent_args = [AGENT_4o_MINI] # agent = AGENT_4o @@ -48,4 +48,5 @@ # run the experiments -run_experiments(n_jobs, exp_args_list, study_dir) +if __name__ == "__main__": + run_experiments(n_jobs, exp_args_list, study_dir) From 600cfcacf17c1703b7855d93d3d95cc786f28121 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Thu, 19 Sep 2024 11:09:50 -0400 Subject: [PATCH 06/15] typo --- main.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/main.py b/main.py index 5bdef316..e6234637 100644 --- a/main.py +++ b/main.py @@ -7,13 +7,11 @@ import logging -from agentlab.agents.generic_agent import (AGENT_CUSTOM, RANDOM_SEARCH_AGENT, - AGENT_4o, AGENT_4o_MINI) +from agentlab.agents.generic_agent import AGENT_CUSTOM, RANDOM_SEARCH_AGENT, AGENT_4o, AGENT_4o_MINI from agentlab.analyze.inspect_results import get_most_recent_folder from agentlab.experiments import study_generators from agentlab.experiments.exp_utils import RESULTS_DIR -from agentlab.experiments.launch_exp import (make_study_dir, relaunch_study, - run_experiments) +from agentlab.experiments.launch_exp import make_study_dir, relaunch_study, run_experiments logging.getLogger().setLevel(logging.INFO) @@ -50,5 +48,4 @@ # run the experiments if __name__ == "__main__": - run_experiments(n_jobs, exp_args_list, study_dir)if __name__ == "__main__": - run_experiments(n_jobs, exp_args_list, study_dir) \ No newline at end of file + run_experiments(n_jobs, exp_args_list, study_dir) From 708bde5719c738e1b62a3489f8d47fb37aa0fa69 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Fri, 20 Sep 2024 10:30:04 -0400 Subject: [PATCH 07/15] tracking is thread safe and mostly tested --- src/agentlab/llm/tracking.py | 52 +++++++++++----- tests/llm/test_tracking.py | 112 +++++++++++++++++++++++++++++++++++ 2 files changed, 149 insertions(+), 15 deletions(-) create mode 100644 tests/llm/test_tracking.py diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 32777904..d964c156 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -1,5 +1,6 @@ import ast import os +import threading from abc import ABC, abstractmethod from contextlib import contextmanager from typing import Any, List, Optional @@ -11,12 +12,14 @@ from agentlab.llm.langchain_utils import _convert_messages_to_dict +TRACKER = threading.local() + class LLMTracker: def __init__(self): self.input_tokens = 0 self.output_tokens = 0 - self.cost = 0 + self.cost = 0.0 def __call__(self, input_tokens: int, output_tokens: int, cost: float): self.input_tokens += input_tokens @@ -31,20 +34,33 @@ def stats(self): "cost": self.cost, } + def add_tracker(self, tracker: "LLMTracker"): + self(tracker.input_tokens, tracker.output_tokens, tracker.cost) + + def __repr__(self): + return f"LLMTracker(input_tokens={self.input_tokens}, output_tokens={self.output_tokens}, cost={self.cost})" + @contextmanager -def set_tracker(tracker: LLMTracker): - global current_tracker - previous_tracker = globals().get("current_tracker", None) - current_tracker = tracker - yield - current_tracker = previous_tracker +def set_tracker(): + global TRACKER + if not hasattr(TRACKER, "instance"): + TRACKER.instance = None + previous_tracker = TRACKER.instance # type: LLMTracker + TRACKER.instance = LLMTracker() + try: + yield TRACKER.instance + finally: + # If there was a previous tracker, add the current one to it + if isinstance(previous_tracker, LLMTracker): + previous_tracker.add_tracker(TRACKER.instance) + # Restore the previous tracker + TRACKER.instance = previous_tracker def get_action_decorator(get_action): def wrapper(self, obs): - tracker = LLMTracker() - with set_tracker(tracker): + with set_tracker() as tracker: action, agent_info = get_action(self, obs) agent_info.get("stats").update(tracker.stats) return action, agent_info @@ -110,9 +126,8 @@ def __call__(self, messages: List[BaseMessage]) -> str: output_tokens = completion.usage.completion_tokens cost = input_tokens * self.input_cost + output_tokens * self.output_cost - global current_tracker - if "current_tracker" in globals() and isinstance(current_tracker, LLMTracker): - current_tracker(input_tokens, output_tokens, cost) + if isinstance(TRACKER.instance, LLMTracker): + TRACKER.instance(input_tokens, output_tokens, cost) return AIMessage(content=completion.choices[0].message.content) @@ -174,7 +189,7 @@ def __init__( self, model_name, api_key=None, - endpoint=None, + deployment_name=None, temperature=0.5, max_tokens=100, ): @@ -182,7 +197,11 @@ def __init__( self.temperature = temperature self.max_tokens = max_tokens - api_key = api_key or os.getenv("OPENAI_API_KEY") + api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY") + + # AZURE_OPENAI_ENDPOINT has to be defined in the environment + endpoint = os.getenv("AZURE_OPENAI_ENDPOINT") + assert endpoint, "AZURE_OPENAI_ENDPOINT has to be defined in the environment" pricings = get_pricing(api="openai") @@ -190,5 +209,8 @@ def __init__( self.output_cost = float(pricings[model_name]["completion"]) self.client = AzureOpenAI( - api_key=api_key, azure_endpoint=endpoint, api_version="2024-02-01" + api_key=api_key, + azure_deployment=deployment_name, + azure_endpoint=endpoint, + api_version="2024-02-01", ) diff --git a/tests/llm/test_tracking.py b/tests/llm/test_tracking.py new file mode 100644 index 00000000..5f12cd68 --- /dev/null +++ b/tests/llm/test_tracking.py @@ -0,0 +1,112 @@ +import os +import time +from functools import partial + +import pytest + +import agentlab.llm.tracking as tracking + + +def test_get_action_decorator(): + action, agent_info = tracking.get_action_decorator(lambda x, y: call_llm())(None, None) + assert action == "action" + assert agent_info["stats"] == { + "input_tokens": 1, + "output_tokens": 1, + "cost": 1.0, + } + + +OPENROUTER_API_KEY_AVAILABLE = os.environ.get("OPENROUTER_API_KEY") is not None + +OPENROUTER_MODELS = ( + "openai/o1-mini-2024-09-12", + "openai/o1-preview-2024-09-12", + "openai/gpt-4o-2024-08-06", + "openai/gpt-4o-2024-05-13", + "anthropic/claude-3.5-sonnet:beta", + "anthropic/claude-3.5-sonnet", + "meta-llama/llama-3.1-405b-instruct", + "meta-llama/llama-3.1-70b-instruct", + "meta-llama/llama-3.1-8b-instruct", + "google/gemini-pro-1.5", + "qwen/qwen-2-vl-72b-instruct", +) + + +@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available") +def test_get_pricing_openrouter(): + pricing = tracking.get_pricing(api="openrouter", api_key=os.environ["OPENROUTER_API_KEY"]) + assert isinstance(pricing, dict) + assert all(isinstance(v, dict) for v in pricing.values()) + for model in OPENROUTER_MODELS: + assert model in pricing + assert isinstance(pricing[model], dict) + assert all(isinstance(v, float) for v in pricing[model].values()) + + +def test_get_pricing_openai(): + pricing = tracking.get_pricing(api="openai") + assert isinstance(pricing, dict) + assert all("prompt" in pricing[model] and "completion" in pricing[model] for model in pricing) + assert all(isinstance(pricing[model]["prompt"], float) for model in pricing) + assert all(isinstance(pricing[model]["completion"], float) for model in pricing) + + +def call_llm(): + if isinstance(tracking.TRACKER.instance, tracking.LLMTracker): + tracking.TRACKER.instance(1, 1, 1) + return "action", {"stats": {}} + + +def test_tracker(): + with tracking.set_tracker() as tracker: + _, _ = call_llm() + + assert tracker.stats["cost"] == 1 + + +def test_imbricate_trackers(): + with tracking.set_tracker() as tracker4: + with tracking.set_tracker() as tracker1: + _, _ = call_llm() + with tracking.set_tracker() as tracker3: + _, _ = call_llm() + _, _ = call_llm() + with tracking.set_tracker() as tracker1bis: + _, _ = call_llm() + + assert tracker1.stats["cost"] == 1 + assert tracker1bis.stats["cost"] == 1 + assert tracker3.stats["cost"] == 3 + assert tracker4.stats["cost"] == 4 + + +def test_threaded_trackers(): + """thread_2 occurs in the middle of thread_1, results should be separate.""" + import threading + + def thread_1(results=None): + with tracking.set_tracker() as tracker: + time.sleep(1) + _, _ = call_llm() + time.sleep(1) + results[0] = tracker.stats + + def thread_2(results=None): + time.sleep(1) + with tracking.set_tracker() as tracker: + _, _ = call_llm() + results[1] = tracker.stats + + results = [None] * 2 + threads = [ + threading.Thread(target=partial(thread_1, results=results)), + threading.Thread(target=partial(thread_2, results=results)), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert all(result["cost"] == 1 for result in results) From 7c209459e067ac76e21543fc2f72bf70db900c39 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Fri, 20 Sep 2024 10:48:38 -0400 Subject: [PATCH 08/15] added pricy tests for ChatModels --- tests/llm/test_tracking.py | 68 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/tests/llm/test_tracking.py b/tests/llm/test_tracking.py index 5f12cd68..a8141975 100644 --- a/tests/llm/test_tracking.py +++ b/tests/llm/test_tracking.py @@ -110,3 +110,71 @@ def thread_2(results=None): thread.join() assert all(result["cost"] == 1 for result in results) + + +OPENAI_API_KEY_AVAILABLE = os.environ.get("OPENAI_API_KEY") is not None + + +@pytest.mark.pricy +@pytest.mark.skipif(not OPENAI_API_KEY_AVAILABLE, reason="OpenAI API key is not available") +def test_openai_chat_model(): + chat_model = tracking.OpenAIChatModel("gpt-4o-mini") + assert chat_model.input_cost > 0 + assert chat_model.output_cost > 0 + + from langchain.schema import HumanMessage, SystemMessage + + messages = [ + SystemMessage(content="You are an helpful virtual assistant"), + HumanMessage(content="Give the third prime number"), + ] + with tracking.set_tracker() as tracker: + answer = chat_model.invoke(messages) + assert "5" in answer.content + assert tracker.stats["cost"] > 0 + + +AZURE_OPENAI_API_KEY_AVAILABLE = ( + os.environ.get("AZURE_OPENAI_API_KEY") is not None + and os.environ.get("AZURE_OPENAI_ENDPOINT") is not None +) + + +@pytest.mark.pricy +@pytest.mark.skipif( + not AZURE_OPENAI_API_KEY_AVAILABLE, reason="Azure OpenAI API key is not available" +) +def test_azure_chat_model(): + chat_model = tracking.AzureChatModel(model_name="gpt-35-turbo", deployment_name="gpt-35-turbo") + assert chat_model.input_cost > 0 + assert chat_model.output_cost > 0 + + from langchain.schema import HumanMessage, SystemMessage + + messages = [ + SystemMessage(content="You are an helpful virtual assistant"), + HumanMessage(content="Give the third prime number"), + ] + with tracking.set_tracker() as tracker: + answer = chat_model.invoke(messages) + assert "5" in answer.content + assert tracker.stats["cost"] > 0 + + +@pytest.mark.pricy +@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available") +def test_openrouter_chat_model(): + chat_model = tracking.OpenRouterChatModel("openai/gpt-4o-mini") + assert chat_model.input_cost > 0 + assert chat_model.output_cost > 0 + + from langchain.schema import HumanMessage, SystemMessage + + messages = [ + SystemMessage(content="You are an helpful virtual assistant"), + HumanMessage(content="Give the third prime number"), + ] + with tracking.set_tracker() as tracker: + answer = chat_model.invoke(messages) + assert "5" in answer.content + assert tracker.stats["cost"] > 0 From 18a45e0ba1361daafc0f83a8c670a3e95a27fad0 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Fri, 20 Sep 2024 10:51:44 -0400 Subject: [PATCH 09/15] separating get_pricing function --- src/agentlab/llm/tracking.py | 70 ++++++++++++++++++------------------ 1 file changed, 36 insertions(+), 34 deletions(-) diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index d964c156..54ed871d 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -68,37 +68,39 @@ def wrapper(self, obs): return wrapper -def get_pricing(api: str = "openrouter", api_key: str = None): - if api == "openrouter": - assert api_key, "OpenRouter API key is required" - # query api to get model metadata - url = "https://openrouter.ai/api/v1/models" - headers = {"Authorization": f"Bearer {api_key}"} - response = requests.get(url, headers=headers) - - if response.status_code != 200: - raise ValueError("Failed to get model metadata") - - model_metadata = response.json() - return { - model["id"]: {k: float(v) for k, v in model["pricing"].items()} - for model in model_metadata["data"] - } - elif api == "openai": - cost_dict = MODEL_COST_PER_1K_TOKENS - cost_dict = {k: v / 1000 for k, v in cost_dict.items()} - res = {} - for k in cost_dict: - if k.endswith("-completion"): - continue - prompt_key = k - completion_key = k + "-completion" - if completion_key in cost_dict: - res[k] = { - "prompt": cost_dict[prompt_key], - "completion": cost_dict[completion_key], - } - return res +def get_pricing_openrouter(): + api_key = os.getenv("OPENROUTER_API_KEY") + assert api_key, "OpenRouter API key is required" + # query api to get model metadata + url = "https://openrouter.ai/api/v1/models" + headers = {"Authorization": f"Bearer {api_key}"} + response = requests.get(url, headers=headers) + + if response.status_code != 200: + raise ValueError("Failed to get model metadata") + + model_metadata = response.json() + return { + model["id"]: {k: float(v) for k, v in model["pricing"].items()} + for model in model_metadata["data"] + } + + +def get_pricing_openai(): + cost_dict = MODEL_COST_PER_1K_TOKENS + cost_dict = {k: v / 1000 for k, v in cost_dict.items()} + res = {} + for k in cost_dict: + if k.endswith("-completion"): + continue + prompt_key = k + completion_key = k + "-completion" + if completion_key in cost_dict: + res[k] = { + "prompt": cost_dict[prompt_key], + "completion": cost_dict[completion_key], + } + return res class ChatModel(ABC): @@ -149,7 +151,7 @@ def __init__( api_key = api_key or os.getenv("OPENROUTER_API_KEY") - pricings = get_pricing(api="openrouter", api_key=api_key) + pricings = get_pricing_openrouter() self.input_cost = pricings[model_name]["prompt"] self.output_cost = pricings[model_name]["completion"] @@ -174,7 +176,7 @@ def __init__( api_key = api_key or os.getenv("OPENAI_API_KEY") - pricings = get_pricing(api="openai") + pricings = get_pricing_openai() self.input_cost = float(pricings[model_name]["prompt"]) self.output_cost = float(pricings[model_name]["completion"]) @@ -203,7 +205,7 @@ def __init__( endpoint = os.getenv("AZURE_OPENAI_ENDPOINT") assert endpoint, "AZURE_OPENAI_ENDPOINT has to be defined in the environment" - pricings = get_pricing(api="openai") + pricings = get_pricing_openai() self.input_cost = float(pricings[model_name]["prompt"]) self.output_cost = float(pricings[model_name]["completion"]) From 9d12cdfa5136095240a90c43812a9b0495bf461a Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Fri, 20 Sep 2024 10:57:11 -0400 Subject: [PATCH 10/15] updating function names --- src/agentlab/llm/tracking.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 54ed871d..8a530e73 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -3,7 +3,6 @@ import threading from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import Any, List, Optional import requests from langchain.schema import AIMessage, BaseMessage @@ -116,11 +115,10 @@ def __init__(self, model_name, api_key=None, temperature=0.5, max_tokens=100): self.input_cost = 0.0 self.output_cost = 0.0 - def __call__(self, messages: List[BaseMessage]) -> str: - messages_formated = _convert_messages_to_dict(messages) + def __call__(self, messages: list[dict]) -> dict: completion = self.client.chat.completions.create( model=self.model_name, - messages=messages_formated, + messages=messages, temperature=self.temperature, max_tokens=self.max_tokens, ) @@ -131,9 +129,9 @@ def __call__(self, messages: List[BaseMessage]) -> str: if isinstance(TRACKER.instance, LLMTracker): TRACKER.instance(input_tokens, output_tokens, cost) - return AIMessage(content=completion.choices[0].message.content) + return dict(role="assistant", content=completion.choices[0].message.content) - def invoke(self, messages: List[BaseMessage]) -> AIMessage: + def invoke(self, messages: list[dict]) -> dict: return self(messages) From 8e5e5f939030df6ba8147e823a6cda44d0e7ad0e Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Fri, 20 Sep 2024 11:01:20 -0400 Subject: [PATCH 11/15] updating function names --- tests/llm/test_tracking.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/llm/test_tracking.py b/tests/llm/test_tracking.py index a8141975..4672cacd 100644 --- a/tests/llm/test_tracking.py +++ b/tests/llm/test_tracking.py @@ -36,7 +36,7 @@ def test_get_action_decorator(): @pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available") def test_get_pricing_openrouter(): - pricing = tracking.get_pricing(api="openrouter", api_key=os.environ["OPENROUTER_API_KEY"]) + pricing = tracking.get_pricing_openrouter() assert isinstance(pricing, dict) assert all(isinstance(v, dict) for v in pricing.values()) for model in OPENROUTER_MODELS: @@ -46,7 +46,7 @@ def test_get_pricing_openrouter(): def test_get_pricing_openai(): - pricing = tracking.get_pricing(api="openai") + pricing = tracking.get_pricing_openai() assert isinstance(pricing, dict) assert all("prompt" in pricing[model] and "completion" in pricing[model] for model in pricing) assert all(isinstance(pricing[model]["prompt"], float) for model in pricing) From 17e8ff8dbcf2217b7871e09d89f5d03621ec9240 Mon Sep 17 00:00:00 2001 From: Thibault Le Sellier de Chezelles Date: Fri, 20 Sep 2024 11:21:35 -0400 Subject: [PATCH 12/15] ciao retry_parallel --- src/agentlab/llm/llm_utils.py | 63 ----------------------------------- tests/llm/test_llm_utils.py | 34 ------------------- 2 files changed, 97 deletions(-) diff --git a/src/agentlab/llm/llm_utils.py b/src/agentlab/llm/llm_utils.py index 1a8d8b70..5c16ef0d 100644 --- a/src/agentlab/llm/llm_utils.py +++ b/src/agentlab/llm/llm_utils.py @@ -177,69 +177,6 @@ def retry_raise( raise RetryError(f"Could not parse a valid value after {n_retry} retries.") -def retry_parallel(chat: "BaseChatModel", messages, n_retry, parser): - """Retry querying the chat models with the response from the parser until it returns a valid value. - - It will stop after `n_retry`. It assuemes that chat will generate n_parallel answers for each message. - The best answer is selected according to the score returned by the parser. If no answer is valid, the - it will retry with the best answer so far and append to the chat the retry message. If there is a - single parallel generation, it behaves like retry. - - This function is, in principle, more robust than retry. The speed and cost overhead is minimal with - the prompt is large and the length of the generated message is small. - - Args: - chat (BaseChatModel): a langchain BaseChatModel taking a list of messages and - returning a list of answers. - messages (list): the list of messages so far. - n_retry (int): the maximum number of sequential retries. - parser (function): a function taking a message and returning a tuple - with the following fields: - value : the parsed value, - valid : a boolean indicating if the value is valid, - retry_message : a message to send to the chat if the value is not valid - - Returns: - dict: the parsed value, with a string at key "action". - - Raises: - ValueError: if the parser could not parse a valid value after n_retry retries. - BadRequestError: if the message is too long - """ - - for i in range(n_retry): - try: - answers = chat.generate([messages]).generations[0] # chat.n parallel completions - except BadRequestError as e: - # most likely, the added messages triggered a message too long error - # we thus retry without the last two messages - if i == 0: - raise e - msg = f"BadRequestError, most likely the message is too long retrying with previous query." - warn(msg) - messages = messages[:-2] - answers = chat.generate([messages]).generations[0] - - values, valids, retry_messages, scores = zip( - *[parser(answer.message.content) for answer in answers] - ) - idx = np.argmax(scores) - value = values[idx] - valid = valids[idx] - retry_message = retry_messages[idx] - answer = answers[idx].message - - if valid: - return value - - msg = f"Query failed. Retrying {i+1}/{n_retry}.\n[LLM]:\n{answer.content}\n[User]:\n{retry_message}" - warn(msg) - messages.append(answer) # already of type AIMessage - messages.append(SystemMessage(content=retry_message)) - - raise ValueError(f"Could not parse a valid value after {n_retry} retries.") - - def truncate_tokens(text, max_tokens=8000, start=0, model_name="gpt-4"): """Use tiktoken to truncate a text to a maximum number of tokens.""" enc = tiktoken.encoding_for_model(model_name) diff --git a/tests/llm/test_llm_utils.py b/tests/llm/test_llm_utils.py index 1bdbdacb..41dca395 100644 --- a/tests/llm/test_llm_utils.py +++ b/tests/llm/test_llm_utils.py @@ -94,40 +94,6 @@ def test_compress_string(): assert compressed_text == expected_output -@pytest.mark.pricy -def test_retry_parallel(): - chat = AzureChatOpenAI( - model_name="gpt-35-turbo", azure_deployment="gpt-35-turbo", temperature=0.2, n=3 - ) - prompt = """List primes from 1 to 10.""" - messages = [ - SystemMessage(content=prompt), - ] - - global n_call - n_call = 0 - - def parser(message): - global n_call - n_call += 1 - - if n_call <= 3: # First 3 calls, just answer the new prompt - return ( - None, - False, - "I changed my mind. List primes up to 15. Just answer the json list, nothing else.", - 0, - ) - elif n_call == 5: - return "success", True, "", 10 - else: - return "bad", True, "", 1 + np.random.rand() - - value = llm_utils.retry_parallel(chat, messages, parser=parser, n_retry=2) - assert value == "success" - assert n_call == 6 # 2*3 calls - - # Mock ChatOpenAI class class MockChatOpenAI: def invoke(self, messages): From d62357f4e60ce04175ac300e947d741f797e1677 Mon Sep 17 00:00:00 2001 From: Thibault LSDC <78021491+ThibaultLSDC@users.noreply.github.com> Date: Fri, 20 Sep 2024 13:01:30 -0400 Subject: [PATCH 13/15] fixing case when the context isnt used --- src/agentlab/llm/tracking.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 8a530e73..554def65 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -126,7 +126,7 @@ def __call__(self, messages: list[dict]) -> dict: output_tokens = completion.usage.completion_tokens cost = input_tokens * self.input_cost + output_tokens * self.output_cost - if isinstance(TRACKER.instance, LLMTracker): + if hasattr(TRACKER, "instance") and isinstance(TRACKER.instance, LLMTracker): TRACKER.instance(input_tokens, output_tokens, cost) return dict(role="assistant", content=completion.choices[0].message.content) From e3808f9c32d02c162c2b9829e8d48736fa111005 Mon Sep 17 00:00:00 2001 From: ThibaultLSDC Date: Fri, 20 Sep 2024 15:45:08 -0400 Subject: [PATCH 14/15] moving ChatModels to chat_api --- src/agentlab/llm/chat_api.py | 119 ++++++++++++++++++++++++++++++++++- src/agentlab/llm/tracking.py | 114 --------------------------------- 2 files changed, 116 insertions(+), 117 deletions(-) diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index 035838aa..8372111c 100644 --- a/src/agentlab/llm/chat_api.py +++ b/src/agentlab/llm/chat_api.py @@ -5,14 +5,13 @@ from typing import TYPE_CHECKING from langchain.schema import AIMessage -from langchain_openai import AzureChatOpenAI, ChatOpenAI +import agentlab.llm.tracking as tracking from agentlab.llm.langchain_utils import ( ChatOpenRouter, HuggingFaceAPIChatModel, HuggingFaceURLChatModel, ) -from agentlab.llm.tracking import OpenAIChatModel, OpenRouterChatModel if TYPE_CHECKING: from langchain_core.language_models.chat_models import BaseChatModel @@ -127,7 +126,7 @@ class AzureModelArgs(BaseModelArgs): deployment_name: str = None def make_model(self): - return AzureChatOpenAI( + return AzureChatModel( model_name=self.model_name, temperature=self.temperature, max_tokens=self.max_new_tokens, @@ -195,3 +194,117 @@ def __post_init__(self): def make_model(self): pass + + +class ChatModel(ABC): + + @abstractmethod + def __init__(self, model_name, api_key=None, temperature=0.5, max_tokens=100): + self.model_name = model_name + self.temperature = temperature + self.max_tokens = max_tokens + + self.client = tracking.OpenAI() + + self.input_cost = 0.0 + self.output_cost = 0.0 + + def __call__(self, messages: list[dict]) -> dict: + completion = self.client.chat.completions.create( + model=self.model_name, + messages=messages, + temperature=self.temperature, + max_tokens=self.max_tokens, + ) + input_tokens = completion.usage.prompt_tokens + output_tokens = completion.usage.completion_tokens + cost = input_tokens * self.input_cost + output_tokens * self.output_cost + + if isinstance(tracking.TRACKER.instance, tracking.LLMTracker): + tracking.TRACKER.instance(input_tokens, output_tokens, cost) + + return dict(role="assistant", content=completion.choices[0].message.content) + + def invoke(self, messages: list[dict]) -> dict: + return self(messages) + + +class OpenRouterChatModel(ChatModel): + def __init__( + self, + model_name, + api_key=None, + temperature=0.5, + max_tokens=100, + ): + self.model_name = model_name + self.temperature = temperature + self.max_tokens = max_tokens + + api_key = api_key or os.getenv("OPENROUTER_API_KEY") + + pricings = tracking.get_pricing_openrouter() + + self.input_cost = pricings[model_name]["prompt"] + self.output_cost = pricings[model_name]["completion"] + + self.client = tracking.OpenAI( + base_url="https://openrouter.ai/api/v1", + api_key=api_key, + ) + + +class OpenAIChatModel(ChatModel): + def __init__( + self, + model_name, + api_key=None, + temperature=0.5, + max_tokens=100, + ): + self.model_name = model_name + self.temperature = temperature + self.max_tokens = max_tokens + + api_key = api_key or os.getenv("OPENAI_API_KEY") + + pricings = tracking.get_pricing_openai() + + self.input_cost = float(pricings[model_name]["prompt"]) + self.output_cost = float(pricings[model_name]["completion"]) + + self.client = tracking.OpenAI( + api_key=api_key, + ) + + +class AzureChatModel(ChatModel): + def __init__( + self, + model_name, + api_key=None, + deployment_name=None, + temperature=0.5, + max_tokens=100, + ): + self.model_name = model_name + self.temperature = temperature + self.max_tokens = max_tokens + + api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY") + + # AZURE_OPENAI_ENDPOINT has to be defined in the environment + endpoint = os.getenv("AZURE_OPENAI_ENDPOINT") + assert endpoint, "AZURE_OPENAI_ENDPOINT has to be defined in the environment" + + pricings = tracking.get_pricing_openai() + + self.input_cost = float(pricings[model_name]["prompt"]) + self.output_cost = float(pricings[model_name]["completion"]) + + self.client = tracking.AzureOpenAI( + api_key=api_key, + azure_deployment=deployment_name, + azure_endpoint=endpoint, + api_version="2024-02-01", + ) diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 8a530e73..9cc0e0bf 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -100,117 +100,3 @@ def get_pricing_openai(): "completion": cost_dict[completion_key], } return res - - -class ChatModel(ABC): - - @abstractmethod - def __init__(self, model_name, api_key=None, temperature=0.5, max_tokens=100): - self.model_name = model_name - self.temperature = temperature - self.max_tokens = max_tokens - - self.client = OpenAI() - - self.input_cost = 0.0 - self.output_cost = 0.0 - - def __call__(self, messages: list[dict]) -> dict: - completion = self.client.chat.completions.create( - model=self.model_name, - messages=messages, - temperature=self.temperature, - max_tokens=self.max_tokens, - ) - input_tokens = completion.usage.prompt_tokens - output_tokens = completion.usage.completion_tokens - cost = input_tokens * self.input_cost + output_tokens * self.output_cost - - if isinstance(TRACKER.instance, LLMTracker): - TRACKER.instance(input_tokens, output_tokens, cost) - - return dict(role="assistant", content=completion.choices[0].message.content) - - def invoke(self, messages: list[dict]) -> dict: - return self(messages) - - -class OpenRouterChatModel(ChatModel): - def __init__( - self, - model_name, - api_key=None, - temperature=0.5, - max_tokens=100, - ): - self.model_name = model_name - self.temperature = temperature - self.max_tokens = max_tokens - - api_key = api_key or os.getenv("OPENROUTER_API_KEY") - - pricings = get_pricing_openrouter() - - self.input_cost = pricings[model_name]["prompt"] - self.output_cost = pricings[model_name]["completion"] - - self.client = OpenAI( - base_url="https://openrouter.ai/api/v1", - api_key=api_key, - ) - - -class OpenAIChatModel(ChatModel): - def __init__( - self, - model_name, - api_key=None, - temperature=0.5, - max_tokens=100, - ): - self.model_name = model_name - self.temperature = temperature - self.max_tokens = max_tokens - - api_key = api_key or os.getenv("OPENAI_API_KEY") - - pricings = get_pricing_openai() - - self.input_cost = float(pricings[model_name]["prompt"]) - self.output_cost = float(pricings[model_name]["completion"]) - - self.client = OpenAI( - api_key=api_key, - ) - - -class AzureChatModel(ChatModel): - def __init__( - self, - model_name, - api_key=None, - deployment_name=None, - temperature=0.5, - max_tokens=100, - ): - self.model_name = model_name - self.temperature = temperature - self.max_tokens = max_tokens - - api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY") - - # AZURE_OPENAI_ENDPOINT has to be defined in the environment - endpoint = os.getenv("AZURE_OPENAI_ENDPOINT") - assert endpoint, "AZURE_OPENAI_ENDPOINT has to be defined in the environment" - - pricings = get_pricing_openai() - - self.input_cost = float(pricings[model_name]["prompt"]) - self.output_cost = float(pricings[model_name]["completion"]) - - self.client = AzureOpenAI( - api_key=api_key, - azure_deployment=deployment_name, - azure_endpoint=endpoint, - api_version="2024-02-01", - ) From cfe01a0c228cef8c82d7d10cdaebade5798c2fe3 Mon Sep 17 00:00:00 2001 From: ThibaultLSDC Date: Fri, 20 Sep 2024 15:45:41 -0400 Subject: [PATCH 15/15] renaming get_action decorator --- src/agentlab/agents/generic_agent/generic_agent.py | 4 ++-- src/agentlab/agents/most_basic_agent/most_basic_agent.py | 2 ++ src/agentlab/llm/chat_api.py | 6 ++++-- src/agentlab/llm/tracking.py | 2 +- tests/llm/test_tracking.py | 9 +++++---- 5 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/agentlab/agents/generic_agent/generic_agent.py b/src/agentlab/agents/generic_agent/generic_agent.py index eb3b0d87..e2696923 100644 --- a/src/agentlab/agents/generic_agent/generic_agent.py +++ b/src/agentlab/agents/generic_agent/generic_agent.py @@ -10,7 +10,7 @@ from agentlab.agents.utils import openai_monitored_agent from agentlab.llm.chat_api import BaseModelArgs from agentlab.llm.llm_utils import RetryError, retry_raise -from agentlab.llm.tracking import get_action_decorator +from agentlab.llm.tracking import cost_tracker_decorator from .generic_agent_prompt import GenericPromptFlags, MainPrompt @@ -66,7 +66,7 @@ def __init__( def obs_preprocessor(self, obs: dict) -> dict: return self._obs_preprocessor(obs) - @get_action_decorator + @cost_tracker_decorator def get_action(self, obs): self.obs_history.append(obs) diff --git a/src/agentlab/agents/most_basic_agent/most_basic_agent.py b/src/agentlab/agents/most_basic_agent/most_basic_agent.py index 0c8d7023..4b7aaaff 100644 --- a/src/agentlab/agents/most_basic_agent/most_basic_agent.py +++ b/src/agentlab/agents/most_basic_agent/most_basic_agent.py @@ -11,6 +11,7 @@ from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT from agentlab.llm.llm_utils import ParseError, extract_code_blocks, retry_raise +from agentlab.llm.tracking import cost_tracker_decorator if TYPE_CHECKING: from agentlab.llm.chat_api import BaseModelArgs @@ -48,6 +49,7 @@ def __init__( self.action_set = HighLevelActionSet(["bid"], multiaction=False) + @cost_tracker_decorator def get_action(self, obs: Any) -> tuple[str, dict]: system_prompt = f""" You are a web assistant. diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index 8372111c..0fa02994 100644 --- a/src/agentlab/llm/chat_api.py +++ b/src/agentlab/llm/chat_api.py @@ -11,6 +11,7 @@ ChatOpenRouter, HuggingFaceAPIChatModel, HuggingFaceURLChatModel, + _convert_messages_to_dict, ) if TYPE_CHECKING: @@ -210,9 +211,10 @@ def __init__(self, model_name, api_key=None, temperature=0.5, max_tokens=100): self.output_cost = 0.0 def __call__(self, messages: list[dict]) -> dict: + messages_formatted = _convert_messages_to_dict(messages) completion = self.client.chat.completions.create( model=self.model_name, - messages=messages, + messages=messages_formatted, temperature=self.temperature, max_tokens=self.max_tokens, ) @@ -223,7 +225,7 @@ def __call__(self, messages: list[dict]) -> dict: if isinstance(tracking.TRACKER.instance, tracking.LLMTracker): tracking.TRACKER.instance(input_tokens, output_tokens, cost) - return dict(role="assistant", content=completion.choices[0].message.content) + return AIMessage(content=completion.choices[0].message.content) def invoke(self, messages: list[dict]) -> dict: return self(messages) diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 9cc0e0bf..9ce16785 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -57,7 +57,7 @@ def set_tracker(): TRACKER.instance = previous_tracker -def get_action_decorator(get_action): +def cost_tracker_decorator(get_action): def wrapper(self, obs): with set_tracker() as tracker: action, agent_info = get_action(self, obs) diff --git a/tests/llm/test_tracking.py b/tests/llm/test_tracking.py index 4672cacd..e883e458 100644 --- a/tests/llm/test_tracking.py +++ b/tests/llm/test_tracking.py @@ -5,10 +5,11 @@ import pytest import agentlab.llm.tracking as tracking +from agentlab.llm.chat_api import AzureChatModel, OpenAIChatModel, OpenRouterChatModel def test_get_action_decorator(): - action, agent_info = tracking.get_action_decorator(lambda x, y: call_llm())(None, None) + action, agent_info = tracking.cost_tracker_decorator(lambda x, y: call_llm())(None, None) assert action == "action" assert agent_info["stats"] == { "input_tokens": 1, @@ -118,7 +119,7 @@ def thread_2(results=None): @pytest.mark.pricy @pytest.mark.skipif(not OPENAI_API_KEY_AVAILABLE, reason="OpenAI API key is not available") def test_openai_chat_model(): - chat_model = tracking.OpenAIChatModel("gpt-4o-mini") + chat_model = OpenAIChatModel("gpt-4o-mini") assert chat_model.input_cost > 0 assert chat_model.output_cost > 0 @@ -145,7 +146,7 @@ def test_openai_chat_model(): not AZURE_OPENAI_API_KEY_AVAILABLE, reason="Azure OpenAI API key is not available" ) def test_azure_chat_model(): - chat_model = tracking.AzureChatModel(model_name="gpt-35-turbo", deployment_name="gpt-35-turbo") + chat_model = AzureChatModel(model_name="gpt-35-turbo", deployment_name="gpt-35-turbo") assert chat_model.input_cost > 0 assert chat_model.output_cost > 0 @@ -164,7 +165,7 @@ def test_azure_chat_model(): @pytest.mark.pricy @pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available") def test_openrouter_chat_model(): - chat_model = tracking.OpenRouterChatModel("openai/gpt-4o-mini") + chat_model = OpenRouterChatModel("openai/gpt-4o-mini") assert chat_model.input_cost > 0 assert chat_model.output_cost > 0