liqiang
b119135836
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
200 lines
7.5 KiB
Python
200 lines
7.5 KiB
Python
# coding: utf-8
|
|
import requests
|
|
import json
|
|
from io import BytesIO
|
|
import os
|
|
from typing import Any, Optional, Type
|
|
import base64
|
|
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
BaseMessage,
|
|
HumanMessage,
|
|
SystemMessage,
|
|
ToolMessage,
|
|
)
|
|
|
|
from aworld.logs.util import logger
|
|
|
|
|
|
def extract_json_from_model_output(content: str) -> dict:
|
|
"""Extract JSON from model output, handling both plain JSON and code-block-wrapped JSON."""
|
|
try:
|
|
# If content is wrapped in code blocks, extract just the JSON part
|
|
if '```' in content:
|
|
# Find the JSON content between code blocks
|
|
content = content.split('```')[1]
|
|
# Remove language identifier if present (e.g., 'json\n')
|
|
if '\n' in content:
|
|
content = content.split('\n', 1)[1]
|
|
# Parse the cleaned content
|
|
return json.loads(content)
|
|
except json.JSONDecodeError as e:
|
|
logger.warning(f'Failed to parse model output: {content} {str(e)}')
|
|
raise ValueError('Could not parse response.')
|
|
|
|
|
|
def convert_input_messages(input_messages: list[BaseMessage], model_name: Optional[str]) -> list[BaseMessage]:
|
|
"""Convert input messages to a format that is compatible with the planner model"""
|
|
if model_name is None:
|
|
return input_messages
|
|
if model_name == 'deepseek-reasoner' or model_name.startswith('deepseek-r1'):
|
|
converted_input_messages = _convert_messages_for_non_function_calling_models(input_messages)
|
|
merged_input_messages = _merge_successive_messages(converted_input_messages, HumanMessage)
|
|
merged_input_messages = _merge_successive_messages(merged_input_messages, AIMessage)
|
|
return merged_input_messages
|
|
return input_messages
|
|
|
|
|
|
def _convert_messages_for_non_function_calling_models(input_messages: list[BaseMessage]) -> list[BaseMessage]:
|
|
"""Convert messages for non-function-calling models"""
|
|
output_messages = []
|
|
for message in input_messages:
|
|
if isinstance(message, HumanMessage):
|
|
output_messages.append(message)
|
|
elif isinstance(message, SystemMessage):
|
|
output_messages.append(message)
|
|
elif isinstance(message, ToolMessage):
|
|
output_messages.append(HumanMessage(content=message.content))
|
|
elif isinstance(message, AIMessage):
|
|
# check if tool_calls is a valid JSON object
|
|
if message.tool_calls:
|
|
tool_calls = json.dumps(message.tool_calls)
|
|
output_messages.append(AIMessage(content=tool_calls))
|
|
else:
|
|
output_messages.append(message)
|
|
else:
|
|
raise ValueError(f'Unknown message type: {type(message)}')
|
|
return output_messages
|
|
|
|
|
|
def _merge_successive_messages(messages: list[BaseMessage], class_to_merge: Type[BaseMessage]) -> list[BaseMessage]:
|
|
"""Some models like deepseek-reasoner dont allow multiple human messages in a row. This function merges them into one."""
|
|
merged_messages = []
|
|
streak = 0
|
|
for message in messages:
|
|
if isinstance(message, class_to_merge):
|
|
streak += 1
|
|
if streak > 1:
|
|
if isinstance(message.content, list):
|
|
merged_messages[-1].content += message.content[0]['text'] # type:ignore
|
|
else:
|
|
merged_messages[-1].content += message.content
|
|
else:
|
|
merged_messages.append(message)
|
|
else:
|
|
merged_messages.append(message)
|
|
streak = 0
|
|
return merged_messages
|
|
|
|
|
|
def save_conversation(input_messages: list[BaseMessage], response: Any, target: str,
|
|
encoding: Optional[str] = None) -> None:
|
|
"""Save conversation history to file."""
|
|
|
|
# create folders if not exists
|
|
os.makedirs(os.path.dirname(target), exist_ok=True)
|
|
|
|
with open(
|
|
target,
|
|
'w',
|
|
encoding=encoding,
|
|
) as f:
|
|
_write_messages_to_file(f, input_messages)
|
|
_write_response_to_file(f, response)
|
|
|
|
|
|
def _write_messages_to_file(f: Any, messages: list[BaseMessage]) -> None:
|
|
"""Write messages to conversation file"""
|
|
for message in messages:
|
|
f.write(f' {message.__class__.__name__} \n')
|
|
|
|
if isinstance(message.content, list):
|
|
for item in message.content:
|
|
if isinstance(item, dict) and item.get('type') == 'text':
|
|
f.write(item['text'].strip() + '\n')
|
|
elif isinstance(message.content, str):
|
|
try:
|
|
content = json.loads(message.content)
|
|
f.write(json.dumps(content, indent=2) + '\n')
|
|
except json.JSONDecodeError:
|
|
f.write(message.content.strip() + '\n')
|
|
|
|
f.write('\n')
|
|
|
|
|
|
def _write_response_to_file(f: Any, response: Any) -> None:
|
|
"""Write model response to conversation file"""
|
|
f.write(' RESPONSE\n')
|
|
f.write(json.dumps(json.loads(response.model_dump_json(exclude_unset=True)), indent=2))
|
|
|
|
|
|
# Add token counting related functions
|
|
# Note: These functions have been moved from memory.py and agent.py to utils.py, removing the dependency on MessageManager class
|
|
|
|
def estimate_text_tokens(text: str, estimated_characters_per_token: int = 3) -> int:
|
|
"""Roughly estimate token count in text
|
|
|
|
Args:
|
|
text: The text to estimate tokens for
|
|
estimated_characters_per_token: Estimated characters per token, default is 3
|
|
|
|
Returns:
|
|
Estimated token count
|
|
"""
|
|
if not text:
|
|
return 0
|
|
# Use character count divided by average characters per token to estimate tokens
|
|
return len(text) // estimated_characters_per_token
|
|
|
|
|
|
def estimate_message_tokens(message: BaseMessage, image_tokens: int = 800,
|
|
estimated_characters_per_token: int = 3) -> int:
|
|
"""Roughly estimate token count for a single message
|
|
|
|
Args:
|
|
message: The message to estimate tokens for
|
|
image_tokens: Estimated tokens per image, default is 800
|
|
estimated_characters_per_token: Estimated characters per token, default is 3
|
|
|
|
Returns:
|
|
Estimated token count
|
|
"""
|
|
tokens = 0
|
|
# Handle tuple case
|
|
if isinstance(message, tuple):
|
|
# Convert to string and estimate tokens
|
|
message_str = str(message)
|
|
return estimate_text_tokens(message_str, estimated_characters_per_token)
|
|
|
|
if isinstance(message.content, list):
|
|
for item in message.content:
|
|
if 'image_url' in item:
|
|
tokens += image_tokens
|
|
elif isinstance(item, dict) and 'text' in item:
|
|
tokens += estimate_text_tokens(item['text'], estimated_characters_per_token)
|
|
else:
|
|
msg = message.content
|
|
if hasattr(message, 'tool_calls'):
|
|
msg += str(message.tool_calls) # type: ignore
|
|
tokens += estimate_text_tokens(msg, estimated_characters_per_token)
|
|
return tokens
|
|
|
|
|
|
def estimate_messages_tokens(messages: list[BaseMessage], image_tokens: int = 800,
|
|
estimated_characters_per_token: int = 3) -> int:
|
|
"""Roughly estimate total token count for a list of messages
|
|
|
|
Args:
|
|
messages: The list of messages to estimate tokens for
|
|
image_tokens: Estimated tokens per image, default is 800
|
|
estimated_characters_per_token: Estimated characters per token, default is 3
|
|
|
|
Returns:
|
|
Estimated total token count
|
|
"""
|
|
total_tokens = 0
|
|
for msg in messages:
|
|
total_tokens += estimate_message_tokens(msg, image_tokens, estimated_characters_per_token)
|
|
return total_tokens
|