186 lines
5.2 KiB
Python
186 lines
5.2 KiB
Python
"""
|
|
Echo Agent - 简单的回显测试代理
|
|
用于测试 Agent Manager 的部署功能
|
|
"""
|
|
import os
|
|
import sys
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import Optional, Any
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel, Field
|
|
import uvicorn
|
|
|
|
# 添加 common 模块路径
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# 导入回调工具
|
|
try:
|
|
from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
|
|
CALLBACK_ENABLED = True
|
|
except ImportError:
|
|
CALLBACK_ENABLED = False
|
|
AgentCallbackHandler = None
|
|
CallbackContextManager = None
|
|
|
|
# 配置日志
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 环境变量
|
|
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
|
|
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8000"))
|
|
POD_NAME = os.getenv("POD_NAME", "echo-agent")
|
|
USER_ID = os.getenv("USER_ID", "")
|
|
|
|
# FastAPI 应用
|
|
app = FastAPI(
|
|
title="Echo Agent",
|
|
description="简单的回显测试代理，用于验证 Agent Manager 部署功能",
|
|
version="1.0.0"
|
|
)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# 回调处理器
|
|
callback_handler: Optional[AgentCallbackHandler] = None
|
|
|
|
|
|
# ==================== 请求/响应模型 ====================
|
|
|
|
class EchoRequest(BaseModel):
|
|
"""回显请求"""
|
|
message: str = Field(..., description="要回显的消息")
|
|
user_id: Optional[str] = Field(None, description="用户ID（用于计费回调）")
|
|
metadata: Optional[dict] = Field(None, description="附加元数据")
|
|
|
|
|
|
class EchoResponse(BaseModel):
|
|
"""回显响应"""
|
|
message: str
|
|
echo: str
|
|
pod_name: str
|
|
metadata: Optional[dict] = None
|
|
timestamp: str
|
|
|
|
|
|
class HealthResponse(BaseModel):
|
|
"""健康检查响应"""
|
|
status: str
|
|
pod_name: str
|
|
version: str
|
|
callback_enabled: bool
|
|
timestamp: str
|
|
|
|
|
|
# ==================== 生命周期 ====================
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
"""应用启动时初始化回调处理器"""
|
|
global callback_handler
|
|
|
|
if CALLBACK_ENABLED:
|
|
callback_handler = AgentCallbackHandler(
|
|
agent_name=POD_NAME,
|
|
user_id=USER_ID
|
|
)
|
|
logger.info(f"回调处理器已初始化: callback_url={callback_handler.callback_url}")
|
|
else:
|
|
logger.warning("回调模块未加载，计费回调功能不可用")
|
|
|
|
|
|
# ==================== API 端点 ====================
|
|
|
|
@app.get("/health", response_model=HealthResponse)
|
|
@app.get("/", response_model=HealthResponse)
|
|
async def health_check():
|
|
"""健康检查"""
|
|
return HealthResponse(
|
|
status="healthy",
|
|
pod_name=POD_NAME,
|
|
version="1.0.0",
|
|
callback_enabled=CALLBACK_ENABLED,
|
|
timestamp=datetime.utcnow().isoformat()
|
|
)
|
|
|
|
|
|
@app.post("/echo", response_model=EchoResponse)
|
|
async def echo(request: EchoRequest):
|
|
"""回显消息"""
|
|
logger.info(f"Echo: {request.message}")
|
|
|
|
# 使用回调上下文管理器
|
|
if CALLBACK_ENABLED and callback_handler and request.user_id:
|
|
with CallbackContextManager(
|
|
handler=callback_handler,
|
|
user_id=request.user_id,
|
|
request_id=f"echo-{int(datetime.utcnow().timestamp())}"
|
|
) as ctx:
|
|
ctx.add_tool("echo")
|
|
|
|
response = EchoResponse(
|
|
message=request.message,
|
|
echo=f"[Echo from {POD_NAME}] {request.message}",
|
|
pod_name=POD_NAME,
|
|
metadata=request.metadata,
|
|
timestamp=datetime.utcnow().isoformat()
|
|
)
|
|
else:
|
|
response = EchoResponse(
|
|
message=request.message,
|
|
echo=f"[Echo from {POD_NAME}] {request.message}",
|
|
pod_name=POD_NAME,
|
|
metadata=request.metadata,
|
|
timestamp=datetime.utcnow().isoformat()
|
|
)
|
|
|
|
return response
|
|
|
|
|
|
@app.get("/echo")
|
|
async def echo_get(message: str = "Hello", user_id: Optional[str] = None):
|
|
"""GET 方式回显"""
|
|
return await echo(EchoRequest(message=message, user_id=user_id))
|
|
|
|
|
|
@app.get("/info")
|
|
async def get_info():
|
|
"""获取 Agent 信息"""
|
|
return {
|
|
"agent_name": "Echo Agent",
|
|
"pod_name": POD_NAME,
|
|
"version": "1.0.0",
|
|
"callback_enabled": CALLBACK_ENABLED,
|
|
"callback_url": callback_handler.callback_url if callback_handler else None,
|
|
"capabilities": ["echo", "health_check"],
|
|
"environment": {
|
|
"SERVICE_HOST": SERVICE_HOST,
|
|
"SERVICE_PORT": SERVICE_PORT
|
|
},
|
|
"timestamp": datetime.utcnow().isoformat()
|
|
}
|
|
|
|
|
|
# ==================== 主入口 ====================
|
|
|
|
def main():
|
|
"""主函数"""
|
|
logger.info(f"启动 Echo Agent - {POD_NAME}")
|
|
logger.info(f"服务地址: {SERVICE_HOST}:{SERVICE_PORT}")
|
|
logger.info(f"回调功能: {'已启用' if CALLBACK_ENABLED else '未启用'}")
|
|
|
|
uvicorn.run(app, host=SERVICE_HOST, port=SERVICE_PORT, log_level="info")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|