zhanggangyong/agent_management
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
zhanggangyong
2026-03-02 17:11:11 +00:00
parent a34d5a081e
commit ba7e4f3a30
3 changed files with 400 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent_templates/agents/ad_creator_agent/ad_creator_agent.Dockerfile
+2 -4
View File
@@ -11,7 +11,8 @@ RUN pip install --no-cache-dir \
uvicorn[standard]==0.27.0 \
pydantic==2.5.3 \
aiohttp>=3.9.0 \
python-multipart>=0.0.6
python-multipart>=0.0.6 \
azure-storage-blob>=12.19.0
COPY common/agent_callback_utils.py /app/common/
RUN touch /app/common/__init__.py
@@ -21,12 +22,9 @@ COPY agents/ad_creator_agent/ad_creator_agent.py /app/
ENV PYTHONUNBUFFERED=1
ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8000
ENV OUTPUT_DIR=/app/outputs
ENV AGENT_CALLBACK_URL=http://mcp-server:8002/api/v1/billing/agent-callback
RUN mkdir -p /app/outputs/images /app/outputs/videos
HEALTHCHECK --interval=30s --timeout=10s --start-period=20s --retries=3 \
CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health').read()" || exit 1
agent_templates/agents/ad_creator_agent/ad_creator_agent.py
+397 -89
View File
@@ -2,6 +2,7 @@
Ad Creator Agent - 多模态广告创意生成 Agent
通过素材（文字描述/参考图片）生成广告图片或视频
支持模型：Gemini 3 Pro Image / GPT Image 1 / DALL-E 3 / Sora 2
生成文件上传至 Azure Blob Storage，返回带 SAS token 的公开访问 URL
"""
import os
import sys
@@ -10,17 +11,20 @@ import uuid
import json
import base64
import logging
import asyncio
import aiohttp
from typing import Optional, List, Dict, Any
from typing import Optional, List, Dict, Any, AsyncGenerator
from datetime import datetime
from pathlib import Path
from enum import Enum
from io import BytesIO
from fastapi import FastAPI, HTTPException, Header, Depends, UploadFile, File, Form, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, StreamingResponse
from pydantic import BaseModel, Field
import uvicorn
from azure.storage.blob import BlobServiceClient, ContentSettings
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -41,7 +45,6 @@ SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8000"))
POD_NAME = os.getenv("POD_NAME", "ad-creator-agent")
USER_ID = os.getenv("USER_ID", "")
OUTPUT_DIR = os.getenv("OUTPUT_DIR", "/app/outputs")
LLM_BASE_URL = os.getenv(
"LLM_BASE_URL",
@@ -53,9 +56,82 @@ DEFAULT_IMAGE_MODEL = os.getenv("DEFAULT_IMAGE_MODEL", "taiji/gemini-3-pro-image
DEFAULT_TEXT_MODEL = os.getenv("DEFAULT_TEXT_MODEL", "taiji/gpt-4o-mini")
DEFAULT_VIDEO_MODEL = os.getenv("DEFAULT_VIDEO_MODEL", "taiji/sora-2")
Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
Path(f"{OUTPUT_DIR}/images").mkdir(parents=True, exist_ok=True)
Path(f"{OUTPUT_DIR}/videos").mkdir(parents=True, exist_ok=True)
AZURE_STORAGE_CONNECTION_STRING = os.getenv(
"AZURE_STORAGE_CONNECTION_STRING",
"DefaultEndpointsProtocol=https;AccountName=agnettool;AccountKey=BCjWGrpArS35FThjW8wUBU8Bs/cqxsovRBsnuk/pE//R2p09EBcvuV8PuW8Klgh2bmTVjqeaDppB+AStWkcTOA==;EndpointSuffix=core.windows.net"
)
AZURE_BLOB_CONTAINER = os.getenv("AZURE_BLOB_CONTAINER", "multimodal")
AZURE_BLOB_SAS_TOKEN = os.getenv(
"AZURE_BLOB_SAS_TOKEN",
"sp=r&st=2026-03-02T15:55:34Z&se=2028-03-02T00:10:34Z&sv=2024-11-04&sr=c&sig=hv3949MK%2FBajcgvWFUnGzx3jZ4gz3A%2FDALvQzv9mGPQ%3D"
)
# ==================== Azure Blob Storage ====================
class BlobStorage:
"""Azure Blob Storage 管理器"""
def __init__(self):
self._client: Optional[BlobServiceClient] = None
if AZURE_STORAGE_CONNECTION_STRING:
try:
self._client = BlobServiceClient.from_connection_string(AZURE_STORAGE_CONNECTION_STRING)
container_client = self._client.get_container_client(AZURE_BLOB_CONTAINER)
if not container_client.exists():
container_client.create_container()
logger.info(f"Blob Storage 已连接: container={AZURE_BLOB_CONTAINER}")
except Exception as e:
logger.error(f"Blob Storage 连接失败: {e}")
self._client = None
@property
def enabled(self) -> bool:
return self._client is not None
def _content_type(self, filename: str) -> str:
ext = filename.rsplit(".", 1)[-1].lower()
return {
"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg",
"webp": "image/webp", "gif": "image/gif", "mp4": "video/mp4",
}.get(ext, "application/octet-stream")
def upload(self, data: bytes, blob_name: str) -> str:
"""上传二进制数据到 Blob，返回带 SAS 的公开 URL"""
container_client = self._client.get_container_client(AZURE_BLOB_CONTAINER)
content_settings = ContentSettings(content_type=self._content_type(blob_name))
container_client.upload_blob(
name=blob_name, data=data,
overwrite=True, content_settings=content_settings,
)
account_name = self._client.account_name
base_url = f"https://{account_name}.blob.core.windows.net/{AZURE_BLOB_CONTAINER}/{blob_name}"
if AZURE_BLOB_SAS_TOKEN:
return f"{base_url}?{AZURE_BLOB_SAS_TOKEN}"
return base_url
def list_blobs(self, prefix: str = None) -> List[dict]:
"""列出 Blob"""
container_client = self._client.get_container_client(AZURE_BLOB_CONTAINER)
blobs = []
for blob in container_client.list_blobs(name_starts_with=prefix):
account_name = self._client.account_name
base_url = f"https://{account_name}.blob.core.windows.net/{AZURE_BLOB_CONTAINER}/{blob.name}"
url = f"{base_url}?{AZURE_BLOB_SAS_TOKEN}" if AZURE_BLOB_SAS_TOKEN else base_url
blobs.append({
"filename": blob.name,
"url": url,
"size_bytes": blob.size,
"created_at": blob.last_modified.isoformat() if blob.last_modified else "",
})
return blobs
def delete_blob(self, blob_name: str):
container_client = self._client.get_container_client(AZURE_BLOB_CONTAINER)
container_client.delete_blob(blob_name)
blob_storage = BlobStorage()
# ==================== 模型枚举 ====================
@@ -172,7 +248,8 @@ async def startup_event():
logger.info(f"回调处理器已初始化: {callback_handler.callback_url}")
else:
logger.warning("回调模块未加载")
logger.info(f"Ad Creator Agent 启动: port={SERVICE_PORT}, output={OUTPUT_DIR}")
storage_mode = f"azure_blob({AZURE_BLOB_CONTAINER})" if blob_storage.enabled else "local"
logger.info(f"Ad Creator Agent 启动: port={SERVICE_PORT}, storage={storage_mode}")
logger.info(f"默认模型: image={DEFAULT_IMAGE_MODEL}, text={DEFAULT_TEXT_MODEL}, video={DEFAULT_VIDEO_MODEL}")
@@ -236,22 +313,21 @@ async def generate_image_gemini(
img_format = match.group(1).replace("+xml", "")
ext = "jpg" if img_format == "jpeg" else img_format
b64_data = match.group(2).replace("\n", "").replace(" ", "")
image_bytes = base64.b64decode(b64_data)
ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
filename = f"ad_gemini_{ts}.{ext}"
file_path = os.path.join(OUTPUT_DIR, "images", filename)
with open(file_path, "wb") as f:
f.write(base64.b64decode(b64_data))
logger.info(f"Gemini 图片已生成: {file_path} ({os.path.getsize(file_path)} bytes)")
return {
"success": True,
"file_path": file_path,
"filename": filename,
"url": f"/api/v1/files/{filename}",
"model": DEFAULT_IMAGE_MODEL,
}
if blob_storage.enabled:
blob_url = blob_storage.upload(image_bytes, filename)
logger.info(f"Gemini 图片已上传 Blob: {filename} ({len(image_bytes)} bytes)")
return {"success": True, "filename": filename, "url": blob_url, "model": DEFAULT_IMAGE_MODEL}
else:
file_path = os.path.join("/tmp", filename)
with open(file_path, "wb") as f:
f.write(image_bytes)
logger.info(f"Gemini 图片已生成(本地): {file_path} ({len(image_bytes)} bytes)")
return {"success": True, "filename": filename, "url": f"/api/v1/files/{filename}", "model": DEFAULT_IMAGE_MODEL}
async def generate_image_openai(
@@ -294,28 +370,27 @@ async def generate_image_openai(
ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
model_tag = model.split("/")[-1].replace("-", "")
filename = f"ad_{model_tag}_{ts}.png"
file_path = os.path.join(OUTPUT_DIR, "images", filename)
if b64_data:
with open(file_path, "wb") as f:
f.write(base64.b64decode(b64_data))
image_bytes = base64.b64decode(b64_data)
elif image_url:
async with session.get(image_url, timeout=aiohttp.ClientTimeout(total=30)) as dl_resp:
if dl_resp.status != 200:
return {"success": False, "error": f"下载图片失败: HTTP {dl_resp.status}"}
with open(file_path, "wb") as f:
f.write(await dl_resp.read())
image_bytes = await dl_resp.read()
else:
return {"success": False, "error": "API 响应中无图片数据"}
logger.info(f"OpenAI 图片已生成: {file_path} ({os.path.getsize(file_path)} bytes)")
return {
"success": True,
"file_path": file_path,
"filename": filename,
"url": f"/api/v1/files/{filename}",
"model": model,
}
if blob_storage.enabled:
blob_url = blob_storage.upload(image_bytes, filename)
logger.info(f"OpenAI 图片已上传 Blob: {filename} ({len(image_bytes)} bytes)")
return {"success": True, "filename": filename, "url": blob_url, "model": model}
else:
file_path = os.path.join("/tmp", filename)
with open(file_path, "wb") as f:
f.write(image_bytes)
logger.info(f"OpenAI 图片已生成(本地): {file_path} ({len(image_bytes)} bytes)")
return {"success": True, "filename": filename, "url": f"/api/v1/files/{filename}", "model": model}
async def generate_image_dispatch(
@@ -475,28 +550,27 @@ async def generate_video_sora(
ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
filename = f"ad_video_{ts}.mp4"
file_path = os.path.join(OUTPUT_DIR, "videos", filename)
if b64_data:
with open(file_path, "wb") as f:
f.write(base64.b64decode(b64_data))
video_bytes = base64.b64decode(b64_data)
elif video_url:
async with session.get(video_url, timeout=aiohttp.ClientTimeout(total=120)) as dl_resp:
if dl_resp.status != 200:
return {"success": False, "error": f"下载视频失败: HTTP {dl_resp.status}"}
with open(file_path, "wb") as f:
f.write(await dl_resp.read())
video_bytes = await dl_resp.read()
else:
return {"success": False, "error": "Sora 响应中无视频数据"}
logger.info(f"视频已生成: {file_path} ({os.path.getsize(file_path)} bytes)")
return {
"success": True,
"file_path": file_path,
"filename": filename,
"url": f"/api/v1/files/{filename}",
"model": model,
}
if blob_storage.enabled:
blob_url = blob_storage.upload(video_bytes, filename)
logger.info(f"视频已上传 Blob: {filename} ({len(video_bytes)} bytes)")
return {"success": True, "filename": filename, "url": blob_url, "model": model}
else:
file_path = os.path.join("/tmp", filename)
with open(file_path, "wb") as f:
f.write(video_bytes)
logger.info(f"视频已生成(本地): {file_path} ({len(video_bytes)} bytes)")
return {"success": True, "filename": filename, "url": f"/api/v1/files/{filename}", "model": model}
# ==================== API 端点 ====================
@@ -508,6 +582,8 @@ async def health_check():
"status": "healthy",
"service": "Ad Creator Agent",
"pod_name": POD_NAME,
"storage": "azure_blob" if blob_storage.enabled else "local",
"blob_container": AZURE_BLOB_CONTAINER if blob_storage.enabled else None,
"models": {
"image": DEFAULT_IMAGE_MODEL,
"text": DEFAULT_TEXT_MODEL,
@@ -520,13 +596,19 @@ async def health_check():
@app.get("/status")
async def status():
images = list(Path(f"{OUTPUT_DIR}/images").glob("*"))
videos = list(Path(f"{OUTPUT_DIR}/videos").glob("*"))
if blob_storage.enabled:
images = blob_storage.list_blobs(prefix="ad_")
img_count = sum(1 for b in images if not b["filename"].startswith("ad_video_"))
vid_count = sum(1 for b in images if b["filename"].startswith("ad_video_"))
else:
img_count = 0
vid_count = 0
return {
"status": "running",
"pod_name": POD_NAME,
"generated_images": len(images),
"generated_videos": len(videos),
"storage": "azure_blob" if blob_storage.enabled else "local",
"generated_images": img_count,
"generated_videos": vid_count,
"timestamp": datetime.utcnow().isoformat(),
}
@@ -674,60 +756,55 @@ async def api_generate_video(request: GenerateVideoRequest, api_key: str = Depen
@app.get("/api/v1/files/{filename}")
async def download_file(filename: str):
"""下载生成的文件"""
for subdir in ["images", "videos"]:
path = os.path.join(OUTPUT_DIR, subdir, filename)
if os.path.exists(path):
ext = filename.rsplit(".", 1)[-1].lower()
media_types = {
"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg",
"webp": "image/webp", "mp4": "video/mp4", "gif": "image/gif",
}
return FileResponse(path, media_type=media_types.get(ext, "application/octet-stream"), filename=filename)
"""获取文件（Blob 模式下 302 跳转到 Blob URL）"""
if blob_storage.enabled:
account_name = blob_storage._client.account_name
base_url = f"https://{account_name}.blob.core.windows.net/{AZURE_BLOB_CONTAINER}/{filename}"
url = f"{base_url}?{AZURE_BLOB_SAS_TOKEN}" if AZURE_BLOB_SAS_TOKEN else base_url
return RedirectResponse(url=url)
path = os.path.join("/tmp", filename)
if os.path.exists(path):
ext = filename.rsplit(".", 1)[-1].lower()
media_types = {
"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg",
"webp": "image/webp", "mp4": "video/mp4", "gif": "image/gif",
}
return FileResponse(path, media_type=media_types.get(ext, "application/octet-stream"), filename=filename)
raise HTTPException(status_code=404, detail="文件不存在")
@app.get("/api/v1/list-files")
async def list_files(file_type: str = "all"):
"""列出已生成的文件"""
"""列出已生成的文件（从 Blob Storage 列出）"""
result = {"images": [], "videos": []}
if file_type in ("all", "image"):
img_dir = Path(f"{OUTPUT_DIR}/images")
for f in sorted(img_dir.glob("*"), key=lambda p: p.stat().st_mtime, reverse=True):
result["images"].append({
"filename": f.name,
"url": f"/api/v1/files/{f.name}",
"size_bytes": f.stat().st_size,
"created_at": datetime.fromtimestamp(f.stat().st_mtime).isoformat(),
})
if file_type in ("all", "video"):
vid_dir = Path(f"{OUTPUT_DIR}/videos")
for f in sorted(vid_dir.glob("*"), key=lambda p: p.stat().st_mtime, reverse=True):
result["videos"].append({
"filename": f.name,
"url": f"/api/v1/files/{f.name}",
"size_bytes": f.stat().st_size,
"created_at": datetime.fromtimestamp(f.stat().st_mtime).isoformat(),
})
if blob_storage.enabled:
all_blobs = blob_storage.list_blobs(prefix="ad_")
for b in all_blobs:
if b["filename"].startswith("ad_video_"):
if file_type in ("all", "video"):
result["videos"].append(b)
else:
if file_type in ("all", "image"):
result["images"].append(b)
return result
@app.post("/api/v1/cleanup")
async def cleanup_files(max_age_hours: int = 24):
"""清理超过指定时间的旧文件"""
import time
cutoff = time.time() - max_age_hours * 3600
from datetime import timezone, timedelta
cutoff = datetime.now(timezone.utc) - timedelta(hours=max_age_hours)
deleted = 0
for subdir in ["images", "videos"]:
d = Path(f"{OUTPUT_DIR}/{subdir}")
for f in d.glob("*"):
if f.stat().st_mtime < cutoff:
f.unlink()
deleted += 1
if blob_storage.enabled:
all_blobs = blob_storage.list_blobs(prefix="ad_")
for b in all_blobs:
if b["created_at"] and datetime.fromisoformat(b["created_at"]) < cutoff:
try:
blob_storage.delete_blob(b["filename"])
deleted += 1
except Exception as e:
logger.warning(f"删除 blob {b['filename']} 失败: {e}")
return {"deleted": deleted, "max_age_hours": max_age_hours}
@@ -781,6 +858,237 @@ async def chat(request: ChatRequest, api_key: str = Depends(get_api_key)):
}
# ==================== MCP 端点 ====================
SERVER_NAME = "Ad Creator Agent"
MCP_TOOL_MAP = {
"generate_ad_image": None,
"generate_ad_copy": None,
"generate_full_ad": None,
"list_generated_files": None,
}
MCP_TOOL_LIST = [
{
"name": "generate_ad_image",
"description": "生成广告图片。支持 Gemini / GPT Image / DALL-E 模型，可指定风格、宽高比和品牌名。返回图片的公开 URL。",
"inputSchema": {
"type": "object",
"properties": {
"prompt": {"type": "string", "description": "广告图片描述（英文效果更好）"},
"model": {"type": "string", "description": "模型: taiji/gemini-3-pro-image-preview, taiji/gpt-image-1, taiji/dall-e-3"},
"aspect_ratio": {"type": "string", "description": "宽高比: 1:1, 16:9, 9:16, 4:3, 3:4"},
"style": {"type": "string", "description": "风格: modern, minimalist, luxury, playful, tech, vintage"},
"brand_name": {"type": "string", "description": "品牌名称"},
},
"required": ["prompt"],
},
},
{
"name": "generate_ad_copy",
"description": "生成广告文案方案，包含标题、正文、CTA、hashtags 以及配图 prompt。",
"inputSchema": {
"type": "object",
"properties": {
"product": {"type": "string", "description": "产品/服务描述"},
"target_audience": {"type": "string", "description": "目标受众"},
"tone": {"type": "string", "description": "语气: professional, casual, humorous, urgent, luxury"},
"platform": {"type": "string", "description": "投放平台: instagram, facebook, tiktok, billboard, general"},
"language": {"type": "string", "description": "语言: zh, en, ja"},
},
"required": ["product"],
},
},
{
"name": "generate_full_ad",
"description": "一键生成完整广告：先生成文案，再根据文案自动生成配图。",
"inputSchema": {
"type": "object",
"properties": {
"product": {"type": "string", "description": "产品/服务描述"},
"style": {"type": "string", "description": "广告风格"},
"brand_name": {"type": "string", "description": "品牌名称"},
"target_audience": {"type": "string", "description": "目标受众"},
"tone": {"type": "string", "description": "语气"},
"platform": {"type": "string", "description": "投放平台"},
"language": {"type": "string", "description": "语言: zh, en, ja"},
},
"required": ["product"],
},
},
{
"name": "list_generated_files",
"description": "列出已生成的广告素材文件（图片和视频）。",
"inputSchema": {
"type": "object",
"properties": {
"file_type": {"type": "string", "description": "类型: all, image, video"},
},
},
},
]
async def _mcp_generate_ad_image(api_key: str, **kwargs) -> str:
prompt = kwargs.get("prompt", "")
if kwargs.get("style"):
prompt = f"[{kwargs['style']} style] {prompt}"
if kwargs.get("brand_name"):
prompt = f"{prompt}. Brand: {kwargs['brand_name']}"
result = await generate_image_dispatch(
prompt=prompt, api_key=api_key,
model=kwargs.get("model"), aspect_ratio=kwargs.get("aspect_ratio", "1:1"),
)
return json.dumps(result, ensure_ascii=False, indent=2)
async def _mcp_generate_ad_copy(api_key: str, **kwargs) -> str:
result = await generate_ad_copy(
product=kwargs["product"], api_key=api_key,
target_audience=kwargs.get("target_audience"),
tone=kwargs.get("tone", "professional"),
platform=kwargs.get("platform", "general"),
language=kwargs.get("language", "zh"),
)
return json.dumps(result, ensure_ascii=False, indent=2)
async def _mcp_generate_full_ad(api_key: str, **kwargs) -> str:
copy_result = await generate_ad_copy(
product=kwargs["product"], api_key=api_key,
target_audience=kwargs.get("target_audience"),
tone=kwargs.get("tone", "professional"),
platform=kwargs.get("platform", "general"),
language=kwargs.get("language", "zh"),
)
image_prompt = copy_result.get("image_prompt", "") or f"Advertisement for: {kwargs['product']}"
if kwargs.get("style"):
image_prompt = f"[{kwargs['style']} style] {image_prompt}"
if kwargs.get("brand_name"):
image_prompt = f"{image_prompt}. Brand: {kwargs['brand_name']}"
image_result = await generate_image_dispatch(prompt=image_prompt, api_key=api_key)
return json.dumps({"success": True, "copy": copy_result, "image": image_result}, ensure_ascii=False, indent=2)
async def _mcp_list_files(api_key: str, **kwargs) -> str:
result = {"images": [], "videos": []}
if blob_storage.enabled:
all_blobs = blob_storage.list_blobs(prefix="ad_")
ft = kwargs.get("file_type", "all")
for b in all_blobs:
if b["filename"].startswith("ad_video_"):
if ft in ("all", "video"):
result["videos"].append(b)
else:
if ft in ("all", "image"):
result["images"].append(b)
return json.dumps(result, ensure_ascii=False, indent=2)
_MCP_HANDLERS = {
"generate_ad_image": _mcp_generate_ad_image,
"generate_ad_copy": _mcp_generate_ad_copy,
"generate_full_ad": _mcp_generate_full_ad,
"list_generated_files": _mcp_list_files,
}
sessions: Dict[str, Dict] = {}
def _get_api_key_from_request(request: Request) -> Optional[str]:
api_key = request.headers.get("api-key") or request.headers.get("api_key")
if not api_key:
auth = request.headers.get("Authorization")
if auth:
api_key = auth[7:] if auth.startswith("Bearer ") else auth
return api_key or LLM_API_KEY or None
async def _handle_mcp_request(data: Dict, session_id: str = None, api_key: str = None) -> Dict:
method = data.get("method")
params = data.get("params", {})
req_id = data.get("id")
if method == "tools/call" and not api_key:
return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32001, "message": "缺少 API Key"}}
try:
if method == "initialize":
session_id = session_id or str(uuid.uuid4())
sessions[session_id] = {"initialized": True}
return {
"jsonrpc": "2.0", "id": req_id,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": SERVER_NAME, "version": "1.0.0"},
},
}
elif method == "tools/list":
return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": MCP_TOOL_LIST}}
elif method == "tools/call":
tool_name = params.get("name")
args = params.get("arguments", {})
handler = _MCP_HANDLERS.get(tool_name)
if not handler:
raise ValueError(f"Unknown tool: {tool_name}")
result = await handler(api_key=api_key, **args)
return {"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": result}]}}
elif method == "ping":
return {"jsonrpc": "2.0", "id": req_id, "result": {}}
else:
raise ValueError(f"Unknown method: {method}")
except Exception as e:
return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32603, "message": str(e)}}
@app.post("/mcp")
async def mcp_endpoint(request: Request):
"""MCP HTTP 端点"""
try:
body = await request.json()
session_id = request.headers.get("x-mcp-session-id")
api_key = _get_api_key_from_request(request)
response = await _handle_mcp_request(body, session_id, api_key)
return JSONResponse(content=response, headers={"x-mcp-session-id": session_id or ""})
except Exception as e:
return JSONResponse(status_code=400, content={"jsonrpc": "2.0", "error": {"code": -32700, "message": str(e)}})
@app.get("/mcp/sse")
async def mcp_sse(request: Request):
"""MCP SSE 端点"""
session_id = request.headers.get("x-mcp-session-id") or str(uuid.uuid4())
async def stream() -> AsyncGenerator[str, None]:
yield f"data: {json.dumps({'type': 'connection', 'sessionId': session_id})}\n\n"
while True:
await asyncio.sleep(30)
yield f"data: {json.dumps({'type': 'ping'})}\n\n"
return StreamingResponse(stream(), media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "x-mcp-session-id": session_id})
@app.post("/mcp/sse")
async def mcp_sse_post(request: Request):
"""MCP SSE POST 端点"""
try:
body = await request.json()
session_id = request.headers.get("x-mcp-session-id") or str(uuid.uuid4())
api_key = _get_api_key_from_request(request)
async def stream() -> AsyncGenerator[str, None]:
response = await _handle_mcp_request(body, session_id, api_key)
yield f"data: {json.dumps(response)}\n\n"
return StreamingResponse(stream(), media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "x-mcp-session-id": session_id})
except Exception as e:
return JSONResponse(status_code=400, content={"jsonrpc": "2.0", "error": {"code": -32700, "message": str(e)}})
# ==================== 主入口 ====================
def main():
agent_templates/agents/search_agent/search_agent_main.py
+1 -1
View File
@@ -49,7 +49,7 @@ callback_handler: Optional[AgentCallbackHandler] = None
# 环境变量
USER_ID = os.getenv("USER_ID", "")
# FastAPI应用
app = FastAPI(
title="Intelligent Search AI Agent",