Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ dependencies = [
"browsergym>=0.7.1",
"joblib>=1.2.0",
"openai>=1.7,<2",
"langchain_community",
"tiktoken",
"huggingface_hub",
"contexttimer",
Expand Down
13 changes: 10 additions & 3 deletions src/agentlab/llm/llm_utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import base64
import collections
import importlib
import io
import json
import logging
Expand All @@ -16,11 +17,17 @@
import openai
import tiktoken
import yaml
from langchain.schema import BaseMessage as LangchainBaseMessage
from langchain_community.adapters.openai import convert_message_to_dict
from PIL import Image
from transformers import AutoModel, AutoTokenizer

langchain_community = importlib.util.find_spec("langchain_community")
if langchain_community is not None:
from langchain.schema import BaseMessage as LangchainBaseMessage
from langchain_community.adapters.openai import convert_message_to_dict
else:
LangchainBaseMessage = None
convert_message_to_dict = None

if TYPE_CHECKING:
from agentlab.llm.chat_api import ChatModel

Expand All @@ -32,7 +39,7 @@ def messages_to_dict(messages: list[dict] | list[LangchainBaseMessage]) -> dict:
new_messages.add_message(m)
elif isinstance(m, str):
new_messages.add_message({"role": "<unknown role>", "content": m})
elif isinstance(m, LangchainBaseMessage):
elif LangchainBaseMessage is not None and isinstance(m, LangchainBaseMessage):
new_messages.add_message(convert_message_to_dict(m))
else:
raise ValueError(f"Unknown message type: {type(m)}")
Expand Down
29 changes: 25 additions & 4 deletions src/agentlab/llm/tracking.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import importlib
import logging
import os
import re
Expand All @@ -9,7 +10,13 @@
from typing import Optional

import requests
from langchain_community.callbacks import bedrock_anthropic_callback, openai_info

langchain_community = importlib.util.find_spec("langchain_community")
if langchain_community is not None:
from langchain_community.callbacks import bedrock_anthropic_callback, openai_info
else:
bedrock_anthropic_callback = None
openai_info = None
from litellm import completion_cost, get_model_info

TRACKER = threading.local()
Expand Down Expand Up @@ -103,7 +110,14 @@ def get_pricing_openrouter():

def get_pricing_openai():
"""Returns a dictionary of model pricing for OpenAI models."""
cost_dict = openai_info.MODEL_COST_PER_1K_TOKENS
try:
cost_dict = openai_info.MODEL_COST_PER_1K_TOKENS
except Exception as e:
logging.warning(
f"Failed to get OpenAI pricing: {e}. "
"Please install langchain-community or use LiteLLM API for pricing information."
)
return {}
cost_dict = {k: v / 1000 for k, v in cost_dict.items()}
res = {}
for k in cost_dict:
Expand All @@ -126,8 +140,15 @@ def _remove_version_suffix(model_name):

def get_pricing_anthropic():
"""Returns a dictionary of model pricing for Anthropic models."""
input_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_INPUT_TOKENS
output_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_OUTPUT_TOKENS
try:
input_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_INPUT_TOKENS
output_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_OUTPUT_TOKENS
except Exception as e:
Comment on lines +143 to +146
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Over-broad exception handling category Error Handling

Tell me more
What is the issue?

Using a broad Exception catch that could mask specific errors like AttributeError when bedrock_anthropic_callback is None.

Why this matters

Different error types require different handling strategies. Catching specific exceptions helps identify and handle issues appropriately.

Suggested change ∙ Feature Preview
try:
    if bedrock_anthropic_callback is None:
        raise ImportError("langchain_community not installed")
    input_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_INPUT_TOKENS
    output_cost_dict = bedrock_anthropic_callback.MODEL_COST_PER_1K_OUTPUT_TOKENS
except (ImportError, AttributeError) as e:
    logging.warning(
        f"Failed to get Anthropic pricing: {e}. "
        "Please install langchain-community or use LiteLLM API for pricing information."
    )
    return {}
Provide feedback to improve future suggestions

Nice Catch Incorrect Not in Scope Not in coding standard Other

💬 Looking for more details? Reply to this comment to chat with Korbit.

logging.warning(
f"Failed to get Anthropic pricing: {e}. "
"Please install langchain-community or use LiteLLM API for pricing information."
)
return {}

res = {}
for k, v in input_cost_dict.items():
Expand Down
Loading
Loading