From ba7e4f3a305f9b8ca6fc39c22bd700586a476f4b Mon Sep 17 00:00:00 2001 From: zhanggangyong Date: Mon, 2 Mar 2026 17:11:11 +0000 Subject: [PATCH] update --- .../ad_creator_agent.Dockerfile | 6 +- .../ad_creator_agent/ad_creator_agent.py | 486 ++++++++++++++---- .../agents/search_agent/search_agent_main.py | 2 +- 3 files changed, 400 insertions(+), 94 deletions(-) diff --git a/agent_templates/agents/ad_creator_agent/ad_creator_agent.Dockerfile b/agent_templates/agents/ad_creator_agent/ad_creator_agent.Dockerfile index b809906..606f55c 100644 --- a/agent_templates/agents/ad_creator_agent/ad_creator_agent.Dockerfile +++ b/agent_templates/agents/ad_creator_agent/ad_creator_agent.Dockerfile @@ -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 diff --git a/agent_templates/agents/ad_creator_agent/ad_creator_agent.py b/agent_templates/agents/ad_creator_agent/ad_creator_agent.py index df862a2..9f5403a 100644 --- a/agent_templates/agents/ad_creator_agent/ad_creator_agent.py +++ b/agent_templates/agents/ad_creator_agent/ad_creator_agent.py @@ -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(): diff --git a/agent_templates/agents/search_agent/search_agent_main.py b/agent_templates/agents/search_agent/search_agent_main.py index 134efcb..753e7a9 100644 --- a/agent_templates/agents/search_agent/search_agent_main.py +++ b/agent_templates/agents/search_agent/search_agent_main.py @@ -49,7 +49,7 @@ callback_handler: Optional[AgentCallbackHandler] = None # 环境变量 USER_ID = os.getenv("USER_ID", "") - + # FastAPI应用 app = FastAPI( title="Intelligent Search AI Agent",