From db34819ab69dcfce2aeb703ebd2d31b9155ef5a7 Mon Sep 17 00:00:00 2001 From: jinyuan Date: Tue, 23 Jun 2026 20:25:39 +0100 Subject: [PATCH 1/9] add util functions in utils/utils.py --- evoagentx/utils/utils.py | 504 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 496 insertions(+), 8 deletions(-) diff --git a/evoagentx/utils/utils.py b/evoagentx/utils/utils.py index a3f127ab..e44cf7e8 100644 --- a/evoagentx/utils/utils.py +++ b/evoagentx/utils/utils.py @@ -1,10 +1,17 @@ +import contextvars +import json import os import re import time +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from math import tanh from typing import ( TYPE_CHECKING, Any, + Dict, List, + Optional, Set, Type, Union, @@ -14,12 +21,21 @@ import regex import requests +from openai.types.chat.chat_completion_message_tool_call import ( + ChatCompletionMessageToolCall, +) +from pydantic import BaseModel, ValidationError +from pydantic_core import PydanticUndefined from tqdm import tqdm +from ..core.base_config import Parameter from ..core.logging import logger +from ..core.registry import MODULE_REGISTRY # Import for type hints (avoiding circular imports with TYPE_CHECKING) if TYPE_CHECKING: + from ..agents import Agent + from ..models import LLMConfig from ..tools.tool import Tool, Toolkit @@ -43,13 +59,31 @@ def safe_remove(data: Union[List[Any], Set[Any]], remove_value: Any): def generate_dynamic_class_name(base_name: str) -> str: base_name = base_name.strip() - + cleaned_name = re.sub(r'[^a-zA-Z0-9\s]', ' ', base_name) components = cleaned_name.split() class_name = ''.join(x.capitalize() for x in components) return class_name if class_name else 'DefaultClassName' + +def get_unique_class_name(candidate_name: str) -> str: + """ + Get a unique class name by checking if it already exists in the registry. + If it does, append "Vx" to make it unique. + """ + if not MODULE_REGISTRY.has_module(candidate_name): + return candidate_name + + i = 1 + while True: + unique_name = f"{candidate_name}V{i}" + if not MODULE_REGISTRY.has_module(unique_name): + break + i += 1 + return unique_name + + def normalize_text(s: str) -> str: def remove_articles(text): @@ -77,7 +111,7 @@ def download_file(url: str, save_file: str, max_retries=3, timeout=10): resume_byte_pos = 0 if os.path.exists(save_file): resume_byte_pos = os.path.getsize(save_file) - + response_head = requests.head(url=url) total_size = int(response_head.headers.get("content-length", 0)) @@ -91,13 +125,13 @@ def download_file(url: str, save_file: str, max_retries=3, timeout=10): # total_size = int(response.headers.get("content-length", 0)) mode = 'ab' if resume_byte_pos else 'wb' progress_bar = tqdm(total=total_size, unit="iB", unit_scale=True, initial=resume_byte_pos) - + with open(save_file, mode) as file: for chunk_data in response.iter_content(chunk_size=1024): if chunk_data: size = file.write(chunk_data) progress_bar.update(size) - + progress_bar.close() if os.path.getsize(save_file) >= (total_size + resume_byte_pos): @@ -119,6 +153,296 @@ def download_file(url: str, save_file: str, max_retries=3, timeout=10): raise RuntimeError(error_message) +def recursive_remove(data: Any, keys: List[str]) -> Any: + """ + Recursively removes specified keys from dictionaries and their nested structures within a + dictionary or list, if an object is not a list or dictionary return as is. + + Args: + data (Any): Specified keys will be removed from `data` if it is a dictionary or a list containing dictionaries. + keys (List[str]): A list of string keys to be removed. + """ + if isinstance(data, dict): + new_dict = {} + for k, v in data.items(): + if k not in keys: + new_dict[k] = recursive_remove(v, keys) + return new_dict + elif isinstance(data, list): + new_list = [recursive_remove(item, keys) for item in data] + return new_list + else: + return data + + +def tool_names_to_tools( + tool_names: Optional[List[str]] = None, + tools: Optional[List] = None, +) -> Optional[List]: + + if not tool_names: + return None + + if not tools: + raise ValueError(f"Must provide the following tools: {tool_names}") + + tool_map = {tool.name: tool for tool in tools} + + tool_list = [] + for tool_name in tool_names: + if tool_name not in tool_map: + raise ValueError(f"'{tool_name}' not found in provided tools") + tool_list.append(tool_map[tool_name]) + return tool_list + + +def add_llm_config_to_agent_dict(agent_dict: Dict, llm_config: Optional['LLMConfig'] = None) -> Dict: + """Assign the llm_config to agent_dict, overwriting any existing value. + If `llm` exists, it will be overwritten by `llm_config` to prevent conflicts. + If `is_human` is True, llm_config will not be added. + """ + + agent_dict_copy = agent_dict.copy() + + if agent_dict_copy.get("is_human", False): + return agent_dict_copy + + agent_llm_config = agent_dict_copy.get("llm_config", None) + agent_llm = agent_dict_copy.get("llm", None) + + if llm_config is None and agent_llm_config is None and agent_llm is None: + raise ValueError("Must provide `llm_config` or `llm` for agent") + + if llm_config is not None: + agent_dict_copy.pop("llm", None) + agent_dict_copy["llm_config"] = llm_config + return agent_dict_copy + + +def create_agent_from_dict( + agent_dict: Dict, + llm_config: Optional['LLMConfig'] = None, + tools: Optional[List] = None, + agents: Optional[List] = None, + **kwargs +) -> 'Agent': + + agent_class_name = agent_dict.pop("class_name", None) + + if agent_class_name is None: + agent_class_name = "CustomizeAgent" + + cls = MODULE_REGISTRY.get_module(agent_class_name) + agent = cls.from_dict(data=agent_dict, llm_config=llm_config, tools=tools, agents=agents, **kwargs) + return agent + + +def pydantic_to_parameters(base_model: Type[BaseModel], ignore: List[str] = []) -> List[Parameter]: + """ + Converts a Pydantic BaseModel class into a list of Parameter instances. + + Args: + model: A Pydantic BaseModel class. + + Returns: + A list of Parameter objects, where each object corresponds to a field + in the input BaseModel. + """ + parameters = [] + for field_name, field_info in base_model.model_fields.items(): + if field_name in ignore: + continue + + # Determine the description + description = field_info.description if field_info.description else field_name + + # Determine if the field is required + # A field is considered required if it doesn't have a default value + # and isn't Optional. + required = field_info.is_required() + + field_type = python_to_json_type[extract_type(field_info.annotation)] + + # Create the Parameter instance + param = Parameter( + name=field_name, + type=field_type, + description=description, + required=required, + json_schema=field_info.json_schema_extra, + ) + parameters.append(param) + return parameters + + +def validate_param( + required_param: Parameter, + actual_param: Parameter, + required_params_name: str, + actual_params_name: str, +): + """ + Checks if `actual_param` has the same type, required, description and json_schema value as `required_param`. + """ + + def format_error_msg( + attr_name: str, + required_value: Any, + actual_value: Any, + ) -> str: + return f"Mismatch for '{required_param.name}': {required_params_name} ({attr_name}={required_value}) vs. {actual_params_name} ({attr_name}={actual_value})" + + try: + actual_type = string_to_python_type[actual_param.type] + required_type = string_to_python_type[required_param.type] + except KeyError as e: + logger.warning(f"Unsupported type in {actual_param.name}: {e}") + actual_type = actual_param.type + required_type = required_param.type + + if required_type != actual_type: + raise ValueError(format_error_msg("type", required_param.type, actual_param.type)) + + if required_param.required != actual_param.required: + raise ValueError(format_error_msg("required", required_param.required, actual_param.required)) + + if required_param.description != actual_param.description: + raise ValueError(format_error_msg("description", required_param.description, actual_param.description)) + + if required_param.json_schema != actual_param.json_schema: + raise ValueError(format_error_msg("json_schema", required_param.json_schema, actual_param.json_schema)) + + +def format_validation_error(error: ValidationError) -> str: + """ + Formats a Pydantic ValidationError into a nicely formatted string. + + Args: + error: The Pydantic ValidationError object. + + Returns: + A formatted string containing all error details. + """ + formatted_messages: List[str] = [] + + for e in error.errors(): + path_parts = [] + for item in e['loc']: + if isinstance(item, int): + path_parts[-1] += f"[{item}]" + else: + path_parts.append(str(item)) + + error_location_str = ".".join(path_parts) + + if error_location_str: + formatted_message = ( + f"Location: {error_location_str}\n" + f"{e['msg']}\n" + ) + else: + formatted_message = f"{e['msg']}\n" + + formatted_messages.append(formatted_message) + + return "\n".join(formatted_messages) + + +def params_to_json(params: List[Parameter], ignore: List[str] = []) -> str: + params_dict = [param.to_dict(ignore=ignore) for param in params] + params_json = json.dumps(params_dict, indent=4, ensure_ascii=False) + return params_json + + +def fix_property_name(object: Any, json_schema: Dict) -> Any: + """ + Recursively fixes the property names of `object` to match the provided JSON schema. + """ + if object is None: + return object + + if json_schema["type"] == "array" and json_schema["items"]["type"] == "object": + return [fix_property_name(item, json_schema["items"]) for item in object] + + elif json_schema["type"] == "object": + fixed_object = dict() + properties = json_schema.get("properties") + + if properties is None: + return object + + for property_name, property_schema in properties.items(): + + if property_schema["type"] == "array": + property = object.get(property_name, None) + if property is not None: + fixed_object[property_name] = [fix_property_name(item, property_schema["items"]) for item in property] + + elif property_schema["type"] == "object": + property = object.get(property_name, None) + if property is not None: + fixed_object[property_name] = fix_property_name(property, property_schema) + + else: + object_properties_lower = {name.lower(): name for name in object} + schema_properties_lower = {name.lower(): name for name in properties} + + for name in object_properties_lower: + if name in schema_properties_lower: + fixed_object[schema_properties_lower[name]] = object[object_properties_lower[name]] + else: + fixed_object[object_properties_lower[name]] = object[object_properties_lower[name]] + + return fixed_object + + else: + return object + + +def resolve_json_schema_ref(json_schema: Any, root_schema: Optional[Dict] = None) -> Any: + """ + Recursively resolve all $ref in a JSON schema. + + Parameters: + json_schema (Any): The current schema to resolve. + root_schema (Optional[Dict]): The root schema used to resolve references. If not provided, it will be set to `json_schema`. + + Returns: + Any: A new schema with all $ref replaced by their actual definitions. + """ + if root_schema is None: + if not isinstance(json_schema, dict): + raise ValueError("`root_schema` must be a dictionary") + root_schema = json_schema + + if isinstance(json_schema, dict): + if "$ref" in json_schema: + ref_path = json_schema["$ref"] + + # Only support internal references (starting with "#/") + if not ref_path.startswith("#/"): + raise ValueError(f"External references not supported: {ref_path}") + + # Navigate the path + parts = ref_path.lstrip("#/").split("/") + target = root_schema + + for part in parts: + target = target[part] + + resolved = resolve_json_schema_ref(deepcopy(target), root_schema) + return resolved + + # Recurse into dict values + return {k: resolve_json_schema_ref(v, root_schema) for k, v in json_schema.items()} + + elif isinstance(json_schema, list): + return [resolve_json_schema_ref(item, root_schema) for item in json_schema] + + else: + return json_schema + + def remove_none(obj): """ Recursively removes all keys where the value is None. @@ -152,14 +476,88 @@ def extract_type(annotation: Type) -> Type: return annotation +def get_name_to_value_map(names: List[str], lookup_dict: Dict[str, Any]) -> Dict[str, Any]: + """ + Takes a list of names and a dictionary, and returns a dictionary mapping each name to its corresponding value in the lookup dictionary. + + Args: + names (List[str]): The list of names to look up. + lookup_dict (Dict[str, Any]): The dictionary to look up the names in. + + Returns: + Dict[str, Any]: A dictionary mapping each name to its corresponding value in the lookup dictionary. + """ + return {name: lookup_dict.get(name, None) for name in names} + + +def compute_score( + score: float, + improvement_score: float, + min_score: float = 1., + max_score: float = 10., + decimal_places: int = 2 +) -> float: + new_score = score + (max_score - score) * tanh(improvement_score) + return round(min(max_score, max(min_score, new_score)), decimal_places) + + +def transform_score(old_score: float, old_min: float, old_max: float, new_min: float, new_max: float) -> float: + """Transforms a score from one scale to another using linear scaling.""" + assert old_min < old_max, "`old_min` must be less than `old_max`" + assert new_min < new_max, "`new_min` must be less than `new_max`" + + transformed_score = ((old_score - old_min) / (old_max - old_min)) * (new_max - new_min) + new_min + return transformed_score + + +def format_execution_data(input_data: List[Dict], output_data: List[Dict]) -> List[Dict]: + execution_data = [ + { + "execution_input": input, + "execution_output": output + } + for input, output in zip(input_data, output_data, strict=True) + ] + + return execution_data + + +def compose_decorators(*decorators): + def combined(func): + wrapped = func + for decorator in decorators: + wrapped = decorator(wrapped) + return wrapped + return combined + + +def add_dict(a: Dict[str, Union[float, int]], b: Dict[str, Union[float, int]]) -> Dict[str, Union[float, int]]: + """ + Adds the values from two dict together if they share the same key. + Also keeps the values that don't share keys in the final output. + """ + if not a: + return b + + if not b: + return a + + dict_sum = a.copy() + + for key, value in b.items(): + dict_sum[key] = dict_sum.get(key, 0) + value + + return dict_sum + + + def compile_tool_schemas(tools: List[Union['Tool', 'Toolkit']]) -> List[dict]: """ Compiles the schemas of a list of tools or toolkits. - + Args: tools: A list of tools or toolkits - extra_description: Whether to include extra description in the schema - + Returns: A list of dictionaries containing the schemas of the tools or toolkits """ @@ -179,6 +577,96 @@ def compile_tool_schemas(tools: List[Union['Tool', 'Toolkit']]) -> List[dict]: return schemas +def format_tool_calls(tool_calls: List[ChatCompletionMessageToolCall]) -> List[Dict]: + """ + Formats a list of tool calls into a EAX format. + + Args: + tool_calls: A list of tool calls + + Returns: + A string containing the formatted tool calls + """ + formatted_tool_calls = [] + + for tool_call in tool_calls: + tool_name = tool_call.function.name + try: + tool_args = json.loads(tool_call.function.arguments) + except Exception: + logger.error(f"Failed to parse tool call arguments for `{tool_name}`:\n{tool_call.function.arguments}") + continue + + formatted_tool_calls.append( + { + "id": tool_call.id, + "function_name": tool_name, + "function_args": tool_args, + } + ) + + return formatted_tool_calls + + +def get_field_default(model: Type[BaseModel], field_name: str) -> Optional[Any]: + """ + Retrieves the default value for a specified field in a Pydantic model. + + If the field has a default value, returns it. If the field does not have a default value but + has a default factory, calls the factory to generate and return the default value. + Returns None if no default or factory is defined. + + Args: + model: The Pydantic model class. + field_name: The name of the field to check. + + Returns: + The default value of the field or None if no default is defined. + """ + + field = model.model_fields.get(field_name) + + if field is not None: + if field.default is not PydanticUndefined: + return field.default + + if field.default_factory is not None: + return field.default_factory() + + return None + + +def to_params(items: List[Union[Parameter, dict]]) -> List[Parameter]: + """ + Convert a list of dictionaries or Parameter objects into a list of Parameter objects. + """ + params: List[Parameter] = [] + for item in items: + if isinstance(item, dict): + params.append(Parameter(**item)) + elif isinstance(item, Parameter): + params.append(item) + else: + raise TypeError(f"Expects dict or Parameter, but got {type(item).__name__}") + return params + + + +class ContextualThreadPoolExecutor(ThreadPoolExecutor): + """ThreadPoolExecutor that preserves context variables""" + + def __init__(self, max_workers: Optional[int] = None, **kwargs): + super().__init__(max_workers=max_workers, **kwargs) + + def submit(self, fn, *args, **kwargs): + current_context = contextvars.copy_context() + + def wrapped_fn(*args, **kwargs): + return current_context.run(fn, *args, **kwargs) + + return super().submit(wrapped_fn, *args, **kwargs) + + string_to_python_type = { "string": str, "integer": int, @@ -219,7 +707,7 @@ def compile_tool_schemas(tools: List[Union['Tool', 'Toolkit']]) -> List[dict]: "number": "number", "boolean": "boolean", "object": "object", - "array": "array", + "array": "array", "str": "string", "int": "integer", "float": "number", From 2dacfb743f2752c7f2c0f0e252011c04bd488e41 Mon Sep 17 00:00:00 2001 From: jinyuan Date: Tue, 23 Jun 2026 20:30:41 +0100 Subject: [PATCH 2/9] add _validate_type_and_schema in Parameter class --- evoagentx/core/base_config.py | 35 +++++++++++++++---- tests/src/core/test_base_config.py | 55 ++++++++++++++++++++++++++++-- 2 files changed, 80 insertions(+), 10 deletions(-) diff --git a/evoagentx/core/base_config.py b/evoagentx/core/base_config.py index c7ea0a3e..9159f1c2 100644 --- a/evoagentx/core/base_config.py +++ b/evoagentx/core/base_config.py @@ -1,7 +1,11 @@ -# from pydantic import BaseModel -from typing import Optional, List +from typing import List, Optional + +from jsonschema import Draft7Validator +from pydantic import model_validator + from .module import BaseModule + class BaseConfig(BaseModule): """ @@ -53,15 +57,32 @@ def get_set_params(self, ignore: List[str] = []) -> dict: class Parameter(BaseModule): """Parameter class used to define configuration parameters. - + Attributes: name: Parameter name - type: Parameter type + type: Parameter type, support json & python type. if type is `object` or `array`, then schema is required. description: Parameter description required: Whether the parameter is required, defaults to True + json_schema: the json schema of the parameter, required when type is `object` or `array`. """ name: str - type: str - description: str - required: Optional[bool] = True + type: str + description: str + required: Optional[bool] = True + json_schema: Optional[dict] = None + + @model_validator(mode="after") + def _validate_type_and_schema(self): + from ..utils.utils import string_to_python_type + if self.type not in string_to_python_type: + raise ValueError(f"Invalid `type`: {self.type}. Allowed: {list(string_to_python_type.keys())}") + if self.type in {"object", "array"} and not self.json_schema: + raise ValueError("`json_schema` is required when `type` is `object` or `array`.") + if self.json_schema is not None: + try: + Draft7Validator.check_schema(self.json_schema) + except Exception as e: + raise ValueError(f"Invalid `json_schema` for '{self.name}': {self.json_schema}.") from e + assert self.type == self.json_schema.get("type"), f"`type` and `json_schema.type` must be the same if `json_schema` is provided. But got `type`: {self.type}, `json_schema.type`: {self.json_schema.get('type')}" + return self diff --git a/tests/src/core/test_base_config.py b/tests/src/core/test_base_config.py index 28c0668a..00ae3e6c 100644 --- a/tests/src/core/test_base_config.py +++ b/tests/src/core/test_base_config.py @@ -1,15 +1,18 @@ import unittest from typing import List -from evoagentx.core.base_config import BaseConfig + +import pytest + +from evoagentx.core.base_config import BaseConfig, Parameter class ToyConfig(BaseConfig): - var1: str + var1: str var2: List[str] var3: int = 111 -class TestModule(unittest.TestCase): +class TestBaseConfig(unittest.TestCase): def test_base_config(self): @@ -25,5 +28,51 @@ def test_base_config(self): self.assertEqual(set_params["var1"], "test") +class TestParameter(unittest.TestCase): + + def test_basic_parameter(self): + param = Parameter(name="x", type="string", description="a string param") + self.assertEqual(param.name, "x") + self.assertEqual(param.type, "string") + self.assertTrue(param.required) + self.assertIsNone(param.json_schema) + + def test_python_type_aliases(self): + for type_str in ("str", "int", "float", "bool", "dict", "list"): + param = Parameter(name="p", type=type_str, description="test") + self.assertEqual(param.type, type_str) + + def test_invalid_type_raises(self): + with self.assertRaises(Exception): + Parameter(name="p", type="invalid_type", description="bad type") + + def test_object_type_requires_json_schema(self): + with self.assertRaises(Exception): + Parameter(name="p", type="object", description="missing schema") + + def test_array_type_requires_json_schema(self): + with self.assertRaises(Exception): + Parameter(name="p", type="array", description="missing schema") + + def test_object_with_valid_json_schema(self): + schema = {"type": "object", "properties": {"key": {"type": "string"}}} + param = Parameter(name="p", type="object", description="an object", json_schema=schema) + self.assertEqual(param.json_schema, schema) + + def test_array_with_valid_json_schema(self): + schema = {"type": "array", "items": {"type": "string"}} + param = Parameter(name="p", type="array", description="an array", json_schema=schema) + self.assertEqual(param.json_schema, schema) + + def test_json_schema_type_mismatch_raises(self): + schema = {"type": "object", "properties": {}} + with self.assertRaises(Exception): + Parameter(name="p", type="string", description="mismatch", json_schema=schema) + + def test_invalid_json_schema_raises(self): + with self.assertRaises(Exception): + Parameter(name="p", type="object", description="bad schema", json_schema={"type": "object", "properties": "not_a_dict"}) + + if __name__ == "__main__": unittest.main() From 7fc9d44ecc5e47fcd2de0b5dd40ffaba7f04e332 Mon Sep 17 00:00:00 2001 From: jinyuan Date: Tue, 23 Jun 2026 20:39:08 +0100 Subject: [PATCH 3/9] fix parse_data_from_text typing and simplify JSON parsing --- evoagentx/core/module_utils.py | 93 ++++++++++++++++------------------ 1 file changed, 44 insertions(+), 49 deletions(-) diff --git a/evoagentx/core/module_utils.py b/evoagentx/core/module_utils.py index 954c11df..4701cd0c 100644 --- a/evoagentx/core/module_utils.py +++ b/evoagentx/core/module_utils.py @@ -1,7 +1,8 @@ import json import os from datetime import date, datetime -from typing import Any, Dict, List, Optional, Type, Union, get_args, get_origin +from types import UnionType +from typing import Any, Dict, List, Type, Union, get_args, get_origin from uuid import uuid4 import regex @@ -84,29 +85,6 @@ def save_json(data, path: str, type: str="json", use_indent: bool=True) -> str: return path -def extract_fenced_blocks(text: str, labels: Optional[List[str]] = None) -> List[str]: - """ - Extract fenced code blocks from the given text. - - Args: - text (str): The text to extract fenced code blocks from. - labels (List[str]): The labels to extract fenced code blocks for. - - Returns: - List[str]: Code blocks with specified labels. - """ - # Pattern to match fenced blocks: ```label\ncode\n``` - pattern = r"```([a-zA-Z0-9_\-\+]*)\s*\n*(.*?)\n*```" - matches = regex.findall(pattern, text, regex.DOTALL) - - if labels: - # Normalize labels for case-insensitive matching - labels_lower = {label.lower() for label in labels} - return [code.strip() for lang, code in matches if lang.strip().lower() in labels_lower] - - return [code.strip() for _, code in matches] - - def escape_json_values(string: str) -> str: def escape_value(match): @@ -196,30 +174,25 @@ def _replacer(match) -> str: def fix_json(string: str) -> str: string = remove_json_comments(string) - string = fix_json_booleans(string) + # string = fix_json_booleans(string) string = escape_json_values(string) return string def parse_json_from_text(text: str) -> List[str]: """ - Autoregressively extract JSON object from text + Autoregressively extract JSON object from text + + Args: + text (str): a text that includes JSON data - Args: - text (str): a text that includes JSON data - Returns: List[str]: a list of parsed JSON data """ - fenced_blocks = extract_fenced_blocks(text) - if fenced_blocks: - matches = fenced_blocks - else: - json_pattern = r"""(?:\{(?:[^{}]*|(?R))*\}|\[(?:[^\[\]]*|(?R))*\])""" - pattern = regex.compile(json_pattern, regex.VERBOSE) - matches = pattern.findall(text) - - matches = [fix_json(m) for m in matches] + json_pattern = r"""(?:\{(?:[^{}]*|(?R))*\}|\[(?:[^\[\]]*|(?R))*\])""" + pattern = regex.compile(json_pattern, regex.VERBOSE) + matches = pattern.findall(text) + matches = [fix_json(match) for match in matches] return matches @@ -231,20 +204,42 @@ def parse_xml_from_text(text: str, label: str) -> List[str]: values = [match.strip() for match in matches] return values -def parse_data_from_text(text: str, datatype: str): - - if datatype == "str": +def parse_data_from_text(text: str, datatype: Type): + if datatype is str: data = text - elif datatype == "int": + + elif datatype is int: data = int(text) - elif datatype == "float": + + elif datatype is float: data = float(text) - elif datatype == "bool": + + elif datatype is bool: data = text.lower() in ("true", "yes", "1", "on", "True") - elif datatype == "list": - data = eval(text) - elif datatype == "dict": - data = eval(text) + + elif datatype is list: + try: + data = json.loads(text) + except json.JSONDecodeError: + data = [item.strip() for item in text.split(",")] + type_args = get_args(datatype) + if len(type_args) == 1: + data = [parse_data_from_text(item, type_args[0]) for item in data] + + elif datatype is dict: + data = json.loads(text) + + elif get_origin(datatype) is Union or get_origin(datatype) is UnionType: + type_args = get_args(datatype) + for i, type_arg in enumerate(type_args): + try: + data = parse_data_from_text(text, type_arg) + break + except Exception: + if i == len(type_args) - 1: + data = text + continue + else: # raise ValueError( # f"Invalid value '{datatype}' is detected for `datatype`. " @@ -252,7 +247,7 @@ def parse_data_from_text(text: str, datatype: str): # ) # logger.warning(f"Unknown datatype '{datatype}' is detected for `datatype`. Return the raw text instead.") # failed to parse the data, return the raw text - return text + return text return data def parse_json_from_llm_output(text: str) -> dict: From 7583895cb1da71bbb411a332aba99c14a377bfbb Mon Sep 17 00:00:00 2001 From: jinyuan Date: Tue, 23 Jun 2026 20:56:40 +0100 Subject: [PATCH 4/9] add json schema auto fix in LLMOutputParser --- evoagentx/models/base_model.py | 133 ++++++++++++++++++++++++++++-- evoagentx/models/model_configs.py | 2 + 2 files changed, 129 insertions(+), 6 deletions(-) diff --git a/evoagentx/models/base_model.py b/evoagentx/models/base_model.py index 386bf42b..6f3b1700 100644 --- a/evoagentx/models/base_model.py +++ b/evoagentx/models/base_model.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable from copy import copy, deepcopy -from typing import Any, Dict, List, Optional, Type, Union +from typing import Any, ClassVar, Dict, List, Optional, Type, Union import yaml from jsonschema import Draft7Validator @@ -14,6 +14,7 @@ from pydantic import Field, model_validator from pydantic_core import PydanticUndefined +from ..core.logging import logger from ..core.module_utils import ( extract_code_blocks, get_type_name, @@ -39,6 +40,7 @@ class LLMOutputParser(Parser): content: The raw text generated by the LLM. """ content: str = Field(default=None, exclude=True, description=RAW_LLM_OUTPUT_DESCRIPTION) + fix_json_schema_error: ClassVar[bool] = False def init_module(self): if "_raw_llm_output" in self.kwargs: @@ -60,8 +62,10 @@ def json_schema_validation(cls, data: dict) -> dict: final_data = remove_none(final_data) validator.validate(final_data) except JSONSchemaValidationError as e: - raise ValueError(e) - + if not cls.fix_json_schema_error: + raise ValueError(e) + final_data = LLMOutputParser.fix_data_on_validation_fail(validator, final_data) + else: for field_name, field_info in cls.model_fields.items(): @@ -73,10 +77,128 @@ def json_schema_validation(cls, data: dict) -> dict: try: validator.validate(field_value) except JSONSchemaValidationError as e: - raise ValueError(e) + if not cls.fix_json_schema_error: + raise ValueError(e) + field_value = LLMOutputParser.fix_data_on_validation_fail(validator, field_value) + final_data[field_name] = field_value return final_data - + + + @staticmethod + def fix_data_on_validation_fail(validator, data: dict) -> dict: + """Attempts to fix JSON schema validation errors by modifying the data. + + Args: + validator: The JSON schema validator. + data: The data to fix. + + Returns: + The modified data. + """ + fixed_data = deepcopy(data) + + try: + fixed_data = LLMOutputParser._recursive_fix(fixed_data, validator.schema) + except Exception as e: + logger.exception(f"Failed to fix data on JSON schema validation fail. {e}") + pass + + try: + validator.validate(fixed_data) + except JSONSchemaValidationError as e: + raise ValueError(e) + + return fixed_data + + + @staticmethod + def _recursive_fix(data: dict, schema: dict) -> dict: + """Recursively fixes data against schema.""" + if schema is None: + return data + + # 1. Fix children first (Bottom-Up) + if isinstance(data, dict) and "properties" in schema: + for k, sub_schema in schema["properties"].items(): + if k in data: + data[k] = LLMOutputParser._recursive_fix(data[k], sub_schema) + + elif isinstance(data, list) and "items" in schema: + items_schema = schema["items"] + for i in range(len(data)): + data[i] = LLMOutputParser._recursive_fix(data[i], items_schema) + + # 2. Validate and fix current level + # Loop because fixing one error may introduce new ones or require re-checking. + max_iter = 10 + validator = Draft7Validator(schema) + + for _ in range(max_iter): + errors = sorted(validator.iter_errors(data), key=lambda e: len(e.path), reverse=True) + + if not errors: + break + + try: + data = LLMOutputParser.fix_validation_error(errors[0], data, inplace=True) + except (IndexError, KeyError, TypeError): + pass + + return data + + + @staticmethod + def fix_validation_error(error: JSONSchemaValidationError, data: dict, inplace: bool = False) -> dict: + """Attempts to fix a single JSON schema validation error by modifying the data. + + Modifications: + - ENUM violation: Set value to the first enum value + - String length: Truncate string to max length or add spaces to reach min length + - Array length: Truncate array to max length or add elements to reach min length + - Numeric range: Set value to minimum or maximum + + Args: + error: The JSON schema validation error. + data: The data to fix. + inplace: Whether to fix the data in-place. + + Returns: + The modified data. + """ + fixed_data = data if inplace else deepcopy(data) + + def _get_parent_and_key(data, path): + parent = data + for p in path[:-1]: + parent = parent[p] + return parent, path[-1] + + parent, key = _get_parent_and_key(fixed_data, list(error.path)) + value = parent[key] + schema = error.schema + + if error.validator == "enum": + parent[key] = schema["enum"][0] + elif error.validator == "maxLength" and isinstance(value, str): + parent[key] = value[:schema["maxLength"]] + elif error.validator == "minLength" and isinstance(value, str): + needed = schema["minLength"] - len(value) + parent[key] = value + (" " * needed) + elif error.validator == "maxItems" and isinstance(value, list): + parent[key] = value[:schema["maxItems"]] + elif error.validator == "minItems" and isinstance(value, list): + needed_count = schema["minItems"] - len(value) + if len(value) > 0: + extension = (value * (needed_count // len(value) + 1))[:needed_count] + parent[key] = value + extension + elif error.validator == "minimum" and isinstance(value, (int, float)): + parent[key] = schema["minimum"] + elif error.validator == "maximum" and isinstance(value, (int, float)): + parent[key] = schema["maximum"] + + return fixed_data + @classmethod def _is_content_defined_in_subclass(cls) -> bool: @@ -823,7 +945,6 @@ def _process_messages_for_multimodal(self, messages: List[List[dict]]) -> List[L else: model_type = "openai" # Default to OpenAI format - from ..core.logging import logger logger.debug(f"Processing multimodal content: llm_type={llm_type}, model_type={model_type}") # Convert multimodal content to appropriate format diff --git a/evoagentx/models/model_configs.py b/evoagentx/models/model_configs.py index f9045542..3f6bc68f 100644 --- a/evoagentx/models/model_configs.py +++ b/evoagentx/models/model_configs.py @@ -171,6 +171,8 @@ class OpenRouterConfig(LLMConfig): tool_choice: Optional[Union[str, dict]] = Field(default=None, description="Controls which tool is called by model. Can be 'none', 'auto', 'required', or specific tool configuration.") stream: Optional[bool] = Field(default=None, description="If set to true, it sends partial message deltas. Tokens will be sent as they become available, with the stream terminated by a [DONE] message.") + extra_body: Optional[dict] = Field(default=None, description="Additional request body parameters for provider-specific features.") + def __str__(self): return self.model From 3a4082b1d91527ef118f5ecee4d2b0fe658fb1fa Mon Sep 17 00:00:00 2001 From: jinyuan Date: Wed, 24 Jun 2026 00:09:16 +0100 Subject: [PATCH 5/9] update cost manager and OpenRouterLLM --- evoagentx/core/metadata.py | 61 ++++++++ evoagentx/models/model_utils.py | 119 +++++++++------- evoagentx/models/openrouter_model.py | 205 ++++++++++++++++----------- evoagentx/prompts/tool_calling.py | 7 + 4 files changed, 254 insertions(+), 138 deletions(-) create mode 100644 evoagentx/core/metadata.py diff --git a/evoagentx/core/metadata.py b/evoagentx/core/metadata.py new file mode 100644 index 00000000..dbf7f5b5 --- /dev/null +++ b/evoagentx/core/metadata.py @@ -0,0 +1,61 @@ +from typing import Any + +from .module import BaseModule + + +class Metadata(BaseModule): + + def __add__(self, other: "Metadata") -> "Metadata": + if not isinstance(other, Metadata): + raise TypeError(f"Cannot add {type(other)} to `Metadata`") + + def merge_values(v1: Any, v2: Any, field_name: str) -> Any: + if isinstance(v1, (int, float)) and isinstance(v2, (int, float)): + return v1 + v2 + + elif isinstance(v1, list) and isinstance(v2, list): + return v1 + v2 + + elif isinstance(v1, set) and isinstance(v2, set): + return v1.union(v2) + + elif isinstance(v1, dict) and isinstance(v2, dict): + merged = dict(v1) + for key, val2 in v2.items(): + if key in merged: + merged[key] = merge_values(merged[key], val2, key) + else: + merged[key] = val2 + return merged + + elif isinstance(v1, str) and isinstance(v2, str): + if v1 == v2: + return v1 + else: + raise ValueError(f"Cannot add strings '{v1}' and '{v2}' for field '{field_name}'") + + elif isinstance(v1, bool) and isinstance(v2, bool): + return v1 or v2 + + elif isinstance(v1, Metadata) and isinstance(v2, Metadata): + return v1 + v2 + + else: + raise TypeError(f"Cannot add values of types {type(v1)} and {type(v2)} for field '{field_name}'") + + merged_data = self.model_dump() + merged_data.pop("class_name", None) + other_data = other.model_dump() + other_data.pop("class_name", None) + + for key, val2 in other_data.items(): + if key in merged_data: + merged_data[key] = merge_values(merged_data[key], val2, key) + else: + merged_data[key] = val2 + + if type(self) is type(other): + return self.__class__(**merged_data) + else: + return Metadata(**merged_data) + diff --git a/evoagentx/models/model_utils.py b/evoagentx/models/model_utils.py index 7fbe1e34..065bc41d 100644 --- a/evoagentx/models/model_utils.py +++ b/evoagentx/models/model_utils.py @@ -1,13 +1,15 @@ import threading +from collections import defaultdict +from typing import Optional + import pandas as pd -from dataclasses import dataclass from ..core.logging import logger from ..core.decorators import atomic_method from ..core.callbacks import suppress_cost_logs from ..core.registry import MODEL_REGISTRY from .model_configs import LLMConfig -from ..models.base_model import BaseLLM +from ..models.base_model import BaseLLM def get_openai_model_cost() -> dict: import json @@ -39,98 +41,107 @@ def infer_litellm_company_from_model(model: str) -> str: return company -@dataclass class Cost: - input_tokens: int - output_tokens: int - input_cost: float - output_cost: float + + def __init__( + self, + input_tokens: Optional[int] = None, + output_tokens: Optional[int] = None, + input_cost: Optional[float] = None, + output_cost: Optional[float] = None, + cost: Optional[float] = None + ): + self.input_tokens = input_tokens + self.output_tokens = output_tokens + # Keep input_cost/output_cost only as temporary constructor compatibility. + self.cost = cost if cost is not None else (input_cost or 0.0) + (output_cost or 0.0) + + @property + def cost(self) -> float: + return self._cost + + @cost.setter + def cost(self, value: float): + self._cost = value or 0.0 class CostManager: def __init__(self): - self.total_input_tokens = {} - self.total_output_tokens = {} - self.total_tokens = {} + self.input_tokens = defaultdict(int) + self.output_tokens = defaultdict(int) + self.total_tokens = defaultdict(int) - self.total_input_cost = {} - self.total_output_cost = {} - self.total_cost = {} + self.cost_per_model = defaultdict(float) self._lock = threading.Lock() - def compute_total_cost(self): - total_tokens, total_cost = 0, 0.0 - for _, value in self.total_tokens.items(): - total_tokens += value - for _, value in self.total_cost.items(): - total_cost += value - return total_tokens, total_cost + @property + def total_llm_cost(self) -> float: + return sum(self.cost_per_model.values()) + + @property + def total_llm_tokens(self) -> int: + return sum(self.total_tokens.values()) + + @property + def total_cost(self) -> float: + return self.total_llm_cost + + def add_llm_cost(self, cost: Cost, model: str): + self.input_tokens[model] += (cost.input_tokens or 0) + self.output_tokens[model] += (cost.output_tokens or 0) + self.total_tokens[model] += (cost.input_tokens or 0) + (cost.output_tokens or 0) + + self.cost_per_model[model] += cost.cost @atomic_method def update_cost(self, cost: Cost, model: str): + self.add_llm_cost(cost, model) + + total_tokens = self.total_llm_tokens + total_llm_cost = self.total_llm_cost + current_llm_cost = cost.cost + current_total_tokens = (cost.input_tokens or 0) + (cost.output_tokens or 0) - self.total_input_tokens[model] = self.total_input_tokens.get(model, 0) + cost.input_tokens - self.total_output_tokens[model] = self.total_output_tokens.get(model, 0) + cost.output_tokens - current_total_tokens = cost.input_tokens + cost.output_tokens - self.total_tokens[model] = self.total_tokens.get(model, 0) + current_total_tokens - - self.total_input_cost[model] = self.total_input_cost.get(model, 0.0) + cost.input_cost - self.total_output_cost[model] = self.total_output_cost.get(model, 0.0) + cost.output_cost - current_total_cost = cost.input_cost + cost.output_cost - self.total_cost[model] = self.total_cost.get(model, 0.0) + current_total_cost - - total_tokens, total_cost = self.compute_total_cost() if not suppress_cost_logs.get(): - logger.info(f"Total cost: ${total_cost:.3f} | Total tokens: {total_tokens} | Current cost: ${current_total_cost:.3f} | Current tokens: {current_total_tokens}") + logger.info(f"Total LLM cost: ${total_llm_cost:.3f} | Total tokens: {total_tokens} | Current LLM cost: ${current_llm_cost:.3f} | Current tokens: {current_total_tokens}") def display_cost(self): data = { "Model": [], - "Total Cost (USD)": [], - "Total Input Cost (USD)": [], - "Total Output Cost (USD)": [], - "Total Tokens": [], - "Total Input Tokens": [], + "Total Cost (USD)": [], + "Total Tokens": [], + "Total Input Tokens": [], "Total Output Tokens": [], } for model in self.total_tokens.keys(): data["Model"].append(model) - data["Total Cost (USD)"].append(round(self.total_cost[model], 4)) - data["Total Input Cost (USD)"].append(round(self.total_input_cost[model], 4)) - data["Total Output Cost (USD)"].append(round(self.total_output_cost[model], 4)) + data["Total Cost (USD)"].append(round(self.cost_per_model[model], 4)) data["Total Tokens"].append(self.total_tokens[model]) - data["Total Input Tokens"].append(self.total_input_tokens[model]) - data["Total Output Tokens"].append(self.total_output_tokens[model]) - - # Convert to DataFrame for display + data["Total Input Tokens"].append(self.input_tokens[model]) + data["Total Output Tokens"].append(self.output_tokens[model]) + df = pd.DataFrame(data) if len(df) > 1: summary = { "Model": "TOTAL", "Total Cost (USD)": df["Total Cost (USD)"].sum(), - "Total Input Cost (USD)": df["Total Input Cost (USD)"].sum(), - "Total Output Cost (USD)": df["Total Output Cost (USD)"].sum(), "Total Tokens": df["Total Tokens"].sum(), "Total Input Tokens": df["Total Input Tokens"].sum(), "Total Output Tokens": df["Total Output Tokens"].sum(), } - df = df._append(summary, ignore_index=True) - + df = pd.concat([df, pd.DataFrame([summary])], ignore_index=True) + print(df.to_string(index=False)) - def get_total_cost(self): - - total_cost = 0.0 - for model in self.total_cost.keys(): - total_cost += self.total_cost[model] - return total_cost + def get_total_cost(self) -> float: + return self.total_llm_cost cost_manager = CostManager() diff --git a/evoagentx/models/openrouter_model.py b/evoagentx/models/openrouter_model.py index 02ec1c98..b8520bcd 100644 --- a/evoagentx/models/openrouter_model.py +++ b/evoagentx/models/openrouter_model.py @@ -1,17 +1,20 @@ -import asyncio -import requests +import json +from typing import Dict, List, Optional, Union + +from openai import AsyncOpenAI, OpenAI, Stream +from openai.types.chat import ChatCompletion, ChatCompletionChunk from tenacity import ( retry, stop_after_attempt, wait_random_exponential, ) -from openai import OpenAI, Stream -from openai.types.chat import ChatCompletion -from typing import Optional, List -from litellm import token_counter + +from ..core.logging import logger from ..core.registry import register_model -from .model_configs import OpenRouterConfig +from ..prompts.tool_calling import TOOL_CALL_FORMAT +from ..utils.utils import format_tool_calls from .base_model import BaseLLM +from .model_configs import OpenRouterConfig from .model_utils import Cost, cost_manager @@ -24,11 +27,10 @@ def init_model(self): self._default_ignore_fields = ["llm_type", "openrouter_key", "openrouter_base", "openrouter_model_base", "output_response"] def _init_client(self, config: OpenRouterConfig): - client = OpenAI( - api_key=config.openrouter_key, - base_url=config.openrouter_base - ) - return client + return OpenAI(api_key=config.openrouter_key, base_url=config.openrouter_base) + + def _init_async_client(self, config: OpenRouterConfig): + return AsyncOpenAI(api_key=config.openrouter_key, base_url=config.openrouter_base) def formulate_messages(self, prompts: List[str], system_messages: Optional[List[str]] = None) -> List[List[dict]]: if system_messages: @@ -63,34 +65,123 @@ def get_completion_params(self, **kwargs): def get_stream_output(self, response: Stream, output_response: bool=True) -> str: output = "" + tool_calls_accum: Dict[int, dict] = {} + usage_chunk = None for chunk in response: - content = chunk.choices[0].delta.content - if content: + if chunk.usage is not None: + usage_chunk = chunk + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if delta.content: if output_response: - print(content, end="", flush=True) - output += content + print(delta.content, end="", flush=True) + output += delta.content + if delta.tool_calls: + self._accumulate_tool_calls(delta.tool_calls, tool_calls_accum) if output_response: print("") + if tool_calls_accum: + formatted = self._format_streamed_tool_calls(tool_calls_accum) + if formatted: + tool_call_str = TOOL_CALL_FORMAT.format(tool_calls=json.dumps(formatted, indent=4, ensure_ascii=False)) + output += tool_call_str + if output_response: + print(tool_call_str) + if usage_chunk is not None: + self._update_cost(usage_chunk) + else: + logger.warning("[OpenRouterLLM] No usage data in stream response; cost will not be recorded.") return output - + async def get_stream_output_async(self, response, output_response: bool = False) -> str: output = "" + tool_calls_accum: Dict[int, dict] = {} + usage_chunk = None async for chunk in response: - content = chunk.choices[0].delta.content - if content: + if chunk.usage is not None: + usage_chunk = chunk + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if delta.content: if output_response: - print(content, end="", flush=True) - output += content + print(delta.content, end="", flush=True) + output += delta.content + if delta.tool_calls: + self._accumulate_tool_calls(delta.tool_calls, tool_calls_accum) if output_response: print("") + if tool_calls_accum: + formatted = self._format_streamed_tool_calls(tool_calls_accum) + if formatted: + tool_call_str = TOOL_CALL_FORMAT.format(tool_calls=json.dumps(formatted, indent=4, ensure_ascii=False)) + output += tool_call_str + if output_response: + print(tool_call_str) + if usage_chunk is not None: + self._update_cost(usage_chunk) + else: + logger.warning("[OpenRouterLLM] No usage data in stream response; cost will not be recorded.") return output def get_completion_output(self, response: ChatCompletion, output_response: bool=True) -> str: - output = response.choices[0].message.content + output = response.choices[0].message.content or "" + tool_calls = getattr(response.choices[0].message, "tool_calls", None) + if tool_calls: + formatted = format_tool_calls(tool_calls) + output += TOOL_CALL_FORMAT.format(tool_calls=json.dumps(formatted, indent=4, ensure_ascii=False)) if output_response: print(output) + self._update_cost(response) return output + @staticmethod + def _accumulate_tool_calls(delta_tool_calls, accum: Dict[int, dict]): + for tc in delta_tool_calls: + idx = tc.index + if idx not in accum: + accum[idx] = {"id": "", "function": {"name": "", "arguments": ""}} + if tc.id: + accum[idx]["id"] = tc.id + if tc.function: + if tc.function.name: + accum[idx]["function"]["name"] += tc.function.name + if tc.function.arguments: + accum[idx]["function"]["arguments"] += tc.function.arguments + + @staticmethod + def _format_streamed_tool_calls(accum: Dict[int, dict]) -> List[dict]: + formatted = [] + for idx in sorted(accum.keys()): + tc = accum[idx] + try: + args = json.loads(tc["function"]["arguments"]) + except Exception: + logger.error(f"Failed to parse streaming tool call arguments for `{tc['function']['name']}`:\n{tc['function']['arguments']}") + continue + formatted.append({"id": tc["id"], "function_name": tc["function"]["name"], "function_args": args}) + return formatted + + def _update_cost(self, response: Union[ChatCompletion, ChatCompletionChunk]): + usage = response.usage + if usage is None: + logger.warning(f"[OpenRouterLLM] usage is None in response (id={response.id}); cost will not be recorded.") + return + cost_value = getattr(usage, "cost", None) + if cost_value is None: + logger.warning( + f"[OpenRouterLLM] usage.cost not present in response (id={response.id}); " + "cost will be recorded as 0. Check OpenRouter dashboard for actual spend." + ) + cost_value = 0.0 + cost = Cost( + input_tokens=usage.prompt_tokens, + output_tokens=usage.completion_tokens, + cost=cost_value, + ) + cost_manager.update_cost(cost, model=self.config.model) + @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5)) def single_generate(self, messages: List[dict], **kwargs) -> str: stream = kwargs.get("stream", self.config.stream) @@ -101,16 +192,13 @@ def single_generate(self, messages: List[dict], **kwargs) -> str: response = self._client.chat.completions.create(messages=messages, **completion_params) if stream: output = self.get_stream_output(response, output_response=output_response) - cost = self._stream_cost(messages=messages, output=output) else: output: str = self.get_completion_output(response=response, output_response=output_response) - cost = self._completion_cost(response) - self._update_cost(cost=cost) except Exception as e: raise RuntimeError(f"Error during single_generate of OpenRouterLLM: {str(e)}") - + return output - + def batch_generate(self, batch_messages: List[List[dict]], **kwargs) -> List[str]: return [self.single_generate(messages=one_messages, **kwargs) for one_messages in batch_messages] @@ -119,68 +207,17 @@ async def single_generate_async(self, messages: List[dict], **kwargs) -> str: output_response = kwargs.get("output_response", self.config.output_response) try: - isolated_client = self._init_client(self.config) + async_client = self._init_async_client(self.config) completion_params = self.get_completion_params(**kwargs) - - loop = asyncio.get_event_loop() - response = await loop.run_in_executor( - None, - lambda: isolated_client.chat.completions.create( - messages=messages, - **completion_params - ) + response = await async_client.chat.completions.create( + messages=messages, **completion_params ) - if stream: - if hasattr(response, "__aiter__"): - output = await self.get_stream_output_async(response, output_response=output_response) - else: - output = self.get_stream_output(response, output_response=output_response) - cost = self._stream_cost(messages=messages, output=output) + output = await self.get_stream_output_async(response, output_response=output_response) else: output: str = self.get_completion_output(response=response, output_response=output_response) - cost = self._completion_cost(response) - self._update_cost(cost=cost) - + except Exception as e: raise RuntimeError(f"Error during single_generate_async of OpenRouterLLM: {str(e)}") - return output - - def _completion_cost(self, response: ChatCompletion) -> Cost: - input_tokens = response.usage.prompt_tokens - output_tokens = response.usage.completion_tokens - return self._compute_cost(input_tokens=input_tokens, output_tokens=output_tokens) - - def _stream_cost(self, messages: List[dict], output: str) -> Cost: - model: str = self.config.model - input_tokens = token_counter(model=model, messages=messages) - output_tokens = token_counter(model=model, text=output) - return self._compute_cost(input_tokens=input_tokens, output_tokens=output_tokens) - - def _compute_cost(self, input_tokens: int, output_tokens: int) -> Cost: - - input_cost_per_token, output_cost_per_token = self._get_cost() - input_cost = input_tokens * input_cost_per_token - output_cost = output_tokens * output_cost_per_token - - cost = Cost(input_tokens=input_tokens, output_tokens=output_tokens, input_cost=input_cost, output_cost=output_cost) - return cost - - def _update_cost(self, cost: Cost): - cost_manager.update_cost(cost=cost, model=self.config.model) - - - def _get_cost(self): - url = self.config.openrouter_model_base - response = requests.get(url) - data = response.json() - - for model in data['data']: - if model['id'] == self.config.model: - pricing = model.get('pricing',{}) - input_cost = float(pricing.get('prompt', 0)) - output_cost = float(pricing.get('completion', 0)) - return input_cost, output_cost - - return 0, 0 \ No newline at end of file + return output \ No newline at end of file diff --git a/evoagentx/prompts/tool_calling.py b/evoagentx/prompts/tool_calling.py index e6ec495d..187a15e5 100644 --- a/evoagentx/prompts/tool_calling.py +++ b/evoagentx/prompts/tool_calling.py @@ -1,3 +1,10 @@ +TOOL_CALL_FORMAT = """ + +{tool_calls} + +""" + + OUTPUT_EXTRACTION_PROMPT = """ You are given the following text: {text} From c188b1802eaac1fb827ccecb3c94da318e31a05a Mon Sep 17 00:00:00 2001 From: jinyuan Date: Wed, 24 Jun 2026 15:16:58 +0100 Subject: [PATCH 6/9] update OpenAILLM and OpenRouterLLM, as well as add unit tests for them --- evoagentx/models/openai_model.py | 250 +++++++++++------ evoagentx/models/openrouter_model.py | 31 ++- evoagentx/prompts/tool_calling.py | 7 +- pyproject.toml | 8 +- pytest.ini | 1 + requirements.txt | 4 +- tests/src/models/mock_response.py | 316 +++++++++++++++++++++- tests/src/models/test_openai_model.py | 267 ++++++++++++++---- tests/src/models/test_openrouter_model.py | 221 +++++++++++++++ 9 files changed, 947 insertions(+), 158 deletions(-) create mode 100644 tests/src/models/test_openrouter_model.py diff --git a/evoagentx/models/openai_model.py b/evoagentx/models/openai_model.py index 04864400..fe5187c7 100644 --- a/evoagentx/models/openai_model.py +++ b/evoagentx/models/openai_model.py @@ -1,47 +1,74 @@ -import asyncio +import json +from typing import Dict, List, Optional, Union + +from openai import AsyncOpenAI, OpenAI, Stream +from openai.types.chat import ChatCompletion, ChatCompletionChunk +from openai.types.completion_usage import CompletionUsage from tenacity import ( retry, stop_after_attempt, wait_random_exponential, ) -from openai import OpenAI, Stream -from openai.types.chat import ChatCompletion -from typing import Optional, List -from litellm import token_counter, cost_per_token +from litellm import cost_per_token +from litellm.types.utils import Usage + +from ..core.logging import logger from ..core.registry import register_model +from ..prompts.tool_calling import TOOL_CALL_FORMAT +from ..utils.utils import format_tool_calls from .model_configs import OpenAILLMConfig from .base_model import BaseLLM -from .model_utils import Cost, cost_manager, get_openai_model_cost +from .model_utils import Cost, cost_manager, get_openai_model_cost @register_model(config_cls=OpenAILLMConfig, alias=["openai_llm"]) class OpenAILLM(BaseLLM): def init_model(self): - config: OpenAILLMConfig = self.config - self._client = self._init_client(config) # OpenAI(api_key=config.openai_key) + self._client = None + self._async_client = None self._default_ignore_fields = [ - "llm_type", "output_response", "openai_key", "deepseek_key", "anthropic_key", - "gemini_key", "meta_llama_key", "openrouter_key", "openrouter_base", "perplexity_key", + "llm_type", "output_response", "openai_key", "deepseek_key", "anthropic_key", + "gemini_key", "meta_llama_key", "openrouter_key", "openrouter_base", "perplexity_key", "groq_key" - ] # parameters in OpenAILLMConfig that are not OpenAI models' input parameters + ] # parameters in OpenAILLMConfig that are not OpenAI models' input parameters if self.config.model not in get_openai_model_cost(): raise KeyError(f"'{self.config.model}' is not a valid OpenAI model name!") - + def _init_client(self, config: OpenAILLMConfig): - client = OpenAI(api_key=config.openai_key) - return client + return OpenAI(api_key=config.openai_key) + + def _init_async_client(self, config: OpenAILLMConfig): + return AsyncOpenAI(api_key=config.openai_key) + + def ensure_client(self): + if self._client is None or self._client.is_closed(): + self._client = self._init_client(self.config) + return self._client + + def close_client(self): + if self._client is not None and not self._client.is_closed(): + self._client.close() + + def ensure_async_client(self): + if self._async_client is None or self._async_client.is_closed(): + self._async_client = self._init_async_client(self.config) + return self._async_client + + async def close_async_client(self): + if self._async_client is not None and not self._async_client.is_closed(): + await self._async_client.close() def formulate_messages(self, prompts: List[str], system_messages: Optional[List[str]] = None) -> List[List[dict]]: - + if system_messages: assert len(prompts) == len(system_messages), f"the number of prompts ({len(prompts)}) is different from the number of system_messages ({len(system_messages)})" else: system_messages = [None] * len(prompts) - - messages_list = [] + + messages_list = [] for prompt, system_message in zip(prompts, system_messages): - messages = [] + messages = [] if system_message: messages.append({"role": "system", "content": system_message}) messages.append({"role": "user", "content": prompt}) @@ -62,8 +89,15 @@ def update_completion_params(self, params1: dict, params2: dict) -> dict: def get_completion_params(self, **kwargs): completion_params = self.config.get_set_params(ignore=self._default_ignore_fields) completion_params = self.update_completion_params(completion_params, kwargs) + # automatically set stream_options to include usage if stream is True, + # as OpenAI's streaming response does not include usage by default, + # which is needed for cost tracking. + if completion_params.get("stream"): + stream_options = dict(completion_params.get("stream_options") or {}) + stream_options.setdefault("include_usage", True) + completion_params["stream_options"] = stream_options return completion_params - + def get_stream_output(self, response: Stream, output_response: bool=True) -> str: """ Process stream response and return the complete output. @@ -71,50 +105,119 @@ def get_stream_output(self, response: Stream, output_response: bool=True) -> str Args: response: The stream response from OpenAI output_response: Whether to print the response in real-time - + Returns: str: The complete output text """ output = "" + tool_calls_accum: Dict[int, dict] = {} + usage_chunk = None for chunk in response: - content = chunk.choices[0].delta.content - if content: + if getattr(chunk, "usage", None) is not None: + usage_chunk = chunk + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if delta.content: if output_response: - print(content, end="", flush=True) - output += content + print(delta.content, end="", flush=True) + output += delta.content + if delta.tool_calls: + self._accumulate_tool_calls(delta.tool_calls, tool_calls_accum) if output_response: print("") + if tool_calls_accum: + formatted = self._format_streamed_tool_calls(tool_calls_accum) + if formatted: + tool_call_str = TOOL_CALL_FORMAT.format(tool_calls=json.dumps(formatted, indent=4, ensure_ascii=False)) + output += tool_call_str + if output_response: + print(tool_call_str) + if usage_chunk is not None: + self._update_cost(usage_chunk) + else: + logger.warning("[OpenAILLM] No usage data in stream response; cost will not be recorded. Set stream_options={'include_usage': True} to enable cost tracking.") return output - + async def get_stream_output_async(self, response, output_response: bool = False) -> str: """ Process async stream response and return the complete output. - + Args: response (AsyncIterator[ChatCompletionChunk]): The async stream response from OpenAI output_response (bool): Whether to print the response in real-time - - + Returns: str: The complete output text """ output = "" + tool_calls_accum: Dict[int, dict] = {} + usage_chunk = None async for chunk in response: - content = chunk.choices[0].delta.content - if content: + if getattr(chunk, "usage", None) is not None: + usage_chunk = chunk + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if delta.content: if output_response: - print(content, end="", flush=True) - output += content + print(delta.content, end="", flush=True) + output += delta.content + if delta.tool_calls: + self._accumulate_tool_calls(delta.tool_calls, tool_calls_accum) if output_response: print("") + if tool_calls_accum: + formatted = self._format_streamed_tool_calls(tool_calls_accum) + if formatted: + tool_call_str = TOOL_CALL_FORMAT.format(tool_calls=json.dumps(formatted, indent=4, ensure_ascii=False)) + output += tool_call_str + if output_response: + print(tool_call_str) + if usage_chunk is not None: + self._update_cost(usage_chunk) + else: + logger.warning("[OpenAILLM] No usage data in stream response; cost will not be recorded. Set stream_options={'include_usage': True} to enable cost tracking.") return output def get_completion_output(self, response: ChatCompletion, output_response: bool=True) -> str: - output = response.choices[0].message.content + output = response.choices[0].message.content or "" + tool_calls = getattr(response.choices[0].message, "tool_calls", None) + if tool_calls: + formatted = format_tool_calls(tool_calls) + output += TOOL_CALL_FORMAT.format(tool_calls=json.dumps(formatted, indent=4, ensure_ascii=False)) if output_response: print(output) + self._update_cost(response) return output + @staticmethod + def _accumulate_tool_calls(delta_tool_calls, accum: Dict[int, dict]): + for tc in delta_tool_calls: + idx = tc.index + if idx not in accum: + accum[idx] = {"id": "", "function": {"name": "", "arguments": ""}} + if tc.id: + accum[idx]["id"] = tc.id + if tc.function: + if tc.function.name: + accum[idx]["function"]["name"] += tc.function.name + if tc.function.arguments: + accum[idx]["function"]["arguments"] += tc.function.arguments + + @staticmethod + def _format_streamed_tool_calls(accum: Dict[int, dict]) -> List[dict]: + formatted = [] + for idx in sorted(accum.keys()): + tc = accum[idx] + try: + args = json.loads(tc["function"]["arguments"]) + except Exception: + logger.error(f"Failed to parse streaming tool call arguments for `{tc['function']['name']}`:\n{tc['function']['arguments']}") + continue + formatted.append({"id": tc["id"], "function_name": tc["function"]["name"], "function_args": args}) + return formatted + @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5)) def single_generate(self, messages: List[dict], **kwargs) -> str: @@ -122,20 +225,18 @@ def single_generate(self, messages: List[dict], **kwargs) -> str: output_response = kwargs["output_response"] if "output_response" in kwargs else self.config.output_response try: + client = self.ensure_client() completion_params = self.get_completion_params(**kwargs) - response = self._client.chat.completions.create(messages=messages, **completion_params) + response = client.chat.completions.create(messages=messages, **completion_params) if stream: output = self.get_stream_output(response, output_response=output_response) - cost = self._stream_cost(messages=messages, output=output) else: output: str = self.get_completion_output(response=response, output_response=output_response) - cost = self._completion_cost(response) # calculate completion cost - self._update_cost(cost=cost) except Exception as e: raise RuntimeError(f"Error during single_generate of OpenAILLM: {str(e)}") - + return output - + def batch_generate(self, batch_messages: List[List[dict]], **kwargs) -> List[str]: return [self.single_generate(messages=one_messages, **kwargs) for one_messages in batch_messages] @@ -145,59 +246,36 @@ async def single_generate_async(self, messages: List[dict], **kwargs) -> str: output_response = kwargs.get("output_response", self.config.output_response) try: - # Create a completely new client instance to avoid thread-local storage issues - # This is a more aggressive approach than using a lock - # isolated_client = OpenAI(api_key=self.config.openai_key) - isolated_client = self._init_client(self.config) + async_client = self.ensure_async_client() completion_params = self.get_completion_params(**kwargs) - - # Use synchronous client in async context to avoid issues - loop = asyncio.get_running_loop() - response = await loop.run_in_executor( - None, - lambda: isolated_client.chat.completions.create( - messages=messages, - **completion_params - ) - ) - + response = await async_client.chat.completions.create(messages=messages, **completion_params) if stream: - if hasattr(response, "__aiter__"): - output = await self.get_stream_output_async(response, output_response=output_response) - else: - output = self.get_stream_output(response, output_response=output_response) - cost = self._stream_cost(messages=messages, output=output) + output = await self.get_stream_output_async(response, output_response=output_response) else: + # The network I/O is already awaited above; the response is fully in memory here, + # so this synchronous parsing/cost call does not block the event loop. output: str = self.get_completion_output(response=response, output_response=output_response) - cost = self._completion_cost(response) # calculate completion cost - self._update_cost(cost=cost) - except Exception as e: raise RuntimeError(f"Error during single_generate_async of OpenAILLM: {str(e)}") return output - - def _completion_cost(self, response: ChatCompletion) -> Cost: - input_tokens = response.usage.prompt_tokens - output_tokens = response.usage.completion_tokens - return self._compute_cost(input_tokens=input_tokens, output_tokens=output_tokens) - - def _stream_cost(self, messages: List[dict], output: str) -> Cost: - model: str = self.config.model - input_tokens = token_counter(model=model, messages=messages) - output_tokens = token_counter(model=model, text=output) - return self._compute_cost(input_tokens=input_tokens, output_tokens=output_tokens) - - def _compute_cost(self, input_tokens: int, output_tokens: int) -> Cost: - # use LiteLLM to compute cost, require the model name to be a valid model name in LiteLLM. - input_cost, output_cost = cost_per_token( - model=self.config.model, - prompt_tokens=input_tokens, - completion_tokens=output_tokens, + + def _compute_cost(self, usage: CompletionUsage) -> Cost: + # Pass the full usage object to LiteLLM so it can apply the correct rates for + # cached / reasoning tokens (cached input tokens are billed at a lower rate). + # LiteLLM expects its own Usage type, so convert from OpenAI's CompletionUsage. + usage_object = usage if isinstance(usage, Usage) else Usage(**usage.model_dump()) + input_cost, output_cost = cost_per_token(model=self.config.model, usage_object=usage_object) + return Cost( + input_tokens=usage.prompt_tokens, + output_tokens=usage.completion_tokens, + input_cost=input_cost, + output_cost=output_cost, ) - cost = Cost(input_tokens=input_tokens, output_tokens=output_tokens, input_cost=input_cost, output_cost=output_cost) - return cost - - def _update_cost(self, cost: Cost): - cost_manager.update_cost(cost=cost, model=self.config.model) - \ No newline at end of file + + def _update_cost(self, response: Union[ChatCompletion, ChatCompletionChunk]): + usage = getattr(response, "usage", None) + if usage is None: + logger.warning(f"[OpenAILLM] usage is None in response (id={getattr(response, 'id', '?')}); cost will not be recorded.") + return + cost_manager.update_cost(cost=self._compute_cost(usage), model=self.config.model) diff --git a/evoagentx/models/openrouter_model.py b/evoagentx/models/openrouter_model.py index b8520bcd..79ee3c9b 100644 --- a/evoagentx/models/openrouter_model.py +++ b/evoagentx/models/openrouter_model.py @@ -22,8 +22,8 @@ class OpenRouterLLM(BaseLLM): def init_model(self): - config: OpenRouterConfig = self.config - self._client = self._init_client(config) + self._client = None + self._async_client = None self._default_ignore_fields = ["llm_type", "openrouter_key", "openrouter_base", "openrouter_model_base", "output_response"] def _init_client(self, config: OpenRouterConfig): @@ -32,6 +32,24 @@ def _init_client(self, config: OpenRouterConfig): def _init_async_client(self, config: OpenRouterConfig): return AsyncOpenAI(api_key=config.openrouter_key, base_url=config.openrouter_base) + def ensure_client(self): + if self._client is None or self._client.is_closed(): + self._client = self._init_client(self.config) + return self._client + + def close_client(self): + if self._client is not None and not self._client.is_closed(): + self._client.close() + + def ensure_async_client(self): + if self._async_client is None or self._async_client.is_closed(): + self._async_client = self._init_async_client(self.config) + return self._async_client + + async def close_async_client(self): + if self._async_client is not None and not self._async_client.is_closed(): + await self._async_client.close() + def formulate_messages(self, prompts: List[str], system_messages: Optional[List[str]] = None) -> List[List[dict]]: if system_messages: assert len(prompts) == len(system_messages), f"the number of prompts ({len(prompts)}) is different from the number of system_messages ({len(system_messages)})" @@ -188,8 +206,9 @@ def single_generate(self, messages: List[dict], **kwargs) -> str: output_response = kwargs.get("output_response", self.config.output_response) try: + client = self.ensure_client() completion_params = self.get_completion_params(**kwargs) - response = self._client.chat.completions.create(messages=messages, **completion_params) + response = client.chat.completions.create(messages=messages, **completion_params) if stream: output = self.get_stream_output(response, output_response=output_response) else: @@ -207,7 +226,7 @@ async def single_generate_async(self, messages: List[dict], **kwargs) -> str: output_response = kwargs.get("output_response", self.config.output_response) try: - async_client = self._init_async_client(self.config) + async_client = self.ensure_async_client() completion_params = self.get_completion_params(**kwargs) response = await async_client.chat.completions.create( messages=messages, **completion_params @@ -215,9 +234,11 @@ async def single_generate_async(self, messages: List[dict], **kwargs) -> str: if stream: output = await self.get_stream_output_async(response, output_response=output_response) else: + # The network I/O is already awaited above; the response is fully in memory here, + # so this synchronous parsing/cost call does not block the event loop. output: str = self.get_completion_output(response=response, output_response=output_response) except Exception as e: raise RuntimeError(f"Error during single_generate_async of OpenRouterLLM: {str(e)}") - return output \ No newline at end of file + return output diff --git a/evoagentx/prompts/tool_calling.py b/evoagentx/prompts/tool_calling.py index 187a15e5..21e9a29a 100644 --- a/evoagentx/prompts/tool_calling.py +++ b/evoagentx/prompts/tool_calling.py @@ -1,7 +1,8 @@ +# todo: Switch back to once CustomizeAction supports the canonical tool-call tag. TOOL_CALL_FORMAT = """ - + {tool_calls} - + """ @@ -125,4 +126,4 @@ --- Invalid Output --- {text} --- End --- -""" \ No newline at end of file +""" diff --git a/pyproject.toml b/pyproject.toml index dffd7631..6ece3c90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "python-dotenv>=1.0.0", "requests>=2.28.0", "openai>=1.55.3", - "litellm>=1.75.2", + "litellm>=1.83.10", "dashscope>=1.23.4", "tree_sitter", "tree_sitter_python", @@ -54,6 +54,8 @@ rag = [ "llama-index-embeddings-azure-openai", "faiss-cpu==1.8.0.post1", "sentence-transformers", + "transformers>=4.41,<5", + "cryptography<49", "neo4j", "ollama", "docx2txt", @@ -67,6 +69,7 @@ tools = [ "selenium", "html2text", "fastmcp>=2.2.0,<3.0", + "cryptography<49", "PyPDF2", "Pillow", "exa-py>=2.0.0", @@ -111,6 +114,8 @@ all = [ "llama-index-embeddings-azure-openai", "faiss-cpu==1.8.0.post1", "sentence-transformers", + "transformers>=4.41,<5", + "cryptography<49", "neo4j", "ollama", "docx2txt", @@ -122,6 +127,7 @@ all = [ "selenium", "html2text", "fastmcp>=2.2.0,<3.0", + "cryptography<49", "PyPDF2", "Pillow", "exa-py>=2.0.0", diff --git a/pytest.ini b/pytest.ini index da636be2..cda836bc 100644 --- a/pytest.ini +++ b/pytest.ini @@ -2,3 +2,4 @@ testpaths = tests python_files = test_*.py addopts = -ra +asyncio_mode = auto diff --git a/requirements.txt b/requirements.txt index 6cfe5a71..bdfc4357 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,7 +18,7 @@ loguru>=0.7.3 python-dotenv>=1.0.0 requests>=2.28.0 openai>=1.55.3 -litellm>=1.75.2 +litellm>=1.83.10 dashscope>=1.23.4 tree_sitter tree_sitter_python @@ -34,6 +34,8 @@ llama-index-graph-stores-neo4j llama-index-embeddings-azure-openai faiss-cpu==1.8.0.post1 sentence-transformers +transformers>=4.41,<5 +cryptography<49 neo4j ollama docx2txt diff --git a/tests/src/models/mock_response.py b/tests/src/models/mock_response.py index a1ae225c..e5625a3b 100644 --- a/tests/src/models/mock_response.py +++ b/tests/src/models/mock_response.py @@ -1,12 +1,20 @@ +import json +from unittest.mock import MagicMock + from openai.types.completion_usage import CompletionUsage -from openai.types.chat.chat_completion_chunk import ChoiceDelta +from openai.types.chat.chat_completion_chunk import ChoiceDelta, ChoiceDeltaToolCall, ChoiceDeltaToolCallFunction from openai.types.chat.chat_completion_chunk import Choice as AChoice from openai.types.chat.chat_completion_chunk import ChatCompletionChunk from openai.types.chat.chat_completion import ChatCompletion, Choice, ChatCompletionMessage +from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function + + +# --------------------------------------------------------------------------- +# OpenAI — non-streaming helpers +# --------------------------------------------------------------------------- def get_openai_chat_completion() -> ChatCompletion: - - openai_chat_completion = ChatCompletion( + return ChatCompletion( id="xxxx", model="model_name", object="chat.completion", @@ -21,14 +29,44 @@ def get_openai_chat_completion() -> ChatCompletion: ], usage=CompletionUsage(completion_tokens=1, prompt_tokens=22, total_tokens=23), ) - return openai_chat_completion -def get_openai_chat_completion_chunk(usage_as_dict: bool = False) -> ChatCompletionChunk: +def get_openai_tool_call_completion() -> ChatCompletion: + return ChatCompletion( + id="tool_call_xxxx", + model="model_name", + object="chat.completion", + created=11111, + choices=[ + Choice( + finish_reason="tool_calls", + index=0, + message=ChatCompletionMessage( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_abc123", + type="function", + function=Function(name="get_weather", arguments='{"city": "Tokyo"}'), + ) + ], + ), + logprobs=None, + ) + ], + usage=CompletionUsage(completion_tokens=10, prompt_tokens=20, total_tokens=30), + ) + +# --------------------------------------------------------------------------- +# OpenAI — streaming helpers +# --------------------------------------------------------------------------- + +def get_openai_chat_completion_chunk(usage_as_dict: bool = False) -> ChatCompletionChunk: usage = CompletionUsage(completion_tokens=1, prompt_tokens=22, total_tokens=23) usage = usage if not usage_as_dict else usage.model_dump() - openai_chat_completion_chunk = ChatCompletionChunk( + return ChatCompletionChunk( id="xxxx", model="model_name", object="chat.completion.chunk", @@ -43,17 +81,269 @@ def get_openai_chat_completion_chunk(usage_as_dict: bool = False) -> ChatComplet ], usage=usage, ) - return openai_chat_completion_chunk +def get_openai_stream_chunks() -> list: + """Content chunks followed by a usage-only final chunk.""" + content_chunk = ChatCompletionChunk( + id="stream_xxxx", + model="model_name", + object="chat.completion.chunk", + created=11111, + choices=[ + AChoice( + delta=ChoiceDelta(role="assistant", content="Paris"), + finish_reason=None, + index=0, + logprobs=None, + ) + ], + ) + usage_chunk = ChatCompletionChunk( + id="stream_xxxx", + model="model_name", + object="chat.completion.chunk", + created=11111, + choices=[], + usage=CompletionUsage(completion_tokens=1, prompt_tokens=22, total_tokens=23), + ) + return [content_chunk, usage_chunk] + + +def get_openai_tool_call_chunks() -> list: + """Streaming tool-call: name chunk → arguments chunk → usage chunk.""" + name_chunk = ChatCompletionChunk( + id="tool_stream_xxxx", + model="model_name", + object="chat.completion.chunk", + created=11111, + choices=[ + AChoice( + delta=ChoiceDelta( + role="assistant", + content=None, + tool_calls=[ + ChoiceDeltaToolCall( + index=0, + id="call_abc123", + type="function", + function=ChoiceDeltaToolCallFunction(name="get_weather", arguments=""), + ) + ], + ), + finish_reason=None, + index=0, + logprobs=None, + ) + ], + ) + args_chunk = ChatCompletionChunk( + id="tool_stream_xxxx", + model="model_name", + object="chat.completion.chunk", + created=11111, + choices=[ + AChoice( + delta=ChoiceDelta( + tool_calls=[ + ChoiceDeltaToolCall( + index=0, + function=ChoiceDeltaToolCallFunction(arguments='{"city": "Tokyo"}'), + ) + ], + ), + finish_reason="tool_calls", + index=0, + logprobs=None, + ) + ], + ) + usage_chunk = ChatCompletionChunk( + id="tool_stream_xxxx", + model="model_name", + object="chat.completion.chunk", + created=11111, + choices=[], + usage=CompletionUsage(completion_tokens=10, prompt_tokens=20, total_tokens=30), + ) + return [name_chunk, args_chunk, usage_chunk] + + +# --------------------------------------------------------------------------- +# Async iterator wrapper (for async streaming mocks) +# --------------------------------------------------------------------------- + +class AsyncChunkIterator: + def __init__(self, chunks: list): + self._chunks = chunks + + def __aiter__(self): + return self._gen() + + async def _gen(self): + for chunk in self._chunks: + yield chunk + + +# --------------------------------------------------------------------------- +# OpenAI sync/async mock factories +# --------------------------------------------------------------------------- + default_resp = get_openai_chat_completion() default_resp_chunk = get_openai_chat_completion_chunk() -def mock_openai_completions_create(self, stream: bool=False, **kwargs): + +def mock_openai_completions_create(self, stream: bool = False, **kwargs): if stream: - class Iterator(object): - def __iter__(self): + class SyncIterator: + def __iter__(self_inner): yield default_resp_chunk - return Iterator() - else: - return default_resp \ No newline at end of file + return SyncIterator() + return default_resp + + +def mock_openai_stream_completions_create(self, stream: bool = False, **kwargs): + """Returns proper multi-chunk stream with final usage chunk.""" + if stream: + chunks = get_openai_stream_chunks() + class SyncIterator: + def __iter__(self_inner): + for c in chunks: + yield c + return SyncIterator() + return get_openai_chat_completion() + + +def mock_openai_tool_call_create(self, stream: bool = False, **kwargs): + if stream: + chunks = get_openai_tool_call_chunks() + class SyncIterator: + def __iter__(self_inner): + for c in chunks: + yield c + return SyncIterator() + return get_openai_tool_call_completion() + + +async def mock_async_openai_create(self, stream: bool = False, **kwargs): + if stream: + return AsyncChunkIterator(get_openai_stream_chunks()) + return get_openai_chat_completion() + + +async def mock_async_openai_tool_call_create(self, stream: bool = False, **kwargs): + if stream: + return AsyncChunkIterator(get_openai_tool_call_chunks()) + return get_openai_tool_call_completion() + + +# --------------------------------------------------------------------------- +# OpenRouter mock helpers (MagicMock-based to attach usage.cost) +# --------------------------------------------------------------------------- + +def _make_or_usage(prompt_tokens=22, completion_tokens=1, cost=0.000015): + usage = MagicMock() + usage.prompt_tokens = prompt_tokens + usage.completion_tokens = completion_tokens + usage.total_tokens = prompt_tokens + completion_tokens + usage.cost = cost + return usage + + +def get_openrouter_chat_completion(content="Paris", cost=0.000015) -> MagicMock: + resp = MagicMock() + resp.id = "or_xxxx" + resp.usage = _make_or_usage(cost=cost) + resp.choices[0].message.content = content + resp.choices[0].message.tool_calls = None + return resp + + +def get_openrouter_tool_call_completion(cost=0.000015) -> MagicMock: + resp = MagicMock() + resp.id = "or_tool_xxxx" + resp.usage = _make_or_usage(prompt_tokens=20, completion_tokens=10, cost=cost) + resp.choices[0].message.content = None + # Build a minimal tool-call object compatible with format_tool_calls + tc = MagicMock() + tc.id = "call_or123" + tc.function.name = "get_weather" + tc.function.arguments = '{"city": "Tokyo"}' + resp.choices[0].message.tool_calls = [tc] + return resp + + +def get_openrouter_stream_chunks(content="Paris", cost=0.000015) -> list: + content_chunk = MagicMock() + content_chunk.usage = None + content_chunk.choices = [MagicMock()] + content_chunk.choices[0].delta.content = content + content_chunk.choices[0].delta.tool_calls = None + + usage_chunk = MagicMock() + usage_chunk.usage = _make_or_usage(cost=cost) + usage_chunk.choices = [] + return [content_chunk, usage_chunk] + + +def get_openrouter_tool_call_stream_chunks(cost=0.000015) -> list: + name_chunk = MagicMock() + name_chunk.usage = None + name_chunk.choices = [MagicMock()] + name_chunk.choices[0].delta.content = None + tc_name = MagicMock() + tc_name.index = 0 + tc_name.id = "call_or123" + tc_name.function.name = "get_weather" + tc_name.function.arguments = "" + name_chunk.choices[0].delta.tool_calls = [tc_name] + + args_chunk = MagicMock() + args_chunk.usage = None + args_chunk.choices = [MagicMock()] + args_chunk.choices[0].delta.content = None + tc_args = MagicMock() + tc_args.index = 0 + tc_args.id = None + tc_args.function.name = None + tc_args.function.arguments = '{"city": "Tokyo"}' + args_chunk.choices[0].delta.tool_calls = [tc_args] + + usage_chunk = MagicMock() + usage_chunk.usage = _make_or_usage(prompt_tokens=20, completion_tokens=10, cost=cost) + usage_chunk.choices = [] + return [name_chunk, args_chunk, usage_chunk] + + +def mock_openrouter_completions_create(self, stream: bool = False, **kwargs): + if stream: + chunks = get_openrouter_stream_chunks() + class SyncIterator: + def __iter__(self_inner): + for c in chunks: + yield c + return SyncIterator() + return get_openrouter_chat_completion() + + +def mock_openrouter_tool_call_create(self, stream: bool = False, **kwargs): + if stream: + chunks = get_openrouter_tool_call_stream_chunks() + class SyncIterator: + def __iter__(self_inner): + for c in chunks: + yield c + return SyncIterator() + return get_openrouter_tool_call_completion() + + +async def mock_async_openrouter_create(self, stream: bool = False, **kwargs): + if stream: + return AsyncChunkIterator(get_openrouter_stream_chunks()) + return get_openrouter_chat_completion() + + +async def mock_async_openrouter_tool_call_create(self, stream: bool = False, **kwargs): + if stream: + return AsyncChunkIterator(get_openrouter_tool_call_stream_chunks()) + return get_openrouter_tool_call_completion() diff --git a/tests/src/models/test_openai_model.py b/tests/src/models/test_openai_model.py index 2985e713..7f342276 100644 --- a/tests/src/models/test_openai_model.py +++ b/tests/src/models/test_openai_model.py @@ -1,50 +1,219 @@ -from evoagentx.models import OpenAILLMConfig, OpenAILLM -from evoagentx.models import LLMOutputParser -from evoagentx.models import cost_manager - -from tests.src.models.mock_response import mock_openai_completions_create - - -def test_openai_generation(mocker): - - mocker.patch("openai.resources.chat.completions.Completions.create", mock_openai_completions_create) - - model_name = "gpt-4o-mini" - config = OpenAILLMConfig(model=model_name, openai_key="mock_openai_key", output_response=False) - model = OpenAILLM(config) - - prompt = "what is the capital city of China. Only output the answer." - system_prompt = "You are an expert in geography" - - # test different input formats - output = model.generate(prompt=prompt, system_message=system_prompt) - assert isinstance(output, LLMOutputParser) - assert output.content == "Beijing" - assert str(output) == "Beijing" - assert cost_manager.total_tokens[model_name] == 23 - - output = model.generate(prompt=[prompt], system_message=[system_prompt]) - assert isinstance(output, list) and isinstance(output[0], LLMOutputParser) - assert output[0].content == "Beijing" - assert str(output[0]) == "Beijing" - assert cost_manager.total_tokens[model_name] == 23*2 - - output = model.generate(messages=[{'role': 'system', 'content': system_prompt}, {'role': 'user', 'content': prompt}]) - assert isinstance(output, LLMOutputParser) - assert output.content == "Beijing" - assert str(output) == "Beijing" - assert cost_manager.total_tokens[model_name] == 23*3 - - output = model.generate(messages=[[{'role': 'system', 'content': system_prompt}, {'role': 'user', 'content': prompt}]]) - assert isinstance(output, list) and isinstance(output[0], LLMOutputParser) - assert output[0].content == "Beijing" - assert str(output[0]) == "Beijing" - assert cost_manager.total_tokens[model_name] == 23*4 - - # test stream output - output = model.generate(prompt=prompt, system_message=system_prompt, stream=True) - assert isinstance(output, LLMOutputParser) - assert output.content == "Beijing" - assert str(output) == "Beijing" - assert cost_manager.total_tokens[model_name] > 23*4 +import pytest +from evoagentx.models import OpenAILLMConfig, OpenAILLM, LLMOutputParser +from evoagentx.models.model_utils import cost_manager + +from tests.src.models.mock_response import ( + mock_openai_completions_create, + mock_openai_stream_completions_create, + mock_openai_tool_call_create, + mock_async_openai_create, + mock_async_openai_tool_call_create, +) + +OPENAI_MODEL = "gpt-4o-mini" +SYNC_PATCH = "openai.resources.chat.completions.Completions.create" +ASYNC_PATCH = "openai.resources.chat.completions.AsyncCompletions.create" + +GET_WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a given city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["city"], + }, + }, +} + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def reset_cost_manager(): + cost_manager.input_tokens.clear() + cost_manager.output_tokens.clear() + cost_manager.total_tokens.clear() + cost_manager.cost_per_model.clear() + yield + + +def _make_llm(**kwargs) -> OpenAILLM: + config = OpenAILLMConfig( + model=OPENAI_MODEL, + openai_key="mock_openai_key", + output_response=False, + **kwargs, + ) + return OpenAILLM(config=config) + + +def _assert_cost_updated(model: str = OPENAI_MODEL): + assert cost_manager.total_tokens[model] > 0, "No tokens recorded for cost tracking" + + +# --------------------------------------------------------------------------- +# 1. Sync — non-streaming (input format variants) +# --------------------------------------------------------------------------- + +def test_sync_non_stream_prompt(mocker): + mocker.patch(SYNC_PATCH, mock_openai_completions_create) + llm = _make_llm(stream=False) + prompt = "What is the capital of China?" + system = "You are a geography expert." + + out = llm.generate(prompt=prompt, system_message=system) + assert isinstance(out, LLMOutputParser) + assert out.content == "Beijing" + + out = llm.generate(prompt=[prompt], system_message=[system]) + assert isinstance(out, list) + assert out[0].content == "Beijing" + + out = llm.generate(messages=[{"role": "system", "content": system}, {"role": "user", "content": prompt}]) + assert isinstance(out, LLMOutputParser) + assert out.content == "Beijing" + + out = llm.generate(messages=[[{"role": "system", "content": system}, {"role": "user", "content": prompt}]]) + assert isinstance(out, list) + assert out[0].content == "Beijing" + + _assert_cost_updated() + assert cost_manager.total_tokens[OPENAI_MODEL] == 23 * 4 + + +# --------------------------------------------------------------------------- +# 2. Sync — streaming +# --------------------------------------------------------------------------- + +def test_sync_stream(mocker): + mocker.patch(SYNC_PATCH, mock_openai_stream_completions_create) + llm = _make_llm(stream=True) + out = llm.generate(prompt="What is the capital of France?") + assert isinstance(out, LLMOutputParser) + assert out.content == "Paris" + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 3. Sync — tool call (non-streaming) +# --------------------------------------------------------------------------- + +def test_sync_tool_call_non_stream(mocker): + mocker.patch(SYNC_PATCH, mock_openai_tool_call_create) + llm = _make_llm(stream=False, tools=[GET_WEATHER_TOOL], tool_choice="auto") + out = llm.generate(prompt="What is the weather in Tokyo?") + assert isinstance(out, LLMOutputParser) + assert "" in out.content + assert "get_weather" in out.content + assert "Tokyo" in out.content + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 4. Sync — tool call (streaming) +# --------------------------------------------------------------------------- + +def test_sync_tool_call_stream(mocker): + mocker.patch(SYNC_PATCH, mock_openai_tool_call_create) + llm = _make_llm(stream=True, tools=[GET_WEATHER_TOOL], tool_choice="auto") + out = llm.generate(prompt="What is the weather in Tokyo?") + assert isinstance(out, LLMOutputParser) + assert "" in out.content + assert "get_weather" in out.content + assert "Tokyo" in out.content + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 5. Async — non-streaming +# --------------------------------------------------------------------------- + +async def test_async_non_stream(mocker): + mocker.patch(ASYNC_PATCH, mock_async_openai_create) + llm = _make_llm(stream=False) + out = await llm.async_generate(prompt="What is the capital of China?") + assert isinstance(out, LLMOutputParser) + assert out.content == "Beijing" + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 6. Async — streaming +# --------------------------------------------------------------------------- + +async def test_async_stream(mocker): + mocker.patch(ASYNC_PATCH, mock_async_openai_create) + llm = _make_llm(stream=True) + out = await llm.async_generate(prompt="What is the capital of France?") + assert isinstance(out, LLMOutputParser) + assert out.content == "Paris" + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 7. Async — tool call (non-streaming) +# --------------------------------------------------------------------------- + +async def test_async_tool_call_non_stream(mocker): + mocker.patch(ASYNC_PATCH, mock_async_openai_tool_call_create) + llm = _make_llm(stream=False, tools=[GET_WEATHER_TOOL], tool_choice="auto") + out = await llm.async_generate(prompt="What is the weather in Tokyo?") + assert isinstance(out, LLMOutputParser) + assert "" in out.content + assert "get_weather" in out.content + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 8. Async — tool call (streaming) +# --------------------------------------------------------------------------- + +async def test_async_tool_call_stream(mocker): + mocker.patch(ASYNC_PATCH, mock_async_openai_tool_call_create) + llm = _make_llm(stream=True, tools=[GET_WEATHER_TOOL], tool_choice="auto") + out = await llm.async_generate(prompt="What is the weather in Tokyo?") + assert isinstance(out, LLMOutputParser) + assert "" in out.content + assert "get_weather" in out.content + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 9. stream_options include_usage explicitly set +# --------------------------------------------------------------------------- + +def test_stream_with_include_usage(mocker): + mocker.patch(SYNC_PATCH, mock_openai_stream_completions_create) + llm = _make_llm(stream=True, stream_options={"include_usage": True}) + out = llm.generate(prompt="What is the capital of France?") + assert isinstance(out, LLMOutputParser) + assert out.content == "Paris" + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 10. Cost tracking — verify token counts are accumulated correctly +# --------------------------------------------------------------------------- + +def test_cost_accumulation(mocker): + mocker.patch(SYNC_PATCH, mock_openai_completions_create) + llm = _make_llm(stream=False) + prompt = "Test prompt" + + llm.generate(prompt=prompt) + tokens_after_1 = cost_manager.total_tokens[OPENAI_MODEL] + + llm.generate(prompt=prompt) + tokens_after_2 = cost_manager.total_tokens[OPENAI_MODEL] + + assert tokens_after_2 == tokens_after_1 * 2 + assert cost_manager.input_tokens[OPENAI_MODEL] > 0 + assert cost_manager.output_tokens[OPENAI_MODEL] > 0 diff --git a/tests/src/models/test_openrouter_model.py b/tests/src/models/test_openrouter_model.py new file mode 100644 index 00000000..1a27e9d3 --- /dev/null +++ b/tests/src/models/test_openrouter_model.py @@ -0,0 +1,221 @@ +import pytest + +from evoagentx.models.model_configs import OpenRouterConfig +from evoagentx.models.openrouter_model import OpenRouterLLM +from evoagentx.models import LLMOutputParser +from evoagentx.models.model_utils import cost_manager + +from tests.src.models.mock_response import ( + mock_openrouter_completions_create, + mock_openrouter_tool_call_create, + mock_async_openrouter_create, + mock_async_openrouter_tool_call_create, +) + +OPENROUTER_MODEL = "openai/gpt-4o-mini" +SYNC_PATCH = "openai.resources.chat.completions.Completions.create" +ASYNC_PATCH = "openai.resources.chat.completions.AsyncCompletions.create" + +GET_WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a given city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["city"], + }, + }, +} + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def reset_cost_manager(): + cost_manager.input_tokens.clear() + cost_manager.output_tokens.clear() + cost_manager.total_tokens.clear() + cost_manager.cost_per_model.clear() + yield + + +def _make_llm(**kwargs) -> OpenRouterLLM: + config = OpenRouterConfig( + model=OPENROUTER_MODEL, + openrouter_key="mock_or_key", + output_response=False, + **kwargs, + ) + return OpenRouterLLM(config=config) + + +def _assert_cost_updated(model: str = OPENROUTER_MODEL): + assert cost_manager.total_tokens[model] > 0, "No tokens recorded" + assert cost_manager.cost_per_model[model] > 0, "No cost recorded" + + +# --------------------------------------------------------------------------- +# 1. Sync — non-streaming +# --------------------------------------------------------------------------- + +def test_sync_non_stream(mocker): + mocker.patch(SYNC_PATCH, mock_openrouter_completions_create) + llm = _make_llm(stream=False) + out = llm.generate(prompt="What is the capital of France?") + assert isinstance(out, LLMOutputParser) + assert out.content == "Paris" + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 2. Sync — streaming +# --------------------------------------------------------------------------- + +def test_sync_stream(mocker): + mocker.patch(SYNC_PATCH, mock_openrouter_completions_create) + llm = _make_llm(stream=True) + out = llm.generate(prompt="What is the capital of France?") + assert isinstance(out, LLMOutputParser) + assert out.content == "Paris" + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 3. Sync — tool call (non-streaming) +# --------------------------------------------------------------------------- + +def test_sync_tool_call_non_stream(mocker): + mocker.patch(SYNC_PATCH, mock_openrouter_tool_call_create) + llm = _make_llm(stream=False, tools=[GET_WEATHER_TOOL], tool_choice="auto") + out = llm.generate(prompt="What is the weather in Tokyo?") + assert isinstance(out, LLMOutputParser) + assert "" in out.content + assert "get_weather" in out.content + assert "Tokyo" in out.content + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 4. Sync — tool call (streaming) +# --------------------------------------------------------------------------- + +def test_sync_tool_call_stream(mocker): + mocker.patch(SYNC_PATCH, mock_openrouter_tool_call_create) + llm = _make_llm(stream=True, tools=[GET_WEATHER_TOOL], tool_choice="auto") + out = llm.generate(prompt="What is the weather in Tokyo?") + assert isinstance(out, LLMOutputParser) + assert "" in out.content + assert "get_weather" in out.content + assert "Tokyo" in out.content + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 5. Async — non-streaming +# --------------------------------------------------------------------------- + +async def test_async_non_stream(mocker): + mocker.patch(ASYNC_PATCH, mock_async_openrouter_create) + llm = _make_llm(stream=False) + out = await llm.async_generate(prompt="What is the capital of France?") + assert isinstance(out, LLMOutputParser) + assert out.content == "Paris" + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 6. Async — streaming +# --------------------------------------------------------------------------- + +async def test_async_stream(mocker): + mocker.patch(ASYNC_PATCH, mock_async_openrouter_create) + llm = _make_llm(stream=True) + out = await llm.async_generate(prompt="What is the capital of France?") + assert isinstance(out, LLMOutputParser) + assert out.content == "Paris" + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 7. Async — tool call (non-streaming) +# --------------------------------------------------------------------------- + +async def test_async_tool_call_non_stream(mocker): + mocker.patch(ASYNC_PATCH, mock_async_openrouter_tool_call_create) + llm = _make_llm(stream=False, tools=[GET_WEATHER_TOOL], tool_choice="auto") + out = await llm.async_generate(prompt="What is the weather in Tokyo?") + assert isinstance(out, LLMOutputParser) + assert "" in out.content + assert "get_weather" in out.content + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 8. Async — tool call (streaming) +# --------------------------------------------------------------------------- + +async def test_async_tool_call_stream(mocker): + mocker.patch(ASYNC_PATCH, mock_async_openrouter_tool_call_create) + llm = _make_llm(stream=True, tools=[GET_WEATHER_TOOL], tool_choice="auto") + out = await llm.async_generate(prompt="What is the weather in Tokyo?") + assert isinstance(out, LLMOutputParser) + assert "" in out.content + assert "get_weather" in out.content + _assert_cost_updated() + + +# --------------------------------------------------------------------------- +# 9. Cost — missing usage.cost logs warning and records 0 +# --------------------------------------------------------------------------- + +def test_missing_cost_warns(mocker): + """When usage.cost is absent, cost is recorded as 0 without raising.""" + from unittest.mock import MagicMock + + def mock_no_cost_create(self, stream=False, **kwargs): + resp = MagicMock() + resp.id = "or_no_cost" + usage = MagicMock() + usage.prompt_tokens = 10 + usage.completion_tokens = 5 + usage.total_tokens = 15 + del usage.cost # force getattr to fall back to MagicMock default + usage.cost = None # explicitly None → triggers the warning path + resp.usage = usage + resp.choices[0].message.content = "Hello" + resp.choices[0].message.tool_calls = None + return resp + + mocker.patch(SYNC_PATCH, mock_no_cost_create) + llm = _make_llm(stream=False) + out = llm.generate(prompt="Hello") + assert isinstance(out, LLMOutputParser) + assert cost_manager.cost_per_model[OPENROUTER_MODEL] == 0.0 + assert cost_manager.total_tokens[OPENROUTER_MODEL] == 15 + + +# --------------------------------------------------------------------------- +# 10. Cost accumulation across multiple calls +# --------------------------------------------------------------------------- + +def test_cost_accumulation(mocker): + mocker.patch(SYNC_PATCH, mock_openrouter_completions_create) + llm = _make_llm(stream=False) + + llm.generate(prompt="Call 1") + tokens_1 = cost_manager.total_tokens[OPENROUTER_MODEL] + cost_1 = cost_manager.cost_per_model[OPENROUTER_MODEL] + + llm.generate(prompt="Call 2") + tokens_2 = cost_manager.total_tokens[OPENROUTER_MODEL] + cost_2 = cost_manager.cost_per_model[OPENROUTER_MODEL] + + assert tokens_2 == tokens_1 * 2 + assert cost_2 == pytest.approx(cost_1 * 2) From f2e23f8f856cfbf7e3645701b88cdabd034ac74c Mon Sep 17 00:00:00 2001 From: jinyuan Date: Wed, 24 Jun 2026 15:40:34 +0100 Subject: [PATCH 7/9] update LiteLLM to be compatible with OpenAILLM --- evoagentx/models/litellm_model.py | 108 ++++++++++++++---------------- 1 file changed, 52 insertions(+), 56 deletions(-) diff --git a/evoagentx/models/litellm_model.py b/evoagentx/models/litellm_model.py index 73d574e5..1a3b5a7e 100644 --- a/evoagentx/models/litellm_model.py +++ b/evoagentx/models/litellm_model.py @@ -7,6 +7,7 @@ ) from litellm import completion, acompletion from typing import List +from ..core.logging import logger from ..core.registry import register_model from .model_configs import LiteLLMConfig from .openai_model import OpenAILLM @@ -87,102 +88,97 @@ def init_model(self): "groq_key", "api_base", "is_local", "azure_endpoint", "azure_key", "api_version", "api_key" ] # parameters in LiteLLMConfig that are not LiteLLM models' input parameters - def _compute_cost(self, input_tokens: int, output_tokens: int) -> Cost: + def _apply_provider_params(self, completion_params: dict) -> dict: + """Inject provider-specific routing parameters (local / Azure) into the + LiteLLM completion params. OpenAI and the remaining providers are routed + purely through the environment variables set in ``init_model``.""" + company = infer_litellm_company_from_model(self.model) + if self.config.is_local or company == "local": # route local model through its api_base + completion_params["api_base"] = self.api_base + completion_params["api_key"] = self.api_key + elif company == "azure": # Add Azure OpenAI specific parameters + completion_params["api_base"] = self.config.azure_endpoint + completion_params["api_version"] = self.config.api_version + completion_params["api_key"] = self.config.azure_key + return completion_params + + def _compute_cost(self, usage) -> Cost: + input_tokens = getattr(usage, "prompt_tokens", 0) or 0 + output_tokens = getattr(usage, "completion_tokens", 0) or 0 + # Local models are free; LiteLLM has no pricing for them. if self.config.is_local: return Cost(input_tokens=input_tokens, output_tokens=output_tokens, input_cost=0.0, output_cost=0.0) - return super()._compute_cost(input_tokens, output_tokens) + try: + return super()._compute_cost(usage) + except Exception as e: + # Unlike OpenAILLM (which validates the model against the price map at + # init), LiteLLM accepts arbitrary provider/model names, and + # ``litellm.cost_per_token`` raises for any model missing from LiteLLM's + # price map. Since cost is computed inside the generation path, a pricing + # gap must not abort an otherwise-successful (and possibly retried) call — + # record the tokens and fall back to zero cost. + logger.warning( + f"[LiteLLM] Could not compute cost for model '{self.config.model}': {e}. " + "Recording tokens with zero cost." + ) + return Cost(input_tokens=input_tokens, output_tokens=output_tokens, input_cost=0.0, output_cost=0.0) @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5)) def single_generate(self, messages: List[dict], **kwargs) -> str: """ - Generate a single response using the completion function. + Generate a single response using the LiteLLM completion function. - Args: + Args: messages (List[dict]): A list of dictionaries representing the conversation history. **kwargs (Any): Additional parameters to be passed to the `completion` function. - - Returns: + + Returns: str: A string containing the model's response. """ - stream = kwargs["stream"] if "stream" in kwargs else self.config.stream - output_response = kwargs["output_response"] if "output_response" in kwargs else self.config.output_response + stream = kwargs.get("stream", self.config.stream) + output_response = kwargs.get("output_response", self.config.output_response) try: completion_params = self.get_completion_params(**kwargs) - company = infer_litellm_company_from_model(self.model) - if self.config.is_local or company == "local": # update save api_base for local model - completion_params["api_base"] = self.api_base - elif company == "azure": # Add Azure OpenAI specific parameters - completion_params["api_base"] = self.config.azure_endpoint - completion_params["api_version"] = self.config.api_version - completion_params["api_key"] = self.config.azure_key + self._apply_provider_params(completion_params) response = completion(messages=messages, **completion_params) + # get_stream_output / get_completion_output record cost internally via _update_cost. if stream: output = self.get_stream_output(response, output_response=output_response) - cost = self._stream_cost(messages=messages, output=output) else: output: str = self.get_completion_output(response=response, output_response=output_response) - cost = self._completion_cost(response=response) - self._update_cost(cost=cost) - except Exception as e: - raise RuntimeError(f"Error during single_generate: {str(e)}") - + raise RuntimeError(f"Error during single_generate of LiteLLM: {str(e)}") + return output - - def batch_generate(self, batch_messages: List[List[dict]], **kwargs) -> List[str]: - """ - Generate responses for a batch of messages. - Args: - batch_messages (List[List[dict]]): A list of message lists, where each sublist represents a conversation. - **kwargs (Any): Additional parameters to be passed to the `completion` function. - - Returns: - List[str]: A list of responses for each conversation. - """ - results = [] - for messages in batch_messages: - response = self.single_generate(messages, **kwargs) - results.append(response) - return results - async def single_generate_async(self, messages: List[dict], **kwargs) -> str: """ - Generate a single response using the async completion function. + Generate a single response using the async LiteLLM completion function. - Args: + Args: messages (List[dict]): A list of dictionaries representing the conversation history. **kwargs (Any): Additional parameters to be passed to the `completion` function. - - Returns: + + Returns: str: A string containing the model's response. """ - stream = kwargs["stream"] if "stream" in kwargs else self.config.stream - output_response = kwargs["output_response"] if "output_response" in kwargs else self.config.output_response + stream = kwargs.get("stream", self.config.stream) + output_response = kwargs.get("output_response", self.config.output_response) try: completion_params = self.get_completion_params(**kwargs) - company = infer_litellm_company_from_model(self.model) - if self.config.is_local or company == "local": # add api base for local model - completion_params["api_base"] = self.api_base - elif company == "azure": # Add Azure OpenAI specific parameters - completion_params["api_base"] = self.config.azure_endpoint - completion_params["api_version"] = self.config.api_version - completion_params["api_key"] = self.config.azure_key + self._apply_provider_params(completion_params) response = await acompletion(messages=messages, **completion_params) if stream: if hasattr(response, "__aiter__"): output = await self.get_stream_output_async(response, output_response=output_response) else: output = self.get_stream_output(response, output_response=output_response) - cost = self._stream_cost(messages=messages, output=output) else: output: str = self.get_completion_output(response=response, output_response=output_response) - cost = self._completion_cost(response=response) - self._update_cost(cost=cost) except Exception as e: - raise RuntimeError(f"Error during single_generate_async: {str(e)}") - + raise RuntimeError(f"Error during single_generate_async of LiteLLM: {str(e)}") + return output From ffc42362824bee3984d0bb2a002951be30dbef9f Mon Sep 17 00:00:00 2001 From: jinyuan Date: Wed, 24 Jun 2026 16:24:22 +0100 Subject: [PATCH 8/9] update SiliconFlowLLM to be compatible with OpenAILLM --- evoagentx/models/siliconflow_model.py | 216 +++++---------------- evoagentx/models/siliconflow_model_cost.py | 132 ------------- 2 files changed, 50 insertions(+), 298 deletions(-) delete mode 100644 evoagentx/models/siliconflow_model_cost.py diff --git a/evoagentx/models/siliconflow_model.py b/evoagentx/models/siliconflow_model.py index 33ec14a0..12ea71c4 100644 --- a/evoagentx/models/siliconflow_model.py +++ b/evoagentx/models/siliconflow_model.py @@ -1,179 +1,63 @@ -import asyncio -from tenacity import ( - retry, - stop_after_attempt, - wait_random_exponential, -) -from typing import List, Tuple +from typing import Union + +from openai import AsyncOpenAI, OpenAI +from openai.types.chat import ChatCompletion, ChatCompletionChunk from .openai_model import OpenAILLM from .model_configs import SiliconFlowConfig +from ..core.logging import logger from ..core.registry import register_model -from openai import OpenAI, Stream -# from loguru import logger from .model_utils import Cost, cost_manager -from openai.types.chat import ChatCompletion -from .siliconflow_model_cost import model_cost - -@register_model(config_cls=SiliconFlowConfig, alias=["siliconflow"]) -class SiliconFlowLLM(OpenAILLM): - - def init_model(self): - config: SiliconFlowConfig = self.config - self._client = self._init_client(config) # OpenAI(api_key=config.siliconflow_key, base_url="https://api.siliconflow.cn/v1") - self._default_ignore_fields = ["llm_type", "siliconflow_key", "output_response"] # parameters in SiliconFlowConfig that are not OpenAI models' input parameters - self._last_response = None - - def _init_client(self, config: SiliconFlowConfig): - client = OpenAI(api_key=config.siliconflow_key, base_url="https://api.siliconflow.cn/v1") - return client - @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5)) - def single_generate(self, messages: List[dict], **kwargs) -> str: +# SiliconFlow exposes an OpenAI-compatible API. Use the international endpoint +# (api.siliconflow.com), not the mainland one (api.siliconflow.cn). +SILICONFLOW_BASE_URL = "https://api.siliconflow.com/v1" - stream = kwargs["stream"] if "stream" in kwargs else self.config.stream - output_response = kwargs["output_response"] if "output_response" in kwargs else self.config.output_response - try: - completion_params = self.get_completion_params(**kwargs) - response = self._client.chat.completions.create( - messages=messages, - **completion_params - ) - if stream: - output, stream_response = self._get_stream_output_with_response(response, output_response=output_response) - cost = self._completion_cost(stream_response) - else: - output: str = response.choices[0].message.content - cost = self._completion_cost(response) - self._last_response = response - if output_response: - print(output) - self._update_cost(cost=cost) - except Exception as e: - if "account balance is insufficient" in str(e): - print("Warning: Account balance insufficient. Please recharge your account.") - return "" - raise RuntimeError(f"Error during single_generate of OpenAILLM: {str(e)}") - - return output +@register_model(config_cls=SiliconFlowConfig, alias=["siliconflow"]) +class SiliconFlowLLM(OpenAILLM): + """SiliconFlow LLM client. - async def single_generate_async(self, messages: List[dict], **kwargs) -> str: + SiliconFlow speaks the OpenAI chat-completions protocol, so this reuses all of + ``OpenAILLM``'s generation/streaming/tool-call logic and only overrides client + construction and cost handling. - stream = kwargs.get("stream", self.config.stream) - output_response = kwargs.get("output_response", self.config.output_response) + Unlike OpenRouter, SiliconFlow does not return ``usage.cost`` in its responses, + and LiteLLM has no pricing data for SiliconFlow-hosted models, so the dollar + cost cannot be computed or approximated. Token counts are still tracked; the + per-model cost is recorded as 0. A warning is emitted once at init time. + """ - try: - isolated_client = self._init_client(self.config) - completion_params = self.get_completion_params(**kwargs) + def init_model(self): + self._client = None + self._async_client = None + # parameters in SiliconFlowConfig that are not SiliconFlow models' input parameters + self._default_ignore_fields = ["llm_type", "siliconflow_key", "output_response"] + logger.warning( + "[SiliconFlowLLM] SiliconFlow does not report usage.cost and LiteLLM has no " + "pricing data for SiliconFlow models, so dollar cost cannot be computed. " + "Token usage will be tracked, but cost will be recorded as 0." + ) - loop = asyncio.get_event_loop() - response = await loop.run_in_executor( - None, - lambda: isolated_client.chat.completions.create( - messages=messages, - **completion_params - ) + def _init_client(self, config: SiliconFlowConfig): + return OpenAI(api_key=config.siliconflow_key, base_url=SILICONFLOW_BASE_URL) + + def _init_async_client(self, config: SiliconFlowConfig): + return AsyncOpenAI(api_key=config.siliconflow_key, base_url=SILICONFLOW_BASE_URL) + + def _update_cost(self, response: Union[ChatCompletion, ChatCompletionChunk]): + # Override OpenAILLM's LiteLLM-based cost computation: only record token + # counts and leave cost at 0 (see class docstring). + usage = getattr(response, "usage", None) + if usage is None: + logger.warning( + f"[SiliconFlowLLM] usage is None in response (id={getattr(response, 'id', '?')}); " + "tokens will not be recorded." ) - - if stream: - output, stream_response = self.get_stream_output(response, output_response=output_response) - cost = self._completion_cost(stream_response) - else: - output: str = response.choices[0].message.content - cost = self._completion_cost(response) - self._last_response = response - if output_response: - print(output) - self._update_cost(cost=cost) - except Exception as e: - if "account balance is insufficient" in str(e): - print("Warning: Account balance insufficient. Please recharge your account.") - return "" - raise RuntimeError(f"Error during single_generate_async of SiliconFlowLLM: {str(e)}") - - return output - - - def _completion_cost(self, response: ChatCompletion) -> Cost: - input_tokens = response.usage.prompt_tokens - output_tokens = response.usage.completion_tokens - return self._compute_cost(input_tokens=input_tokens, output_tokens=output_tokens) - - - def _compute_cost(self, input_tokens: int, output_tokens: int) -> Cost: - model: str = self.config.model - # total_tokens = input_tokens + output_tokens - if model not in model_cost: - return Cost(input_tokens=input_tokens, output_tokens=output_tokens, input_cost=0.0, output_cost=0.0) - - if "token_cost" in model_cost[model]: - # total_cost = total_tokens * model_cost[model]["token_cost"] / 1e6 - input_cost = input_tokens * model_cost[model]["token_cost"] / 1e6 - output_cost = output_tokens * model_cost[model]["token_cost"] / 1e6 - else: - # total_cost = input_tokens * model_cost[model]["input_token_cost"] / 1e6 + output_tokens * model_cost[model]["output_token_cost"] / 1e6 - input_cost = input_tokens * model_cost[model]["input_token_cost"] / 1e6 - output_cost = output_tokens * model_cost[model]["output_token_cost"] / 1e6 - - return Cost(input_tokens=input_tokens, output_tokens=output_tokens, input_cost=input_cost, output_cost=output_cost) - - - def get_cost(self) -> dict: - cost_info = {} - if self._last_response is None: - cost_info["error"] = "No response available yet — no generation has been performed" - return cost_info - try: - tokens = self._last_response.usage - if tokens.prompt_tokens == -1: - cost_info["note"] = "Token counts not available in stream mode" - cost_info["prompt_tokens"] = 0 - cost_info["completion_tokens"] = 0 - cost_info["total_tokens"] = 0 - else: - cost_info["prompt_tokens"] = tokens.prompt_tokens - cost_info["completion_tokens"] = tokens.completion_tokens - cost_info["total_tokens"] = tokens.total_tokens - except Exception as e: - print(f"Error during get_cost of SiliconFlow: {str(e)}") - cost_info["error"] = str(e) - - return cost_info - - def get_stream_output(self, response: Stream, output_response: bool=True) -> str: - output, _ = self._get_stream_output_with_response(response, output_response=output_response) - return output - - def _get_stream_output_with_response(self, response: Stream, output_response: bool=True) -> Tuple[str, object]: - output = "" - last_chunk = None - for chunk in response: - content = chunk.choices[0].delta.content - if content: - if output_response: - print(content, end="", flush=True) - output += content - last_chunk = chunk - - if output_response: - print("") - - # Build response object from the last chunk's usage info - if last_chunk is not None and hasattr(last_chunk, 'usage') and last_chunk.usage is not None: - stream_response = last_chunk - else: - # Create a placeholder response object for stream mode - stream_response = type('StreamResponse', (), { - 'usage': type('StreamUsage', (), { - 'prompt_tokens': -1, - 'completion_tokens': -1, - 'total_tokens': -1 - }) - }) - - self._last_response = stream_response - return output, stream_response - - def _update_cost(self, cost: Cost): - cost_manager.update_cost(cost=cost, model=self.config.model) \ No newline at end of file + return + cost = Cost( + input_tokens=usage.prompt_tokens, + output_tokens=usage.completion_tokens, + cost=0.0, + ) + cost_manager.update_cost(cost=cost, model=self.config.model) diff --git a/evoagentx/models/siliconflow_model_cost.py b/evoagentx/models/siliconflow_model_cost.py deleted file mode 100644 index 9355871f..00000000 --- a/evoagentx/models/siliconflow_model_cost.py +++ /dev/null @@ -1,132 +0,0 @@ -model_cost = { - "deepseek-ai/DeepSeek-R1": { - "input_token_cost": 4, - "output_token_cost": 16, - }, - "deepseek-ai/DeepSeek-V3": { - "input_token_cost": 2, - "output_token_cost": 8, - }, - "Pro/deepseek-ai/DeepSeek-R1":{ - "input_token_cost": 4, - "output_token_cost": 16, - }, - "Pro/deepseek-ai/DeepSeek-V3": { - "input_token_cost": 2, - "output_token_cost": 8, - }, - "deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { - "token_cost": 4.13 - }, - "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { - "token_cost": 0.7 - }, - "Qwen/QVQ-72B-Preview": { - "token_cost": 9.9 - }, - "deepseek-ai/DeepSeek-V2.5": { - "token_cost": 1.33 - }, - "meta-llama/Llama-3.3-70B-Instruct": { - "token_cost": 4.13, - }, - "Qwen/QwQ-32B-Preview": { - "token_cost": 1.26 - }, - "Qwen/Qwen2.5-Coder-32B-Instruct": { - "token_cost": 1.26 - }, - "Qwen/Qwen2-VL-72B-Instruct": { - "token_cost": 4.13 - }, - "OpenGVLab/InternVL2-26B": { - "token_cost": 1 - }, - "Qwen/Qwen2.5-72B-Instruct-128K": { - "token_cost": 4.13 - }, - "deepseek-ai/deepseek-vl2": { - "token_cost": 0.99 - }, - "Qwen/Qwen2.5-72B-Instruct": { - "token_cost": 4.13 - }, - "Qwen/Qwen2.5-32B-Instruct": { - "token_cost": 1.26 - }, - "Qwen/Qwen2.5-14B-Instruct": { - "token_cost": 0.7 - }, - "TeleAI/TeleChat2": { - "token_cost": 1.33 - }, - "internlm/internlm2_5-20b-chat": { - "token_cost": 1 - }, - "meta-llama/Meta-Llama-3.1-405B-Instruct": { - "token_cost": 21 - }, - "meta-llama/Meta-Llama-3.1-70B-Instruct": { - "token_cost": 4.13 - }, - "01-ai/Yi-1.5-34B-Chat-16K": { - "token_cost": 1.26 - }, - "google/gemma-2-27b-it": { - "token_cost": 1.26 - }, - "LoRA/meta-llama/Meta-Llama-3.1-8B-Instruct": { - "token_cost": 0.63 - }, - "LoRA/Qwen/Qwen2.5-32B-Instruct": { - "token_cost": 0.63 - }, - "LoRA/Qwen/Qwen2.5-14B-Instruct": { - "token_cost": 1.05 - }, - "Vendor-A/Qwen/Qwen2.5-72B-Instruct": { - "token_cost": 1 - }, - "Pro/deepseek-ai/DeepSeek-R1-Distill-Llama-8B": { - "token_cost": 0.42 - }, - "Pro/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B": { - "token_cost": 0.35 - }, - "Pro/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": { - "token_cost": 0.14 - }, - "Pro/Qwen/Qwen2.5-Coder-7B-Instruct": { - "token_cost": 0.35 - }, - "Pro/Qwen/Qwen2-VL-7B-Instruct": { - "token_cost": 0.35 - }, - "Pro/OpenGVLab/InternVL2-8B": { - "token_cost": 0.35 - }, - "Pro/Qwen/Qwen2.5-7B-Instruct": { - "token_cost": 0.35 - }, - "Pro/meta-llama/Meta-Llama-3.1-8B-Instruct": { - "token_cost": 0.42 - }, - "LoRA/Qwen/Qwen2.5-72B-Instruct": { - "token_cost": 6.2 - }, - "Pro/Qwen/Qwen2-7B-Instruct": { - "token_cost": 0.35 - }, - "Pro/Qwen/Qwen2-1.5B-Instruct": { - "token_cost": 0.14 - }, - "LoRA/Qwen/Qwen2.5-7B-Instruct": { - "token_cost": 0.53 - }, - "Pro/THUDM/glm-4-9b-chat": { - "token_cost": 0.6 - }, - "Pro/google/gemma-2-9b-it": { - "token_cost": 0.6 - } -} \ No newline at end of file From 7dbd71c88f639d63492a99bc34e2573be9e125eb Mon Sep 17 00:00:00 2001 From: jinyuan Date: Wed, 24 Jun 2026 17:08:10 +0100 Subject: [PATCH 9/9] update AliyunLLM --- evoagentx/models/aliyun_model.py | 440 +++++------------------------- evoagentx/models/model_configs.py | 10 +- 2 files changed, 76 insertions(+), 374 deletions(-) diff --git a/evoagentx/models/aliyun_model.py b/evoagentx/models/aliyun_model.py index 5f951940..62c57fb8 100644 --- a/evoagentx/models/aliyun_model.py +++ b/evoagentx/models/aliyun_model.py @@ -1,386 +1,90 @@ -import asyncio -from typing import Optional, List, Any -from tenacity import ( - retry, - stop_after_attempt, - wait_random_exponential, -) -from dashscope import Generation # aliyun DashScope SDK -import dashscope +from openai import AsyncOpenAI, OpenAI +from openai.types.completion_usage import CompletionUsage +from litellm import cost_per_token +from litellm.types.utils import Usage +from ..core.logging import logger from ..core.registry import register_model +from .openai_model import OpenAILLM from .model_configs import AliyunLLMConfig -from .base_model import BaseLLM -from .model_utils import Cost, cost_manager -from ..core.logging import logger -import os +from .model_utils import Cost, get_openai_model_cost + @register_model(config_cls=AliyunLLMConfig, alias=["aliyun_llm"]) -class AliyunLLM(BaseLLM): +class AliyunLLM(OpenAILLM): + """Aliyun Bailian (DashScope) LLM client. + + Bailian exposes an OpenAI-compatible endpoint (``.../compatible-mode/v1``), so this + reuses all of ``OpenAILLM``'s generation/streaming/tool-call logic and only overrides + client construction and cost handling. + + Two things differ from plain OpenAI: + + 1. Besides the API key (``aliyun_api_key``, i.e. the ``DASHSCOPE_API_KEY``), a + ``aliyun_base_url`` is required because the endpoint is workspace-specific (the + URL embeds the user's WorkspaceId). + 2. The compatible-mode responses carry standard ``usage`` token counts but no + ``usage.cost``. Dollar cost is recovered through LiteLLM, whose pricing table + keys DashScope models under the ``dashscope/`` prefix (e.g. ``dashscope/qwen-plus``). + If LiteLLM has no pricing for the configured model, tokens are still tracked and + cost is recorded as 0 (a warning is emitted once at init time). + """ + def init_model(self): - """ - Initialize the DashScope Generation client. - """ config: AliyunLLMConfig = self.config if not config.aliyun_api_key: raise ValueError("Aliyun API key is required. You should set `aliyun_api_key` in AliyunLLMConfig") - - # API key - os.environ["DASHSCOPE_API_KEY"] = config.aliyun_api_key - dashscope.api_key = config.aliyun_api_key - - # model - self._client = Generation() + if not config.aliyun_base_url: + raise ValueError( + "Aliyun base URL is required. You should set `aliyun_base_url` in AliyunLLMConfig " + "(it is workspace-specific, e.g. " + "'https://{WorkspaceId}.cn-beijing.maas.aliyuncs.com/compatible-mode/v1')." + ) + self._client = None + self._async_client = None + # parameters in AliyunLLMConfig that are not OpenAI-compatible request params self._default_ignore_fields = [ - "llm_type", "output_response", "aliyun_api_key", "aliyun_access_key_id", - "aliyun_access_key_secret", "model_name" + "llm_type", "output_response", "aliyun_api_key", "aliyun_base_url", ] - - def formulate_messages(self, prompts: List[str], system_messages: Optional[List[str]] = None) -> List[List[dict]]: - """ - Format messages for the Aliyun model. - - Args: - prompts (List[str]): List of user prompts. - system_messages (Optional[List[str]]): Optional list of system messages. - - Returns: - List[List[dict]]: Formatted messages for the model. - """ - if system_messages: - assert len(prompts) == len(system_messages), f"the number of prompts ({len(prompts)}) is different from the number of system_messages ({len(system_messages)})" + # LiteLLM keys DashScope models under the "dashscope/" prefix for pricing lookups. + self._litellm_model = self._to_litellm_model(config.model) + self._has_pricing = self._litellm_model in get_openai_model_cost() + if self._has_pricing: + logger.warning( + f"[AliyunLLM] Aliyun does not report usage.cost, so dollar cost for '{config.model}' " + f"is estimated from LiteLLM's public pricing for '{self._litellm_model}'. The figure " + "may differ from Aliyun's official billing." + ) else: - system_messages = [None] * len(prompts) - - messages_list = [] - for prompt, system_message in zip(prompts, system_messages): - messages = [] - if system_message: - messages.append({"role": "system", "content": system_message}) - messages.append({"role": "user", "content": prompt}) - messages_list.append(messages) - return messages_list - - def update_completion_params(self, params1: dict, params2: dict) -> dict: - """ - Update completion parameters with new values. - - Args: - params1 (dict): Base parameters. - params2 (dict): New parameters to update with. - - Returns: - dict: Updated parameters. - """ - config_params: list = self.config.get_config_params() - for key, value in params2.items(): - if key in self._default_ignore_fields: - continue - if key not in config_params: - continue - params1[key] = value - return params1 - - def get_completion_params(self, **kwargs): - """ - Get completion parameters for the model. - - Returns: - dict: Parameters for model completion. - """ - completion_params = self.config.get_set_params(ignore=self._default_ignore_fields) - completion_params = self.update_completion_params(completion_params, kwargs) - completion_params["model"] = self.config.model - return completion_params - - def get_stream_output(self, response: Any, output_response: bool = True) -> str: - """ - Process streaming response from the model. - - Args: - response: The streaming response from the model. - output_response (bool): Whether to print the response. - - Returns: - str: The complete response text. - """ - output = "" - try: - for chunk in response: - if not hasattr(chunk, 'output') or chunk.output is None: - error_msg = getattr(chunk, 'message', 'Invalid chunk format from model') - raise ValueError(f"Model stream chunk error: {error_msg}") - if hasattr(chunk.output, 'text'): - content = chunk.output.text - elif hasattr(chunk.output, 'choices') and chunk.output.choices: - content = chunk.output.choices[0].message.content - else: - continue - if content: - if output_response: - print(content, end="", flush=True) - output += content - except Exception as e: - print(f"Error processing stream: {str(e)}") - if not output: - raise RuntimeError(f"Failed to process stream response: {str(e)}") - if output_response: - print("") - return output - - async def get_stream_output_async(self, response: Any, output_response: bool = False) -> str: - """ - Process streaming response asynchronously. - - Args: - response: The streaming response from the model. - output_response (bool): Whether to print the response. - - Returns: - str: The complete response text. - """ - output = "" - try: - async for chunk in response: - if not hasattr(chunk, 'output') or chunk.output is None: - error_msg = getattr(chunk, 'message', 'Invalid chunk format from model') - raise ValueError(f"Model stream chunk error: {error_msg}") - if hasattr(chunk.output, 'text'): - content = chunk.output.text - elif hasattr(chunk.output, 'choices') and chunk.output.choices: - content = chunk.output.choices[0].message.content - else: - continue - if content: - if output_response: - print(content, end="", flush=True) - output += content - except Exception as e: - print(f"Error processing async stream: {str(e)}") - if not output: - raise RuntimeError(f"Failed to process async stream response: {str(e)}") - if output_response: - print("") - return output - - def get_completion_output(self, response: Any, output_response: bool = True) -> str: - """ - Process non-streaming response from the model. - - Args: - response: The response from the model. - output_response (bool): Whether to print the response. - - Returns: - str: The complete response text. - """ - try: - if not hasattr(response, 'output') or response.output is None: - error_msg = getattr(response, 'message', 'Invalid response format from model') - raise ValueError(f"Model response error: {error_msg}") - - if hasattr(response.output, 'text'): - output = response.output.text - elif hasattr(response.output, 'choices') and response.output.choices: - output = response.output.choices[0].message.content - else: - raise ValueError("Unexpected response format") - - if output_response: - print(output) - return output - except Exception as e: - raise RuntimeError(f"Error processing completion response: {str(e)}") - - @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(5)) - def single_generate(self, messages: List[dict], **kwargs) -> str: - """ - Generate a single response from the model. - - Args: - messages (List[dict]): The conversation history. - **kwargs: Additional parameters for generation. - - Returns: - str: The generated response. - """ - stream = kwargs.get("stream", self.config.stream) - output_response = kwargs.get("output_response", self.config.output_response) - - try: - completion_params = self.get_completion_params(**kwargs) - response = self._client.call(messages=messages, **completion_params) - - if response is None: - raise RuntimeError("Received empty response from model") - - if stream: - output = self.get_stream_output(response, output_response=output_response) - cost = self._stream_cost(response) - else: - output = self.get_completion_output(response=response, output_response=output_response) - cost = self._completion_cost(response) - self._update_cost(cost=cost) - return output - except Exception as e: - raise RuntimeError(f"Error during single_generate of AliyunLLM: {str(e)}") + logger.warning( + f"[AliyunLLM] LiteLLM has no pricing for '{self._litellm_model}', so dollar cost " + "cannot be computed. Token usage will be tracked, but cost will be recorded as 0." + ) - def batch_generate(self, batch_messages: List[List[dict]], **kwargs) -> List[str]: - """ - Generate responses for a batch of messages. - - Args: - batch_messages (List[List[dict]]): List of conversation histories. - **kwargs: Additional parameters for generation. - - Returns: - List[str]: List of generated responses. - """ - if not isinstance(batch_messages, list) or not batch_messages: - raise ValueError("batch_messages must be a non-empty list of message lists") - return [self.single_generate(messages=one_messages, **kwargs) for one_messages in batch_messages] + @staticmethod + def _to_litellm_model(model: str) -> str: + return model if model.startswith("dashscope/") else f"dashscope/{model}" - async def single_generate_async(self, messages: List[dict], **kwargs) -> str: - """ - Asynchronously generate a single response. - - Args: - messages (List[dict]): The conversation history. - **kwargs: Additional parameters for the generation. - - Returns: - str: The generated response. - """ - stream = kwargs.get("stream", self.config.stream) - output_response = kwargs.get("output_response", self.config.output_response) - - try: - completion_params = self.get_completion_params(**kwargs) - loop = asyncio.get_event_loop() - response = await loop.run_in_executor( - None, - lambda: self._client.call(messages=messages, **completion_params) - ) - - if stream: - output = await self.get_stream_output_async(response, output_response=output_response) - cost = self._stream_cost(response) - else: - output = self.get_completion_output(response=response, output_response=output_response) - cost = self._completion_cost(response) - - self._update_cost(cost=cost) - return output - - except Exception as e: - raise RuntimeError(f"Error during single_generate_async of AliyunLLM: {str(e)}") + def _init_client(self, config: AliyunLLMConfig): + return OpenAI(api_key=config.aliyun_api_key, base_url=config.aliyun_base_url) - def _completion_cost(self, response: Any) -> Cost: - """cost""" - try: - if not response: - return Cost(input_tokens=0, output_tokens=0, input_cost=0.0, output_cost=0.0) - - # tokens number - input_tokens = 0 - output_tokens = 0 - - if hasattr(response, 'usage'): - usage = response.usage - if hasattr(usage, 'input_tokens'): - input_tokens = usage.input_tokens - elif hasattr(usage, 'prompt_tokens'): - input_tokens = usage.prompt_tokens - - if hasattr(usage, 'output_tokens'): - output_tokens = usage.output_tokens - elif hasattr(usage, 'completion_tokens'): - output_tokens = usage.completion_tokens - - # - if input_tokens == 0 and output_tokens == 0 and hasattr(response, 'output'): - if hasattr(response.output, 'text'): - output_tokens = len(response.output.text.split()) * 1.3 - elif hasattr(response.output, 'choices') and response.output.choices: - output_tokens = len(response.output.choices[0].message.content.split()) * 1.3 - - total_cost = self._estimate_cost(input_tokens, output_tokens) - return Cost( - input_tokens=input_tokens, - output_tokens=output_tokens, - input_cost=total_cost * 0.4, # - output_cost=total_cost * 0.6 # - ) - except Exception as e: - logger.warning(f"Error computing completion cost: {str(e)}") - return Cost(input_tokens=0, output_tokens=0, input_cost=0.0, output_cost=0.0) + def _init_async_client(self, config: AliyunLLMConfig): + return AsyncOpenAI(api_key=config.aliyun_api_key, base_url=config.aliyun_base_url) - def _stream_cost(self, response: Any) -> Cost: - """cost""" - try: - if not response: - return Cost(input_tokens=0, output_tokens=0, input_cost=0.0, output_cost=0.0) - - # - input_tokens = 0 - output_tokens = 0 - - if hasattr(response, 'usage'): - usage = response.usage - if hasattr(usage, 'input_tokens'): - input_tokens = usage.input_tokens - elif hasattr(usage, 'prompt_tokens'): - input_tokens = usage.prompt_tokens - - if hasattr(usage, 'output_tokens'): - output_tokens = usage.output_tokens - elif hasattr(usage, 'completion_tokens'): - output_tokens = usage.completion_tokens - - # - if input_tokens == 0 and output_tokens == 0 and hasattr(response, 'output'): - if hasattr(response.output, 'text'): - output_tokens = len(response.output.text.split()) * 1.3 # - elif hasattr(response.output, 'choices') and response.output.choices: - output_tokens = len(response.output.choices[0].message.content.split()) * 1.3 - - total_cost = self._estimate_cost(input_tokens, output_tokens) + def _compute_cost(self, usage: CompletionUsage) -> Cost: + # Aliyun's compatible-mode usage has no `cost` field; price it via LiteLLM using + # the "dashscope/"-prefixed model name. Fall back to token-only when unpriced. + if not self._has_pricing: return Cost( - input_tokens=input_tokens, - output_tokens=output_tokens, - input_cost=total_cost * 0.4, # - output_cost=total_cost * 0.6 # + input_tokens=usage.prompt_tokens, + output_tokens=usage.completion_tokens, + cost=0.0, ) - except Exception as e: - logger.warning(f"Error computing stream cost: {str(e)}") - return Cost(input_tokens=0, output_tokens=0, input_cost=0.0, output_cost=0.0) - - def _estimate_cost(self, input_tokens: int, output_tokens: int) -> float: - """cost - - """ - model = self.config.model.lower() - if "turbo" in model: - input_cost = (input_tokens / 1000) * 0.0005 - output_cost = (output_tokens / 1000) * 0.001 - elif "max" in model: - input_cost = (input_tokens / 1000) * 0.002 - output_cost = (output_tokens / 1000) * 0.004 - else: # default - input_cost = (input_tokens / 1000) * 0.001 - output_cost = (output_tokens / 1000) * 0.002 - - return input_cost + output_cost - - def _update_cost(self, cost: Cost): - """ - Update the cost manager with the new cost. - - Args: - cost (Cost): The cost to update. - """ - try: - cost_manager.update_cost(cost=cost, model=self.config.model) - except Exception as e: - logger.warning(f"Error updating cost: {str(e)}") - + usage_object = usage if isinstance(usage, Usage) else Usage(**usage.model_dump()) + input_cost, output_cost = cost_per_token(model=self._litellm_model, usage_object=usage_object) + return Cost( + input_tokens=usage.prompt_tokens, + output_tokens=usage.completion_tokens, + input_cost=input_cost, + output_cost=output_cost, + ) diff --git a/evoagentx/models/model_configs.py b/evoagentx/models/model_configs.py index 3f6bc68f..80e60f00 100644 --- a/evoagentx/models/model_configs.py +++ b/evoagentx/models/model_configs.py @@ -179,10 +179,9 @@ def __str__(self): class AliyunLLMConfig(LLMConfig): llm_type: str = "AliyunLLM" - aliyun_api_key: Optional[str] = Field(default=None, description="The API key used to authenticate Aliyun requests") - aliyun_access_key_id: Optional[str] = Field(default=None, description="The Access Key ID for Aliyun authentication") - aliyun_access_key_secret: Optional[str] = Field(default=None, description="The Access Key Secret for Aliyun authentication") - + aliyun_api_key: Optional[str] = Field(default=None, description="The API key used to authenticate Aliyun requests (i.e. the DASHSCOPE_API_KEY)") + aliyun_base_url: Optional[str] = Field(default=None, description="The OpenAI-compatible base URL for the Aliyun Bailian endpoint. It is workspace-specific, e.g. 'https://{WorkspaceId}.cn-beijing.maas.aliyuncs.com/compatible-mode/v1'") + # generation parameters temperature: Optional[float] = Field(default=None, description="The temperature used to control randomness in generation. Higher values increase diversity.") top_p: Optional[float] = Field(default=None, description="Nucleus sampling parameter. Only sample from tokens with cumulative probability greater than top_p.") @@ -197,9 +196,8 @@ class AliyunLLMConfig(LLMConfig): tool_choice: Optional[str] = Field(default=None, description="Controls whether the model should call a tool. Options include 'none' (no tool call), 'auto' (model decides), or a specific tool name.") # model-specific parameters - model_name: Optional[str] = Field(default=None, description="The name of the Aliyun model to use, e.g., 'qwen-max', 'qwen-turbo'.") enable_search: Optional[bool] = Field(default=None, description="Whether to enable web search augmentation for the model, if supported.") - + # output format response_format: Optional[Union[BaseModel, dict]] = Field(default=None, description="Specifies the format of the model output, e.g., JSON schema for structured responses.") output_modalities: Optional[List] = Field(default=None, description="Output types the model should generate, e.g., ['text', 'image'] for multimodal models.")