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
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:
@@ -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 can’t 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()
|
||||
+794
@@ -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
|
||||
Reference in New Issue
Block a user