chenchen/pingtai_agent
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

将search改为最新版

Browse Source
This commit is contained in:
zhanggangyong
2026-02-27 09:48:46 +00:00
parent dd4caf42ee
commit f27f32dfb5
74 changed files with 12871 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
code_agent/Dockerfile
+20
View File
@@ -0,0 +1,20 @@
FROM python:3.12-slim
WORKDIR /app
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
RUN apt-get update && apt-get install -y gcc curl && rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
EXPOSE 8000
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
CMD curl -f http://localhost:8000/health || exit 1
CMD ["python", "run_api_server.py"]
code_agent/README.md
+235
View File
@@ -0,0 +1,235 @@
# Code Agent
根据 specs_agent 生成的任务清单自动生成 Agent 项目代码。
## 功能特性
- **任务清单解析**：解析 specs_agent 生成的 Markdown 格式任务清单
- **代码自动生成**：根据任务清单和项目规范生成完整的 Agent 项目
- **模板 + AI 混合生成**：标准文件使用模板，业务逻辑使用 AI 生成
- **规范约束**：支持 steering_agent 提供的项目规范
## 工作流程
```
specs_agent 任务清单 + steering_agent 项目规范
↓
code_agent
↓
完整 Agent 项目代码
```
## 快速开始
### 环境变量
```bash
# 必需
export OPENAI_API_KEY=your-api-key
# 可选
export OPENAI_BASE_URL=https://litellm.taiji.io/v1
export MODEL_NAME=taiji/gpt-4o-mini
export API_KEY=your-api-key # API 访问密钥
export API_HOST=0.0.0.0
export API_PORT=8000
```
### 本地运行
```bash
# 安装依赖
pip install -r requirements.txt
# 启动服务
python run_api_server.py
```
### Docker 运行
```bash
# 构建镜像
docker build -t code-agent .
# 运行容器
docker run -p 8000:8000 \
-e OPENAI_API_KEY=your-key \
code-agent
```
## API 端点
| 端点 | 方法 | 描述 |
|------|------|------|
| `/` | GET | 服务状态 |
| `/health` | GET | 健康检查 |
| `/mcp` | POST | MCP JSON-RPC |
| `/mcp/sse` | GET/POST | MCP SSE 流式 |
| `/api/v1/parse` | POST | 解析任务清单 |
| `/api/v1/generate` | POST | 生成完整项目 |
| `/api/v1/tools` | GET | 获取工具列表 |
## MCP 工具
| 工具名称 | 功能 |
|---------|------|
| `parse_task_list` | 解析 specs_agent 生成的任务清单 |
| `generate_project` | 根据任务清单生成完整的 Agent 项目 |
## 使用示例
### 1. 解析任务清单
```bash
curl -X POST http://localhost:8000/api/v1/parse \
-H "Content-Type: application/json" \
-H "api-key: your-key" \
-d '{
"task_list_markdown": "# 任务清单: 数据去重 Agent\n\n## 任务列表\n\n### Task-001: 创建项目结构\n**类型**: 创建文件\n**文件**: `src/__init__.py`\n**依赖**: 无"
}'
```
### 2. 生成完整项目
```bash
curl -X POST http://localhost:8000/api/v1/generate \
-H "Content-Type: application/json" \
-H "api-key: your-key" \
-d '{
"task_list": "# 任务清单: 数据去重 Agent\n...",
"project_context": "{\"naming_conventions\": {\"files\": \"snake_case\"}}"
}'
```
### 3. 完整流程示例
```python
import requests
import json
# 1. 从 specs_agent 获取任务清单
task_list = """
# 任务清单: 数据去重 Agent
## 任务列表
### Task-001: 创建项目目录结构
**类型**: 创建文件
**文件**:
- `src/__init__.py`
- `src/server/__init__.py`
**依赖**: 无
### Task-002: 实现 MCP 工具
**类型**: 创建文件
**文件**: `src/server/mcp_server.py`
**依赖**: Task-001
**描述**: 实现数据去重的 MCP 工具
**代码要点**:
```python
@server.tool()
async def deduplicate(data: List[str]) -> str:
...
```
## 执行顺序
1. Task-001 → Task-002
"""
# 2. 从 steering_agent 获取项目规范（可选）
project_context = {
"naming_conventions": {
"files": "snake_case",
"classes": "PascalCase",
"functions": "snake_case"
},
"code_patterns": {
"async_required": True,
"error_handling": "try-except with JSON response"
}
}
# 3. 调用 code_agent 生成项目
response = requests.post(
"http://localhost:8000/api/v1/generate",
headers={"api-key": "your-key"},
json={
"task_list": task_list,
"project_context": json.dumps(project_context)
}
)
result = response.json()
if result["success"]:
print(f"✅ 生成成功！共 {len(result['files'])} 个文件")
for file in result["files"]:
print(f" - {file['path']} ({file['type']})")
else:
print(f"❌ 生成失败: {result['error']}")
```
## 生成结果格式
```json
{
"success": true,
"project_name": "dedup_agent",
"files": [
{
"path": "dedup_agent/src/__init__.py",
"content": "\"\"\"数据去重 Agent 源代码包\"\"\"",
"type": "template",
"task_id": "Task-001"
},
{
"path": "dedup_agent/src/server/mcp_server.py",
"content": "...(完整代码)...",
"type": "ai_generated",
"task_id": "Task-002"
}
],
"tasks_completed": 2,
"tasks_total": 2,
"error": null
}
```
## 与其他 Agent 的协作
```
用户 → specs_agent → 任务清单
↓
用户 → steering_agent → 项目规范
↓
code_agent
↓
完整项目代码
```
## 项目结构
```
code_agent/
├── Dockerfile
├── README.md
├── USAGE.md
├── requirements.txt
├── run_api_server.py
└── src/
├── __init__.py
└── server/
├── __init__.py
├── api_server.py # FastAPI + MCP HTTP
├── mcp_server.py # MCP 工具定义
├── parser.py # 任务清单解析器
├── generator.py # 代码生成器
└── templates/ # 代码模板
├── __init__.py
├── base_templates.py
└── prompts.py
```
## 许可证
MIT License
code_agent/USAGE.md
+341
View File
@@ -0,0 +1,341 @@
# Code Agent 使用指南
## 概述
Code Agent 是一个代码生成工具，用于根据 specs_agent 生成的任务清单自动生成完整的 Agent 项目代码。
## MCP 工具详解
### 1. parse_task_list
解析 specs_agent 生成的 Markdown 格式任务清单。
#### 输入参数
| 参数 | 类型 | 必需 | 描述 |
|------|------|------|------|
| `task_list_markdown` | string | 是 | specs_agent 生成的任务清单（Markdown 格式） |
#### 输出格式
```json
{
"success": true,
"project_name": "xxx_agent",
"total_tasks": 6,
"tasks": [
{
"id": "Task-001",
"title": "创建项目目录结构",
"type": "create_file",
"files": ["src/__init__.py", "src/server/__init__.py"],
"dependencies": [],
"description": "创建基础目录结构和初始化文件",
"acceptance_criteria": ["目录结构正确", "文件可导入"],
"code_hints": null
}
],
"execution_order": [["Task-001"], ["Task-002", "Task-003"], ["Task-004"]]
}
```
#### 使用示例
```bash
curl -X POST http://localhost:8000/api/v1/parse \
-H "Content-Type: application/json" \
-H "api-key: your-key" \
-d '{
"task_list_markdown": "# 任务清单: 数据去重 Agent\n\n## 任务列表\n\n### Task-001: 创建项目结构\n**类型**: 创建文件\n**文件**: `src/__init__.py`\n**依赖**: 无"
}'
```
---
### 2. generate_project
根据任务清单生成完整的 Agent 项目。
#### 输入参数
| 参数 | 类型 | 必需 | 描述 |
|------|------|------|------|
| `task_list` | string | 是 | specs_agent 生成的任务清单（Markdown 格式） |
| `project_context` | string | 否 | steering_agent 的项目规范（JSON 格式） |
#### 输出格式
成功时：
```json
{
"success": true,
"project_name": "xxx_agent",
"files": [
{
"path": "xxx_agent/src/__init__.py",
"content": "\"\"\"xxx Agent 源代码包\"\"\"",
"type": "template",
"task_id": "Task-001"
},
{
"path": "xxx_agent/src/server/mcp_server.py",
"content": "...(完整代码)...",
"type": "ai_generated",
"task_id": "Task-002"
}
],
"tasks_completed": 6,
"tasks_total": 6,
"error": null
}
```
失败时：
```json
{
"success": false,
"project_name": "xxx_agent",
"files": [...],
"tasks_completed": 3,
"tasks_total": 6,
"error": "Task-004 生成失败: ..."
}
```
#### 使用示例
```bash
curl -X POST http://localhost:8000/api/v1/generate \
-H "Content-Type: application/json" \
-H "api-key: your-key" \
-d '{
"task_list": "# 任务清单: 数据去重 Agent\n...",
"project_context": "{\"naming_conventions\": {\"files\": \"snake_case\"}}"
}'
```
---
## 任务清单格式
code_agent 期望的任务清单格式如下：
```markdown
# 任务清单: [项目名称]
## 任务总览
- 总任务数: X
- 预计文件数: X
## 任务列表
### Task-001: [任务标题]
**类型**: 创建文件 / 修改文件 / 配置
**文件**: `path/to/file.py`
**依赖**: 无 / Task-XXX
**描述**:
[详细描述要做什么]
**验收标准**:
- [ ] 标准1
- [ ] 标准2
**代码要点**:
```python
# 关键代码片段
```
---
### Task-002: [任务标题]
...
## 执行顺序
1. Task-001 → Task-002 → Task-003
2. Task-004 (可并行)
```
---
## 项目规范格式
steering_agent 提供的项目规范格式：
```json
{
"project_context": {
"name": "pingtai_agent",
"type": "agent_platform",
"tech_stack": {
"language": "Python 3.12",
"framework": ["FastAPI", "Pydantic AI", "FastMCP"],
"dependencies": ["pydantic-ai>=0.0.14", "fastapi>=0.109.0"]
}
},
"naming_conventions": {
"files": "snake_case",
"classes": "PascalCase",
"functions": "snake_case",
"constants": "UPPER_SNAKE_CASE"
},
"code_patterns": {
"async_required": true,
"error_handling": "try-except with JSON response",
"logging": "print statements with emoji prefix"
},
"rules": {
"must": ["所有 MCP 工具必须返回 JSON 格式"],
"must_not": ["禁止硬编码 API Key"],
"prefer": ["优先使用 Pydantic 模型"]
}
}
```
---
## 生成的文件类型
### 模板文件（template）
使用固定模板生成，仅替换变量：
| 文件 | 说明 |
|------|------|
| `src/__init__.py` | 源代码包初始化 |
| `src/server/__init__.py` | 服务器模块初始化 |
| `Dockerfile` | Docker 配置 |
| `requirements.txt` | Python 依赖 |
| `run_api_server.py` | 启动脚本 |
### AI 生成文件（ai_generated）
使用 AI 根据任务描述和项目规范生成：
| 文件 | 说明 |
|------|------|
| `src/server/mcp_server.py` | MCP 工具定义 |
| `src/server/api_server.py` | FastAPI 服务器 |
| `README.md` | 项目说明文档 |
| `USAGE.md` | 使用指南 |
---
## 完整使用流程
### 步骤 1: 使用 specs_agent 生成任务清单
```python
import requests
# 生成需求文档
req_resp = requests.post(
"http://specs-agent:8000/api/v1/requirements",
headers={"api-key": "your-key"},
json={"brief_description": "创建一个数据去重Agent"}
)
requirements_doc = req_resp.json()["document"]
# 生成设计文档
design_resp = requests.post(
"http://specs-agent:8000/api/v1/design",
headers={"api-key": "your-key"},
json={"requirements_doc": requirements_doc}
)
design_doc = design_resp.json()["document"]
# 生成任务清单
tasks_resp = requests.post(
"http://specs-agent:8000/api/v1/tasks",
headers={"api-key": "your-key"},
json={"design_doc": design_doc}
)
task_list = tasks_resp.json()["document"]
```
### 步骤 2: 使用 steering_agent 获取项目规范
```python
# 提取项目知识
extract_resp = requests.post(
"http://steering-agent:8000/api/v1/extract",
headers={"api-key": "your-key"},
json={"project_path": "/path/to/reference/project"}
)
# 获取项目上下文
context_resp = requests.get(
"http://steering-agent:8000/api/v1/context?format=json"
)
project_context = json.dumps(context_resp.json()["context"])
```
### 步骤 3: 使用 code_agent 生成项目
```python
# 生成完整项目
generate_resp = requests.post(
"http://code-agent:8000/api/v1/generate",
headers={"api-key": "your-key"},
json={
"task_list": task_list,
"project_context": project_context
}
)
result = generate_resp.json()
if result["success"]:
# 将生成的文件写入磁盘
import os
for file_info in result["files"]:
file_path = file_info["path"]
os.makedirs(os.path.dirname(file_path), exist_ok=True)
with open(file_path, "w") as f:
f.write(file_info["content"])
print(f"✅ 已创建: {file_path}")
else:
print(f"❌ 生成失败: {result['error']}")
```
---
## 常见问题
### Q: 如果不提供 project_context 会怎样？
A: code_agent 会使用默认的 agent_templates 规范，包括：
- 命名规范：snake_case 文件名、PascalCase 类名
- 代码模式：async 函数、try-except 错误处理
- 标准规则：MCP 工具返回 JSON、禁止硬编码密钥
### Q: 生成失败后如何处理？
A: 生成失败时会立即停止，返回已生成的文件和错误信息。你可以：
1. 检查错误信息，修复任务清单中的问题
2. 重新调用 generate_project 生成完整项目
### Q: 如何自定义生成的代码风格？
A: 通过 project_context 参数传入 steering_agent 的项目规范，包括：
- naming_conventions：命名规范
- code_patterns：代码模式
- rules：必须/禁止/偏好规则
### Q: 生成的文件如何保存？
A: code_agent 只返回文件内容，不写入磁盘。你需要自己处理文件保存：
```python
import os
for file_info in result["files"]:
file_path = file_info["path"]
os.makedirs(os.path.dirname(file_path), exist_ok=True)
with open(file_path, "w") as f:
f.write(file_info["content"])
```
### Q: 支持增量生成吗？
A: 不支持。每次调用 generate_project 都会生成完整项目。如果只需要重新生成某个文件，需要修改任务清单后重新生成整个项目。
code_agent/requirements.txt
+13
View File
@@ -0,0 +1,13 @@
# Pydantic AI
pydantic-ai>=0.0.14
# MCP
mcp>=0.9.0
fastmcp>=0.1.0
# FastAPI
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
# HTTP Client
aiohttp>=3.9.0
code_agent/run_api_server.py
+17
View File
@@ -0,0 +1,17 @@
#!/usr/bin/env python
"""启动 Code Agent API 服务器"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
if __name__ == '__main__':
from src.server.api_server import app
import uvicorn
import os
host = os.getenv('API_HOST', '0.0.0.0')
port = int(os.getenv('API_PORT', '8000'))
print(f"🚀 启动 Code Agent API: http://{host}:{port}")
uvicorn.run(app, host=host, port=port, log_level="info")
code_agent/src/__init__.py
+4
View File
@@ -0,0 +1,4 @@
"""Code Agent 源代码包
根据 specs_agent 生成的任务清单自动生成 Agent 项目代码。
"""
code_agent/src/server/__init__.py
+1
View File
@@ -0,0 +1 @@
"""服务器模块"""
code_agent/src/server/api_server.py
+292
View File
@@ -0,0 +1,292 @@
"""Code Agent API 服务器
提供 REST API 和 MCP HTTP/SSE 端点。
"""
import os
import json
import uuid
from typing import Optional, Dict, Any, AsyncGenerator
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Header, Request, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel
from .mcp_server import TOOL_MAP, TOOL_LIST
from .parser import parse_task_list
from .generator import generate_project
# ==================== 配置 ====================
SERVER_NAME = "Code Agent API"
# ==================== FastAPI 应用 ====================
@asynccontextmanager
async def lifespan(app: FastAPI):
"""应用生命周期管理"""
print(f"🚀 {SERVER_NAME} 启动")
yield
print(f"🛑 {SERVER_NAME} 关闭")
app = FastAPI(
title=SERVER_NAME,
description="根据任务清单生成 Agent 项目代码",
version="1.0.0",
lifespan=lifespan
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# ==================== API Key 验证 ====================
async def verify_api_key(
api_key: Optional[str] = Header(None, alias="api-key"),
authorization: Optional[str] = Header(None)
) -> str:
"""验证 API Key"""
if api_key and api_key.strip() and api_key.strip() != "sk":
return api_key.strip()
if authorization:
key = authorization[7:].strip() if authorization.startswith("Bearer ") else authorization.strip()
if key and key != "sk":
return key
raise HTTPException(status_code=401, detail="缺少 API Key")
def get_api_key_from_request(request: Request) -> Optional[str]:
"""从请求头提取 API Key（不验证）"""
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
# ==================== 请求模型 ====================
class ParseRequest(BaseModel):
"""解析任务清单请求"""
task_list_markdown: str
class GenerateRequest(BaseModel):
"""生成项目请求"""
task_list: str
project_context: Optional[str] = None
# ==================== 健康检查 ====================
@app.get("/")
async def root():
"""服务状态"""
return {
"service": SERVER_NAME,
"status": "running",
"version": "1.0.0",
"tools": list(TOOL_MAP.keys())
}
@app.get("/health")
async def health():
"""健康检查"""
return {"status": "healthy", "service": SERVER_NAME}
# ==================== MCP 端点 ====================
sessions: Dict[str, Dict] = {}
async def handle_mcp_request(data: Dict, session_id: str = None, api_key: str = None) -> Dict:
"""处理 MCP JSON-RPC 请求"""
method = data.get("method")
params = data.get("params", {})
req_id = data.get("id")
# tools/call 需要验证 API Key
if method == "tools/call" and (not api_key or api_key == "sk"):
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": TOOL_LIST}}
elif method == "tools/call":
tool_name = params.get("name")
args = params.get("arguments", {})
if tool_name not in TOOL_MAP:
raise ValueError(f"Unknown tool: {tool_name}")
# 设置 API Key 到环境变量
old_key = os.environ.get('OPENAI_API_KEY')
if api_key:
os.environ['OPENAI_API_KEY'] = api_key
try:
result = await TOOL_MAP[tool_name](**args)
finally:
if old_key:
os.environ['OPENAI_API_KEY'] = old_key
return {
"jsonrpc": "2.0", "id": req_id,
"result": {"content": [{"type": "text", "text": str(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"
import asyncio
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)}})
# ==================== REST API 端点 ====================
@app.post("/api/v1/parse")
async def api_parse(request: ParseRequest):
"""
解析任务清单。
将 specs_agent 生成的 Markdown 任务清单解析为结构化 JSON。
"""
try:
result = parse_task_list(request.task_list_markdown)
return result
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/v1/generate")
async def api_generate(
request: GenerateRequest,
api_key: str = Depends(verify_api_key)
):
"""
生成完整项目。
根据任务清单和项目规范生成完整的 Agent 项目代码。
这是 code_agent 的主要端点。
"""
try:
# 解析项目上下文
context = None
if request.project_context:
try:
context = json.loads(request.project_context)
except json.JSONDecodeError:
pass
# 设置 API Key 到环境变量
old_key = os.environ.get('OPENAI_API_KEY')
os.environ['OPENAI_API_KEY'] = api_key
try:
# 生成项目
result = await generate_project(request.task_list, context)
finally:
if old_key:
os.environ['OPENAI_API_KEY'] = old_key
if not result.get('success'):
return JSONResponse(
status_code=400,
content=result
)
return result
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/v1/tools")
async def api_tools():
"""获取可用工具列表"""
return {
"tools": {name: "MCP 工具" for name in TOOL_MAP.keys()},
"count": len(TOOL_MAP)
}
# 导出
__all__ = ['app']
code_agent/src/server/generator.py
+341
View File
@@ -0,0 +1,341 @@
"""代码生成器
根据任务清单和项目规范生成 Agent 项目代码。
"""
import os
import json
import re
from typing import List, Dict, Optional, Any
from pydantic_ai import Agent
from .parser import TaskListParser, Task
from .templates.base_templates import BaseTemplates
from .templates.prompts import PromptBuilder
# LiteLLM Gateway 配置
_BASE_URL = os.getenv('OPENAI_BASE_URL',
os.getenv('LLM_BASE_URL', 'https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1'))
_API_KEY = os.getenv('OPENAI_API_KEY', 'sk')
os.environ.setdefault('OPENAI_API_KEY', _API_KEY)
os.environ.setdefault('OPENAI_BASE_URL', _BASE_URL)
# 模型名称（pydantic_ai 需要 openai: 前缀）
def _get_model_name() -> str:
model = os.getenv('MODEL_NAME', os.getenv('LITELLM_MODEL', 'taiji/gpt-4o-mini'))
return model if ':' in model else f'openai:{model}'
MODEL_NAME = _get_model_name()
class CodeGenerator:
"""代码生成器"""
def __init__(self, project_context: dict = None):
"""
初始化代码生成器。
Args:
project_context: steering_agent 提供的项目上下文
"""
self.context = project_context or {}
self.parser = TaskListParser()
self.prompt_builder = PromptBuilder(project_context)
self.templates = BaseTemplates()
async def generate_project(self, task_list_markdown: str) -> dict:
"""
根据任务清单生成完整项目。
Args:
task_list_markdown: specs_agent 生成的任务清单
Returns:
生成结果，包含所有文件
"""
# 解析任务清单
parsed = self.parser.parse(task_list_markdown)
if not parsed['success']:
return {
"success": False,
"project_name": None,
"files": [],
"tasks_completed": 0,
"tasks_total": 0,
"error": parsed.get('error', '任务清单解析失败')
}
project_name = parsed['project_name']
tasks = parsed['tasks']
execution_order = parsed['execution_order']
# 提取项目描述
project_description = self._extract_project_description(task_list_markdown, project_name)
# 收集所有工具信息（用于生成 api_server 和文档）
tools_info = self._collect_tools_info(tasks)
# 按执行顺序生成文件
generated_files = []
tasks_completed = 0
try:
for task_group in execution_order:
for task_id in task_group:
# 找到对应的任务
task = self._find_task(tasks, task_id)
if not task:
continue
# 生成该任务的所有文件
files = await self._generate_task_files(
task=task,
project_name=project_name,
project_description=project_description,
tools_info=tools_info,
generated_files=generated_files
)
generated_files.extend(files)
tasks_completed += 1
return {
"success": True,
"project_name": project_name,
"files": generated_files,
"tasks_completed": tasks_completed,
"tasks_total": len(tasks),
"error": None
}
except Exception as e:
return {
"success": False,
"project_name": project_name,
"files": generated_files,
"tasks_completed": tasks_completed,
"tasks_total": len(tasks),
"error": f"生成失败: {str(e)}"
}
def _extract_project_description(self, markdown: str, project_name: str) -> str:
"""提取项目描述"""
# 尝试从任务清单中提取描述
match = re.search(r'##\s*项目背景\s*\n(.+?)(?=\n##|$)', markdown, re.DOTALL)
if match:
return match.group(1).strip()
# 使用项目名称生成描述
display_name = project_name.replace('_', ' ').title()
return f"{display_name} - 一个基于 Pydantic AI 的智能 Agent"
def _collect_tools_info(self, tasks: List[dict]) -> dict:
"""收集所有工具信息"""
tools = []
for task in tasks:
# 从任务描述和代码要点中提取工具信息
description = task.get('description', '')
code_hints = task.get('code_hints', '')
title = task.get('title', '')
# 检查是否是 MCP 工具相关任务
if 'mcp' in title.lower() or 'tool' in title.lower() or '工具' in title:
tools.append({
'task_id': task['id'],
'title': title,
'description': description,
'code_hints': code_hints
})
return {
'tools': tools,
'tools_description': self._format_tools_description(tools),
'tool_list': self._format_tool_list(tools)
}
def _format_tools_description(self, tools: List[dict]) -> str:
"""格式化工具描述"""
if not tools:
return "根据项目需求实现相应的 MCP 工具"
lines = []
for tool in tools:
lines.append(f"### {tool['title']}")
if tool['description']:
lines.append(tool['description'])
if tool['code_hints']:
lines.append(f"代码要点:\n```\n{tool['code_hints']}\n```")
lines.append("")
return "\n".join(lines)
def _format_tool_list(self, tools: List[dict]) -> str:
"""格式化工具列表"""
if not tools:
return "- 根据项目需求定义的 MCP 工具"
lines = []
for tool in tools:
lines.append(f"- {tool['title']}: {tool['description'][:100] if tool['description'] else '待实现'}")
return "\n".join(lines)
def _find_task(self, tasks: List[dict], task_id: str) -> Optional[dict]:
"""查找任务"""
for task in tasks:
if task['id'] == task_id:
return task
return None
async def _generate_task_files(
self,
task: dict,
project_name: str,
project_description: str,
tools_info: dict,
generated_files: List[dict]
) -> List[dict]:
"""生成任务对应的文件"""
files = []
for file_path in task.get('files', []):
# 添加项目名前缀
full_path = f"{project_name}/{file_path}"
# 检查是否为模板文件
if self.templates.is_template_file(file_path):
content = self.templates.get_template_content(
file_path,
project_name
)
files.append({
"path": full_path,
"content": content,
"type": "template",
"task_id": task['id']
})
else:
# 使用 AI 生成
content = await self._generate_file_with_ai(
file_path=file_path,
task=task,
project_name=project_name,
project_description=project_description,
tools_info=tools_info,
generated_files=generated_files
)
files.append({
"path": full_path,
"content": content,
"type": "ai_generated",
"task_id": task['id']
})
return files
async def _generate_file_with_ai(
self,
file_path: str,
task: dict,
project_name: str,
project_description: str,
tools_info: dict,
generated_files: List[dict]
) -> str:
"""使用 AI 生成文件内容"""
normalized = file_path.replace('\\', '/').lower()
# 根据文件类型选择提示词
if 'mcp_server' in normalized:
prompt = self.prompt_builder.build_mcp_server_prompt(
project_name=project_name,
project_description=project_description,
tools_description=tools_info['tools_description'],
code_hints=task.get('code_hints', '')
)
elif 'api_server' in normalized:
# 获取已生成的 mcp_server 内容作为参考
mcp_content = self._get_generated_file_content(generated_files, 'mcp_server')
tool_list = tools_info['tool_list']
if mcp_content:
tool_list += f"\n\n已生成的 mcp_server.py 参考:\n```python\n{mcp_content[:2000]}\n```"
prompt = self.prompt_builder.build_api_server_prompt(
project_name=project_name,
project_description=project_description,
tool_list=tool_list
)
elif 'readme' in normalized:
prompt = self.prompt_builder.build_readme_prompt(
project_name=project_name,
project_description=project_description,
tool_list=tools_info['tool_list']
)
elif 'usage' in normalized:
prompt = self.prompt_builder.build_usage_prompt(
project_name=project_name,
project_description=project_description,
tool_definitions=tools_info['tools_description']
)
else:
# 通用生成提示词
prompt = f"""请生成 {file_path} 文件。
项目名称：{project_name}
项目描述：{project_description}
任务描述：
{task.get('description', '')}
代码要点：
{task.get('code_hints', '')}
验收标准：
{chr(10).join('- ' + c for c in task.get('acceptance_criteria', []))}
输出完整的文件内容，不要有其他解释。
"""
# 调用 AI 生成
agent = Agent(
MODEL_NAME,
system_prompt="你是一个专业的 Python 开发者，专注于生成高质量的代码。只输出代码或文档内容，不要有其他解释。"
)
result = await agent.run(prompt)
content = result.output
# 提取代码块（如果有）
content = self._extract_code_content(content, file_path)
return content
def _get_generated_file_content(self, generated_files: List[dict], keyword: str) -> Optional[str]:
"""获取已生成文件的内容"""
for file in generated_files:
if keyword in file['path'].lower():
return file['content']
return None
def _extract_code_content(self, content: str, file_path: str) -> str:
"""提取代码内容"""
# 如果内容被代码块包裹，提取代码块内容
code_match = re.search(r'```(?:python|dockerfile|markdown|md)?\n(.+?)```', content, re.DOTALL)
if code_match:
return code_match.group(1).strip()
# 如果是 Markdown 文件，可能有多个代码块，保留原样
if file_path.lower().endswith('.md'):
return content.strip()
return content.strip()
# 便捷函数
async def generate_project(task_list: str, project_context: dict = None) -> dict:
"""生成项目"""
generator = CodeGenerator(project_context)
return await generator.generate_project(task_list)
code_agent/src/server/mcp_server.py
+165
View File
@@ -0,0 +1,165 @@
"""Code Agent MCP 服务器
提供代码生成相关的 MCP 工具。
"""
import os
import json
from typing import Optional
from mcp.server.fastmcp import FastMCP
from .parser import parse_task_list as _parse_task_list
from .generator import generate_project as _generate_project
# ==================== 配置 ====================
# LiteLLM Gateway 配置
_BASE_URL = os.getenv('OPENAI_BASE_URL',
os.getenv('LLM_BASE_URL', 'https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1'))
_API_KEY = os.getenv('OPENAI_API_KEY', 'sk')
os.environ.setdefault('OPENAI_API_KEY', _API_KEY)
os.environ.setdefault('OPENAI_BASE_URL', _BASE_URL)
# ==================== MCP 服务器 ====================
server = FastMCP('Code Agent')
# 系统提示词
SYSTEM_PROMPT = """你是 Code Agent，一个专业的代码生成助手。
你的核心能力：
1. 解析 specs_agent 生成的任务清单
2. 根据任务清单和项目规范生成完整的 Agent 项目代码
3. 使用模板生成标准文件，使用 AI 生成业务逻辑文件
工作原则：
- 严格按照任务清单的依赖顺序生成代码
- 遵循 steering_agent 提供的项目规范
- 生成的代码必须完整可运行
- 失败时立即停止并返回错误信息
"""
# ==================== MCP 工具定义 ====================
@server.tool()
async def parse_task_list(
task_list_markdown: str
) -> str:
"""
解析 specs_agent 生成的任务清单。
将 Markdown 格式的任务清单解析为结构化的 JSON，
便于后续的代码生成处理。
Args:
task_list_markdown: specs_agent 生成的任务清单（Markdown 格式）
Returns:
JSON 格式的解析结果，包含项目名称、任务列表和执行顺序
"""
try:
result = _parse_task_list(task_list_markdown)
return json.dumps(result, ensure_ascii=False, indent=2)
except Exception as e:
return json.dumps({
"success": False,
"error": f"解析失败: {str(e)}"
}, ensure_ascii=False)
@server.tool()
async def generate_project(
task_list: str,
project_context: Optional[str] = None
) -> str:
"""
根据任务清单生成完整的 Agent 项目。
这是 code_agent 的核心工具，接收 specs_agent 的任务清单
和 steering_agent 的项目规范，生成完整的项目代码。
Args:
task_list: specs_agent 生成的任务清单（Markdown 格式）
project_context: steering_agent 的项目规范（JSON 格式，可选）
如果不提供，将使用默认的 agent_templates 规范
Returns:
JSON 格式的生成结果，包含所有文件路径和内容
"""
try:
# 解析项目上下文
context = None
if project_context:
try:
context = json.loads(project_context)
except json.JSONDecodeError:
pass
# 生成项目
result = await _generate_project(task_list, context)
return json.dumps(result, ensure_ascii=False, indent=2)
except Exception as e:
return json.dumps({
"success": False,
"error": f"生成失败: {str(e)}",
"project_name": None,
"files": [],
"tasks_completed": 0,
"tasks_total": 0
}, ensure_ascii=False)
# ==================== 工具映射（供 API 使用）====================
TOOL_MAP = {
'parse_task_list': parse_task_list,
'generate_project': generate_project,
}
TOOL_LIST = [
{
"name": "parse_task_list",
"description": "解析 specs_agent 生成的任务清单",
"inputSchema": {
"type": "object",
"properties": {
"task_list_markdown": {
"type": "string",
"description": "specs_agent 生成的任务清单（Markdown 格式）"
}
},
"required": ["task_list_markdown"]
}
},
{
"name": "generate_project",
"description": "根据任务清单生成完整的 Agent 项目",
"inputSchema": {
"type": "object",
"properties": {
"task_list": {
"type": "string",
"description": "specs_agent 生成的任务清单（Markdown 格式）"
},
"project_context": {
"type": "string",
"description": "steering_agent 的项目规范（JSON 格式，可选）"
}
},
"required": ["task_list"]
}
}
]
# 导出
__all__ = ['server', 'SYSTEM_PROMPT', 'TOOL_MAP', 'TOOL_LIST']
if __name__ == '__main__':
server.run()
code_agent/src/server/parser.py
+329
View File
@@ -0,0 +1,329 @@
"""任务清单解析器
解析 specs_agent 生成的 Markdown 格式任务清单。
"""
import re
import json
from typing import List, Optional
from dataclasses import dataclass, asdict
@dataclass
class Task:
"""单个任务"""
id: str # Task-001
title: str # 任务标题
type: str # create_file / modify_file / config
files: List[str] # 涉及的文件列表
dependencies: List[str] # 依赖的任务 ID
description: str # 详细描述
acceptance_criteria: List[str] # 验收标准
code_hints: Optional[str] # 代码要点
@dataclass
class TaskList:
"""任务清单"""
project_name: str # 项目名称
total_tasks: int # 总任务数
tasks: List[Task] # 任务列表
execution_order: List[List[str]] # 执行顺序（支持并行）
class TaskListParser:
"""任务清单解析器"""
def parse(self, markdown: str) -> dict:
"""
解析 Markdown 格式的任务清单。
Args:
markdown: specs_agent 生成的任务清单
Returns:
解析后的任务清单（字典格式）
"""
try:
# 提取项目名称
project_name = self._extract_project_name(markdown)
# 提取任务列表
tasks = self._extract_tasks(markdown)
# 提取执行顺序
execution_order = self._extract_execution_order(markdown, tasks)
result = {
"success": True,
"project_name": project_name,
"total_tasks": len(tasks),
"tasks": [asdict(task) for task in tasks],
"execution_order": execution_order
}
return result
except Exception as e:
return {
"success": False,
"error": f"解析失败: {str(e)}",
"project_name": None,
"total_tasks": 0,
"tasks": [],
"execution_order": []
}
def _extract_project_name(self, markdown: str) -> str:
"""提取项目名称"""
# 匹配 "# 任务清单: xxx" 或 "# 任务清单: xxx Agent"
match = re.search(r'#\s*任务清单[:：]\s*(.+?)(?:\n|$)', markdown)
if match:
name = match.group(1).strip()
# 转换为 snake_case 格式
name = self._to_snake_case(name)
# 确保以 _agent 结尾
if not name.endswith('_agent'):
name = name + '_agent'
return name
# 尝试其他格式
match = re.search(r'项目名称[:：]\s*(.+?)(?:\n|$)', markdown)
if match:
return self._to_snake_case(match.group(1).strip())
return "new_agent"
def _to_snake_case(self, name: str) -> str:
"""转换为 snake_case"""
# 移除 "Agent" 后缀（稍后会添加 _agent）
name = re.sub(r'\s*[Aa]gent\s*$', '', name)
# 中文转拼音或保留
# 替换空格和特殊字符
name = re.sub(r'[\s\-]+', '_', name)
# 转小写
name = name.lower()
# 移除非字母数字下划线
name = re.sub(r'[^a-z0-9_\u4e00-\u9fff]', '', name)
return name
def _extract_tasks(self, markdown: str) -> List[Task]:
"""提取所有任务"""
tasks = []
# 匹配任务块: ### Task-XXX: 标题
task_pattern = r'###\s*(Task-\d+)[:：]\s*(.+?)(?=\n###\s*Task-|\n##\s+执行顺序|\n---\s*$|$)'
matches = re.findall(task_pattern, markdown, re.DOTALL)
for task_id, content in matches:
task = self._parse_task_block(task_id, content)
tasks.append(task)
return tasks
def _parse_task_block(self, task_id: str, content: str) -> Task:
"""解析单个任务块"""
lines = content.strip().split('\n')
# 提取标题（第一行）
title = lines[0].strip() if lines else ""
# 提取类型
task_type = self._extract_field(content, r'\*\*类型\*\*[:：]\s*(.+?)(?:\n|$)')
task_type = self._normalize_task_type(task_type)
# 提取文件列表
files = self._extract_files(content)
# 提取依赖
dependencies = self._extract_dependencies(content)
# 提取描述
description = self._extract_description(content)
# 提取验收标准
acceptance_criteria = self._extract_acceptance_criteria(content)
# 提取代码要点
code_hints = self._extract_code_hints(content)
return Task(
id=task_id,
title=title,
type=task_type,
files=files,
dependencies=dependencies,
description=description,
acceptance_criteria=acceptance_criteria,
code_hints=code_hints
)
def _extract_field(self, content: str, pattern: str) -> str:
"""提取字段值"""
match = re.search(pattern, content, re.IGNORECASE)
return match.group(1).strip() if match else ""
def _normalize_task_type(self, task_type: str) -> str:
"""标准化任务类型"""
task_type = task_type.lower()
if '创建' in task_type or 'create' in task_type:
return 'create_file'
elif '修改' in task_type or 'modify' in task_type:
return 'modify_file'
elif '配置' in task_type or 'config' in task_type:
return 'config'
return 'create_file'
def _extract_files(self, content: str) -> List[str]:
"""提取文件列表"""
files = []
# 匹配 **文件**: `xxx` 或 **文件**: \n- `xxx`
file_section = re.search(
r'\*\*文件\*\*[:：]\s*(.+?)(?=\n\*\*|\n###|\n##|$)',
content,
re.DOTALL
)
if file_section:
section = file_section.group(1)
# 提取反引号中的文件路径
file_matches = re.findall(r'`([^`]+)`', section)
files.extend(file_matches)
# 如果没有反引号，尝试提取列表项
if not files:
list_matches = re.findall(r'-\s*(.+?)(?:\n|$)', section)
files.extend([f.strip() for f in list_matches])
return files
def _extract_dependencies(self, content: str) -> List[str]:
"""提取依赖任务"""
deps = []
dep_match = re.search(r'\*\*依赖\*\*[:：]\s*(.+?)(?:\n|$)', content)
if dep_match:
dep_str = dep_match.group(1).strip()
if dep_str.lower() in ['无', 'none', '-', '']:
return []
# 提取 Task-XXX 格式
deps = re.findall(r'Task-\d+', dep_str)
return deps
def _extract_description(self, content: str) -> str:
"""提取描述"""
# 匹配 **描述**: 后的内容
desc_match = re.search(
r'\*\*描述\*\*[:：]\s*\n?(.+?)(?=\n\*\*验收|\n\*\*代码要点|\n###|\n##|$)',
content,
re.DOTALL
)
if desc_match:
return desc_match.group(1).strip()
return ""
def _extract_acceptance_criteria(self, content: str) -> List[str]:
"""提取验收标准"""
criteria = []
# 匹配验收标准部分
criteria_match = re.search(
r'\*\*验收标准\*\*[:：]?\s*\n?(.+?)(?=\n\*\*代码要点|\n###|\n##|$)',
content,
re.DOTALL
)
if criteria_match:
section = criteria_match.group(1)
# 提取列表项 - [ ] xxx 或 - xxx
items = re.findall(r'-\s*\[?\s*\]?\s*(.+?)(?:\n|$)', section)
criteria = [item.strip() for item in items if item.strip()]
return criteria
def _extract_code_hints(self, content: str) -> Optional[str]:
"""提取代码要点"""
# 匹配代码要点部分
hints_match = re.search(
r'\*\*代码要点\*\*[:：]?\s*\n?(.+?)(?=\n###|\n##|$)',
content,
re.DOTALL
)
if hints_match:
hints = hints_match.group(1).strip()
# 提取代码块
code_match = re.search(r'```[\w]*\n(.+?)```', hints, re.DOTALL)
if code_match:
return code_match.group(1).strip()
return hints if hints else None
return None
def _extract_execution_order(self, markdown: str, tasks: List[Task]) -> List[List[str]]:
"""提取执行顺序"""
order = []
# 匹配执行顺序部分
order_match = re.search(
r'##\s*执行顺序\s*\n(.+?)(?=\n##|$)',
markdown,
re.DOTALL
)
if order_match:
section = order_match.group(1)
# 解析每一行
lines = section.strip().split('\n')
for line in lines:
if not line.strip():
continue
# 提取该行的所有 Task-XXX
task_ids = re.findall(r'Task-\d+', line)
if task_ids:
order.append(task_ids)
# 如果没有找到执行顺序，根据依赖关系计算
if not order:
order = self._calculate_execution_order(tasks)
return order
def _calculate_execution_order(self, tasks: List[Task]) -> List[List[str]]:
"""根据依赖关系计算执行顺序"""
if not tasks:
return []
# 构建依赖图
task_map = {task.id: task for task in tasks}
remaining = set(task.id for task in tasks)
completed = set()
order = []
while remaining:
# 找出所有依赖已满足的任务
ready = []
for task_id in remaining:
task = task_map[task_id]
if all(dep in completed for dep in task.dependencies):
ready.append(task_id)
if not ready:
# 有循环依赖，将剩余任务全部加入
ready = list(remaining)
order.append(sorted(ready))
completed.update(ready)
remaining -= set(ready)
return order
# 便捷函数
def parse_task_list(markdown: str) -> dict:
"""解析任务清单"""
parser = TaskListParser()
return parser.parse(markdown)
code_agent/src/server/templates/__init__.py
+1
View File
@@ -0,0 +1 @@
"""代码模板模块"""
code_agent/src/server/templates/base_templates.py
+142
View File
@@ -0,0 +1,142 @@
"""固定代码模板
用于生成 Agent 项目的标准文件。
"""
# src/__init__.py 模板
SRC_INIT_TEMPLATE = '''"""{agent_name} 源代码包"""
'''
# src/server/__init__.py 模板
SERVER_INIT_TEMPLATE = '''"""服务器模块"""
'''
# Dockerfile 模板
DOCKERFILE_TEMPLATE = '''FROM python:3.12-slim
WORKDIR /app
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
RUN apt-get update && apt-get install -y gcc curl && rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
EXPOSE 8000
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \\
CMD curl -f http://localhost:8000/health || exit 1
CMD ["python", "run_api_server.py"]
'''
# requirements.txt 模板
REQUIREMENTS_TEMPLATE = '''# Pydantic AI
pydantic-ai>=0.0.14
# MCP
mcp>=0.9.0
fastmcp>=0.1.0
# FastAPI
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
# HTTP Client
aiohttp>=3.9.0
'''
# run_api_server.py 模板
RUN_API_SERVER_TEMPLATE = '''#!/usr/bin/env python
"""启动 {agent_display_name} API 服务器"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
if __name__ == '__main__':
from src.server.api_server import app
import uvicorn
import os
host = os.getenv('API_HOST', '0.0.0.0')
port = int(os.getenv('API_PORT', '8000'))
print(f"🚀 启动 {agent_display_name} API: http://{{host}}:{{port}}")
uvicorn.run(app, host=host, port=port, log_level="info")
'''
class BaseTemplates:
"""基础模板管理器"""
@staticmethod
def get_src_init(agent_name: str) -> str:
"""获取 src/__init__.py 内容"""
display_name = agent_name.replace('_', ' ').title()
return SRC_INIT_TEMPLATE.format(agent_name=display_name)
@staticmethod
def get_server_init() -> str:
"""获取 src/server/__init__.py 内容"""
return SERVER_INIT_TEMPLATE
@staticmethod
def get_dockerfile() -> str:
"""获取 Dockerfile 内容"""
return DOCKERFILE_TEMPLATE
@staticmethod
def get_requirements(extra_deps: list = None) -> str:
"""获取 requirements.txt 内容"""
content = REQUIREMENTS_TEMPLATE
if extra_deps:
content += "\n# Additional dependencies\n"
content += "\n".join(extra_deps) + "\n"
return content
@staticmethod
def get_run_api_server(agent_name: str) -> str:
"""获取 run_api_server.py 内容"""
display_name = agent_name.replace('_', ' ').title()
return RUN_API_SERVER_TEMPLATE.format(agent_display_name=display_name)
@staticmethod
def is_template_file(file_path: str) -> bool:
"""判断是否为模板文件"""
template_files = [
'src/__init__.py',
'src/server/__init__.py',
'Dockerfile',
'requirements.txt',
'run_api_server.py'
]
# 标准化路径
normalized = file_path.replace('\\', '/')
# 移除项目名前缀
for tf in template_files:
if normalized.endswith(tf):
return True
return False
@staticmethod
def get_template_content(file_path: str, agent_name: str, extra_deps: list = None) -> str:
"""根据文件路径获取模板内容"""
normalized = file_path.replace('\\', '/')
if normalized.endswith('src/__init__.py'):
return BaseTemplates.get_src_init(agent_name)
elif normalized.endswith('src/server/__init__.py'):
return BaseTemplates.get_server_init()
elif normalized.endswith('Dockerfile'):
return BaseTemplates.get_dockerfile()
elif normalized.endswith('requirements.txt'):
return BaseTemplates.get_requirements(extra_deps)
elif normalized.endswith('run_api_server.py'):
return BaseTemplates.get_run_api_server(agent_name)
return None
code_agent/src/server/templates/prompts.py
+275
View File
@@ -0,0 +1,275 @@
"""AI 生成提示词
用于生成 Agent 项目的业务逻辑文件。
"""
# mcp_server.py 生成提示词
MCP_SERVER_PROMPT = """你是一个专业的 Python 开发者，专注于 Agent 开发。
## 任务
请根据以下信息生成 mcp_server.py 文件。
## 项目信息
- 项目名称：{project_name}
- 项目描述：{project_description}
## 需要实现的 MCP 工具
{tools_description}
## 代码要点
{code_hints}
## 项目规范
### 技术栈
{tech_stack}
### 命名规范
- 文件名: {naming_files}
- 类名: {naming_classes}
- 函数名: {naming_functions}
- 常量: {naming_constants}
### 代码模式
- 异步要求: {async_required}
- 错误处理: {error_handling}
- 日志方式: {logging_style}
### 必须遵循的规则
{must_rules}
### 禁止事项
{must_not_rules}
## 输出要求
1. 使用 FastMCP 框架定义工具
2. 使用 Pydantic AI 调用 LLM（如需要）
3. 所有工具函数必须是 async
4. 所有工具必须返回 JSON 格式（使用 json.dumps）
5. 包含完整的 TOOL_MAP 和 TOOL_LIST 常量
6. 包含适当的 SYSTEM_PROMPT
7. 严格遵循上述命名规范和代码模式
输出完整的 Python 代码，不要省略任何部分。只输出代码，不要有其他解释。
"""
# api_server.py 生成提示词
API_SERVER_PROMPT = """你是一个专业的 Python 开发者，专注于 FastAPI 开发。
## 任务
请根据以下 MCP 工具定义生成 api_server.py 文件。
## 项目信息
- 项目名称：{project_name}
- 项目描述：{project_description}
## MCP 工具列表
{tool_list}
## 项目规范
### 命名规范
- 函数名: {naming_functions}
- 类名: {naming_classes}
### 代码模式
- 异步要求: {async_required}
- 错误处理: {error_handling}
### 必须遵循的规则
{must_rules}
## 输出要求
1. 使用 FastAPI 框架
2. 实现 MCP HTTP/SSE 端点（/mcp, /mcp/sse）
3. 为每个 MCP 工具实现对应的 REST API 端点
4. 实现 API Key 验证（从 api-key header 获取）
5. 实现健康检查端点（/health）
6. 实现根路径状态端点（/）
7. 严格遵循上述规范
输出完整的 Python 代码，不要省略任何部分。只输出代码，不要有其他解释。
"""
# README.md 生成提示词
README_PROMPT = """请为以下 Agent 项目生成 README.md 文件。
## 项目信息
- 项目名称：{project_name}
- 项目描述：{project_description}
## MCP 工具列表
{tool_list}
## 要求
1. 包含项目简介
2. 包含功能列表
3. 包含快速开始指南（Docker 和本地运行）
4. 包含 API 端点说明
5. 包含环境变量配置说明
6. 使用 Markdown 格式
输出完整的 README.md 内容。只输出 Markdown 内容，不要有其他解释。
"""
# USAGE.md 生成提示词
USAGE_PROMPT = """请为以下 Agent 项目生成 USAGE.md 使用文档。
## 项目信息
- 项目名称：{project_name}
- 项目描述：{project_description}
## MCP 工具详细定义
{tool_definitions}
## 要求
1. 为每个工具提供详细的使用说明
2. 包含输入参数说明
3. 包含输出格式说明
4. 包含使用示例（curl 命令）
5. 包含常见问题解答
输出完整的 USAGE.md 内容。只输出 Markdown 内容，不要有其他解释。
"""
class PromptBuilder:
"""提示词构建器"""
def __init__(self, project_context: dict = None):
"""
初始化提示词构建器。
Args:
project_context: steering_agent 提供的项目上下文
"""
self.context = project_context or {}
def _get_naming_conventions(self) -> dict:
"""获取命名规范"""
naming = self.context.get('naming_conventions', {})
return {
'files': naming.get('files', 'snake_case'),
'classes': naming.get('classes', 'PascalCase'),
'functions': naming.get('functions', 'snake_case'),
'constants': naming.get('constants', 'UPPER_SNAKE_CASE')
}
def _get_code_patterns(self) -> dict:
"""获取代码模式"""
patterns = self.context.get('code_patterns', {})
return {
'async_required': patterns.get('async_required', True),
'error_handling': patterns.get('error_handling', 'try-except with JSON response'),
'logging_style': patterns.get('logging', 'print statements with emoji prefix')
}
def _get_rules(self) -> dict:
"""获取规则"""
rules = self.context.get('rules', {})
must = rules.get('must', [
'所有 MCP 工具必须返回 JSON 格式',
'所有 API 端点必须有 docstring',
'必须实现 /health 健康检查端点'
])
must_not = rules.get('must_not', [
'禁止硬编码 API Key',
'禁止使用同步阻塞操作',
'禁止在工具函数中直接抛出异常'
])
return {
'must': '\n'.join(f'- {r}' for r in must),
'must_not': '\n'.join(f'- {r}' for r in must_not)
}
def _get_tech_stack(self) -> str:
"""获取技术栈信息"""
project = self.context.get('project_context', {})
tech = project.get('tech_stack', {})
language = tech.get('language', 'Python 3.12')
frameworks = tech.get('framework', ['FastAPI', 'Pydantic AI', 'FastMCP'])
if isinstance(frameworks, list):
frameworks_str = ', '.join(frameworks)
else:
frameworks_str = str(frameworks)
return f"- 语言: {language}\n- 框架: {frameworks_str}"
def build_mcp_server_prompt(
self,
project_name: str,
project_description: str,
tools_description: str,
code_hints: str = ""
) -> str:
"""构建 mcp_server.py 生成提示词"""
naming = self._get_naming_conventions()
patterns = self._get_code_patterns()
rules = self._get_rules()
return MCP_SERVER_PROMPT.format(
project_name=project_name,
project_description=project_description,
tools_description=tools_description,
code_hints=code_hints or "无特殊要求",
tech_stack=self._get_tech_stack(),
naming_files=naming['files'],
naming_classes=naming['classes'],
naming_functions=naming['functions'],
naming_constants=naming['constants'],
async_required=patterns['async_required'],
error_handling=patterns['error_handling'],
logging_style=patterns['logging_style'],
must_rules=rules['must'],
must_not_rules=rules['must_not']
)
def build_api_server_prompt(
self,
project_name: str,
project_description: str,
tool_list: str
) -> str:
"""构建 api_server.py 生成提示词"""
naming = self._get_naming_conventions()
patterns = self._get_code_patterns()
rules = self._get_rules()
return API_SERVER_PROMPT.format(
project_name=project_name,
project_description=project_description,
tool_list=tool_list,
naming_functions=naming['functions'],
naming_classes=naming['classes'],
async_required=patterns['async_required'],
error_handling=patterns['error_handling'],
must_rules=rules['must']
)
def build_readme_prompt(
self,
project_name: str,
project_description: str,
tool_list: str
) -> str:
"""构建 README.md 生成提示词"""
return README_PROMPT.format(
project_name=project_name,
project_description=project_description,
tool_list=tool_list
)
def build_usage_prompt(
self,
project_name: str,
project_description: str,
tool_definitions: str
) -> str:
"""构建 USAGE.md 生成提示词"""
return USAGE_PROMPT.format(
project_name=project_name,
project_description=project_description,
tool_definitions=tool_definitions
)
docs/code_agent_plan.md
+893
View File
@@ -0,0 +1,893 @@
# Code Agent 规划文档
> ✅ **状态**: 实施完成
> 📅 **完成时间**: 2026-02-24
## 1. 项目概述
### 1.1 定位
**Code Agent** 是一个代码生成Agent，用于：
- 根据 specs_agent 生成的任务清单自动生成代码
- 与 steering_agent 配合获取项目规范
- 生成完整的 Agent 项目文件
### 1.2 核心价值
| 价值 | 说明 |
|------|------|
| **自动化代码生成** | 根据任务清单自动生成符合规范的代码 |
| **规范一致性** | 从 steering_agent 获取规范，确保代码风格一致 |
| **完整项目输出** | 生成完整的 Agent 项目，包括所有必要文件 |
### 1.3 工作流程
```mermaid
flowchart TB
subgraph Input["输入"]
A[specs_agent 任务清单]
B[steering_agent 项目规范]
end
subgraph Process["Code Agent 处理"]
C[解析任务清单]
D[解析项目规范]
E[按依赖顺序生成代码]
end
subgraph Output["输出"]
G[完整 Agent 项目文件 - JSON格式]
end
A --> C
B --> D
C --> E
D --> E
E --> G
```
### 1.4 与其他 Agent 的协作
```mermaid
sequenceDiagram
participant User as 用户
participant Specs as specs_agent
participant Steering as steering_agent
participant Code as code_agent
User->>Specs: 1. 一句话需求
Specs-->>User: 需求文档
User->>Specs: 确认
Specs-->>User: 设计文档
User->>Specs: 确认
Specs-->>User: 任务清单
User->>Steering: 2. 提取项目知识
Steering-->>User: 项目上下文 JSON
User->>Code: 3. 生成项目
Note over Code: 输入: 任务清单 + 项目上下文
Code-->>User: 4. 完整项目文件 JSON
Note over User: 用户决定如何处理生成的文件
```
### 1.5 在 Agent 开发流程中的位置
```mermaid
flowchart LR
subgraph Phase1["阶段1: 规范"]
A[specs_agent] --> B[需求文档]
B --> C[设计文档]
C --> D[任务清单]
end
subgraph Phase2["阶段2: 约束"]
E[steering_agent] --> F[项目知识提取]
F --> G[规则定义]
G --> H[项目上下文]
end
subgraph Phase3["阶段3: 生成"]
D --> I[code_agent]
H --> I
I --> J[完整项目代码]
end
style I fill:#f96,stroke:#333
```
---
## 2. 功能设计
### 2.1 输入格式定义
#### 2.1.1 specs_agent 任务清单格式
code_agent 需要解析 specs_agent 生成的任务清单，格式如下：
```markdown
# 任务清单: [项目名称]
## 任务总览
- 总任务数: X
- 预计文件数: X
## 任务列表
### Task-001: [任务标题]
**类型**: 创建文件 / 修改文件 / 配置
**文件**: `path/to/file.py`
**依赖**: 无 / Task-XXX
**描述**:
[详细描述要做什么]
**验收标准**:
- [ ] 标准1
- [ ] 标准2
**代码要点**:
[关键代码片段或伪代码]
---
### Task-002: [任务标题]
...
## 执行顺序
1. Task-001 → Task-002 → Task-003
2. Task-004 (可并行)
```
#### 2.1.2 steering_agent 项目上下文格式
code_agent 使用 steering_agent 提供的 JSON 上下文：
```json
{
"project_context": {
"name": "pingtai_agent",
"type": "agent_platform",
"tech_stack": {
"language": "Python 3.12",
"framework": ["FastAPI", "Pydantic AI", "FastMCP"],
"dependencies": ["pydantic-ai>=0.0.14", "fastapi>=0.109.0"]
},
"structure": {
"pattern": "agent_templates",
"directories": ["src/", "src/server/"],
"key_files": ["mcp_server.py", "api_server.py", "run_api_server.py"]
}
},
"naming_conventions": {
"files": "snake_case",
"classes": "PascalCase",
"functions": "snake_case",
"constants": "UPPER_SNAKE_CASE"
},
"code_patterns": {
"async_required": true,
"error_handling": "try-except with JSON response",
"logging": "print statements with emoji prefix"
},
"rules": {
"must": ["所有 MCP 工具必须返回 JSON 格式", "..."],
"must_not": ["禁止硬编码 API Key", "..."],
"prefer": ["优先使用 Pydantic 模型定义请求/响应", "..."]
}
}
```
### 2.2 任务清单解析
解析后的内部数据结构：
```python
@dataclass
class Task:
id: str # Task-001
title: str # 任务标题
type: str # create_file / modify_file / config
files: List[str] # 涉及的文件列表
dependencies: List[str] # 依赖的任务 ID
description: str # 详细描述
acceptance_criteria: List[str] # 验收标准
code_hints: Optional[str] # 代码要点
@dataclass
class TaskList:
project_name: str # 项目名称
total_tasks: int # 总任务数
tasks: List[Task] # 任务列表
execution_order: List[List[str]] # 执行顺序（支持并行）
```
### 2.3 代码生成策略
#### 2.3.1 基于模板生成（固定文件）
对于标准文件，使用 agent_templates 作为基础：
| 文件 | 生成策略 | 变量替换 |
|------|---------|---------|
| `src/__init__.py` | 固定模板 | `{agent_name}` |
| `src/server/__init__.py` | 固定模板 | 无 |
| `Dockerfile` | 固定模板 | 无 |
| `requirements.txt` | 固定模板 | 可追加额外依赖 |
| `run_api_server.py` | 固定模板 | `{agent_name}` |
#### 2.3.2 基于 AI 生成（业务文件）
对于业务逻辑文件，使用 AI 生成：
| 文件 | 生成依据 | 上下文来源 |
|------|---------|-----------|
| `src/server/mcp_server.py` | 任务描述 + 代码要点 | steering_agent 规范 |
| `src/server/api_server.py` | MCP 工具定义 | steering_agent 规范 |
| `README.md` | 项目信息 + 工具列表 | 任务清单 |
| `USAGE.md` | 工具定义 + 使用示例 | 任务清单 |
### 2.4 规范约束应用
从 steering_agent 获取的规范用于 AI 生成提示词：
| 规范类型 | 应用方式 |
|---------|---------|
| **命名规范** | 在提示词中明确要求遵循 |
| **代码模式** | 提供示例代码片段 |
| **必须规则** | 作为硬性约束列出 |
| **禁止规则** | 作为禁止事项列出 |
| **偏好规则** | 作为建议列出 |
---
## 3. MCP 工具设计
### 3.1 工具清单
| 工具名称 | 功能 | 说明 |
|---------|------|------|
| `parse_task_list` | 解析任务清单 | 将 specs_agent 的 Markdown 任务清单解析为结构化 JSON |
| `generate_project` | 生成完整项目 | 根据任务清单和项目规范生成所有文件 |
> **设计原则**：
> - 与 specs_agent 和 steering_agent 保持相同的工具设计风格
> - 所有工具返回 JSON 格式（使用 `json.dumps`）
> - 所有工具函数使用 `async`
> - 失败时返回包含 `success: false` 和 `error` 字段的 JSON
### 3.2 工具详细设计
#### 3.2.1 parse_task_list
```python
@server.tool()
async def parse_task_list(
task_list_markdown: str
) -> str:
"""
解析 specs_agent 生成的任务清单。
将 Markdown 格式的任务清单解析为结构化的 JSON，
便于后续的代码生成处理。
Args:
task_list_markdown: specs_agent 生成的任务清单（Markdown 格式）
Returns:
JSON 格式的解析结果：
{
"success": true,
"project_name": "xxx_agent",
"total_tasks": 6,
"tasks": [
{
"id": "Task-001",
"title": "创建项目目录结构",
"type": "create_file",
"files": ["src/__init__.py", "src/server/__init__.py"],
"dependencies": [],
"description": "创建基础目录结构和初始化文件",
"acceptance_criteria": ["目录结构正确", "文件可导入"],
"code_hints": null
},
...
],
"execution_order": [["Task-001"], ["Task-002", "Task-003"], ["Task-004"]]
}
"""
```
#### 3.2.2 generate_project
```python
@server.tool()
async def generate_project(
task_list: str,
project_context: Optional[str] = None
) -> str:
"""
根据任务清单生成完整的 Agent 项目。
这是 code_agent 的核心工具，接收 specs_agent 的任务清单
和 steering_agent 的项目规范，生成完整的项目代码。
Args:
task_list: specs_agent 生成的任务清单（Markdown 格式）
project_context: steering_agent 的项目规范（JSON 格式，可选）
如果不提供，将使用默认的 agent_templates 规范
Returns:
JSON 格式的生成结果：
{
"success": true,
"project_name": "xxx_agent",
"files": [
{
"path": "xxx_agent/src/__init__.py",
"content": "\"\"\"xxx Agent 源代码包\"\"\"",
"type": "template",
"task_id": "Task-001"
},
{
"path": "xxx_agent/src/server/mcp_server.py",
"content": "...(完整代码)...",
"type": "ai_generated",
"task_id": "Task-002"
}
],
"tasks_completed": 6,
"tasks_total": 6,
"error": null
}
失败时返回：
{
"success": false,
"project_name": "xxx_agent",
"files": [...已生成的文件...],
"tasks_completed": 3,
"tasks_total": 6,
"error": "Task-004 生成失败: ..."
}
"""
```
### 3.3 内部处理流程
```mermaid
flowchart TD
A[接收任务清单] --> B[解析 Markdown]
B --> C{解析成功?}
C -->|否| D[返回解析错误]
C -->|是| E[解析项目规范]
E --> F[计算任务依赖顺序]
F --> G[开始生成循环]
G --> H{还有任务?}
H -->|否| I[返回成功结果]
H -->|是| J[获取下一个任务]
J --> K{是模板文件?}
K -->|是| L[使用模板生成]
K -->|否| M[使用 AI 生成]
L --> N{生成成功?}
M --> N
N -->|是| O[添加到结果]
N -->|否| P[返回失败结果]
O --> H
```
---
## 4. 代码生成模板
### 4.1 固定模板
#### src/__init__.py
```python
"""[Agent名称] 源代码包"""
```
#### src/server/__init__.py
```python
"""服务器模块"""
```
#### Dockerfile
```dockerfile
FROM python:3.12-slim
WORKDIR /app
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
RUN apt-get update && apt-get install -y gcc curl && rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
EXPOSE 8000
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
CMD curl -f http://localhost:8000/health || exit 1
CMD ["python", "run_api_server.py"]
```
#### requirements.txt
```
# Pydantic AI
pydantic-ai>=0.0.14
# MCP
mcp>=0.9.0
fastmcp>=0.1.0
# FastAPI
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
# HTTP Client
aiohttp>=3.9.0
```
#### run_api_server.py
```python
#!/usr/bin/env python
"""启动 [Agent名称] API 服务器"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
if __name__ == '__main__':
from src.server.api_server import app
import uvicorn
import os
host = os.getenv('API_HOST', '0.0.0.0')
port = int(os.getenv('API_PORT', '8000'))
print(f"🚀 启动 [Agent名称] API: http://{host}:{port}")
uvicorn.run(app, host=host, port=port, log_level="info")
```
### 4.2 AI 生成提示词
AI 生成提示词需要整合 steering_agent 提供的项目规范，确保生成的代码符合项目约束。
#### 4.2.1 mcp_server.py 生成提示词
```python
MCP_SERVER_PROMPT = """
你是一个专业的 Python 开发者，专注于 Agent 开发。
## 任务
请根据以下信息生成 mcp_server.py 文件。
## 项目信息
- 项目名称：{project_name}
- 项目描述：{project_description}
## 需要实现的 MCP 工具
{tools_description}
## 代码要点
{code_hints}
## 项目规范（来自 steering_agent）
### 技术栈
{tech_stack}
### 命名规范
- 文件名: {naming_files}
- 类名: {naming_classes}
- 函数名: {naming_functions}
- 常量: {naming_constants}
### 代码模式
- 异步要求: {async_required}
- 错误处理: {error_handling}
- 日志方式: {logging_style}
### 必须遵循的规则
{must_rules}
### 禁止事项
{must_not_rules}
## 输出要求
1. 使用 FastMCP 框架定义工具
2. 使用 Pydantic AI 调用 LLM（如需要）
3. 所有工具函数必须是 async
4. 所有工具必须返回 JSON 格式（使用 json.dumps）
5. 包含完整的 TOOL_MAP 和 TOOL_LIST 常量
6. 包含适当的 SYSTEM_PROMPT
7. 严格遵循上述命名规范和代码模式
输出完整的 Python 代码，不要省略任何部分。
"""
```
#### 4.2.2 api_server.py 生成提示词
```python
API_SERVER_PROMPT = """
你是一个专业的 Python 开发者，专注于 FastAPI 开发。
## 任务
请根据以下 MCP 工具定义生成 api_server.py 文件。
## 项目信息
- 项目名称：{project_name}
## MCP 工具列表
{tool_list}
## 项目规范（来自 steering_agent）
### 命名规范
- 函数名: {naming_functions}
- 类名: {naming_classes}
### 代码模式
- 异步要求: {async_required}
- 错误处理: {error_handling}
### 必须遵循的规则
{must_rules}
## 输出要求
1. 使用 FastAPI 框架
2. 实现 MCP HTTP/SSE 端点（/mcp, /mcp/sse）
3. 为每个 MCP 工具实现对应的 REST API 端点
4. 实现 API Key 验证（从 api-key header 获取）
5. 实现健康检查端点（/health）
6. 实现根路径状态端点（/）
7. 严格遵循上述规范
输出完整的 Python 代码，不要省略任何部分。
"""
```
#### 4.2.3 README.md 生成提示词
```python
README_PROMPT = """
请为以下 Agent 项目生成 README.md 文件。
## 项目信息
- 项目名称：{project_name}
- 项目描述：{project_description}
## MCP 工具列表
{tool_list}
## 要求
1. 包含项目简介
2. 包含功能列表
3. 包含快速开始指南
4. 包含 API 端点说明
5. 包含环境变量配置说明
6. 使用 Markdown 格式
输出完整的 README.md 内容。
"""
```
#### 4.2.4 USAGE.md 生成提示词
```python
USAGE_PROMPT = """
请为以下 Agent 项目生成 USAGE.md 使用文档。
## 项目信息
- 项目名称：{project_name}
## MCP 工具详细定义
{tool_definitions}
## 要求
1. 为每个工具提供详细的使用说明
2. 包含输入参数说明
3. 包含输出格式说明
4. 包含使用示例（curl 命令）
5. 包含常见问题解答
输出完整的 USAGE.md 内容。
"""
```
---
## 5. 项目结构
```
code_agent/
├── Dockerfile
├── README.md
├── USAGE.md
├── requirements.txt
├── run_api_server.py
└── src/
├── __init__.py
└── server/
├── __init__.py
├── api_server.py # FastAPI + MCP HTTP
├── mcp_server.py # MCP 工具定义
├── parser.py # 任务清单解析器
├── generator.py # 代码生成器
└── templates/ # 代码模板
├── __init__.py
├── base_templates.py # 固定模板
└── prompts.py # AI 生成提示词
```
---
## 6. API 端点设计
| 端点 | 方法 | 描述 |
|------|------|------|
| `/` | GET | 服务状态 |
| `/health` | GET | 健康检查 |
| `/mcp` | POST | MCP JSON-RPC |
| `/mcp/sse` | GET/POST | MCP SSE 流式 |
| `/api/v1/parse` | POST | 解析任务清单 |
| `/api/v1/generate` | POST | 生成完整项目（主要端点） |
> **注意**：简化为两个核心业务端点，与 MCP 工具一一对应。
---
## 7. 使用示例
### 7.1 完整流程示例
```python
import requests
import json
import os
# ============================================
# 阶段 1: 使用 specs_agent 生成规范文档
# ============================================
# 1.1 生成需求文档
specs_resp = requests.post(
"http://specs-agent:8000/api/v1/requirements",
headers={"api-key": "your-key"},
json={"brief_description": "创建一个数据去重Agent"}
)
requirements_doc = specs_resp.json()["document"]
print("✅ 需求文档已生成，请确认...")
# 1.2 用户确认后，生成设计文档
design_resp = requests.post(
"http://specs-agent:8000/api/v1/design",
headers={"api-key": "your-key"},
json={"requirements_doc": requirements_doc}
)
design_doc = design_resp.json()["document"]
print("✅ 设计文档已生成，请确认...")
# 1.3 用户确认后，生成任务清单
tasks_resp = requests.post(
"http://specs-agent:8000/api/v1/tasks",
headers={"api-key": "your-key"},
json={"design_doc": design_doc}
)
task_list = tasks_resp.json()["document"]
print("✅ 任务清单已生成")
# ============================================
# 阶段 2: 使用 steering_agent 获取项目规范
# ============================================
# 2.1 提取项目知识（可选，如果有参考项目）
extract_resp = requests.post(
"http://steering-agent:8000/api/v1/extract",
headers={"api-key": "your-key"},
json={"project_path": "/path/to/reference/project"}
)
# 2.2 获取项目上下文
context_resp = requests.get(
"http://steering-agent:8000/api/v1/context?format=json"
)
project_context = json.dumps(context_resp.json()["context"])
print("✅ 项目规范已获取")
# ============================================
# 阶段 3: 使用 code_agent 生成代码
# ============================================
# 3.1 生成完整项目
generate_resp = requests.post(
"http://code-agent:8000/api/v1/generate",
headers={"api-key": "your-key"},
json={
"task_list": task_list,
"project_context": project_context
}
)
result = generate_resp.json()
if result["success"]:
print(f"✅ 生成成功！共 {len(result['files'])} 个文件")
# 3.2 将生成的文件写入磁盘（由调用方决定）
output_dir = "./output/dedup_agent"
for file_info in result["files"]:
file_path = os.path.join(output_dir, file_info["path"])
os.makedirs(os.path.dirname(file_path), exist_ok=True)
with open(file_path, "w") as f:
f.write(file_info["content"])
print(f" 📄 {file_info['path']} ({file_info['type']})")
else:
print(f"❌ 生成失败: {result['error']}")
print(f" 已完成 {result['tasks_completed']}/{result['tasks_total']} 个任务")
```
### 7.2 仅解析任务清单
```python
# 如果只需要解析任务清单，不生成代码
parse_resp = requests.post(
"http://code-agent:8000/api/v1/parse",
headers={"api-key": "your-key"},
json={"task_list_markdown": task_list}
)
parsed = parse_resp.json()
print(f"项目名称: {parsed['project_name']}")
print(f"任务数量: {parsed['total_tasks']}")
for task in parsed["tasks"]:
print(f" - {task['id']}: {task['title']}")
```
### 7.3 生成结果示例
```json
{
"success": true,
"project_name": "dedup_agent",
"files": [
{
"path": "dedup_agent/src/__init__.py",
"content": "\"\"\"数据去重 Agent 源代码包\"\"\"",
"type": "template",
"task_id": "Task-001"
},
{
"path": "dedup_agent/src/server/__init__.py",
"content": "\"\"\"服务器模块\"\"\"",
"type": "template",
"task_id": "Task-001"
},
{
"path": "dedup_agent/src/server/mcp_server.py",
"content": "...(AI 生成的完整代码)...",
"type": "ai_generated",
"task_id": "Task-002"
},
{
"path": "dedup_agent/src/server/api_server.py",
"content": "...(AI 生成的完整代码)...",
"type": "ai_generated",
"task_id": "Task-003"
},
{
"path": "dedup_agent/Dockerfile",
"content": "...",
"type": "template",
"task_id": "Task-004"
},
{
"path": "dedup_agent/requirements.txt",
"content": "...",
"type": "template",
"task_id": "Task-004"
},
{
"path": "dedup_agent/run_api_server.py",
"content": "...",
"type": "template",
"task_id": "Task-004"
},
{
"path": "dedup_agent/README.md",
"content": "...(AI 生成)...",
"type": "ai_generated",
"task_id": "Task-005"
},
{
"path": "dedup_agent/USAGE.md",
"content": "...(AI 生成)...",
"type": "ai_generated",
"task_id": "Task-006"
}
],
"tasks_completed": 6,
"tasks_total": 6,
"error": null
}
```
---
## 8. 实施计划
### 8.1 开发任务清单
- [x] 创建项目目录结构
- [x] 实现任务清单解析器 (`parser.py`)
- [x] 实现代码模板 (`templates/`)
- [x] 固定模板 (`base_templates.py`)
- [x] AI 生成提示词 (`prompts.py`)
- [x] 实现代码生成器 (`generator.py`)
- [x] 实现 MCP 工具 (`mcp_server.py`)
- [x] `parse_task_list`
- [x] `generate_project`
- [x] 实现 API 端点 (`api_server.py`)
- [x] 编写文档 (`README.md`, `USAGE.md`)
- [ ] 测试验证
---
## 9. 设计决策（已确认）
| 决策项 | 决定 | 说明 |
|--------|------|------|
| **输出方式** | 仅返回内容 | 不写入磁盘，由调用方决定如何处理 |
| **错误处理** | 失败时停止 | 任何任务失败立即停止，返回错误信息 |
| **增量生成** | 不支持 | 每次生成完整项目，不支持部分文件生成 |
| **代码验证** | 不需要 | 不进行语法检查，由调用方负责验证 |
---
## 10. 风险与注意事项
1. **AI 生成质量**：AI 生成的代码可能不完美，需要人工审核
2. **依赖顺序**：任务之间有依赖关系，需要正确处理执行顺序
3. **上下文长度**：完整的任务清单和项目规范可能很长，需要注意 token 限制
4. **与 steering_agent 的集成**：需要正确解析 steering_agent 返回的 JSON 格式
---
## 11. 已创建文件清单
```
code_agent/
├── Dockerfile ✅
├── README.md ✅
├── USAGE.md ✅
├── requirements.txt ✅
├── run_api_server.py ✅
└── src/
├── __init__.py ✅
└── server/
├── __init__.py ✅
├── api_server.py ✅
├── mcp_server.py ✅
├── parser.py ✅
├── generator.py ✅
└── templates/
├── __init__.py ✅
├── base_templates.py ✅
└── prompts.py ✅
```
docs/steering_agent_plan.md
+1025
View File
File diff suppressed because it is too large Load Diff
plans/search_agent_migration_plan.md
+551
View File
@@ -0,0 +1,551 @@
# Search Agent 迁移计划
> 将 `search_agent/` 项目按照 `agent_templates/` 模板结构进行重构
## 架构变更图
### 迁移前后目录结构对比
```mermaid
graph LR
subgraph 迁移前
A1[search_agent/]
A1 --> A2[api.py]
A1 --> A3[config.py]
A1 --> A4[main.py]
A1 --> A5[agent/]
A1 --> A6[models/]
A1 --> A7[modules/]
A1 --> A8[tools/]
A1 --> A9[utils/]
end
subgraph 迁移后
B1[search_agent/]
B1 --> B2[Dockerfile]
B1 --> B3[run_api_server.py]
B1 --> B4[src/]
B4 --> B5[server/]
B5 --> B6[api_server.py]
B5 --> B7[mcp_server.py]
B5 --> B8[core/]
B8 --> B9[config.py]
B8 --> B10[agent.py]
B8 --> B11[modules/]
B8 --> B12[tools/]
B8 --> B13[utils/]
end
```
### API 端点变更
```mermaid
graph TB
subgraph 原 API
O1[GET /]
O2[GET /health]
O3[POST /search]
O4[GET /config]
end
subgraph 新 API
N1[GET /]
N2[GET /health]
N3[POST /mcp]
N4[GET /mcp/sse]
N5[POST /mcp/sse]
N6[POST /api/v1/search]
N7[POST /api/v1/quick_search]
end
O1 -.-> N1
O2 -.-> N2
O3 -.-> N6
O4 -.删除.-> X[X]
```
### MCP 工具调用流程
```mermaid
sequenceDiagram
participant Client
participant API as api_server.py
participant MCP as mcp_server.py
participant Agent as SearchAgent
participant Tools as External APIs
Client->>API: POST /mcp tools/call search
API->>MCP: TOOL_MAP search
MCP->>Agent: agent.search query
Agent->>Tools: Serper/Jina APIs
Tools-->>Agent: 搜索结果
Agent-->>MCP: AgentResponse
MCP-->>API: JSON result
API-->>Client: MCP Response
```
---
## 概述
将 `search_agent/` 项目按照 `agent_templates/` 模板结构进行重构，使其符合统一的 Agent 项目规范。
## 当前结构 vs 目标结构
### 当前 search_agent 结构
```
search_agent/
├── .env # 环境变量配置
├── api.py # FastAPI 服务（独立文件）
├── config.py # 配置管理
├── main.py # CLI 入口
├── requirements.txt # 依赖
├── agent/
│ ├── search_agent.py # 主 Agent 类
│ └── prompts.py # Prompt 模板
├── models/
│ └── schemas.py # 数据模型
├── modules/ # 7 个功能模块
│ ├── query_analyzer.py
│ ├── search_planner.py
│ ├── search_executor.py
│ ├── content_extractor.py
│ ├── result_processor.py
│ ├── answer_generator.py
│ └── reflector.py
├── tools/ # 外部 API 封装
│ ├── serper.py
│ ├── jina_reader.py
│ └── jina_reranker.py
└── utils/
├── llm_client.py
└── helpers.py
```
### 目标结构（按模板）
```
search_agent/
├── Dockerfile # 新增：Docker 配置
├── README.md # 更新：按模板格式
├── requirements.txt # 更新：添加 pydantic-ai, mcp, fastmcp
├── run_api_server.py # 新增：统一启动脚本
└── src/
├── __init__.py # 新增
└── server/
├── __init__.py # 新增
├── api_server.py # 重构：按模板格式
└── mcp_server.py # 新增：MCP 工具定义
└── core/ # 新增：业务逻辑目录
├── __init__.py
├── config.py # 移动：配置管理
├── agent.py # 移动：SearchAgent 类
├── prompts.py # 移动：Prompt 模板
├── schemas.py # 移动：数据模型
├── modules/ # 移动：功能模块
│ ├── __init__.py
│ ├── query_analyzer.py
│ ├── search_planner.py
│ ├── search_executor.py
│ ├── content_extractor.py
│ ├── result_processor.py
│ ├── answer_generator.py
│ └── reflector.py
├── tools/ # 移动：外部 API 封装
│ ├── __init__.py
│ ├── serper.py
│ ├── jina_reader.py
│ └── jina_reranker.py
└── utils/ # 移动：工具函数
├── __init__.py
├── llm_client.py
└── helpers.py
```
## 主要变更点
### 1. 项目结构变更
| 变更类型 | 原路径 | 新路径 |
|---------|--------|--------|
| 新增 | - | `Dockerfile` |
| 新增 | - | `run_api_server.py` |
| 新增 | - | `src/__init__.py` |
| 新增 | - | `src/server/__init__.py` |
| 重构 | `api.py` | `src/server/api_server.py` |
| 新增 | - | `src/server/mcp_server.py` |
| 移动 | `config.py` | `src/server/core/config.py` |
| 移动 | `agent/search_agent.py` | `src/server/core/agent.py` |
| 移动 | `agent/prompts.py` | `src/server/core/prompts.py` |
| 移动 | `models/schemas.py` | `src/server/core/schemas.py` |
| 移动 | `modules/*` | `src/server/core/modules/*` |
| 移动 | `tools/*` | `src/server/core/tools/*` |
| 移动 | `utils/*` | `src/server/core/utils/*` |
| 删除 | `main.py` | - （功能合并到 api_server） |
| 删除 | `.env` | - （环境变量由外部提供） |
| 删除 | 多个 README 文件 | 合并为一个 `README.md` |
### 2. 依赖变更
**新增依赖**（参考 `agent_templates/requirements.txt`）：
```
pydantic-ai>=0.0.14
mcp>=0.9.0
fastmcp>=0.1.0
```
**保留依赖**：
```
aiohttp>=3.9.0
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
pydantic>=2.5.0
python-dotenv>=1.0.0
loguru>=0.7.0
```
**可移除依赖**：
```
requests>=2.31.0 # 改用 aiohttp
orjson>=3.9.0 # 可选
typing-extensions>=4.9.0 # Python 3.12 内置
asyncio-throttle>=1.0.2 # 未使用
```
### 3. 配置变更
**原配置方式**（`config.py`）：
- 使用 `dataclass` 定义 `Config` 类
- 从 `.env` 文件加载配置
- 包含 LLM、Serper、Jina 等多个 API 配置
**新配置方式**（按模板）：
- 使用环境变量直接读取
- 主要配置项：
- `OPENAI_BASE_URL` / `LLM_BASE_URL`：LiteLLM Gateway URL
- `OPENAI_API_KEY`：API Key（运行时传入）
- `MODEL_NAME` / `LITELLM_MODEL`：模型名称
- `SERPER_API_KEY`：Serper API Key
- `JINA_API_KEY`：Jina API Key
- `API_PORT`：服务端口
### 4. API 变更
**原 API 端点**：
- `GET /` - 健康检查
- `GET /health` - 详细健康检查
- `POST /search` - 执行搜索
- `GET /config` - 获取配置
**新 API 端点**（按模板）：
- `GET /` - 服务状态
- `GET /health` - 健康检查
- `POST /mcp` - MCP HTTP 端点
- `GET /mcp/sse` - MCP SSE 端点
- `POST /mcp/sse` - MCP SSE POST 端点
- `POST /api/v1/search` - 业务 API（搜索）
- `POST /api/v1/quick_search` - 业务 API（快速搜索）
### 5. MCP 工具定义
新增 `mcp_server.py`，定义以下 MCP 工具：
```python
@server.tool()
async def search(query: str, max_iterations: int = 3) -> str:
"""
执行智能搜索
Args:
query: 搜索查询问题
max_iterations: 最大迭代次数（1-10）
Returns:
JSON 格式的搜索结果
"""
@server.tool()
async def quick_search(query: str) -> str:
"""
快速搜索（单次迭代）
Args:
query: 搜索查询问题
Returns:
JSON 格式的搜索结果
"""
```
---
## 迁移任务清单
### Phase 1: 创建基础结构
- [ ] 1.1 创建 `src/` 目录结构
- [ ] 1.2 创建 `src/__init__.py`
- [ ] 1.3 创建 `src/server/__init__.py`
- [ ] 1.4 创建 `src/server/core/` 目录结构
### Phase 2: 迁移核心业务代码
- [ ] 2.1 迁移 `config.py` → `src/server/core/config.py`
- 更新导入路径
- 适配新的环境变量命名
- [ ] 2.2 迁移 `models/schemas.py` → `src/server/core/schemas.py`
- 更新导入路径
- [ ] 2.3 迁移 `agent/prompts.py` → `src/server/core/prompts.py`
- 更新导入路径
- [ ] 2.4 迁移 `utils/` → `src/server/core/utils/`
- 更新 `llm_client.py` 导入路径
- 更新 `helpers.py` 导入路径
- [ ] 2.5 迁移 `tools/` → `src/server/core/tools/`
- 更新 `serper.py` 导入路径
- 更新 `jina_reader.py` 导入路径
- 更新 `jina_reranker.py` 导入路径
- [ ] 2.6 迁移 `modules/` → `src/server/core/modules/`
- 更新所有模块的导入路径
- [ ] 2.7 迁移 `agent/search_agent.py` → `src/server/core/agent.py`
- 更新导入路径
### Phase 3: 创建新的服务器文件
- [ ] 3.1 创建 `src/server/mcp_server.py`
- 定义 `search` 工具
- 定义 `quick_search` 工具
- 导出 `TOOL_MAP` 和 `TOOL_LIST`
- [ ] 3.2 创建 `src/server/api_server.py`
- 按模板格式重构
- 添加 MCP 端点
- 添加业务 API 端点
- 移除原有的 `LogCollector` 和 `AgentWrapper`（简化）
### Phase 4: 创建项目配置文件
- [ ] 4.1 创建 `Dockerfile`
- 基于 `agent_templates/Dockerfile`
- [ ] 4.2 创建 `run_api_server.py`
- 基于 `agent_templates/run_api_server.py`
- [ ] 4.3 更新 `requirements.txt`
- 添加 pydantic-ai, mcp, fastmcp
- 移除不需要的依赖
- [ ] 4.4 更新 `README.md`
- 按模板格式重写
### Phase 5: 清理和测试
- [ ] 5.1 删除旧文件
- `api.py`
- `main.py`
- `config.py`（根目录）
- `agent/` 目录
- `models/` 目录
- `modules/` 目录
- `tools/` 目录
- `utils/` 目录
- 多余的 README 文件
- [ ] 5.2 测试服务启动
- `python run_api_server.py`
- [ ] 5.3 测试 API 端点
- 健康检查
- MCP 端点
- 搜索 API
---
## 关键代码变更示例
### mcp_server.py 核心代码
```python
"""Search Agent MCP 服务器"""
import os
import json
from typing import Optional
from mcp.server.fastmcp import FastMCP
from .core.config import Config
from .core.agent import SearchAgent
# 配置
_BASE_URL = os.getenv('OPENAI_BASE_URL',
os.getenv('LLM_BASE_URL', 'https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1'))
_API_KEY = os.getenv('OPENAI_API_KEY', 'sk')
os.environ.setdefault('OPENAI_API_KEY', _API_KEY)
os.environ.setdefault('OPENAI_BASE_URL', _BASE_URL)
server = FastMCP('Search Agent')
SYSTEM_PROMPT = '''你是一个智能搜索助手。
能够理解用户查询意图、自动规划搜索策略、从多个来源获取信息，
并生成高质量、有来源引用的答案。'''
# 全局 Agent 实例
_agent: Optional[SearchAgent] = None
def get_agent() -> SearchAgent:
"""获取或创建 Agent 实例"""
global _agent
if _agent is None:
config = Config.from_env()
_agent = SearchAgent(config)
return _agent
@server.tool()
async def search(
query: str,
max_iterations: int = 3
) -> str:
"""
执行智能搜索
Args:
query: 搜索查询问题
max_iterations: 最大迭代次数（1-10）
Returns:
JSON 格式的搜索结果
"""
try:
agent = get_agent()
# 临时修改迭代次数
original = agent.config.max_iterations
agent.config.max_iterations = min(max(1, max_iterations), 10)
try:
response = await agent.search(query)
finally:
agent.config.max_iterations = original
return json.dumps({
"success": True,
"answer": response.answer.to_dict(),
"statistics": {
"iterations": response.iterations,
"total_sources_consulted": response.total_sources_consulted,
"search_queries_used": response.search_queries_used
}
}, ensure_ascii=False, indent=2)
except Exception as e:
return json.dumps({
"success": False,
"error": str(e)
}, ensure_ascii=False)
@server.tool()
async def quick_search(query: str) -> str:
"""
快速搜索（单次迭代）
Args:
query: 搜索查询问题
Returns:
JSON 格式的搜索结果
"""
try:
agent = get_agent()
answer = await agent.quick_search(query)
return json.dumps({
"success": True,
"answer": answer.to_dict()
}, ensure_ascii=False, indent=2)
except Exception as e:
return json.dumps({
"success": False,
"error": str(e)
}, ensure_ascii=False)
TOOL_MAP = {
'search': search,
'quick_search': quick_search,
}
TOOL_LIST = [
{
"name": "search",
"description": "执行智能搜索，支持多轮迭代和自我反思",
"inputSchema": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "搜索查询问题"},
"max_iterations": {"type": "integer", "description": "最大迭代次数", "default": 3}
},
"required": ["query"]
}
},
{
"name": "quick_search",
"description": "快速搜索，单次迭代，适合简单问题",
"inputSchema": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "搜索查询问题"}
},
"required": ["query"]
}
}
]
```
### 导入路径变更对照表
| 原导入 | 新导入 |
|--------|--------|
| `from config import Config` | `from .core.config import Config` |
| `from models.schemas import ...` | `from .core.schemas import ...` |
| `from agent.search_agent import SearchAgent` | `from .core.agent import SearchAgent` |
| `from agent.prompts import ...` | `from .core.prompts import ...` |
| `from modules.xxx import ...` | `from .core.modules.xxx import ...` |
| `from tools.xxx import ...` | `from .core.tools.xxx import ...` |
| `from utils.xxx import ...` | `from .core.utils.xxx import ...` |
---
## 环境变量配置
迁移后需要设置的环境变量：
| 变量名 | 必需 | 说明 | 默认值 |
|--------|------|------|--------|
| `OPENAI_BASE_URL` 或 `LLM_BASE_URL` | 是 | LiteLLM Gateway URL | - |
| `OPENAI_API_KEY` | 是 | API Key（运行时传入） | sk |
| `MODEL_NAME` 或 `LITELLM_MODEL` | 否 | 模型名称 | taiji/gpt-4o-mini |
| `SERPER_API_KEY` | 是 | Serper API Key | - |
| `JINA_API_KEY` | 是 | Jina API Key | - |
| `API_PORT` | 否 | 服务端口 | 8000 |
| `MAX_ITERATIONS` | 否 | 默认最大迭代次数 | 3 |
| `MAX_RESULTS_PER_QUERY` | 否 | 每次搜索结果数 | 10 |
| `CONTENT_MAX_LENGTH` | 否 | 内容最大长度 | 5000 |
| `LOG_LEVEL` | 否 | 日志级别 | INFO |
| `TIMEOUT` | 否 | 请求超时时间 | 30 |
---
## 注意事项
1. **保持业务逻辑不变**：迁移过程中不修改核心搜索逻辑，只调整项目结构和导入路径
2. **API 兼容性**：新增 MCP 端点的同时，保留原有的 `/search` 端点（改为 `/api/v1/search`）
3. **配置兼容**：支持原有的环境变量命名，同时支持模板的命名方式
4. **测试覆盖**：迁移后需要测试所有功能模块是否正常工作
5. **文档更新**：合并多个 README 文件为一个统一的文档
plans/search_agent_parallel_optimization.md
+384
View File
@@ -0,0 +1,384 @@
# Search Agent 并行优化计划
## 1. 当前架构分析
### 1.1 执行流程图
```mermaid
flowchart TD
A[用户查询] --> B[查询分析 QueryAnalyzer]
B --> C[搜索规划 SearchPlanner]
C --> D[搜索执行 SearchExecutor]
D --> E[内容提取 ContentExtractor]
E --> F[结果处理 ResultProcessor]
F --> G[答案生成 AnswerGenerator]
G --> H[反思评估 Reflector]
H -->|需要补充搜索| C
H -->|完成| I[返回结果]
style B fill:#f9f,stroke:#333
style C fill:#f9f,stroke:#333
style D fill:#9f9,stroke:#333
style E fill:#9f9,stroke:#333
style F fill:#f9f,stroke:#333
style G fill:#f9f,stroke:#333
style H fill:#f9f,stroke:#333
```
**图例说明：**
- 🟪 粉色：串行执行的LLM调用步骤
- 🟩 绿色：已实现并行的步骤
### 1.2 当前已有的并行化
| 模块 | 文件位置 | 并行方式 | 说明 |
|------|----------|----------|------|
| 搜索执行 | `search_executor.py:51-64` | `asyncio.gather` | 多个搜索任务并行执行 |
| 内容提取 | `jina_reader.py:115-148` | `asyncio.gather` + Semaphore | 批量URL并行提取，限制并发数为5 |
### 1.3 时间消耗分析
基于代码分析，各步骤的预估耗时：
| 步骤 | 类型 | 预估耗时 | 说明 |
|------|------|----------|------|
| 查询分析 | LLM调用 | 1-3秒 | 单次LLM调用 |
| 搜索规划 | 本地计算 | <0.1秒 | 纯逻辑处理 |
| 搜索执行 | API调用 | 2-5秒 | 并行调用Serper API |
| 内容提取 | API调用 | 5-15秒 | 并行调用Jina Reader，最多10个URL |
| 结果处理 | API调用 | 1-3秒 | 调用Jina Reranker |
| 答案生成 | LLM调用 | 3-8秒 | 单次LLM调用，timeout=120s |
| 反思评估 | LLM调用 | 1-3秒 | 单次LLM调用 |
**单次迭代总耗时：约 13-37 秒**
---
## 2. 优化方案
### 2.1 方案一：流水线并行（Pipeline Parallelism）
**核心思想：** 在不破坏数据依赖的前提下，让可以并行的步骤同时执行。
#### 优化点1：搜索执行与部分内容提取并行
当前流程：
```
搜索执行(全部完成) → 内容提取(全部完成)
```
优化后：
```
搜索任务1完成 → 立即开始提取任务1的URL
搜索任务2完成 → 立即开始提取任务2的URL
...
```
**实现方式：** 使用 `asyncio.as_completed` 或流式处理
```python
async def execute_and_extract_streaming(self, plan: SearchPlan) -> List[Document]:
# 创建搜索任务
search_tasks = [self._execute_task(task) for task in plan.tasks]
all_documents = []
extraction_tasks = []
# 使用 as_completed 流式处理
for coro in asyncio.as_completed(search_tasks):
search_results = await coro
# 搜索结果一出来就开始提取
if search_results:
extraction_task = self.content_extractor.extract_batch(search_results)
extraction_tasks.append(extraction_task)
# 等待所有提取任务完成
if extraction_tasks:
results = await asyncio.gather(*extraction_tasks)
for docs in results:
all_documents.extend(docs)
return all_documents
```
**预期收益：** 减少 2-5 秒等待时间
#### 优化点2：答案生成与反思评估并行准备
当前流程：
```
答案生成(完成) → 反思评估(开始)
```
优化后：
```
答案生成(完成) → 同时启动反思评估 + 准备下一轮搜索计划
```
**注意：** 这个优化收益有限，因为反思评估必须等待答案完成。
### 2.2 方案二：批量并行（Batch Parallelism）
**核心思想：** 增加单次迭代的并行度，减少迭代次数。
#### 优化点1：增加搜索任务并发数
当前配置：
- 搜索任务数：3-4个（原始查询 + 2个扩展查询 + 可选新闻）
- 内容提取并发：5个
优化建议：
- 增加扩展查询数量到 4-5 个
- 增加内容提取并发数到 8-10 个
**配置修改：**
```python
# config.py
MAX_EXPANDED_QUERIES = 4 # 从2增加到4
MAX_CONCURRENT_EXTRACTION = 10 # 从5增加到10
```
#### 优化点2：提前终止策略
在内容提取阶段，如果已经获取到足够高质量的内容，可以提前终止其他提取任务。
```python
async def extract_batch_with_early_stop(
self,
search_results: List[SearchResult],
min_docs: int = 5,
quality_threshold: float = 0.8
) -> List[Document]:
# 使用 as_completed 并在达到阈值时取消剩余任务
...
```
### 2.3 方案三：跳过非必要步骤
#### 优化点1：简单查询跳过反思
对于简单查询（如事实查询），如果答案置信度为 high，可以跳过反思评估。
```python
# agent.py 修改
if answer.confidence == "high" and analysis.intent == Intent.FACT_CHECK:
# 跳过反思，直接返回
break
```
#### 优化点2：缓存查询分析结果
对于相似的查询，可以复用之前的查询分析结果。
### 2.4 方案四：异步流式响应（Streaming Response）
**核心思想：** 不等待全部完成，逐步返回结果给用户。
```mermaid
sequenceDiagram
participant User
participant Agent
participant Search
participant Extract
participant LLM
User->>Agent: 发起搜索
Agent->>Search: 开始搜索
Agent-->>User: 状态: 正在搜索...
Search-->>Agent: 搜索结果
Agent->>Extract: 开始提取
Agent-->>User: 状态: 正在提取内容...
Extract-->>Agent: 文档内容
Agent->>LLM: 生成答案
Agent-->>User: 状态: 正在生成答案...
LLM-->>Agent: 答案
Agent-->>User: 最终答案
```
**实现方式：** 使用 SSE（Server-Sent Events）或 WebSocket
---
## 3. 推荐实施方案
### 3.1 优先级排序
| 优先级 | 优化项 | 预期收益 | 实现复杂度 | 风险 |
|--------|--------|----------|------------|------|
| P0 | 搜索+提取流水线并行 | 高（减少3-8秒） | 中 | 低 |
| P0 | 增加内容提取并发数 | 中（减少2-5秒） | 低 | 低 |
| P1 | 简单查询跳过反思 | 中（减少1-3秒） | 低 | 低 |
| P1 | 流式响应 | 高（用户体验） | 高 | 中 |
| P2 | 提前终止策略 | 中 | 中 | 中 |
| P2 | 查询分析缓存 | 低 | 中 | 低 |
### 3.2 第一阶段实施（P0）
#### 修改文件清单
1. **`search_agent/src/server/core/agent.py`**
- 合并搜索执行和内容提取为流水线模式
- 添加新方法 `_search_and_extract_pipeline`
2. **`search_agent/src/server/core/modules/search_executor.py`**
- 添加流式搜索方法 `execute_streaming`
- 返回 `AsyncIterator[List[SearchResult]]`
3. **`search_agent/src/server/core/tools/jina_reader.py`**
- 增加 `max_concurrent` 配置到 10
- 添加提前终止支持
4. **`search_agent/src/server/core/config.py`**
- 添加新配置项 `MAX_CONCURRENT_EXTRACTION`
---
## 4. 详细实施步骤
### 4.1 步骤1：修改配置文件
```python
# config.py 添加
MAX_CONCURRENT_EXTRACTION: int = 10 # 内容提取最大并发数
ENABLE_PIPELINE_MODE: bool = True # 启用流水线模式
```
### 4.2 步骤2：实现流水线搜索+提取
```python
# agent.py 新增方法
async def _search_and_extract_pipeline(
self,
plan: SearchPlan
) -> List[Document]:
"""
流水线模式：搜索完成后立即开始提取
"""
all_documents: List[Document] = []
extraction_tasks: List[asyncio.Task] = []
seen_urls: set = set()
# 创建搜索协程
search_coros = [
self.search_executor._execute_task(task)
for task in plan.tasks
]
# 使用 as_completed 流式处理搜索结果
for coro in asyncio.as_completed(search_coros):
try:
search_results = await coro
if not search_results:
continue
# 去重
new_results = []
for r in search_results:
if r.url not in seen_urls:
seen_urls.add(r.url)
new_results.append(r)
# 立即启动内容提取
if new_results:
task = asyncio.create_task(
self.content_extractor.extract_batch(new_results, max_urls=5)
)
extraction_tasks.append(task)
except Exception as e:
logger.warning(f"搜索任务失败: {e}")
# 等待所有提取任务完成
if extraction_tasks:
results = await asyncio.gather(*extraction_tasks, return_exceptions=True)
for result in results:
if isinstance(result, list):
all_documents.extend(result)
return all_documents
```
### 4.3 步骤3：增加内容提取并发数
```python
# jina_reader.py 修改
def __init__(
self,
api_key: str,
timeout: int = 30,
max_concurrent: int = 10, # 从5改为10
max_content_length: int = 5000
):
```
### 4.4 步骤4：添加简单查询快速路径
```python
# agent.py 修改 search 方法
# 在答案生成后添加
if answer.confidence == "high" and analysis.intent in [Intent.FACT_CHECK, Intent.HOW_TO]:
logger.info("高置信度答案，跳过反思评估")
break
```
---
## 5. 预期效果
### 5.1 优化前后对比
| 场景 | 优化前耗时 | 优化后耗时 | 提升 |
|------|------------|------------|------|
| 简单事实查询 | 15-25秒 | 8-15秒 | 40-50% |
| 复杂研究查询 | 30-60秒 | 20-40秒 | 30-40% |
| 多轮迭代查询 | 60-120秒 | 40-80秒 | 30-35% |
### 5.2 架构优化后流程图
```mermaid
flowchart TD
A[用户查询] --> B[查询分析]
B --> C[搜索规划]
C --> D[流水线执行]
subgraph D[流水线执行]
D1[搜索任务1] --> E1[提取1]
D2[搜索任务2] --> E2[提取2]
D3[搜索任务3] --> E3[提取3]
end
D --> F[结果处理]
F --> G[答案生成]
G --> H{置信度检查}
H -->|高置信度| I[返回结果]
H -->|需要补充| J[反思评估]
J -->|继续| C
J -->|完成| I
style D fill:#9f9,stroke:#333
style H fill:#ff9,stroke:#333
```
---
## 6. 待办事项清单
- [x] 修改 `config.py` 添加新配置项
- [x] 修改 `jina_reader.py` 增加并发数到10
- [x] 在 `agent.py` 实现 `_search_and_extract_pipeline` 方法
- [x] 修改 `agent.py` 的 `search` 方法使用流水线模式
- [x] 添加高置信度快速返回逻辑
- [x] 修改 `content_extractor.py` 使用配置的并发数
- [x] 修改 `quick_search` 方法支持流水线模式
- [x] 更新 `.env` 文件添加新配置项
- [x] 更新 README 文档说明新配置项
- [ ] 添加单元测试验证并行执行
- [ ] 性能测试对比优化前后耗时
---
## 7. 风险与注意事项
1. **API 限流风险**：增加并发可能触发 Jina/Serper API 的限流，需要监控错误率
2. **内存占用**：并行任务增加会占用更多内存，需要监控
3. **错误处理**：流水线模式下错误处理更复杂，需要确保异常不会导致整体失败
4. **测试覆盖**：需要添加并行场景的测试用例
search_agent/.env
+27
View File
@@ -0,0 +1,27 @@
# Search Agent 环境变量配置
# LLM 配置
OPENAI_BASE_URL=https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1
OPENAI_API_KEY=sk
# Serper 搜索 API
SERPER_API_KEY=8253b4f240b520194065312f90e85f9be0fa205f
# Jina API
JINA_API_KEY=jina_e26dc30420a44a1e859216528065b203TkMRmsoz-FgMDQC5FZX9jr5oF2CI
# Agent 配置
MAX_ITERATIONS=3
MAX_RESULTS_PER_QUERY=10
CONTENT_MAX_LENGTH=5000
# 日志级别
LOG_LEVEL=INFO
# 并行优化配置
# 内容提取最大并发数（默认10，提升并行效率）
MAX_CONCURRENT_EXTRACTION=10
# 启用流水线模式：搜索完成后立即开始内容提取（默认true）
ENABLE_PIPELINE_MODE=true
# 启用高置信度快速返回：简单查询跳过反思评估（默认true）
ENABLE_FAST_RETURN=true
search_agent/Dockerfile
+20
View File
@@ -0,0 +1,20 @@
FROM python:3.12-slim
WORKDIR /app
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
RUN apt-get update && apt-get install -y gcc curl && rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
EXPOSE 8000
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
CMD curl -f http://localhost:8000/health || exit 1
CMD ["python", "run_api_server.py"]
search_agent/README.md
+230
View File
@@ -0,0 +1,230 @@
# Search Agent
基于 **Pydantic AI** 的智能搜索 Agent，能够理解用户查询意图、自动规划搜索策略、从多个来源获取信息，并生成高质量、有来源引用的答案。
## 功能特点
| 能力 | 描述 |
|------|------|
| 🧠 查询理解 | 分析用户意图，提取关键实体，生成扩展查询 |
| 📋 搜索规划 | 智能分解问题，制定搜索策略 |
| 🔎 多源搜索 | 支持Web搜索和新闻搜索 |
| 📄 内容提取 | 智能提取网页核心内容 |
| 🎯 结果排序 | 基于相关性重排搜索结果 |
| ✍️ 答案生成 | 综合信息生成结构化回答 |
| 🔄 自我反思 | 评估答案质量，决定是否迭代 |
| ⚡ 并行优化 | 流水线执行，高置信度快速返回 |
## 性能优化
本 Agent 实现了多项并行优化，显著提升响应速度：
### 流水线模式（Pipeline Mode）
传统模式下，搜索和内容提取是串行执行的：
```
搜索任务1 → 搜索任务2 → 搜索任务3 → 等待全部完成 → 内容提取
```
流水线模式下，搜索完成后立即开始提取：
```
搜索任务1完成 → 立即开始提取1
搜索任务2完成 → 立即开始提取2
搜索任务3完成 → 立即开始提取3
```
**预期收益：减少 3-8 秒等待时间**
### 高置信度快速返回
对于简单查询（事实查询、操作指南），如果答案置信度为 `high`，则跳过反思评估步骤，直接返回结果。
**预期收益：减少 1-3 秒 LLM 调用时间**
### 性能对比
| 场景 | 优化前耗时 | 优化后耗时 | 提升 |
|------|------------|------------|------|
| 简单事实查询 | 15-25秒 | 8-15秒 | 40-50% |
| 复杂研究查询 | 30-60秒 | 20-40秒 | 30-40% |
## 快速开始
### 1. 安装依赖
```bash
cd search_agent
pip install -r requirements.txt
```
### 2. 配置环境变量
```bash
# LLM 配置
export OPENAI_BASE_URL=https://your-litellm-gateway/v1
export OPENAI_API_KEY=your-api-key
export MODEL_NAME=taiji/gpt-4o-mini
# Serper API（Google搜索）
export SERPER_API_KEY=your-serper-key
# Jina API（内容提取和重排序）
export JINA_API_KEY=your-jina-key
# 可选配置
export MAX_ITERATIONS=3
export MAX_RESULTS_PER_QUERY=10
export CONTENT_MAX_LENGTH=5000
export API_PORT=8000
```
### 3. 本地测试
```bash
python run_api_server.py
```
### 4. 构建镜像
```bash
docker build -t search-agent:latest .
```
## 项目结构
```
search_agent/
├── Dockerfile
├── requirements.txt
├── run_api_server.py # 启动脚本
└── src/
├── __init__.py
└── server/
├── __init__.py
├── api_server.py # FastAPI + MCP HTTP
├── mcp_server.py # MCP 工具定义
└── core/ # 业务逻辑
├── config.py # 配置管理
├── agent.py # SearchAgent 主类
├── schemas.py # 数据模型
├── prompts.py # Prompt 模板
├── modules/ # 功能模块
│ ├── query_analyzer.py
│ ├── search_planner.py
│ ├── search_executor.py
│ ├── content_extractor.py
│ ├── result_processor.py
│ ├── answer_generator.py
│ └── reflector.py
├── tools/ # 外部 API 封装
│ ├── serper.py
│ ├── jina_reader.py
│ └── jina_reranker.py
└── utils/ # 工具函数
├── llm_client.py
└── helpers.py
```
## API 端点
### 健康检查
- `GET /` - 服务状态
- `GET /health` - 健康检查
### MCP 端点
- `POST /mcp` - MCP HTTP 端点
- `GET /mcp/sse` - MCP SSE 端点
- `POST /mcp/sse` - MCP SSE POST 端点
### REST API
- `POST /api/v1/search` - 智能搜索
- `POST /api/v1/quick_search` - 快速搜索
- `GET /api/v1/tools` - 获取工具列表
## MCP 工具
### search
执行智能搜索，支持多轮迭代和自我反思。
```json
{
"query": "什么是Kubernetes？",
"max_iterations": 3
}
```
### quick_search
快速搜索，单次迭代，适合简单问题。
```json
{
"query": "Python是什么？"
}
```
## 使用示例
### Python
```python
import requests
# 智能搜索
response = requests.post(
"http://localhost:8000/api/v1/search",
headers={"api-key": "your-api-key"},
json={"query": "什么是Kubernetes？", "max_iterations": 3},
timeout=120
)
result = response.json()
print(result["answer"]["content"])
```
### cURL
```bash
curl -X POST "http://localhost:8000/api/v1/search" \
-H "api-key: your-api-key" \
-H "Content-Type: application/json" \
-d '{"query": "什么是Kubernetes？"}'
```
## 环境变量
| 变量 | 必需 | 说明 | 默认值 |
|------|------|------|--------|
| OPENAI_BASE_URL | 是 | LiteLLM Gateway URL | - |
| OPENAI_API_KEY | 是 | API Key | sk |
| MODEL_NAME | 否 | 模型名称 | taiji/gpt-4o-mini |
| SERPER_API_KEY | 是 | Serper API Key | - |
| JINA_API_KEY | 是 | Jina API Key | - |
| MAX_ITERATIONS | 否 | 最大迭代次数 | 3 |
| MAX_RESULTS_PER_QUERY | 否 | 每次搜索结果数 | 10 |
| CONTENT_MAX_LENGTH | 否 | 内容最大长度 | 5000 |
| API_PORT | 否 | 服务端口 | 8000 |
| LOG_LEVEL | 否 | 日志级别 | INFO |
| TIMEOUT | 否 | 请求超时时间 | 30 |
| **并行优化配置** | | | |
| MAX_CONCURRENT_EXTRACTION | 否 | 内容提取最大并发数 | 10 |
| ENABLE_PIPELINE_MODE | 否 | 启用流水线模式 | true |
| ENABLE_FAST_RETURN | 否 | 启用高置信度快速返回 | true |
## 注册到 Agent Manager
在 `k8s_manager.py` 中添加：
```python
# TEMPLATE_PORTS
"search_agent": 8000,
# image_map
"search_agent": "agnettaiji.azurecr.io/ai-agents/search-agent:latest",
```
在 `app.py` 的 `valid_templates` 中添加 `"search_agent"`。
search_agent/USAGE.md
+376
View File
@@ -0,0 +1,376 @@
# 智能搜索 Agent
---
本 Agent 提供基于 **多源搜索与AI生成** 的智能问答能力，通过 **HTTP API** 与 **MCP（Model Context Protocol）** 对外提供服务。
核心能力：
- **查询理解**：分析用户意图，提取关键实体，生成扩展查询
- **多源搜索**：支持 Web 搜索和新闻搜索
- **内容提取**：智能提取网页核心内容
- **结果排序**：基于相关性重排搜索结果
- **答案生成**：综合信息生成结构化回答
- **自我反思**：评估答案质量，决定是否迭代
---
## 功能概览
提供智能搜索的 **查询分析、多源检索、内容提取、答案生成** 能力，返回带来源引用的结构化答案。
支持能力：
- 多轮迭代搜索与自我反思
- 多语言查询扩展
- 网页内容智能提取
- 基于相关性的结果重排
- 带来源引用的答案生成
---
## 1⃣ search — 智能搜索
### 功能说明
执行 **多轮迭代搜索**：分析查询意图、规划搜索策略、从多个来源获取信息、生成高质量答案，并通过自我反思评估答案质量决定是否继续迭代。
---
### REST API 调用
```
POST /api/v1/search
Content-Type: application/json
api-key: {your-api-key}
```
```json
{
"query": "2024年诺贝尔物理学奖得主是谁",
"max_iterations": 3
}
```
---
### MCP 调用
```json
{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "search",
"arguments": {
"query": "2024年诺贝尔物理学奖得主是谁",
"max_iterations": 3
}
}
}
```
---
### 参数说明
| 参数 | 类型 | 必需 | 默认值 | 说明 |
| --- | --- | --- | --- | --- |
| query | string | ✅ | - | 搜索查询问题 |
| max_iterations | integer | ❌ | 3 | 最大迭代次数（1-10） |
---
### 搜索流程
| 阶段 | 模块 | 说明 |
| --- | --- | --- |
| 1 | 查询分析 | 分析用户意图，提取关键实体，生成扩展查询 |
| 2 | 搜索规划 | 根据查询类型规划搜索任务（Web/新闻） |
| 3 | 搜索执行 | 并行执行多个搜索任务 |
| 4 | 内容提取 | 从搜索结果中提取网页核心内容 |
| 5 | 结果处理 | 基于相关性重排搜索结果 |
| 6 | 答案生成 | 综合信息生成带来源引用的答案 |
| 7 | 自我反思 | 评估答案质量，决定是否继续迭代 |
---
### 返回结果
```json
{
"success": true,
"answer": {
"content": "# 2024年诺贝尔物理学奖得主\n\n2024年诺贝尔物理学奖授予了...\n\n---\n**参考来源**：\n1. [诺贝尔奖官网](https://www.nobelprize.org/)\n2. [维基百科](https://zh.wikipedia.org/)",
"sources": [
{
"index": 1,
"title": "诺贝尔奖官网",
"url": "https://www.nobelprize.org/",
"relevance": "high",
"relevance_score": 0.892
},
{
"index": 2,
"title": "维基百科",
"url": "https://zh.wikipedia.org/",
"relevance": "high",
"relevance_score": 0.856
}
],
"confidence": "high",
"all_sources": {
"high": [
{
"index": 1,
"title": "诺贝尔奖官网",
"url": "https://www.nobelprize.org/",
"relevance": "high",
"relevance_score": 0.892
},
{
"index": 2,
"title": "维基百科",
"url": "https://zh.wikipedia.org/",
"relevance": "high",
"relevance_score": 0.856
}
],
"medium": [
{
"index": 3,
"title": "新华网",
"url": "https://www.xinhuanet.com/",
"relevance": "medium",
"relevance_score": 0.623
}
],
"low": [
{
"index": 4,
"title": "某博客",
"url": "https://example.com/blog",
"relevance": "low",
"relevance_score": 0.312
}
]
}
},
"statistics": {
"iterations": 1,
"total_sources_consulted": 10,
"search_queries_used": [
"2024 Nobel Prize in Physics winner",
"2024年诺贝尔物理学奖得主是谁"
]
}
}
```
---
### 返回字段说明
| 字段 | 类型 | 说明 |
| --- | --- | --- |
| success | boolean | 处理是否成功 |
| answer.content | string | Markdown 格式的答案内容 |
| answer.sources | array | 答案中引用的来源列表 |
| answer.sources[].relevance | string | 来源相关性等级（high/medium/low） |
| answer.sources[].relevance_score | number | 来源相关性分数（0-1） |
| answer.confidence | string | 答案置信度（high/medium/low） |
| answer.all_sources | object | **所有搜索来源（按相关性分组）** |
| answer.all_sources.high | array | 高相关性来源（分数 >= 0.7） |
| answer.all_sources.medium | array | 中相关性来源（0.4 <= 分数 < 0.7） |
| answer.all_sources.low | array | 低相关性来源（分数 < 0.4） |
| statistics.iterations | integer | 实际迭代次数 |
| statistics.total_sources_consulted | integer | 咨询的来源总数 |
| statistics.search_queries_used | array | 使用的搜索查询列表 |
---
## 2️⃣ quick_search — 快速搜索
### 功能说明
执行 **单次迭代搜索**，适合简单问题，执行速度更快。限制搜索任务数量，快速生成答案。
---
### REST API 调用
```
POST /api/v1/quick_search
Content-Type: application/json
api-key: {your-api-key}
```
```json
{
"query": "什么是人工智能"
}
```
---
### MCP 调用
```json
{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "quick_search",
"arguments": {
"query": "什么是人工智能"
}
}
}
```
---
### 参数说明
| 参数 | 类型 | 必需 | 默认值 | 说明 |
| --- | --- | --- | --- | --- |
| query | string | ✅ | - | 搜索查询问题 |
---
### 返回结果
```json
{
"success": true,
"answer": {
"content": "# 什么是人工智能\n\n人工智能（AI）是一个广泛的计算机科学领域...\n\n---\n**参考来源**：\n1. [IBM](https://www.ibm.com/)\n2. [Google Cloud](https://cloud.google.com/)",
"sources": [
{
"index": 1,
"title": "IBM - 人工智能",
"url": "https://www.ibm.com/cn-zh/think/topics/artificial-intelligence",
"relevance": "high",
"relevance_score": 0.912
},
{
"index": 2,
"title": "Google Cloud - 什么是人工智能",
"url": "https://cloud.google.com/learn/what-is-artificial-intelligence",
"relevance": "high",
"relevance_score": 0.878
}
],
"confidence": "high",
"all_sources": {
"high": [
{
"index": 1,
"title": "IBM - 人工智能",
"url": "https://www.ibm.com/cn-zh/think/topics/artificial-intelligence",
"relevance": "high",
"relevance_score": 0.912
},
{
"index": 2,
"title": "Google Cloud - 什么是人工智能",
"url": "https://cloud.google.com/learn/what-is-artificial-intelligence",
"relevance": "high",
"relevance_score": 0.878
}
],
"medium": [
{
"index": 3,
"title": "百度百科 - 人工智能",
"url": "https://baike.baidu.com/item/人工智能",
"relevance": "medium",
"relevance_score": 0.654
}
],
"low": []
}
}
}
```
---
## 统一错误格式
成功：
```json
{
"success": true,
"answer": {},
"statistics": {}
}
```
失败：
```json
{
"success": false,
"error": "错误描述"
}
```
---
## 服务端点
| 端点 | 方法 | 说明 |
| --- | --- | --- |
| / | GET | 服务状态 |
| /health | GET | 健康检查 |
| /mcp | POST | MCP JSON-RPC |
| /mcp/sse | GET/POST | MCP SSE 流式 |
| /api/v1/search | POST | 智能搜索（多轮迭代） |
| /api/v1/quick_search | POST | 快速搜索（单次迭代） |
| /api/v1/tools | GET | 获取可用工具列表 |
---
## 认证方式
API 调用需要通过以下方式之一传递 API Key：
1. **请求头 api-key**：
```
api-key: your-api-key
```
2. **Authorization Bearer**：
```
Authorization: Bearer your-api-key
```
---
## 部署信息
| 配置项 | 值 |
| --- | --- |
| 镜像名称 | search-agent:latest |
| 服务端口 | 8000 |
| 健康检查 | /health |
---
## 环境变量
| 变量名 | 必需 | 默认值 | 说明 |
| --- | --- | --- | --- |
| OPENAI_BASE_URL | ❌ | LiteLLM Gateway | LLM API 基础 URL |
| OPENAI_API_KEY | ❌ | sk | LLM API 密钥（可通过请求头传递） |
| MODEL_NAME | ❌ | taiji/gpt-4o-mini | 模型名称 |
| SERPER_API_KEY | ❌ | 内置 | Serper 搜索 API 密钥 |
| JINA_API_KEY | ❌ | 内置 | Jina API 密钥 |
| MAX_ITERATIONS | ❌ | 3 | 默认最大迭代次数 |
| MAX_RESULTS_PER_QUERY | ❌ | 10 | 每次查询最大结果数 |
| CONTENT_MAX_LENGTH | ❌ | 5000 | 内容最大长度 |
search_agent/requirements.txt
+20
View File
@@ -0,0 +1,20 @@
# Pydantic AI
pydantic-ai>=0.0.14
# MCP
mcp>=0.9.0
fastmcp>=0.1.0
# FastAPI
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
pydantic>=2.5.0
# HTTP Client
aiohttp>=3.9.0
# 环境变量
python-dotenv>=1.0.0
# 日志
loguru>=0.7.0
search_agent/run_api_server.py
+17
View File
@@ -0,0 +1,17 @@
#!/usr/bin/env python
"""启动 Search Agent API 服务器"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
if __name__ == '__main__':
from src.server.api_server import app
import uvicorn
import os
host = os.getenv('API_HOST', '0.0.0.0')
port = int(os.getenv('API_PORT', '8000'))
print(f"🚀 启动 Search Agent API: http://{host}:{port}")
uvicorn.run(app, host=host, port=port, log_level="info")
search_agent/src/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Search Agent 源代码包"""
search_agent/src/__pycache__/__init__.cpython-312.pyc
BIN
View File
Binary file not shown.
search_agent/src/server/__init__.py
+1
View File
@@ -0,0 +1 @@
"""服务器模块"""
search_agent/src/server/__pycache__/__init__.cpython-312.pyc
BIN
View File
Binary file not shown.
search_agent/src/server/__pycache__/api_server.cpython-312.pyc
BIN
View File
Binary file not shown.
search_agent/src/server/api_server.py
+293
View File
@@ -0,0 +1,293 @@
"""
Search Agent API 服务器
提供 REST API 和 MCP HTTP/SSE 端点。
"""
import os
import json
import uuid
from typing import Optional, Dict, Any, AsyncGenerator
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Header, Request, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field
from .mcp_server import TOOL_MAP, TOOL_LIST
# ==================== 配置 ====================
SERVER_NAME = "Search Agent API"
# ==================== FastAPI 应用 ====================
@asynccontextmanager
async def lifespan(app: FastAPI):
"""应用生命周期管理"""
print(f"🚀 {SERVER_NAME} 启动")
yield
print(f"🛑 {SERVER_NAME} 关闭")
app = FastAPI(
title=SERVER_NAME,
description="智能搜索服务，基于多源搜索和AI生成答案",
version="1.0.0",
lifespan=lifespan
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# ==================== API Key 验证 ====================
async def verify_api_key(
api_key: Optional[str] = Header(None, alias="api-key"),
authorization: Optional[str] = Header(None)
) -> str:
"""验证 API Key"""
if api_key and api_key.strip() and api_key.strip() != "sk":
return api_key.strip()
if authorization:
key = authorization[7:].strip() if authorization.startswith("Bearer ") else authorization.strip()
if key and key != "sk":
return key
raise HTTPException(status_code=401, detail="缺少 API Key")
def get_api_key_from_request(request: Request) -> Optional[str]:
"""从请求头提取 API Key（不验证）"""
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
# ==================== 请求模型 ====================
class SearchRequest(BaseModel):
"""搜索请求"""
query: str = Field(..., description="搜索查询问题", min_length=1, max_length=500)
max_iterations: Optional[int] = Field(3, description="最大迭代次数", ge=1, le=10)
class QuickSearchRequest(BaseModel):
"""快速搜索请求"""
query: str = Field(..., description="搜索查询问题", min_length=1, max_length=500)
# ==================== 健康检查 ====================
@app.get("/")
async def root():
"""服务状态"""
return {
"service": SERVER_NAME,
"status": "running",
"version": "1.0.0",
"tools": list(TOOL_MAP.keys())
}
@app.get("/health")
async def health():
"""健康检查"""
return {"status": "healthy", "service": SERVER_NAME}
# ==================== MCP 端点 ====================
sessions: Dict[str, Dict] = {}
async def handle_mcp_request(data: Dict, session_id: str = None, api_key: str = None) -> Dict:
"""处理 MCP JSON-RPC 请求"""
method = data.get("method")
params = data.get("params", {})
req_id = data.get("id")
# tools/call 需要验证 API Key
if method == "tools/call" and (not api_key or api_key == "sk"):
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": TOOL_LIST}}
elif method == "tools/call":
tool_name = params.get("name")
args = params.get("arguments", {})
if tool_name not in TOOL_MAP:
raise ValueError(f"Unknown tool: {tool_name}")
# 设置 API Key 到环境变量
old_key = os.environ.get('OPENAI_API_KEY')
if api_key:
os.environ['OPENAI_API_KEY'] = api_key
try:
result = await TOOL_MAP[tool_name](**args)
finally:
if old_key:
os.environ['OPENAI_API_KEY'] = old_key
return {
"jsonrpc": "2.0", "id": req_id,
"result": {"content": [{"type": "text", "text": str(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"
import asyncio
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)}})
# ==================== REST API 端点 ====================
@app.post("/api/v1/search")
async def api_search(
request: SearchRequest,
api_key: str = Depends(verify_api_key)
):
"""
执行智能搜索
支持多轮迭代和自我反思，能够：
- 分析查询意图
- 自动规划搜索策略
- 从多个来源获取信息
- 生成高质量答案
"""
try:
# 设置 API Key 到环境变量
old_key = os.environ.get('OPENAI_API_KEY')
os.environ['OPENAI_API_KEY'] = api_key
try:
result = await TOOL_MAP['search'](
query=request.query,
max_iterations=request.max_iterations or 3
)
return json.loads(result)
finally:
if old_key:
os.environ['OPENAI_API_KEY'] = old_key
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/v1/quick_search")
async def api_quick_search(
request: QuickSearchRequest,
api_key: str = Depends(verify_api_key)
):
"""
快速搜索（单次迭代）
适合简单问题，执行速度更快。
"""
try:
# 设置 API Key 到环境变量
old_key = os.environ.get('OPENAI_API_KEY')
os.environ['OPENAI_API_KEY'] = api_key
try:
result = await TOOL_MAP['quick_search'](query=request.query)
return json.loads(result)
finally:
if old_key:
os.environ['OPENAI_API_KEY'] = old_key
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/v1/tools")
async def api_tools():
"""获取可用工具列表"""
return {
"tools": {name: "MCP 工具" for name in TOOL_MAP.keys()},
"count": len(TOOL_MAP)
}
# 导出
__all__ = ['app']
search_agent/src/server/core/__init__.py
+1
View File
@@ -0,0 +1 @@
"""核心业务逻辑模块"""
search_agent/src/server/core/__pycache__/config.cpython-312.pyc
BIN
View File
Binary file not shown.
search_agent/src/server/core/__pycache__/prompts.cpython-312.pyc
BIN
View File
Binary file not shown.
search_agent/src/server/core/__pycache__/schemas.cpython-312.pyc
BIN
View File
Binary file not shown.
search_agent/src/server/core/agent.py
+342
View File
@@ -0,0 +1,342 @@
"""
搜索Agent主类
协调各模块执行智能搜索
优化特性：
- 流水线并行：搜索完成后立即开始内容提取，不等待所有搜索完成
- 高置信度快速返回：简单查询跳过反思评估
- 可配置的并发数：支持更高的内容提取并发
"""
import asyncio
from typing import List, Optional, Set
from loguru import logger
from .config import Config
from .schemas import (
QueryAnalysis,
SearchPlan,
SearchTask,
SearchResult,
Document,
RankedDocument,
Answer,
AgentResponse,
Intent,
)
from .modules.query_analyzer import QueryAnalyzer
from .modules.search_planner import SearchPlanner
from .modules.search_executor import SearchExecutor
from .modules.content_extractor import ContentExtractor
from .modules.result_processor import ResultProcessor
from .modules.answer_generator import AnswerGenerator
from .modules.reflector import Reflector
class SearchAgent:
"""智能搜索Agent"""
def __init__(self, config: Config):
"""
初始化搜索Agent
Args:
config: 配置对象
"""
self.config = config
# 初始化各模块
self.query_analyzer = QueryAnalyzer(config)
self.search_planner = SearchPlanner(config)
self.search_executor = SearchExecutor(config)
self.content_extractor = ContentExtractor(config)
self.result_processor = ResultProcessor(config)
self.answer_generator = AnswerGenerator(config)
self.reflector = Reflector(config)
logger.info("SearchAgent 初始化完成")
async def search(self, query: str) -> AgentResponse:
"""
执行智能搜索
Args:
query: 用户查询
Returns:
AgentResponse对象
"""
logger.info(f"="*60)
logger.info(f"开始搜索: {query}")
logger.info(f"流水线模式: {self.config.enable_pipeline_mode}")
logger.info(f"快速返回: {self.config.enable_fast_return}")
logger.info(f"="*60)
iteration = 0
all_documents: List[Document] = []
all_queries: List[str] = []
seen_urls: Set[str] = set() # 用于流水线模式的URL去重
# 1. 查询理解
analysis = await self.query_analyzer.analyze(query)
logger.info(f"查询分析完成: intent={analysis.intent.value}")
answer: Optional[Answer] = None
while iteration < self.config.max_iterations:
iteration += 1
logger.info(f"\n--- 迭代 {iteration}/{self.config.max_iterations} ---")
# 2. 搜索规划
if iteration == 1:
plan = await self.search_planner.plan(analysis)
else:
# 后续迭代使用建议的补充查询
plan = self.search_planner.plan_supplementary(
query,
analysis.expanded_queries
)
all_queries.extend([t.query for t in plan.tasks])
logger.info(f"搜索计划: {len(plan.tasks)} 个任务")
# 3+4. 搜索执行 + 内容提取（根据配置选择模式）
if self.config.enable_pipeline_mode:
# 流水线模式：搜索完成后立即开始提取
documents = await self._search_and_extract_pipeline(
plan, seen_urls
)
else:
# 传统模式：先完成所有搜索，再提取
search_results = await self.search_executor.execute(plan)
logger.info(f"搜索结果: {len(search_results)} 条")
if not search_results:
logger.warning("没有搜索结果")
if answer is None:
answer = self.answer_generator._empty_answer()
break
documents = await self.content_extractor.extract_batch(
search_results,
max_urls=10
)
all_documents.extend(documents)
logger.info(f"提取文档: {len(documents)} 个")
if not documents:
logger.warning("没有成功提取到文档内容")
if answer is None:
answer = self.answer_generator._empty_answer()
break
# 5. 结果处理（去重+重排序），获取所有排序结果
ranked_docs, all_ranked_docs = await self.result_processor.process_all(
query=query,
documents=all_documents,
top_k=5
)
logger.info(f"排序结果: top {len(ranked_docs)} 个，共 {len(all_ranked_docs)} 个")
if not ranked_docs:
logger.warning("没有有效的排序结果")
continue
# 6. 生成答案（传递所有排序文档用于展示完整来源列表）
answer = await self.answer_generator.generate(
query=query,
documents=ranked_docs,
all_ranked_documents=all_ranked_docs
)
logger.info(f"答案生成完成: confidence={answer.confidence}")
# 7. 高置信度快速返回（跳过反思评估）
if self.config.enable_fast_return and self._should_fast_return(answer, analysis):
logger.info("高置信度答案，跳过反思评估，快速返回")
break
# 8. 反思评估
assessment = await self.reflector.assess(query, answer)
# 9. 判断是否继续迭代
if not self.reflector.should_continue(assessment, iteration):
break
# 更新分析，准备下一轮搜索
if assessment.suggested_queries:
analysis.expanded_queries = assessment.suggested_queries
logger.info(f"补充搜索: {assessment.suggested_queries}")
# 确保有答案返回
if answer is None:
answer = self.answer_generator._empty_answer()
# 去重统计
unique_urls = set(d.url for d in all_documents)
response = AgentResponse(
answer=answer,
iterations=iteration,
total_sources_consulted=len(unique_urls),
search_queries_used=list(set(all_queries))
)
logger.info(f"\n{'='*60}")
logger.info(f"搜索完成!")
logger.info(f"迭代次数: {iteration}")
logger.info(f"参考来源: {len(unique_urls)}")
logger.info(f"搜索查询: {len(response.search_queries_used)}")
logger.info(f"{'='*60}\n")
return response
async def _search_and_extract_pipeline(
self,
plan: SearchPlan,
seen_urls: Set[str]
) -> List[Document]:
"""
流水线模式：搜索完成后立即开始内容提取
优化原理：不等待所有搜索任务完成，每个搜索任务完成后立即开始提取其结果
Args:
plan: 搜索计划
seen_urls: 已处理的URL集合（用于去重）
Returns:
提取的文档列表
"""
logger.info(f"[流水线模式] 开始执行 {len(plan.tasks)} 个搜索任务")
all_documents: List[Document] = []
extraction_tasks: List[asyncio.Task] = []
# 创建搜索协程
search_coros = [
self.search_executor._execute_task(task)
for task in plan.tasks
]
# 使用 as_completed 流式处理搜索结果
for coro in asyncio.as_completed(search_coros):
try:
search_results = await coro
if not search_results:
continue
# 去重：只处理新的URL
new_results = []
for r in search_results:
if r.url not in seen_urls:
seen_urls.add(r.url)
new_results.append(r)
logger.debug(f"[流水线] 搜索返回 {len(search_results)} 条，新增 {len(new_results)} 条")
# 立即启动内容提取（提取所有新URL，由并发控制限制速率）
if new_results:
task = asyncio.create_task(
self.content_extractor.extract_batch(new_results, max_urls=len(new_results))
)
extraction_tasks.append(task)
except Exception as e:
logger.warning(f"[流水线] 搜索任务失败: {e}")
# 等待所有提取任务完成
if extraction_tasks:
logger.info(f"[流水线] 等待 {len(extraction_tasks)} 个提取任务完成")
results = await asyncio.gather(*extraction_tasks, return_exceptions=True)
for result in results:
if isinstance(result, list):
all_documents.extend(result)
elif isinstance(result, Exception):
logger.warning(f"[流水线] 提取任务失败: {result}")
logger.info(f"[流水线模式] 完成，共提取 {len(all_documents)} 个文档")
return all_documents
def _should_fast_return(self, answer: Answer, analysis: QueryAnalysis) -> bool:
"""
判断是否应该快速返回（跳过反思评估）
条件：
1. 答案置信度为 high
2. 查询意图为简单类型（事实查询、操作指南）
Args:
answer: 生成的答案
analysis: 查询分析结果
Returns:
是否应该快速返回
"""
# 简单查询类型
simple_intents = {Intent.FACT_CHECK, Intent.HOW_TO}
return (
answer.confidence == "high" and
analysis.intent in simple_intents
)
async def quick_search(self, query: str) -> Answer:
"""
快速搜索（单次迭代）
优化：支持流水线模式，提升响应速度
Args:
query: 用户查询
Returns:
Answer对象，包含分级的来源列表
"""
logger.info(f"[快速搜索] 开始: {query}")
# 简化分析
analysis = await self.query_analyzer.analyze(query)
# 只执行一次搜索
plan = await self.search_planner.plan(analysis)
plan.tasks = plan.tasks[:2] # 限制搜索任务数量
# 根据配置选择执行模式
if self.config.enable_pipeline_mode:
# 流水线模式
seen_urls: Set[str] = set()
documents = await self._search_and_extract_pipeline(plan, seen_urls)
else:
# 传统模式
search_results = await self.search_executor.execute(plan)
if not search_results:
return self.answer_generator._empty_answer()
documents = await self.content_extractor.extract_batch(
search_results,
max_urls=5
)
if not documents:
return self.answer_generator._empty_answer()
# 处理结果，获取所有排序文档
ranked_docs, all_ranked_docs = await self.result_processor.process_all(
query=query,
documents=documents,
top_k=3
)
if not ranked_docs:
return self.answer_generator._empty_answer()
# 生成答案（传递所有排序文档用于展示完整来源列表）
answer = await self.answer_generator.generate(
query=query,
documents=ranked_docs,
all_ranked_documents=all_ranked_docs
)
logger.info(f"[快速搜索] 完成: confidence={answer.confidence}")
return answer
search_agent/src/server/core/config.py
+118
View File
@@ -0,0 +1,118 @@
"""
配置管理模块
负责加载和管理所有配置项
"""
import os
from dataclasses import dataclass
from typing import Optional
from dotenv import load_dotenv
@dataclass
class Config:
"""Agent配置类"""
# LLM配置
llm_base_url: str
llm_api_key: str # 可以通过请求头或body传递
llm_model: str
# Serper配置
serper_api_key: str
# Jina配置
jina_api_key: str
# Agent配置
max_iterations: int
max_results_per_query: int
content_max_length: int
# 可选配置
log_level: str = "INFO"
timeout: int = 30
# 并行优化配置
max_concurrent_extraction: int = 10 # 内容提取最大并发数
enable_pipeline_mode: bool = True # 启用流水线模式（搜索完成后立即提取）
enable_fast_return: bool = True # 高置信度答案快速返回（跳过反思）
@classmethod
def from_env(cls, env_path: Optional[str] = None, api_key: Optional[str] = None) -> "Config":
"""从环境变量加载配置
Args:
env_path: 可选的.env文件路径
api_key: 可选的API密钥，通过请求头或body传递时使用
"""
if env_path:
load_dotenv(env_path)
else:
load_dotenv()
# LLM配置 - 与模板 mcp_server.py 保持一致的读取方式
llm_base_url = os.getenv("OPENAI_BASE_URL",
os.getenv("LLM_BASE_URL", "https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1"))
llm_api_key = api_key or os.getenv("OPENAI_API_KEY", "sk")
llm_model = os.getenv("MODEL_NAME", os.getenv("LITELLM_MODEL", "taiji/gpt-4o-mini"))
return cls(
# LLM配置
llm_base_url=llm_base_url,
llm_api_key=llm_api_key,
llm_model=llm_model,
# Serper配置 - 内置默认值
serper_api_key=os.getenv("SERPER_API_KEY", "8253b4f240b520194065312f90e85f9be0fa205f"),
# Jina配置 - 内置默认值
jina_api_key=os.getenv("JINA_API_KEY", "jina_e26dc30420a44a1e859216528065b203TkMRmsoz-FgMDQC5FZX9jr5oF2CI"),
# Agent配置
max_iterations=int(os.getenv("MAX_ITERATIONS", "3")),
max_results_per_query=int(os.getenv("MAX_RESULTS_PER_QUERY", "10")),
content_max_length=int(os.getenv("CONTENT_MAX_LENGTH", "5000")),
# 可选配置
log_level=os.getenv("LOG_LEVEL", "INFO"),
timeout=int(os.getenv("TIMEOUT", "30")),
# 并行优化配置
max_concurrent_extraction=int(os.getenv("MAX_CONCURRENT_EXTRACTION", "10")),
enable_pipeline_mode=os.getenv("ENABLE_PIPELINE_MODE", "true").lower() == "true",
enable_fast_return=os.getenv("ENABLE_FAST_RETURN", "true").lower() == "true"
)
def with_api_key(self, api_key: str) -> "Config":
"""返回一个使用新API密钥的配置副本"""
return Config(
llm_base_url=self.llm_base_url,
llm_api_key=api_key,
llm_model=self.llm_model,
serper_api_key=self.serper_api_key,
jina_api_key=self.jina_api_key,
max_iterations=self.max_iterations,
max_results_per_query=self.max_results_per_query,
content_max_length=self.content_max_length,
log_level=self.log_level,
timeout=self.timeout,
max_concurrent_extraction=self.max_concurrent_extraction,
enable_pipeline_mode=self.enable_pipeline_mode,
enable_fast_return=self.enable_fast_return
)
def validate(self) -> bool:
"""验证配置是否完整"""
required_fields = [
("llm_base_url", self.llm_base_url),
("serper_api_key", self.serper_api_key),
("jina_api_key", self.jina_api_key),
]
missing = [name for name, value in required_fields if not value]
if missing:
raise ValueError(f"缺少必要的配置项: {', '.join(missing)}")
return True
search_agent/src/server/core/modules/__init__.py
+21
View File
@@ -0,0 +1,21 @@
"""
核心模块
"""
from .query_analyzer import QueryAnalyzer
from .search_planner import SearchPlanner
from .search_executor import SearchExecutor
from .content_extractor import ContentExtractor
from .result_processor import ResultProcessor
from .answer_generator import AnswerGenerator
from .reflector import Reflector
__all__ = [
"QueryAnalyzer",
"SearchPlanner",
"SearchExecutor",
"ContentExtractor",
"ResultProcessor",
"AnswerGenerator",
"Reflector",
]
search_agent/src/server/core/modules/answer_generator.py
+260
View File
@@ -0,0 +1,260 @@
"""
答案生成模块
综合多个来源的信息生成结构化答案
优化特性：
- 展示所有搜索来源，按相关性分为高、中、低三级
- 提供更专业的来源引用展示
"""
from typing import List, Optional
from loguru import logger
from ..config import Config
from ..schemas import RankedDocument, Answer, Source, SourcesByRelevance
from ..utils.llm_client import LLMClient
from ..utils.helpers import format_documents_for_prompt
# 相关性分数阈值
RELEVANCE_HIGH_THRESHOLD = 0.7 # >= 0.7 为高相关性
RELEVANCE_MEDIUM_THRESHOLD = 0.4 # >= 0.4 为中相关性，< 0.4 为低相关性
# 答案生成Prompt
ANSWER_GENERATION_PROMPT = """你是一个专业的信息整合专家。根据以下搜索结果，回答用户的问题。
## 要求
1. 综合多个来源的信息，给出全面准确的回答
2. 使用清晰的结构组织答案（标题、列表、重点标注等）
3. 在答案中标注信息来源，格式：[来源1]、[来源2]
4. 如果信息有冲突，说明不同观点
5. 如果信息不足以完整回答问题，明确指出缺失的部分
6. 回答使用中文
7. 优先使用高相关性来源的信息
## 输出JSON格式
{
"answer": "结构化的答案（Markdown格式，包含来源引用）",
"sources": [
{"index": 1, "title": "来源标题", "url": "来源URL"},
{"index": 2, "title": "来源标题", "url": "来源URL"}
],
"confidence": "high/medium/low，基于信息质量和一致性判断"
}"""
class AnswerGenerator:
"""答案生成模块"""
def __init__(self, config: Config):
"""
初始化答案生成器
Args:
config: 配置对象
"""
self.config = config
self.llm = LLMClient(
base_url=config.llm_base_url,
api_key=config.llm_api_key,
model=config.llm_model,
timeout=120 # 答案生成可能需要更长时间
)
async def generate(
self,
query: str,
documents: List[RankedDocument],
all_ranked_documents: Optional[List[RankedDocument]] = None
) -> Answer:
"""
根据文档生成答案
Args:
query: 用户查询
documents: 排序后的文档列表（用于生成答案）
all_ranked_documents: 所有排序后的文档（用于展示完整来源列表）
Returns:
Answer对象，包含分级的来源列表
"""
if not documents:
return self._empty_answer()
logger.info(f"开始生成答案，使用 {len(documents)} 个文档")
# 格式化文档（包含相关性分数信息）
formatted_docs = self._format_documents_with_relevance(documents)
user_message = f"""## 用户问题
{query}
## 搜索结果（按相关性排序）
{formatted_docs}"""
try:
result = await self.llm.chat_json(
system_prompt=ANSWER_GENERATION_PROMPT,
user_message=user_message,
temperature=0.5
)
# 解析答案中引用的来源
cited_sources = self._parse_cited_sources(result.get("sources", []), documents)
# 构建所有来源的分级列表
all_docs = all_ranked_documents if all_ranked_documents else documents
all_sources = self._build_sources_by_relevance(all_docs)
answer = Answer(
content=result.get("answer", ""),
sources=cited_sources,
confidence=result.get("confidence", "medium"),
all_sources=all_sources
)
logger.info(
f"答案生成完成，置信度: {answer.confidence}, "
f"来源统计: 高={len(all_sources.high)}, 中={len(all_sources.medium)}, 低={len(all_sources.low)}"
)
return answer
except Exception as e:
logger.error(f"答案生成失败: {e}")
return self._fallback_answer(query, documents)
def _format_documents_with_relevance(self, documents: List[RankedDocument]) -> str:
"""格式化文档，包含相关性信息"""
parts = []
for i, doc in enumerate(documents, 1):
relevance_level = self._get_relevance_level(doc.relevance_score)
parts.append(f"### 来源 [{i}] - 相关性: {relevance_level} ({doc.relevance_score:.2f})")
parts.append(f"**标题**: {doc.document.title}")
parts.append(f"**URL**: {doc.document.url}")
# 限制内容长度
content = doc.document.content
max_len = self.config.content_max_length // len(documents)
if len(content) > max_len:
content = content[:max_len] + "..."
parts.append(f"**内容**: {content}\n")
return "\n".join(parts)
def _get_relevance_level(self, score: float) -> str:
"""根据分数获取相关性等级"""
if score >= RELEVANCE_HIGH_THRESHOLD:
return "高"
elif score >= RELEVANCE_MEDIUM_THRESHOLD:
return "中"
else:
return "低"
def _get_relevance_level_en(self, score: float) -> str:
"""根据分数获取相关性等级（英文）"""
if score >= RELEVANCE_HIGH_THRESHOLD:
return "high"
elif score >= RELEVANCE_MEDIUM_THRESHOLD:
return "medium"
else:
return "low"
def _parse_cited_sources(
self,
source_list: List[dict],
documents: List[RankedDocument]
) -> List[Source]:
"""解析答案中引用的来源，添加相关性信息"""
# 构建URL到文档的映射
url_to_doc = {doc.document.url: doc for doc in documents}
sources = []
for i, s in enumerate(source_list):
url = s.get("url", "")
doc = url_to_doc.get(url)
relevance = "medium"
relevance_score = None
if doc:
relevance = self._get_relevance_level_en(doc.relevance_score)
relevance_score = doc.relevance_score
sources.append(Source(
index=s.get("index", i + 1),
title=s.get("title", ""),
url=url,
relevance=relevance,
relevance_score=relevance_score
))
return sources
def _build_sources_by_relevance(self, documents: List[RankedDocument]) -> SourcesByRelevance:
"""构建按相关性分组的来源列表"""
high_sources = []
medium_sources = []
low_sources = []
for i, doc in enumerate(documents, 1):
source = Source(
index=i,
title=doc.document.title,
url=doc.document.url,
relevance=self._get_relevance_level_en(doc.relevance_score),
relevance_score=doc.relevance_score
)
if doc.relevance_score >= RELEVANCE_HIGH_THRESHOLD:
high_sources.append(source)
elif doc.relevance_score >= RELEVANCE_MEDIUM_THRESHOLD:
medium_sources.append(source)
else:
low_sources.append(source)
return SourcesByRelevance(
high=high_sources,
medium=medium_sources,
low=low_sources
)
def _empty_answer(self) -> Answer:
"""生成空答案（无文档时）"""
return Answer(
content="抱歉，未能找到相关信息来回答您的问题。",
sources=[],
confidence="low",
all_sources=SourcesByRelevance(high=[], medium=[], low=[])
)
def _fallback_answer(
self,
query: str,
documents: List[RankedDocument]
) -> Answer:
"""后备答案生成（LLM失败时）"""
# 简单汇总文档内容
content_parts = [f"关于「{query}」，以下是搜索到的相关信息：\n"]
sources = []
for i, doc in enumerate(documents[:5], 1):
actual_doc = doc.document
relevance = self._get_relevance_level_en(doc.relevance_score)
content_parts.append(f"### 来源 [{i}] (相关性: {relevance}): {actual_doc.title}\n")
content_parts.append(f"{actual_doc.content[:500]}...\n\n")
sources.append(Source(
index=i,
title=actual_doc.title,
url=actual_doc.url,
relevance=relevance,
relevance_score=doc.relevance_score
))
# 构建分级来源列表
all_sources = self._build_sources_by_relevance(documents)
return Answer(
content="".join(content_parts),
sources=sources,
confidence="low",
all_sources=all_sources
)
search_agent/src/server/core/modules/content_extractor.py
+102
View File
@@ -0,0 +1,102 @@
"""
内容提取模块
使用Jina Reader提取网页内容
"""
from typing import List, Optional
from loguru import logger
from ..config import Config
from ..schemas import Document, SearchResult, SearchSource
from ..tools.jina_reader import JinaReaderClient
class ContentExtractor:
"""内容提取模块"""
def __init__(self, config: Config):
"""
初始化内容提取器
Args:
config: 配置对象
"""
self.config = config
self.jina_reader = JinaReaderClient(
api_key=config.jina_api_key,
timeout=config.timeout,
max_concurrent=config.max_concurrent_extraction, # 使用配置的并发数
max_content_length=config.content_max_length
)
async def extract(self, search_result: SearchResult) -> Optional[Document]:
"""
从搜索结果提取内容
Args:
search_result: 搜索结果
Returns:
Document对象，如果提取失败则返回None
"""
return await self.jina_reader.extract_content(
url=search_result.url,
source=search_result.source
)
async def extract_batch(
self,
search_results: List[SearchResult],
max_urls: int = 10
) -> List[Document]:
"""
批量提取内容
Args:
search_results: 搜索结果列表
max_urls: 最大提取URL数量
Returns:
Document列表
"""
# 去重并限制数量
seen_urls = set()
unique_results = []
for result in search_results:
if result.url not in seen_urls and len(unique_results) < max_urls:
seen_urls.add(result.url)
unique_results.append(result)
logger.info(f"开始提取 {len(unique_results)} 个URL的内容")
# 提取内容
urls = [r.url for r in unique_results]
# 保存source信息以便后续使用
url_to_source = {r.url: r.source for r in unique_results}
documents = await self.jina_reader.extract_batch(urls)
# 更新document的source信息
for doc in documents:
if doc.url in url_to_source:
doc.source = url_to_source[doc.url]
return documents
async def extract_urls(
self,
urls: List[str],
source: SearchSource = SearchSource.WEB
) -> List[Document]:
"""
直接从URL列表提取内容
Args:
urls: URL列表
source: 来源类型
Returns:
Document列表
"""
return await self.jina_reader.extract_batch(urls, source)
search_agent/src/server/core/modules/query_analyzer.py
+116
View File
@@ -0,0 +1,116 @@
"""
查询理解模块
负责分析用户查询意图、提取关键实体、生成扩展查询
"""
from typing import Optional
from loguru import logger
from ..config import Config
from ..schemas import QueryAnalysis, Intent
from ..utils.llm_client import LLMClient
# 查询分析Prompt
QUERY_ANALYSIS_PROMPT = """你是一个查询分析专家。分析用户的搜索查询，提取以下信息。
请输出JSON格式：
{
"intent": "查询意图，必须是以下之一: fact_check(事实核查), comparison(对比分析), how_to(操作指南), news(新闻资讯), research(深度研究)",
"entities": ["关键实体列表，提取查询中的核心概念、人名、产品名等"],
"expanded_queries": ["扩展查询1", "扩展查询2", "扩展查询3"],
"need_news": true或false,
"time_filter": "时间过滤器，null表示不限时间，qdr:d(过去24小时), qdr:w(过去一周), qdr:m(过去一月), qdr:y(过去一年)"
}
扩展查询要求：
1. 生成2-4个扩展查询，包含不同角度或同义表达
2. 至少包含一个英文查询（如果原查询是中文）
3. 保持查询的核心意图
时间过滤器选择规则：
- 查询涉及"最新"、"近期"、"今年"等时效性词语 → 设置相应的时间过滤器
- 查询涉及具体年份（如"2024年"） → qdr:y
- 一般性查询 → null"""
class QueryAnalyzer:
"""查询理解模块"""
def __init__(self, config: Config):
"""
初始化查询分析器
Args:
config: 配置对象
"""
self.config = config
self.llm = LLMClient(
base_url=config.llm_base_url,
api_key=config.llm_api_key,
model=config.llm_model
)
async def analyze(self, query: str) -> QueryAnalysis:
"""
分析用户查询
Args:
query: 用户查询字符串
Returns:
QueryAnalysis对象
"""
logger.info(f"开始分析查询: {query}")
try:
result = await self.llm.chat_json(
system_prompt=QUERY_ANALYSIS_PROMPT,
user_message=f"用户查询: {query}",
temperature=0.3
)
# 解析意图
intent_str = result.get("intent", "research")
intent = self._parse_intent(intent_str)
# 构建分析结果
analysis = QueryAnalysis(
original_query=query,
intent=intent,
entities=result.get("entities", []),
expanded_queries=result.get("expanded_queries", [query]),
need_news=result.get("need_news", False),
time_filter=result.get("time_filter")
)
logger.info(f"查询分析完成: intent={intent.value}, entities={analysis.entities}")
return analysis
except Exception as e:
logger.error(f"查询分析失败: {e}")
# 返回默认分析结果
return self._default_analysis(query)
def _parse_intent(self, intent_str: str) -> Intent:
"""解析意图字符串为枚举"""
intent_mapping = {
"fact_check": Intent.FACT_CHECK,
"comparison": Intent.COMPARISON,
"how_to": Intent.HOW_TO,
"news": Intent.NEWS,
"research": Intent.RESEARCH
}
return intent_mapping.get(intent_str.lower(), Intent.RESEARCH)
def _default_analysis(self, query: str) -> QueryAnalysis:
"""生成默认的查询分析结果"""
return QueryAnalysis(
original_query=query,
intent=Intent.RESEARCH,
entities=[],
expanded_queries=[query],
need_news=False,
time_filter=None
)
search_agent/src/server/core/modules/reflector.py
+165
View File
@@ -0,0 +1,165 @@
"""
反思迭代模块
评估答案质量，决定是否需要补充搜索
"""
from typing import List
from loguru import logger
from ..config import Config
from ..schemas import Answer, QualityAssessment
from ..utils.llm_client import LLMClient
# 反思评估Prompt
REFLECTION_PROMPT = """你是一个质量评估专家。评估以下答案是否充分回答了用户的问题。
## 评估维度
1. **完整性**: 答案是否覆盖了问题的所有方面？
2. **准确性**: 答案内容是否有明确的来源支持？
3. **深度**: 答案是否提供了足够的细节和解释？
## 输出JSON格式
{
"completeness": 0.0-1.0,
"missing_aspects": ["如果有缺失，列出缺失的方面"],
"needs_more_search": true或false,
"suggested_queries": ["如果需要补充搜索，建议的搜索词"]
}
## 判断标准
- completeness >= 0.8 且没有重要信息缺失 → needs_more_search = false
- completeness < 0.8 或有重要信息缺失 → needs_more_search = true
- 建议的搜索词应该针对缺失的方面"""
class Reflector:
"""反思迭代模块"""
# 质量阈值
COMPLETENESS_THRESHOLD = 0.8
def __init__(self, config: Config):
"""
初始化反思器
Args:
config: 配置对象
"""
self.config = config
self.llm = LLMClient(
base_url=config.llm_base_url,
api_key=config.llm_api_key,
model=config.llm_model
)
async def assess(
self,
query: str,
answer: Answer
) -> QualityAssessment:
"""
评估答案质量
Args:
query: 原始查询
answer: 生成的答案
Returns:
QualityAssessment对象
"""
logger.info("开始评估答案质量")
# 如果答案置信度已经很低，直接建议补充搜索
if answer.confidence == "low" and not answer.content:
return QualityAssessment(
completeness=0.0,
missing_aspects=["缺少相关信息"],
needs_more_search=True,
suggested_queries=[query]
)
user_message = f"""## 用户问题
{query}
## 生成的答案
{answer.content}
## 答案的来源数量
{len(answer.sources)} 个来源
## 答案的置信度
{answer.confidence}"""
try:
result = await self.llm.chat_json(
system_prompt=REFLECTION_PROMPT,
user_message=user_message,
temperature=0.3
)
assessment = QualityAssessment(
completeness=float(result.get("completeness", 0.5)),
missing_aspects=result.get("missing_aspects", []),
needs_more_search=result.get("needs_more_search", False),
suggested_queries=result.get("suggested_queries", [])
)
logger.info(
f"质量评估: completeness={assessment.completeness:.2f}, "
f"needs_more_search={assessment.needs_more_search}"
)
return assessment
except Exception as e:
logger.error(f"质量评估失败: {e}")
return self._default_assessment(answer)
def _default_assessment(self, answer: Answer) -> QualityAssessment:
"""默认评估结果"""
# 根据答案置信度估计完整性
confidence_score = {
"high": 0.9,
"medium": 0.7,
"low": 0.4
}.get(answer.confidence, 0.5)
return QualityAssessment(
completeness=confidence_score,
missing_aspects=[],
needs_more_search=confidence_score < self.COMPLETENESS_THRESHOLD,
suggested_queries=[]
)
def should_continue(
self,
assessment: QualityAssessment,
current_iteration: int
) -> bool:
"""
判断是否应该继续迭代
Args:
assessment: 质量评估结果
current_iteration: 当前迭代次数
Returns:
是否继续迭代
"""
# 达到最大迭代次数
if current_iteration >= self.config.max_iterations:
logger.info(f"达到最大迭代次数 ({self.config.max_iterations})，停止迭代")
return False
# 完整性达标
if assessment.completeness >= self.COMPLETENESS_THRESHOLD:
logger.info(f"完整性达标 ({assessment.completeness:.2f})，停止迭代")
return False
# 没有建议的补充搜索
if not assessment.suggested_queries:
logger.info("没有建议的补充搜索，停止迭代")
return False
return assessment.needs_more_search
search_agent/src/server/core/modules/result_processor.py
+133
View File
@@ -0,0 +1,133 @@
"""
结果处理模块
负责结果去重、相关性排序、筛选
优化特性：
- 支持返回所有排序后的文档（用于展示完整来源列表）
- 按相关性分数排序
"""
from typing import List, Tuple
from loguru import logger
from ..config import Config
from ..schemas import Document, RankedDocument
from ..tools.jina_reranker import JinaRerankerClient
from ..utils.helpers import deduplicate_by_url
class ResultProcessor:
"""结果处理模块"""
def __init__(self, config: Config):
"""
初始化结果处理器
Args:
config: 配置对象
"""
self.config = config
self.reranker = JinaRerankerClient(
api_key=config.jina_api_key,
timeout=config.timeout
)
async def process(
self,
query: str,
documents: List[Document],
top_k: int = 5
) -> List[RankedDocument]:
"""
处理文档：去重 + 重排序 + 筛选
Args:
query: 原始查询
documents: 文档列表
top_k: 返回前k个结果
Returns:
排序后的RankedDocument列表（仅返回top_k个）
"""
top_docs, _ = await self.process_all(query, documents, top_k)
return top_docs
async def process_all(
self,
query: str,
documents: List[Document],
top_k: int = 5
) -> Tuple[List[RankedDocument], List[RankedDocument]]:
"""
处理文档并返回所有排序结果
Args:
query: 原始查询
documents: 文档列表
top_k: 用于答案生成的前k个结果
Returns:
(top_k文档列表, 所有排序后的文档列表)
"""
if not documents:
logger.warning("没有文档需要处理")
return [], []
logger.info(f"开始处理 {len(documents)} 个文档")
# 1. 去重
unique_docs = self._deduplicate(documents)
logger.debug(f"去重后: {len(unique_docs)} 个文档")
# 2. 过滤空内容
valid_docs = [d for d in unique_docs if d.content and len(d.content.strip()) > 50]
logger.debug(f"有效文档: {len(valid_docs)} 个")
if not valid_docs:
logger.warning("没有有效文档")
return [], []
# 3. 重排序（获取所有文档的排序结果）
all_ranked_docs = await self.reranker.rerank(
query=query,
documents=valid_docs,
top_k=len(valid_docs), # 获取所有文档的排序
content_max_length=self.config.content_max_length // 5
)
# 分离 top_k 和全部结果
top_docs = all_ranked_docs[:top_k]
logger.info(f"处理完成，返回 top {len(top_docs)} 个，共 {len(all_ranked_docs)} 个排序结果")
return top_docs, all_ranked_docs
def _deduplicate(self, documents: List[Document]) -> List[Document]:
"""去重文档"""
return deduplicate_by_url(documents, "url")
async def process_without_rerank(
self,
documents: List[Document],
top_k: int = 5
) -> List[RankedDocument]:
"""
处理文档（不进行重排序）
Args:
documents: 文档列表
top_k: 返回前k个结果
Returns:
RankedDocument列表（按原始顺序）
"""
unique_docs = self._deduplicate(documents)
valid_docs = [d for d in unique_docs if d.content and len(d.content.strip()) > 50]
return [
RankedDocument(
document=doc,
relevance_score=1.0 - (i * 0.1),
rank=i + 1
)
for i, doc in enumerate(valid_docs[:top_k])
]
search_agent/src/server/core/modules/search_executor.py
+88
View File
@@ -0,0 +1,88 @@
"""
搜索执行模块
执行搜索计划，调用Serper API
"""
import asyncio
from typing import List
from loguru import logger
from ..config import Config
from ..schemas import SearchPlan, SearchTask, SearchResult
from ..tools.serper import SerperClient
class SearchExecutor:
"""搜索执行模块"""
def __init__(self, config: Config):
"""
初始化搜索执行器
Args:
config: 配置对象
"""
self.config = config
self.serper = SerperClient(
api_key=config.serper_api_key,
timeout=config.timeout
)
async def execute(self, plan: SearchPlan) -> List[SearchResult]:
"""
执行搜索计划
Args:
plan: 搜索计划
Returns:
搜索结果列表
"""
logger.info(f"开始执行搜索计划: {len(plan.tasks)} 个任务")
if plan.strategy == "parallel":
results = await self._execute_parallel(plan.tasks)
else:
results = await self._execute_sequential(plan.tasks)
logger.info(f"搜索完成，共获取 {len(results)} 条结果")
return results
async def _execute_parallel(self, tasks: List[SearchTask]) -> List[SearchResult]:
"""并行执行搜索任务"""
coroutines = [self._execute_task(task) for task in tasks]
results_list = await asyncio.gather(*coroutines, return_exceptions=True)
# 合并结果
all_results = []
for results in results_list:
if isinstance(results, list):
all_results.extend(results)
elif isinstance(results, Exception):
logger.warning(f"搜索任务失败: {results}")
return all_results
async def _execute_sequential(self, tasks: List[SearchTask]) -> List[SearchResult]:
"""串行执行搜索任务"""
all_results = []
for task in tasks:
try:
results = await self._execute_task(task)
all_results.extend(results)
except Exception as e:
logger.warning(f"搜索任务失败: {e}")
return all_results
async def _execute_task(self, task: SearchTask) -> List[SearchResult]:
"""执行单个搜索任务"""
logger.debug(f"执行搜索: {task.query} [{task.source.value}]")
return await self.serper.search(
query=task.query,
source=task.source,
num_results=task.num_results,
time_filter=task.time_filter
)
search_agent/src/server/core/modules/search_planner.py
+136
View File
@@ -0,0 +1,136 @@
"""
搜索规划模块
根据查询分析结果制定搜索计划
"""
from typing import List
from loguru import logger
from ..config import Config
from ..schemas import (
QueryAnalysis,
SearchPlan,
SearchTask,
SearchSource,
Intent
)
class SearchPlanner:
"""搜索规划模块"""
def __init__(self, config: Config):
"""
初始化搜索规划器
Args:
config: 配置对象
"""
self.config = config
self.max_results = config.max_results_per_query
async def plan(self, analysis: QueryAnalysis) -> SearchPlan:
"""
根据查询分析制定搜索计划
Args:
analysis: 查询分析结果
Returns:
SearchPlan对象
"""
logger.info(f"开始制定搜索计划: intent={analysis.intent.value}")
tasks = []
# 根据意图确定搜索策略
strategy = self._determine_strategy(analysis)
# 构建搜索任务
tasks.extend(self._create_web_tasks(analysis))
if analysis.need_news:
tasks.extend(self._create_news_tasks(analysis))
plan = SearchPlan(
tasks=tasks,
strategy=strategy
)
logger.info(f"搜索计划: {len(tasks)} 个任务, 策略={strategy}")
return plan
def _determine_strategy(self, analysis: QueryAnalysis) -> str:
"""确定执行策略"""
# 大多数情况使用并行策略
if analysis.intent == Intent.COMPARISON:
# 对比类查询可能需要串行以获取更相关的结果
return "parallel"
return "parallel"
def _create_web_tasks(self, analysis: QueryAnalysis) -> List[SearchTask]:
"""创建Web搜索任务"""
tasks = []
# 原始查询
tasks.append(SearchTask(
query=analysis.original_query,
source=SearchSource.WEB,
time_filter=analysis.time_filter,
num_results=self.max_results
))
# 扩展查询（限制数量避免过多请求）
for query in analysis.expanded_queries[:2]:
if query != analysis.original_query:
tasks.append(SearchTask(
query=query,
source=SearchSource.WEB,
time_filter=analysis.time_filter,
num_results=self.max_results
))
return tasks
def _create_news_tasks(self, analysis: QueryAnalysis) -> List[SearchTask]:
"""创建新闻搜索任务"""
tasks = []
# 新闻搜索使用原始查询
tasks.append(SearchTask(
query=analysis.original_query,
source=SearchSource.NEWS,
time_filter=analysis.time_filter or "qdr:m", # 默认过去一个月
num_results=self.max_results
))
return tasks
def plan_supplementary(
self,
original_query: str,
suggested_queries: List[str]
) -> SearchPlan:
"""
创建补充搜索计划
Args:
original_query: 原始查询
suggested_queries: 建议的补充查询
Returns:
SearchPlan对象
"""
tasks = []
for query in suggested_queries[:3]: # 限制补充搜索数量
tasks.append(SearchTask(
query=query,
source=SearchSource.WEB,
num_results=self.max_results
))
return SearchPlan(
tasks=tasks,
strategy="parallel"
)
search_agent/src/server/core/prompts.py
+125
View File
@@ -0,0 +1,125 @@
"""
Prompt模板汇总
集中管理所有LLM Prompt模板
"""
# ==================== 查询分析 Prompt ====================
QUERY_ANALYSIS_PROMPT = """你是一个查询分析专家。分析用户的搜索查询，提取以下信息。
请输出JSON格式：
{
"intent": "查询意图，必须是以下之一: fact_check(事实核查), comparison(对比分析), how_to(操作指南), news(新闻资讯), research(深度研究)",
"entities": ["关键实体列表，提取查询中的核心概念、人名、产品名等"],
"expanded_queries": ["扩展查询1", "扩展查询2", "扩展查询3"],
"need_news": true或false,
"time_filter": "时间过滤器，null表示不限时间，qdr:d(过去24小时), qdr:w(过去一周), qdr:m(过去一月), qdr:y(过去一年)"
}
扩展查询要求：
1. 生成2-4个扩展查询，包含不同角度或同义表达
2. 至少包含一个英文查询（如果原查询是中文）
3. 保持查询的核心意图
时间过滤器选择规则：
- 查询涉及"最新"、"近期"、"今年"等时效性词语 → 设置相应的时间过滤器
- 查询涉及具体年份（如"2024年"） → qdr:y
- 一般性查询 → null"""
# ==================== 搜索规划 Prompt ====================
SEARCH_PLANNING_PROMPT = """你是一个搜索规划专家。根据查询分析结果，制定搜索计划。
输入信息：
- 原始查询
- 查询意图
- 关键实体
- 是否需要新闻
输出搜索任务列表，每个任务包含：
- query: 搜索词
- source: web 或 news
- time_filter: 时间过滤器（可选）
搜索策略规则：
1. 简单事实查询 → 单次Web搜索
2. 时效性查询 → Web搜索 + 新闻搜索
3. 复杂分析查询 → 多个扩展查询
4. 对比类查询 → 分别搜索各对比对象"""
# ==================== 答案生成 Prompt ====================
ANSWER_GENERATION_PROMPT = """你是一个专业的信息整合专家。根据以下搜索结果，回答用户的问题。
## 要求
1. 综合多个来源的信息，给出全面准确的回答
2. 使用清晰的结构组织答案（标题、列表、重点标注等）
3. 在答案中标注信息来源，格式：[来源1]、[来源2]
4. 如果信息有冲突，说明不同观点
5. 如果信息不足以完整回答问题，明确指出缺失的部分
6. 回答使用中文
## 输出JSON格式
{
"answer": "结构化的答案（Markdown格式，包含来源引用）",
"sources": [
{"index": 1, "title": "来源标题", "url": "来源URL"},
{"index": 2, "title": "来源标题", "url": "来源URL"}
],
"confidence": "high/medium/low，基于信息质量和一致性判断"
}"""
# ==================== 反思评估 Prompt ====================
REFLECTION_PROMPT = """你是一个质量评估专家。评估以下答案是否充分回答了用户的问题。
## 评估维度
1. **完整性**: 答案是否覆盖了问题的所有方面？
2. **准确性**: 答案内容是否有明确的来源支持？
3. **深度**: 答案是否提供了足够的细节和解释？
## 输出JSON格式
{
"completeness": 0.0-1.0,
"missing_aspects": ["如果有缺失，列出缺失的方面"],
"needs_more_search": true或false,
"suggested_queries": ["如果需要补充搜索，建议的搜索词"]
}
## 判断标准
- completeness >= 0.8 且没有重要信息缺失 → needs_more_search = false
- completeness < 0.8 或有重要信息缺失 → needs_more_search = true
- 建议的搜索词应该针对缺失的方面"""
# ==================== 工具函数 ====================
def format_query_analysis_prompt(query: str) -> str:
"""格式化查询分析Prompt"""
return f"{QUERY_ANALYSIS_PROMPT}\n\n用户查询: {query}"
def format_answer_generation_prompt(query: str, documents: str) -> str:
"""格式化答案生成Prompt"""
return f"""{ANSWER_GENERATION_PROMPT}
## 用户问题
{query}
## 搜索结果
{documents}"""
def format_reflection_prompt(query: str, answer: str, sources_count: int, confidence: str) -> str:
"""格式化反思评估Prompt"""
return f"""{REFLECTION_PROMPT}
## 用户问题
{query}
## 生成的答案
{answer}
## 答案的来源数量
{sources_count} 个来源
## 答案的置信度
{confidence}"""
search_agent/src/server/core/schemas.py
+239
View File
@@ -0,0 +1,239 @@
"""
数据模型定义
定义Agent使用的所有数据结构
"""
from dataclasses import dataclass, field
from typing import List, Optional
from enum import Enum
class SearchSource(Enum):
"""搜索来源枚举"""
WEB = "web"
NEWS = "news"
class Intent(Enum):
"""查询意图枚举"""
FACT_CHECK = "fact_check" # 事实核查
COMPARISON = "comparison" # 对比分析
HOW_TO = "how_to" # 操作指南
NEWS = "news" # 新闻资讯
RESEARCH = "research" # 深度研究
@dataclass
class QueryAnalysis:
"""查询分析结果"""
original_query: str # 原始查询
intent: Intent # 查询意图
entities: List[str] # 关键实体
expanded_queries: List[str] # 扩展查询列表
need_news: bool # 是否需要新闻搜索
time_filter: Optional[str] = None # 时间过滤器
def to_dict(self) -> dict:
"""转换为字典"""
return {
"original_query": self.original_query,
"intent": self.intent.value,
"entities": self.entities,
"expanded_queries": self.expanded_queries,
"need_news": self.need_news,
"time_filter": self.time_filter
}
@dataclass
class SearchTask:
"""搜索任务"""
query: str # 搜索查询
source: SearchSource # 搜索来源
time_filter: Optional[str] = None # 时间过滤器
num_results: int = 10 # 结果数量
def to_dict(self) -> dict:
"""转换为字典"""
return {
"query": self.query,
"source": self.source.value,
"time_filter": self.time_filter,
"num_results": self.num_results
}
@dataclass
class SearchPlan:
"""搜索计划"""
tasks: List[SearchTask] # 搜索任务列表
strategy: str = "parallel" # 执行策略: parallel/sequential
def to_dict(self) -> dict:
"""转换为字典"""
return {
"tasks": [t.to_dict() for t in self.tasks],
"strategy": self.strategy
}
@dataclass
class SearchResult:
"""搜索结果"""
title: str # 标题
url: str # URL
snippet: str # 摘要
source: SearchSource # 来源类型
position: int # 排名位置
date: Optional[str] = None # 日期（新闻）
def to_dict(self) -> dict:
"""转换为字典"""
return {
"title": self.title,
"url": self.url,
"snippet": self.snippet,
"source": self.source.value,
"position": self.position,
"date": self.date
}
@dataclass
class Document:
"""提取的文档内容"""
url: str # URL
title: str # 标题
content: str # 内容
source: SearchSource # 来源类型
def to_dict(self) -> dict:
"""转换为字典"""
return {
"url": self.url,
"title": self.title,
"content": self.content,
"source": self.source.value
}
@dataclass
class RankedDocument:
"""排序后的文档"""
document: Document # 文档
relevance_score: float # 相关性分数
rank: int # 排名
def to_dict(self) -> dict:
"""转换为字典"""
return {
"document": self.document.to_dict(),
"relevance_score": self.relevance_score,
"rank": self.rank
}
class SourceRelevance(Enum):
"""来源相关性等级"""
HIGH = "high" # 高相关性（相关性分数 >= 0.7）
MEDIUM = "medium" # 中相关性（0.4 <= 相关性分数 < 0.7）
LOW = "low" # 低相关性（相关性分数 < 0.4）
@dataclass
class Source:
"""来源引用"""
index: int # 索引
title: str # 标题
url: str # URL
relevance: str = "medium" # 相关性等级: high/medium/low
relevance_score: Optional[float] = None # 相关性分数 0-1
def to_dict(self) -> dict:
"""转换为字典"""
result = {
"index": self.index,
"title": self.title,
"url": self.url,
"relevance": self.relevance
}
if self.relevance_score is not None:
result["relevance_score"] = round(self.relevance_score, 3)
return result
@dataclass
class SourcesByRelevance:
"""按相关性分组的来源"""
high: List[Source] # 高相关性来源
medium: List[Source] # 中相关性来源
low: List[Source] # 低相关性来源
def to_dict(self) -> dict:
"""转换为字典"""
return {
"high": [s.to_dict() for s in self.high],
"medium": [s.to_dict() for s in self.medium],
"low": [s.to_dict() for s in self.low]
}
@property
def total_count(self) -> int:
"""总来源数量"""
return len(self.high) + len(self.medium) + len(self.low)
@dataclass
class Answer:
"""生成的答案"""
content: str # Markdown格式的答案内容
sources: List[Source] # 答案中引用的来源列表
confidence: str # 置信度: high/medium/low
all_sources: Optional[SourcesByRelevance] = None # 所有搜索来源（按相关性分组）
def to_dict(self) -> dict:
"""转换为字典"""
result = {
"content": self.content,
"sources": [s.to_dict() for s in self.sources],
"confidence": self.confidence
}
if self.all_sources:
result["all_sources"] = self.all_sources.to_dict()
return result
@dataclass
class QualityAssessment:
"""质量评估"""
completeness: float # 完整性 0-1
missing_aspects: List[str] # 缺失的方面
needs_more_search: bool # 是否需要更多搜索
suggested_queries: List[str] # 建议的补充搜索
def to_dict(self) -> dict:
"""转换为字典"""
return {
"completeness": self.completeness,
"missing_aspects": self.missing_aspects,
"needs_more_search": self.needs_more_search,
"suggested_queries": self.suggested_queries
}
@dataclass
class AgentResponse:
"""Agent最终响应"""
answer: Answer # 答案
iterations: int # 迭代次数
total_sources_consulted: int # 参考来源总数
search_queries_used: List[str] # 使用的搜索查询
def to_dict(self) -> dict:
"""转换为字典"""
return {
"answer": self.answer.to_dict(),
"iterations": self.iterations,
"total_sources_consulted": self.total_sources_consulted,
"search_queries_used": self.search_queries_used
}
search_agent/src/server/core/tools/__init__.py
+13
View File
@@ -0,0 +1,13 @@
"""
外部API工具封装模块
"""
from .serper import SerperClient
from .jina_reader import JinaReaderClient
from .jina_reranker import JinaRerankerClient
__all__ = [
"SerperClient",
"JinaReaderClient",
"JinaRerankerClient",
]
search_agent/src/server/core/tools/jina_reader.py
+179
View File
@@ -0,0 +1,179 @@
"""
Jina Reader API封装
提供网页内容提取功能
"""
import asyncio
from typing import List, Optional
import aiohttp
from loguru import logger
from ..schemas import Document, SearchSource
class JinaReaderClient:
"""Jina Reader API客户端"""
BASE_URL = "https://r.jina.ai"
def __init__(
self,
api_key: str,
timeout: int = 30,
max_concurrent: int = 10,
max_content_length: int = 5000
):
"""
初始化Jina Reader客户端
Args:
api_key: Jina API密钥
timeout: 请求超时时间（秒）
max_concurrent: 最大并发请求数（默认10，提升并行效率）
max_content_length: 最大内容长度
"""
self.api_key = api_key
self.timeout = timeout
self.max_concurrent = max_concurrent
self.max_content_length = max_content_length
self._semaphore = asyncio.Semaphore(max_concurrent)
async def extract_content(
self,
url: str,
source: SearchSource = SearchSource.WEB
) -> Optional[Document]:
"""
提取单个URL的内容
Args:
url: 要提取的网页URL
source: 来源类型
Returns:
Document对象，如果提取失败则返回None
"""
reader_url = f"{self.BASE_URL}/{url}"
headers = {
"Authorization": f"Bearer {self.api_key}",
"Accept": "application/json"
}
try:
async with self._semaphore:
async with aiohttp.ClientSession() as session:
async with session.get(
reader_url,
headers=headers,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
logger.warning(f"Jina Reader提取失败 [{response.status}]: {url}")
return None
# Jina Reader可能返回JSON或纯文本
content_type = response.headers.get("Content-Type", "")
if "application/json" in content_type:
result = await response.json()
# 处理嵌套的data字段
if "data" in result:
result = result["data"]
content = result.get("content", "")
title = result.get("title", "")
else:
# 纯文本响应（Markdown格式）
content = await response.text()
# 从内容中提取标题（第一行通常是标题）
lines = content.strip().split("\n")
title = lines[0].lstrip("#").strip() if lines else ""
# 限制内容长度
if len(content) > self.max_content_length:
content = content[:self.max_content_length]
logger.debug(f"提取成功: {url[:50]}... 内容长度: {len(content)}")
return Document(
url=url,
title=title,
content=content,
source=source
)
except aiohttp.ClientError as e:
logger.warning(f"Jina Reader网络错误 [{url}]: {e}")
return None
except asyncio.TimeoutError:
logger.warning(f"Jina Reader超时: {url}")
return None
except Exception as e:
logger.warning(f"Jina Reader异常 [{url}]: {e}")
return None
async def extract_batch(
self,
urls: List[str],
source: SearchSource = SearchSource.WEB
) -> List[Document]:
"""
批量提取多个URL的内容
Args:
urls: URL列表
source: 来源类型
Returns:
成功提取的Document列表
"""
logger.info(f"批量提取 {len(urls)} 个URL的内容")
tasks = [
self.extract_content(url, source)
for url in urls
]
results = await asyncio.gather(*tasks, return_exceptions=True)
# 过滤掉失败的结果
documents = []
for result in results:
if isinstance(result, Document):
documents.append(result)
elif isinstance(result, Exception):
logger.warning(f"提取异常: {result}")
logger.info(f"成功提取 {len(documents)}/{len(urls)} 个文档")
return documents
async def extract_with_retry(
self,
url: str,
source: SearchSource = SearchSource.WEB,
max_retries: int = 2,
retry_delay: float = 1.0
) -> Optional[Document]:
"""
带重试的内容提取
Args:
url: 要提取的网页URL
source: 来源类型
max_retries: 最大重试次数
retry_delay: 重试延迟（秒）
Returns:
Document对象，如果最终失败则返回None
"""
for attempt in range(max_retries + 1):
result = await self.extract_content(url, source)
if result is not None:
return result
if attempt < max_retries:
logger.debug(f"重试提取 [{attempt + 1}/{max_retries}]: {url}")
await asyncio.sleep(retry_delay)
return None
search_agent/src/server/core/tools/jina_reranker.py
+190
View File
@@ -0,0 +1,190 @@
"""
Jina Reranker API封装
提供搜索结果重排序功能
"""
from typing import List, Tuple
import aiohttp
from loguru import logger
from ..schemas import Document, RankedDocument
class JinaRerankerClient:
"""Jina Reranker API客户端"""
BASE_URL = "https://api.jina.ai/v1/rerank"
MODEL = "jina-reranker-v2-base-multilingual"
def __init__(self, api_key: str, timeout: int = 30):
"""
初始化Jina Reranker客户端
Args:
api_key: Jina API密钥
timeout: 请求超时时间（秒）
"""
self.api_key = api_key
self.timeout = timeout
async def rerank(
self,
query: str,
documents: List[Document],
top_k: int = 5,
content_max_length: int = 1000
) -> List[RankedDocument]:
"""
对文档进行相关性重排序
Args:
query: 查询字符串
documents: 文档列表
top_k: 返回前k个结果
content_max_length: 用于排序的内容最大长度
Returns:
排序后的RankedDocument列表
"""
if not documents:
return []
# 准备文档内容（截断到合适长度）
doc_contents = [
doc.content[:content_max_length] if doc.content else doc.title
for doc in documents
]
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
payload = {
"model": self.MODEL,
"query": query,
"documents": doc_contents,
"top_n": min(top_k, len(documents))
}
try:
async with aiohttp.ClientSession() as session:
async with session.post(
self.BASE_URL,
headers=headers,
json=payload,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
error_text = await response.text()
logger.error(f"Jina Reranker API错误: {response.status} - {error_text}")
# 如果重排序失败，返回原始顺序
return self._fallback_ranking(documents, top_k)
result = await response.json()
return self._parse_rerank_results(documents, result, top_k)
except aiohttp.ClientError as e:
logger.error(f"Jina Reranker网络错误: {e}")
return self._fallback_ranking(documents, top_k)
except Exception as e:
logger.error(f"Jina Reranker异常: {e}")
return self._fallback_ranking(documents, top_k)
def _parse_rerank_results(
self,
documents: List[Document],
response: dict,
top_k: int
) -> List[RankedDocument]:
"""解析重排序结果"""
results = []
reranked = response.get("results", [])
for rank, item in enumerate(reranked[:top_k], 1):
index = item.get("index", 0)
score = item.get("relevance_score", 0.0)
if 0 <= index < len(documents):
ranked_doc = RankedDocument(
document=documents[index],
relevance_score=score,
rank=rank
)
results.append(ranked_doc)
logger.debug(f"重排序返回 {len(results)} 个结果")
return results
def _fallback_ranking(
self,
documents: List[Document],
top_k: int
) -> List[RankedDocument]:
"""后备排序：保持原始顺序"""
logger.warning("使用后备排序（原始顺序）")
return [
RankedDocument(
document=doc,
relevance_score=1.0 - (i * 0.1), # 模拟递减分数
rank=i + 1
)
for i, doc in enumerate(documents[:top_k])
]
async def rerank_texts(
self,
query: str,
texts: List[str],
top_k: int = 5
) -> List[Tuple[int, float]]:
"""
对纯文本列表进行重排序
Args:
query: 查询字符串
texts: 文本列表
top_k: 返回前k个结果
Returns:
(原始索引, 相关性分数) 的列表
"""
if not texts:
return []
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
payload = {
"model": self.MODEL,
"query": query,
"documents": texts,
"top_n": min(top_k, len(texts))
}
try:
async with aiohttp.ClientSession() as session:
async with session.post(
self.BASE_URL,
headers=headers,
json=payload,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
logger.error(f"Reranker API错误: {response.status}")
return [(i, 1.0 - i * 0.1) for i in range(min(top_k, len(texts)))]
result = await response.json()
return [
(item["index"], item["relevance_score"])
for item in result.get("results", [])[:top_k]
]
except Exception as e:
logger.error(f"Reranker异常: {e}")
return [(i, 1.0 - i * 0.1) for i in range(min(top_k, len(texts)))]
search_agent/src/server/core/tools/serper.py
+211
View File
@@ -0,0 +1,211 @@
"""
Serper API封装
提供Google搜索和新闻搜索功能
"""
from typing import List, Optional, Dict, Any
import aiohttp
from loguru import logger
from ..schemas import SearchResult, SearchSource
class SerperClient:
"""Serper API客户端"""
BASE_URL = "https://google.serper.dev"
ENDPOINTS = {
"web": "/search",
"news": "/news"
}
def __init__(self, api_key: str, timeout: int = 30):
"""
初始化Serper客户端
Args:
api_key: Serper API密钥
timeout: 请求超时时间（秒）
"""
self.api_key = api_key
self.timeout = timeout
async def _request(
self,
endpoint: str,
payload: Dict[str, Any]
) -> Dict[str, Any]:
"""
发送请求到Serper API
Args:
endpoint: API端点
payload: 请求体
Returns:
API响应
"""
url = f"{self.BASE_URL}{endpoint}"
headers = {
"X-API-KEY": self.api_key,
"Content-Type": "application/json"
}
try:
async with aiohttp.ClientSession() as session:
async with session.post(
url,
headers=headers,
json=payload,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
error_text = await response.text()
logger.error(f"Serper API错误: {response.status} - {error_text}")
raise Exception(f"Serper API请求失败: {response.status}")
return await response.json()
except aiohttp.ClientError as e:
logger.error(f"Serper请求网络错误: {e}")
raise
async def search_web(
self,
query: str,
num_results: int = 10,
gl: str = "cn",
hl: str = "zh-cn",
time_filter: Optional[str] = None
) -> List[SearchResult]:
"""
执行Web搜索
Args:
query: 搜索查询
num_results: 返回结果数量
gl: 地区代码
hl: 语言代码
time_filter: 时间过滤器 (qdr:d/qdr:w/qdr:m/qdr:y)
Returns:
搜索结果列表
"""
payload = {
"q": query,
"num": num_results,
"gl": gl,
"hl": hl
}
if time_filter:
payload["tbs"] = time_filter
logger.info(f"执行Web搜索: {query}")
result = await self._request(self.ENDPOINTS["web"], payload)
return self._parse_web_results(result)
async def search_news(
self,
query: str,
num_results: int = 10,
gl: str = "cn",
hl: str = "zh-cn",
time_filter: Optional[str] = None
) -> List[SearchResult]:
"""
执行新闻搜索
Args:
query: 搜索查询
num_results: 返回结果数量
gl: 地区代码
hl: 语言代码
time_filter: 时间过滤器
Returns:
搜索结果列表
"""
payload = {
"q": query,
"num": num_results,
"gl": gl,
"hl": hl
}
if time_filter:
payload["tbs"] = time_filter
logger.info(f"执行新闻搜索: {query}")
result = await self._request(self.ENDPOINTS["news"], payload)
return self._parse_news_results(result)
def _parse_web_results(self, response: Dict[str, Any]) -> List[SearchResult]:
"""解析Web搜索结果"""
results = []
organic = response.get("organic", [])
for item in organic:
result = SearchResult(
title=item.get("title", ""),
url=item.get("link", ""),
snippet=item.get("snippet", ""),
source=SearchSource.WEB,
position=item.get("position", 0),
date=None
)
results.append(result)
logger.debug(f"Web搜索返回 {len(results)} 条结果")
return results
def _parse_news_results(self, response: Dict[str, Any]) -> List[SearchResult]:
"""解析新闻搜索结果"""
results = []
news = response.get("news", [])
for i, item in enumerate(news, 1):
result = SearchResult(
title=item.get("title", ""),
url=item.get("link", ""),
snippet=item.get("snippet", ""),
source=SearchSource.NEWS,
position=i,
date=item.get("date")
)
results.append(result)
logger.debug(f"新闻搜索返回 {len(results)} 条结果")
return results
async def search(
self,
query: str,
source: SearchSource,
num_results: int = 10,
time_filter: Optional[str] = None
) -> List[SearchResult]:
"""
统一搜索接口
Args:
query: 搜索查询
source: 搜索来源类型
num_results: 返回结果数量
time_filter: 时间过滤器
Returns:
搜索结果列表
"""
if source == SearchSource.NEWS:
return await self.search_news(query, num_results, time_filter=time_filter)
else:
return await self.search_web(query, num_results, time_filter=time_filter)
search_agent/src/server/core/utils/__init__.py
+21
View File
@@ -0,0 +1,21 @@
"""
工具函数模块
"""
from .llm_client import LLMClient
from .helpers import (
flatten,
deduplicate_by_url,
truncate_text,
extract_json_from_text,
format_documents_for_prompt,
)
__all__ = [
"LLMClient",
"flatten",
"deduplicate_by_url",
"truncate_text",
"extract_json_from_text",
"format_documents_for_prompt",
]
search_agent/src/server/core/utils/__pycache__/helpers.cpython-312.pyc
BIN
View File
Binary file not shown.
search_agent/src/server/core/utils/helpers.py
+196
View File
@@ -0,0 +1,196 @@
"""
通用工具函数
"""
import re
import json
from typing import List, TypeVar, Optional, Dict, Any
T = TypeVar('T')
def flatten(nested_list: List[List[T]]) -> List[T]:
"""
将嵌套列表展平为一维列表
Args:
nested_list: 嵌套列表
Returns:
展平后的一维列表
"""
return [item for sublist in nested_list for item in sublist]
def deduplicate_by_url(items: List[Any], url_attr: str = "url") -> List[Any]:
"""
根据URL去重
Args:
items: 包含URL属性的对象列表
url_attr: URL属性名
Returns:
去重后的列表
"""
seen_urls = set()
unique_items = []
for item in items:
url = getattr(item, url_attr, None) or item.get(url_attr)
if url and url not in seen_urls:
seen_urls.add(url)
unique_items.append(item)
return unique_items
def truncate_text(text: str, max_length: int, suffix: str = "...") -> str:
"""
截断文本到指定长度
Args:
text: 原始文本
max_length: 最大长度
suffix: 截断后缀
Returns:
截断后的文本
"""
if len(text) <= max_length:
return text
return text[:max_length - len(suffix)] + suffix
def extract_json_from_text(text: str) -> Optional[Dict[str, Any]]:
"""
从文本中提取JSON对象
Args:
text: 可能包含JSON的文本
Returns:
提取的JSON字典，如果提取失败则返回None
"""
# 尝试直接解析
try:
return json.loads(text)
except json.JSONDecodeError:
pass
# 尝试提取```json ... ```块
json_block_pattern = r'```(?:json)?\s*([\s\S]*?)```'
matches = re.findall(json_block_pattern, text)
for match in matches:
try:
return json.loads(match.strip())
except json.JSONDecodeError:
continue
# 尝试提取{ ... }块
brace_pattern = r'\{[\s\S]*\}'
matches = re.findall(brace_pattern, text)
for match in matches:
try:
return json.loads(match)
except json.JSONDecodeError:
continue
return None
def format_documents_for_prompt(documents: List[Any], max_length: int = 2000) -> str:
"""
格式化文档列表为Prompt中使用的文本
Args:
documents: 文档列表（RankedDocument或Document对象）
max_length: 每个文档的最大内容长度
Returns:
格式化后的文本
"""
formatted_parts = []
for i, doc in enumerate(documents, 1):
# 支持RankedDocument和Document两种类型
if hasattr(doc, 'document'):
# RankedDocument
actual_doc = doc.document
score = f" (相关性: {doc.relevance_score:.2f})"
else:
# Document
actual_doc = doc
score = ""
content = truncate_text(actual_doc.content, max_length)
part = f"""### 来源 [{i}]{score}
**标题**: {actual_doc.title}
**URL**: {actual_doc.url}
**内容**:
{content}
"""
formatted_parts.append(part)
return "\n---\n".join(formatted_parts)
def clean_url(url: str) -> str:
"""
清理和标准化URL
Args:
url: 原始URL
Returns:
清理后的URL
"""
# 移除末尾的斜杠
url = url.rstrip("/")
# 移除锚点
if "#" in url:
url = url.split("#")[0]
return url
def is_valid_url(url: str) -> bool:
"""
验证URL是否有效
Args:
url: URL字符串
Returns:
是否有效
"""
url_pattern = re.compile(
r'^https?://' # http:// or https://
r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|' # domain
r'localhost|' # localhost
r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})' # IP
r'(?::\d+)?' # optional port
r'(?:/?|[/?]\S+)$', re.IGNORECASE)
return bool(url_pattern.match(url))
def merge_dicts(base: Dict, override: Dict) -> Dict:
"""
合并两个字典，override中的值会覆盖base中的值
Args:
base: 基础字典
override: 覆盖字典
Returns:
合并后的字典
"""
result = base.copy()
result.update(override)
return result
search_agent/src/server/core/utils/llm_client.py
+158
View File
@@ -0,0 +1,158 @@
"""
LLM客户端模块
封装与LLM的交互（支持Azure OpenAI风格API）
"""
import json
from typing import Optional, List, Dict, Any
import aiohttp
from loguru import logger
class LLMClient:
"""LLM客户端，用于与LLM API交互"""
# API版本
API_VERSION = "2024-10-21"
def __init__(
self,
base_url: str,
api_key: str,
model: str = "taiji/gpt-4o-mini",
timeout: int = 60
):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
self.timeout = timeout
async def chat(
self,
messages: List[Dict[str, str]],
temperature: float = 0.7,
max_tokens: int = 4096,
response_format: Optional[Dict[str, str]] = None
) -> str:
"""
发送聊天请求到LLM
Args:
messages: 消息列表，格式 [{"role": "user", "content": "..."}]
temperature: 温度参数
max_tokens: 最大token数
response_format: 响应格式（如 {"type": "json_object"}）
Returns:
LLM的响应文本
"""
# Azure OpenAI 风格的URL
url = f"{self.base_url}/chat/completions?api-version={self.API_VERSION}"
# Azure OpenAI 使用 api-key 头
headers = {
"api-key": self.api_key,
"Content-Type": "application/json"
}
payload = {
"model": self.model,
"messages": messages,
"temperature": temperature,
"max_completion_tokens": max_tokens # 新版API使用 max_completion_tokens
}
if response_format:
payload["response_format"] = response_format
try:
async with aiohttp.ClientSession() as session:
async with session.post(
url,
headers=headers,
json=payload,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
error_text = await response.text()
logger.error(f"LLM API错误: {response.status} - {error_text}")
raise Exception(f"LLM API请求失败: {response.status}")
result = await response.json()
return result["choices"][0]["message"]["content"]
except aiohttp.ClientError as e:
logger.error(f"LLM请求网络错误: {e}")
raise
except Exception as e:
logger.error(f"LLM请求异常: {e}")
raise
async def chat_with_system(
self,
system_prompt: str,
user_message: str,
temperature: float = 0.7,
max_tokens: int = 4096,
response_format: Optional[Dict[str, str]] = None
) -> str:
"""
使用系统提示和用户消息进行对话
Args:
system_prompt: 系统提示
user_message: 用户消息
temperature: 温度参数
max_tokens: 最大token数
response_format: 响应格式
Returns:
LLM的响应文本
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_message}
]
return await self.chat(
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
response_format=response_format
)
async def chat_json(
self,
system_prompt: str,
user_message: str,
temperature: float = 0.3
) -> Dict[str, Any]:
"""
请求JSON格式的响应
Args:
system_prompt: 系统提示
user_message: 用户消息
temperature: 温度参数（JSON响应建议使用较低温度）
Returns:
解析后的JSON字典
"""
from .helpers import extract_json_from_text
response = await self.chat_with_system(
system_prompt=system_prompt,
user_message=user_message,
temperature=temperature,
response_format={"type": "json_object"}
)
try:
return json.loads(response)
except json.JSONDecodeError:
# 尝试从文本中提取JSON
extracted = extract_json_from_text(response)
if extracted:
return extracted
logger.error(f"无法解析LLM响应为JSON: {response[:200]}")
raise ValueError("LLM响应不是有效的JSON格式")
search_agent/src/server/mcp_server.py
+198
View File
@@ -0,0 +1,198 @@
"""
Search Agent MCP 服务器
提供智能搜索相关的 MCP 工具。
"""
import os
import json
from typing import Optional
from mcp.server.fastmcp import FastMCP
from .core.config import Config
from .core.agent import SearchAgent
# ==================== 配置 ====================
# LiteLLM Gateway 配置
_BASE_URL = os.getenv('OPENAI_BASE_URL',
os.getenv('LLM_BASE_URL', 'https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1'))
_API_KEY = os.getenv('OPENAI_API_KEY', 'sk')
os.environ.setdefault('OPENAI_API_KEY', _API_KEY)
os.environ.setdefault('OPENAI_BASE_URL', _BASE_URL)
# ==================== MCP 服务器 ====================
server = FastMCP('Search Agent')
# 系统提示词
SYSTEM_PROMPT = '''你是一个智能搜索助手。
能够理解用户查询意图、自动规划搜索策略、从多个来源获取信息，
并生成高质量、有来源引用的答案。
核心能力：
1. 查询理解 - 分析用户意图，提取关键实体，生成扩展查询
2. 多源搜索 - 支持Web搜索和新闻搜索
3. 内容提取 - 智能提取网页核心内容
4. 结果排序 - 基于相关性重排搜索结果
5. 答案生成 - 综合信息生成结构化回答
6. 自我反思 - 评估答案质量，决定是否迭代
'''
# 全局 Agent 实例
_agent: Optional[SearchAgent] = None
def get_agent() -> SearchAgent:
"""获取或创建 Agent 实例"""
global _agent
if _agent is None:
config = Config.from_env()
_agent = SearchAgent(config)
return _agent
def reset_agent():
"""重置 Agent 实例（配置变更时调用）"""
global _agent
_agent = None
# ==================== MCP 工具定义 ====================
@server.tool()
async def search(
query: str,
max_iterations: int = 3
) -> str:
"""
执行智能搜索
支持多轮迭代和自我反思，能够：
- 分析查询意图
- 自动规划搜索策略
- 从多个来源获取信息
- 生成高质量答案
Args:
query: 搜索查询问题
max_iterations: 最大迭代次数（1-10），默认3次
Returns:
JSON 格式的搜索结果，包含答案、来源和统计信息
"""
try:
agent = get_agent()
# 临时修改迭代次数
original = agent.config.max_iterations
agent.config.max_iterations = min(max(1, max_iterations), 10)
try:
response = await agent.search(query)
finally:
agent.config.max_iterations = original
return json.dumps({
"success": True,
"answer": response.answer.to_dict(),
"statistics": {
"iterations": response.iterations,
"total_sources_consulted": response.total_sources_consulted,
"search_queries_used": response.search_queries_used
}
}, ensure_ascii=False, indent=2)
except Exception as e:
return json.dumps({
"success": False,
"error": str(e)
}, ensure_ascii=False)
@server.tool()
async def quick_search(query: str) -> str:
"""
快速搜索（单次迭代）
适合简单问题，执行速度更快：
- 单次搜索迭代
- 限制搜索任务数量
- 快速生成答案
Args:
query: 搜索查询问题
Returns:
JSON 格式的搜索结果
"""
try:
agent = get_agent()
answer = await agent.quick_search(query)
return json.dumps({
"success": True,
"answer": answer.to_dict()
}, ensure_ascii=False, indent=2)
except Exception as e:
return json.dumps({
"success": False,
"error": str(e)
}, ensure_ascii=False)
# ==================== 工具映射（供 API 使用）====================
TOOL_MAP = {
'search': search,
'quick_search': quick_search,
}
TOOL_LIST = [
{
"name": "search",
"description": "执行智能搜索，支持多轮迭代和自我反思",
"inputSchema": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "搜索查询问题"
},
"max_iterations": {
"type": "integer",
"description": "最大迭代次数（1-10）",
"default": 3,
"minimum": 1,
"maximum": 10
}
},
"required": ["query"]
}
},
{
"name": "quick_search",
"description": "快速搜索，单次迭代，适合简单问题",
"inputSchema": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "搜索查询问题"
}
},
"required": ["query"]
}
}
]
# 导出
__all__ = ['server', 'SYSTEM_PROMPT', 'TOOL_MAP', 'TOOL_LIST']
if __name__ == '__main__':
server.run()
steering_agent/Dockerfile
+24
View File
@@ -0,0 +1,24 @@
FROM python:3.12-slim
WORKDIR /app
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
ENV DATA_DIR=/app/data
RUN apt-get update && apt-get install -y gcc curl && rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
# 创建数据目录用于持久化存储
RUN mkdir -p /app/data
EXPOSE 8000
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
CMD curl -f http://localhost:8000/health || exit 1
CMD ["python", "run_api_server.py"]
steering_agent/README.md
+176
View File
@@ -0,0 +1,176 @@
# Steering Agent
项目约束管理Agent，提供持久的项目知识，保证一致的代码生成。
## 核心功能
**只做一件事**：对整个项目进行约束，让AI始终在预定规则下生成代码。
### 核心价值
| 价值 | 说明 |
|------|------|
| **持久的项目知识** | 自动提取并维护项目的技术栈、规范、模式等信息 |
| **一致的代码生成** | 确保 AI 生成的代码符合项目既有风格和规范 |
| **减少重复解释** | 一次定义规则，所有代码生成自动遵循 |
![alt text](image.png)
### 提供能力
| 工具 | 功能 |
|------|------|
| `extract_project_knowledge` | 从代码库提取项目知识 |
| `add_rule` | 添加用户定义规则 |
| `remove_rule` | 移除规则 |
| `list_rules` | 列出所有规则 |
| `get_context` | 获取完整项目上下文 |
| `check_compliance` | 检查代码合规性 |
| `generate_steering_doc` | 生成规范文档 |
## 快速开始
### 本地运行
```bash
cd steering_agent
pip install -r requirements.txt
python run_api_server.py
```
### Docker 运行
```bash
docker build -t steering-agent:latest .
docker run -p 8000:8000 steering-agent:latest
```
## 使用示例
### 1. 提取项目知识
```bash
curl -X POST http://localhost:8000/api/v1/extract \
-H "Content-Type: application/json" \
-d '{"project_path": "/path/to/your/project"}'
```
### 2. 添加规则
```bash
curl -X POST http://localhost:8000/api/v1/rules \
-H "Content-Type: application/json" \
-d '{
"rule_type": "must",
"rule_content": "所有 MCP 工具必须返回 JSON 格式",
"category": "output"
}'
```
### 3. 获取项目上下文
```bash
curl http://localhost:8000/api/v1/context?format=json
```
### 4. 检查代码合规性
```bash
curl -X POST http://localhost:8000/api/v1/check \
-H "Content-Type: application/json" \
-d '{
"code": "def my_tool():\n return \"result\"",
"file_type": "python"
}'
```
## 知识提取内容
| 类别 | 提取内容 |
|------|---------|
| **技术栈** | 语言、框架、版本、依赖 |
| **目录结构** | 项目布局、关键文件、项目模式 |
| **命名规范** | 文件命名、类名、函数名、变量名 |
| **代码模式** | 异步使用、错误处理、日志方式、文档字符串 |
| **配置规范** | 环境变量、配置文件 |
## 规则类型
| 类型 | 说明 | 示例 |
|------|------|------|
| `must` | 必须遵循 | "所有函数必须有 docstring" |
| `must_not` | 禁止事项 | "禁止硬编码 API Key" |
| `prefer` | 推荐做法 | "优先使用 async/await" |
| `security` | 安全规则 | "禁止使用 eval()" |
| `architecture` | 架构规则 | "MCP 工具必须定义在 mcp_server.py" |
## 与其他 Agent 配合
### 与 specs_agent 配合
1. 使用 `get_context` 获取项目上下文
2. 将上下文传递给 specs_agent
3. specs_agent 生成符合项目规范的需求/设计文档
### 与 format_police_agent 配合
1. 使用 `check_compliance` 检查代码
2. 调用 format_police 检查输出格式
3. 确保所有输出符合规范
## 环境变量
| 变量 | 必需 | 说明 |
|------|------|------|
| API_PORT | 否 | 服务端口，默认 8000 |
| API_HOST | 否 | 服务主机，默认 0.0.0.0 |
## 项目结构
```
steering_agent/
├── Dockerfile
├── README.md
├── USAGE.md
├── requirements.txt
├── run_api_server.py
└── src/
├── __init__.py
└── server/
├── __init__.py
├── api_server.py
├── mcp_server.py
├── extractors/
│ ├── __init__.py
│ ├── tech_stack.py
│ ├── structure.py
│ ├── naming.py
│ ├── patterns.py
│ └── config.py
└── rules/
├── __init__.py
├── rule_store.py
└── checker.py
```
## 服务端点
| 端点 | 方法 | 说明 |
|------|------|------|
| `/` | GET | 服务状态 |
| `/health` | GET | 健康检查 |
| `/mcp` | POST | MCP JSON-RPC |
| `/mcp/sse` | GET/POST | MCP SSE 流式 |
| `/api/v1/extract` | POST | 提取项目知识 |
| `/api/v1/rules` | GET | 列出规则 |
| `/api/v1/rules` | POST | 添加规则 |
| `/api/v1/rules/{id}` | DELETE | 删除规则 |
| `/api/v1/context` | GET | 获取项目上下文 |
| `/api/v1/check` | POST | 检查代码合规性 |
| `/api/v1/doc` | GET | 生成规范文档 |
## 部署信息
| 配置项 | 值 |
|--------|-----|
| 镜像地址 | agnettaiji.azurecr.io/ai-agents/steering-agent:latest |
| 服务端口 | 8000 |
| 健康检查 | /health |
steering_agent/USAGE.md
+478
View File
@@ -0,0 +1,478 @@
# Steering Agent 使用指南
## 概述
Steering Agent 是一个项目约束管理工具，用于：
- 从代码库自动提取项目知识
- 管理用户定义的规则
- 在代码生成前提供上下文
- 在代码生成后进行合规检查
## 工作流程
```
┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐
│ 知识提取 │ ──▶ │ 规则定义 │ ──▶ │ 使用场景 │
│ │ │ │ │ │
│ extract_ │ │ add_rule │ │ get_context │
│ project_ │ │ remove_rule │ │ check_ │
│ knowledge │ │ list_rules │ │ compliance │
└─────────────────┘ └─────────────────┘ └─────────────────┘
```
## API 使用
---
## 1. 知识提取
### 提取项目知识
**POST** `/api/v1/extract`
```bash
curl -X POST http://localhost:8000/api/v1/extract \
-H "Content-Type: application/json" \
-d '{
"project_path": "/home/user/my_project"
}'
```
**请求参数**:
| 参数 | 类型 | 必需 | 说明 |
|------|------|------|------|
| project_path | string | 是 | 项目根目录路径 |
| include_patterns | array | 否 | 包含的文件模式 |
| exclude_patterns | array | 否 | 排除的文件模式 |
**响应示例**:
```json
{
"success": true,
"knowledge": {
"project_path": "/home/user/my_project",
"extracted_at": "2024-01-01T12:00:00",
"tech_stack": {
"language": {"name": "Python", "version": "3.12"},
"framework": [{"name": "FastAPI", "version": ">=0.109.0"}],
"dependencies": {"production": ["fastapi>=0.109.0", "..."]}
},
"structure": {
"pattern": {"name": "agent_templates", "confidence": 0.95},
"key_files": {"entry_point": "run_api_server.py", "...": "..."}
},
"naming": {
"files": {"dominant_pattern": "snake_case", "consistency": 1.0},
"classes": {"dominant_pattern": "PascalCase", "consistency": 1.0}
},
"patterns": {
"async_usage": {"async_required": true, "async_function_ratio": 0.85},
"error_handling": {"style": "try-except-json-response"}
},
"config": {
"env_variables": [
{"name": "OPENAI_API_KEY", "required": true, "description": "OpenAI API Key"}
]
}
}
}
```
---
## 2. 规则管理
### 添加规则
**POST** `/api/v1/rules`
```bash
curl -X POST http://localhost:8000/api/v1/rules \
-H "Content-Type: application/json" \
-d '{
"rule_type": "must",
"rule_content": "所有 MCP 工具必须返回 JSON 格式",
"category": "output",
"priority": "high"
}'
```
**请求参数**:
| 参数 | 类型 | 必需 | 说明 |
|------|------|------|------|
| rule_type | string | 是 | 规则类型: must/must_not/prefer/security/architecture |
| rule_content | string | 是 | 规则内容 |
| category | string | 否 | 分类标签 |
| priority | string | 否 | 优先级: high/normal/low |
**响应示例**:
```json
{
"success": true,
"rule": {
"id": "a1b2c3d4",
"rule_type": "must",
"content": "所有 MCP 工具必须返回 JSON 格式",
"category": "output",
"priority": "high",
"created_at": "2024-01-01T12:00:00"
}
}
```
### 列出规则
**GET** `/api/v1/rules`
```bash
# 列出所有规则
curl http://localhost:8000/api/v1/rules
# 按类型过滤
curl "http://localhost:8000/api/v1/rules?rule_type=must"
# 按分类过滤
curl "http://localhost:8000/api/v1/rules?category=output"
```
**响应示例**:
```json
{
"success": true,
"rules": [
{
"id": "a1b2c3d4",
"rule_type": "must",
"content": "所有 MCP 工具必须返回 JSON 格式",
"category": "output",
"priority": "high"
}
],
"statistics": {
"total": 5,
"by_type": {"must": 2, "must_not": 2, "prefer": 1},
"by_priority": {"high": 1, "normal": 4}
}
}
```
### 删除规则
**DELETE** `/api/v1/rules/{rule_id}`
```bash
curl -X DELETE http://localhost:8000/api/v1/rules/a1b2c3d4
```
---
## 3. 获取项目上下文
### 获取 JSON 格式上下文
**GET** `/api/v1/context`
```bash
curl "http://localhost:8000/api/v1/context?format=json"
```
**响应示例**:
```json
{
"success": true,
"format": "json",
"context": {
"project_context": {
"name": "my_project",
"tech_stack": {"language": {"name": "Python", "version": "3.12"}},
"structure": {"pattern": {"name": "agent_templates"}}
},
"naming_conventions": {
"files": "snake_case",
"classes": "PascalCase",
"functions": "snake_case"
},
"code_patterns": {
"async_required": true,
"error_handling": "try-except-json-response",
"logging": "print"
},
"rules": {
"must": ["所有 MCP 工具必须返回 JSON 格式"],
"must_not": ["禁止硬编码 API Key"],
"prefer": [],
"security": [],
"architecture": []
}
}
}
```
### 获取 Markdown 格式上下文
```bash
curl "http://localhost:8000/api/v1/context?format=markdown"
```
---
## 4. 合规检查
### 检查代码合规性
**POST** `/api/v1/check`
```bash
curl -X POST http://localhost:8000/api/v1/check \
-H "Content-Type: application/json" \
-d '{
"code": "@server.tool()\nasync def my_tool():\n return \"result\"",
"file_type": "python",
"strict_mode": false
}'
```
**请求参数**:
| 参数 | 类型 | 必需 | 说明 |
|------|------|------|------|
| code | string | 是 | 待检查的代码 |
| file_type | string | 否 | 文件类型 |
| strict_mode | boolean | 否 | 严格模式（警告视为错误） |
**响应示例**:
```json
{
"success": true,
"compliant": false,
"violations": [
{
"rule": "所有 MCP 工具必须返回 JSON 格式",
"rule_type": "must",
"severity": "error",
"suggestion": "MCP 工具应使用 json.dumps() 返回 JSON 格式"
}
],
"warnings": [],
"summary": {
"total_violations": 1,
"total_warnings": 0,
"rules_checked": 5
}
}
```
---
## 5. 生成规范文档
### 生成 Markdown 文档
**GET** `/api/v1/doc`
```bash
curl "http://localhost:8000/api/v1/doc?format=markdown&include_examples=true"
```
**响应示例**:
```json
{
"success": true,
"document": "# my_project - 项目规范文档\n\n> 生成时间: 2024-01-01 12:00:00\n..."
}
```
---
## MCP 协议使用
### 工具列表
```bash
curl -X POST http://localhost:8000/mcp \
-H "Content-Type: application/json" \
-d '{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/list"
}'
```
### 调用工具
```bash
curl -X POST http://localhost:8000/mcp \
-H "Content-Type: application/json" \
-d '{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "extract_project_knowledge",
"arguments": {
"project_path": "/path/to/project"
}
}
}'
```
---
## 完整使用流程示例
```python
import requests
import json
BASE_URL = "http://localhost:8000"
# 1. 提取项目知识
print("=== 1. 提取项目知识 ===")
resp = requests.post(
f"{BASE_URL}/api/v1/extract",
json={"project_path": "/path/to/my_project"}
)
knowledge = resp.json()
print(f"技术栈: {knowledge['knowledge']['tech_stack']['language']}")
# 2. 添加规则
print("\n=== 2. 添加规则 ===")
rules_to_add = [
{"rule_type": "must", "rule_content": "所有 MCP 工具必须返回 JSON 格式"},
{"rule_type": "must_not", "rule_content": "禁止硬编码 API Key"},
{"rule_type": "prefer", "rule_content": "优先使用 async/await"},
]
for rule in rules_to_add:
resp = requests.post(f"{BASE_URL}/api/v1/rules", json=rule)
print(f"添加规则: {rule['rule_content']}")
# 3. 获取项目上下文
print("\n=== 3. 获取项目上下文 ===")
resp = requests.get(f"{BASE_URL}/api/v1/context?format=json")
context = resp.json()
print(f"规则数量: {len(context['context']['rules']['must'])} must, "
f"{len(context['context']['rules']['must_not'])} must_not")
# 4. 检查代码合规性
print("\n=== 4. 检查代码合规性 ===")
code_to_check = '''
@server.tool()
async def my_tool():
api_key = "sk-1234567890abcdef"
return "result"
'''
resp = requests.post(
f"{BASE_URL}/api/v1/check",
json={"code": code_to_check, "file_type": "python"}
)
result = resp.json()
print(f"合规: {result['compliant']}")
for violation in result['violations']:
print(f" - {violation['rule']}: {violation['suggestion']}")
# 5. 生成规范文档
print("\n=== 5. 生成规范文档 ===")
resp = requests.get(f"{BASE_URL}/api/v1/doc")
doc = resp.json()
# 保存到文件
with open("project_spec.md", "w") as f:
f.write(doc["document"])
print("规范文档已保存到 project_spec.md")
```
---
## 与其他 Agent 集成
### 与 specs_agent 集成
```python
# 1. 从 steering_agent 获取上下文
context_resp = requests.get("http://steering-agent:8000/api/v1/context?format=json")
context = context_resp.json()["context"]
# 2. 将上下文传递给 specs_agent
specs_resp = requests.post(
"http://specs-agent:8000/api/v1/requirements",
headers={"api-key": "your-key"},
json={
"brief_description": "创建一个数据去重Agent",
"context": json.dumps(context) # 传递项目上下文
}
)
```
### 与 format_police_agent 集成
```python
# 1. 检查代码合规性
check_resp = requests.post(
"http://steering-agent:8000/api/v1/check",
json={"code": generated_code}
)
# 2. 如果有输出，使用 format_police 检查格式
if check_resp.json()["compliant"]:
format_resp = requests.post(
"http://format-police-agent:8000/api/v1/check",
json={"content": output}
)
```
---
## 常见规则示例
### 必须遵循 (must)
```json
{"rule_type": "must", "rule_content": "所有 MCP 工具必须返回 JSON 格式"}
{"rule_type": "must", "rule_content": "所有函数必须有 docstring"}
{"rule_type": "must", "rule_content": "必须实现 /health 健康检查端点"}
{"rule_type": "must", "rule_content": "所有 API 端点必须有类型提示"}
```
### 禁止事项 (must_not)
```json
{"rule_type": "must_not", "rule_content": "禁止硬编码 API Key"}
{"rule_type": "must_not", "rule_content": "禁止使用 print 进行日志输出"}
{"rule_type": "must_not", "rule_content": "禁止使用同步阻塞操作"}
{"rule_type": "must_not", "rule_content": "禁止在工具函数中直接抛出异常"}
```
### 安全规则 (security)
```json
{"rule_type": "security", "rule_content": "禁止使用 eval() 或 exec()"}
{"rule_type": "security", "rule_content": "所有用户输入必须验证"}
{"rule_type": "security", "rule_content": "敏感信息必须使用环境变量"}
```
### 架构规则 (architecture)
```json
{"rule_type": "architecture", "rule_content": "MCP 工具必须定义在 mcp_server.py"}
{"rule_type": "architecture", "rule_content": "API 端点必须定义在 api_server.py"}
{"rule_type": "architecture", "rule_content": "遵循 agent_templates 目录结构"}
```
---
## 错误处理
所有 API 在失败时返回：
```json
{
"success": false,
"error": "错误信息"
}
```
常见错误：
- `项目路径不存在`: 检查 project_path 是否正确
- `无效的规则类型`: rule_type 必须是 must/must_not/prefer/security/architecture
- `规则已存在`: 相同内容的规则已添加
- `规则不存在`: 要删除的规则 ID 不存在
steering_agent/image.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 244 KiB

steering_agent/requirements.txt
+13
View File
@@ -0,0 +1,13 @@
# Pydantic AI
pydantic-ai>=0.0.14
# MCP
mcp>=0.9.0
fastmcp>=0.1.0
# FastAPI
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
# HTTP Client
aiohttp>=3.9.0
steering_agent/run_api_server.py
+23
View File
@@ -0,0 +1,23 @@
#!/usr/bin/env python
"""启动 Steering Agent API 服务器"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
if __name__ == '__main__':
from src.server.api_server import app
import uvicorn
import os
host = os.getenv('API_HOST', '0.0.0.0')
port = int(os.getenv('API_PORT', '8000'))
print(f"🚀 启动 Steering Agent API: http://{host}:{port}")
print("📋 核心功能:")
print(" - extract_project_knowledge: 提取项目知识")
print(" - add_rule / list_rules: 规则管理")
print(" - get_context: 获取项目上下文")
print(" - check_compliance: 合规检查")
print(" - generate_steering_doc: 生成规范文档")
uvicorn.run(app, host=host, port=port, log_level="info")
steering_agent/src/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Steering Agent 源代码包"""
steering_agent/src/server/__init__.py
+1
View File
@@ -0,0 +1 @@
"""服务器模块"""
steering_agent/src/server/api_server.py
+353
View File
@@ -0,0 +1,353 @@
"""
Steering Agent HTTP API 服务器
提供 REST API 和 MCP HTTP/SSE 端点。
"""
import json
import uuid
import os
from typing import Optional, Dict, Any, AsyncGenerator, List
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Request, Header, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field
from .mcp_server import TOOL_MAP, TOOL_LIST
# ==================== 配置 ====================
SERVER_NAME = "Steering Agent API"
# ==================== FastAPI 应用 ====================
@asynccontextmanager
async def lifespan(app: FastAPI):
print(f"🚀 {SERVER_NAME} 启动")
yield
print(f"🛑 {SERVER_NAME} 关闭")
app = FastAPI(
title=SERVER_NAME,
description="项目约束管理Agent，提供持久的项目知识，保证一致的代码生成",
version="1.0.0",
lifespan=lifespan
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# ==================== API Key 验证 ====================
async def verify_api_key(
api_key: Optional[str] = Header(None, alias="api-key"),
authorization: Optional[str] = Header(None)
) -> str:
"""验证 API Key"""
if api_key and api_key.strip() and api_key.strip() != "sk":
return api_key.strip()
if authorization:
key = authorization[7:].strip() if authorization.startswith("Bearer ") else authorization.strip()
if key and key != "sk":
return key
raise HTTPException(status_code=401, detail="缺少 API Key")
def get_api_key_from_request(request: Request) -> Optional[str]:
"""从请求头提取 API Key（不验证）"""
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
# ==================== 健康检查 ====================
@app.get("/")
async def root():
return {
"service": SERVER_NAME,
"status": "running",
"description": "项目约束管理Agent",
"tools": list(TOOL_MAP.keys()),
"capabilities": [
"extract_project_knowledge - 提取项目知识",
"add_rule / remove_rule / list_rules - 规则管理",
"get_context - 获取项目上下文",
"check_compliance - 合规检查",
"generate_steering_doc - 生成规范文档"
]
}
@app.get("/health")
async def health():
return {"status": "healthy", "service": SERVER_NAME}
# ==================== MCP 端点 ====================
sessions: Dict[str, Dict] = {}
async def handle_mcp_request(data: Dict, session_id: str = None, api_key: str = None) -> Dict:
"""处理 MCP JSON-RPC 请求"""
method = data.get("method")
params = data.get("params", {})
req_id = data.get("id")
# tools/call 需要验证 API Key
if method == "tools/call" and (not api_key or api_key == "sk"):
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": TOOL_LIST}}
elif method == "tools/call":
tool_name = params.get("name")
args = params.get("arguments", {})
if tool_name not in TOOL_MAP:
raise ValueError(f"Unknown tool: {tool_name}")
# 设置 API Key 到环境变量
old_key = os.environ.get('OPENAI_API_KEY')
if api_key:
os.environ['OPENAI_API_KEY'] = api_key
try:
result = await TOOL_MAP[tool_name](**args)
finally:
if old_key:
os.environ['OPENAI_API_KEY'] = old_key
return {
"jsonrpc": "2.0", "id": req_id,
"result": {"content": [{"type": "text", "text": str(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"
import asyncio
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)}})
# ==================== 业务 API ====================
# --- 知识提取 ---
class ExtractRequest(BaseModel):
"""知识提取请求"""
project_path: str = Field(..., description="项目根目录路径")
include_patterns: Optional[List[str]] = Field(None, description="包含的文件模式")
exclude_patterns: Optional[List[str]] = Field(None, description="排除的文件模式")
@app.post("/api/v1/extract")
async def api_extract(request: ExtractRequest, api_key: str = Depends(verify_api_key)):
"""提取项目知识"""
try:
result = await TOOL_MAP['extract_project_knowledge'](
project_path=request.project_path,
include_patterns=request.include_patterns,
exclude_patterns=request.exclude_patterns
)
return json.loads(result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# --- 规则管理 ---
class AddRuleRequest(BaseModel):
"""添加规则请求"""
rule_type: str = Field(..., description="规则类型")
rule_content: str = Field(..., description="规则内容")
category: Optional[str] = Field(None, description="分类标签")
priority: Optional[str] = Field("normal", description="优先级")
@app.post("/api/v1/rules")
async def api_add_rule(request: AddRuleRequest, api_key: str = Depends(verify_api_key)):
"""添加规则"""
try:
result = await TOOL_MAP['add_rule'](
rule_type=request.rule_type,
rule_content=request.rule_content,
category=request.category,
priority=request.priority
)
return json.loads(result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/v1/rules")
async def api_list_rules(
rule_type: Optional[str] = None,
category: Optional[str] = None,
api_key: str = Depends(verify_api_key)
):
"""列出规则"""
try:
result = await TOOL_MAP['list_rules'](
rule_type=rule_type,
category=category
)
return json.loads(result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.delete("/api/v1/rules/{rule_id}")
async def api_remove_rule(rule_id: str, api_key: str = Depends(verify_api_key)):
"""删除规则"""
try:
result = await TOOL_MAP['remove_rule'](rule_id=rule_id)
return json.loads(result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# --- 上下文获取 ---
@app.get("/api/v1/context")
async def api_get_context(
format: Optional[str] = "json",
include_rules: Optional[bool] = True,
include_patterns: Optional[bool] = True,
api_key: str = Depends(verify_api_key)
):
"""获取项目上下文"""
try:
result = await TOOL_MAP['get_context'](
output_format=format,
include_rules=include_rules,
include_patterns=include_patterns
)
return json.loads(result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# --- 合规检查 ---
class CheckRequest(BaseModel):
"""合规检查请求"""
code: str = Field(..., description="待检查的代码")
file_type: Optional[str] = Field(None, description="文件类型")
strict_mode: Optional[bool] = Field(False, description="严格模式")
@app.post("/api/v1/check")
async def api_check_compliance(request: CheckRequest, api_key: str = Depends(verify_api_key)):
"""检查代码合规性"""
try:
result = await TOOL_MAP['check_compliance'](
code=request.code,
file_type=request.file_type,
strict_mode=request.strict_mode
)
return json.loads(result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# --- 文档生成 ---
@app.get("/api/v1/doc")
async def api_generate_doc(
format: Optional[str] = "markdown",
include_examples: Optional[bool] = True,
api_key: str = Depends(verify_api_key)
):
"""生成规范文档"""
try:
result = await TOOL_MAP['generate_steering_doc'](
output_format=format,
include_examples=include_examples
)
return json.loads(result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == '__main__':
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
steering_agent/src/server/extractors/__init__.py
+15
View File
@@ -0,0 +1,15 @@
"""知识提取器模块"""
from .tech_stack import TechStackExtractor
from .structure import StructureExtractor
from .naming import NamingExtractor
from .patterns import PatternExtractor
from .config import ConfigExtractor
__all__ = [
'TechStackExtractor',
'StructureExtractor',
'NamingExtractor',
'PatternExtractor',
'ConfigExtractor'
]
steering_agent/src/server/extractors/config.py
+222
View File
@@ -0,0 +1,222 @@
"""
配置规范提取器
提取项目的配置规范：环境变量、配置文件、密钥处理等。
"""
import os
import re
from typing import Dict, List, Any, Optional
class ConfigExtractor:
"""配置规范提取器"""
def extract(self, project_path: str) -> Dict[str, Any]:
"""
提取项目的配置规范
Args:
project_path: 项目根目录路径
Returns:
配置规范信息字典
"""
python_files = self._find_python_files(project_path)
return {
"env_variables": self._extract_env_variables(python_files),
"config_files": self._analyze_config_files(project_path),
"secrets_handling": self._analyze_secrets_handling(python_files)
}
def _find_python_files(self, project_path: str) -> List[str]:
"""查找所有 Python 文件"""
python_files = []
ignore_dirs = {'__pycache__', '.git', '.venv', 'venv', 'node_modules', '.pytest_cache'}
for root, dirs, files in os.walk(project_path):
dirs[:] = [d for d in dirs if d not in ignore_dirs]
for f in files:
if f.endswith('.py'):
python_files.append(os.path.join(root, f))
return python_files
def _read_file_content(self, file_path: str) -> Optional[str]:
"""读取文件内容"""
try:
with open(file_path, 'r', encoding='utf-8') as f:
return f.read()
except Exception:
return None
def _extract_env_variables(self, python_files: List[str]) -> List[Dict[str, Any]]:
"""从代码中提取使用的环境变量"""
env_vars = {}
# 匹配模式
patterns = [
# os.getenv('VAR', 'default')
(r"os\.getenv\s*\(\s*['\"](\w+)['\"]\s*(?:,\s*['\"]([^'\"]*)['\"])?\s*\)", 'getenv'),
# os.environ.get('VAR', 'default')
(r"os\.environ\.get\s*\(\s*['\"](\w+)['\"]\s*(?:,\s*['\"]([^'\"]*)['\"])?\s*\)", 'environ_get'),
# os.environ['VAR']
(r"os\.environ\s*\[\s*['\"](\w+)['\"]\s*\]", 'environ_bracket'),
# os.environ.setdefault('VAR', 'default')
(r"os\.environ\.setdefault\s*\(\s*['\"](\w+)['\"]\s*,\s*['\"]([^'\"]*)['\"]", 'setdefault'),
]
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
for pattern, pattern_type in patterns:
matches = re.findall(pattern, content)
for match in matches:
if isinstance(match, tuple):
var_name = match[0]
default_value = match[1] if len(match) > 1 else None
else:
var_name = match
default_value = None
if var_name not in env_vars:
env_vars[var_name] = {
"name": var_name,
"required": pattern_type == 'environ_bracket',
"default": default_value if default_value else None,
"description": self._guess_env_description(var_name)
}
elif default_value and not env_vars[var_name]["default"]:
env_vars[var_name]["default"] = default_value
return list(env_vars.values())
def _guess_env_description(self, var_name: str) -> str:
"""根据变量名猜测描述"""
descriptions = {
"OPENAI_API_KEY": "OpenAI API Key",
"OPENAI_BASE_URL": "OpenAI API Base URL",
"API_KEY": "API Key",
"API_PORT": "服务端口",
"API_HOST": "服务主机地址",
"MODEL_NAME": "模型名称",
"LITELLM_MODEL": "LiteLLM 模型名称",
"LITELLM_GATEWAY_URL": "LiteLLM Gateway URL",
"LLM_BASE_URL": "LLM API Base URL",
"DATABASE_URL": "数据库连接 URL",
"REDIS_URL": "Redis 连接 URL",
"SECRET_KEY": "密钥",
"DEBUG": "调试模式",
"LOG_LEVEL": "日志级别",
}
if var_name in descriptions:
return descriptions[var_name]
# 根据命名模式猜测
if "KEY" in var_name or "SECRET" in var_name or "TOKEN" in var_name:
return "密钥/令牌"
elif "URL" in var_name or "HOST" in var_name:
return "服务地址"
elif "PORT" in var_name:
return "端口号"
elif "PATH" in var_name or "DIR" in var_name:
return "路径"
elif "NAME" in var_name:
return "名称"
elif "TIMEOUT" in var_name:
return "超时时间"
return "配置项"
def _analyze_config_files(self, project_path: str) -> List[Dict[str, Any]]:
"""分析配置文件"""
config_files = []
# 配置文件类型定义
config_types = {
"Dockerfile": ("docker", "Docker 容器配置"),
"docker-compose.yml": ("docker-compose", "Docker Compose 配置"),
"docker-compose.yaml": ("docker-compose", "Docker Compose 配置"),
"requirements.txt": ("python_deps", "Python 依赖"),
"requirements-dev.txt": ("python_deps_dev", "Python 开发依赖"),
"pyproject.toml": ("python_project", "Python 项目配置"),
"setup.py": ("python_setup", "Python 包配置"),
"setup.cfg": ("python_setup", "Python 包配置"),
"package.json": ("node_deps", "Node.js 依赖"),
".env": ("env", "环境变量"),
".env.example": ("env_example", "环境变量示例"),
"config.json": ("json_config", "JSON 配置"),
"config.yaml": ("yaml_config", "YAML 配置"),
"config.yml": ("yaml_config", "YAML 配置"),
".gitignore": ("git", "Git 忽略规则"),
"Makefile": ("make", "Make 构建配置"),
}
for filename, (file_type, purpose) in config_types.items():
file_path = os.path.join(project_path, filename)
if os.path.exists(file_path):
config_files.append({
"file": filename,
"type": file_type,
"purpose": purpose
})
return config_files
def _analyze_secrets_handling(self, python_files: List[str]) -> Dict[str, Any]:
"""分析密钥处理方式"""
hardcoded_secrets = False
env_based = False
# 检测硬编码密钥的模式
hardcoded_patterns = [
r'api_key\s*=\s*["\'][a-zA-Z0-9_-]{20,}["\']',
r'secret\s*=\s*["\'][a-zA-Z0-9_-]{20,}["\']',
r'password\s*=\s*["\'][^"\']{8,}["\']',
r'token\s*=\s*["\'][a-zA-Z0-9_-]{20,}["\']',
]
# 检测环境变量获取密钥的模式
env_patterns = [
r'os\.getenv\s*\(\s*["\'].*(?:KEY|SECRET|TOKEN|PASSWORD)',
r'os\.environ\.get\s*\(\s*["\'].*(?:KEY|SECRET|TOKEN|PASSWORD)',
r'os\.environ\s*\[\s*["\'].*(?:KEY|SECRET|TOKEN|PASSWORD)',
]
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
# 检查硬编码
for pattern in hardcoded_patterns:
if re.search(pattern, content, re.IGNORECASE):
# 排除示例和测试
if 'example' not in file_path.lower() and 'test' not in file_path.lower():
hardcoded_secrets = True
break
# 检查环境变量方式
for pattern in env_patterns:
if re.search(pattern, content, re.IGNORECASE):
env_based = True
break
# 确定处理模式
if env_based and not hardcoded_secrets:
pattern = "os.getenv with fallback"
elif hardcoded_secrets:
pattern = "hardcoded (not recommended)"
else:
pattern = "unknown"
return {
"hardcoded_secrets": hardcoded_secrets,
"env_based": env_based,
"pattern": pattern,
"recommendation": "使用环境变量管理敏感信息" if hardcoded_secrets else None
}
steering_agent/src/server/extractors/naming.py
+235
View File
@@ -0,0 +1,235 @@
"""
命名规范提取器
使用 AST 解析 Python 代码，分析类名、函数名、变量名的命名模式。
"""
import os
import re
import ast
from typing import Dict, List, Any, Optional
from pathlib import Path
from collections import Counter
class NamingExtractor:
"""命名规范提取器"""
def extract(self, project_path: str) -> Dict[str, Any]:
"""
提取项目的命名规范
Args:
project_path: 项目根目录路径
Returns:
命名规范信息字典
"""
# 收集所有 Python 文件中的命名
names = self._collect_names(project_path)
return {
"files": self._analyze_file_naming(project_path),
"classes": self._analyze_pattern(names["classes"], "classes"),
"functions": self._analyze_pattern(names["functions"], "functions"),
"variables": self._analyze_pattern(names["variables"], "variables"),
"constants": self._analyze_pattern(names["constants"], "constants"),
"summary": self._generate_summary(names)
}
def _find_python_files(self, project_path: str) -> List[str]:
"""查找所有 Python 文件"""
python_files = []
ignore_dirs = {'__pycache__', '.git', '.venv', 'venv', 'node_modules', '.pytest_cache'}
for root, dirs, files in os.walk(project_path):
# 过滤忽略的目录
dirs[:] = [d for d in dirs if d not in ignore_dirs]
for f in files:
if f.endswith('.py'):
python_files.append(os.path.join(root, f))
return python_files
def _collect_names(self, project_path: str) -> Dict[str, List[str]]:
"""使用 AST 收集代码中的命名"""
names = {
"classes": [],
"functions": [],
"variables": [],
"constants": []
}
for py_file in self._find_python_files(project_path):
try:
with open(py_file, 'r', encoding='utf-8') as f:
content = f.read()
tree = ast.parse(content)
for node in ast.walk(tree):
# 类名
if isinstance(node, ast.ClassDef):
names["classes"].append(node.name)
# 函数名（包括方法）
elif isinstance(node, ast.FunctionDef) or isinstance(node, ast.AsyncFunctionDef):
# 排除魔术方法
if not (node.name.startswith('__') and node.name.endswith('__')):
names["functions"].append(node.name)
# 变量和常量（模块级别的赋值）
elif isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Name):
name = target.id
# 排除私有变量
if not name.startswith('_'):
# 全大写视为常量
if name.isupper() or (name.upper() == name and '_' in name):
names["constants"].append(name)
else:
names["variables"].append(name)
except Exception:
continue
return names
def _detect_naming_pattern(self, name: str) -> Optional[str]:
"""检测单个名称的命名模式"""
if not name:
return None
# UPPER_SNAKE_CASE (常量)
if re.match(r'^[A-Z][A-Z0-9_]*$', name):
return "UPPER_SNAKE_CASE"
# PascalCase (类名)
if re.match(r'^[A-Z][a-zA-Z0-9]*$', name):
return "PascalCase"
# snake_case
if re.match(r'^[a-z][a-z0-9_]*$', name):
return "snake_case"
# camelCase
if re.match(r'^[a-z][a-zA-Z0-9]*$', name):
return "camelCase"
# kebab-case (通常用于文件名)
if re.match(r'^[a-z][a-z0-9-]*$', name):
return "kebab-case"
return "mixed"
def _analyze_pattern(self, names: List[str], category: str) -> Dict[str, Any]:
"""分析命名模式"""
if not names:
return {
"dominant_pattern": None,
"distribution": {},
"examples": [],
"consistency": 0.0
}
# 统计各模式的数量
patterns = Counter()
for name in names:
pattern = self._detect_naming_pattern(name)
if pattern:
patterns[pattern] += 1
if not patterns:
return {
"dominant_pattern": None,
"distribution": {},
"examples": names[:5],
"consistency": 0.0
}
# 找出主导模式
dominant = patterns.most_common(1)[0][0]
total = sum(patterns.values())
consistency = patterns[dominant] / total if total > 0 else 0.0
# 获取示例
examples = []
for name in names:
if self._detect_naming_pattern(name) == dominant and name not in examples:
examples.append(name)
if len(examples) >= 5:
break
return {
"dominant_pattern": dominant,
"distribution": dict(patterns),
"examples": examples,
"consistency": round(consistency, 2)
}
def _analyze_file_naming(self, project_path: str) -> Dict[str, Any]:
"""分析文件命名规范"""
python_files = []
ignore_dirs = {'__pycache__', '.git', '.venv', 'venv', 'node_modules'}
for root, dirs, files in os.walk(project_path):
dirs[:] = [d for d in dirs if d not in ignore_dirs]
for f in files:
if f.endswith('.py') and not f.startswith('__'):
# 去掉扩展名
name = f[:-3]
python_files.append(name)
if not python_files:
return {
"dominant_pattern": None,
"examples": [],
"consistency": 0.0
}
# 分析文件名模式
patterns = Counter()
for name in python_files:
pattern = self._detect_naming_pattern(name)
if pattern:
patterns[pattern] += 1
if not patterns:
return {
"dominant_pattern": None,
"examples": python_files[:5],
"consistency": 0.0
}
dominant = patterns.most_common(1)[0][0]
total = sum(patterns.values())
consistency = patterns[dominant] / total if total > 0 else 0.0
# 获取示例（加回 .py 扩展名）
examples = [f"{name}.py" for name in python_files[:5]]
return {
"dominant_pattern": dominant,
"examples": examples,
"consistency": round(consistency, 2)
}
def _generate_summary(self, names: Dict[str, List[str]]) -> Dict[str, str]:
"""生成命名规范摘要"""
summary = {}
# 分析各类别的主导模式
for category, name_list in names.items():
if name_list:
patterns = Counter()
for name in name_list:
pattern = self._detect_naming_pattern(name)
if pattern:
patterns[pattern] += 1
if patterns:
dominant = patterns.most_common(1)[0][0]
summary[category] = dominant
return summary
steering_agent/src/server/extractors/patterns.py
+390
View File
@@ -0,0 +1,390 @@
"""
代码模式提取器
分析代码中的常见模式：异步使用、错误处理、日志方式、文档字符串等。
"""
import os
import re
import ast
from typing import Dict, List, Any, Optional
from collections import Counter
class PatternExtractor:
"""代码模式提取器"""
def extract(self, project_path: str) -> Dict[str, Any]:
"""
提取项目的代码模式
Args:
project_path: 项目根目录路径
Returns:
代码模式信息字典
"""
python_files = self._find_python_files(project_path)
return {
"async_usage": self._analyze_async_usage(python_files),
"error_handling": self._analyze_error_handling(python_files),
"logging": self._analyze_logging(python_files),
"imports": self._analyze_import_style(python_files),
"docstrings": self._analyze_docstrings(python_files),
"type_hints": self._analyze_type_hints(python_files),
"common_patterns": self._detect_common_patterns(python_files)
}
def _find_python_files(self, project_path: str) -> List[str]:
"""查找所有 Python 文件"""
python_files = []
ignore_dirs = {'__pycache__', '.git', '.venv', 'venv', 'node_modules', '.pytest_cache'}
for root, dirs, files in os.walk(project_path):
dirs[:] = [d for d in dirs if d not in ignore_dirs]
for f in files:
if f.endswith('.py'):
python_files.append(os.path.join(root, f))
return python_files
def _read_file_content(self, file_path: str) -> Optional[str]:
"""读取文件内容"""
try:
with open(file_path, 'r', encoding='utf-8') as f:
return f.read()
except Exception:
return None
def _analyze_async_usage(self, python_files: List[str]) -> Dict[str, Any]:
"""分析异步使用情况"""
total_functions = 0
async_functions = 0
await_count = 0
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
try:
tree = ast.parse(content)
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef):
total_functions += 1
elif isinstance(node, ast.AsyncFunctionDef):
total_functions += 1
async_functions += 1
# 统计 await 使用
await_count += len(re.findall(r'\bawait\b', content))
except Exception:
continue
async_ratio = async_functions / total_functions if total_functions > 0 else 0.0
return {
"async_required": async_ratio > 0.5,
"async_function_ratio": round(async_ratio, 2),
"total_functions": total_functions,
"async_functions": async_functions,
"await_usage": "consistent" if await_count > async_functions else "minimal"
}
def _analyze_error_handling(self, python_files: List[str]) -> Dict[str, Any]:
"""分析错误处理模式"""
try_except_count = 0
json_error_response = 0
raise_count = 0
http_exception_count = 0
patterns = []
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
# 统计 try-except
try_except_count += len(re.findall(r'\btry\s*:', content))
# 检查是否返回 JSON 格式的错误
if re.search(r'return\s+json\.dumps\s*\(\s*\{[^}]*["\']error["\']\s*:', content):
json_error_response += 1
# 统计 raise
raise_count += len(re.findall(r'\braise\b', content))
# 检查 HTTPException
if 'HTTPException' in content:
http_exception_count += 1
# 确定主要的错误处理风格
if json_error_response > 0:
style = "try-except-json-response"
pattern = '返回包含 success/error 字段的 JSON'
elif http_exception_count > 0:
style = "http-exception"
pattern = '使用 HTTPException 抛出 HTTP 错误'
elif try_except_count > 0:
style = "try-except-raise"
pattern = '捕获异常后重新抛出'
else:
style = "minimal"
pattern = '最小化错误处理'
return {
"style": style,
"pattern": pattern,
"try_except_count": try_except_count,
"json_error_responses": json_error_response,
"http_exceptions": http_exception_count
}
def _analyze_logging(self, python_files: List[str]) -> Dict[str, Any]:
"""分析日志方式"""
print_count = 0
logging_count = 0
emoji_prefix = False
print_examples = []
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
# 统计 print 使用
print_matches = re.findall(r'print\s*\(["\']([^"\']+)["\']', content)
print_count += len(print_matches)
# 检查是否使用 emoji 前缀
for match in print_matches:
if re.match(r'^[\U0001F300-\U0001F9FF]', match):
emoji_prefix = True
if len(print_examples) < 3:
print_examples.append(match[:50])
# 统计 logging 使用
logging_count += len(re.findall(r'logging\.(debug|info|warning|error|critical)', content))
logging_count += len(re.findall(r'logger\.(debug|info|warning|error|critical)', content))
# 确定主要的日志方式
if logging_count > print_count:
method = "logging"
elif print_count > 0:
method = "print"
else:
method = "none"
return {
"method": method,
"format": "emoji_prefix" if emoji_prefix else "plain",
"print_count": print_count,
"logging_count": logging_count,
"examples": print_examples
}
def _analyze_import_style(self, python_files: List[str]) -> Dict[str, Any]:
"""分析导入风格"""
absolute_imports = 0
relative_imports = 0
common_imports = Counter()
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
try:
tree = ast.parse(content)
for node in ast.walk(tree):
if isinstance(node, ast.Import):
absolute_imports += 1
for alias in node.names:
module = alias.name.split('.')[0]
common_imports[module] += 1
elif isinstance(node, ast.ImportFrom):
if node.level > 0:
relative_imports += 1
else:
absolute_imports += 1
if node.module:
module = node.module.split('.')[0]
common_imports[module] += 1
except Exception:
continue
total = absolute_imports + relative_imports
relative_ratio = relative_imports / total if total > 0 else 0.0
return {
"style": "relative" if relative_ratio > 0.3 else "absolute",
"relative_import_ratio": round(relative_ratio, 2),
"absolute_imports": absolute_imports,
"relative_imports": relative_imports,
"common_imports": [imp for imp, _ in common_imports.most_common(10)]
}
def _analyze_docstrings(self, python_files: List[str]) -> Dict[str, Any]:
"""分析文档字符串"""
total_functions = 0
functions_with_docstring = 0
docstring_styles = Counter()
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
try:
tree = ast.parse(content)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
total_functions += 1
# 检查是否有 docstring
if (node.body and
isinstance(node.body[0], ast.Expr) and
isinstance(node.body[0].value, ast.Constant) and
isinstance(node.body[0].value.value, str)):
functions_with_docstring += 1
docstring = node.body[0].value.value
# 检测 docstring 风格
if 'Args:' in docstring or 'Returns:' in docstring:
docstring_styles['google'] += 1
elif ':param' in docstring or ':return:' in docstring:
docstring_styles['sphinx'] += 1
elif 'Parameters' in docstring or 'Returns' in docstring:
docstring_styles['numpy'] += 1
else:
docstring_styles['simple'] += 1
except Exception:
continue
coverage = functions_with_docstring / total_functions if total_functions > 0 else 0.0
dominant_style = docstring_styles.most_common(1)[0][0] if docstring_styles else "none"
return {
"coverage": round(coverage, 2),
"style": dominant_style,
"total_functions": total_functions,
"documented_functions": functions_with_docstring,
"style_distribution": dict(docstring_styles)
}
def _analyze_type_hints(self, python_files: List[str]) -> Dict[str, Any]:
"""分析类型提示"""
total_functions = 0
functions_with_return_type = 0
total_parameters = 0
parameters_with_type = 0
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
try:
tree = ast.parse(content)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
total_functions += 1
# 检查返回类型
if node.returns is not None:
functions_with_return_type += 1
# 检查参数类型
for arg in node.args.args:
if arg.arg != 'self' and arg.arg != 'cls':
total_parameters += 1
if arg.annotation is not None:
parameters_with_type += 1
except Exception:
continue
return_coverage = functions_with_return_type / total_functions if total_functions > 0 else 0.0
param_coverage = parameters_with_type / total_parameters if total_parameters > 0 else 0.0
return {
"coverage": round((return_coverage + param_coverage) / 2, 2),
"return_type_coverage": round(return_coverage, 2),
"parameter_type_coverage": round(param_coverage, 2),
"total_functions": total_functions,
"total_parameters": total_parameters
}
def _detect_common_patterns(self, python_files: List[str]) -> List[Dict[str, Any]]:
"""检测常见代码模式"""
patterns = []
mcp_tool_files = []
fastapi_endpoint_files = []
pydantic_model_files = []
singleton_agent_files = []
for file_path in python_files:
content = self._read_file_content(file_path)
if not content:
continue
rel_path = os.path.basename(file_path)
# MCP 工具模式
if '@server.tool()' in content:
mcp_tool_files.append(rel_path)
# FastAPI 端点模式
if re.search(r'@app\.(get|post|put|delete|patch)', content):
fastapi_endpoint_files.append(rel_path)
# Pydantic 模型模式
if re.search(r'class\s+\w+\s*\(\s*BaseModel\s*\)', content):
pydantic_model_files.append(rel_path)
# Agent 单例模式
if re.search(r'def\s+get_agent\s*\([^)]*\)\s*->\s*Agent', content):
singleton_agent_files.append(rel_path)
if mcp_tool_files:
patterns.append({
"name": "mcp_tool_pattern",
"description": "MCP 工具定义模式",
"occurrences": len(mcp_tool_files),
"files": mcp_tool_files[:5]
})
if fastapi_endpoint_files:
patterns.append({
"name": "fastapi_endpoint_pattern",
"description": "FastAPI 端点模式",
"occurrences": len(fastapi_endpoint_files),
"files": fastapi_endpoint_files[:5]
})
if pydantic_model_files:
patterns.append({
"name": "pydantic_model_pattern",
"description": "Pydantic 模型模式",
"occurrences": len(pydantic_model_files),
"files": pydantic_model_files[:5]
})
if singleton_agent_files:
patterns.append({
"name": "singleton_agent_pattern",
"description": "Agent 单例模式",
"occurrences": len(singleton_agent_files),
"files": singleton_agent_files[:5]
})
return patterns
steering_agent/src/server/extractors/structure.py
+307
View File
@@ -0,0 +1,307 @@
"""
目录结构提取器
扫描项目文件系统，提取目录结构、关键文件、项目模式等信息。
"""
import os
from typing import Dict, List, Any, Optional
from pathlib import Path
import fnmatch
# 忽略的目录和文件模式
IGNORE_PATTERNS = [
"__pycache__",
".git",
".venv",
"venv",
"env",
".env",
"node_modules",
".idea",
".vscode",
"*.pyc",
"*.pyo",
"*.pyd",
".DS_Store",
"*.egg-info",
"dist",
"build",
".pytest_cache",
".mypy_cache",
"*.log",
".coverage",
"htmlcov",
]
# 项目模式定义
PROJECT_PATTERNS = {
"agent_templates": {
"description": "基于 Pydantic AI 的 Agent 模板结构",
"required_files": [
"src/server/mcp_server.py",
"src/server/api_server.py",
"run_api_server.py"
],
"optional_files": [
"Dockerfile",
"requirements.txt",
"README.md"
]
},
"fastapi_standard": {
"description": "标准 FastAPI 项目结构",
"required_files": ["main.py"],
"alternative_files": ["app.py", "app/main.py"],
"optional_files": ["requirements.txt", "Dockerfile"]
},
"python_package": {
"description": "Python 包项目结构",
"required_files": [],
"alternative_files": ["setup.py", "pyproject.toml"],
"optional_files": ["README.md", "LICENSE"]
}
}
class StructureExtractor:
"""项目结构提取器"""
def __init__(self, max_depth: int = 5):
self.max_depth = max_depth
def extract(self, project_path: str) -> Dict[str, Any]:
"""
提取项目结构信息
Args:
project_path: 项目根目录路径
Returns:
项目结构信息字典
"""
return {
"root_files": self._get_root_files(project_path),
"directories": self._get_directory_tree(project_path),
"key_files": self._identify_key_files(project_path),
"pattern": self._detect_project_pattern(project_path),
"statistics": self._get_statistics(project_path)
}
def _should_ignore(self, name: str) -> bool:
"""检查是否应该忽略该文件/目录"""
for pattern in IGNORE_PATTERNS:
if fnmatch.fnmatch(name, pattern):
return True
return False
def _get_root_files(self, project_path: str) -> List[str]:
"""获取根目录下的文件列表"""
files = []
try:
for item in os.listdir(project_path):
item_path = os.path.join(project_path, item)
if os.path.isfile(item_path) and not self._should_ignore(item):
files.append(item)
except Exception:
pass
return sorted(files)
def _get_directory_tree(self, project_path: str, current_depth: int = 0) -> Dict[str, Any]:
"""获取目录树结构"""
if current_depth >= self.max_depth:
return {"...": "max_depth_reached"}
tree = {}
try:
for item in sorted(os.listdir(project_path)):
if self._should_ignore(item):
continue
item_path = os.path.join(project_path, item)
if os.path.isfile(item_path):
tree[item] = "file"
elif os.path.isdir(item_path):
subtree = self._get_directory_tree(item_path, current_depth + 1)
if subtree: # 只添加非空目录
tree[item] = subtree
except Exception:
pass
return tree
def _identify_key_files(self, project_path: str) -> Dict[str, Optional[str]]:
"""识别关键文件"""
key_files = {
"entry_point": None,
"mcp_tools": None,
"api_routes": None,
"dependencies": None,
"container": None,
"readme": None,
"config": None
}
# 入口点检测
entry_candidates = [
"run_api_server.py",
"main.py",
"app.py",
"run.py",
"server.py"
]
for candidate in entry_candidates:
if os.path.exists(os.path.join(project_path, candidate)):
key_files["entry_point"] = candidate
break
# MCP 工具文件
mcp_candidates = [
"src/server/mcp_server.py",
"mcp_server.py",
"tools.py"
]
for candidate in mcp_candidates:
if os.path.exists(os.path.join(project_path, candidate)):
key_files["mcp_tools"] = candidate
break
# API 路由文件
api_candidates = [
"src/server/api_server.py",
"api_server.py",
"api.py",
"routes.py",
"app/api.py"
]
for candidate in api_candidates:
if os.path.exists(os.path.join(project_path, candidate)):
key_files["api_routes"] = candidate
break
# 依赖文件
dep_candidates = [
"requirements.txt",
"pyproject.toml",
"setup.py",
"package.json"
]
for candidate in dep_candidates:
if os.path.exists(os.path.join(project_path, candidate)):
key_files["dependencies"] = candidate
break
# 容器配置
if os.path.exists(os.path.join(project_path, "Dockerfile")):
key_files["container"] = "Dockerfile"
# README
readme_candidates = ["README.md", "README.rst", "README.txt", "README"]
for candidate in readme_candidates:
if os.path.exists(os.path.join(project_path, candidate)):
key_files["readme"] = candidate
break
return key_files
def _detect_project_pattern(self, project_path: str) -> Dict[str, Any]:
"""检测项目模式"""
best_match = None
best_confidence = 0.0
for pattern_name, pattern_def in PROJECT_PATTERNS.items():
confidence = self._calculate_pattern_confidence(project_path, pattern_def)
if confidence > best_confidence:
best_confidence = confidence
best_match = pattern_name
if best_match and best_confidence > 0.5:
return {
"name": best_match,
"description": PROJECT_PATTERNS[best_match]["description"],
"confidence": round(best_confidence, 2)
}
return {
"name": "unknown",
"description": "未识别的项目结构",
"confidence": 0.0
}
def _calculate_pattern_confidence(self, project_path: str, pattern_def: Dict) -> float:
"""计算项目与模式的匹配度"""
required_files = pattern_def.get("required_files", [])
alternative_files = pattern_def.get("alternative_files", [])
optional_files = pattern_def.get("optional_files", [])
# 检查必需文件
required_found = 0
for f in required_files:
if os.path.exists(os.path.join(project_path, f)):
required_found += 1
# 如果有必需文件但没有全部找到，返回低置信度
if required_files and required_found < len(required_files):
return 0.0
# 检查替代文件（至少需要一个）
alternative_found = False
if alternative_files:
for f in alternative_files:
if os.path.exists(os.path.join(project_path, f)):
alternative_found = True
break
if not alternative_found:
return 0.0
# 计算可选文件匹配度
optional_found = 0
for f in optional_files:
if os.path.exists(os.path.join(project_path, f)):
optional_found += 1
# 计算总体置信度
total_weight = len(required_files) * 2 + len(optional_files)
if total_weight == 0:
return 0.5 if alternative_found else 0.0
score = (required_found * 2 + optional_found) / total_weight
return min(1.0, score + (0.3 if alternative_found else 0))
def _get_statistics(self, project_path: str) -> Dict[str, int]:
"""获取项目统计信息"""
stats = {
"total_files": 0,
"total_directories": 0,
"python_files": 0,
"javascript_files": 0,
"markdown_files": 0,
"config_files": 0
}
try:
for root, dirs, files in os.walk(project_path):
# 过滤忽略的目录
dirs[:] = [d for d in dirs if not self._should_ignore(d)]
stats["total_directories"] += len(dirs)
for f in files:
if self._should_ignore(f):
continue
stats["total_files"] += 1
if f.endswith('.py'):
stats["python_files"] += 1
elif f.endswith(('.js', '.ts', '.jsx', '.tsx')):
stats["javascript_files"] += 1
elif f.endswith('.md'):
stats["markdown_files"] += 1
elif f.endswith(('.json', '.yaml', '.yml', '.toml', '.ini', '.cfg')):
stats["config_files"] += 1
except Exception:
pass
return stats
steering_agent/src/server/extractors/tech_stack.py
+240
View File
@@ -0,0 +1,240 @@
"""
技术栈提取器
从项目文件中提取技术栈信息：语言、框架、依赖、运行时等。
"""
import os
import re
from typing import Dict, List, Optional, Any
from pathlib import Path
# 框架检测模式
FRAMEWORK_PATTERNS = {
# Python 框架
"fastapi": "FastAPI",
"flask": "Flask",
"django": "Django",
"pydantic-ai": "Pydantic AI",
"pydantic_ai": "Pydantic AI",
"fastmcp": "FastMCP",
"mcp": "MCP",
"uvicorn": "Uvicorn",
"aiohttp": "aiohttp",
"httpx": "httpx",
"requests": "Requests",
"sqlalchemy": "SQLAlchemy",
"celery": "Celery",
# Node.js 框架
"express": "Express",
"next": "Next.js",
"react": "React",
"vue": "Vue.js",
"angular": "Angular",
}
class TechStackExtractor:
"""技术栈信息提取器"""
def extract(self, project_path: str) -> Dict[str, Any]:
"""
提取项目的技术栈信息
Args:
project_path: 项目根目录路径
Returns:
技术栈信息字典
"""
return {
"language": self._detect_language(project_path),
"framework": self._detect_frameworks(project_path),
"dependencies": self._extract_dependencies(project_path),
"runtime": self._detect_runtime(project_path)
}
def _detect_language(self, project_path: str) -> Dict[str, Any]:
"""检测主要编程语言"""
result = {
"name": "Unknown",
"version": None,
"source": None
}
# 检查 Dockerfile 中的 Python 版本
dockerfile_path = os.path.join(project_path, "Dockerfile")
if os.path.exists(dockerfile_path):
try:
with open(dockerfile_path, 'r', encoding='utf-8') as f:
content = f.read()
# 匹配 FROM python:3.12-slim 等格式
match = re.search(r'FROM\s+python:(\d+\.?\d*)', content, re.IGNORECASE)
if match:
result = {
"name": "Python",
"version": match.group(1),
"source": "Dockerfile"
}
except Exception:
pass
# 检查 requirements.txt 存在
if result["name"] == "Unknown":
if os.path.exists(os.path.join(project_path, "requirements.txt")):
result = {
"name": "Python",
"version": None,
"source": "requirements.txt"
}
# 检查 package.json
if result["name"] == "Unknown":
package_json = os.path.join(project_path, "package.json")
if os.path.exists(package_json):
try:
import json
with open(package_json, 'r', encoding='utf-8') as f:
data = json.load(f)
engines = data.get("engines", {})
node_version = engines.get("node")
result = {
"name": "Node.js",
"version": node_version,
"source": "package.json"
}
except Exception:
result = {
"name": "Node.js",
"version": None,
"source": "package.json"
}
return result
def _detect_frameworks(self, project_path: str) -> List[Dict[str, str]]:
"""检测使用的框架"""
frameworks = []
seen = set()
# 从 requirements.txt 检测
req_path = os.path.join(project_path, "requirements.txt")
if os.path.exists(req_path):
try:
with open(req_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if not line or line.startswith('#'):
continue
# 解析包名和版本
match = re.match(r'^([a-zA-Z0-9_-]+)\[?[^\]]*\]?(.*)$', line)
if match:
pkg_name = match.group(1).lower()
version_spec = match.group(2).strip()
if pkg_name in FRAMEWORK_PATTERNS:
framework_name = FRAMEWORK_PATTERNS[pkg_name]
if framework_name not in seen:
seen.add(framework_name)
frameworks.append({
"name": framework_name,
"version": version_spec if version_spec else None
})
except Exception:
pass
# 从 package.json 检测
package_json = os.path.join(project_path, "package.json")
if os.path.exists(package_json):
try:
import json
with open(package_json, 'r', encoding='utf-8') as f:
data = json.load(f)
all_deps = {}
all_deps.update(data.get("dependencies", {}))
all_deps.update(data.get("devDependencies", {}))
for pkg_name, version in all_deps.items():
pkg_lower = pkg_name.lower()
if pkg_lower in FRAMEWORK_PATTERNS:
framework_name = FRAMEWORK_PATTERNS[pkg_lower]
if framework_name not in seen:
seen.add(framework_name)
frameworks.append({
"name": framework_name,
"version": version
})
except Exception:
pass
return frameworks
def _extract_dependencies(self, project_path: str) -> Dict[str, List[str]]:
"""提取项目依赖"""
dependencies = {
"production": [],
"development": []
}
# Python 依赖
req_path = os.path.join(project_path, "requirements.txt")
if os.path.exists(req_path):
try:
with open(req_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if line and not line.startswith('#'):
dependencies["production"].append(line)
except Exception:
pass
# 开发依赖
dev_req_path = os.path.join(project_path, "requirements-dev.txt")
if os.path.exists(dev_req_path):
try:
with open(dev_req_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if line and not line.startswith('#'):
dependencies["development"].append(line)
except Exception:
pass
# Node.js 依赖
package_json = os.path.join(project_path, "package.json")
if os.path.exists(package_json):
try:
import json
with open(package_json, 'r', encoding='utf-8') as f:
data = json.load(f)
for pkg, ver in data.get("dependencies", {}).items():
dependencies["production"].append(f"{pkg}@{ver}")
for pkg, ver in data.get("devDependencies", {}).items():
dependencies["development"].append(f"{pkg}@{ver}")
except Exception:
pass
return dependencies
def _detect_runtime(self, project_path: str) -> Dict[str, Any]:
"""检测运行时环境"""
runtime = {
"container": None,
"base_image": None
}
dockerfile_path = os.path.join(project_path, "Dockerfile")
if os.path.exists(dockerfile_path):
runtime["container"] = "Docker"
try:
with open(dockerfile_path, 'r', encoding='utf-8') as f:
content = f.read()
match = re.search(r'FROM\s+([^\s]+)', content, re.IGNORECASE)
if match:
runtime["base_image"] = match.group(1)
except Exception:
pass
return runtime
steering_agent/src/server/mcp_server.py
+770
View File
@@ -0,0 +1,770 @@
"""
Steering Agent MCP 服务器
核心功能：项目约束管理，提供持久的项目知识，保证一致的代码生成。
- 从代码库自动提取项目知识
- 管理用户定义的规则
- 在代码生成前提供上下文
- 在代码生成后进行合规检查
"""
import json
import os
from typing import Optional, List
from datetime import datetime
from pathlib import Path
from mcp.server.fastmcp import FastMCP
from .extractors import (
TechStackExtractor,
StructureExtractor,
NamingExtractor,
PatternExtractor,
ConfigExtractor
)
from .rules import RuleStore, ComplianceChecker
# ==================== 配置 ====================
# LiteLLM Gateway 配置
_BASE_URL = os.getenv('OPENAI_BASE_URL',
os.getenv('LLM_BASE_URL', 'https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1'))
_API_KEY = os.getenv('OPENAI_API_KEY', 'sk')
os.environ.setdefault('OPENAI_API_KEY', _API_KEY)
os.environ.setdefault('OPENAI_BASE_URL', _BASE_URL)
# 数据存储目录
DATA_DIR = Path(os.getenv('DATA_DIR', '/app/data'))
KNOWLEDGE_FILE = DATA_DIR / 'project_knowledge.json'
# ==================== MCP 服务器 ====================
server = FastMCP('Steering Agent')
# 全局状态
_project_knowledge = {}
_rule_store = RuleStore()
_compliance_checker = ComplianceChecker(_rule_store)
# 提取器实例
_tech_stack_extractor = TechStackExtractor()
_structure_extractor = StructureExtractor()
_naming_extractor = NamingExtractor()
_pattern_extractor = PatternExtractor()
_config_extractor = ConfigExtractor()
def _load_knowledge_from_file():
"""从文件加载项目知识"""
global _project_knowledge
try:
if KNOWLEDGE_FILE.exists():
with open(KNOWLEDGE_FILE, 'r', encoding='utf-8') as f:
data = json.load(f)
_project_knowledge = data.get("knowledge", {})
print(f"📂 已加载项目知识从 {KNOWLEDGE_FILE}")
except Exception as e:
print(f"⚠️ 加载项目知识文件失败: {e}")
_project_knowledge = {}
def _save_knowledge_to_file():
"""保存项目知识到文件"""
try:
DATA_DIR.mkdir(parents=True, exist_ok=True)
data = {
"version": "1.0",
"updated_at": datetime.now().isoformat(),
"knowledge": _project_knowledge
}
with open(KNOWLEDGE_FILE, 'w', encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=2)
except Exception as e:
print(f"⚠️ 保存项目知识文件失败: {e}")
# 启动时加载知识
_load_knowledge_from_file()
# ==================== 知识提取工具 ====================
@server.tool()
async def extract_project_knowledge(
project_path: str,
include_patterns: Optional[List[str]] = None,
exclude_patterns: Optional[List[str]] = None
) -> str:
"""
从代码库自动提取项目知识。
提取内容包括：
- 技术栈信息（语言、框架、依赖）
- 目录结构和项目模式
- 命名规范
- 代码模式
- 配置规范
Args:
project_path: 项目根目录路径
include_patterns: 包含的文件模式（可选）
exclude_patterns: 排除的文件模式（可选）
Returns:
提取的项目知识（JSON 格式）
"""
global _project_knowledge
if not project_path or not os.path.exists(project_path):
return json.dumps({
"success": False,
"error": f"项目路径不存在: {project_path}"
}, ensure_ascii=False)
try:
# 提取各类知识
knowledge = {
"project_path": project_path,
"extracted_at": datetime.now().isoformat(),
"tech_stack": _tech_stack_extractor.extract(project_path),
"structure": _structure_extractor.extract(project_path),
"naming": _naming_extractor.extract(project_path),
"patterns": _pattern_extractor.extract(project_path),
"config": _config_extractor.extract(project_path)
}
# 保存到全局状态
_project_knowledge = knowledge
# 持久化到文件
_save_knowledge_to_file()
return json.dumps({
"success": True,
"knowledge": knowledge
}, ensure_ascii=False, indent=2)
except Exception as e:
return json.dumps({
"success": False,
"error": str(e)
}, ensure_ascii=False)
# ==================== 规则管理工具 ====================
@server.tool()
async def add_rule(
rule_type: str,
rule_content: str,
category: Optional[str] = None,
priority: Optional[str] = "normal"
) -> str:
"""
添加用户定义的规则。
Args:
rule_type: 规则类型 (must/must_not/prefer/security/architecture)
rule_content: 规则内容
category: 分类标签（可选）
priority: 优先级 high/normal/low（默认 normal）
Returns:
添加结果（JSON 格式）
"""
result = _rule_store.add_rule(
rule_type=rule_type,
content=rule_content,
category=category,
priority=priority or "normal"
)
return json.dumps(result, ensure_ascii=False, indent=2)
@server.tool()
async def remove_rule(rule_id: str) -> str:
"""
移除规则。
Args:
rule_id: 规则 ID
Returns:
移除结果（JSON 格式）
"""
result = _rule_store.remove_rule(rule_id)
return json.dumps(result, ensure_ascii=False, indent=2)
@server.tool()
async def list_rules(
rule_type: Optional[str] = None,
category: Optional[str] = None
) -> str:
"""
列出所有规则。
Args:
rule_type: 按类型过滤（可选）
category: 按分类过滤（可选）
Returns:
规则列表（JSON 格式）
"""
rules = _rule_store.list_rules(rule_type=rule_type, category=category)
stats = _rule_store.get_statistics()
return json.dumps({
"success": True,
"rules": rules,
"statistics": stats
}, ensure_ascii=False, indent=2)
# ==================== 上下文获取工具 ====================
@server.tool()
async def get_context(
output_format: Optional[str] = "json",
include_rules: Optional[bool] = True,
include_patterns: Optional[bool] = True
) -> str:
"""
获取完整的项目上下文，用于提供给其他 AI Agent。
Args:
output_format: 输出格式 json/markdown（默认 json）
include_rules: 是否包含规则（默认 True）
include_patterns: 是否包含代码模式（默认 True）
Returns:
项目上下文（JSON 或 Markdown 格式）
"""
context = {
"project_context": {}
}
# 添加项目知识
if _project_knowledge:
tech_stack = _project_knowledge.get("tech_stack", {})
structure = _project_knowledge.get("structure", {})
naming = _project_knowledge.get("naming", {})
context["project_context"] = {
"name": os.path.basename(_project_knowledge.get("project_path", "unknown")),
"tech_stack": tech_stack,
"structure": {
"pattern": structure.get("pattern", {}),
"key_files": structure.get("key_files", {})
}
}
# 命名规范
context["naming_conventions"] = naming.get("summary", {})
# 代码模式
if include_patterns:
patterns = _project_knowledge.get("patterns", {})
context["code_patterns"] = {
"async_required": patterns.get("async_usage", {}).get("async_required", False),
"error_handling": patterns.get("error_handling", {}).get("style", "unknown"),
"logging": patterns.get("logging", {}).get("method", "unknown"),
"docstring_style": patterns.get("docstrings", {}).get("style", "unknown")
}
# 添加规则
if include_rules:
context["rules"] = _rule_store.get_rules_by_type()
# 根据格式输出
if output_format == "markdown":
markdown = _generate_markdown_context(context)
return json.dumps({
"success": True,
"format": "markdown",
"context": markdown
}, ensure_ascii=False, indent=2)
else:
return json.dumps({
"success": True,
"format": "json",
"context": context
}, ensure_ascii=False, indent=2)
def _generate_markdown_context(context: dict) -> str:
"""生成 Markdown 格式的上下文"""
lines = []
# 项目信息
project = context.get("project_context", {})
lines.append(f"# 项目规范文档: {project.get('name', 'Unknown')}")
lines.append("")
# 技术栈
tech_stack = project.get("tech_stack", {})
if tech_stack:
lines.append("## 技术栈")
lang = tech_stack.get("language", {})
if lang:
lines.append(f"- **语言**: {lang.get('name', 'Unknown')} {lang.get('version', '')}")
frameworks = tech_stack.get("framework", [])
if frameworks:
fw_list = ", ".join([f"{f.get('name', '')} {f.get('version', '')}" for f in frameworks])
lines.append(f"- **框架**: {fw_list}")
lines.append("")
# 命名规范
naming = context.get("naming_conventions", {})
if naming:
lines.append("## 命名规范")
for category, pattern in naming.items():
lines.append(f"- **{category}**: {pattern}")
lines.append("")
# 代码模式
patterns = context.get("code_patterns", {})
if patterns:
lines.append("## 代码模式")
lines.append(f"- **异步要求**: {'是' if patterns.get('async_required') else '否'}")
lines.append(f"- **错误处理**: {patterns.get('error_handling', 'unknown')}")
lines.append(f"- **日志方式**: {patterns.get('logging', 'unknown')}")
lines.append(f"- **文档风格**: {patterns.get('docstring_style', 'unknown')}")
lines.append("")
# 规则
rules = context.get("rules", {})
if rules:
if rules.get("must"):
lines.append("## 必须遵循的规则")
for i, rule in enumerate(rules["must"], 1):
lines.append(f"{i}. {rule}")
lines.append("")
if rules.get("must_not"):
lines.append("## 禁止事项")
for i, rule in enumerate(rules["must_not"], 1):
lines.append(f"{i}. {rule}")
lines.append("")
if rules.get("prefer"):
lines.append("## 推荐做法")
for i, rule in enumerate(rules["prefer"], 1):
lines.append(f"{i}. {rule}")
lines.append("")
if rules.get("security"):
lines.append("## 安全规则")
for i, rule in enumerate(rules["security"], 1):
lines.append(f"{i}. {rule}")
lines.append("")
if rules.get("architecture"):
lines.append("## 架构规则")
for i, rule in enumerate(rules["architecture"], 1):
lines.append(f"{i}. {rule}")
lines.append("")
return "\n".join(lines)
# ==================== 合规检查工具 ====================
@server.tool()
async def check_compliance(
code: str,
file_type: Optional[str] = None,
strict_mode: Optional[bool] = False
) -> str:
"""
检查代码是否符合项目规则。
Args:
code: 待检查的代码
file_type: 文件类型（可选，如 python）
strict_mode: 严格模式，所有警告视为错误（默认 False）
Returns:
检查结果（JSON 格式），包含：
- 是否通过
- 违反的规则列表
- 修改建议
"""
if not code or not code.strip():
return json.dumps({
"success": False,
"error": "代码不能为空"
}, ensure_ascii=False)
try:
result = _compliance_checker.check_compliance(
code=code,
file_type=file_type,
strict_mode=strict_mode or False
)
return json.dumps({
"success": True,
**result
}, ensure_ascii=False, indent=2)
except Exception as e:
return json.dumps({
"success": False,
"error": str(e)
}, ensure_ascii=False)
# ==================== 文档生成工具 ====================
@server.tool()
async def generate_steering_doc(
output_format: Optional[str] = "markdown",
include_examples: Optional[bool] = True
) -> str:
"""
生成人类可读的项目规范文档。
Args:
output_format: 输出格式（默认 markdown）
include_examples: 是否包含示例（默认 True）
Returns:
规范文档（Markdown 格式）
"""
lines = []
# 标题
project_name = "Unknown"
if _project_knowledge:
project_name = os.path.basename(_project_knowledge.get("project_path", "Unknown"))
lines.append(f"# {project_name} - 项目规范文档")
lines.append("")
lines.append(f"> 生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
lines.append("> 生成工具: Steering Agent")
lines.append("")
lines.append("---")
lines.append("")
# 技术栈
if _project_knowledge:
tech_stack = _project_knowledge.get("tech_stack", {})
lines.append("## 1. 技术栈")
lines.append("")
lang = tech_stack.get("language", {})
lines.append(f"### 语言")
lines.append(f"- **名称**: {lang.get('name', 'Unknown')}")
lines.append(f"- **版本**: {lang.get('version', 'Unknown')}")
lines.append("")
frameworks = tech_stack.get("framework", [])
if frameworks:
lines.append("### 框架")
lines.append("| 框架 | 版本 |")
lines.append("|------|------|")
for fw in frameworks:
lines.append(f"| {fw.get('name', '')} | {fw.get('version', '')} |")
lines.append("")
deps = tech_stack.get("dependencies", {})
if deps.get("production"):
lines.append("### 依赖")
lines.append("```")
for dep in deps["production"][:10]:
lines.append(dep)
if len(deps["production"]) > 10:
lines.append(f"... 共 {len(deps['production'])} 个依赖")
lines.append("```")
lines.append("")
# 目录结构
if _project_knowledge:
structure = _project_knowledge.get("structure", {})
lines.append("## 2. 项目结构")
lines.append("")
pattern = structure.get("pattern", {})
if pattern:
lines.append(f"### 项目模式")
lines.append(f"- **模式**: {pattern.get('name', 'unknown')}")
lines.append(f"- **描述**: {pattern.get('description', '')}")
lines.append(f"- **置信度**: {pattern.get('confidence', 0)}")
lines.append("")
key_files = structure.get("key_files", {})
if key_files:
lines.append("### 关键文件")
lines.append("| 类型 | 文件 |")
lines.append("|------|------|")
for file_type, file_path in key_files.items():
if file_path:
lines.append(f"| {file_type} | `{file_path}` |")
lines.append("")
# 命名规范
if _project_knowledge:
naming = _project_knowledge.get("naming", {})
lines.append("## 3. 命名规范")
lines.append("")
lines.append("| 类别 | 规范 | 一致性 | 示例 |")
lines.append("|------|------|--------|------|")
for category in ["files", "classes", "functions", "variables", "constants"]:
info = naming.get(category, {})
pattern = info.get("dominant_pattern", "N/A")
consistency = info.get("consistency", 0)
examples = info.get("examples", [])[:2]
examples_str = ", ".join(examples) if examples else "N/A"
lines.append(f"| {category} | {pattern} | {consistency:.0%} | {examples_str} |")
lines.append("")
# 代码模式
if _project_knowledge:
patterns = _project_knowledge.get("patterns", {})
lines.append("## 4. 代码模式")
lines.append("")
async_info = patterns.get("async_usage", {})
lines.append("### 异步使用")
lines.append(f"- **异步要求**: {'是' if async_info.get('async_required') else '否'}")
lines.append(f"- **异步函数比例**: {async_info.get('async_function_ratio', 0):.0%}")
lines.append("")
error_info = patterns.get("error_handling", {})
lines.append("### 错误处理")
lines.append(f"- **风格**: {error_info.get('style', 'unknown')}")
lines.append(f"- **模式**: {error_info.get('pattern', 'unknown')}")
lines.append("")
docstring_info = patterns.get("docstrings", {})
lines.append("### 文档字符串")
lines.append(f"- **覆盖率**: {docstring_info.get('coverage', 0):.0%}")
lines.append(f"- **风格**: {docstring_info.get('style', 'unknown')}")
lines.append("")
# 规则
rules = _rule_store.get_rules_by_type()
lines.append("## 5. 项目规则")
lines.append("")
if rules.get("must"):
lines.append("### 必须遵循")
for rule in rules["must"]:
lines.append(f"- ✅ {rule}")
lines.append("")
if rules.get("must_not"):
lines.append("### 禁止事项")
for rule in rules["must_not"]:
lines.append(f"- ❌ {rule}")
lines.append("")
if rules.get("prefer"):
lines.append("### 推荐做法")
for rule in rules["prefer"]:
lines.append(f"- 💡 {rule}")
lines.append("")
if rules.get("security"):
lines.append("### 安全规则")
for rule in rules["security"]:
lines.append(f"- 🔒 {rule}")
lines.append("")
if rules.get("architecture"):
lines.append("### 架构规则")
for rule in rules["architecture"]:
lines.append(f"- 🏗️ {rule}")
lines.append("")
# 配置规范
if _project_knowledge:
config = _project_knowledge.get("config", {})
env_vars = config.get("env_variables", [])
if env_vars:
lines.append("## 6. 环境变量")
lines.append("")
lines.append("| 变量 | 必需 | 默认值 | 描述 |")
lines.append("|------|------|--------|------|")
for var in env_vars:
required = "是" if var.get("required") else "否"
default = var.get("default") or "-"
lines.append(f"| `{var.get('name')}` | {required} | {default} | {var.get('description', '')} |")
lines.append("")
doc = "\n".join(lines)
return json.dumps({
"success": True,
"document": doc
}, ensure_ascii=False, indent=2)
# ==================== 工具映射（供 API 使用）====================
TOOL_MAP = {
'extract_project_knowledge': extract_project_knowledge,
'add_rule': add_rule,
'remove_rule': remove_rule,
'list_rules': list_rules,
'get_context': get_context,
'check_compliance': check_compliance,
'generate_steering_doc': generate_steering_doc,
}
TOOL_LIST = [
{
"name": "extract_project_knowledge",
"description": "从代码库自动提取项目知识（技术栈、结构、命名规范、代码模式）",
"inputSchema": {
"type": "object",
"properties": {
"project_path": {
"type": "string",
"description": "项目根目录路径"
},
"include_patterns": {
"type": "array",
"items": {"type": "string"},
"description": "包含的文件模式（可选）"
},
"exclude_patterns": {
"type": "array",
"items": {"type": "string"},
"description": "排除的文件模式（可选）"
}
},
"required": ["project_path"]
}
},
{
"name": "add_rule",
"description": "添加用户定义的规则",
"inputSchema": {
"type": "object",
"properties": {
"rule_type": {
"type": "string",
"enum": ["must", "must_not", "prefer", "security", "architecture"],
"description": "规则类型"
},
"rule_content": {
"type": "string",
"description": "规则内容"
},
"category": {
"type": "string",
"description": "分类标签（可选）"
},
"priority": {
"type": "string",
"enum": ["high", "normal", "low"],
"description": "优先级（默认 normal）"
}
},
"required": ["rule_type", "rule_content"]
}
},
{
"name": "remove_rule",
"description": "移除规则",
"inputSchema": {
"type": "object",
"properties": {
"rule_id": {
"type": "string",
"description": "规则 ID"
}
},
"required": ["rule_id"]
}
},
{
"name": "list_rules",
"description": "列出所有规则",
"inputSchema": {
"type": "object",
"properties": {
"rule_type": {
"type": "string",
"description": "按类型过滤（可选）"
},
"category": {
"type": "string",
"description": "按分类过滤（可选）"
}
}
}
},
{
"name": "get_context",
"description": "获取完整的项目上下文，用于提供给其他 AI Agent",
"inputSchema": {
"type": "object",
"properties": {
"output_format": {
"type": "string",
"enum": ["json", "markdown"],
"description": "输出格式（默认 json）"
},
"include_rules": {
"type": "boolean",
"description": "是否包含规则（默认 True）"
},
"include_patterns": {
"type": "boolean",
"description": "是否包含代码模式（默认 True）"
}
}
}
},
{
"name": "check_compliance",
"description": "检查代码是否符合项目规则",
"inputSchema": {
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "待检查的代码"
},
"file_type": {
"type": "string",
"description": "文件类型（可选）"
},
"strict_mode": {
"type": "boolean",
"description": "严格模式（默认 False）"
}
},
"required": ["code"]
}
},
{
"name": "generate_steering_doc",
"description": "生成人类可读的项目规范文档",
"inputSchema": {
"type": "object",
"properties": {
"output_format": {
"type": "string",
"description": "输出格式（默认 markdown）"
},
"include_examples": {
"type": "boolean",
"description": "是否包含示例（默认 True）"
}
}
}
}
]
if __name__ == '__main__':
server.run()
steering_agent/src/server/rules/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""规则管理模块"""
from .rule_store import RuleStore
from .checker import ComplianceChecker
__all__ = ['RuleStore', 'ComplianceChecker']
steering_agent/src/server/rules/checker.py
+308
View File
@@ -0,0 +1,308 @@
"""
合规检查器
检查代码是否符合项目规则。
"""
import re
import ast
from typing import Dict, List, Any, Optional
from .rule_store import RuleStore
class ComplianceChecker:
"""合规检查器"""
def __init__(self, rule_store: RuleStore):
self.rule_store = rule_store
# 内置检查规则映射
self._builtin_checks = {
# must 规则
"所有 MCP 工具必须返回 JSON 格式": self._check_json_return,
"所有函数必须有 docstring": self._check_docstring,
"必须使用 async/await 进行异步操作": self._check_async_usage,
"所有 API 端点必须有 docstring": self._check_docstring,
# must_not 规则
"禁止硬编码 API Key": self._check_no_hardcoded_secrets,
"禁止硬编码密钥": self._check_no_hardcoded_secrets,
"禁止使用 print 进行日志输出": self._check_no_print,
"禁止使用同步阻塞操作": self._check_no_sync_blocking,
"禁止在工具函数中直接抛出异常": self._check_no_direct_raise,
}
def check_compliance(
self,
code: str,
file_type: Optional[str] = None,
strict_mode: bool = False
) -> Dict[str, Any]:
"""
检查代码是否符合规则
Args:
code: 待检查的代码
file_type: 文件类型 (python, javascript, etc.)
strict_mode: 严格模式（所有警告视为错误）
Returns:
检查结果
"""
violations = []
warnings = []
# 获取所有规则
rules = self.rule_store.get_rules_by_type()
# 检查 must 规则
for rule_content in rules.get("must", []):
result = self._check_rule(code, rule_content, "must")
if result:
violations.append(result)
# 检查 must_not 规则
for rule_content in rules.get("must_not", []):
result = self._check_rule(code, rule_content, "must_not")
if result:
violations.append(result)
# 检查 security 规则
for rule_content in rules.get("security", []):
result = self._check_rule(code, rule_content, "security")
if result:
violations.append(result)
# 检查 prefer 规则（作为警告）
for rule_content in rules.get("prefer", []):
result = self._check_rule(code, rule_content, "prefer")
if result:
if strict_mode:
violations.append(result)
else:
warnings.append(result)
# 检查 architecture 规则
for rule_content in rules.get("architecture", []):
result = self._check_rule(code, rule_content, "architecture")
if result:
violations.append(result)
compliant = len(violations) == 0
return {
"compliant": compliant,
"violations": violations,
"warnings": warnings,
"summary": {
"total_violations": len(violations),
"total_warnings": len(warnings),
"rules_checked": sum(len(v) for v in rules.values())
}
}
def _check_rule(
self,
code: str,
rule_content: str,
rule_type: str
) -> Optional[Dict[str, Any]]:
"""
检查单个规则
Returns:
违规信息，如果符合规则则返回 None
"""
# 尝试使用内置检查
if rule_content in self._builtin_checks:
check_func = self._builtin_checks[rule_content]
is_violated, suggestion = check_func(code, rule_type)
if is_violated:
return {
"rule": rule_content,
"rule_type": rule_type,
"severity": "error" if rule_type in ["must", "must_not", "security"] else "warning",
"suggestion": suggestion
}
else:
# 使用通用关键词检查
result = self._generic_check(code, rule_content, rule_type)
if result:
return result
return None
def _generic_check(
self,
code: str,
rule_content: str,
rule_type: str
) -> Optional[Dict[str, Any]]:
"""通用规则检查（基于关键词）"""
# 提取规则中的关键词
keywords = self._extract_keywords(rule_content)
if not keywords:
return None
# must_not 规则：检查是否包含禁止的内容
if rule_type == "must_not":
for keyword in keywords:
if keyword.lower() in code.lower():
return {
"rule": rule_content,
"rule_type": rule_type,
"severity": "error",
"suggestion": f"代码中包含禁止的内容: {keyword}"
}
return None
def _extract_keywords(self, rule_content: str) -> List[str]:
"""从规则内容中提取关键词"""
# 简单的关键词提取
keywords = []
# 提取引号中的内容
quoted = re.findall(r'["\']([^"\']+)["\']', rule_content)
keywords.extend(quoted)
# 提取代码相关的关键词
code_keywords = re.findall(r'\b(print|eval|exec|os\.system|subprocess)\b', rule_content)
keywords.extend(code_keywords)
return keywords
# ==================== 内置检查函数 ====================
def _check_json_return(self, code: str, rule_type: str) -> tuple:
"""检查是否返回 JSON 格式"""
# 检查是否有 @server.tool() 装饰器
if '@server.tool()' not in code:
return False, None
# 检查是否使用 json.dumps 返回
if 'json.dumps' in code:
return False, None
# 检查是否直接返回字符串
if re.search(r'return\s+["\'][^"\']+["\']', code):
return True, "MCP 工具应使用 json.dumps() 返回 JSON 格式"
return False, None
def _check_docstring(self, code: str, rule_type: str) -> tuple:
"""检查函数是否有 docstring"""
try:
tree = ast.parse(code)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
# 跳过私有函数和魔术方法
if node.name.startswith('_'):
continue
# 检查是否有 docstring
if not (node.body and
isinstance(node.body[0], ast.Expr) and
isinstance(node.body[0].value, ast.Constant) and
isinstance(node.body[0].value.value, str)):
return True, f"函数 '{node.name}' 缺少 docstring"
return False, None
except Exception:
return False, None
def _check_async_usage(self, code: str, rule_type: str) -> tuple:
"""检查是否使用 async/await"""
# 如果代码中有异步调用但没有使用 await
if 'async def' in code:
return False, None
# 检查是否有需要异步的操作但没有使用 async
async_indicators = ['aiohttp', 'httpx', 'asyncio', 'await']
for indicator in async_indicators:
if indicator in code and 'async def' not in code:
return True, "检测到异步操作，建议使用 async/await"
return False, None
def _check_no_hardcoded_secrets(self, code: str, rule_type: str) -> tuple:
"""检查是否硬编码密钥"""
patterns = [
(r'api_key\s*=\s*["\'][a-zA-Z0-9_-]{20,}["\']', "API Key"),
(r'secret\s*=\s*["\'][a-zA-Z0-9_-]{20,}["\']', "Secret"),
(r'password\s*=\s*["\'][^"\']{8,}["\']', "Password"),
(r'token\s*=\s*["\'][a-zA-Z0-9_-]{20,}["\']', "Token"),
]
for pattern, secret_type in patterns:
if re.search(pattern, code, re.IGNORECASE):
return True, f"检测到硬编码的 {secret_type}，请使用环境变量"
return False, None
def _check_no_print(self, code: str, rule_type: str) -> tuple:
"""检查是否使用 print"""
# 排除 docstring 和注释中的 print
lines = code.split('\n')
in_docstring = False
for line in lines:
stripped = line.strip()
# 检测 docstring
if '"""' in stripped or "'''" in stripped:
in_docstring = not in_docstring
continue
if in_docstring:
continue
# 跳过注释
if stripped.startswith('#'):
continue
# 检查 print 调用
if re.search(r'\bprint\s*\(', line):
return True, "使用 logging 模块替代 print"
return False, None
def _check_no_sync_blocking(self, code: str, rule_type: str) -> tuple:
"""检查是否使用同步阻塞操作"""
blocking_patterns = [
(r'\brequests\.(get|post|put|delete|patch)\s*\(', "requests 库是同步的，使用 aiohttp 或 httpx"),
(r'\btime\.sleep\s*\(', "使用 asyncio.sleep 替代 time.sleep"),
(r'\burllib\.request\b', "使用 aiohttp 替代 urllib"),
]
for pattern, suggestion in blocking_patterns:
if re.search(pattern, code):
return True, suggestion
return False, None
def _check_no_direct_raise(self, code: str, rule_type: str) -> tuple:
"""检查工具函数中是否直接抛出异常"""
# 检查是否在 @server.tool() 装饰的函数中直接 raise
if '@server.tool()' not in code:
return False, None
# 简单检查：如果有 raise 但没有 try-except 包裹
if 'raise ' in code:
# 检查是否在 try-except 中
try:
tree = ast.parse(code)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
# 检查函数体中是否有未被 try 包裹的 raise
for child in ast.walk(node):
if isinstance(child, ast.Raise):
# 检查是否在 try 块中
# 这是简化的检查
pass
except Exception:
pass
return False, None
steering_agent/src/server/rules/rule_store.py
+283
View File
@@ -0,0 +1,283 @@
"""
规则存储
文件存储用户定义的规则，支持持久化。
"""
import json
import os
import uuid
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field, asdict
from datetime import datetime
from pathlib import Path
# 数据存储目录
DATA_DIR = Path(os.getenv('DATA_DIR', '/app/data'))
RULES_FILE = DATA_DIR / 'rules.json'
@dataclass
class Rule:
"""规则数据类"""
id: str
rule_type: str # must, must_not, prefer, security, architecture
content: str
category: Optional[str] = None
priority: str = "normal" # high, normal, low
created_at: str = field(default_factory=lambda: datetime.now().isoformat())
def to_dict(self) -> Dict[str, Any]:
return {
"id": self.id,
"rule_type": self.rule_type,
"content": self.content,
"category": self.category,
"priority": self.priority,
"created_at": self.created_at
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> 'Rule':
"""从字典创建规则"""
return cls(
id=data["id"],
rule_type=data["rule_type"],
content=data["content"],
category=data.get("category"),
priority=data.get("priority", "normal"),
created_at=data.get("created_at", datetime.now().isoformat())
)
class RuleStore:
"""规则存储（文件持久化）"""
# 有效的规则类型
VALID_RULE_TYPES = ["must", "must_not", "prefer", "security", "architecture"]
# 有效的优先级
VALID_PRIORITIES = ["high", "normal", "low"]
def __init__(self, storage_path: Optional[Path] = None):
"""
初始化规则存储
Args:
storage_path: 存储文件路径，默认使用 DATA_DIR/rules.json
"""
self._storage_path = storage_path or RULES_FILE
self._rules: Dict[str, Rule] = {}
self._load_from_file()
def _ensure_data_dir(self):
"""确保数据目录存在"""
self._storage_path.parent.mkdir(parents=True, exist_ok=True)
def _load_from_file(self):
"""从文件加载规则"""
try:
if self._storage_path.exists():
with open(self._storage_path, 'r', encoding='utf-8') as f:
data = json.load(f)
for rule_data in data.get("rules", []):
rule = Rule.from_dict(rule_data)
self._rules[rule.id] = rule
print(f"📂 已加载 {len(self._rules)} 条规则从 {self._storage_path}")
except Exception as e:
print(f"⚠️ 加载规则文件失败: {e}")
self._rules = {}
def _save_to_file(self):
"""保存规则到文件"""
try:
self._ensure_data_dir()
data = {
"version": "1.0",
"updated_at": datetime.now().isoformat(),
"rules": [rule.to_dict() for rule in self._rules.values()]
}
with open(self._storage_path, 'w', encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=2)
except Exception as e:
print(f"⚠️ 保存规则文件失败: {e}")
def add_rule(
self,
rule_type: str,
content: str,
category: Optional[str] = None,
priority: str = "normal"
) -> Dict[str, Any]:
"""
添加规则
Args:
rule_type: 规则类型 (must, must_not, prefer, security, architecture)
content: 规则内容
category: 分类标签
priority: 优先级 (high, normal, low)
Returns:
添加结果
"""
# 验证规则类型
if rule_type not in self.VALID_RULE_TYPES:
return {
"success": False,
"error": f"无效的规则类型: {rule_type}，有效类型: {self.VALID_RULE_TYPES}"
}
# 验证优先级
if priority not in self.VALID_PRIORITIES:
return {
"success": False,
"error": f"无效的优先级: {priority}，有效优先级: {self.VALID_PRIORITIES}"
}
# 检查是否已存在相同内容的规则
for rule in self._rules.values():
if rule.content == content and rule.rule_type == rule_type:
return {
"success": False,
"error": "规则已存在",
"existing_rule_id": rule.id
}
# 创建规则
rule_id = str(uuid.uuid4())[:8]
rule = Rule(
id=rule_id,
rule_type=rule_type,
content=content,
category=category,
priority=priority
)
self._rules[rule_id] = rule
self._save_to_file() # 持久化
return {
"success": True,
"rule": rule.to_dict()
}
def remove_rule(self, rule_id: str) -> Dict[str, Any]:
"""
移除规则
Args:
rule_id: 规则 ID
Returns:
移除结果
"""
if rule_id not in self._rules:
return {
"success": False,
"error": f"规则不存在: {rule_id}"
}
removed_rule = self._rules.pop(rule_id)
self._save_to_file() # 持久化
return {
"success": True,
"removed_rule": removed_rule.to_dict()
}
def get_rule(self, rule_id: str) -> Optional[Rule]:
"""获取单个规则"""
return self._rules.get(rule_id)
def list_rules(
self,
rule_type: Optional[str] = None,
category: Optional[str] = None,
priority: Optional[str] = None
) -> List[Dict[str, Any]]:
"""
列出规则
Args:
rule_type: 按类型过滤
category: 按分类过滤
priority: 按优先级过滤
Returns:
规则列表
"""
rules = list(self._rules.values())
# 过滤
if rule_type:
rules = [r for r in rules if r.rule_type == rule_type]
if category:
rules = [r for r in rules if r.category == category]
if priority:
rules = [r for r in rules if r.priority == priority]
# 按优先级排序
priority_order = {"high": 0, "normal": 1, "low": 2}
rules.sort(key=lambda r: priority_order.get(r.priority, 1))
return [r.to_dict() for r in rules]
def get_rules_by_type(self) -> Dict[str, List[str]]:
"""
按类型分组获取规则内容
Returns:
按类型分组的规则内容
"""
grouped = {
"must": [],
"must_not": [],
"prefer": [],
"security": [],
"architecture": []
}
for rule in self._rules.values():
if rule.rule_type in grouped:
grouped[rule.rule_type].append(rule.content)
return grouped
def clear_all(self) -> Dict[str, Any]:
"""清除所有规则"""
count = len(self._rules)
self._rules.clear()
self._save_to_file() # 持久化
return {
"success": True,
"cleared_count": count
}
def get_statistics(self) -> Dict[str, Any]:
"""获取规则统计"""
stats = {
"total": len(self._rules),
"by_type": {},
"by_priority": {},
"by_category": {}
}
for rule in self._rules.values():
# 按类型统计
stats["by_type"][rule.rule_type] = stats["by_type"].get(rule.rule_type, 0) + 1
# 按优先级统计
stats["by_priority"][rule.priority] = stats["by_priority"].get(rule.priority, 0) + 1
# 按分类统计
if rule.category:
stats["by_category"][rule.category] = stats["by_category"].get(rule.category, 0) + 1
return stats
def reload(self):
"""重新从文件加载规则"""
self._rules.clear()
self._load_from_file()