121 lines
4.1 KiB
Python
121 lines
4.1 KiB
Python
"""
|
||
Dev Agent API Server
|
||
使用 ADK 官方 ApiServer 构建 REST API 服务，支持：
|
||
- REST API 调用 agent（/run、/run_sse）
|
||
- 会话管理（创建/获取/删除，SQLite 持久化）
|
||
- Swagger UI 交互式文档（/docs）
|
||
- 上下文自动压缩
|
||
- 长期记忆（InMemory，后续可换向量库）
|
||
|
||
启动方式：
|
||
python api_server.py
|
||
|
||
主要端点：
|
||
GET /list-apps 列出所有 agent
|
||
POST /run 同步运行 agent
|
||
POST /run_sse 流式运行 agent（SSE）
|
||
GET /apps/{app}/users/{user}/sessions/{session} 获取会话
|
||
POST /apps/{app}/users/{user}/sessions/{session} 创建会话
|
||
GET /docs Swagger UI
|
||
"""
|
||
import os
|
||
import sys
|
||
|
||
# 确保项目根目录在 path 里
|
||
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||
if PROJECT_ROOT not in sys.path:
|
||
sys.path.insert(0, PROJECT_ROOT)
|
||
|
||
from dotenv import load_dotenv
|
||
load_dotenv(os.path.join(PROJECT_ROOT, "my_agent", ".env"))
|
||
|
||
# 强制 UTF-8
|
||
os.environ["PYTHONUTF8"] = "1"
|
||
|
||
import uvicorn
|
||
from google.adk.cli.api_server import ApiServer
|
||
from google.adk.cli.utils.base_agent_loader import BaseAgentLoader
|
||
from google.adk.sessions.sqlite_session_service import SqliteSessionService
|
||
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
|
||
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||
from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
|
||
from google.adk.evaluation.in_memory_eval_sets_manager import InMemoryEvalSetsManager
|
||
from google.adk.evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
|
||
from my_agent.app import dev_app
|
||
|
||
|
||
# 配置
|
||
HOST = os.getenv("API_SERVER_HOST", "0.0.0.0")
|
||
PORT = int(os.getenv("API_SERVER_PORT", "8001"))
|
||
|
||
# 数据目录
|
||
DATA_DIR = os.path.join(PROJECT_ROOT, "data")
|
||
os.makedirs(DATA_DIR, exist_ok=True)
|
||
|
||
|
||
class DevAgentLoader(BaseAgentLoader):
|
||
"""自定义 agent 加载器，直接返回我们的 App 对象（带 compaction 配置）"""
|
||
|
||
def load_agent(self, agent_name: str):
|
||
if agent_name == dev_app.name:
|
||
return dev_app
|
||
raise ValueError(f"Agent not found: {agent_name}")
|
||
|
||
def list_agents(self) -> list[str]:
|
||
return [dev_app.name]
|
||
|
||
|
||
def create_api_server() -> ApiServer:
|
||
"""构造 ApiServer 实例"""
|
||
# 会话服务：SQLite 持久化
|
||
session_service = SqliteSessionService(
|
||
db_path=os.path.join(DATA_DIR, "sessions.db")
|
||
)
|
||
|
||
# 记忆服务：长期记忆（先用内存版）
|
||
memory_service = InMemoryMemoryService()
|
||
|
||
# 工件服务
|
||
artifact_service = InMemoryArtifactService()
|
||
|
||
# 认证服务（暂不需要，内存版占位）
|
||
credential_service = InMemoryCredentialService()
|
||
|
||
# 评测集管理（暂不需要，占位）
|
||
eval_sets_manager = InMemoryEvalSetsManager()
|
||
eval_set_results_manager = LocalEvalSetResultsManager(agents_dir=DATA_DIR)
|
||
|
||
return ApiServer(
|
||
agent_loader=DevAgentLoader(),
|
||
session_service=session_service,
|
||
memory_service=memory_service,
|
||
artifact_service=artifact_service,
|
||
credential_service=credential_service,
|
||
eval_sets_manager=eval_sets_manager,
|
||
eval_set_results_manager=eval_set_results_manager,
|
||
agents_dir=os.path.join(PROJECT_ROOT, "my_agent"),
|
||
auto_create_session=True,
|
||
)
|
||
|
||
|
||
def main():
|
||
api_server = create_api_server()
|
||
fastapi_app = api_server.get_fast_api_app()
|
||
|
||
print("=" * 60)
|
||
print("Dev Agent API Server 启动中...")
|
||
print(f" 监听地址: http://{HOST}:{PORT}")
|
||
print(f" Swagger UI: http://{HOST}:{PORT}/docs")
|
||
print(f" 同步运行: POST http://{HOST}:{PORT}/run")
|
||
print(f" 流式运行: POST http://{HOST}:{PORT}/run_sse")
|
||
print(f" 列出agent: GET http://{HOST}:{PORT}/list-apps")
|
||
print(f" 会话持久化: SQLite ({DATA_DIR}/sessions.db)")
|
||
print(f" 上下文压缩: 已启用")
|
||
print("=" * 60)
|
||
|
||
uvicorn.run(fastapi_app, host=HOST, port=PORT, log_level="info")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|