ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
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

This commit is contained in:
2026-08-20 13:12:50 +00:00
commit b119135836
10275 changed files with 3284984 additions and 0 deletions
@@ -0,0 +1,59 @@
import logging
import os
from aworld.config import ModelConfig
from aworld.config.conf import AgentConfig, ClientType
from pydantic import BaseModel
from aworldspace.base_agent import AworldBaseAgent
from aworldspace.utils.mcp_utils import load_all_mcp_config
SYSTEM_PROMPT = f"""You are an helpful AI assistant, aimed at solving any task presented by the user. """
class Pipeline(AworldBaseAgent):
class Valves(BaseModel):
pass
def __init__(self):
self.valves = self.Valves()
logging.info("default init success")
async def get_agent_config(self, body):
default_llm_provider = os.environ.get("LLM_PROVIDER")
llm_model_name = os.environ.get("LLM_MODEL_NAME")
llm_api_key = os.environ.get("LLM_API_KEY")
llm_base_url = os.environ.get("LLM_BASE_URL")
task = await self.get_task_from_body(body)
logging.info(f"task llm config is: {task.llm_provider}, {task.llm_model_name},{task.llm_base_url}")
llm_config = ModelConfig(
llm_provider=task.llm_provider if task and task.llm_provider else default_llm_provider,
llm_model_name=task.llm_model_name if task and task.llm_model_name else llm_model_name,
llm_api_key=task.llm_api_key if task and task.llm_api_key else llm_api_key,
llm_base_url=task.llm_base_url if task and task.llm_base_url else llm_base_url,
max_retries=task.max_retries if task and task.max_retries else 3
)
return AgentConfig(
name=self.agent_name(),
llm_config=llm_config,
system_prompt=task.task_system_prompt if task and task.task_system_prompt else SYSTEM_PROMPT
)
def agent_name(self) -> str:
return "DefaultAgent"
async def get_mcp_servers(self, body) -> list[str]:
task = await self.get_task_from_body(body)
if task.mcp_servers:
logging.info(f"mcp_servers from task: {task.mcp_servers}")
return task.mcp_servers
return [
"ms-playwright"
]
async def load_mcp_config(self) -> dict:
return load_all_mcp_config()
@@ -0,0 +1,264 @@
import logging
import os
import re
from pathlib import Path
from typing import Dict, Any, List, Optional
from aworld.config import ModelConfig
from aworld.config.conf import AgentConfig, TaskConfig, ClientType
from aworld.core.task import Task
from aworld.output import Outputs, Output, StreamingOutputs
from aworld.utils.common import get_local_ip
from datasets import load_dataset, concatenate_datasets
from pydantic import BaseModel, Field
from aworldspace.base_agent import AworldBaseAgent
from aworldspace.utils.mcp_utils import load_all_mcp_config
from aworldspace.utils.utils import question_scorer
GAIA_SYSTEM_PROMPT = f"""You are an all-capable AI assistant, aimed at solving any task presented by the user. You have various tools at your disposal that you can call upon to efficiently complete complex requests. Whether it's programming, information retrieval, file processing, or web browsing, you can handle it all.
Please note that the task may be complex. Do not attempt to solve it all at once. You should break the task down and use different tools step by step to solve it. After using each tool, clearly explain the execution results and suggest the next steps.
Please utilize appropriate tools for the task, analyze the results obtained from these tools, and provide your reasoning. Always use available tools such as browser, calcutor, etc. to verify correctness rather than relying on your internal knowledge.
If you believe the problem has been solved, please output the `final answer`. The `final answer` should be given in <answer></answer> format, while your other thought process should be output in <think></think> tags.
Your `final answer` should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Here are some tips to help you give better instructions:
<tips>
1. Do not use any tools outside of the provided tools list.
2. Even if the task is complex, there is always a solution. If you cant find the answer using one method, try another approach or use different tools to find the solution.
3. When using browser `playwright_click` tool, you need to check if the element exists and is clickable before clicking it.
4. Before providing the `final answer`, carefully reflect on whether the task has been fully solved. If you have not solved the task, please provide your reasoning and suggest the next steps.
5. Due to context length limitations, always try to complete browser-based tasks with the minimal number of steps possible.
6. When providing the `final answer`, answer the user's question directly and precisely. For example, if asked "what animal is x?" and x is a monkey, simply answer "monkey" rather than "x is a monkey".
7. When you need to process excel file, prioritize using the `excel` tool instead of writing custom code with `terminal-controller` tool.
8. If you need to download a file, please use the `terminal-controller` tool to download the file and save it to the specified path.
9. The browser doesn't support direct searching on www.google.com. Use the `google-search` to get the relevant website URLs or contents instead of `ms-playwright` directly.
10. Always use only one tool at a time in each step of your execution.
11. Using `mcp__ms-playwright__browser_pdf_save` tool to save the pdf file of URLs to the specified path.
12. Using `mcp__terminal-controller__execute_command` tool to set the timeout to 300 seconds when downloading large files such as pdf.
13. Using `mcp__ms-playwright__browser_take_screenshot` tool to save the screenshot of URLs to the specified path when you need to understand the gif / jpg of the URLs.
14. When there are questions related to YouTube video comprehension, use tools in `youtube_download_server` and `video_server` to analyze the video content by the given question.
</tips>
Now, here is the task. Stay focused and complete it carefully using the appropriate tools!
"""
class Pipeline(AworldBaseAgent):
class Valves(BaseModel):
llm_provider: Optional[str] = Field(default=None, description="llm_model_name")
llm_model_name: Optional[str] = Field(default=None, description="llm_model_name")
llm_base_url: Optional[str] = Field(default=None,description="llm_base_urly")
llm_api_key: Optional[str] = Field(default=None,description="llm api key" )
system_prompt: str = Field(default=GAIA_SYSTEM_PROMPT,description="system_prompt")
history_messages: int = Field(default=100, description="rounds of history messages")
def __init__(self):
self.valves = self.Valves()
self.gaia_files = os.path.abspath(os.path.join(os.path.curdir, "aworldspace", "datasets", "gaia_dataset"))
logging.info(f"gaia_files path {self.gaia_files}")
self.full_dataset = load_dataset(
os.path.join(self.gaia_files, "GAIA.py"),
name="2023_all",
trust_remote_code=True
)
self.full_dataset = concatenate_datasets([self.full_dataset['validation'], self.full_dataset['test']])
# Create task_id to index mapping for improved lookup performance
self.task_id_to_index = {}
for i, task in enumerate(self.full_dataset):
self.task_id_to_index[task['task_id']] = i
logging.info(f"Loaded {len(self.full_dataset)} tasks, created task_id mapping")
logging.info("gaia_agent init success")
async def get_custom_input(self, user_message: str, model_id: str, messages: List[dict], body: dict) -> Any:
task = await self.get_gaia_task(user_message)
logging.info(f"🌈 -----------------------------------------------")
logging.info(f"🚀 Start to process: gaia_task_{task['task_id']}")
logging.info(f"📝 Detail: {task}")
logging.info(f"❓ Question: {task['Question']}")
logging.info(f"⭐ Level: {task['Level']}")
logging.info(f"🛠️ Tools: {task['Annotator Metadata']['Tools']}")
logging.info(f"🌈 -----------------------------------------------")
return task['Question']
async def get_agent_config(self, body):
default_llm_provider = self.valves.llm_provider if self.valves.llm_provider else os.environ.get("LLM_PROVIDER")
llm_model_name = self.valves.llm_model_name if self.valves.llm_model_name else os.environ.get("LLM_MODEL_NAME")
llm_api_key = self.valves.llm_api_key if self.valves.llm_api_key else os.environ.get("LLM_API_KEY")
llm_base_url = self.valves.llm_base_url if self.valves.llm_base_url else os.environ.get("LLM_BASE_URL")
system_prompt = self.valves.system_prompt if self.valves.system_prompt else GAIA_SYSTEM_PROMPT
task = await self.get_task_from_body(body)
if task:
logging.info(f"task llm config is: {task.llm_provider}, {task.llm_model_name}, {task.llm_api_key}, {task.llm_base_url}")
llm_config = ModelConfig(
llm_provider=task.llm_provider if task and task.llm_provider else default_llm_provider,
llm_model_name=task.llm_model_name if task and task.llm_model_name else llm_model_name,
llm_api_key=task.llm_api_key if task and task.llm_api_key else llm_api_key,
llm_base_url=task.llm_base_url if task and task.llm_base_url else llm_base_url,
max_retries=task.max_retries if task and task.max_retries else 3
)
return AgentConfig(
name=self.agent_name(),
llm_config=llm_config,
system_prompt=task.task_system_prompt if task and task.task_system_prompt else system_prompt
)
def agent_name(self) -> str:
return "GaiaAgent"
async def get_mcp_servers(self, body) -> list[str]:
task = await self.get_task_from_body(body)
if task and task.mcp_servers:
logging.info(f"mcp_servers from task: {task.mcp_servers}")
return task.mcp_servers
return [
"e2b-server",
"terminal-controller",
"excel",
"calculator",
"ms-playwright",
"audio_server",
"image_server",
"video_server",
"search_server",
"download_server",
"document_server",
"youtube_server",
"reasoning_server",
]
async def get_gaia_task(self, task_id: str) -> dict:
"""
Get GAIA task by task_id
Args:
task_id: Unique identifier of the task
Returns:
Corresponding task dictionary
"""
# Search by task_id
if task_id in self.task_id_to_index:
index = self.task_id_to_index[task_id]
gaia_task = self.full_dataset[index]
else:
raise ValueError(f"Task with task_id '{task_id}' not found in dataset")
return self.add_file_path(gaia_task)
def get_all_task_ids(self) -> List[str]:
"""
Get list of all available task_ids
Returns:
List of all task_ids
"""
return list(self.task_id_to_index.keys())
def get_task_count(self) -> int:
"""
Get total number of tasks
Returns:
Total task count
"""
return len(self.full_dataset)
def get_task_index_by_id(self, task_id: str) -> int:
"""
Get task index in dataset by task_id
Args:
task_id: Unique identifier of the task
Returns:
Index of the task in the dataset
"""
if task_id in self.task_id_to_index:
return self.task_id_to_index[task_id]
else:
raise ValueError(f"Task with task_id '{task_id}' not found in dataset")
async def custom_output_before_task(self, outputs: Outputs, chat_id: str, task: Task) -> None:
task_config:TaskConfig = task.conf
gaia_task = await self.get_gaia_task(task_config.ext['origin_message'])
result = f"\n\n`{get_local_ip()}` execute `GAIA TASK#{task_config.ext['origin_message']}`:\n\n---\n\n"
result += f"**Question**: {gaia_task['Question']}\n"
result += f"**Answer**: {gaia_task['Final answer']}\n"
result += f"**Level**: {gaia_task['Level']}\n"
result += f"**Tools**: \n {gaia_task['Annotator Metadata']['Tools']}\n"
result += f"\n\n-----\n\n"
await outputs.add_output(Output(data = result))
async def custom_output_after_task(self, outputs: Outputs, chat_id: str, task: Task):
"""
check gaia task output
Args:
outputs:
chat_id:
task:
Returns:
"""
task_config: TaskConfig = task.conf
gaia_task_id = task_config['ext']['origin_message']
gaia_task = await self.get_gaia_task(gaia_task_id)
agent_result = ""
if isinstance(outputs, StreamingOutputs):
agent_result = await outputs._visited_outputs[-2].get_finished_response() # read llm result
match = re.search(r"<answer>(.*?)</answer>", agent_result)
answer = agent_result
if match:
answer = match.group(1)
logging.info(f"🤖 Agent answer: {answer}")
logging.info(f"👨‍🏫 Correct answer: {gaia_task['Final answer']}")
is_correct = question_scorer(answer, gaia_task["Final answer"])
if is_correct:
logging.info(f"📝Question {gaia_task_id} Correct! 🎉")
result = f"\n\n📝 **Question: {gaia_task_id} -> Agent Answer:[{answer}] is `Correct`**"
else:
logging.info(f"📝Question {gaia_task_id} Incorrect! ❌")
result = f"\n\n📝 **Question: {gaia_task_id} -> Agent Answer:`{answer}` != Correct answer: `{gaia_task['Final answer']}` is `Incorrect` ❌**"
metadata = await outputs.get_metadata()
if not metadata:
await outputs.set_metadata({})
metadata = await outputs.get_metadata()
metadata['gaia_correct'] = is_correct
metadata['gaia_result'] = result
metadata['agent_answer'] = answer
metadata['correct_answer'] = gaia_task['Final answer']
return result
def add_file_path(self, task: Dict[str, Any]
):
split = "validation" if task["Annotator Metadata"]["Steps"] != "" else "test"
if task["file_name"]:
file_path = Path(f"{self.gaia_files}/2023/{split}/" + task["file_name"])
if file_path.suffix in [".pdf", ".docx", ".doc", ".txt"]:
task["Question"] += f" Here are the necessary document files: {file_path}"
elif file_path.suffix in [".jpg", ".jpeg", ".png"]:
task["Question"] += f" Here are the necessary image files: {file_path}"
elif file_path.suffix in [".xlsx", "xls", ".csv"]:
task[
"Question"
] += f" Here are the necessary table files: {file_path}, for processing excel file, you can use the excel tool or write python code to process the file step-by-step and get the information."
elif file_path.suffix in [".py"]:
task["Question"] += f" Here are the necessary python files: {file_path}"
else:
task["Question"] += f" Here are the necessary files: {file_path}"
return task
async def load_mcp_config(self) -> dict:
return load_all_mcp_config()
@@ -0,0 +1,794 @@
import logging
import os
import traceback
from typing import Dict, Any, List, Union
from typing import Optional
from aworld.core.event.base import Message
from aworldspace.base_agent import AworldBaseAgent
from pydantic import BaseModel, Field
import aworld.trace as trace
from aworld.config.conf import AgentConfig, ConfigDict
from aworld.config.conf import TaskConfig
from aworld.agents.llm_agent import Agent
from aworld.core.common import Observation, ActionModel
from aworld.core.memory import MemoryItem
from aworld.core.task import Task
from aworld.logs.util import logger
from aworld.models.llm import acall_llm_model
from aworld.models.model_response import ToolCall, Function
from aworld.output import Output, StreamingOutputs
from aworld.output import Outputs
from aworld.output.base import MessageOutput
from aworld.utils.common import sync_exec
BROWSER_SYSTEM_PROMPT = """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task.
## Output Format
```
Thought: ...
Action: ...
```
## Action Space
navigate(website='xxx') #Open the target website, usually the first action to open browser.
click(start_box='[x1, y1, x2, y2]')
left_double(start_box='[x1, y1, x2, y2]')
right_single(start_box='[x1, y1, x2, y2]')
drag(start_box='[x1, y1, x2, y2]', end_box='[x3, y3, x4, y4]')
hotkey(key='')
type(content='') #If you want to submit your input, use "\n" at the end of `content`.
scroll(direction='down or up or right or left')
wait() #Sleep for 5s and take a screenshot to check for any changes.
finished(content='xxx') # Use escape characters \\', \\", and \\n in content part to ensure we can parse the content in normal python string format.
## Note
- only one action per step.
- Use Chinese in `Thought` part.
- Write a small plan and finally summarize your next action (with its target element) in one sentence in `Thought` part.
## User Instruction
"""
import json
import re
MAX_IMAGE = 50
def parse_action_output(output_text):
# 提取Thought部分
logger.info(f"{output_text=}")
thought_match = re.search(r'Thought:(.*?)\nAction:', output_text, re.DOTALL)
thought = thought_match.group(1).strip() if thought_match else ""
# 提取Action部分
action_match = re.search(r'Action:(.*?)(?:\n|$)', output_text, re.DOTALL)
action_text = action_match.group(1).strip() if action_match else ""
# 初始化结果字典
result = {
"thought": thought,
"action": "",
"key": None,
"content": None,
"start_box": None,
"end_box": None,
"direction": None,
"website": None,
}
if not action_text:
return json.dumps(result, ensure_ascii=False)
# tmp 兼容ui-tars1.5-7b
action_text = action_text.replace("'(","'[").replace(")'","]'")
# 解析action类型
action_parts = action_text.split('(')
action_type = action_parts[0]
result["action"] = action_type
# 解析参数
if len(action_parts) > 1:
params_text = action_parts[1].rstrip(')')
params = {}
# gpt-4o兼容
if 'start_box' in params_text:
params_text = params_text.replace(", ", " ").replace(",", " ")
if 'end_box' in params_text:
params_text = params_text.replace(" end_box", ", end_box")
# 处理键值对参数
for param in params_text.split(','):
param = param.strip()
if '=' in param:
key, value = param.split('=', 1)
key = key.strip()
value = value.strip().strip('\'"')
# 处理bbox格式
if 'box' in key:
print(value)
# 提取坐标数字
numbers = re.findall(r'\d+', value)
print(numbers)
if numbers:
coords = [int(num) for num in numbers]
if len(coords) == 4:
if key == 'start_box':
result["start_box"] = coords
elif key == 'end_box':
result["end_box"] = coords
if len(coords) == 2:
if key == 'start_box':
result["start_box"] = [coords[0], coords[1], coords[0], coords[1]]
elif key == 'end_box':
result["end_box"] = [coords[0], coords[1], coords[0], coords[1]]
elif key == 'key':
result["key"] = value.replace("pagedown", "PageDown").replace("pageup", "PageUp").replace("enter","Enter")
elif key == 'content':
# 处理转义字符
value = value.replace('\\n', '\n').replace('\\"', '"').replace("\\'", "'")
result["content"] = value
elif key == 'website':
result["website"] = value
elif key == 'direction':
result["direction"] = value
return result, thought, action_text
def parse_tool_call(line):
# 提取 Action和param
result, thought, action_text = parse_action_output(line)
action = result['action']
# 映射到实际函数名和参数
if action == 'navigate':
func_name = 'mcp__ms-playwright__browser_navigate'
content = {'url': result['website']}
elif action == 'click':
func_name = 'mcp__ms-playwright__browser_screen_click'
x = int((result["start_box"][0] + result["start_box"][2]) / 2)
y = int((result["start_box"][1] + result["start_box"][3]) / 2)
content = {'element': '', 'x': x, 'y': y}
elif action == 'right_single':
func_name = 'mcp__ms-playwright__browser_screen_click'
x = int((result["start_box"][0] + result["start_box"][2]) / 2)
y = int((result["start_box"][1] + result["start_box"][3]) / 2)
content = {'element': 'right click target', 'x': x, 'y': y, 'button': 'right'}
elif action == 'drag':
func_name = 'mcp__ms-playwright__browser_screen_drag'
x1 = int((result["start_box"][0] + result["start_box"][2]) / 2)
y1 = int((result["start_box"][1] + result["start_box"][3]) / 2)
x2 = int((result["end_box"][0] + result["end_box"][2]) / 2)
y2 = int((result["end_box"][1] + result["end_box"][3]) / 2)
content = {
'element': f'drag from [{x1},{y1}] to [{x2},{y2}]',
'startX': x1,
'startY': y1,
'endX': x2,
'endY': y2
}
elif action == 'hotkey':
func_name = 'mcp__ms-playwright__browser_press_key'
content = {'key': result["key"]}
elif action == 'type':
func_name = 'mcp__ms-playwright__browser_screen_type'
content = {'text': result['content']}
elif action == 'scroll':
# 暂时使用presskey代替scroll
func_name = 'mcp__ms-playwright__browser_press_key'
direction = result['direction']
key_map = {
'up': 'PageUp',
'down': 'PageDown',
'left': 'ArrowLeft',
'right': 'ArrowRight'
}
key = key_map.get(direction, 'ArrowDown')
content = {'key': key}
elif action == 'wait':
func_name = 'mcp__ms-playwright__browser_wait_for'
content = {'time': 5}
elif action == 'finished':
func_name = "finished"
content = result['content']
else:
return ""
return Function(name=func_name, arguments=json.dumps(content)), thought, action_text, result
# eval code start
def identify_key_points(task):
system_msg = """You are an expert tasked with analyzing a given task to identify the key points explicitly stated in the task description.
**Objective**: Carefully analyze the task description and extract the critical elements explicitly mentioned in the task for achieving its goal.
**Instructions**:
1. Read the task description carefully.
2. Identify and extract **key points** directly stated in the task description.
- A **key point** is a critical element, condition, or step explicitly mentioned in the task description.
- Do not infer or add any unstated elements.
- Words such as "best," "highest," "cheapest," "latest," "most recent," "lowest," "closest," "highest-rated," "largest," and "newest" must go through the sort function(e.g., the key point should be "Filter by highest").
**Respond with**:
- **Key Points**: A numbered list of the explicit key points for completing this task, one per line, without explanations or additional details."""
prompt = """Task: {task}"""
text = prompt.format(task=task)
messages = [
{"role": "system", "content": system_msg},
{
"role": "user",
"content": [
{"type": "text", "text": text}
],
}
]
return messages
def judge_image(task, image_path, key_points):
system_msg = """You are an expert evaluator tasked with determining whether an image contains information about the necessary steps to complete a task.
**Objective**: Analyze the provided image and decide if it shows essential steps or evidence required for completing the task. Use your reasoning to explain your decision before assigning a score.
**Instructions**:
1. Provide a detailed description of the image, including its contents, visible elements, text (if any), and any notable features.
2. Carefully examine the image and evaluate whether it contains necessary steps or evidence crucial to task completion:
- Identify key points that could be relevant to task completion, such as actions, progress indicators, tool usage, applied filters, or step-by-step instructions.
- Does the image show actions, progress indicators, or critical information directly related to completing the task?
- Is this information indispensable for understanding or ensuring task success?
- If the image contains partial but relevant information, consider its usefulness rather than dismissing it outright.
3. Provide your response in the following format:
- **Reasoning**: Explain your thought process and observations. Mention specific elements in the image that indicate necessary steps, evidence, or lack thereof.
- **Score**: Assign a score based on the reasoning, using the following scale:
- **1**: The image does not contain any necessary steps or relevant information.
- **2**: The image contains minimal or ambiguous information, unlikely to be essential.
- **3**: The image includes some relevant steps or hints but lacks clarity or completeness.
- **4**: The image contains important steps or evidence that are highly relevant but not fully comprehensive.
- **5**: The image clearly displays necessary steps or evidence crucial for completing the task.
Respond with:
1. **Reasoning**: [Your explanation]
2. **Score**: [1-5]"""
# jpg_base64_str = encode_image(Image.open(image_path))
prompt = """**Task**: {task}
**Key Points for Task Completion**: {key_points}
The snapshot of the web page is shown in the image."""
text = prompt.format(task=task, key_points=key_points)
messages = [
{"role": "system", "content": system_msg},
{
"role": "user",
"content": [
{"type": "text", "text": text},
{
"type": "image_url",
"image_url": {"url": image_path, "detail": "high"},
},
],
}
]
return messages
def WebJudge_Online_Mind2Web_eval(task, last_actions, images_path, image_responses, key_points, score_threshold):
system_msg = """You are an expert in evaluating the performance of a web navigation agent. The agent is designed to help a human user navigate a website to complete a task. Given the user's task, the agent's action history, key points for task completion, some potentially important web pages in the agent's trajectory and their reasons, your goal is to determine whether the agent has completed the task and achieved all requirements.
Your response must strictly follow the following evaluation criteria!
*Important Evaluation Criteria*:
1: The filtered results must be displayed correctly. If filters were not properly applied (i.e., missing selection, missing confirmation, or no visible effect in results), the task is not considered successful.
2: You must carefully check whether these snapshots and action history meet these key points. Ensure that specific filter conditions, such as "best," "highest," "cheapest," "latest," "most recent," "lowest," "closest," "highest-rated," "largest," and "newest" are correctly applied using the filter function(e.g., sort function).
3: Certain key points or requirements should be applied by the filter. Otherwise, a search with all requirements as input will be deemed a failure since it cannot guarantee that all results meet the requirements!
4: If the task requires filtering by a specific range of money, years, or the number of beds and bathrooms, the applied filter must exactly match the given requirement. Any deviation results in failure. To ensure the task is successful, the applied filter must precisely match the specified range without being too broad or too narrow.
Examples of Failure Cases:
- If the requirement is less than $50, but the applied filter is less than $25, it is a failure.
- If the requirement is $1500-$2500, but the applied filter is $2000-$2500, it is a failure.
- If the requirement is $25-$200, but the applied filter is $0-$200, it is a failure.
- If the required years are 2004-2012, but the filter applied is 2001-2012, it is a failure.
- If the required years are before 2015, but the applied filter is 2000-2014, it is a failure.
- If the task requires exactly 2 beds, but the filter applied is 2+ beds, it is a failure.
5: Some tasks require a submission action or a display of results to be considered successful.
6: If the retrieved information is invalid or empty(e.g., No match was found), but the agent has correctly performed the required action, it should still be considered successful.
7: If the current page already displays all available items, then applying a filter is not necessary. As long as the agent selects items that meet the requirements (e.g., the cheapest or lowest price), the task is still considered successful.
*IMPORTANT*
Format your response into two lines as shown below:
Thoughts: <your thoughts and reasoning process based on double-checking each key points and the evaluation criteria>
Status: "success" or "failure"
"""
prompt = """User Task: {task}
Key Points: {key_points}
Action History:
{last_actions}
The potentially important snapshots of the webpage in the agent's trajectory and their reasons:
{thoughts}"""
whole_content_img = []
whole_thoughts = []
record = []
pattern = r"[1-5]"
for response, image_path in zip(image_responses, images_path):
try:
score_text = response.split("Score")[1]
thought = response.split("**Reasoning**:")[-1].strip().lstrip("\n").split("\n\n")[0].replace('\n', ' ')
score = re.findall(pattern, score_text)[0]
record.append({"Response": response, "Score": int(score)})
except Exception as e:
print(f"Error processing response: {e}")
score = 0
record.append({"Response": response, "Score": 0})
if int(score) >= score_threshold:
# jpg_base64_str = encode_image(Image.open(image_path))
whole_content_img.append(
{
'type': 'image_url',
"image_url": {"url": image_path, "detail": "high"},
}
)
if thought != "":
whole_thoughts.append(thought)
whole_content_img = whole_content_img[:MAX_IMAGE]
whole_thoughts = whole_thoughts[:MAX_IMAGE]
if len(whole_content_img) == 0:
prompt = """User Task: {task}
Key Points: {key_points}
Action History:
{last_actions}"""
text = prompt.format(task=task,
last_actions="\n".join(f"{i + 1}. {action}" for i, action in enumerate(last_actions)),
key_points=key_points,
thoughts="\n".join(f"{i + 1}. {thought}" for i, thought in enumerate(whole_thoughts)))
messages = [
{"role": "system", "content": system_msg},
{
"role": "user",
"content": [
{"type": "text", "text": text}]
+ whole_content_img
}
]
return messages, text, system_msg, record
# eval code end
class PlayWrightAgent(Agent):
def __init__(self, conf: Union[Dict[str, Any], ConfigDict, AgentConfig], **kwargs):
self.screen_capture = True
self.step_images = []
self.step_thoughts = []
self.step_actions = []
self.step_results = []
self.success = False
super().__init__(conf, **kwargs)
async def async_policy(self, observation: Observation, info: Dict[str, Any] = {}, message: Message = None,
**kwargs) -> Union[
List[ActionModel], None]:
"""The strategy of an agent can be to decide which tools to use in the environment, or to delegate tasks to other agents.
Args:
observation: The state observed from tools in the environment.
info: Extended information is used to assist the agent to decide a policy.
Returns:
ActionModel sequence from agent policy
"""
outputs = None
if kwargs.get("outputs") and isinstance(kwargs.get("outputs"), Outputs):
outputs = kwargs.get("outputs")
# Get current step information for trace recording
step = kwargs.get("step", 0)
exp_id = kwargs.get("exp_id", None)
source_span = trace.get_current_span()
if hasattr(observation, 'context') and observation.context:
self.task_histories = observation.context
self._finished = False
await self.async_desc_transform(message.context)
self.tools = None
if "data:image/jpeg;base64," in observation.content:
logger.info("transfer base64 content to image")
observation.image = observation.content
observation.content = "observation:"
self.step_images.append(observation.image)
images = observation.images if self.conf.use_vision else None
if self.conf.use_vision and not images and observation.image:
images = [observation.image]
messages = self.messages_transform(content=observation.content,
image_urls=images,
sys_prompt=self.system_prompt,
agent_prompt=self.agent_prompt)
self._log_messages(messages)
if isinstance(messages[-1]['content'], list):
messages[-1]['role'] = 'user' # 有image的话必须使用user请求,而且不写入历史对话
# self.memory.add(MemoryItem(
# content=messages[-1]['content'],
# metadata={
# "role": messages[-1]['role'],
# "agent_name": self.name(),
# }
# ))
else:
self.memory.add(MemoryItem(
content=messages[-1]['content'],
metadata={
"role": messages[-1]['role'],
"agent_name": self.name(),
}
))
llm_response = None
span_name = f"llm_call_{exp_id}"
with trace.span(span_name) as llm_span:
llm_span.set_attributes({
"exp_id": exp_id,
"step": step,
"messages": json.dumps([str(m) for m in messages], ensure_ascii=False)
})
if source_span:
source_span.set_attribute("messages", json.dumps([str(m) for m in messages], ensure_ascii=False))
try:
llm_response = await acall_llm_model(
self.llm,
messages=messages,
model=self.model_name,
# temperature=self.conf.llm_config.llm_temperature,
temperature=0.0,
tools=self.tools if not self.use_tools_in_prompt and self.tools else None,
stream=kwargs.get("stream", False)
)
# Record LLM response
llm_span.set_attributes({
"llm_response": json.dumps(llm_response.to_dict(), ensure_ascii=False),
"tool_calls": json.dumps([tool_call.model_dump() for tool_call in
llm_response.tool_calls] if llm_response.tool_calls else [],
ensure_ascii=False),
"error": llm_response.error if llm_response.error else ""
})
except Exception as e:
logger.warn(traceback.format_exc())
llm_span.set_attribute("error", str(e))
raise e
finally:
if llm_response:
use_tools = self.use_tool_list(llm_response)
is_use_tool_prompt = len(use_tools) > 0
if llm_response.error:
logger.info(f"llm result error: {llm_response.error}")
else:
self.memory.add(MemoryItem(
content=llm_response.content,
metadata={
"role": "assistant",
"agent_name": self.name(),
"tool_calls": llm_response.tool_calls if not self.use_tools_in_prompt else use_tools,
"is_use_tool_prompt": is_use_tool_prompt if not self.use_tools_in_prompt else False
}
))
function, origin_thought, origin_action, origin_result = parse_tool_call(
llm_response.message['content'])
self.step_thoughts.append(origin_thought)
self.step_actions.append(origin_action)
self.step_results.append(origin_result)
if function.name == "finished":
self._finished = True
llm_response.content = "<answer>" + llm_response.content + "</answer>"
llm_response.tool_calls = None
else:
llm_response.content = None
tool_call = ToolCall(
id="tooluse_mock",
type="function",
function=function,
)
screen_capture = ToolCall(
id="screen_capture",
type="function",
function=Function(
name="mcp__ms-playwright__browser_screen_capture",
arguments="{}"
)
)
llm_response.tool_calls = [tool_call, screen_capture]
else:
logger.error(f"{self.name()} failed to get LLM response")
raise RuntimeError(f"{self.name()} failed to get LLM response")
if outputs and isinstance(outputs, Outputs):
await outputs.add_output(MessageOutput(source=llm_response, json_parse=False))
agent_result = await self.model_output_parser.parse(llm_response, agent_id=self.id())
if not agent_result.is_call_tool:
self._finished = True
logger.info(self.step_thoughts)
logger.info(self.step_actions)
# now is eval code:
logger.info(f"step:{step}")
if self.finished or step >= 20: # 暂时写死,这里应该是max_step
task = self.task.split("Please first navigate to the target")[0]
key_points_messages = identify_key_points(task)
# eval_model_name = "shangshu.gpt-4o"
eval_model_name = self.model_name
tmp_llm_response = await acall_llm_model(
self.llm,
messages=key_points_messages,
model=eval_model_name,
temperature=0
)
key_points = tmp_llm_response.content
key_points = key_points.replace("\n\n", "\n")
try:
key_points = key_points.split("**Key Points**:")[1]
key_points = "\n".join(line.lstrip() for line in key_points.splitlines())
except:
key_points = key_points.split("Key Points:")[-1]
key_points = "\n".join(line.lstrip() for line in key_points.splitlines())
logger.info(f"key_points: {key_points}")
tasks_messages = [judge_image(task, image_path, key_points) for image_path in self.step_images]
# 这里暂时使用串行执行的写法
image_responses = []
for task_messages in tasks_messages:
logger.info(task_messages)
image_response = await acall_llm_model(
self.llm, # 假设这是你传给函数的第一个参数
messages=task_messages, # 每个请求的消息内容
model=eval_model_name, # 模型名称
temperature=0 # 温度参数
)
image_responses.append(image_response)
image_responses = [i.content for i in image_responses]
logger.info(f"image_responses: {image_responses}")
eval_messages, text, system_msg, record = WebJudge_Online_Mind2Web_eval(
self.task, self.step_actions, self.step_images, image_responses, key_points, 3)
response = await acall_llm_model(
self.llm,
messages=eval_messages,
model=eval_model_name,
temperature=0
)
eval_response = response.content
logger.info(f"eval_response: {eval_response}")
if "success" in eval_response.lower().split('status:')[1]:
self.success = True
# now is saving code:
result_dict = {
'task': task,
'images': self.step_images,
'actions': self.step_actions,
'thoughts': self.step_thoughts,
'results': self.step_results,
'success': self.success,
'final_answer': llm_response.content,
'eval_response': eval_response,
'is_done': self.finished,
'done_step': step,
}
result_dict = json.dumps(result_dict, ensure_ascii=False)
agent_result.actions[0].policy_info = result_dict
agent_result.actions[0].tool_name = None
agent_result.actions[0].action_name = None
agent_result.actions[0].agent_name = self.name()
# saving is over...
return agent_result.actions
class Pipeline(AworldBaseAgent):
class Valves(BaseModel):
llm_provider: Optional[str] = Field(default=None, description="llm_model_name")
llm_model_name: Optional[str] = Field(default=None, description="llm_model_name")
llm_base_url: Optional[str] = Field(default=None, description="llm_base_urly")
llm_api_key: Optional[str] = Field(default=None, description="llm api key")
system_prompt: str = Field(default=BROWSER_SYSTEM_PROMPT, description="system_prompt")
history_messages: int = Field(default=100, description="rounds of history messages")
def __init__(self):
self.valves = self.Valves()
self.agent_config = AgentConfig(
name=self.agent_name(),
llm_provider=self.valves.llm_provider if self.valves.llm_provider else os.environ.get("LLM_PROVIDER"),
llm_model_name=self.valves.llm_model_name if self.valves.llm_model_name else os.environ.get(
"LLM_MODEL_NAME"),
llm_api_key=self.valves.llm_api_key if self.valves.llm_api_key else os.environ.get("LLM_API_KEY"),
llm_base_url=self.valves.llm_base_url if self.valves.llm_base_url else os.environ.get("LLM_BASE_URL"),
system_prompt=self.valves.system_prompt if self.valves.system_prompt else BROWSER_SYSTEM_PROMPT
)
self.m2w_files = os.path.abspath(os.path.join(os.path.curdir, "aworldspace", "datasets", "online-mind2web"))
logging.info(f"m2w_files path {self.m2w_files}")
file_path = os.path.join(self.m2w_files, "Online_Mind2Web.json")
with open(file_path, 'r') as file:
self.full_dataset = json.load(file)
logging.info("playwright_agent init success")
# 重写build_agent
async def build_agent(self, body: dict):
agent_config = await self.get_agent_config(body)
mcp_servers = await self.get_mcp_servers(body)
agent = PlayWrightAgent(
conf=agent_config,
name=agent_config.name,
system_prompt=agent_config.system_prompt,
mcp_servers=mcp_servers,
mcp_config=await self.load_mcp_config(),
history_messages=await self.get_history_messages(body)
)
return agent
async def get_custom_input(self, user_message: str, model_id: str, messages: List[dict], body: dict) -> Any:
task = await self.get_m2w_task(int(user_message))
return task['Task']
async def get_agent_config(self, body):
default_llm_provider = self.valves.llm_provider if self.valves.llm_provider else os.environ.get("LLM_PROVIDER")
llm_model_name = self.valves.llm_model_name if self.valves.llm_model_name else os.environ.get("LLM_MODEL_NAME")
llm_api_key = self.valves.llm_api_key if self.valves.llm_api_key else os.environ.get("LLM_API_KEY")
llm_base_url = self.valves.llm_base_url if self.valves.llm_base_url else os.environ.get("LLM_BASE_URL")
system_prompt = self.valves.system_prompt if self.valves.system_prompt else BROWSER_SYSTEM_PROMPT
task = await self.get_task_from_body(body)
logging.info(
f"task llm config is: {task.llm_provider}, {task.llm_model_name}, {task.llm_api_key}, {task.llm_base_url}")
return AgentConfig(
name=self.agent_name(),
llm_provider=task.llm_provider if task and task.llm_provider else default_llm_provider,
llm_model_name=task.llm_model_name if task and task.llm_model_name else llm_model_name,
llm_api_key=task.llm_api_key if task and task.llm_api_key else llm_api_key,
llm_base_url=task.llm_base_url if task and task.llm_base_url else llm_base_url,
system_prompt=task.task_system_prompt if task and task.task_system_prompt else system_prompt
)
def agent_name(self) -> str:
return "PlaywrightAgent"
async def get_mcp_servers(self, body) -> list[str]:
task = await self.get_task_from_body(body)
if task.mcp_servers:
logging.info(f"mcp_servers from task: {task.mcp_servers}")
return task.mcp_servers
return [
"ms-playwright"
]
async def get_m2w_task(self, index) -> dict:
logging.info(f"Start to process: m2w_task_{index}")
m2w_task = self.full_dataset[index]
logging.info(f"Detail: {m2w_task}")
logging.info(f"Task: {m2w_task['confirmed_task']}")
logging.info(f"Level: {m2w_task['level']}")
logging.info(f"Website: {m2w_task['website']}")
return self.add_file_path(m2w_task)
async def custom_output_before_task(self, outputs: Outputs, chat_id: str, task: Task) -> None:
task_config: TaskConfig = task.conf
m2w_task = await self.get_m2w_task(int(task_config.ext['origin_message']))
result = f"\n\n`Web TASK#{task_config.ext['origin_message']}`\n\n---\n\n"
result += f"**Task**: {m2w_task['Task']}\n"
result += f"**Level**: {m2w_task['level']}\n"
result += f"**Website**: \n {m2w_task['website']}\n"
result += f"\n\n-----\n\n"
await outputs.add_output(Output(data=result))
async def custom_output_after_task(self, outputs: Outputs, chat_id: str, task: Task):
"""
check gaia task output
Args:
outputs:
chat_id:
task:
Returns:
"""
task_config: TaskConfig = task.conf
web_task_id = int(task_config['ext']['origin_message'])
web_task = await self.get_m2w_task(web_task_id)
agent_result = ""
if isinstance(outputs, StreamingOutputs):
agent_result = await outputs._visited_outputs[-2].get_finished_response() # read llm result
# match = re.search(r"<answer>(.*?)</answer>", agent_result)
result = ""
# if match:
# answer = match.group(1)
logging.info(f"Agent answer: {agent_result}")
metadata = await outputs.get_metadata()
if not metadata:
await outputs.set_metadata({})
metadata = await outputs.get_metadata()
metadata['web_task'] = web_task
return result
def add_file_path(self, task: Dict[str, Any]
):
task["Task"] = "Task: " + task['confirmed_task'] + '\n' + "Please first navigate to the target " + "Website: " + \
task['website']
return task
async def load_mcp_config(self) -> dict:
return {
"mcpServers": {
"ms-playwright": {
"command": "npx",
"args": [
"@playwright/mcp@0.0.27",
"--vision",
"--no-sandbox",
"--headless",
"--isolated"
],
"env": {
"PLAYWRIGHT_TIMEOUT": "120000",
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
}
}
}
@@ -0,0 +1,32 @@
from typing import Optional
from pydantic import BaseModel, Field
from aworldspace.base_agent import AworldBaseAgent
"""
Agent Space
"""
class AgentMeta(BaseModel):
name: str = None
desc: str = None
class AgentSpace(BaseModel):
agent_modules: Optional[dict] = Field(default_factory=dict, description="agent module")
agents_meta: Optional[dict] = Field(default_factory=dict, description="agents meta")
def register(self, agent_name: str, agent_instance: AworldBaseAgent, metadata: dict=None):
# Register agent metadata and instance
self.agent_modules[agent_name] = agent_instance
async def get_agent_modules(self):
return self.agent_modules
async def get_agents_meta(self):
return self.agents_meta
AGENT_SPACE = AgentSpace()
@@ -0,0 +1,242 @@
import json
import logging
import os
import traceback
import uuid
from abc import abstractmethod
from typing import List, AsyncGenerator, Any
from aworld.config import AgentConfig, TaskConfig, ContextRuleConfig, OptimizationConfig
from aworld.agents.llm_agent import Agent
from aworld.core.task import Task
from aworld.output import WorkSpace, AworldUI, Outputs
from aworld.output.ui.markdown_aworld_ui import MarkdownAworldUI
from aworld.output.utils import load_workspace
from aworld.runner import Runners
from client.aworld_client import AworldTask
class AworldBaseAgent:
def pipes(self) -> list[dict]:
return [{"id": self.agent_name(), "name": self.agent_name()}]
@abstractmethod
def agent_name(self) -> str:
pass
async def pipe(
self,
user_message: str,
model_id: str,
messages: List[dict],
body: dict
):
try:
logging.info(f"🤖{self.agent_name()} received user_message is {user_message}, form-data = {body}")
task = await self.get_task_from_body(body)
if task:
logging.info(f"🤖{self.agent_name()} received task is {task.task_id}_{task.client_id}_{task.user_id}")
task_id = task.task_id
else:
task_id = str(uuid.uuid4())
session_id = task_id
if body.get('metadata'):
# user_id = body.get('metadata').get('user_id')
session_id = body.get('metadata').get('chat_id', task_id)
task_id = body.get('metadata').get('message_id', task_id)
user_input = await self.get_custom_input(user_message, model_id, messages, body)
if task and task.llm_custom_input:
user_input = task.llm_custom_input
logging.info(f"🤖{self.agent_name()} call llm input is [{user_input}]")
# build agent task read from config
swarm = await self.build_swarm(body=body)
agent = None
if not swarm:
# build single agent task read from config
agent = await self.build_agent(body=body)
logging.info(f"🤖{self.agent_name()} build agent finished")
# return task
task = await self.build_task(agent=agent, task_id=task_id, user_input=user_input, user_message=user_message, body=body)
logging.info(f"🤖{self.agent_name()} build task finished, task_id is {task_id}")
workspace_type = os.environ.get("WORKSPACE_TYPE", "local")
workspace_path = os.environ.get("WORKSPACE_PATH", "./data/workspaces")
workspace = await load_workspace(session_id, workspace_type, workspace_path)
# render output
async_generator = await self.parse_task_output(session_id, task, workspace)
return async_generator()
except Exception as e:
return await self._format_exception(e)
async def _format_exception(self, e: Exception) -> str:
traceback.print_exc()
# tb_lines = traceback.format_exception(type(e), e, e.__traceback__)
# detailed_error = "".join(tb_lines)
# logging.error(e)
# return json.dumps({"error": detailed_error}, ensure_ascii=False)
return "💥💥💥process failed💥💥💥"
async def _format_error(self, status_code: int, error: bytes) -> str:
if isinstance(error, str):
error_str = error
else:
error_str = error.decode(errors="ignore")
try:
err_msg = json.loads(error_str).get("message", error_str)[:200]
except Exception:
err_msg = error_str[:200]
return json.dumps(
{"error": f"HTTP {status_code}: {err_msg}"}, ensure_ascii=False
)
async def get_custom_input(self, user_message: str,
model_id: str,
messages: List[dict],
body: dict) -> Any:
user_input = body["messages"][-1]["content"]
return user_input
@abstractmethod
async def get_history_messages(self, body) -> int:
task = await self.get_task_from_body(body)
if task:
return task.history_messages
return 100
@abstractmethod
async def get_agent_config(self, body) -> AgentConfig:
pass
@abstractmethod
async def get_mcp_servers(self, body) -> list[str]:
pass
async def build_agent(self, body: dict):
agent_config =await self.get_agent_config(body)
mcp_servers = await self.get_mcp_servers(body)
agent = Agent(
conf=agent_config,
name=agent_config.name,
system_prompt=agent_config.system_prompt,
mcp_servers=mcp_servers,
mcp_config=await self.load_mcp_config(),
history_messages=await self.get_history_messages(body),
context_rule=ContextRuleConfig(
optimization_config=OptimizationConfig(
enabled=False,
)
)
)
return agent
async def build_task(self, agent, task_id, user_input, user_message, body):
aworld_task = await self.get_task_from_body(body)
task = Task(
id=task_id,
name=task_id,
input=user_input,
agent=agent,
conf=TaskConfig(
task_id=task_id,
stream=False,
ext={
"origin_message": user_message
},
max_steps=aworld_task.max_steps if aworld_task else 100
)
)
return task
async def parse_task_output(self, chat_id, task: Task, workspace: WorkSpace):
_SENTINEL = object()
async def async_generator():
from asyncio import Queue
queue = Queue()
async def consume_all():
openwebui_ui = MarkdownAworldUI(
session_id=chat_id,
workspace=workspace
)
# get outputs
outputs = Runners.streamed_run_task(task)
# output hooks
await self.custom_output_before_task(outputs, chat_id, task)
# render output
try:
async for output in outputs.stream_events():
res = await AworldUI.parse_output(output, openwebui_ui)
if res:
if isinstance(res, AsyncGenerator):
async for item in res:
await queue.put(item)
else:
await queue.put(res)
custom_output = await self.custom_output_after_task(outputs, chat_id, task)
if custom_output:
await queue.put(custom_output)
await queue.put(task)
finally:
await queue.put(_SENTINEL)
# Start the consumer in the background
import asyncio
consumer_task = asyncio.create_task(consume_all())
while True:
item = await queue.get()
if item is _SENTINEL:
break
yield item
await consumer_task
logging.info(f"🤖{self.agent_name()} task#{task.id} output finished🔚🔚🔚")
return async_generator
async def custom_output_before_task(self, outputs: Outputs, chat_id: str, task: Task) -> str | None:
return None
async def custom_output_after_task(self, outputs: Outputs, chat_id: str, task: Task):
pass
async def get_task_from_body(self, body: dict) -> AworldTask | None:
try:
if not body.get("user") or not body.get("user").get("aworld_task"):
return None
return AworldTask.model_validate_json(body.get("user").get("aworld_task"))
except Exception as err:
logging.error(f"Error parsing AworldTask: {err}; data: {body.get('user_message')}")
traceback.print_exc()
return None
@abstractmethod
async def load_mcp_config(self) -> dict:
pass
async def build_swarm(self, body):
return None
@@ -0,0 +1,210 @@
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Optional
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from base import AworldTask, AworldTaskResult
from aworldspace.db.models import (
Base, AworldTaskModel, AworldTaskResultModel,
orm_to_pydantic_task, pydantic_to_orm_task,
orm_to_pydantic_result, pydantic_to_orm_result
)
class AworldTaskDB(ABC):
@abstractmethod
async def query_task_by_id(self, task_id: str) -> AworldTask:
pass
@abstractmethod
async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
pass
@abstractmethod
async def insert_task(self, task: AworldTask):
pass
@abstractmethod
async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
pass
@abstractmethod
async def update_task(self, task: AworldTask):
pass
@abstractmethod
async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
pass
@abstractmethod
async def save_task_result(self, result: AworldTaskResult):
pass
class SqliteTaskDB(AworldTaskDB):
def __init__(self, db_path: str):
self.engine = create_engine(db_path, echo=False, future=True)
Base.metadata.create_all(self.engine)
self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)
async def query_task_by_id(self, task_id: str) -> Optional[AworldTask]:
with self.Session() as session:
orm_task = session.query(AworldTaskModel).filter_by(task_id=task_id).first()
return orm_to_pydantic_task(orm_task) if orm_task else None
async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
with self.Session() as session:
orm_result = (
session.query(AworldTaskResultModel)
.filter_by(task_id=task_id)
.order_by(AworldTaskResultModel.created_at.desc())
.first()
)
return orm_to_pydantic_result(orm_result) if orm_result else None
async def insert_task(self, task: AworldTask):
with self.Session() as session:
orm_task = pydantic_to_orm_task(task)
session.add(orm_task)
session.commit()
async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
with self.Session() as session:
orm_tasks = (
session.query(AworldTaskModel)
.filter_by(status=status)
.limit(nums)
.all()
)
return [orm_to_pydantic_task(t) for t in orm_tasks]
async def update_task(self, task: AworldTask):
with self.Session() as session:
orm_task = session.query(AworldTaskModel).filter_by(task_id=task.task_id).first()
if orm_task:
for k, v in task.model_dump().items():
setattr(orm_task, k, v)
orm_task.updated_at = datetime.utcnow()
session.commit()
async def save_task_result(self, result: AworldTaskResult):
with self.Session() as session:
orm_task = pydantic_to_orm_result(result)
session.add(orm_task)
session.commit()
async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
with self.Session() as session:
query = session.query(AworldTaskModel)
# Handle special filters for time ranges
start_time = filter.pop('start_time', None)
end_time = filter.pop('end_time', None)
# Apply regular filters
for k, v in filter.items():
if hasattr(AworldTaskModel, k):
query = query.filter(getattr(AworldTaskModel, k) == v)
# Apply time range filters
if start_time:
query = query.filter(AworldTaskModel.created_at >= start_time)
if end_time:
query = query.filter(AworldTaskModel.created_at <= end_time)
total = query.count()
orm_tasks = query.offset((page_num - 1) * page_size).limit(page_size).all()
items = [orm_to_pydantic_task(t) for t in orm_tasks]
return {
"total": total,
"page_num": page_num,
"page_size": page_size,
"items": items
}
class PostgresTaskDB(AworldTaskDB):
def __init__(self, db_url: str):
# db_url example: 'postgresql+psycopg2://user:password@host:port/dbname'
self.engine = create_engine(db_url, echo=False, future=True)
Base.metadata.create_all(self.engine)
self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)
async def query_task_by_id(self, task_id: str) -> Optional[AworldTask]:
with self.Session() as session:
orm_task = session.query(AworldTaskModel).filter_by(task_id=task_id).first()
return orm_to_pydantic_task(orm_task) if orm_task else None
async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
with self.Session() as session:
orm_result = (
session.query(AworldTaskResultModel)
.filter_by(task_id=task_id)
.order_by(AworldTaskResultModel.created_at.desc())
.first()
)
return orm_to_pydantic_result(orm_result) if orm_result else None
async def insert_task(self, task: AworldTask):
with self.Session() as session:
orm_task = pydantic_to_orm_task(task)
session.add(orm_task)
session.commit()
async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
with self.Session() as session:
orm_tasks = (
session.query(AworldTaskModel)
.filter_by(status=status)
.limit(nums)
.all()
)
return [orm_to_pydantic_task(t) for t in orm_tasks]
async def update_task(self, task: AworldTask):
with self.Session() as session:
orm_task = session.query(AworldTaskModel).filter_by(task_id=task.task_id).first()
if orm_task:
for k, v in task.model_dump().items():
setattr(orm_task, k, v)
orm_task.updated_at = datetime.utcnow()
session.commit()
async def save_task_result(self, result: AworldTaskResult):
with self.Session() as session:
orm_task = pydantic_to_orm_result(result)
session.add(orm_task)
session.commit()
async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
with self.Session() as session:
query = session.query(AworldTaskModel)
# Handle special filters for time ranges
start_time = filter.pop('start_time', None)
end_time = filter.pop('end_time', None)
# Apply regular filters
for k, v in filter.items():
if hasattr(AworldTaskModel, k):
query = query.filter(getattr(AworldTaskModel, k) == v)
# Apply time range filters
if start_time:
query = query.filter(AworldTaskModel.created_at >= start_time)
if end_time:
query = query.filter(AworldTaskModel.created_at <= end_time)
total = query.count()
orm_tasks = query.offset((page_num - 1) * page_size).limit(page_size).all()
items = [orm_to_pydantic_task(t) for t in orm_tasks]
return {
"total": total,
"page_num": page_num,
"page_size": page_size,
"items": items
}
@@ -0,0 +1,66 @@
from sqlalchemy import Column, String, Integer, Text, DateTime, JSON, create_engine
from sqlalchemy.orm import declarative_base
from datetime import datetime
from typing import Optional
from base import AworldTask, AworldTaskResult
Base = declarative_base()
class AworldTaskModel(Base):
__tablename__ = 'aworld_tasks'
task_id = Column(String, primary_key=True)
agent_id = Column(String)
agent_input = Column(Text)
session_id = Column(String)
user_id = Column(String)
llm_provider = Column(String)
llm_model_name = Column(String)
llm_api_key = Column(String)
llm_base_url = Column(String)
llm_custom_input = Column(Text)
task_system_prompt = Column(Text)
mcp_servers = Column(JSON)
node_id = Column(String)
client_id = Column(String)
status = Column(String, default='INIT')
history_messages = Column(Integer, default=100)
max_steps = Column(Integer, default=100)
max_retries = Column(Integer, default=5)
ext_info = Column(JSON, default=dict)
created_at = Column(DateTime, default=datetime.utcnow)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
class AworldTaskResultModel(Base):
__tablename__ = 'aworld_tasks_results'
task_result_id = Column(Integer, primary_key=True, autoincrement=True)
task_id = Column(String)
server_host = Column(String)
data = Column(JSON)
created_at = Column(DateTime, default=datetime.utcnow)
def orm_to_pydantic_task(orm_obj: AworldTaskModel) -> AworldTask:
return AworldTask(**{c.name: getattr(orm_obj, c.name) for c in orm_obj.__table__.columns})
def pydantic_to_orm_task(pydantic_obj: AworldTask) -> AworldTaskModel:
return AworldTaskModel(**pydantic_obj.model_dump())
def orm_to_pydantic_result(orm_obj: AworldTaskResultModel) -> AworldTaskResult:
return AworldTaskResult(
server_host=orm_obj.server_host,
data=orm_obj.data
)
def pydantic_to_orm_result(pydantic_obj: AworldTaskResult) -> AworldTaskResultModel:
return AworldTaskResultModel(
task_id=pydantic_obj.task.task_id if pydantic_obj.task else None,
server_host=pydantic_obj.server_host,
data=pydantic_obj.data
)
@@ -0,0 +1,439 @@
import json
import os
import time
from datetime import datetime
from typing import AsyncGenerator, Optional, List
from aworld.utils.common import get_local_ip
from fastapi import APIRouter, Query, Response
from fastapi.responses import StreamingResponse
import logging
import traceback
from asyncio import Queue
import asyncio
from aworld.models.model_response import ModelResponse
from pydantic import BaseModel, Field, PrivateAttr
from aworldspace.db.db import AworldTaskDB, SqliteTaskDB, PostgresTaskDB
from aworldspace.utils.job import generate_openai_chat_completion, call_pipeline
from aworldspace.utils.log import task_logger
from base import AworldTask, AworldTaskResult, OpenAIChatCompletionForm, OpenAIChatMessage, AworldTaskForm
from config import ROOT_DIR
__STOP_TASK__ = object()
class AworldTaskExecutor(BaseModel):
"""
task executor
- load task from db and execute task in a loop
- use semaphore to limit concurrent tasks
"""
_task_db: AworldTaskDB = PrivateAttr()
_tasks: Queue = PrivateAttr()
max_concurrent: int = Field(default=os.environ.get("AWORLD_MAX_CONCURRENT_TASKS", 2), description="max concurrent tasks")
def __init__(self, task_db: AworldTaskDB):
super().__init__()
self._task_db = task_db
self._tasks = Queue()
self._semaphore = asyncio.BoundedSemaphore(self.max_concurrent)
async def start(self):
"""
execute task in a loop
"""
await asyncio.sleep(5)
logging.info(f"🚀[task executor] start, max concurrent is {self.max_concurrent}")
while True:
# load task if queue is empty and semaphore is not full
if self._tasks.empty():
await self.load_task()
task = await self._tasks.get()
if not task:
logging.info("task is none")
continue
if task == __STOP_TASK__:
logging.info("✅[task executor] stop, all tasks finished")
break
# acquire semaphore
await self._semaphore.acquire()
asyncio.create_task(self._run_task_and_release_semaphore(task))
async def stop(self):
logging.info("🛑 task executor stop, wait for all tasks to finish")
await self._tasks.put(__STOP_TASK__)
async def _run_task_and_release_semaphore(self, task: AworldTask):
"""
execute task and release semaphore when done
"""
start_time = time.time()
logging.info(f"🚀[task executor] execute task#{task.task_id} start, lock acquired")
try:
await self.execute_task(task)
finally:
# release semaphore
self._semaphore.release()
logging.info(f"✅[task executor] execute task#{task.task_id} success, use time {time.time() - start_time:.2f}s")
async def load_task(self):
interval = os.environ.get("AWORLD_TASK_LOAD_INTERVAL", 10)
# calculate the number of tasks to load
need_load = self._semaphore._value
if need_load <= 0:
logging.info(f"🔍[task executor] runner is busy, wait {interval}s and retry")
await asyncio.sleep(interval)
return await self.load_task()
tasks = await self._task_db.query_tasks_by_status(status="INIT", nums=need_load)
logging.info(f"🔍[task executor] load {len(tasks)} tasks from db (need {need_load})")
if not tasks or len(tasks) == 0:
logging.info(f"🔍[task executor] no task to load, wait {interval}s and retry")
await asyncio.sleep(interval)
return await self.load_task()
for task in tasks:
task.mark_running()
await self._task_db.update_task(task)
await self._tasks.put(task)
return True
async def execute_task(self, task: AworldTask):
"""
execute task
"""
try:
result = await self._execute_task(task)
task.mark_success()
await self._task_db.update_task(task)
await self._task_db.save_task_result(result)
task_logger.log_task_submission(task, "execute_finished", task_result=result)
except Exception as err:
task.mark_failed()
await self._task_db.update_task(task)
traceback.print_exc()
task_logger.log_task_submission(task, "execute_failed", details=f"err is {err}")
async def _execute_task(self, task: AworldTask):
# build params
messages = [
OpenAIChatMessage(role="user", content=task.agent_input)
]
# call_llm_model
form_data = OpenAIChatCompletionForm(
model=task.agent_id,
messages=messages,
stream=True,
user={
"user_id": task.user_id,
"session_id": task.session_id,
"task_id": task.task_id,
"aworld_task": task.model_dump_json()
}
)
data = await generate_openai_chat_completion(form_data)
task_result = {}
task.node_id = get_local_ip()
items = []
md_file = ""
if data.body_iterator:
if isinstance(data.body_iterator, AsyncGenerator):
async for item_content in data.body_iterator:
async def parse_item(_item_content) -> Optional[ModelResponse]:
if item_content == "data: [DONE]":
return None
return ModelResponse.from_openai_stream_chunk(json.loads(item_content.replace("data:", "")))
# if isinstance(item, ModelResponse)
item = await parse_item(item_content)
items.append(item)
if not item:
continue
if item.content:
md_file = task_logger.log_task_result(task, item)
logging.info(f"task#{task.task_id} response data chunk is: {item}"[:500])
if item.raw_response and item.raw_response and isinstance(item.raw_response, dict) and item.raw_response.get('task_output_meta'):
task_result = item.raw_response.get('task_output_meta')
data = {
"task_result": task_result,
"md_file": md_file,
"replays_file": f"trace_data/{datetime.now().strftime('%Y%m%d')}/{get_local_ip()}/replays/task_replay_{task.task_id}.json"
}
result = AworldTaskResult(task=task, server_host=get_local_ip(), data=data)
return result
class AworldTaskManager(BaseModel):
_task_db: AworldTaskDB = PrivateAttr()
_task_executor: AworldTaskExecutor = PrivateAttr()
def __init__(self, task_db: AworldTaskDB):
super().__init__()
self._task_db = task_db
self._task_executor = AworldTaskExecutor(task_db=self._task_db)
async def start_task_executor(self):
asyncio.create_task(self._task_executor.start())
async def stop_task_executor(self):
self._task_executor.tasks.put_nowait(None)
async def submit_task(self, task: AworldTask):
# save to db
await self._task_db.insert_task(task)
# log it
task_logger.log_task_submission(task, status="init")
return AworldTaskResult(task = task)
async def load_one_unfinished_task(self) -> Optional[AworldTask]:
tasks = await self._task_db.query_tasks_by_status(status="INIT", nums=1)
if not tasks or len(tasks) == 0:
return None
cur_task = tasks[0]
cur_task.mark_running()
await self._task_db.update_task(cur_task)
# from db load one task by locked and mark task running
return cur_task
async def get_task_result(self, task_id: str) -> Optional[AworldTaskResult]:
task = await self._task_db.query_task_by_id(task_id)
if task:
task_result = await self._task_db.query_latest_task_result_by_id(task_id)
if task_result:
return task_result
return AworldTaskResult(task=task)
async def get_batch_task_results(self, task_ids: List[str]) -> List[dict]:
"""
Batch retrieve task results, returns dictionary format
Each dict contains: task (required) and task_result (may be None)
"""
results = []
for task_id in task_ids:
task = await self._task_db.query_task_by_id(task_id)
if task:
task_result = await self._task_db.query_latest_task_result_by_id(task_id)
result_dict = {
"task": task,
"task_result": task_result # May be None
}
results.append(result_dict)
return results
async def query_and_download_task_results(
self,
start_time: Optional[datetime] = None,
end_time: Optional[datetime] = None,
task_id: Optional[str] = None,
page_size: int = 100
) -> List[dict]:
"""
Query tasks and get results, support time range and task_id filtering
"""
all_results = []
page_num = 1
while True:
# Build query filter conditions
filter_dict = {}
if start_time:
filter_dict['start_time'] = start_time
if end_time:
filter_dict['end_time'] = end_time
if task_id:
filter_dict['task_id'] = task_id
# Page query tasks
page_result = await self._task_db.page_query_tasks(
filter=filter_dict,
page_size=page_size,
page_num=page_num
)
if not page_result['items']:
break
tasks = page_result['items']
for task in tasks:
# Only query task_result (may not exist)
task_result = await self._task_db.query_latest_task_result_by_id(task.task_id)
# Use task information to build results
result_data = {
"task_id": task.task_id,
"agent_id": task.agent_id,
"status": task.status,
"created_at": task.created_at.isoformat() if task.created_at else None,
"updated_at": task.updated_at.isoformat() if task.updated_at else None,
"user_id": task.user_id,
"session_id": task.session_id,
"node_id": task.node_id,
"client_id": task.client_id,
"task_data": task.model_dump(mode='json'),
"has_result": task_result is not None,
"server_host": task_result.server_host if task_result else None,
"result_data": task_result.data if task_result else None,
}
all_results.append(result_data)
if len(page_result['items']) < page_size:
break
page_num += 1
return all_results
########################################################################################
########################### API
########################################################################################
router = APIRouter()
task_db_path = os.environ.get("AWORLD_TASK_DB_PATH", f"sqlite:///{ROOT_DIR}/db/aworld.db")
if task_db_path.startswith("sqlite://"):
task_db = SqliteTaskDB(db_path = task_db_path)
elif task_db_path.startswith("mysql://"):
task_db = None # todo: add mysql task db
elif task_db_path.startswith("postgresql://") or task_db_path.startswith("postgresql+"):
task_db = PostgresTaskDB(db_url=task_db_path)
else:
raise ValueError("❌ task_db_path is not a valid sqlite, mysql or postgresql path")
task_manager = AworldTaskManager(task_db)
@router.post("/submit_task")
async def submit_task(form_data: AworldTaskForm) -> Optional[AworldTaskResult]:
logging.info(f"🚀 submit task#{form_data.task.task_id} start")
if not form_data.task:
raise ValueError("task is empty")
try:
task_result = await task_manager.submit_task(form_data.task)
logging.info(f"✅ submit task#{form_data.task.task_id} success")
return task_result
except Exception as err:
traceback.print_exc()
logging.error(f"❌ submit task#{form_data.task.task_id} failed, err is {err}")
raise ValueError("❌ submit task failed, please see logs for details")
@router.get("/task_result")
async def get_task_result(task_id) -> Optional[AworldTaskResult]:
if not task_id:
raise ValueError("❌ task_id is empty")
logging.info(f"🚀 get task result#{task_id} start")
try:
task_result = await task_manager.get_task_result(task_id)
logging.info(f"✅ get task result#{task_id} success, task result is {task_result}")
return task_result
except Exception as err:
traceback.print_exc()
logging.error(f"❌ get task result#{task_id} failed, err is {err}")
raise ValueError("❌ get task result failed, please see logs for details")
@router.post("/get_batch_task_results")
async def get_batch_task_results(task_ids: List[str]) -> List[dict]:
if not task_ids or len(task_ids) == 0:
raise ValueError("❌ task_ids is empty")
logging.info(f"🚀 get batch task results start, task_ids: {task_ids}")
try:
batch_results = await task_manager.get_batch_task_results(task_ids)
logging.info(f"✅ get batch task results success, found {len(batch_results)} results")
return batch_results
except Exception as err:
traceback.print_exc()
logging.error(f"❌ get batch task results failed, err is {err}")
raise ValueError("❌ get batch task results failed, please see logs for details")
@router.get("/download_task_results")
async def download_task_results(
start_time: Optional[str] = Query(None, description="Start time, format: YYYY-MM-DD HH:MM:SS"),
end_time: Optional[str] = Query(None, description="End time, format: YYYY-MM-DD HH:MM:SS"),
task_id: Optional[str] = Query(None, description="Task ID"),
page_size: int = Query(100, description="Page size, ge=1, le=1000")
) -> StreamingResponse:
"""
Download task results, generate jsonl format file
Query parameters support: time range (based on creation time), task_id
"""
logging.info(f"🚀 download task results start, start_time: {start_time}, end_time: {end_time}, task_id: {task_id}")
try:
start_datetime = None
end_datetime = None
if start_time:
try:
start_datetime = datetime.strptime(start_time, "%Y-%m-%d %H:%M:%S")
except ValueError:
raise ValueError("❌ start_time格式错误,请使用 YYYY-MM-DD HH:MM:SS 格式")
if end_time:
try:
end_datetime = datetime.strptime(end_time, "%Y-%m-%d %H:%M:%S")
except ValueError:
raise ValueError("❌ end_time格式错误,请使用 YYYY-MM-DD HH:MM:SS 格式")
results = await task_manager.query_and_download_task_results(
start_time=start_datetime,
end_time=end_datetime,
task_id=task_id,
page_size=page_size
)
if not results:
logging.info("📄 no task results found")
def generate_empty():
yield ""
return StreamingResponse(
generate_empty(),
media_type="application/jsonl",
headers={"Content-Disposition": "attachment; filename=task_results_empty.jsonl"}
)
# Generate jsonl content
def generate_jsonl():
for result in results:
yield json.dumps(result, ensure_ascii=False) + "\n"
# Generate file name
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"task_results_{timestamp}.jsonl"
logging.info(f"✅ download task results success, total: {len(results)} results")
return StreamingResponse(
generate_jsonl(),
media_type="application/jsonl",
headers={"Content-Disposition": f"attachment; filename={filename}"}
)
except Exception as err:
traceback.print_exc()
logging.error(f"❌ download task results failed, err is {err}")
raise ValueError(f"❌ download task results failed: {str(err)}")
@@ -0,0 +1,259 @@
import inspect
import json
import inspect
import json
import logging
import time
import uuid
from typing import Generator, Iterator, AsyncGenerator, Optional
from aworld.core.task import Task
from aworld.utils.common import get_local_ip
from fastapi import status, HTTPException
from fastapi.concurrency import run_in_threadpool
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from aworldspace.base import AGENT_SPACE
from aworldspace.utils.utils import get_last_user_message
from base import OpenAIChatCompletionForm
async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
messages = [message.model_dump() for message in form_data.messages]
user_message = get_last_user_message(messages)
PIPELINES = await AGENT_SPACE.get_agents_meta()
PIPELINE_MODULES = await AGENT_SPACE.get_agent_modules()
if (
form_data.model not in PIPELINES
or PIPELINES[form_data.model]["type"] == "filter"
):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Pipeline {form_data.model} not found",
)
def job():
pipeline = PIPELINES[form_data.model]
pipeline_id = form_data.model
if pipeline["type"] == "manifold":
manifold_id, pipeline_id = pipeline_id.split(".", 1)
pipe = PIPELINE_MODULES[manifold_id].pipe
else:
pipe = PIPELINE_MODULES[pipeline_id].pipe
def process_line(model, line):
if isinstance(line, Task):
task_output_meta = line.outputs._metadata
line = openai_chat_chunk_message_template(model, "", task_output_meta=task_output_meta)
return f"data: {json.dumps(line)}\n\n"
if isinstance(line, BaseModel):
line = line.model_dump_json()
line = f"data: {line}"
if isinstance(line, dict):
line = f"data: {json.dumps(line)}"
try:
line = line.decode("utf-8")
except Exception:
pass
if line.startswith("data:"):
return f"{line}\n\n"
else:
line = openai_chat_chunk_message_template(model, line)
return f"data: {json.dumps(line)}\n\n"
if form_data.stream:
async def stream_content():
async def execute_pipe(_pipe):
if inspect.iscoroutinefunction(_pipe):
return await _pipe(user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump())
else:
return _pipe(user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump())
try:
res = await execute_pipe(pipe)
# Directly return if the response is a StreamingResponse
if isinstance(res, StreamingResponse):
async for data in res.body_iterator:
yield data
return
if isinstance(res, dict):
yield f"data: {json.dumps(res)}\n\n"
return
except Exception as e:
logging.error(f"Error: {e}")
import traceback
traceback.print_exc()
yield f"data: {json.dumps({'error': {'detail': str(e)}})}\n\n"
return
if isinstance(res, str):
message = openai_chat_chunk_message_template(form_data.model, res)
yield f"data: {json.dumps(message)}\n\n"
if isinstance(res, Iterator):
for line in res:
yield process_line(form_data.model, line)
if isinstance(res, AsyncGenerator):
async for line in res:
yield process_line(form_data.model, line)
logging.info(f"AsyncGenerator end...")
if isinstance(res, str) or isinstance(res, Generator) or isinstance(res, AsyncGenerator):
finish_message = openai_chat_chunk_message_template(
form_data.model, ""
)
finish_message["choices"][0]["finish_reason"] = "stop"
print(f"Pipe-Dataline:::: DONE")
yield f"data: {json.dumps(finish_message)}\n\n"
yield "data: [DONE]"
return StreamingResponse(stream_content(), media_type="text/event-stream")
else:
res = pipe(
user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump(),
)
logging.info(f"stream:false:{res}")
if isinstance(res, dict):
return res
elif isinstance(res, BaseModel):
return res.model_dump()
else:
message = ""
if isinstance(res, str):
message = res
if isinstance(res, Generator):
for stream in res:
message = f"{message}{stream}"
logging.info(f"stream:false:{message}")
return {
"id": f"{form_data.model}-{str(uuid.uuid4())}",
"object": "chat.completion",
"created": int(time.time()),
"model": form_data.model,
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": message,
},
"logprobs": None,
"finish_reason": "stop",
}
],
}
return await run_in_threadpool(job)
async def call_pipeline(form_data: OpenAIChatCompletionForm):
messages = [message.model_dump() for message in form_data.messages]
user_message = get_last_user_message(messages)
PIPELINES = await AGENT_SPACE.get_agents_meta()
PIPELINE_MODULES = await AGENT_SPACE.get_agent_modules()
if (
form_data.model not in PIPELINES
or PIPELINES[form_data.model]["type"] == "filter"
):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Pipeline {form_data.model} not found",
)
pipeline = PIPELINES[form_data.model]
pipeline_id = form_data.model
if pipeline["type"] == "manifold":
manifold_id, pipeline_id = pipeline_id.split(".", 1)
pipe = PIPELINE_MODULES[manifold_id].pipe
else:
pipe = PIPELINE_MODULES[pipeline_id].pipe
if form_data.stream:
async def execute_pipe(_pipe):
if inspect.iscoroutinefunction(_pipe):
return await _pipe(user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump())
else:
return _pipe(user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump())
res = await execute_pipe(pipe)
return res
else:
if not inspect.iscoroutinefunction(pipe):
return await run_in_threadpool(
pipe,
user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump()
)
else:
return await pipe(
user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump()
)
def openai_chat_chunk_message_template(
model: str,
content: Optional[str] = None,
tool_calls: Optional[list[dict]] = None,
usage: Optional[dict] = None,
**kwargs
) -> dict:
template = openai_chat_message_template(model, **kwargs)
template["object"] = "chat.completion.chunk"
template["choices"][0]["index"] = 0
template["choices"][0]["delta"] = {}
if content:
template["choices"][0]["delta"]["content"] = content
if tool_calls:
template["choices"][0]["delta"]["tool_calls"] = tool_calls
if not content and not tool_calls:
template["choices"][0]["finish_reason"] = "stop"
if usage:
template["usage"] = usage
return template
def openai_chat_message_template(model: str, **kwargs):
return {
"id": f"{model}-{str(uuid.uuid4())}",
"created": int(time.time()),
"model": model,
"node_id": get_local_ip(),
"task_output_meta": kwargs.get("task_output_meta"),
"choices": [{"index": 0, "logprobs": None, "finish_reason": None}],
}
@@ -0,0 +1,197 @@
import importlib.util
import json
import logging
import os
import subprocess
import sys
import traceback
from aworldspace.base import AGENT_SPACE
import aworld.trace as trace # noqa
from config import AGENTS_DIR
if not os.path.exists(AGENTS_DIR):
os.makedirs(AGENTS_DIR)
PIPELINES = {}
PIPELINE_MODULES = {}
def get_all_pipelines():
pipelines = {}
for pipeline_id in PIPELINE_MODULES.keys():
pipeline = PIPELINE_MODULES[pipeline_id]
if hasattr(pipeline, "type"):
if pipeline.type == "manifold":
manifold_pipelines = []
# Check if pipelines is a function or a list
if callable(pipeline.pipelines):
manifold_pipelines = pipeline.pipelines()
else:
manifold_pipelines = pipeline.pipelines
for p in manifold_pipelines:
manifold_pipeline_id = f'{pipeline_id}.{p["id"]}'
manifold_pipeline_name = p["name"]
if hasattr(pipeline, "name"):
manifold_pipeline_name = (
f"{pipeline.name}{manifold_pipeline_name}"
)
pipelines[manifold_pipeline_id] = {
"module": pipeline_id,
"type": pipeline.type if hasattr(pipeline, "type") else "pipe",
"id": manifold_pipeline_id,
"name": manifold_pipeline_name,
"valves": (
pipeline.valves if hasattr(pipeline, "valves") else None
),
}
if pipeline.type == "filter":
pipelines[pipeline_id] = {
"module": pipeline_id,
"type": (pipeline.type if hasattr(pipeline, "type") else "pipe"),
"id": pipeline_id,
"name": (
pipeline.name if hasattr(pipeline, "name") else pipeline_id
),
"pipelines": (
pipeline.valves.pipelines
if hasattr(pipeline, "valves")
and hasattr(pipeline.valves, "pipelines")
else []
),
"priority": (
pipeline.valves.priority
if hasattr(pipeline, "valves")
and hasattr(pipeline.valves, "priority")
else 0
),
"valves": pipeline.valves if hasattr(pipeline, "valves") else None,
}
else:
pipelines[pipeline_id] = {
"module": pipeline_id,
"type": (pipeline.type if hasattr(pipeline, "type") else "pipe"),
"id": pipeline_id,
"name": (pipeline.name if hasattr(pipeline, "name") else pipeline_id),
"valves": pipeline.valves if hasattr(pipeline, "valves") else None,
}
return pipelines
def parse_frontmatter(content):
frontmatter = {}
for line in content.split("\n"):
if ":" in line:
key, value = line.split(":", 1)
frontmatter[key.strip().lower()] = value.strip()
return frontmatter
def install_frontmatter_requirements(requirements):
if requirements:
req_list = [req.strip() for req in requirements.split(",")]
for req in req_list:
print(f"Installing requirement: {req}")
subprocess.check_call([sys.executable, "-m", "pip", "install", req])
else:
print("No requirements found in frontmatter.")
async def load_module_from_path(module_name, module_path):
try:
# Read the module content
with open(module_path, "r") as file:
content = file.read()
# Parse frontmatter
frontmatter = {}
if content.startswith('"""'):
end = content.find('"""', 3)
if end != -1:
frontmatter_content = content[3:end]
frontmatter = parse_frontmatter(frontmatter_content)
# Install requirements if specified
if "requirements" in frontmatter:
install_frontmatter_requirements(frontmatter["requirements"])
# Load the module
spec = importlib.util.spec_from_file_location(module_name, module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
logging.info(f"Loaded module start: {module.__name__}")
if hasattr(module, "Pipeline"):
return module.Pipeline()
else:
logging.info(f"Loaded module failed: {module.__name__ } No Pipeline class found")
raise Exception("No Pipeline class found")
except Exception as e:
logging.info(f"Error loading module: {module_name}, error is {e}")
traceback.print_exc()
# Move the file to the error folder
failed_pipelines_folder = os.path.join(AGENTS_DIR, "failed")
if not os.path.exists(failed_pipelines_folder):
os.makedirs(failed_pipelines_folder)
# failed_file_path = os.path.join(failed_pipelines_folder, f"{module_name}.py")
# if module_path.__contains__(PIPELINES_DIR):
# os.rename(module_path, failed_file_path)
print(e)
return None
async def load_modules_from_directory(directory):
logging.info(f"load_modules_from_directory: {directory}")
global PIPELINE_MODULES
for filename in os.listdir(directory):
if filename.endswith(".py"):
module_name = filename[:-3] # Remove the .py extension
module_path = os.path.join(directory, filename)
# Create subfolder matching the filename without the .py extension
subfolder_path = os.path.join(directory, module_name)
if not os.path.exists(subfolder_path):
os.makedirs(subfolder_path)
logging.info(f"Created subfolder: {subfolder_path}")
# Create a valves.json file if it doesn't exist
valves_json_path = os.path.join(subfolder_path, "valves.json")
if not os.path.exists(valves_json_path):
with open(valves_json_path, "w") as f:
json.dump({}, f)
logging.info(f"Created valves.json in: {subfolder_path}")
pipeline = await load_module_from_path(module_name, module_path)
if pipeline:
# Overwrite pipeline.valves with values from valves.json
if os.path.exists(valves_json_path):
with open(valves_json_path, "r") as f:
valves_json = json.load(f)
if hasattr(pipeline, "valves"):
ValvesModel = pipeline.valves.__class__
# Create a ValvesModel instance using default values and overwrite with valves_json
combined_valves = {
**pipeline.valves.model_dump(),
**valves_json,
}
valves = ValvesModel(**combined_valves)
pipeline.valves = valves
logging.info(f"Updated valves for module: {module_name}")
pipeline_id = pipeline.id if hasattr(pipeline, "id") else module_name
PIPELINE_MODULES[pipeline_id] = pipeline
logging.info(f"Loaded module success: {module_name}")
else:
logging.warning(f"No Pipeline class found in {module_name}")
AGENT_SPACE.agent_modules = PIPELINE_MODULES
AGENT_SPACE.agents_meta = get_all_pipelines()
@@ -0,0 +1,75 @@
import logging
import os
from datetime import datetime
from aworld.models.model_response import ModelResponse
from base import AworldTask, AworldTaskResult
from config import ROOT_LOG
class TaskLogger:
"""任务提交日志记录器"""
def __init__(self, log_file: str = "aworld_task_submissions.log"):
self.log_file = os.path.join(ROOT_LOG, 'task_logs' , log_file)
self._ensure_log_file_exists()
def _ensure_log_file_exists(self):
"""确保日志文件存在"""
if not os.path.exists(self.log_file):
os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
with open(self.log_file, 'w', encoding='utf-8') as f:
f.write("# Aworld Task Submission Log\n")
f.write(
"# Format: [timestamp] task_id | agent_id | server | status | agent_answer | correct_answer | is_correct | details\n\n")
def log_task_submission(self, task: AworldTask, status: str, details: str = "",
task_result: AworldTaskResult = None):
"""记录任务提交日志"""
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
log_entry = f"[{timestamp}] {task.task_id} | {task.agent_id} | {task.node_id} | {status} | {task_result.data.get('agent_answer') if task_result and task_result.data else None} | {task_result.data.get('correct_answer') if task_result and task_result.data else None} | {task_result.data.get('gaia_correct') if task_result and task_result.data else None} |{details}\n"
try:
with open(self.log_file, 'a', encoding='utf-8') as f:
f.write(log_entry)
except Exception as e:
logging.error(f"Failed to write task submission log: {e}")
def log_task_result(self, task: AworldTask, result: ModelResponse):
try:
date_str = datetime.now().strftime("%Y%m%d")
result_dir = os.path.join(ROOT_LOG, 'task_logs', 'result', date_str)
os.makedirs(result_dir, exist_ok=True)
md_file = f"{result_dir}/{task.task_id}.md"
content_parts = []
if hasattr(result, 'content') and result.content:
if isinstance(result.content, list):
content_parts.extend(result.content)
else:
content_parts.append(str(result.content))
file_exists = os.path.exists(md_file)
with open(md_file, 'a', encoding='utf-8') as f:
if not file_exists:
f.write(f"# Task Result: {task.task_id}\n\n")
f.write(f"**Agent ID:** {task.agent_id}\n\n")
f.write(f"**Timestamp:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
f.write("## Content\n\n")
if content_parts:
for i, content in enumerate(content_parts, 1):
f.write(f"{content}\n\n")
else:
f.write("No content available.\n\n")
return md_file
except Exception as e:
logging.error(f"Failed to write task result log: {e}")
return None
task_logger = TaskLogger(log_file=f"aworld_task_submissions_{datetime.now().strftime('%Y%m%d')}.log")
@@ -0,0 +1,199 @@
import os
def load_all_mcp_config():
return {
"mcpServers": {
"e2b-server": {
"command": "npx",
"args": [
"-y",
"@e2b/mcp-server"
],
"env": {
"E2B_API_KEY": os.environ["E2B_API_KEY"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"filesystem": {
"command": "npx",
"args": [
"-y",
"@modelcontextprotocol/server-filesystem",
"${FILESYSTEM_SERVER_WORKDIR}"
]
},
"terminal-controller": {
"command": "python",
"args": [
"-m",
"terminal_controller"
],
"env": {
"SESSION_REQUEST_CONNECT_TIMEOUT": "300"
}
},
"calculator": {
"command": "python",
"args": [
"-m",
"mcp_server_calculator"
],
"env": {
"SESSION_REQUEST_CONNECT_TIMEOUT": "20"
}
},
"excel": {
"command": "uvx",
"args": ["excel-mcp-server", "stdio"],
"env": {
"EXCEL_MCP_PAGING_CELLS_LIMIT": "4000",
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"google-search": {
"command": "npx",
"args": [
"-y",
"@adenot/mcp-google-search"
],
"env": {
"GOOGLE_API_KEY": os.environ["GOOGLE_API_KEY"],
"GOOGLE_SEARCH_ENGINE_ID": os.environ["GOOGLE_CSE_ID"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"ms-playwright": {
"command": "npx",
"args": [
"@playwright/mcp@latest",
"--no-sandbox",
"--headless",
"--isolated"
],
"env": {
"PLAYWRIGHT_TIMEOUT": "120000",
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"audio_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.audio_server"
],
"env": {
"AUDIO_LLM_API_KEY": os.environ["AUDIO_LLM_API_KEY"],
"AUDIO_LLM_BASE_URL": os.environ["AUDIO_LLM_BASE_URL"],
"AUDIO_LLM_MODEL_NAME": os.environ["AUDIO_LLM_MODEL_NAME"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"image_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.image_server"
],
"env": {
"LLM_API_KEY": os.environ.get("LLM_API_KEY"),
"LLM_MODEL_NAME": os.environ.get("LLM_MODEL_NAME"),
"LLM_BASE_URL": os.environ.get("LLM_BASE_URL"),
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"youtube_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.youtube_server"
],
"env": {
"CHROME_DRIVER_PATH": os.environ['CHROME_DRIVER_PATH'],
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"video_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.video_server"
],
"env": {
"LLM_API_KEY": os.environ.get("LLM_API_KEY"),
"LLM_MODEL_NAME": os.environ.get("LLM_MODEL_NAME"),
"LLM_BASE_URL": os.environ.get("LLM_BASE_URL"),
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"search_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.search_server"
],
"env": {
"GOOGLE_API_KEY": os.environ["GOOGLE_API_KEY"],
"GOOGLE_CSE_ID": os.environ["GOOGLE_CSE_ID"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"download_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.download_server"
],
"env": {
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"document_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.document_server"
],
"env": {
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"browser_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.browser_server"
],
"env": {
"LLM_API_KEY": os.environ.get("LLM_API_KEY"),
"LLM_MODEL_NAME": os.environ.get("LLM_MODEL_NAME"),
"LLM_BASE_URL": os.environ.get("LLM_BASE_URL"),
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"reasoning_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.reasoning_server"
],
"env": {
"LLM_API_KEY": os.environ.get("LLM_API_KEY"),
"LLM_MODEL_NAME": os.environ.get("LLM_MODEL_NAME"),
"LLM_BASE_URL": os.environ.get("LLM_BASE_URL"),
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"e2b-code-server": {
"command": "python",
"args": [
"-m",
"mcp_servers.e2b_code_server"
],
"env": {
"E2B_API_KEY": os.environ["E2B_API_KEY"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
}
}
@@ -0,0 +1,344 @@
import json
import re
import string
from pathlib import Path
from typing import Any, Dict, List, Optional
from loguru import logger
from tabulate import tabulate
def normalize_str(input_str, remove_punct=True) -> str:
no_spaces = re.sub(r"\s", "", input_str)
if remove_punct:
translator = str.maketrans("", "", string.punctuation)
return no_spaces.lower().translate(translator)
else:
return no_spaces.lower()
def split_string(s: str, char_list: Optional[List[str]] = None) -> list[str]:
if char_list is None:
char_list = [",", ";"]
pattern = f"[{''.join(char_list)}]"
return re.split(pattern, s)
def normalize_number_str(number_str: str) -> float:
for char in ["$", "%", ","]:
number_str = number_str.replace(char, "")
try:
return float(number_str)
except ValueError:
logger.error(f"String {number_str} cannot be normalized to number str.")
return float("inf")
def question_scorer(model_answer: str, ground_truth: str) -> bool:
def is_float(element: Any) -> bool:
try:
float(element)
return True
except ValueError:
return False
try:
if is_float(ground_truth):
logger.info(f"Evaluating {model_answer} as a number.")
normalized_answer = normalize_number_str(model_answer)
return normalized_answer == float(ground_truth)
elif any(char in ground_truth for char in [",", ";"]):
logger.info(f"Evaluating {model_answer} as a comma separated list.")
gt_elems = split_string(ground_truth)
ma_elems = split_string(model_answer)
if len(gt_elems) != len(ma_elems):
logger.warning("Answer lists have different lengths, returning False.")
return False
comparisons = []
for ma_elem, gt_elem in zip(ma_elems, gt_elems):
if is_float(gt_elem):
normalized_ma_elem = normalize_number_str(ma_elem)
comparisons.append(normalized_ma_elem == float(gt_elem))
else:
ma_elem = normalize_str(ma_elem, remove_punct=False)
gt_elem = normalize_str(gt_elem, remove_punct=False)
comparisons.append(ma_elem == gt_elem)
return all(comparisons)
else:
logger.info(f"Evaluating {model_answer} as a string.")
ma_elem = normalize_str(model_answer)
gt_elem = normalize_str(ground_truth)
return ma_elem == gt_elem
except Exception as e:
logger.error(f"Error during evaluation: {e}")
return False
def load_dataset_meta(path: str, split: str = "validation"):
data_dir = Path(path) / split
dataset = []
with open(data_dir / "metadata.jsonl", "r", encoding="utf-8") as metaf:
lines = metaf.readlines()
for line in lines:
data = json.loads(line)
if data["task_id"] == "0-0-0-0-0":
continue
if data["file_name"]:
data["file_name"] = data_dir / data["file_name"]
dataset.append(data)
return dataset
def load_dataset_meta_dict(path: str, split: str = "validation"):
data_dir = Path(path) / split
dataset = {}
with open(data_dir / "metadata.jsonl", "r", encoding="utf-8") as metaf:
lines = metaf.readlines()
for line in lines:
data = json.loads(line)
if data["task_id"] == "0-0-0-0-0":
continue
if data["file_name"]:
data["file_name"] = data_dir / data["file_name"]
dataset[data["task_id"]] = data
return dataset
def add_file_path(
task: Dict[str, Any], file_path: str = "./gaia_dataset", split: str = "validation"
):
if task["file_name"]:
file_path = Path(f"{file_path}/{split}") / task["file_name"]
if file_path.suffix in [".pdf", ".docx", ".doc", ".txt"]:
task["Question"] += f" Here are the necessary document files: {file_path}"
elif file_path.suffix in [".jpg", ".jpeg", ".png"]:
task["Question"] += f" Here are the necessary image files: {file_path}"
elif file_path.suffix in [".xlsx", "xls", ".csv"]:
task["Question"] += (
f" Here are the necessary table files: {file_path}, for processing excel file,"
" you can use the excel tool or write python code to process the file"
" step-by-step and get the information."
)
elif file_path.suffix in [".py"]:
task["Question"] += f" Here are the necessary python files: {file_path}"
else:
task["Question"] += f" Here are the necessary files: {file_path}"
return task
def report_results(entries):
# Initialize counters
total_entries = len(entries)
total_correct = 0
# Initialize level statistics
level_stats = {}
# Process each entry
for entry in entries:
level = entry.get("level")
is_correct = entry.get("is_correct", False)
# Initialize level stats if not already present
if level not in level_stats:
level_stats[level] = {"total": 0, "correct": 0, "accuracy": 0}
# Update counters
level_stats[level]["total"] += 1
if is_correct:
total_correct += 1
level_stats[level]["correct"] += 1
# Calculate accuracy for each level
for level, stats in level_stats.items():
if stats["total"] > 0:
stats["accuracy"] = (stats["correct"] / stats["total"]) * 100
# Print overall statistics with colorful logging
logger.info("Overall Statistics:")
overall_accuracy = (total_correct / total_entries) * 100
# Create overall statistics table
overall_table = [
["Total Entries", total_entries],
["Total Correct", total_correct],
["Overall Accuracy", f"{overall_accuracy:.2f}%"],
]
logger.success(tabulate(overall_table, tablefmt="grid"))
logger.info("")
# Create level statistics table
logger.info("Statistics by Level:")
level_table = []
headers = ["Level", "Total Entries", "Correct Answers", "Accuracy"]
for level in sorted(level_stats.keys()):
stats = level_stats[level]
level_table.append(
[level, stats["total"], stats["correct"], f"{stats['accuracy']:.2f}%"]
)
logger.success(tabulate(level_table, headers=headers, tablefmt="grid"))
import uuid
import time
from typing import List
import inspect
from typing import get_type_hints, Tuple
def stream_message_template(model: str, message: str):
return {
"id": f"{model}-{str(uuid.uuid4())}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [
{
"index": 0,
"delta": {"content": message},
"logprobs": None,
"finish_reason": None,
}
],
}
def get_last_user_message(messages: List[dict]) -> str:
for message in reversed(messages):
if message["role"] == "user":
if isinstance(message["content"], list):
for item in message["content"]:
if item["type"] == "text":
return item["text"]
return message["content"]
return None
def get_last_assistant_message(messages: List[dict]) -> str:
for message in reversed(messages):
if message["role"] == "assistant":
if isinstance(message["content"], list):
for item in message["content"]:
if item["type"] == "text":
return item["text"]
return message["content"]
return None
def get_system_message(messages: List[dict]) -> dict:
for message in messages:
if message["role"] == "system":
return message
return None
def remove_system_message(messages: List[dict]) -> List[dict]:
return [message for message in messages if message["role"] != "system"]
def pop_system_message(messages: List[dict]) -> Tuple[dict, List[dict]]:
return get_system_message(messages), remove_system_message(messages)
def add_or_update_system_message(content: str, messages: List[dict]) -> List[dict]:
"""
Adds a new system message at the beginning of the messages list
or updates the existing system message at the beginning.
:param msg: The message to be added or appended.
:param messages: The list of message dictionaries.
:return: The updated list of message dictionaries.
"""
if messages and messages[0].get("role") == "system":
messages[0]["content"] += f"{content}\n{messages[0]['content']}"
else:
# Insert at the beginning
messages.insert(0, {"role": "system", "content": content})
return messages
def doc_to_dict(docstring):
lines = docstring.split("\n")
description = lines[1].strip()
param_dict = {}
for line in lines:
if ":param" in line:
line = line.replace(":param", "").strip()
param, desc = line.split(":", 1)
param_dict[param.strip()] = desc.strip()
ret_dict = {"description": description, "params": param_dict}
return ret_dict
def get_tools_specs(tools) -> List[dict]:
function_list = [
{"name": func, "function": getattr(tools, func)}
for func in dir(tools)
if callable(getattr(tools, func)) and not func.startswith("__")
]
specs = []
for function_item in function_list:
function_name = function_item["name"]
function = function_item["function"]
function_doc = doc_to_dict(function.__doc__ or function_name)
specs.append(
{
"name": function_name,
# TODO: multi-line desc?
"description": function_doc.get("description", function_name),
"parameters": {
"type": "object",
"properties": {
param_name: {
"type": param_annotation.__name__.lower(),
**(
{
"enum": (
param_annotation.__args__
if hasattr(param_annotation, "__args__")
else None
)
}
if hasattr(param_annotation, "__args__")
else {}
),
"description": function_doc.get("params", {}).get(
param_name, param_name
),
}
for param_name, param_annotation in get_type_hints(
function
).items()
if param_name != "return"
},
"required": [
name
for name, param in inspect.signature(
function
).parameters.items()
if param.default is param.empty
],
},
}
)
return specs