Files
EPEEAIKit/art-agent/backend/app/agent/tools.py

69 lines
2.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Agent 可调用的工具定义和实现。"""
from app.services.image_gen import generate_images
# OpenAI Function Calling 格式的工具定义
TOOL_DEFINITIONS = [
{
"type": "function",
"function": {
"name": "generate_image",
"description": (
"根据文字描述生成图片。prompt 必须是英文。"
"如果用户提供了参考图,会自动传入 ref_image_url 参数。"
"若当前生图模型为 Stable Diffusion XL可在正提示后单独一行写 ---NEGATIVE--- 再写负向提示;"
"不传则服务端会使用该模型的默认负向词。"
),
"parameters": {
"type": "object",
"properties": {
"prompt": {
"type": "string",
"description": "英文图片描述 prompt详细描述要生成的图片内容、风格、颜色等",
},
"num_images": {
"type": "integer",
"description": "生成图片数量1-4 张",
"default": 1,
"minimum": 1,
"maximum": 4,
},
},
"required": ["prompt"],
},
},
}
]
async def execute_tool(
tool_name: str,
arguments: dict,
ref_image_urls: list[str] | None = None,
image_model: str | None = None,
) -> dict:
"""执行工具调用,返回结果。"""
if tool_name == "generate_image":
prompt = arguments["prompt"]
num_images = arguments.get("num_images", 1)
result = await generate_images(
prompt=prompt,
num_images=num_images,
ref_image_urls=ref_image_urls,
model_id=image_model,
)
valid_urls = [u for u in result.urls if not u.startswith("[")]
errors = [u for u in result.urls if u.startswith("[")]
return {
"success": len(valid_urls) > 0,
"images": valid_urls,
"errors": errors,
"prompt_used": prompt,
"effective_prompt": result.effective_prompt,
"negative_prompt": result.negative_prompt,
"model_name": result.model_name,
"model_id": result.model_id,
}
return {"success": False, "error": f"未知工具: {tool_name}"}