feat: Agent 多模式系统、视觉模型集成、LLM 能力分层与 P3 性能收尾
核心架构变更：
1. Agent 多模式系统替代 Coordinator
- 移除 src/agent/coordinator/（Coordinator Agent/Worker/Tools，946 行）
- 新建 src/agent/modes/：声明式模式抽象（AgentMode/ModeConfig/ToolSet）
- 三种内置模式：
- default：通用科研助手，零覆盖保持现有行为
- deep-research：16 步、启用思考、research 权限、系统性调研
- literature-reader：白名单工具、只读沙箱、结构化阅读
- ModeRegistry + ModeConfig 预设 + ToolSet 过滤 + 身份/原则覆盖
- AgentRuntime::with_mode() 统一入口，模式持久化到 session.mode 字段
- GET /api/chat/modes 提供模式列表给前端选择器
2. 视觉模型与图片分析
- 新增 analyze_image 工具（340 行）：本地/URL 图片 → 视觉模型流式分析
- LlmClient::analyze_image_stream()：SSE 增量实时推送
- 配置：LLM_VISION_MODEL / LLM_VISION_API_KEY / LLM_VISION_API_BASE
- 前端：粘贴/选择图片附件，重试时复用文件路径
- Service 层移除 /chat/rag 和 /chat/figure 端点，统一走 Agent SSE
- Body limit 提升至 100MB 适配大图上传
3. LLM 三级能力分层
- Tier 1 (Core) → Tier 2 (Medium) → Tier 3 (Fast)，级联回退
- medium_llm / fast_llm / vision_llm 注入 AppState
- 资产批量翻译 → Medium LLM + Semaphore(3) 并发控制
- 记忆提取/上下文压缩子代理 → Fast LLM
- SubAgentRunner::with_llm_client() 支持注入专用 LLM
4. 数据库与性能优化
- SQLite 启用 WAL + busy_timeout(10s) 处理并发写入
- RAG ingest：DELETE 合并为原子语句 + 批量事务写入
- Meta sync：save_paper_to_db_tx() 事务化批量插入
- 翻译词典：first_words HashSet 预过滤 + next_valid_index 跳跃优化
- read_file 不截断输出 + skip_persist 防止级联磁盘持久化
5. 工具系统增强
- ToolContext 增加 tool_call_id + max_output_chars
- ToolOutput 增加 skip_persist 标记
- TextDelta SSE 携带可选 tool_call_id 支持工具的流式输出
- ChatMessage::text() 辅助方法
This commit is contained in:
parent
cec4b8cf7b
commit
85b6429c30
22
.env.example
22
.env.example
@ -21,6 +21,23 @@ LLM_MODEL=gpt-4o-mini
|
||||
# 备用模型链（逗号分隔，按顺序轮换尝试，可选）
|
||||
# LLM_FALLBACK_CHAIN=gpt-4o-mini,deepseek-chat
|
||||
|
||||
# 中型大语言模型配置（用于文献翻译、RAG，默认级联回退至核心 LLM 配置）
|
||||
# LLM_MEDIUM_API_KEY=your_llm_medium_api_key_here
|
||||
# LLM_MEDIUM_API_BASE=https://api.openai.com/v1
|
||||
# LLM_MEDIUM_MODEL=gpt-4o-mini
|
||||
|
||||
# 快速大语言模型配置（用于记忆提取等后台辅助任务，默认级联回退至中型 LLM 配置）
|
||||
# LLM_FAST_API_KEY=your_llm_fast_api_key_here
|
||||
# LLM_FAST_API_BASE=https://api.openai.com/v1
|
||||
# LLM_FAST_MODEL=gpt-3.5-turbo
|
||||
|
||||
# ── 视觉模型（可选，配置后启用图片分析能力）──
|
||||
# 专用视觉模型名称。主 Agent 可为纯文本模型，图片分析通过 analyze_image 工具
|
||||
# 委托给此模型。用户上传图片时也会先用此模型预处理再交给 Agent。
|
||||
# LLM_VISION_MODEL=gpt-4o
|
||||
# LLM_VISION_API_KEY=your_vision_api_key_here（空时回退到 LLM_API_KEY）
|
||||
# LLM_VISION_API_BASE=https://api.openai.com/v1（空时回退到 LLM_API_BASE）
|
||||
|
||||
# 向量嵌入模型配置（未设置时默认回退到 LLM 的 API Key 和 Base）
|
||||
# EMBEDDING_API_KEY=your_embedding_api_key_here
|
||||
# EMBEDDING_API_BASE=https://api.openai.com/v1
|
||||
@ -92,11 +109,6 @@ LOG_DIR=./logs
|
||||
# 触发 snip 压缩的最大消息数（默认 50）
|
||||
# AGENT_MAX_MESSAGES=50
|
||||
|
||||
# ── P2 Coordinator Mode ──
|
||||
# Worker 最大并发数（默认 4）
|
||||
# AGENT_COORDINATOR_MAX_WORKERS=4
|
||||
# Worker 超时秒数（默认 300）
|
||||
# AGENT_COORDINATOR_WORKER_TIMEOUT=300
|
||||
|
||||
# ── 权限系统 ──
|
||||
# Agent 工具权限规则（逗号分隔，支持内容级匹配 "ToolName(pattern)"）
|
||||
|
||||
1
.gitignore
vendored
1
.gitignore
vendored
@ -20,3 +20,4 @@ library/MinerU/
|
||||
|
||||
libs/
|
||||
.claude
|
||||
.checkpoints
|
||||
|
||||
75
CLAUDE.md
75
CLAUDE.md
@ -69,11 +69,13 @@ src/
|
||||
│ └── helpers.rs # Shared DB helpers, format conversion, path validation
|
||||
├── agent/ # ReAct-based research agent (LLM-driven tool-use loop)
|
||||
│ ├── tools/ # AgentTool trait, ToolRegistry, 25+ tool implementations per domain file
|
||||
│ ├── runtime/ # ReAct loop, streaming, session/context, token budget, error recovery,
|
||||
│ │ # permission checker, file cache, system prompt assembly, circuit breaker
|
||||
│ ├── runtime/ # ReAct loop core: session, context, streaming, executor, token_budget,
|
||||
│ │ # permission, permission_profile, checkpoint, circuit_breaker, error_recovery,
|
||||
│ │ # hardline, partitioner, file_cache, system_prompt, finalize, denial_tracker
|
||||
│ ├── compact/ # Context compression (micro/auto/manual layers + collapse)
|
||||
│ ├── memory/ # Persistent memory manager: extraction, dedup, decay, age, guardrails
|
||||
│ ├── hooks.rs # Lifecycle events (PreToolUse/PostToolUse/Stop/etc.)
|
||||
│ ├── hooks/ # Lifecycle hooks: registry, dispatch, builtins, matcher, traits (PreToolUse/PostToolUse/Stop/etc.)
|
||||
│ ├── modes/ # Agent session modes: default, deep-research, literature-reader (identity/tools/config presets)
|
||||
│ ├── skills.rs # SkillRegistry: hot-loads SKILL.md files from skills/ directory
|
||||
│ ├── subagent.rs # Context-isolated sub-agent runner (subagent tool)
|
||||
│ ├── team/ # Multi-agent team: file-based inbox, lead/teammate coordination
|
||||
@ -121,7 +123,7 @@ The agent (`src/agent/`) implements a **ReAct** (Thought → Action → Observat
|
||||
|
||||
1. **`AgentRuntime`** (`runtime/mod.rs`) orchestrates the loop: session create/resume → context build → ReAct loop → finalize
|
||||
2. **Streaming**: LLM response is streamed via SSE (`AgentStreamEvent`) — thought, tool_call (with `id`), tool_result (with `tool_call_id`), text_delta, usage, error, done. Tool calls execute in parallel.
|
||||
3. **Tools**: Each tool implements `AgentTool` trait (name, description, JSON Schema parameters, execute). Core tools: read_file, grep_files, glob_files, run_bash, file_write, file_edit, search_papers, download_paper, parse_paper, get_paper_content, rag_search, query_target, save_note, todo_write, compress_context, load_skill, subagent, ask_user, save_memory. Plus background tools (bg_task_run, bg_task_check) and team tools (spawn_teammate, send_teammate_message, team_broadcast, check_team_inbox).
|
||||
3. **Tools**: Each tool implements `AgentTool` trait (name, description, JSON Schema parameters, execute). Core tools: read_file, grep_files, glob_files, run_bash, file_write, file_edit, search_papers, download_paper, parse_paper, get_paper_content, rag_search, query_target, save_note, todo_write, compress_context, load_skill, subagent, ask_user, save_memory, analyze_image (vision model, only when `LLM_VISION_MODEL` set). Plus background tools (bg_task_run, bg_task_check) and team tools (spawn_teammate, send_teammate_message, team_broadcast, check_team_inbox).
|
||||
4. **Thinking mode**: `enable_thinking` flag propagates from `AgentChatRequest` → `AgentConfig` → `ToolContext` → `LlmClient::chat_stream`. Only enabled for Qwen/DashScope backends; frontend-controlled via the `thinking` request field.
|
||||
5. **Tool call ID tracking**: LLM may not return tool_call IDs — `LlmClient` generates UUID fallbacks. `ToolCall` and `ToolResult` SSE events carry matching IDs for precise frontend pairing.
|
||||
6. **ToolContext** (`tools/mod.rs`): Injected into every tool execution — holds `app_state`, `session_id`, `sse_tx` (for intermediate events), `enable_thinking`, `read_file_state` (file cache for dedup), `silent` (sub-agents skip permission prompts).
|
||||
@ -162,6 +164,71 @@ The `obscura-inprocess` feature compiles `obscura-browser` and `obscura-net` dir
|
||||
|
||||
A 1MB+ bilingual astronomy terminology file loaded at startup into a Trie tree for longest-match glossary construction, used by the translation service to guide LLM translations with domain-accurate term mappings.
|
||||
|
||||
### Agent Session Modes (`src/agent/modes/`)
|
||||
|
||||
Session-level modes configure the agent's identity, tool access, step limits, thinking, and permissions at creation time. Modes are pure-data constants defined at compile time — adding a new mode requires no logic changes.
|
||||
|
||||
Three built-in modes (registered in `ModeRegistry::builtins()`):
|
||||
|
||||
| Mode | ID | Tools | Max Steps | Thinking | Permission Profile |
|
||||
|------|-----|-------|-----------|----------|---------------------|
|
||||
| 通用科研助手 | `default` | All (unrestricted) | 8 (default) | user-controlled | none |
|
||||
| 深度研究 | `deep-research` | All | 16 | forced on | `research` |
|
||||
| 文献阅读助手 | `literature-reader` | Allowlist (read-only + literature) | 6 | forced off | `readonly` |
|
||||
|
||||
Key types in `modes/mod.rs`:
|
||||
- **`AgentMode`** — static definition struct: id, name, description, icon, identity/principles overrides, extra sections, `ToolSet`, `ModeConfig`
|
||||
- **`ToolSet`** — `All`, `Allowlist(&[&str])`, or `Except(&[&str])` — constrains which tools are available
|
||||
- **`ModeConfig`** — optional overrides for `max_steps`, `enable_thinking`, `tool_timeout_secs`, `permission_profile`
|
||||
- **`ModeRegistry`** — holds `&'static AgentMode` references; `get(id)` for lookup
|
||||
|
||||
The mode is stored in the `agent_sessions.mode` column (migration `20260624000000_add_session_mode.sql`, defaults to `'default'`). The frontend `ResearchAgentPanel` exposes mode selection.
|
||||
|
||||
**Mode vs Skill**: Modes are session-level ("who am I"), Skills are task-level ("how do I do X"). Modes affect initialization; the ReAct loop itself is mode-agnostic.
|
||||
|
||||
### Permission System (`src/agent/runtime/permission.rs`)
|
||||
|
||||
Tool execution is gated by a priority-ordered rule chain: **Deny > Allow > Ask** (first match wins). Rules support content-level pattern matching (e.g., `run_bash(rm *)`).
|
||||
|
||||
**Permission modes** (set via `AGENT_PERMISSION_MODE` env or mode's `permission_profile`):
|
||||
- `default` — full rule chain evaluation
|
||||
- `accept_edits` — auto-allow `file_write`/`file_edit` within the working directory
|
||||
- `bypass` — skip all Ask checks (Deny rules still enforced)
|
||||
- `dont_ask` — convert all Ask to Deny
|
||||
|
||||
**Permission profiles** (`src/agent/runtime/permission_profile.rs`) are named presets loaded by `AGENT_PERMISSION_MODE` or a mode's `permission_profile` field. Profiles `research` and `readonly` are used by deep-research and literature-reader modes respectively.
|
||||
|
||||
**Configuration** (in `.env`):
|
||||
- `AGENT_PERMISSIONS_DENY` — comma-separated deny rules (e.g., `run_bash(rm *),run_bash(sudo *)`)
|
||||
- `AGENT_PERMISSIONS_ALLOW` — comma-separated allow rules
|
||||
- `AGENT_PERMISSIONS_ASK` — comma-separated ask rules
|
||||
- `AGENT_PERMISSION_MODE` — default/accept_edits/bypass/dont_ask
|
||||
|
||||
### Multi-Tier LLM Configuration
|
||||
|
||||
The system supports three LLM tiers with cascade fallback:
|
||||
|
||||
| Tier | Env Prefix | Purpose | Fallback |
|
||||
|------|-----------|---------|----------|
|
||||
| Primary | `LLM_` | Main agent reasoning | — |
|
||||
| Medium | `LLM_MEDIUM_` | Translation, RAG | Primary LLM config |
|
||||
| Fast | `LLM_FAST_` | Memory extraction, background tasks | Medium LLM config |
|
||||
|
||||
Each tier has `_API_KEY`, `_API_BASE`, `_MODEL` variants. Unset tiers cascade to the next tier down.
|
||||
|
||||
**Fallback chain**: `LLM_FALLBACK_CHAIN` (comma-separated model names) provides automatic model rotation on repeated 529 errors. `LLM_FALLBACK_MODEL` is a single backup model tried before the chain.
|
||||
|
||||
### Vision Model (`analyze_image` tool)
|
||||
|
||||
When `LLM_VISION_MODEL` env is set, the `analyze_image` tool (`src/agent/tools/astro/analyze_image.rs`) is registered. It delegates image analysis to a dedicated vision model, enabling the main agent to use a text-only model while still processing images. Supports local paths (relative to library dir) and HTTP(S) URLs. Results stream via SSE `TextDelta` events. Also supports `LLM_VISION_API_KEY` and `LLM_VISION_API_BASE` (fall back to primary LLM config).
|
||||
|
||||
### Auto Memory Extraction
|
||||
|
||||
Controlled by env vars:
|
||||
- `EXTRACT_MEMORY_ENABLED` (default `false`) — enables automatic memory extraction at session end and during compaction
|
||||
- `EXTRACT_MEMORY_THROTTLE_TURNS` (default `3`) — extract every N turns
|
||||
- `EXTRACT_MEMORY_MAX_STEPS` (default `3`) — sub-agent max steps for extraction
|
||||
|
||||
## Code Conventions
|
||||
|
||||
- Use `anyhow` for application errors, `thiserror` for library-style typed errors
|
||||
|
||||
@ -12,7 +12,8 @@ import {
|
||||
Brain, Settings, Eye, CheckCircle2, AlertTriangle,
|
||||
Send, Loader, Plus, Trash2, Compass, Clock, Square,
|
||||
BarChart3, ScrollText, Network, Rewind, RotateCcw,
|
||||
GitBranch, RefreshCw, Search, X, MessageSquare, BookOpen
|
||||
GitBranch, RefreshCw, Search, X, MessageSquare, BookOpen,
|
||||
PanelLeftClose, PanelLeftOpen, Paperclip
|
||||
} from 'lucide-react';
|
||||
import { AskUserQuestionCard } from './AskUserQuestionCard';
|
||||
import { PermissionRequestCard } from './PermissionRequestCard';
|
||||
@ -22,6 +23,7 @@ interface SessionSummary {
|
||||
session_id: string;
|
||||
title: string;
|
||||
model: string;
|
||||
mode: string;
|
||||
turn_count: number;
|
||||
summary?: string | null;
|
||||
created_at: string;
|
||||
@ -72,6 +74,7 @@ type ColorScheme = 'parent' | 'subagent';
|
||||
|
||||
interface ActiveTurn {
|
||||
question: string;
|
||||
imagePath?: string;
|
||||
timeline: TimelineItem[];
|
||||
finalAnswer: string;
|
||||
error?: string;
|
||||
@ -86,6 +89,7 @@ interface ProcessedTurn {
|
||||
turn_index: number;
|
||||
question: string;
|
||||
questionMessageId?: number;
|
||||
imagePath?: string;
|
||||
timeline: TimelineItem[];
|
||||
usage?: {
|
||||
prompt_tokens: number;
|
||||
@ -113,6 +117,7 @@ interface RetryResult {
|
||||
new_turn_index: number;
|
||||
deleted_count: number;
|
||||
session_id: string;
|
||||
image_path?: string | null;
|
||||
}
|
||||
|
||||
const safeSchema = {
|
||||
@ -180,6 +185,23 @@ function getToolDisplayName(name: string): string {
|
||||
}
|
||||
}
|
||||
|
||||
interface AgentMode {
|
||||
id: string;
|
||||
name: string;
|
||||
description: string;
|
||||
icon: string;
|
||||
}
|
||||
|
||||
const ICON_MAP: Record<string, React.ComponentType<{ className?: string }>> = {
|
||||
Brain,
|
||||
Compass,
|
||||
BookOpen
|
||||
};
|
||||
|
||||
function getIconComponent(iconName: string): React.ComponentType<{ className?: string }> {
|
||||
return ICON_MAP[iconName] || Brain;
|
||||
}
|
||||
|
||||
interface ResearchAgentPanelProps {
|
||||
showConfirm?: (message: string, onConfirm: () => void, title?: string) => void;
|
||||
showAlert?: (message: string, title?: string) => void;
|
||||
@ -192,10 +214,41 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
const [activeTurn, setActiveTurn] = useState<ActiveTurn | null>(null);
|
||||
const [streaming, setStreaming] = useState(false);
|
||||
const [input, setInput] = useState('');
|
||||
const [agentMode, setAgentMode] = useState('default');
|
||||
const [agentModes, setAgentModes] = useState<AgentMode[]>([]);
|
||||
const [thinking, setThinking] = useState(false);
|
||||
const [coordinatorMode, setCoordinatorMode] = useState(false);
|
||||
const [pendingImage, setPendingImage] = useState<{ data?: string; path?: string; mime_type: string; name: string } | null>(null);
|
||||
const fileInputRef = useRef<HTMLInputElement>(null);
|
||||
const [loadingSessions, setLoadingSessions] = useState(false);
|
||||
|
||||
// 将 File 转为 base64 并设置为 pendingImage
|
||||
const attachImage = (file: File) => {
|
||||
if (!file.type.startsWith('image/')) return;
|
||||
const reader = new FileReader();
|
||||
reader.onload = () => {
|
||||
const result = reader.result as string;
|
||||
const commaIdx = result.indexOf(',');
|
||||
const data = commaIdx > -1 ? result.substring(commaIdx + 1) : result;
|
||||
setPendingImage({ data, mime_type: file.type, name: file.name });
|
||||
};
|
||||
reader.readAsDataURL(file);
|
||||
};
|
||||
|
||||
// 粘贴事件：从剪贴板获取图片
|
||||
const handlePaste = (e: React.ClipboardEvent) => {
|
||||
const items = e.clipboardData?.items;
|
||||
if (!items) return;
|
||||
for (let i = 0; i < items.length; i++) {
|
||||
if (items[i].type.startsWith('image/')) {
|
||||
e.preventDefault();
|
||||
const file = items[i].getAsFile();
|
||||
if (file) attachImage(file);
|
||||
return;
|
||||
}
|
||||
}
|
||||
};
|
||||
const [loadingHistory, setLoadingHistory] = useState(false);
|
||||
const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
|
||||
|
||||
// 全文搜索历史记录状态
|
||||
const [searchQuery, setSearchQuery] = useState('');
|
||||
@ -266,6 +319,22 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
|
||||
useEffect(() => {
|
||||
fetchSessions(true);
|
||||
|
||||
const fetchModes = async () => {
|
||||
try {
|
||||
const res = await axios.get<AgentMode[]>('/api/chat/modes');
|
||||
setAgentModes(res.data);
|
||||
} catch (e) {
|
||||
console.error('获取 Agent 运行模式列表失败:', e);
|
||||
// Fallback in case of API failure for resilience
|
||||
setAgentModes([
|
||||
{ id: 'default', name: '默认', icon: 'Brain', description: '通用天体物理学研究助手，自主判断搜索、阅读或计算' },
|
||||
{ id: 'deep-research', name: '深度', icon: 'Compass', description: '多来源系统性文献调研与交叉验证，适合撰写综述' },
|
||||
{ id: 'literature-reader', name: '精读', icon: 'BookOpen', description: '专注论文精读、翻译与笔记提炼，强制只读安全模式' },
|
||||
]);
|
||||
}
|
||||
};
|
||||
fetchModes();
|
||||
}, []);
|
||||
|
||||
// 加载选中的会话历史消息
|
||||
@ -278,6 +347,10 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
const res = await axios.get<{ session: SessionSummary; messages: MessageRecord[] }>(`/api/chat/sessions/${sessionId}`);
|
||||
const newMessages = res.data.messages;
|
||||
|
||||
if (res.data.session && res.data.session.mode) {
|
||||
setAgentMode(res.data.session.mode);
|
||||
}
|
||||
|
||||
// 如果有成功回调，在更新 messages 之前立刻执行（如清除 activeTurn），保证在同一个 React 渲染批处理周期内完成
|
||||
if (onSuccess) {
|
||||
onSuccess();
|
||||
@ -531,7 +604,15 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
// 重新加载会话以移除已被硬删除的消息，然后自动触发发送
|
||||
await loadSessionHistory(currentSessionId, true);
|
||||
await fetchSessions();
|
||||
|
||||
|
||||
// 如果原消息附带了图片，通过 path 复用已有文件，避免重新上传
|
||||
if (res.data.image_path) {
|
||||
const ext = res.data.image_path.split('.').pop() || 'png';
|
||||
const mimeType = { jpg: 'image/jpeg', jpeg: 'image/jpeg', gif: 'image/gif', webp: 'image/webp', png: 'image/png' }[ext] || 'image/png';
|
||||
handleSend(questionText, { path: res.data.image_path, mime_type: mimeType, name: res.data.image_path.split('/').pop() || 'image' });
|
||||
return;
|
||||
}
|
||||
|
||||
// 触发自动重发
|
||||
handleSend(questionText);
|
||||
} catch (e: any) {
|
||||
@ -555,7 +636,8 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
};
|
||||
|
||||
// 发送消息与流式响应处理
|
||||
const handleSend = async (questionText: string) => {
|
||||
// imageOverride: 重试时直接传入图片数据，绕过 React 异步状态更新
|
||||
const handleSend = async (questionText: string, imageOverride?: { data?: string; path?: string; mime_type: string; name: string } | null) => {
|
||||
if (!questionText.trim() || streaming) return;
|
||||
|
||||
if (!currentSessionId) {
|
||||
@ -563,11 +645,16 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
}
|
||||
|
||||
setInput('');
|
||||
const image = imageOverride !== undefined ? imageOverride : pendingImage;
|
||||
setPendingImage(null);
|
||||
setStreaming(true);
|
||||
setShouldAutoScroll(true);
|
||||
|
||||
const active: ActiveTurn = {
|
||||
question: questionText,
|
||||
imagePath: image
|
||||
? (image.path ? `/api/files/${image.path}` : `data:${image.mime_type};base64,${image.data}`)
|
||||
: undefined,
|
||||
timeline: [],
|
||||
finalAnswer: '',
|
||||
};
|
||||
@ -582,8 +669,13 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
body: JSON.stringify({
|
||||
question: questionText,
|
||||
session_id: currentSessionId,
|
||||
mode: agentMode,
|
||||
thinking,
|
||||
coordinator_mode: coordinatorMode,
|
||||
...(image ? {
|
||||
image: image.path
|
||||
? { path: image.path, mime_type: image.mime_type }
|
||||
: { data: image.data, mime_type: image.mime_type },
|
||||
} : {}),
|
||||
}),
|
||||
});
|
||||
|
||||
@ -756,6 +848,8 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
});
|
||||
break;
|
||||
case 'tool_result':
|
||||
// 自动展开结果区域
|
||||
setExpandedResults(prev => ({ ...prev, [event.tool_call_id]: true }));
|
||||
setActiveTurn(prev => {
|
||||
if (!prev) return prev;
|
||||
const tcId = event.tool_call_id;
|
||||
@ -788,10 +882,11 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
}
|
||||
}
|
||||
// 父代理工具结果（正常路径）
|
||||
// 始终更新 result，即使已有流式输出的部分内容——最终 tool_result 是完整答案
|
||||
return {
|
||||
...prev,
|
||||
timeline: prev.timeline.map(t => {
|
||||
if (t.type === 'tool_call' && t.id === tcId && !t.result) {
|
||||
if (t.type === 'tool_call' && t.id === tcId) {
|
||||
return {
|
||||
...t,
|
||||
result: { output: event.output, isError: event.is_error },
|
||||
@ -803,8 +898,42 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
});
|
||||
break;
|
||||
case 'text_delta':
|
||||
// 如果带 tool_call_id，自动展开对应工具的结果区域以显示流式输出
|
||||
if (event.tool_call_id) {
|
||||
setExpandedResults(prev => ({ ...prev, [event.tool_call_id]: true }));
|
||||
}
|
||||
setActiveTurn(prev => {
|
||||
if (!prev) return prev;
|
||||
// 如果带有 tool_call_id，流式输出到对应工具的结果区域
|
||||
if (event.tool_call_id) {
|
||||
const tcId = event.tool_call_id;
|
||||
// 检查 timeline 中是否有匹配的 tool_call
|
||||
const hasToolCall = prev.timeline.some(
|
||||
t => t.type === 'tool_call' && 'id' in t && t.id === tcId
|
||||
);
|
||||
if (hasToolCall) {
|
||||
return {
|
||||
...prev,
|
||||
timeline: prev.timeline.map(t => {
|
||||
if (t.type === 'tool_call' && t.id === tcId) {
|
||||
const prevOutput = t.result?.output || '';
|
||||
return {
|
||||
...t,
|
||||
result: {
|
||||
output: prevOutput + event.content,
|
||||
isError: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
return t;
|
||||
}),
|
||||
};
|
||||
}
|
||||
// 工具调用尚未出现（TTL 场景），暂存到 finalAnswer 不处理
|
||||
// 等 tool_call + tool_result 出现后，tool_result 会覆盖
|
||||
return prev;
|
||||
}
|
||||
// 无 tool_call_id → 主回答文本流
|
||||
return {
|
||||
...prev,
|
||||
finalAnswer: prev.finalAnswer + event.content,
|
||||
@ -1065,8 +1194,13 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
</button>
|
||||
</div>
|
||||
{isResultExpanded ? (
|
||||
<div className="font-mono text-[10px] bg-slate-50 border border-slate-200 rounded-lg p-2.5 text-slate-700 overflow-x-auto whitespace-pre-wrap select-text leading-relaxed max-h-96">
|
||||
{item.result!.output}
|
||||
<div className="text-xs bg-slate-50 border border-slate-200 rounded-lg p-3 text-slate-700 overflow-auto select-text leading-relaxed max-h-[60vh] prose prose-sm max-w-none prose-headings:text-slate-800 prose-code:text-sky-700 prose-pre:bg-slate-100 prose-pre:text-xs prose-img:rounded-lg">
|
||||
<ReactMarkdown
|
||||
remarkPlugins={[remarkMath, remarkGfm]}
|
||||
rehypePlugins={[rehypeRaw, [rehypeSanitize, safeSchema], rehypeKatex]}
|
||||
>
|
||||
{item.result!.output}
|
||||
</ReactMarkdown>
|
||||
</div>
|
||||
) : (
|
||||
<div className="text-[10px] text-slate-400 italic font-medium">
|
||||
@ -1111,20 +1245,29 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
<div className="w-full flex-1 flex overflow-hidden bg-slate-100 rounded-2xl border border-slate-200 h-[calc(100vh-130px)] shadow-xs">
|
||||
|
||||
{/* 左侧会话列表侧栏 */}
|
||||
<div className="w-64 bg-slate-50 border-r border-slate-200 flex flex-col justify-between shrink-0 select-none">
|
||||
<div className={`transition-all duration-300 ease-in-out ${sidebarCollapsed ? 'w-0 overflow-hidden opacity-0 border-r-0' : 'w-64 border-r border-slate-200'} bg-slate-50 flex flex-col justify-between shrink-0 select-none`}>
|
||||
<div className="flex flex-col min-h-0 flex-1">
|
||||
<div className="p-4 border-b border-slate-200 bg-white flex items-center justify-between shrink-0">
|
||||
<span className="text-xs font-extrabold text-slate-800 tracking-wider flex items-center gap-1.5">
|
||||
<Clock className="w-3.5 h-3.5 text-sky-600" />
|
||||
<span>智能科研研讨</span>
|
||||
{!sidebarCollapsed && <span>智能科研研讨</span>}
|
||||
</span>
|
||||
<button
|
||||
onClick={handleNewSession}
|
||||
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-600 hover:text-slate-800 transition-all cursor-pointer shadow-2xs hover:scale-105"
|
||||
title="新建会话"
|
||||
>
|
||||
<Plus className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
<div className="flex items-center gap-1 shrink-0">
|
||||
<button
|
||||
onClick={handleNewSession}
|
||||
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-600 hover:text-slate-800 transition-all cursor-pointer shadow-2xs hover:scale-105"
|
||||
title="新建会话"
|
||||
>
|
||||
<Plus className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
<button
|
||||
onClick={() => setSidebarCollapsed(true)}
|
||||
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-500 hover:text-slate-700 transition-all cursor-pointer"
|
||||
title="收起侧栏"
|
||||
>
|
||||
<PanelLeftClose className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* 搜索框区域 */}
|
||||
@ -1274,15 +1417,26 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
<div className="flex-1 flex flex-col overflow-hidden bg-white relative">
|
||||
{/* 会话头部 */}
|
||||
<div className="px-5 py-4 border-b border-slate-200 flex items-center justify-between shrink-0 bg-white">
|
||||
<div className="min-w-0 pr-4">
|
||||
<h3 className="text-xs font-bold text-slate-800 tracking-wide line-clamp-1">
|
||||
{currentSessionId
|
||||
? (sessions.find(s => s.session_id === currentSessionId)?.title || '未命名会话')
|
||||
: '探索性科研研讨'}
|
||||
</h3>
|
||||
<p className="text-[10px] text-slate-400 font-semibold mt-0.5">
|
||||
科研智能体 (ReAct 循环：Thought → Action → Observation)
|
||||
</p>
|
||||
<div className="min-w-0 pr-4 flex items-center gap-2">
|
||||
{sidebarCollapsed && (
|
||||
<button
|
||||
onClick={() => setSidebarCollapsed(false)}
|
||||
className="p-1.5 rounded-lg border border-slate-200 bg-white hover:bg-slate-50 text-slate-500 hover:text-slate-700 transition-all cursor-pointer mr-1 shrink-0"
|
||||
title="展开侧栏"
|
||||
>
|
||||
<PanelLeftOpen className="w-4 h-4" />
|
||||
</button>
|
||||
)}
|
||||
<div>
|
||||
<h3 className="text-xs font-bold text-slate-800 tracking-wide line-clamp-1">
|
||||
{currentSessionId
|
||||
? (sessions.find(s => s.session_id === currentSessionId)?.title || '未命名会话')
|
||||
: '探索性科研研讨'}
|
||||
</h3>
|
||||
<p className="text-[10px] text-slate-400 font-semibold mt-0.5">
|
||||
科研智能体 (ReAct 循环：Thought → Action → Observation)
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
<div className="flex items-center gap-1.5">
|
||||
{/* 指标面板按钮 */}
|
||||
@ -1358,7 +1512,7 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
<div
|
||||
ref={scrollContainerRef}
|
||||
onScroll={handleScroll}
|
||||
className="flex-1 overflow-y-auto p-5 space-y-6 bg-slate-50/50 scrollbar-thin"
|
||||
className="flex-1 overflow-y-auto p-5 pb-36 space-y-6 bg-slate-50/50 scrollbar-thin"
|
||||
>
|
||||
|
||||
{loadingHistory ? (
|
||||
@ -1411,8 +1565,11 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
<Rewind className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
)}
|
||||
<div className="rounded-2xl px-4 py-3 text-xs leading-relaxed font-semibold shadow-2xs border bg-sky-600 text-white border-sky-600 select-text flex-1">
|
||||
{turn.question}
|
||||
<div className="rounded-2xl px-4 py-3 text-xs leading-relaxed font-semibold shadow-2xs border bg-sky-600 text-white border-sky-600 select-text flex-1 space-y-2">
|
||||
{turn.imagePath && (
|
||||
<img src={turn.imagePath} alt="用户上传的图片" className="max-h-48 rounded-lg border border-sky-500/30" />
|
||||
)}
|
||||
<div>{turn.question}</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@ -1439,8 +1596,11 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
{/* 当前提问 */}
|
||||
<div className="flex flex-col items-end space-y-1">
|
||||
<span className="text-[10px] font-bold text-slate-400 px-1">我</span>
|
||||
<div className="max-w-[85%] rounded-2xl px-4 py-3 text-xs leading-relaxed font-semibold shadow-2xs border bg-sky-600 text-white border-sky-600">
|
||||
{activeTurn.question}
|
||||
<div className="max-w-[85%] rounded-2xl px-4 py-3 text-xs leading-relaxed font-semibold shadow-2xs border bg-sky-600 text-white border-sky-600 space-y-2">
|
||||
{activeTurn.imagePath && (
|
||||
<img src={activeTurn.imagePath} alt="用户上传的图片" className="max-h-48 rounded-lg border border-sky-500/30" />
|
||||
)}
|
||||
<div>{activeTurn.question}</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@ -1534,65 +1694,137 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
e.preventDefault();
|
||||
handleSend(input);
|
||||
}}
|
||||
className="p-4 border-t border-slate-200 bg-white shrink-0"
|
||||
className="absolute bottom-0 left-0 right-0 p-4 bg-transparent pointer-events-none shrink-0"
|
||||
>
|
||||
<div className="flex gap-2.5 relative items-center max-w-4xl mx-auto">
|
||||
{/* 思考模式开关 */}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setThinking(!thinking)}
|
||||
disabled={streaming}
|
||||
title={thinking ? '思考模式已开启（启用 LLM 推理过程）' : '思考模式已关闭（点击开启）'}
|
||||
className={`p-2.5 rounded-xl border text-xs font-bold transition-all cursor-pointer flex items-center gap-1.5 shrink-0 ${
|
||||
thinking
|
||||
? 'bg-purple-50 border-purple-300 text-purple-700 shadow-2xs'
|
||||
: 'bg-slate-50 border-slate-250 text-slate-400 hover:text-purple-500 hover:border-purple-200'
|
||||
} disabled:opacity-60`}
|
||||
>
|
||||
<Brain className={`w-4 h-4 ${thinking ? 'text-purple-500' : ''}`} />
|
||||
<span className="hidden sm:inline">{thinking ? '思考中' : '思考'}</span>
|
||||
</button>
|
||||
{/* 协调者模式开关 */}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setCoordinatorMode(!coordinatorMode)}
|
||||
disabled={streaming}
|
||||
title={coordinatorMode ? '协调者模式已开启（委托子智能体执行）' : '协调者模式已关闭（单智能体直接执行）'}
|
||||
className={`p-2.5 rounded-xl border text-xs font-bold transition-all cursor-pointer flex items-center gap-1.5 shrink-0 ${
|
||||
coordinatorMode
|
||||
? 'bg-sky-50 border-sky-300 text-sky-700 shadow-2xs'
|
||||
: 'bg-slate-50 border-slate-250 text-slate-400 hover:text-sky-500 hover:border-sky-200'
|
||||
} disabled:opacity-60`}
|
||||
>
|
||||
<Network className={`w-4 h-4 ${coordinatorMode ? 'text-sky-500' : ''}`} />
|
||||
<span className="hidden sm:inline">{coordinatorMode ? '协调中' : '协调'}</span>
|
||||
</button>
|
||||
<input
|
||||
type="text"
|
||||
<div className="max-w-4xl mx-auto bg-slate-50 border border-slate-200/60 rounded-2xl p-2.5 focus-within:bg-white focus-within:border-sky-500 focus-within:ring-1 focus-within:ring-sky-500/10 transition-all pointer-events-auto shadow-lg">
|
||||
<textarea
|
||||
value={input}
|
||||
onChange={(e) => setInput(e.target.value)}
|
||||
onKeyDown={(e) => {
|
||||
if (e.key === 'Enter' && !e.shiftKey) {
|
||||
e.preventDefault();
|
||||
handleSend(input);
|
||||
}
|
||||
}}
|
||||
disabled={streaming}
|
||||
placeholder="向科研智能体提问（可进行 ADS 检索、文献结构解析、SIMBAD 天体查询等）..."
|
||||
className="flex-1 bg-slate-50 border border-slate-250 rounded-xl text-xs text-slate-900 placeholder-slate-400 pl-4.5 pr-12 py-3.5 focus:outline-none focus:bg-white focus:border-sky-500 focus:ring-1 focus:ring-sky-500/10 leading-relaxed font-semibold transition-all disabled:opacity-60"
|
||||
placeholder="向科研智能体提问（可粘贴或上传图片，配合深度研究模式分析图表）..."
|
||||
rows={2}
|
||||
onPaste={handlePaste}
|
||||
className="w-full bg-transparent resize-none border-none outline-none focus:outline-none focus:ring-0 text-xs text-slate-900 placeholder-slate-400 leading-relaxed font-semibold min-h-[44px] max-h-40"
|
||||
/>
|
||||
{streaming ? (
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleStop}
|
||||
className="absolute right-2 p-2 rounded-xl bg-rose-600 hover:bg-rose-700 text-white transition-colors cursor-pointer flex items-center justify-center"
|
||||
title="手动停止执行"
|
||||
>
|
||||
<Square className="w-4 h-4 fill-white" />
|
||||
</button>
|
||||
) : (
|
||||
<button
|
||||
type="submit"
|
||||
disabled={!input.trim()}
|
||||
className="absolute right-2 p-2 rounded-xl bg-sky-600 hover:bg-sky-700 text-white disabled:bg-slate-200 disabled:text-slate-400 transition-colors cursor-pointer flex items-center justify-center"
|
||||
>
|
||||
<Send className="w-4 h-4" />
|
||||
</button>
|
||||
{/* 图片预览 */}
|
||||
{pendingImage && (
|
||||
<div className="flex items-center gap-2 px-1 pb-2">
|
||||
<img
|
||||
src={pendingImage.path ? `/api/files/${pendingImage.path}` : `data:${pendingImage.mime_type};base64,${pendingImage.data}`}
|
||||
alt={pendingImage.name}
|
||||
className="h-16 rounded-lg border border-slate-200 object-cover"
|
||||
/>
|
||||
<span className="text-[10px] text-slate-500 truncate max-w-[200px]">{pendingImage.name}</span>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setPendingImage(null)}
|
||||
className="p-0.5 rounded-full bg-slate-200 hover:bg-slate-300 text-slate-500 transition-colors"
|
||||
>
|
||||
<X className="w-3 h-3" />
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
<div className="flex justify-between items-center mt-2.5 pt-2 border-t border-slate-100/80">
|
||||
<div className="flex flex-wrap gap-2 items-center">
|
||||
{/* 图片上传按钮 */}
|
||||
<input
|
||||
ref={fileInputRef}
|
||||
type="file"
|
||||
accept="image/*"
|
||||
className="hidden"
|
||||
onChange={(e) => {
|
||||
const file = e.target.files?.[0];
|
||||
if (file) attachImage(file);
|
||||
e.target.value = ''; // 允许重复选择同一文件
|
||||
}}
|
||||
/>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => fileInputRef.current?.click()}
|
||||
disabled={streaming}
|
||||
title="上传或粘贴图片（也可直接 Ctrl+V 粘贴）"
|
||||
className={`p-1.5 rounded-lg border text-[10px] font-bold transition-all cursor-pointer flex items-center gap-1 ${
|
||||
pendingImage
|
||||
? 'bg-sky-50 border-sky-300 text-sky-700'
|
||||
: 'bg-white border-slate-200 text-slate-400 hover:text-sky-600 hover:border-sky-200'
|
||||
} disabled:opacity-60`}
|
||||
>
|
||||
<Paperclip className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
{/* Agent 模式选择器 */}
|
||||
<div className="flex gap-1 bg-slate-100 p-0.5 rounded-lg border border-slate-200">
|
||||
{agentModes.map((m) => {
|
||||
const IconComp = getIconComponent(m.icon);
|
||||
const isSelected = agentMode === m.id;
|
||||
return (
|
||||
<button
|
||||
key={m.id}
|
||||
type="button"
|
||||
onClick={() => setAgentMode(m.id)}
|
||||
disabled={streaming}
|
||||
title={m.description}
|
||||
className={`px-2 py-1.5 rounded-md text-[10px] font-extrabold transition-all border cursor-pointer flex items-center gap-1 shrink-0 ${
|
||||
isSelected
|
||||
? 'bg-white text-indigo-700 shadow-xs border-slate-200/50'
|
||||
: 'bg-transparent border-transparent text-slate-400 hover:text-indigo-600'
|
||||
} disabled:opacity-60`}
|
||||
>
|
||||
<IconComp className={`w-3.5 h-3.5 ${isSelected ? 'text-indigo-500' : ''}`} />
|
||||
<span className="hidden sm:inline">{m.name}</span>
|
||||
</button>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
|
||||
{/* 思考模式开关 — 仅默认模式允许用户选择，其他模式由代码固定 */}
|
||||
{agentMode === 'default' && (
|
||||
<>
|
||||
<div className="h-4 w-px bg-slate-200 mx-1 hidden sm:block" />
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setThinking(!thinking)}
|
||||
disabled={streaming}
|
||||
title={thinking ? '思考模式已开启（启用 LLM 推理过程）' : '思考模式已关闭（点击开启）'}
|
||||
className={`px-2.5 py-1.5 rounded-lg border text-[10px] font-bold transition-all cursor-pointer flex items-center gap-1 shrink-0 ${
|
||||
thinking
|
||||
? 'bg-sky-50 border-sky-205 text-sky-700 shadow-3xs'
|
||||
: 'bg-white border-slate-200 text-slate-400 hover:text-sky-600 hover:border-sky-200'
|
||||
} disabled:opacity-60`}
|
||||
>
|
||||
<Brain className={`w-3.5 h-3.5 ${thinking ? 'text-sky-500' : ''}`} />
|
||||
<span className="hidden sm:inline">思考</span>
|
||||
</button>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div>
|
||||
{streaming ? (
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleStop}
|
||||
className="p-2 rounded-xl bg-rose-600 hover:bg-rose-700 text-white transition-colors cursor-pointer flex items-center justify-center shadow-xs"
|
||||
title="手动停止执行"
|
||||
>
|
||||
<Square className="w-3.5 h-3.5 fill-white" />
|
||||
</button>
|
||||
) : (
|
||||
<button
|
||||
type="submit"
|
||||
disabled={!input.trim()}
|
||||
className="p-2 rounded-xl bg-sky-600 hover:bg-sky-700 text-white disabled:bg-slate-100 disabled:text-slate-350 transition-colors cursor-pointer flex items-center justify-center shadow-xs"
|
||||
>
|
||||
<Send className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</form>
|
||||
</div>
|
||||
@ -1647,6 +1879,10 @@ function groupMessagesIntoTurns(messages: MessageRecord[]): ProcessedTurn[] {
|
||||
if (msg.role === 'user') {
|
||||
turn.question = msg.content;
|
||||
turn.questionMessageId = msg.id;
|
||||
// 提取用户消息中的图片路径
|
||||
if (msg.metadata?.image_path) {
|
||||
turn.imagePath = `/api/files/${msg.metadata.image_path}`;
|
||||
}
|
||||
} else if (msg.role === 'assistant') {
|
||||
const stepNum = msg.step_index;
|
||||
const hasToolCalls = msg.tool_calls && msg.tool_calls.length > 0;
|
||||
@ -1668,7 +1904,7 @@ function groupMessagesIntoTurns(messages: MessageRecord[]): ProcessedTurn[] {
|
||||
parsedArgs = typeof tc.function.arguments === 'string'
|
||||
? JSON.parse(tc.function.arguments)
|
||||
: tc.function.arguments;
|
||||
} catch (e) { /* ignore parse errors */ }
|
||||
} catch { /* ignore parse errors */ }
|
||||
turn.timeline.push({
|
||||
type: 'tool_call',
|
||||
step: stepNum,
|
||||
@ -1749,7 +1985,7 @@ function groupMessagesIntoTurns(messages: MessageRecord[]): ProcessedTurn[] {
|
||||
parsedArgs = typeof tc.function.arguments === 'string'
|
||||
? JSON.parse(tc.function.arguments)
|
||||
: tc.function.arguments;
|
||||
} catch (e) { /* ignore */ }
|
||||
} catch { /* ignore */ }
|
||||
children.push({
|
||||
type: 'tool_call',
|
||||
step: stepNum,
|
||||
|
||||
@ -135,20 +135,18 @@ export function AIAssistantPanel({
|
||||
const handleSend = async (questionText: string) => {
|
||||
if (!questionText.trim() || loading) return;
|
||||
|
||||
const figure = pendingFigure;
|
||||
setError(null);
|
||||
const userMsg: Message = {
|
||||
sender: 'user',
|
||||
text: questionText,
|
||||
imageUrl: pendingFigure?.url
|
||||
imageUrl: figure?.url || undefined
|
||||
};
|
||||
setMessages(prev => [...prev, userMsg]);
|
||||
setInput('');
|
||||
setLoading(true);
|
||||
setShouldAutoScroll(true);
|
||||
|
||||
const isFigureQuery = !!pendingFigure;
|
||||
const currentFigurePath = pendingFigure?.path;
|
||||
|
||||
if (onClearPendingFigure) {
|
||||
onClearPendingFigure();
|
||||
}
|
||||
@ -163,26 +161,57 @@ export function AIAssistantPanel({
|
||||
setMessages(prev => [...prev, aiMsgPlaceholder]);
|
||||
|
||||
try {
|
||||
if (isFigureQuery && currentFigurePath) {
|
||||
// 图表多模态分析依然调用特定接口
|
||||
const res = await axios.post<{ answer: string }>('/api/chat/figure', {
|
||||
bibcode,
|
||||
image_path: currentFigurePath,
|
||||
question: questionText
|
||||
});
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
const last = next[next.length - 1];
|
||||
if (last && last.sender === 'ai') {
|
||||
last.text = res.data.answer;
|
||||
// 如果有待处理的图片，先获取并转为 base64
|
||||
let imagePayload: { data: string; mime_type: string } | undefined = undefined;
|
||||
if (figure) {
|
||||
try {
|
||||
const imgResponse = await fetch(figure.url);
|
||||
if (!imgResponse.ok) {
|
||||
throw new Error(`Failed to fetch image: ${imgResponse.statusText}`);
|
||||
}
|
||||
return next;
|
||||
});
|
||||
setLoading(false);
|
||||
} else {
|
||||
// 调用智能体流式对话接口（带文献上下文提示，确保 Agent 优先分析当前阅读的文献）
|
||||
const contextPrefix = `[当前正在阅读文献: ${bibcode}] `;
|
||||
const requestQuery = questionText.includes(bibcode) ? questionText : `${contextPrefix}${questionText}`;
|
||||
const blob = await imgResponse.blob();
|
||||
|
||||
const base64Res = await new Promise<{ data: string; mimeType: string }>((resolve, reject) => {
|
||||
const reader = new FileReader();
|
||||
reader.readAsDataURL(blob);
|
||||
reader.onloadend = () => {
|
||||
const result = reader.result as string;
|
||||
const commaIdx = result.indexOf(',');
|
||||
if (commaIdx === -1) {
|
||||
reject(new Error('Invalid data URL'));
|
||||
return;
|
||||
}
|
||||
const data = result.substring(commaIdx + 1);
|
||||
const match = result.match(/^data:([^;]+);base64,/);
|
||||
const mimeType = match ? match[1] : blob.type || 'image/png';
|
||||
resolve({ data, mimeType });
|
||||
};
|
||||
reader.onerror = reject;
|
||||
});
|
||||
|
||||
imagePayload = {
|
||||
data: base64Res.data,
|
||||
mime_type: base64Res.mimeType
|
||||
};
|
||||
} catch (imgErr) {
|
||||
console.error('Failed to process image attachment:', imgErr);
|
||||
setError('获取或处理选中的图表失败，请重试');
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
// 移除用户和 AI 占位消息
|
||||
if (next.length >= 2) {
|
||||
return next.slice(0, -2);
|
||||
}
|
||||
return next;
|
||||
});
|
||||
setLoading(false);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// 调用智能体流式对话接口（带文献上下文提示，确保 Agent 优先分析当前阅读的文献）
|
||||
const contextPrefix = `[当前正在阅读文献: ${bibcode}] `;
|
||||
const requestQuery = questionText.includes(bibcode) ? questionText : `${contextPrefix}${questionText}`;
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/chat/agent', {
|
||||
@ -193,6 +222,8 @@ export function AIAssistantPanel({
|
||||
body: JSON.stringify({
|
||||
question: requestQuery,
|
||||
session_id: sessionId,
|
||||
mode: 'literature-reader',
|
||||
image: imagePayload,
|
||||
}),
|
||||
});
|
||||
|
||||
@ -271,7 +302,7 @@ export function AIAssistantPanel({
|
||||
const last = next[next.length - 1];
|
||||
if (last && last.sender === 'ai') {
|
||||
if (!last.steps) last.steps = [];
|
||||
let stepObj = last.steps.find(s => s.step === event.step && s.type === 'tool_call');
|
||||
const stepObj = last.steps.find(s => s.step === event.step && s.type === 'tool_call');
|
||||
if (stepObj) {
|
||||
stepObj.label = `调用: ${getToolDisplayName(event.name)} (已获取观测数据)`;
|
||||
stepObj.isError = event.is_error;
|
||||
@ -337,26 +368,17 @@ export function AIAssistantPanel({
|
||||
setLoading(false);
|
||||
|
||||
} catch (streamErr) {
|
||||
console.warn('智能体流式对话失败，正在尝试回退至旧版本 RAG 单次问答...', streamErr);
|
||||
|
||||
// 降级回退 (Deprecated Fallback) —— 调用原始单次 /api/chat/rag
|
||||
const res = await axios.post<{ answer: string; sources: RetrievalSource[] }>('/api/chat/rag', {
|
||||
question: questionText,
|
||||
top_k: 5
|
||||
});
|
||||
|
||||
console.error('智能体流式对话失败:', streamErr);
|
||||
setError('智能体流式对话失败，请稍后重试');
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
const last = next[next.length - 1];
|
||||
if (last && last.sender === 'ai') {
|
||||
last.text = res.data.answer;
|
||||
last.sources = res.data.sources;
|
||||
if (next.length > 0 && next[next.length - 1].sender === 'ai' && !next[next.length - 1].text) {
|
||||
next.pop();
|
||||
}
|
||||
return next;
|
||||
});
|
||||
setLoading(false);
|
||||
}
|
||||
}
|
||||
} catch (err: any) {
|
||||
console.error('问答请求失败:', err);
|
||||
setError(err.response?.data?.error || err.message || '网络请求错误，请稍后重试');
|
||||
@ -396,15 +418,15 @@ export function AIAssistantPanel({
|
||||
<div
|
||||
ref={scrollContainerRef}
|
||||
onScroll={handleScroll}
|
||||
className="flex-1 overflow-y-auto p-4 space-y-4 min-h-0"
|
||||
className="flex-1 overflow-y-auto p-4 pb-28 space-y-4 min-h-0"
|
||||
>
|
||||
{messages.length === 0 ? (
|
||||
<div className="py-6 space-y-6">
|
||||
<div className="text-center space-y-2 max-w-sm mx-auto">
|
||||
<Compass className="w-10 h-10 mx-auto text-sky-500 opacity-60" />
|
||||
<h3 className="text-xs font-bold text-slate-800">探索馆藏文献知识库</h3>
|
||||
<h3 className="text-xs font-bold text-slate-800">文献阅读助手</h3>
|
||||
<p className="text-[11px] text-slate-500 leading-relaxed font-semibold">
|
||||
基于 ReAct 框架的智能科研体。支持自动读取已解析的 Markdown 文献，在向量库进行 RAG 查询以及天体参数检索。
|
||||
专注文献精读与理解的 AI 助手。支持语义检索、逐段解读、公式分析，严格只读模式保障文献安全。
|
||||
</p>
|
||||
</div>
|
||||
|
||||
@ -582,35 +604,43 @@ export function AIAssistantPanel({
|
||||
e.preventDefault();
|
||||
handleSend(input);
|
||||
}}
|
||||
className="p-3 border-t border-slate-200 bg-white shrink-0"
|
||||
className="absolute bottom-0 left-0 right-0 p-3 bg-transparent pointer-events-none shrink-0"
|
||||
>
|
||||
<div className="flex gap-2 relative items-center">
|
||||
<input
|
||||
type="text"
|
||||
<div className="bg-slate-50 border border-slate-200/60 rounded-2xl p-2.5 focus-within:bg-white focus-within:border-sky-500 focus-within:ring-1 focus-within:ring-sky-500/10 transition-all relative pointer-events-auto shadow-md">
|
||||
<textarea
|
||||
value={input}
|
||||
onChange={(e) => setInput(e.target.value)}
|
||||
onKeyDown={(e) => {
|
||||
if (e.key === 'Enter' && !e.shiftKey) {
|
||||
e.preventDefault();
|
||||
handleSend(input);
|
||||
}
|
||||
}}
|
||||
disabled={loading}
|
||||
placeholder={pendingFigure ? "针对选中图表提问..." : "向 AI 馆藏助手提问..."}
|
||||
className="flex-1 bg-slate-50 border border-slate-200 rounded-xl text-xs text-slate-900 placeholder-slate-400 pl-3 pr-10 py-2.5 focus:outline-none focus:bg-white focus:border-sky-500 focus:ring-1 focus:ring-sky-500/10 leading-relaxed font-semibold transition-all disabled:opacity-60"
|
||||
rows={2}
|
||||
className="w-full bg-transparent resize-none border-none outline-none focus:outline-none focus:ring-0 text-xs text-slate-900 placeholder-slate-400 pr-8 leading-relaxed font-semibold min-h-[36px] max-h-32"
|
||||
/>
|
||||
{loading ? (
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleStop}
|
||||
className="absolute right-1.5 p-1.5 rounded-lg bg-rose-600 hover:bg-rose-700 text-white transition-colors cursor-pointer flex items-center justify-center"
|
||||
title="手动停止执行"
|
||||
>
|
||||
<Square className="w-3.5 h-3.5 fill-white" />
|
||||
</button>
|
||||
) : (
|
||||
<button
|
||||
type="submit"
|
||||
disabled={!input.trim()}
|
||||
className="absolute right-1.5 p-1.5 rounded-lg bg-sky-600 hover:bg-sky-700 text-white disabled:bg-slate-200 disabled:text-slate-400 transition-colors cursor-pointer flex items-center justify-center"
|
||||
>
|
||||
<Send className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
)}
|
||||
<div className="absolute right-2 bottom-2">
|
||||
{loading ? (
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleStop}
|
||||
className="p-1.5 rounded-lg bg-rose-600 hover:bg-rose-700 text-white transition-colors cursor-pointer flex items-center justify-center shadow-xs"
|
||||
title="手动停止执行"
|
||||
>
|
||||
<Square className="w-3.5 h-3.5 fill-white" />
|
||||
</button>
|
||||
) : (
|
||||
<button
|
||||
type="submit"
|
||||
disabled={!input.trim()}
|
||||
className="p-1.5 rounded-lg bg-sky-600 hover:bg-sky-700 text-white disabled:bg-slate-100 disabled:text-slate-350 transition-colors cursor-pointer flex items-center justify-center shadow-xs"
|
||||
>
|
||||
<Send className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</form>
|
||||
</div>
|
||||
|
||||
@ -57,6 +57,7 @@ export interface SessionSummary {
|
||||
session_id: string;
|
||||
title: string;
|
||||
model: string;
|
||||
mode: string;
|
||||
turn_count: number;
|
||||
created_at: string;
|
||||
updated_at: string;
|
||||
|
||||
11
migrations/20260624000000_add_session_mode.sql
Normal file
11
migrations/20260624000000_add_session_mode.sql
Normal file
@ -0,0 +1,11 @@
|
||||
-- Migration: Add Agent Session Mode Support
|
||||
--
|
||||
-- 为代理会话增加运行模式 (mode) 字段，支持多模式 Agent 架构。
|
||||
-- 现有会话自动获得 mode='default'（保持当前行为不变）。
|
||||
--
|
||||
-- mode 取值:
|
||||
-- 'default' — 通用科研助手
|
||||
-- 'deep-research' — 深度研究模式
|
||||
-- 'literature-reader' — 文献阅读助手模式
|
||||
|
||||
ALTER TABLE agent_sessions ADD COLUMN mode TEXT NOT NULL DEFAULT 'default';
|
||||
@ -227,7 +227,7 @@ pub fn rough_estimate_tokens(messages: &[ChatMessage]) -> usize {
|
||||
messages
|
||||
.iter()
|
||||
.map(|m| {
|
||||
let content_len = m.content.as_ref().map_or(0, |c| c.len());
|
||||
let content_len = m.text().map_or(0, |c| c.len());
|
||||
content_len + 4 // 消息结构 overhead 约 4 token
|
||||
})
|
||||
.sum()
|
||||
@ -270,12 +270,13 @@ fn extract_prior_summary(messages: &[ChatMessage]) -> Option<String> {
|
||||
// 从后往前找最近的一份摘要（可能有多次压缩）
|
||||
for msg in messages.iter().rev() {
|
||||
if let Some(ref content) = msg.content {
|
||||
if content.starts_with(SUMMARY_MARKER) || content.starts_with(COMPRESSION_LOG_MARKER) {
|
||||
let text = content.as_str();
|
||||
if text.starts_with(SUMMARY_MARKER) || text.starts_with(COMPRESSION_LOG_MARKER) {
|
||||
// 剥离标记，提取纯摘要内容
|
||||
let summary_body = if let Some(pos) = content.find('\n') {
|
||||
content[pos + 1..].trim().to_string()
|
||||
let summary_body = if let Some(pos) = text.find('\n') {
|
||||
text[pos + 1..].trim().to_string()
|
||||
} else {
|
||||
content.clone()
|
||||
text.to_string()
|
||||
};
|
||||
if !summary_body.is_empty() && summary_body.len() < 2000 {
|
||||
return Some(summary_body);
|
||||
@ -305,7 +306,7 @@ async fn generate_summary(
|
||||
MessageRole::Tool => "工具",
|
||||
_ => return None,
|
||||
};
|
||||
m.content.as_ref().map(|c| {
|
||||
m.text().map(|c| {
|
||||
let preview: String = c.chars().take(200).collect();
|
||||
format!("[{}] {}", role, preview)
|
||||
})
|
||||
@ -564,8 +565,7 @@ pub fn extract_memories_from_compaction(
|
||||
.iter()
|
||||
.filter(|m| m.role != crate::clients::llm::MessageRole::System)
|
||||
.filter_map(|m| {
|
||||
m.content
|
||||
.as_ref()
|
||||
m.text()
|
||||
.map(|c| {
|
||||
let role_label = match m.role {
|
||||
crate::clients::llm::MessageRole::User => "用户",
|
||||
@ -609,7 +609,8 @@ pub fn extract_memories_from_compaction(
|
||||
let runner =
|
||||
crate::agent::subagent::SubAgentRunner::new_with_registry(app.clone(), tool_registry)
|
||||
.with_parent_session(sid.clone())
|
||||
.with_thinking(false);
|
||||
.with_thinking(false)
|
||||
.with_llm_client(app.fast_llm.clone());
|
||||
|
||||
let prompt = format!(
|
||||
"以下内容来自上下文压缩时被丢弃的对话片段。请从中提取值得持久化保存的关键信息。\n\n\
|
||||
@ -677,10 +678,7 @@ mod tests {
|
||||
|
||||
// 工具结果应该被压缩为 [Previous: used search_papers]
|
||||
let tool_msg = &messages[3];
|
||||
assert_eq!(
|
||||
tool_msg.content.as_deref(),
|
||||
Some("[Previous: used search_papers]")
|
||||
);
|
||||
assert_eq!(tool_msg.text(), Some("[Previous: used search_papers]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -695,12 +693,8 @@ mod tests {
|
||||
|
||||
// 应该注入了身份确认块
|
||||
let identity_msg = &messages[2];
|
||||
assert!(identity_msg.content.as_ref().unwrap().contains("身份确认"));
|
||||
assert!(identity_msg
|
||||
.content
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.contains("天体物理学研究助手"));
|
||||
assert!(identity_msg.text().unwrap().contains("身份确认"));
|
||||
assert!(identity_msg.text().unwrap().contains("天体物理学研究助手"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -752,14 +746,11 @@ mod tests {
|
||||
|
||||
// 第一个工具结果应该被压缩
|
||||
let compressed = &messages[3];
|
||||
assert_eq!(
|
||||
compressed.content.as_deref(),
|
||||
Some("[Previous: used search_papers]")
|
||||
);
|
||||
assert_eq!(compressed.text(), Some("[Previous: used search_papers]"));
|
||||
|
||||
// 最近的一个工具结果应该保留
|
||||
let recent = &messages[5];
|
||||
assert_eq!(recent.content.as_deref(), Some("Paper content here..."));
|
||||
assert_eq!(recent.text(), Some("Paper content here..."));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -841,11 +832,7 @@ mod tests {
|
||||
assert!(messages.len() < original_len);
|
||||
assert!(messages.len() <= 51); // HEAD_KEEP + placeholder + tail
|
||||
// 检查占位消息
|
||||
assert!(messages[HEAD_KEEP]
|
||||
.content
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.contains("省略"));
|
||||
assert!(messages[HEAD_KEEP].text().unwrap().contains("省略"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@ -860,7 +847,7 @@ mod tests {
|
||||
// 第一条必须是 system 消息
|
||||
assert_eq!(messages[0].role, MessageRole::System);
|
||||
assert_eq!(
|
||||
messages[0].content.as_deref(),
|
||||
messages[0].text(),
|
||||
Some("You are an astrophysics research assistant.")
|
||||
);
|
||||
}
|
||||
|
||||
@ -1,354 +0,0 @@
|
||||
// src/agent/coordinator/coordinator.rs
|
||||
//
|
||||
// CoordinatorAgent: 运行仅包含元工具的专门化 ReAct 循环。
|
||||
//
|
||||
// Coordinator 的系统提示词明确指导 LLM 采用"委托→检查→合成"的流程。
|
||||
// 循环比主 AgentRuntime 更简单：无压缩、无错误恢复阶梯、无 token 预算——
|
||||
// Coordinator 会话通常很短（3-5 步），这些高级特性暂不需要。
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc::UnboundedSender;
|
||||
use tracing::warn;
|
||||
|
||||
use super::tools::*;
|
||||
use super::worker::WorkerPool;
|
||||
use super::CoordinatorConfig;
|
||||
use crate::agent::hooks::HookRegistry;
|
||||
use crate::agent::runtime::{AgentConfig, AgentStreamEvent};
|
||||
use crate::agent::tools::{ToolContext, ToolRegistry};
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::ChatMessage;
|
||||
|
||||
/// Coordinator 系统提示词
|
||||
fn coordinator_system_prompt() -> String {
|
||||
String::from(
|
||||
"你是一个协调者 Agent。你的职责是分析复杂任务，将其分解为独立的子任务，\
|
||||
委托给 Worker 代理执行，然后合成最终答案。\n\n\
|
||||
## 你的工具\n\
|
||||
- **delegate_task**: 委托子任务给 Worker。\
|
||||
Worker 可以搜索文献、下载论文、解析内容等。每个 Worker 独立运行。\
|
||||
可以同时委托多个任务，它们会并行执行。\n\
|
||||
- **check_task**: 查看已委托任务的状态。如果任务仍在运行，稍后再检查。\n\
|
||||
- **task_stop**: 停止一个运行中的任务（例如中途发现方向错误）。\n\
|
||||
- **synthesize**: 在所有必需任务完成后，收集并合成结果。\n\n\
|
||||
## 工作流程\n\
|
||||
1. 分析用户请求，确定可以并行的独立子任务。\n\
|
||||
2. 使用 delegate_task 将每个子任务委托给 Worker（可一次性全部委托）。\n\
|
||||
3. 使用 check_task 监控进度，必要时等待任务完成。\n\
|
||||
4. 所有任务完成后，使用 synthesize 收集结果。\n\
|
||||
5. 基于合成结果向用户提供最终答案。\n\n\
|
||||
## 重要规则\n\
|
||||
- 最多可同时运行 4 个 Worker。\n\
|
||||
- 先委托所有任务，然后检查进度。不要在每次委托后等待。\n\
|
||||
- 只使用以上列出的工具。你的 Worker 拥有完整工具访问权限。\n\
|
||||
- 使用中文回复。\n\
|
||||
- 如果 synthesize 显示还有未完成的任务，使用 check_task 再检查一次。",
|
||||
)
|
||||
}
|
||||
|
||||
/// CoordinatorAgent 运行协调者 ReAct 循环。
|
||||
pub struct CoordinatorAgent {
|
||||
app_state: Arc<AppState>,
|
||||
config: AgentConfig,
|
||||
coordinator_config: CoordinatorConfig,
|
||||
tool_registry: ToolRegistry,
|
||||
hook_registry: HookRegistry,
|
||||
pool: Arc<WorkerPool>,
|
||||
}
|
||||
|
||||
impl CoordinatorAgent {
|
||||
/// 创建新的 CoordinatorAgent 实例。
|
||||
pub fn new(
|
||||
app_state: Arc<AppState>,
|
||||
agent_config: AgentConfig,
|
||||
coordinator_config: CoordinatorConfig,
|
||||
) -> Self {
|
||||
let pool = Arc::new(WorkerPool::new(
|
||||
app_state.clone(),
|
||||
coordinator_config.clone(),
|
||||
));
|
||||
|
||||
// 构建 coordinator-only 工具注册表
|
||||
let mut tool_registry = ToolRegistry::empty();
|
||||
tool_registry.add_tool(Box::new(DelegateTaskTool::new(
|
||||
pool.clone(),
|
||||
agent_config.enable_thinking,
|
||||
)));
|
||||
tool_registry.add_tool(Box::new(CheckTaskTool::new(pool.clone())));
|
||||
tool_registry.add_tool(Box::new(TaskStopTool::new(pool.clone())));
|
||||
tool_registry.add_tool(Box::new(SynthesizeTool::new(pool.clone())));
|
||||
|
||||
let hook_registry = HookRegistry::with_builtins(
|
||||
app_state.db.clone(),
|
||||
app_state.cancelled_runs.clone(),
|
||||
None,
|
||||
);
|
||||
|
||||
CoordinatorAgent {
|
||||
app_state,
|
||||
config: agent_config,
|
||||
coordinator_config,
|
||||
tool_registry,
|
||||
hook_registry,
|
||||
pool,
|
||||
}
|
||||
}
|
||||
|
||||
/// 运行协调者 ReAct 循环，返回最终文本答案。
|
||||
pub async fn run(
|
||||
&self,
|
||||
session_id: &str,
|
||||
question: &str,
|
||||
turn_index: i32,
|
||||
tx: UnboundedSender<AgentStreamEvent>,
|
||||
) -> anyhow::Result<String> {
|
||||
let llm = &self.app_state.llm;
|
||||
let db = &self.app_state.db;
|
||||
|
||||
// 触发 OnSessionStart
|
||||
self.hook_registry
|
||||
.run_on_session_start(&crate::agent::hooks::SessionStartContext {
|
||||
session_id: session_id.to_string(),
|
||||
turn_index,
|
||||
is_resume: false,
|
||||
})
|
||||
.await;
|
||||
|
||||
let _ = tx.send(AgentStreamEvent::Session {
|
||||
session_id: session_id.to_string(),
|
||||
title: String::new(),
|
||||
});
|
||||
|
||||
// 发送初始 thought 给前端
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: "正在分析任务并规划协调者策略...".into(),
|
||||
step: 0,
|
||||
});
|
||||
|
||||
// 构建消息上下文
|
||||
let mut messages = Vec::new();
|
||||
messages.push(ChatMessage::system(coordinator_system_prompt()));
|
||||
messages.push(ChatMessage::user(question));
|
||||
|
||||
let tool_defs = self.tool_registry.definitions();
|
||||
|
||||
// 保存用户消息到数据库
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, token_count, metadata, agent_name) \
|
||||
VALUES (?, ?, ?, 'user', ?, ?, '{}', 'coordinator')",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(turn_index)
|
||||
.bind(0)
|
||||
.bind(question)
|
||||
.bind(question.len() as i32 / 4)
|
||||
.execute(db)
|
||||
.await;
|
||||
|
||||
let max_steps = self.coordinator_config.worker_max_steps;
|
||||
|
||||
// Coordinator ReAct 循环
|
||||
for step in 1..=max_steps {
|
||||
// 检查取消
|
||||
let is_cancelled = {
|
||||
if let Ok(mut c) = self.app_state.cancelled_runs.lock() {
|
||||
c.remove(session_id)
|
||||
} else {
|
||||
false
|
||||
}
|
||||
};
|
||||
if is_cancelled {
|
||||
warn!("[Coordinator] 用户取消执行, session={}", session_id);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: "用户已手动中止执行。".to_string(),
|
||||
});
|
||||
break;
|
||||
}
|
||||
|
||||
// LLM 流式调用
|
||||
let output = crate::agent::runtime::streaming::process_llm_stream(
|
||||
llm,
|
||||
&messages,
|
||||
&tool_defs,
|
||||
&tx,
|
||||
step,
|
||||
session_id,
|
||||
self.app_state.cancelled_runs.clone(),
|
||||
self.config.enable_thinking,
|
||||
)
|
||||
.await;
|
||||
|
||||
// 发送 usage
|
||||
if let Some(ref usage) = output.usage {
|
||||
let _ = tx.send(AgentStreamEvent::Usage {
|
||||
prompt_tokens: usage.prompt_tokens,
|
||||
completion_tokens: usage.completion_tokens,
|
||||
total_tokens: usage.total_tokens,
|
||||
});
|
||||
}
|
||||
|
||||
let text = match &output.status {
|
||||
crate::agent::runtime::streaming::StreamStatus::Success => output.content.clone(),
|
||||
crate::agent::runtime::streaming::StreamStatus::Cancelled => {
|
||||
break;
|
||||
}
|
||||
crate::agent::runtime::streaming::StreamStatus::Error(e) => {
|
||||
warn!("[Coordinator] LLM 调用失败, step={}: {}", step, e);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("协调者 Agent 调用失败: {}", e),
|
||||
});
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
// 无工具调用 → 协调者直接返回文本答案
|
||||
if output.tool_calls.is_none() || output.tool_calls.as_ref().unwrap().is_empty() {
|
||||
// 保存 assistant 消息
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, token_count, metadata, agent_name) \
|
||||
VALUES (?, ?, ?, 'assistant', ?, ?, '{}', 'coordinator')",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(turn_index)
|
||||
.bind(step as i32)
|
||||
.bind(&text)
|
||||
.bind(text.len() as i32 / 4)
|
||||
.execute(db)
|
||||
.await;
|
||||
|
||||
return Ok(text);
|
||||
}
|
||||
|
||||
// 有工具调用 —— 保存 assistant 消息（含 tool_calls JSON）
|
||||
let tool_calls_json = serde_json::to_string(&output.tool_calls).unwrap_or_default();
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, tool_calls, token_count, metadata, agent_name) \
|
||||
VALUES (?, ?, ?, 'assistant', ?, ?, ?, '{}', 'coordinator')",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(turn_index)
|
||||
.bind(step as i32)
|
||||
.bind(&text)
|
||||
.bind(&tool_calls_json)
|
||||
.bind(text.len() as i32 / 4)
|
||||
.execute(db)
|
||||
.await;
|
||||
|
||||
// 执行工具调用（顺序执行——元工具是轻量级的）
|
||||
let tool_calls = output.tool_calls.unwrap();
|
||||
let mut tool_results = Vec::new();
|
||||
|
||||
for tc in &tool_calls {
|
||||
let tool_name = tc.function.name.clone();
|
||||
let tool_args: serde_json::Value =
|
||||
serde_json::from_str(&tc.function.arguments).unwrap_or_default();
|
||||
|
||||
let tool_call_id = if tc.id.is_empty() {
|
||||
format!("tc_{}", uuid::Uuid::new_v4())
|
||||
} else {
|
||||
tc.id.clone()
|
||||
};
|
||||
|
||||
// 发送 ToolCall SSE
|
||||
let _ = tx.send(AgentStreamEvent::ToolCall {
|
||||
id: tool_call_id.clone(),
|
||||
name: tool_name.clone(),
|
||||
arguments: tool_args.clone(),
|
||||
step,
|
||||
});
|
||||
|
||||
let tool_ctx = ToolContext {
|
||||
app_state: self.app_state.clone(),
|
||||
session_id: session_id.to_string(),
|
||||
silent: false,
|
||||
read_file_state: Arc::new(std::sync::Mutex::new(
|
||||
crate::agent::runtime::file_cache::FileStateCache::new(),
|
||||
)),
|
||||
sse_tx: Some(tx.clone()),
|
||||
enable_thinking: self.config.enable_thinking,
|
||||
additional_allowed_dirs: self.config.additional_allowed_dirs.clone(),
|
||||
};
|
||||
|
||||
let result = match self.tool_registry.get(&tool_name) {
|
||||
Some(tool) => tool.execute(tool_args.clone(), &tool_ctx).await,
|
||||
None => crate::agent::tools::ToolOutput::error(format!(
|
||||
"未知的协调者工具: {}",
|
||||
tool_name
|
||||
)),
|
||||
};
|
||||
|
||||
// 发送 ToolResult SSE
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: tool_call_id.clone(),
|
||||
name: tool_name.clone(),
|
||||
output: result.content.clone(),
|
||||
is_error: result.is_error,
|
||||
metadata: result.metadata.clone(),
|
||||
step,
|
||||
});
|
||||
|
||||
// 保存 tool message
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, tool_call_id, metadata, agent_name) \
|
||||
VALUES (?, ?, ?, 'tool', ?, ?, '{}', 'coordinator')",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(turn_index)
|
||||
.bind(step as i32)
|
||||
.bind(&result.content)
|
||||
.bind(&tool_call_id)
|
||||
.execute(db)
|
||||
.await;
|
||||
|
||||
tool_results.push((tool_name, result));
|
||||
}
|
||||
|
||||
// 将 tool 结果推入消息历史
|
||||
messages.push(ChatMessage::assistant(&text));
|
||||
for (name, result) in &tool_results {
|
||||
messages.push(ChatMessage::user(format!(
|
||||
"[工具结果: {}]\n{}",
|
||||
name,
|
||||
crate::agent::tools::truncate_content(
|
||||
&result.content,
|
||||
self.config.max_tool_output_chars,
|
||||
)
|
||||
)));
|
||||
}
|
||||
|
||||
// OnStepComplete hook
|
||||
self.hook_registry
|
||||
.run_on_step_complete(&crate::agent::hooks::StepCompleteContext {
|
||||
session_id: session_id.to_string(),
|
||||
step,
|
||||
max_steps,
|
||||
messages_count: messages.len(),
|
||||
estimated_tokens: 0,
|
||||
token_limit: self.config.token_hard_limit,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
// 循环耗尽 → 尝试合成剩余结果
|
||||
let completed = self.pool.completed_results().await;
|
||||
if !completed.is_empty() {
|
||||
let mut fallback = format!(
|
||||
"协调者已达到最大步数限制 ({}步)。以下是已完成的任务结果:\n\n",
|
||||
max_steps
|
||||
);
|
||||
for task in &completed {
|
||||
fallback.push_str(&format!(
|
||||
"## {}\n{}\n\n",
|
||||
task.description,
|
||||
task.result.as_deref().unwrap_or("(无结果)")
|
||||
));
|
||||
}
|
||||
Ok(fallback)
|
||||
} else {
|
||||
Ok(format!(
|
||||
"协调者已达到最大步数限制 ({}步)，但未完成任何任务。请检查委托的任务是否合理。",
|
||||
max_steps
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,93 +0,0 @@
|
||||
// src/agent/coordinator/mod.rs
|
||||
//
|
||||
// 协调者模式：层次化多代理编排（P2 特性, 参考 Claude Code coordinatorMode.ts）。
|
||||
//
|
||||
// Coordinator 代理仅有元工具（delegate_task, check_task, task_stop, synthesize），
|
||||
// 将实际工作委托给拥有完整工具访问权限的 Worker 代理。
|
||||
// Worker 通过 SubAgentRunner 实现上下文隔离的独立 ReAct 循环。
|
||||
|
||||
pub mod agent;
|
||||
pub mod tools;
|
||||
pub mod worker;
|
||||
|
||||
/// 协调者任务状态
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum CoordinatorTaskStatus {
|
||||
Pending,
|
||||
Running,
|
||||
Completed,
|
||||
Failed { reason: String },
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
/// Coordinator 维护的任务条目
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CoordinatorTask {
|
||||
pub id: String,
|
||||
pub description: String,
|
||||
pub worker_session_id: Option<String>,
|
||||
pub status: CoordinatorTaskStatus,
|
||||
pub result: Option<String>,
|
||||
pub started_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
pub completed_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
}
|
||||
|
||||
/// Coordinator 配置
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CoordinatorConfig {
|
||||
/// Worker 最大并发数
|
||||
pub max_concurrent_workers: usize,
|
||||
/// Worker 最大推理步数
|
||||
pub worker_max_steps: usize,
|
||||
/// Worker 超时时间（秒）
|
||||
pub worker_timeout_secs: u64,
|
||||
}
|
||||
|
||||
impl Default for CoordinatorConfig {
|
||||
fn default() -> Self {
|
||||
CoordinatorConfig {
|
||||
max_concurrent_workers: 4,
|
||||
worker_max_steps: 5,
|
||||
worker_timeout_secs: 300,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_coordinator_task_lifecycle() {
|
||||
let task = CoordinatorTask {
|
||||
id: "test_1".into(),
|
||||
description: "测试任务".into(),
|
||||
worker_session_id: None,
|
||||
status: CoordinatorTaskStatus::Pending,
|
||||
result: None,
|
||||
started_at: None,
|
||||
completed_at: None,
|
||||
};
|
||||
assert_eq!(task.status, CoordinatorTaskStatus::Pending);
|
||||
assert!(task.result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coordinator_config_defaults() {
|
||||
let config = CoordinatorConfig::default();
|
||||
assert_eq!(config.max_concurrent_workers, 4);
|
||||
assert_eq!(config.worker_max_steps, 5);
|
||||
assert_eq!(config.worker_timeout_secs, 300);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_failed_status_holds_reason() {
|
||||
let status = CoordinatorTaskStatus::Failed {
|
||||
reason: "timeout".into(),
|
||||
};
|
||||
match status {
|
||||
CoordinatorTaskStatus::Failed { reason } => assert_eq!(reason, "timeout"),
|
||||
_ => panic!("expected Failed"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,340 +0,0 @@
|
||||
// src/agent/coordinator/tools.rs
|
||||
//
|
||||
// Coordinator 元工具：delegate_task、check_task、task_stop、synthesize。
|
||||
//
|
||||
// 每个工具实现 AgentTool trait。工具间通过 Arc<WorkerPool> 共享任务板。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tracing::info;
|
||||
|
||||
use super::worker::WorkerPool;
|
||||
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
|
||||
|
||||
// ── delegate_task ──
|
||||
|
||||
pub struct DelegateTaskTool {
|
||||
pool: Arc<WorkerPool>,
|
||||
enable_thinking: bool,
|
||||
}
|
||||
|
||||
impl DelegateTaskTool {
|
||||
pub fn new(pool: Arc<WorkerPool>, enable_thinking: bool) -> Self {
|
||||
DelegateTaskTool {
|
||||
pool,
|
||||
enable_thinking,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for DelegateTaskTool {
|
||||
fn name(&self) -> &str {
|
||||
"delegate_task"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"将子任务委托给 Worker 代理执行。Worker 拥有完整工具访问权限。\
|
||||
使用此工具将独立子任务分发给 Worker，然后使用 check_task 检查进度。\
|
||||
可以先委托所有任务再统一检查，也可以逐个委托立即检查。\
|
||||
适用于：多篇文献综述、并行数据收集、独立计算任务。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"task_description": {
|
||||
"type": "string",
|
||||
"description": "要委托的任务完整描述（Worker 会据此开始工作）"
|
||||
},
|
||||
"context": {
|
||||
"type": "string",
|
||||
"description": "可选：附加上下文信息帮助 Worker 理解任务背景"
|
||||
}
|
||||
},
|
||||
"required": ["task_description"]
|
||||
})
|
||||
}
|
||||
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Cancel
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let task_desc = match args.get("task_description").and_then(|v| v.as_str()) {
|
||||
Some(s) => s.to_string(),
|
||||
None => return ToolOutput::error("缺少 task_description 参数"),
|
||||
};
|
||||
let context = args.get("context").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let full_desc = if context.is_empty() {
|
||||
task_desc
|
||||
} else {
|
||||
format!("{}\n\n附加上下文：{}", task_desc, context)
|
||||
};
|
||||
|
||||
let task_id = self
|
||||
.pool
|
||||
.delegate(&full_desc, &ctx.session_id, self.enable_thinking)
|
||||
.await;
|
||||
|
||||
info!("[Coordinator] 委托任务: id={}", task_id);
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"任务已委托给 Worker。task_id: {}\n使用 check_task(task_id=\"{}\") 查看进度。",
|
||||
task_id, task_id
|
||||
),
|
||||
json!({"task_id": task_id}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// ── check_task ──
|
||||
|
||||
pub struct CheckTaskTool {
|
||||
pool: Arc<WorkerPool>,
|
||||
}
|
||||
|
||||
impl CheckTaskTool {
|
||||
pub fn new(pool: Arc<WorkerPool>) -> Self {
|
||||
CheckTaskTool { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for CheckTaskTool {
|
||||
fn name(&self) -> &str {
|
||||
"check_task"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"查询已委托任务的当前状态和结果。返回 Pending/Running/Completed/Failed/TimedOut。\
|
||||
对于已完成的任务，结果包含在返回中。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"task_id": {
|
||||
"type": "string",
|
||||
"description": "要查询的任务 ID（delegate_task 返回的）"
|
||||
}
|
||||
},
|
||||
"required": ["task_id"]
|
||||
})
|
||||
}
|
||||
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Cancel
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let task_id = match args.get("task_id").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少 task_id 参数"),
|
||||
};
|
||||
|
||||
match self.pool.check(task_id).await {
|
||||
Some(task) => {
|
||||
let status_str = match &task.status {
|
||||
super::CoordinatorTaskStatus::Pending => "Pending",
|
||||
super::CoordinatorTaskStatus::Running => "Running",
|
||||
super::CoordinatorTaskStatus::Completed => "Completed",
|
||||
super::CoordinatorTaskStatus::Failed { .. } => "Failed",
|
||||
super::CoordinatorTaskStatus::TimedOut => "TimedOut",
|
||||
};
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"任务 {}: {}\n描述: {}\n结果: {}",
|
||||
task.id,
|
||||
status_str,
|
||||
task.description,
|
||||
task.result.as_deref().unwrap_or("(尚无结果)")
|
||||
),
|
||||
json!({
|
||||
"task": {
|
||||
"id": task.id,
|
||||
"status": status_str,
|
||||
"description": task.description,
|
||||
"result": task.result,
|
||||
}
|
||||
}),
|
||||
)
|
||||
}
|
||||
None => ToolOutput::error(format!("未找到 task_id: {}", task_id)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── task_stop ──
|
||||
|
||||
pub struct TaskStopTool {
|
||||
pool: Arc<WorkerPool>,
|
||||
}
|
||||
|
||||
impl TaskStopTool {
|
||||
pub fn new(pool: Arc<WorkerPool>) -> Self {
|
||||
TaskStopTool { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for TaskStopTool {
|
||||
fn name(&self) -> &str {
|
||||
"task_stop"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"停止一个正在运行的任务（标记为失败）。如果任务已完成，此操作无效果。\
|
||||
适用于发现 Worker 方向错误时提前终止。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"task_id": {
|
||||
"type": "string",
|
||||
"description": "要停止的任务 ID"
|
||||
}
|
||||
},
|
||||
"required": ["task_id"]
|
||||
})
|
||||
}
|
||||
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Cancel
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let task_id = match args.get("task_id").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少 task_id 参数"),
|
||||
};
|
||||
|
||||
match self.pool.check(task_id).await {
|
||||
Some(_task) => {
|
||||
// 软取消 —— 标记为 Failed。
|
||||
// 完整的取消需要 CancellationToken 传入 Worker，这留待 P3 增强。
|
||||
let mut tasks = self.pool.tasks.write().await;
|
||||
if let Some(t) = tasks.iter_mut().find(|t| t.id == task_id) {
|
||||
if matches!(
|
||||
t.status,
|
||||
super::CoordinatorTaskStatus::Running
|
||||
| super::CoordinatorTaskStatus::Pending
|
||||
) {
|
||||
t.status = super::CoordinatorTaskStatus::Failed {
|
||||
reason: "协调者手动取消".into(),
|
||||
};
|
||||
t.completed_at = Some(chrono::Utc::now());
|
||||
}
|
||||
}
|
||||
ToolOutput::success(
|
||||
format!("任务 {} 已标记为取消。", task_id),
|
||||
json!({"status": "cancelled", "task_id": task_id}),
|
||||
)
|
||||
}
|
||||
None => ToolOutput::error(format!("未找到 task_id: {}", task_id)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── synthesize ──
|
||||
|
||||
pub struct SynthesizeTool {
|
||||
pool: Arc<WorkerPool>,
|
||||
}
|
||||
|
||||
impl SynthesizeTool {
|
||||
pub fn new(pool: Arc<WorkerPool>) -> Self {
|
||||
SynthesizeTool { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for SynthesizeTool {
|
||||
fn name(&self) -> &str {
|
||||
"synthesize"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"收集所有已完成 Worker 任务的结果并合成最终答案。\
|
||||
在所有委托的任务完成后调用此工具，将各 Worker 结果合并为连贯的最终输出。\
|
||||
如果还有未完成的任务，工具会明确提示。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
})
|
||||
}
|
||||
|
||||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||||
InterruptBehavior::Cancel
|
||||
}
|
||||
|
||||
async fn execute(&self, _args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let all = self.pool.all_tasks().await;
|
||||
let completed: Vec<_> = self.pool.completed_results().await;
|
||||
let pending: Vec<_> = all
|
||||
.iter()
|
||||
.filter(|t| {
|
||||
matches!(
|
||||
t.status,
|
||||
super::CoordinatorTaskStatus::Pending | super::CoordinatorTaskStatus::Running
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut output = String::new();
|
||||
|
||||
if !pending.is_empty() {
|
||||
output.push_str(&format!(
|
||||
"⚠ 还有 {} 个任务未完成: {}\n\n",
|
||||
pending.len(),
|
||||
pending
|
||||
.iter()
|
||||
.map(|t| t.id.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
));
|
||||
}
|
||||
|
||||
if completed.is_empty() {
|
||||
output.push_str("尚无已完成的任务结果可供合成。");
|
||||
} else {
|
||||
output.push_str(&format!("已完成 {} 个任务的结果合成:\n\n", completed.len()));
|
||||
for task in &completed {
|
||||
output.push_str(&format!(
|
||||
"## {}\n{}\n\n---\n",
|
||||
task.description.chars().take(120).collect::<String>(),
|
||||
task.result.as_deref().unwrap_or("(无结果)")
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// 收集结构化结果
|
||||
let combined_results: Vec<serde_json::Value> = completed
|
||||
.iter()
|
||||
.map(|t| {
|
||||
json!({
|
||||
"task_id": t.id,
|
||||
"description": t.description,
|
||||
"result": t.result,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
ToolOutput::success(
|
||||
format!(
|
||||
"[协调者合成结果]\n\n{} 个任务已完成。请基于以上结果向用户提供最终答案。",
|
||||
completed.len()
|
||||
),
|
||||
json!({"synthesized": output, "completed_tasks": combined_results}),
|
||||
)
|
||||
}
|
||||
}
|
||||
@ -1,159 +0,0 @@
|
||||
// src/agent/coordinator/worker.rs
|
||||
//
|
||||
// WorkerPool 管理 spawned Worker 子代理的生命周期。
|
||||
//
|
||||
// 每个 Worker 是独立的 SubAgentRunner，拥有完整工具访问权限。
|
||||
// WorkerPool 通过 Semaphore 限制并发数，通过共享任务板跟踪进度。
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{RwLock, Semaphore};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::{CoordinatorConfig, CoordinatorTask, CoordinatorTaskStatus};
|
||||
use crate::agent::subagent::SubAgentRunner;
|
||||
use crate::agent::tools::ToolRegistry;
|
||||
use crate::api::AppState;
|
||||
|
||||
/// Worker 系统提示词 —— 中文，要求 Worker 专注执行并直接返回结果
|
||||
const WORKER_SYSTEM_PROMPT: &str = "\
|
||||
你是一个专注于完成特定子任务的研究助手。你的协调者已经分配了一项任务给你。\
|
||||
请专注于完成该任务，使用所有可用工具，完成后直接返回结果。\
|
||||
不要询问澄清问题——如果你需要做出假设，请记录在结果中。\
|
||||
使用中文回复。";
|
||||
|
||||
/// WorkerPool 管理并发 Worker 子代理。
|
||||
pub struct WorkerPool {
|
||||
app_state: Arc<AppState>,
|
||||
config: CoordinatorConfig,
|
||||
/// 并发信号量 —— 限制同时运行的 worker 数量（Arc 支持在 spawn 中安全共享）
|
||||
concurrency_limiter: Arc<Semaphore>,
|
||||
/// 共享任务板
|
||||
pub(crate) tasks: Arc<RwLock<Vec<CoordinatorTask>>>,
|
||||
}
|
||||
|
||||
impl WorkerPool {
|
||||
pub fn new(app_state: Arc<AppState>, config: CoordinatorConfig) -> Self {
|
||||
let max = config.max_concurrent_workers;
|
||||
WorkerPool {
|
||||
app_state,
|
||||
config,
|
||||
concurrency_limiter: Arc::new(Semaphore::new(max)),
|
||||
tasks: Arc::new(RwLock::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 委托任务给 Worker 子代理。非阻塞 —— 立即返回 task_id。
|
||||
pub async fn delegate(
|
||||
&self,
|
||||
task_description: &str,
|
||||
session_id: &str,
|
||||
enable_thinking: bool,
|
||||
) -> String {
|
||||
let task_id = format!("coord_task_{}", &uuid::Uuid::new_v4().to_string()[..8]);
|
||||
|
||||
let task = CoordinatorTask {
|
||||
id: task_id.clone(),
|
||||
description: task_description.to_string(),
|
||||
worker_session_id: None,
|
||||
status: CoordinatorTaskStatus::Pending,
|
||||
result: None,
|
||||
started_at: None,
|
||||
completed_at: None,
|
||||
};
|
||||
{
|
||||
let mut tasks = self.tasks.write().await;
|
||||
tasks.push(task);
|
||||
}
|
||||
|
||||
let tasks_arc = self.tasks.clone();
|
||||
let limiter = self.concurrency_limiter.clone();
|
||||
let app_state = self.app_state.clone();
|
||||
let worker_max_steps = self.config.worker_max_steps;
|
||||
let worker_timeout = self.config.worker_timeout_secs;
|
||||
let task_id_clone = task_id.clone();
|
||||
let session_id_owned = session_id.to_string();
|
||||
let task_desc_owned = task_description.to_string();
|
||||
|
||||
tokio::spawn(async move {
|
||||
// 获取并发许可
|
||||
let _permit = limiter
|
||||
.acquire()
|
||||
.await
|
||||
.expect("semaphore should not be closed");
|
||||
|
||||
// 更新状态为 Running
|
||||
{
|
||||
let mut tasks = tasks_arc.write().await;
|
||||
if let Some(t) = tasks.iter_mut().find(|t| t.id == task_id_clone) {
|
||||
t.status = CoordinatorTaskStatus::Running;
|
||||
t.started_at = Some(chrono::Utc::now());
|
||||
}
|
||||
}
|
||||
|
||||
// 创建 Worker SubAgentRunner（使用完整 ToolRegistry）
|
||||
let skill_registry = app_state.skill_registry.clone();
|
||||
let tool_registry = ToolRegistry::new(skill_registry);
|
||||
let runner = SubAgentRunner::new_with_registry(app_state, tool_registry)
|
||||
.with_parent_session(session_id_owned.clone())
|
||||
.with_thinking(enable_thinking);
|
||||
|
||||
let result = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(worker_timeout),
|
||||
runner.run(WORKER_SYSTEM_PROMPT, &task_desc_owned, worker_max_steps),
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut tasks = tasks_arc.write().await;
|
||||
if let Some(t) = tasks.iter_mut().find(|t| t.id == task_id_clone) {
|
||||
if let CoordinatorTaskStatus::Failed { reason } = &t.status {
|
||||
if reason == "协调者手动取消" {
|
||||
return;
|
||||
}
|
||||
}
|
||||
match result {
|
||||
Ok(tool_output) if !tool_output.is_error => {
|
||||
t.status = CoordinatorTaskStatus::Completed;
|
||||
t.result = Some(tool_output.content);
|
||||
info!("[Coordinator] Worker 完成: task={}", task_id_clone);
|
||||
}
|
||||
Ok(tool_output) => {
|
||||
t.status = CoordinatorTaskStatus::Failed {
|
||||
reason: tool_output.content.chars().take(200).collect(),
|
||||
};
|
||||
warn!("[Coordinator] Worker 失败: task={}", task_id_clone);
|
||||
}
|
||||
Err(_elapsed) => {
|
||||
t.status = CoordinatorTaskStatus::TimedOut;
|
||||
t.result = Some("Worker 执行超时".to_string());
|
||||
warn!("[Coordinator] Worker 超时: task={}", task_id_clone);
|
||||
}
|
||||
}
|
||||
t.completed_at = Some(chrono::Utc::now());
|
||||
}
|
||||
});
|
||||
|
||||
task_id
|
||||
}
|
||||
|
||||
/// 查询任务状态
|
||||
pub async fn check(&self, task_id: &str) -> Option<CoordinatorTask> {
|
||||
let tasks = self.tasks.read().await;
|
||||
tasks.iter().find(|t| t.id == task_id).cloned()
|
||||
}
|
||||
|
||||
/// 获取所有已完成任务的结果
|
||||
pub async fn completed_results(&self) -> Vec<CoordinatorTask> {
|
||||
let tasks = self.tasks.read().await;
|
||||
tasks
|
||||
.iter()
|
||||
.filter(|t| matches!(t.status, CoordinatorTaskStatus::Completed))
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 获取所有任务的快照
|
||||
pub async fn all_tasks(&self) -> Vec<CoordinatorTask> {
|
||||
let tasks = self.tasks.read().await;
|
||||
tasks.clone()
|
||||
}
|
||||
}
|
||||
@ -174,7 +174,8 @@ pub async fn run_extraction(
|
||||
let tool_registry = build_extraction_tool_registry(memory_manager.clone());
|
||||
|
||||
// 构建子代理
|
||||
let runner = SubAgentRunner::new_with_registry(app_state, tool_registry);
|
||||
let runner = SubAgentRunner::new_with_registry(app_state.clone(), tool_registry)
|
||||
.with_llm_client(app_state.fast_llm.clone());
|
||||
|
||||
// 构建提示词
|
||||
let prompt = build_extraction_prompt(20, &existing_manifest);
|
||||
|
||||
@ -12,9 +12,9 @@
|
||||
pub mod autonomous;
|
||||
pub mod background;
|
||||
pub mod compact;
|
||||
pub mod coordinator;
|
||||
pub mod hooks;
|
||||
pub mod memory;
|
||||
pub mod modes;
|
||||
pub mod runtime;
|
||||
pub mod skills;
|
||||
pub mod subagent;
|
||||
|
||||
50
src/agent/modes/deep_research.rs
Normal file
50
src/agent/modes/deep_research.rs
Normal file
@ -0,0 +1,50 @@
|
||||
// src/agent/modes/deep_research.rs
|
||||
//
|
||||
// 深度研究模式——多来源系统性文献调研，适合撰写综述。
|
||||
|
||||
use super::{AgentMode, ModeConfig, ToolSet};
|
||||
|
||||
pub const DEEP_RESEARCH_MODE: AgentMode = AgentMode {
|
||||
id: "deep-research",
|
||||
name: "深度研究",
|
||||
description: "多来源系统性文献调研与交叉验证，适合撰写综述",
|
||||
icon: "Compass",
|
||||
|
||||
// 身份覆盖：强调系统性调研能力
|
||||
identity_override: Some(
|
||||
"\
|
||||
你是一位专业的天体物理学深度研究员，擅长系统性文献调研与交叉验证。\n\
|
||||
你的目标是穷尽性地覆盖用户的研究问题，不遗漏关键文献，并通过多源交叉验证确保结论的可靠性。",
|
||||
),
|
||||
|
||||
// 原则覆盖：侧重长链推理 + 子代理并行消化 + 任务树管理
|
||||
principles_override: Some(
|
||||
"\
|
||||
# 深度研究原则
|
||||
1. 先广撒网再精读——使用 search_papers 执行 broad search（rows=25），了解领域全貌。
|
||||
2. 按引用数+时效性双层排序筛选核心文献（通常 8-15 篇）。
|
||||
3. 逐篇 download_paper → parse_paper → get_paper_content 深读，提取关键发现和方法论。
|
||||
4. 使用 subagent 并行消化多篇文献，每位子代理负责一篇并返回结构化摘要。
|
||||
5. 用 rag_search 交叉验证本地库中是否有遗漏的关键内容。
|
||||
6. 追踪引文链路——检查核心文献的参考文献和被引情况。
|
||||
7. 使用 todo_write 制定多层级研究任务计划树，每完成一项立即更新状态。
|
||||
8. 对关键结论，检查独立研究的一致/矛盾结果，注意学术派系倾向。
|
||||
9. 遇到工具失败时尝试替代路径，不轻易放弃。
|
||||
10. 输出完整综述：摘要→引言→方法→主要发现→争议与共识→未来方向→参考文献。
|
||||
11. 所有引用使用 ADS bibcode 标注，数学公式使用 LaTeX 格式。
|
||||
12. 用中文回答，保留关键英文术语。",
|
||||
),
|
||||
|
||||
extra_sections: &[],
|
||||
|
||||
// 全量工具（深度研究需要所有工具能力）
|
||||
tool_set: ToolSet::All,
|
||||
|
||||
// 配置：更多步数、启用思考、加载科研权限
|
||||
mode_config: ModeConfig {
|
||||
max_steps: Some(16),
|
||||
enable_thinking: Some(true),
|
||||
tool_timeout_secs: Some(180),
|
||||
permission_profile: Some("research"),
|
||||
},
|
||||
};
|
||||
28
src/agent/modes/default.rs
Normal file
28
src/agent/modes/default.rs
Normal file
@ -0,0 +1,28 @@
|
||||
// src/agent/modes/default.rs
|
||||
//
|
||||
// 默认模式——保持当前 Agent 的全部行为不变。
|
||||
|
||||
use super::{AgentMode, ModeConfig, ToolSet};
|
||||
|
||||
pub const DEFAULT_MODE: AgentMode = AgentMode {
|
||||
id: "default",
|
||||
name: "通用科研助手",
|
||||
description: "通用天体物理学研究助手，自主判断搜索、阅读或计算",
|
||||
icon: "Brain",
|
||||
|
||||
// 使用默认身份和原则
|
||||
identity_override: None,
|
||||
principles_override: None,
|
||||
extra_sections: &[],
|
||||
|
||||
// 全量工具
|
||||
tool_set: ToolSet::All,
|
||||
|
||||
// 不覆盖任何默认配置（全部继承 AgentConfig::default()）
|
||||
mode_config: ModeConfig {
|
||||
max_steps: None,
|
||||
enable_thinking: None,
|
||||
tool_timeout_secs: None,
|
||||
permission_profile: None,
|
||||
},
|
||||
};
|
||||
70
src/agent/modes/literature_reader.rs
Normal file
70
src/agent/modes/literature_reader.rs
Normal file
@ -0,0 +1,70 @@
|
||||
// src/agent/modes/literature_reader.rs
|
||||
//
|
||||
// 文献阅读助手模式——专注论文精读、翻译与理解，强制只读安全沙箱。
|
||||
|
||||
use super::{AgentMode, ModeConfig, ToolSet};
|
||||
|
||||
/// 文献阅读模式允许的工具白名单。
|
||||
///
|
||||
/// 仅保留只读类和文献相关工具，排除所有写入/执行类工具。
|
||||
pub const LITERATURE_READER_TOOLS: &[&str] = &[
|
||||
"read_file",
|
||||
"grep_files",
|
||||
"glob_files",
|
||||
"search_papers",
|
||||
"get_paper_metadata",
|
||||
"get_paper_content",
|
||||
"rag_search",
|
||||
"query_target",
|
||||
"save_note",
|
||||
"todo_write",
|
||||
"load_skill",
|
||||
"ask_user",
|
||||
"compress_context",
|
||||
"search_history",
|
||||
"analyze_image", // 仅视觉模型配置时可用，只读工具
|
||||
];
|
||||
|
||||
pub const LITERATURE_READER_MODE: AgentMode = AgentMode {
|
||||
id: "literature-reader",
|
||||
name: "文献阅读助手",
|
||||
description: "专注论文精读、翻译与笔记提炼，强制只读安全模式",
|
||||
icon: "BookOpen",
|
||||
|
||||
// 身份覆盖：切换为文献解读专家
|
||||
identity_override: Some(
|
||||
"\
|
||||
你是一位严谨的天文文献阅读助手。你的核心任务是帮助用户精读、翻译和提炼学术论文。\n\
|
||||
你专注于解析学术内容、提取核心观点、校验数学公式，并确保翻译中的天文学专有名词绝对精准。\n\
|
||||
你耐心、细致，关注方法论细节和数据可靠性。",
|
||||
),
|
||||
|
||||
// 原则覆盖：结构化阅读、禁止写入、双语输出
|
||||
principles_override: Some(
|
||||
"\
|
||||
# 文献阅读原则
|
||||
1. 使用 get_paper_content 逐段阅读论文内容，按章节结构化分析。
|
||||
2. 采用三段式阅读法：背景与动机 → 方法与数据 → 结论与局限。
|
||||
3. 精准保护 LaTeX 公式占位符和图表引用，翻译时结合天文学 Trie 词典校验专有名词。
|
||||
4. 评估方法论质量：样本量、模型假设、统计方法、系统误差。
|
||||
5. 比较文中结论与领域共识的异同，标注争议点。
|
||||
6. 使用 save_note 保存结构化的阅读笔记（背景/方法/结论/局限/引用）。
|
||||
7. 严禁修改系统文件或执行本地 Shell 命令——你的操作范围仅限于文献阅读。
|
||||
8. 对于数学公式，使用标准 LaTeX 格式保持原貌。
|
||||
9. 用中文解释，同时保留关键英文术语和原始公式。
|
||||
10. 鼓励为重要段落添加高亮标注和批注。",
|
||||
),
|
||||
|
||||
extra_sections: &[],
|
||||
|
||||
// 工具白名单：仅只读 + 文献 + 笔记
|
||||
tool_set: ToolSet::Allowlist(LITERATURE_READER_TOOLS),
|
||||
|
||||
// 配置：较少步数、强制只读权限
|
||||
mode_config: ModeConfig {
|
||||
max_steps: Some(6),
|
||||
enable_thinking: Some(false), // 文献阅读不需要思考模式，节省 token
|
||||
tool_timeout_secs: Some(120),
|
||||
permission_profile: Some("readonly"),
|
||||
},
|
||||
};
|
||||
310
src/agent/modes/mod.rs
Normal file
310
src/agent/modes/mod.rs
Normal file
@ -0,0 +1,310 @@
|
||||
// src/agent/modes/mod.rs
|
||||
//
|
||||
// Agent 多模式系统核心抽象。
|
||||
//
|
||||
// 提供 AgentMode 数据结构、ToolSet 工具集约束、ModeConfig 配置预设
|
||||
// 和 ModeRegistry 注册表。Mode 只在初始化阶段影响 AgentRuntime——
|
||||
// ReAct 循环本身对模式无感知。
|
||||
//
|
||||
// 模式 vs 技能（Skill）：
|
||||
// Mode — 会话级，定义"我是谁"，影响身份/工具集/配置/权限
|
||||
// Skill — 任务级，定义"怎么做"，LLM 按需调用 load_skill
|
||||
|
||||
pub mod deep_research;
|
||||
pub mod default;
|
||||
pub mod literature_reader;
|
||||
|
||||
use deep_research::DEEP_RESEARCH_MODE;
|
||||
use default::DEFAULT_MODE;
|
||||
use literature_reader::LITERATURE_READER_MODE;
|
||||
|
||||
// ── ToolSet ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// 模式对工具集的约束方式。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ToolSet {
|
||||
/// 全量工具可用。
|
||||
All,
|
||||
/// 仅允许列表中的工具（白名单）。
|
||||
Allowlist(&'static [&'static str]),
|
||||
/// 默认工具集中排除列表中的工具（黑名单）。
|
||||
Except(&'static [&'static str]),
|
||||
}
|
||||
|
||||
// ── ModeConfig ───────────────────────────────────────────────────────────
|
||||
|
||||
/// 模式对 AgentConfig 的预设覆盖。
|
||||
///
|
||||
/// 所有字段为 `None` 时表示不覆盖，使用 AgentConfig 的默认值。
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ModeConfig {
|
||||
/// 最大推理步数（None = 使用 AgentConfig 默认值 8）
|
||||
pub max_steps: Option<usize>,
|
||||
/// 是否强制启用思考模式（None = 用户可覆盖）
|
||||
pub enable_thinking: Option<bool>,
|
||||
/// 工具执行超时秒数（None = 使用 AgentConfig 默认值 120）
|
||||
pub tool_timeout_secs: Option<u64>,
|
||||
/// 权限档案名称，对应 permission_profile::load_profile 的参数
|
||||
/// （None = 不额外加载权限档案）
|
||||
pub permission_profile: Option<&'static str>,
|
||||
}
|
||||
|
||||
// ── AgentMode ────────────────────────────────────────────────────────────
|
||||
|
||||
/// 单个 Agent 运行模式的定义。
|
||||
///
|
||||
/// 每个模式是一个纯数据的常量定义，在编译期完全确定。
|
||||
/// 模式之间相互独立，添加新模式不需要修改任何核心逻辑代码。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AgentMode {
|
||||
/// 模式唯一标识（kebab-case，如 "deep-research"）
|
||||
pub id: &'static str,
|
||||
/// 前端显示名称
|
||||
pub name: &'static str,
|
||||
/// 简短描述（~20 字）
|
||||
pub description: &'static str,
|
||||
/// 前端图标名称（对应 lucide-react 图标组件名）
|
||||
pub icon: &'static str,
|
||||
|
||||
// ── 系统提示词覆盖 ──
|
||||
/// 身份 section 覆盖。None = 使用默认 IDENTITY_SECTION。
|
||||
pub identity_override: Option<&'static str>,
|
||||
/// 核心原则 section 覆盖。None = 使用默认 PRINCIPLES_SECTION。
|
||||
pub principles_override: Option<&'static str>,
|
||||
/// 追加的额外 section（追加在静态 section 之后，动态 section 之前）。
|
||||
/// 每个元素为 (section_name, section_content)。
|
||||
pub extra_sections: &'static [(&'static str, &'static str)],
|
||||
|
||||
// ── 工具与配置 ──
|
||||
/// 工具集约束。
|
||||
pub tool_set: ToolSet,
|
||||
/// AgentConfig 预设覆盖。
|
||||
pub mode_config: ModeConfig,
|
||||
}
|
||||
|
||||
// ── ModeRegistry ─────────────────────────────────────────────────────────
|
||||
|
||||
/// 模式注册表——持有所有已注册模式的引用。
|
||||
///
|
||||
/// 使用方式：
|
||||
/// ```ignore
|
||||
/// let registry = ModeRegistry::builtins();
|
||||
/// let mode = registry.get("deep-research");
|
||||
/// ```
|
||||
#[derive(Debug)]
|
||||
pub struct ModeRegistry {
|
||||
modes: Vec<&'static AgentMode>,
|
||||
}
|
||||
|
||||
impl ModeRegistry {
|
||||
/// 创建包含所有内置模式的注册表。
|
||||
pub fn builtins() -> Self {
|
||||
ModeRegistry {
|
||||
modes: vec![&DEFAULT_MODE, &DEEP_RESEARCH_MODE, &LITERATURE_READER_MODE],
|
||||
}
|
||||
}
|
||||
|
||||
/// 按 ID 查找模式。
|
||||
pub fn get(&self, id: &str) -> Option<&&'static AgentMode> {
|
||||
self.modes.iter().find(|m| m.id == id)
|
||||
}
|
||||
|
||||
/// 获取默认模式 ID。
|
||||
pub fn default_id() -> &'static str {
|
||||
"default"
|
||||
}
|
||||
|
||||
/// 列出所有已注册模式。
|
||||
pub fn list(&self) -> &[&'static AgentMode] {
|
||||
&self.modes
|
||||
}
|
||||
|
||||
/// 模式数量。
|
||||
pub fn len(&self) -> usize {
|
||||
self.modes.len()
|
||||
}
|
||||
|
||||
/// 是否为空。
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.modes.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ModeRegistry {
|
||||
fn default() -> Self {
|
||||
Self::builtins()
|
||||
}
|
||||
}
|
||||
|
||||
// ── 辅助函数 ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// 根据 ToolSet 生成工具名称白名单（用于 ToolRegistry::set_definition_filter）。
|
||||
///
|
||||
/// 返回 None 表示 ToolSet::All（无需过滤）。
|
||||
pub fn tool_set_to_filter(tool_set: &ToolSet, all_tool_names: &[String]) -> Option<Vec<String>> {
|
||||
match tool_set {
|
||||
ToolSet::All => None,
|
||||
ToolSet::Allowlist(allowed) => Some(allowed.iter().map(|s| s.to_string()).collect()),
|
||||
ToolSet::Except(excluded) => {
|
||||
let excluded_set: std::collections::HashSet<&str> = excluded.iter().copied().collect();
|
||||
Some(
|
||||
all_tool_names
|
||||
.iter()
|
||||
.filter(|name| !excluded_set.contains(name.as_str()))
|
||||
.cloned()
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── ModeRegistry tests ──
|
||||
|
||||
#[test]
|
||||
fn test_registry_has_all_builtins() {
|
||||
let registry = ModeRegistry::builtins();
|
||||
assert_eq!(registry.len(), 3);
|
||||
assert!(registry.get("default").is_some());
|
||||
assert!(registry.get("deep-research").is_some());
|
||||
assert!(registry.get("literature-reader").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_get_nonexistent() {
|
||||
let registry = ModeRegistry::builtins();
|
||||
assert!(registry.get("nonexistent").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_list_ordered() {
|
||||
let registry = ModeRegistry::builtins();
|
||||
let list = registry.list();
|
||||
// 默认模式排第一
|
||||
assert_eq!(list[0].id, "default");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_default_id() {
|
||||
assert_eq!(ModeRegistry::default_id(), "default");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_is_empty() {
|
||||
assert!(!ModeRegistry::builtins().is_empty());
|
||||
}
|
||||
|
||||
// ── ToolSet tests ──
|
||||
|
||||
#[test]
|
||||
fn test_tool_set_all_is_none() {
|
||||
let all_names: Vec<String> = vec!["a".into(), "b".into(), "c".into()];
|
||||
let filter = tool_set_to_filter(&ToolSet::All, &all_names);
|
||||
assert!(filter.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_set_allowlist() {
|
||||
let all_names: Vec<String> = vec!["a".into(), "b".into(), "c".into()];
|
||||
let filter = tool_set_to_filter(&ToolSet::Allowlist(&["a", "c"]), &all_names);
|
||||
assert_eq!(filter, Some(vec!["a".to_string(), "c".to_string()]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_set_except() {
|
||||
let all_names: Vec<String> = vec!["a".into(), "b".into(), "c".into()];
|
||||
let filter = tool_set_to_filter(&ToolSet::Except(&["b"]), &all_names);
|
||||
assert_eq!(filter, Some(vec!["a".to_string(), "c".to_string()]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_set_except_empty() {
|
||||
let all_names: Vec<String> = vec!["a".into(), "b".into()];
|
||||
let filter = tool_set_to_filter(&ToolSet::Except(&[]), &all_names);
|
||||
assert_eq!(filter, Some(vec!["a".to_string(), "b".to_string()]));
|
||||
}
|
||||
|
||||
// ── ModeConfig tests ──
|
||||
|
||||
#[test]
|
||||
fn test_mode_config_default_all_none() {
|
||||
let config = ModeConfig::default();
|
||||
assert!(config.max_steps.is_none());
|
||||
assert!(config.enable_thinking.is_none());
|
||||
assert!(config.tool_timeout_secs.is_none());
|
||||
assert!(config.permission_profile.is_none());
|
||||
}
|
||||
|
||||
// ── AgentMode field tests ──
|
||||
|
||||
#[test]
|
||||
fn test_default_mode_has_no_overrides() {
|
||||
assert!(DEFAULT_MODE.identity_override.is_none());
|
||||
assert!(DEFAULT_MODE.principles_override.is_none());
|
||||
assert!(DEFAULT_MODE.extra_sections.is_empty());
|
||||
assert_eq!(DEFAULT_MODE.tool_set, ToolSet::All);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deep_research_mode_config() {
|
||||
let mode = &DEEP_RESEARCH_MODE;
|
||||
assert_eq!(mode.mode_config.max_steps, Some(16));
|
||||
assert_eq!(mode.mode_config.enable_thinking, Some(true));
|
||||
assert_eq!(mode.mode_config.permission_profile, Some("research"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_literature_reader_mode_is_readonly() {
|
||||
let mode = &LITERATURE_READER_MODE;
|
||||
// 文献阅读模式应该有只读权限档案
|
||||
assert_eq!(mode.mode_config.permission_profile, Some("readonly"));
|
||||
// 应该是工具白名单模式
|
||||
assert!(matches!(mode.tool_set, ToolSet::Allowlist(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_literature_reader_allowlist_no_write_tools() {
|
||||
let mode = &LITERATURE_READER_MODE;
|
||||
if let ToolSet::Allowlist(tools) = mode.tool_set {
|
||||
// 确保白名单中不包含写入/执行类工具
|
||||
let tools_set: std::collections::HashSet<_> = tools.iter().copied().collect();
|
||||
assert!(!tools_set.contains("run_bash"));
|
||||
assert!(!tools_set.contains("file_write"));
|
||||
assert!(!tools_set.contains("file_edit"));
|
||||
assert!(!tools_set.contains("download_paper"));
|
||||
assert!(!tools_set.contains("parse_paper"));
|
||||
// 确保包含核心阅读工具
|
||||
assert!(tools_set.contains("get_paper_content"));
|
||||
assert!(tools_set.contains("search_papers"));
|
||||
assert!(tools_set.contains("rag_search"));
|
||||
assert!(tools_set.contains("save_note"));
|
||||
} else {
|
||||
panic!("Expected Allowlist tool set");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deep_research_principles_override() {
|
||||
// 深度研究模式应该有原则覆盖（包含"广撒网"或"子代理"等深度研究关键词）
|
||||
let override_text = DEEP_RESEARCH_MODE.principles_override.unwrap();
|
||||
assert!(
|
||||
override_text.contains("广撒网")
|
||||
|| override_text.contains("子代理")
|
||||
|| override_text.contains("综述")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_literature_reader_identity_override() {
|
||||
// 文献阅读模式应该有身份覆盖
|
||||
let identity = LITERATURE_READER_MODE.identity_override.unwrap();
|
||||
assert!(
|
||||
identity.contains("文献") || identity.contains("阅读") || identity.contains("解读")
|
||||
);
|
||||
}
|
||||
}
|
||||
@ -671,7 +671,9 @@ pub async fn execute_parallel(
|
||||
.with_sse_tx(tx.clone())
|
||||
.with_session_id(session_id.to_string())
|
||||
.with_thinking(enable_thinking)
|
||||
.with_additional_dirs(additional_allowed_dirs.clone());
|
||||
.with_additional_dirs(additional_allowed_dirs.clone())
|
||||
.with_tool_call_id(prep.tool_call_id.clone())
|
||||
.with_max_output_chars(max_output_chars);
|
||||
let cancelled = cancelled.clone();
|
||||
let tool_opt = tool_registry.get(&tool_name);
|
||||
|
||||
@ -884,13 +886,18 @@ async fn process_single_result(
|
||||
});
|
||||
|
||||
// 输出处理：小结果直接传递，大结果持久化到磁盘并返回 stub
|
||||
// 但对于已从磁盘读取内容的工具（如 read_file），跳过持久化以防止级联
|
||||
let tool_results_dir = library_dir.join("tool-results");
|
||||
let (processed_content, _persisted_path) = maybe_persist_tool_result(
|
||||
&output.content,
|
||||
tool_call_id,
|
||||
max_output_chars,
|
||||
&tool_results_dir,
|
||||
);
|
||||
let (processed_content, _persisted_path) = if output.skip_persist {
|
||||
(output.content.clone(), None)
|
||||
} else {
|
||||
maybe_persist_tool_result(
|
||||
&output.content,
|
||||
tool_call_id,
|
||||
max_output_chars,
|
||||
&tool_results_dir,
|
||||
)
|
||||
};
|
||||
|
||||
// PostToolUse hook
|
||||
let post_ctx = PostToolUseContext {
|
||||
@ -977,7 +984,7 @@ fn save_tool_message_sync(
|
||||
) {
|
||||
let db_clone = db.clone();
|
||||
let session_id = session_id.to_string();
|
||||
let content = msg.content.as_deref().unwrap_or("").to_string();
|
||||
let content = msg.text().unwrap_or("").to_string();
|
||||
let tool_call_id = msg.tool_call_id.clone();
|
||||
// 提前序列化，避免闭包内的生命周期问题
|
||||
let metadata_str =
|
||||
|
||||
@ -43,6 +43,7 @@ use super::compact;
|
||||
use super::hooks::{
|
||||
HookRegistry, SessionStartContext, StepCompleteContext, UserPromptSubmitContext,
|
||||
};
|
||||
use super::modes::{self, AgentMode, ModeRegistry};
|
||||
use super::terminal::TurnTerminal;
|
||||
use super::tools::ToolRegistry;
|
||||
use crate::api::AppState;
|
||||
@ -91,6 +92,8 @@ pub struct AgentConfig {
|
||||
pub additional_allowed_dirs: Vec<String>,
|
||||
/// 子代理工具白名单（逗号分隔，空=全部工具可用）
|
||||
pub subagent_allowed_tools: Vec<String>,
|
||||
/// Agent 运行模式 ID（"default" / "deep-research" / "literature-reader"）
|
||||
pub mode: String,
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
@ -142,6 +145,7 @@ impl AgentConfig {
|
||||
.unwrap_or(20),
|
||||
additional_allowed_dirs: parse_comma_list("AGENT_ADDITIONAL_DIRS"),
|
||||
subagent_allowed_tools: parse_comma_list("AGENT_SUBAGENT_ALLOWED_TOOLS"),
|
||||
mode: std::env::var("AGENT_MODE").unwrap_or_else(|_| "default".to_string()),
|
||||
};
|
||||
|
||||
// 加载权限档案（AGENT_PERMISSION_PROFILE），追加到现有规则
|
||||
@ -223,9 +227,15 @@ pub enum AgentStreamEvent {
|
||||
metadata: serde_json::Value,
|
||||
step: usize,
|
||||
},
|
||||
/// 文本增量流式输出（最终回答）
|
||||
/// 文本增量流式输出（最终回答或工具流式输出）
|
||||
#[serde(rename = "text_delta")]
|
||||
TextDelta { content: String },
|
||||
TextDelta {
|
||||
content: String,
|
||||
/// 可选：工具调用 ID。当 set 时，此增量属于对应工具的流式输出，
|
||||
/// 前端应将其渲染到工具结果区域而非主文本区。
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tool_call_id: Option<String>,
|
||||
},
|
||||
/// Token 使用统计
|
||||
#[serde(rename = "usage")]
|
||||
Usage {
|
||||
@ -315,14 +325,27 @@ pub struct AgentRuntime {
|
||||
collapse_log: Arc<std::sync::Mutex<compact::collapse::CollapseLog>>,
|
||||
/// Checkpoint 管理器（跨 turn 共享，文件变更操作前自动快照）
|
||||
checkpoint_manager: Arc<checkpoint::CheckpointManager>,
|
||||
/// 是否启用协调者模式（Coordinator delegates to Workers）
|
||||
coordinator_mode: bool,
|
||||
/// 当前运行模式（从 ModeRegistry 解析的静态引用）
|
||||
mode: &'static AgentMode,
|
||||
/// 模式注册表（持有所有已注册模式）
|
||||
mode_registry: ModeRegistry,
|
||||
}
|
||||
|
||||
impl AgentRuntime {
|
||||
/// 创建新的运行时实例
|
||||
pub fn new(app_state: Arc<AppState>) -> Self {
|
||||
let config = AgentConfig::default();
|
||||
let mut config = AgentConfig::default();
|
||||
let mode_registry = ModeRegistry::builtins();
|
||||
let mode = mode_registry.get(&config.mode).copied().unwrap_or_else(|| {
|
||||
tracing::warn!("[AgentRuntime] 未知模式 '{}'，回退到默认模式", config.mode);
|
||||
mode_registry
|
||||
.get(ModeRegistry::default_id())
|
||||
.copied()
|
||||
.unwrap_or(&modes::default::DEFAULT_MODE)
|
||||
});
|
||||
// 合并模式配置预设到 AgentConfig
|
||||
apply_mode_config(&mut config, mode);
|
||||
|
||||
let queue = Arc::new(BgNotificationQueue::new());
|
||||
let metrics_data = Arc::new(std::sync::Mutex::new(super::hooks::MetricsData::default()));
|
||||
let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config));
|
||||
@ -348,6 +371,14 @@ impl AgentRuntime {
|
||||
*session_checker = (*permission_checker).clone();
|
||||
}
|
||||
|
||||
// 视觉模型可用时注册 analyze_image 工具
|
||||
if app_state.vision_llm.is_some() {
|
||||
tool_registry.add_tool(Box::new(crate::agent::tools::astro::AnalyzeImageTool));
|
||||
}
|
||||
|
||||
// ── 应用模式的工具集过滤 ──
|
||||
apply_mode_tool_filter(&mut tool_registry, mode);
|
||||
|
||||
// 初始化 checkpoint 管理器
|
||||
let checkpoint_enabled = std::env::var("AGENT_CHECKPOINT_ENABLED")
|
||||
.unwrap_or_else(|_| "true".to_string())
|
||||
@ -374,12 +405,24 @@ impl AgentRuntime {
|
||||
prompt_cache: std::sync::Mutex::new(SystemPromptCache::new()),
|
||||
collapse_log: Arc::new(std::sync::Mutex::new(compact::collapse::CollapseLog::new())),
|
||||
checkpoint_manager,
|
||||
coordinator_mode: false,
|
||||
mode,
|
||||
mode_registry,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带自定义配置的运行时实例
|
||||
pub fn with_config(app_state: Arc<AppState>, config: AgentConfig) -> Self {
|
||||
pub fn with_config(app_state: Arc<AppState>, mut config: AgentConfig) -> Self {
|
||||
let mode_registry = ModeRegistry::builtins();
|
||||
let mode = mode_registry.get(&config.mode).copied().unwrap_or_else(|| {
|
||||
tracing::warn!("[AgentRuntime] 未知模式 '{}'，回退到默认模式", config.mode);
|
||||
mode_registry
|
||||
.get(ModeRegistry::default_id())
|
||||
.copied()
|
||||
.unwrap_or(&modes::default::DEFAULT_MODE)
|
||||
});
|
||||
// 合并模式配置预设
|
||||
apply_mode_config(&mut config, mode);
|
||||
|
||||
let queue = Arc::new(BgNotificationQueue::new());
|
||||
let metrics_data = Arc::new(std::sync::Mutex::new(super::hooks::MetricsData::default()));
|
||||
let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config));
|
||||
@ -403,6 +446,14 @@ impl AgentRuntime {
|
||||
*session_checker = (*permission_checker).clone();
|
||||
}
|
||||
|
||||
// 视觉模型可用时注册 analyze_image 工具
|
||||
if app_state.vision_llm.is_some() {
|
||||
tool_registry.add_tool(Box::new(crate::agent::tools::astro::AnalyzeImageTool));
|
||||
}
|
||||
|
||||
// ── 应用模式的工具集过滤 ──
|
||||
apply_mode_tool_filter(&mut tool_registry, mode);
|
||||
|
||||
// 初始化 checkpoint 管理器
|
||||
let checkpoint_enabled = std::env::var("AGENT_CHECKPOINT_ENABLED")
|
||||
.unwrap_or_else(|_| "true".to_string())
|
||||
@ -429,7 +480,8 @@ impl AgentRuntime {
|
||||
prompt_cache: std::sync::Mutex::new(SystemPromptCache::new()),
|
||||
collapse_log: Arc::new(std::sync::Mutex::new(compact::collapse::CollapseLog::new())),
|
||||
checkpoint_manager,
|
||||
coordinator_mode: false,
|
||||
mode,
|
||||
mode_registry,
|
||||
}
|
||||
}
|
||||
|
||||
@ -442,45 +494,41 @@ impl AgentRuntime {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// 设置是否启用 LLM 思考模式
|
||||
/// 设置是否启用 LLM 思考模式（向后兼容，优先使用 mode 设置）
|
||||
pub fn with_thinking(mut self, enable: bool) -> Self {
|
||||
self.config.enable_thinking = enable;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置是否启用协调者模式
|
||||
pub fn with_coordinator_mode(mut self, enabled: bool) -> Self {
|
||||
self.coordinator_mode = enabled;
|
||||
/// 设置运行模式（覆盖 AgentConfig 中的 mode 字段）。
|
||||
///
|
||||
/// 调用此方法会重新解析模式并应用对应的配置预设、工具过滤和权限档案。
|
||||
pub fn with_mode(mut self, mode_id: &str) -> Self {
|
||||
self.config.mode = mode_id.to_string();
|
||||
let mode = self.mode_registry.get(mode_id).copied().unwrap_or_else(|| {
|
||||
tracing::warn!(
|
||||
"[AgentRuntime] with_mode: 未知模式 '{}'，回退到默认模式",
|
||||
mode_id
|
||||
);
|
||||
self.mode_registry
|
||||
.get(ModeRegistry::default_id())
|
||||
.copied()
|
||||
.unwrap_or(&modes::default::DEFAULT_MODE)
|
||||
});
|
||||
apply_mode_config(&mut self.config, mode);
|
||||
apply_mode_tool_filter(&mut self.tool_registry, mode);
|
||||
self.mode = mode;
|
||||
self
|
||||
}
|
||||
|
||||
// ── Coordinator Mode ──
|
||||
/// 返回当前模式的 ID。
|
||||
pub fn mode_id(&self) -> &str {
|
||||
self.mode.id
|
||||
}
|
||||
|
||||
/// 运行协调者模式 turn：创建 CoordinatorAgent 并委托执行。
|
||||
async fn run_coordinator_turn(
|
||||
&self,
|
||||
session_info: &session::SessionInfo,
|
||||
question: &str,
|
||||
tx: mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
) -> anyhow::Result<String> {
|
||||
use super::coordinator::agent::CoordinatorAgent;
|
||||
use super::coordinator::CoordinatorConfig;
|
||||
|
||||
let coordinator_config = CoordinatorConfig::default();
|
||||
let coordinator = CoordinatorAgent::new(
|
||||
self.app_state.clone(),
|
||||
self.config.clone(),
|
||||
coordinator_config,
|
||||
);
|
||||
|
||||
coordinator
|
||||
.run(
|
||||
&session_info.session_id,
|
||||
question,
|
||||
session_info.turn_index,
|
||||
tx,
|
||||
)
|
||||
.await
|
||||
/// 返回模式是否强制固定了 thinking。`Some(true/false)` 表示模式已锁定，用户不可覆盖。
|
||||
pub fn mode_fixed_thinking(&self) -> Option<bool> {
|
||||
self.mode.mode_config.enable_thinking
|
||||
}
|
||||
|
||||
// ── Private Helpers ──
|
||||
@ -508,8 +556,7 @@ impl AgentRuntime {
|
||||
};
|
||||
|
||||
// 压缩前捕获消息快照（用于记忆提取桥接，P3）
|
||||
let pre_compact_snapshot: Vec<crate::clients::llm::ChatMessage> =
|
||||
messages.iter().cloned().collect();
|
||||
let pre_compact_snapshot: Vec<crate::clients::llm::ChatMessage> = messages.to_vec();
|
||||
|
||||
compact::compress_context_with_hooks_and_log(
|
||||
messages,
|
||||
@ -563,17 +610,28 @@ impl AgentRuntime {
|
||||
session_id: Option<String>,
|
||||
question: &str,
|
||||
tx: mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
) -> anyhow::Result<String> {
|
||||
self.run_turn_with_image_context(session_id, question, None, None, tx)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 带图片上下文的对话回合。`image_context` 会在用户消息前注入为 system-reminder。
|
||||
/// `image_path` 会存入用户消息的 metadata 以便前端渲染。
|
||||
pub async fn run_turn_with_image_context(
|
||||
&self,
|
||||
session_id: Option<String>,
|
||||
question: &str,
|
||||
image_context: Option<String>,
|
||||
image_path: Option<String>,
|
||||
tx: mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
) -> anyhow::Result<String> {
|
||||
let db = &self.app_state.db;
|
||||
let llm = &self.app_state.llm;
|
||||
|
||||
// Phase 1: 创建或恢复会话
|
||||
let session_info = session::create_or_resume_session(db, session_id.clone(), llm).await?;
|
||||
|
||||
// 协调者模式分支：委托给 CoordinatorAgent
|
||||
if self.coordinator_mode {
|
||||
return self.run_coordinator_turn(&session_info, question, tx).await;
|
||||
}
|
||||
let session_info =
|
||||
session::create_or_resume_session(db, session_id.clone(), llm, &self.config.mode)
|
||||
.await?;
|
||||
|
||||
// 构建 hook 注册表（注入依赖，复用 AgentRuntime 的 metrics_data）
|
||||
let hook_registry = HookRegistry::with_builtins(
|
||||
@ -615,7 +673,19 @@ impl AgentRuntime {
|
||||
)
|
||||
.await?;
|
||||
|
||||
// 保存用户消息到数据库
|
||||
// 图片上下文：在用户消息前注入 system-reminder
|
||||
if let Some(ref img_ctx) = image_context {
|
||||
let reminder = format!("<system-reminder>\n{}\n</system-reminder>", img_ctx);
|
||||
// 插入到倒数第二条位置（用户消息之前）
|
||||
let user_msg = messages.pop().unwrap(); // 用户消息
|
||||
messages.push(ChatMessage::user(&reminder));
|
||||
messages.push(user_msg);
|
||||
}
|
||||
|
||||
// 保存用户消息到数据库（含图片路径元数据，供前端历史渲染）
|
||||
let user_metadata = image_path
|
||||
.as_ref()
|
||||
.map(|p| serde_json::json!({"image_path": p}));
|
||||
self.save_message(
|
||||
db,
|
||||
&session_info.session_id,
|
||||
@ -623,6 +693,7 @@ impl AgentRuntime {
|
||||
0,
|
||||
&ChatMessage::user(question),
|
||||
None,
|
||||
user_metadata,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@ -934,6 +1005,7 @@ impl AgentRuntime {
|
||||
step as i32,
|
||||
&assistant_msg,
|
||||
stream_output.reasoning.as_deref(),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
messages.push(assistant_msg);
|
||||
@ -1008,6 +1080,7 @@ impl AgentRuntime {
|
||||
step as i32,
|
||||
&assistant_msg,
|
||||
stream_output.reasoning.as_deref(),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
messages.push(assistant_msg);
|
||||
@ -1427,6 +1500,7 @@ impl AgentRuntime {
|
||||
/// 1. 所有静态 section 在前 → 内容不变，服务端自然缓存
|
||||
/// 2. 动态 section（environment/tools/skills/memory）在后
|
||||
/// 3. 使用 SystemPromptCache：首次计算后永久复用，/clear 时失效
|
||||
/// 4. 模式（AgentMode）可覆盖 identity 和 principles section
|
||||
fn system_prompt(&self) -> String {
|
||||
use self::system_prompt::{
|
||||
SystemPrompt, IDENTITY_SECTION, PRINCIPLES_SECTION, SAFETY_SECTION,
|
||||
@ -1436,20 +1510,32 @@ impl AgentRuntime {
|
||||
let mut sp = SystemPrompt::new();
|
||||
|
||||
// ═══════ 静态 section（首次计算后永久缓存）═══════
|
||||
// 注意：identity 和 principles 可能被 mode 覆盖
|
||||
|
||||
let mut cache = self.prompt_cache.lock().unwrap_or_else(|e| {
|
||||
tracing::warn!("[SystemPrompt] 缓存锁异常: {:?}", e);
|
||||
e.into_inner()
|
||||
});
|
||||
|
||||
sp.add_section(
|
||||
"identity",
|
||||
cache.get_or_compute("identity", || IDENTITY_SECTION.to_string()),
|
||||
);
|
||||
sp.add_section(
|
||||
"principles",
|
||||
cache.get_or_compute("principles", || PRINCIPLES_SECTION.to_string()),
|
||||
);
|
||||
// Identity：优先使用模式的覆盖，否则使用默认
|
||||
let identity = match self.mode.identity_override {
|
||||
Some(override_text) => override_text.to_string(),
|
||||
None => cache.get_or_compute("identity", || IDENTITY_SECTION.to_string()),
|
||||
};
|
||||
sp.add_section("identity", identity);
|
||||
|
||||
// Principles：优先使用模式的覆盖，否则使用默认
|
||||
let principles = match self.mode.principles_override {
|
||||
Some(override_text) => override_text.to_string(),
|
||||
None => cache.get_or_compute("principles", || PRINCIPLES_SECTION.to_string()),
|
||||
};
|
||||
sp.add_section("principles", principles);
|
||||
|
||||
// 模式的额外 section（追加在静态 section 之后）
|
||||
for (name, content) in self.mode.extra_sections {
|
||||
sp.add_section(name, content.to_string());
|
||||
}
|
||||
|
||||
sp.add_section(
|
||||
"system_context",
|
||||
cache.get_or_compute("system_context", || SYSTEM_CONTEXT_SECTION.to_string()),
|
||||
@ -1472,7 +1558,10 @@ impl AgentRuntime {
|
||||
// 工具目录：使用 tool_catalog() 列出常驻+延迟工具（P3 defer_loading 集成）
|
||||
let tools_section = cache.get_or_compute("tools", || {
|
||||
let catalog = self.tool_registry.tool_catalog();
|
||||
format!("你可以使用以下工具（[deferred] 标记的工具需要通过 load_skill 发现详情）：\n{}", catalog)
|
||||
format!(
|
||||
"你可以使用以下工具（[deferred] 标记的工具需要通过 load_skill 发现详情）：\n{}",
|
||||
catalog
|
||||
)
|
||||
});
|
||||
sp.add_section("tools", tools_section);
|
||||
|
||||
@ -1595,7 +1684,10 @@ impl AgentRuntime {
|
||||
match event {
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated.push_str(&delta);
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta { content: delta });
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: delta,
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
StreamEvent::Usage(u) => {
|
||||
let _ = tx.send(AgentStreamEvent::Usage {
|
||||
@ -1623,6 +1715,7 @@ impl AgentRuntime {
|
||||
step as i32,
|
||||
&assistant_msg,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@ -1630,6 +1723,7 @@ impl AgentRuntime {
|
||||
}
|
||||
|
||||
/// 保存消息到数据库
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn save_message(
|
||||
&self,
|
||||
db: &SqlitePool,
|
||||
@ -1638,9 +1732,19 @@ impl AgentRuntime {
|
||||
step_index: i32,
|
||||
msg: &ChatMessage,
|
||||
thought: Option<&str>,
|
||||
extra_metadata: Option<serde_json::Value>,
|
||||
) -> anyhow::Result<()> {
|
||||
self.save_message_as(db, session_id, turn_index, step_index, msg, thought, "lead")
|
||||
.await
|
||||
self.save_message_as(
|
||||
db,
|
||||
session_id,
|
||||
turn_index,
|
||||
step_index,
|
||||
msg,
|
||||
thought,
|
||||
"lead",
|
||||
extra_metadata,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 保存消息到数据库（指定 agent 身份）
|
||||
@ -1654,6 +1758,7 @@ impl AgentRuntime {
|
||||
msg: &ChatMessage,
|
||||
thought: Option<&str>,
|
||||
agent_name: &str,
|
||||
extra_metadata: Option<serde_json::Value>,
|
||||
) -> anyhow::Result<()> {
|
||||
let role = match msg.role {
|
||||
MessageRole::System => "system",
|
||||
@ -1662,7 +1767,7 @@ impl AgentRuntime {
|
||||
MessageRole::Tool => "tool",
|
||||
};
|
||||
|
||||
let content = msg.content.as_deref().unwrap_or("");
|
||||
let content = msg.content.clone().unwrap_or_default();
|
||||
let tool_calls_json = msg
|
||||
.tool_calls
|
||||
.as_ref()
|
||||
@ -1671,11 +1776,19 @@ impl AgentRuntime {
|
||||
let token_count = content.len() as i32 / 4;
|
||||
|
||||
// metadata: 存储结构化的消息元信息（thought/tool_calls/tool_call_id 等）
|
||||
let metadata = serde_json::json!({
|
||||
let mut metadata = serde_json::json!({
|
||||
"has_thought": thought.is_some(),
|
||||
"has_tool_calls": tool_calls_json.is_some(),
|
||||
"step_index": step_index,
|
||||
});
|
||||
// 合并额外元数据（如图片路径等）
|
||||
if let Some(ref extra) = extra_metadata {
|
||||
if let (Some(base), Some(extra_obj)) = (metadata.as_object_mut(), extra.as_object()) {
|
||||
for (k, v) in extra_obj {
|
||||
base.insert(k.clone(), v.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
let metadata_str = serde_json::to_string(&metadata).unwrap_or_default();
|
||||
|
||||
// raw_json: 存储完整消息的 JSON 序列化（调试/审计用）
|
||||
@ -1718,7 +1831,7 @@ impl AgentRuntime {
|
||||
) {
|
||||
if let Err(e) = self
|
||||
.save_message_as(
|
||||
db, session_id, turn_index, step_index, msg, thought, agent_name,
|
||||
db, session_id, turn_index, step_index, msg, thought, agent_name, None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@ -1726,3 +1839,53 @@ impl AgentRuntime {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Mode helper functions ────────────────────────────────────────────────
|
||||
|
||||
/// 将模式的 AgentConfig 预设合并到给定的 config 中。
|
||||
///
|
||||
/// 仅覆盖 mode_config 中 Some 的字段，None 保持原值不变。
|
||||
fn apply_mode_config(config: &mut AgentConfig, mode: &AgentMode) {
|
||||
if let Some(max_steps) = mode.mode_config.max_steps {
|
||||
config.max_steps = max_steps;
|
||||
}
|
||||
if let Some(enable_thinking) = mode.mode_config.enable_thinking {
|
||||
config.enable_thinking = enable_thinking;
|
||||
}
|
||||
if let Some(tool_timeout_secs) = mode.mode_config.tool_timeout_secs {
|
||||
config.tool_timeout_secs = tool_timeout_secs;
|
||||
}
|
||||
|
||||
// 加载权限档案（模式绑定的 permission_profile）
|
||||
if let Some(profile_name) = mode.mode_config.permission_profile {
|
||||
if let Some(profile) = permission_profile::load_profile(profile_name) {
|
||||
tracing::info!(
|
||||
"[AgentRuntime] 模式 '{}' 加载权限档案: {} — {}",
|
||||
mode.id,
|
||||
profile.name,
|
||||
profile.description
|
||||
);
|
||||
permission_profile::apply_profile_to_config(
|
||||
&profile,
|
||||
&mut config.permission_deny_rules,
|
||||
&mut config.permission_allow_rules,
|
||||
&mut config.permission_ask_rules,
|
||||
&mut config.permission_mode,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 根据模式的 ToolSet 设置工具注册表的定义过滤器。
|
||||
fn apply_mode_tool_filter(tool_registry: &mut ToolRegistry, mode: &AgentMode) {
|
||||
let all_names = tool_registry.tool_names();
|
||||
if let Some(filter) = modes::tool_set_to_filter(&mode.tool_set, &all_names) {
|
||||
tracing::info!(
|
||||
"[AgentRuntime] 模式 '{}' 工具过滤: {} → {} 个工具",
|
||||
mode.id,
|
||||
all_names.len(),
|
||||
filter.len()
|
||||
);
|
||||
tool_registry.set_definition_filter(filter);
|
||||
}
|
||||
}
|
||||
|
||||
@ -13,6 +13,8 @@ use crate::clients::llm::LlmClient;
|
||||
pub struct SessionInfo {
|
||||
pub session_id: String,
|
||||
pub turn_index: i32,
|
||||
/// 会话的运行模式 ID（"default" / "deep-research" / "literature-reader"）
|
||||
pub mode: String,
|
||||
}
|
||||
|
||||
/// 回退操作结果
|
||||
@ -27,10 +29,14 @@ pub struct RewindResult {
|
||||
}
|
||||
|
||||
/// 创建新会话或恢复已有会话。
|
||||
///
|
||||
/// `mode` 指定会话的运行模式（如 "default"、"deep-research"）。
|
||||
/// 新建会话时写入 mode；恢复会话时从 DB 读取 mode（忽略传入的 mode 参数）。
|
||||
pub async fn create_or_resume_session(
|
||||
db: &SqlitePool,
|
||||
session_id: Option<String>,
|
||||
llm: &LlmClient,
|
||||
mode: &str,
|
||||
) -> anyhow::Result<SessionInfo> {
|
||||
match session_id {
|
||||
Some(id) => {
|
||||
@ -46,6 +52,14 @@ pub async fn create_or_resume_session(
|
||||
return Err(anyhow::anyhow!("会话 {} 不存在或已删除", id));
|
||||
}
|
||||
|
||||
// 从 DB 恢复 mode（保持会话创建时的原始模式）
|
||||
let session_mode: String =
|
||||
sqlx::query_scalar("SELECT mode FROM agent_sessions WHERE session_id = ?")
|
||||
.bind(&id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or_else(|_| "default".to_string());
|
||||
|
||||
// 计算当前轮次号（仅统计 active=1 的消息）
|
||||
let turn_index: i32 = sqlx::query_scalar(
|
||||
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages \
|
||||
@ -59,25 +73,44 @@ pub async fn create_or_resume_session(
|
||||
Ok(SessionInfo {
|
||||
session_id: id,
|
||||
turn_index,
|
||||
mode: session_mode,
|
||||
})
|
||||
}
|
||||
None => {
|
||||
let new_id = uuid::Uuid::new_v4().to_string();
|
||||
sqlx::query("INSERT INTO agent_sessions (session_id, title, model) VALUES (?, ?, ?)")
|
||||
.bind(&new_id)
|
||||
.bind("")
|
||||
.bind(llm.model())
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_sessions (session_id, title, model, mode) VALUES (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&new_id)
|
||||
.bind("")
|
||||
.bind(llm.model())
|
||||
.bind(mode)
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
Ok(SessionInfo {
|
||||
session_id: new_id,
|
||||
turn_index: 0,
|
||||
mode: mode.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 从数据库加载会话的运行模式。
|
||||
///
|
||||
/// 返回 None 表示会话不存在或已删除。
|
||||
pub async fn load_session_mode(db: &SqlitePool, session_id: &str) -> Option<String> {
|
||||
sqlx::query_scalar(
|
||||
"SELECT mode FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
}
|
||||
|
||||
/// 加载会话的历史消息（供 LLM 上下文使用）。
|
||||
pub async fn load_history_for_llm(
|
||||
db: &SqlitePool,
|
||||
@ -397,10 +430,13 @@ pub async fn restore_rewound(db: &SqlitePool, session_id: &str) -> anyhow::Resul
|
||||
///
|
||||
/// 返回 `(deleted_message_text, new_turn_index)`。
|
||||
/// 如果没有找到用户消息，返回错误。
|
||||
pub async fn retry_last_turn(db: &SqlitePool, session_id: &str) -> anyhow::Result<(String, i32)> {
|
||||
// 查找最后一条 user 消息（仅 active=1）
|
||||
let last_user: Option<(i64, String, i32)> = sqlx::query_as(
|
||||
"SELECT id, content, turn_index FROM agent_messages \
|
||||
pub async fn retry_last_turn(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<(String, i32, Option<String>)> {
|
||||
// 查找最后一条 user 消息（仅 active=1），同时读取 metadata 中的 image_path
|
||||
let last_user: Option<(i64, String, i32, Option<String>)> = sqlx::query_as(
|
||||
"SELECT id, content, turn_index, metadata FROM agent_messages \
|
||||
WHERE session_id = ? AND role = 'user' AND active = 1 \
|
||||
ORDER BY id DESC LIMIT 1",
|
||||
)
|
||||
@ -408,11 +444,16 @@ pub async fn retry_last_turn(db: &SqlitePool, session_id: &str) -> anyhow::Resul
|
||||
.fetch_optional(db)
|
||||
.await?;
|
||||
|
||||
let (target_id, message_text, _turn) = match last_user {
|
||||
let (target_id, message_text, _turn, metadata_str) = match last_user {
|
||||
Some(t) => t,
|
||||
None => return Err(anyhow::anyhow!("没有找到可重试的用户消息")),
|
||||
};
|
||||
|
||||
// 从 metadata JSON 中提取 image_path
|
||||
let image_path: Option<String> = metadata_str
|
||||
.and_then(|m| serde_json::from_str::<serde_json::Value>(&m).ok())
|
||||
.and_then(|v| v.get("image_path")?.as_str().map(|s| s.to_string()));
|
||||
|
||||
// 硬删除 >= target_id 的所有消息（含 active=0 的历史回退消息）
|
||||
let deleted: i64 = sqlx::query_scalar(
|
||||
"SELECT COUNT(*) FROM agent_messages \
|
||||
@ -443,11 +484,15 @@ pub async fn retry_last_turn(db: &SqlitePool, session_id: &str) -> anyhow::Resul
|
||||
.unwrap_or(0);
|
||||
|
||||
info!(
|
||||
"[Session] 重试: session={}, deleted={} messages from id={}, new_turn={}",
|
||||
session_id, deleted, target_id, new_turn_index
|
||||
"[Session] 重试: session={}, deleted={} messages from id={}, new_turn={}, has_image={}",
|
||||
session_id,
|
||||
deleted,
|
||||
target_id,
|
||||
new_turn_index,
|
||||
image_path.is_some()
|
||||
);
|
||||
|
||||
Ok((message_text, new_turn_index))
|
||||
Ok((message_text, new_turn_index, image_path))
|
||||
}
|
||||
|
||||
/// 分叉结果
|
||||
@ -688,10 +733,7 @@ mod tests {
|
||||
// Actually, rewind_to_message(2) soft-deletes id >= 2
|
||||
// So only system (id=1) remains
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(
|
||||
msgs[0].content.as_deref(),
|
||||
Some("You are a helpful assistant.")
|
||||
);
|
||||
assert_eq!(msgs[0].text(), Some("You are a helpful assistant."));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@ -794,14 +836,14 @@ mod tests {
|
||||
seed_messages(&db, sid).await;
|
||||
|
||||
// Last user message is "Question 3" (id=6)
|
||||
let (msg, new_turn) = retry_last_turn(&db, sid).await.unwrap();
|
||||
let (msg, new_turn, _) = retry_last_turn(&db, sid).await.unwrap();
|
||||
assert_eq!(msg, "Question 3");
|
||||
assert_eq!(new_turn, 2); // turn_index after removing id=6,7
|
||||
|
||||
// Only messages 1-5 should remain
|
||||
let msgs = load_history_for_llm(&db, sid).await.unwrap();
|
||||
assert_eq!(msgs.len(), 5);
|
||||
assert_eq!(msgs.last().unwrap().content.as_deref(), Some("Answer 2"));
|
||||
assert_eq!(msgs.last().unwrap().text(), Some("Answer 2"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@ -829,16 +871,13 @@ mod tests {
|
||||
// Now: ids 1-3 active=1, ids 4-7 active=0
|
||||
|
||||
// Now retry — should hard DELETE from last active user message (id=2)
|
||||
let (msg, _new_turn) = retry_last_turn(&db, sid).await.unwrap();
|
||||
let (msg, _, _) = retry_last_turn(&db, sid).await.unwrap();
|
||||
assert_eq!(msg, "Question 1"); // last active user message
|
||||
|
||||
// Only system message should remain
|
||||
let msgs = load_history_for_llm(&db, sid).await.unwrap();
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(
|
||||
msgs[0].content.as_deref(),
|
||||
Some("You are a helpful assistant.")
|
||||
);
|
||||
assert_eq!(msgs[0].text(), Some("You are a helpful assistant."));
|
||||
|
||||
// Even inactive messages (4-7) should be gone (hard DELETE)
|
||||
let all_count: i64 =
|
||||
@ -900,14 +939,14 @@ mod tests {
|
||||
let branch_msgs = load_history_for_llm(&db, &bid).await.unwrap();
|
||||
let has_branch_msg = branch_msgs
|
||||
.iter()
|
||||
.any(|m| m.content.as_deref() == Some("Branch question"));
|
||||
.any(|m| m.text() == Some("Branch question"));
|
||||
assert!(has_branch_msg);
|
||||
|
||||
// Original does NOT see branch message
|
||||
let orig_msgs = load_history_for_llm(&db, sid).await.unwrap();
|
||||
let has_branch_msg = orig_msgs
|
||||
.iter()
|
||||
.any(|m| m.content.as_deref() == Some("Branch question"));
|
||||
.any(|m| m.text() == Some("Branch question"));
|
||||
assert!(!has_branch_msg);
|
||||
}
|
||||
|
||||
@ -955,7 +994,7 @@ mod tests {
|
||||
|
||||
// Retry: last active user is id=2. Hard DELETE ids >= 2.
|
||||
// This removes everything: 1 (system), 2 (user), 3 (asst), and 4-7 (inactive)
|
||||
let (msg, _) = retry_last_turn(&db, sid).await.unwrap();
|
||||
let (msg, _, _) = retry_last_turn(&db, sid).await.unwrap();
|
||||
assert_eq!(msg, "Question 1");
|
||||
|
||||
// Only system remains
|
||||
|
||||
@ -35,6 +35,7 @@ pub enum StreamStatus {
|
||||
///
|
||||
/// 使用 tokio::select! 在流式读取和取消信号之间竞速。
|
||||
/// 实时发送 Thought/TextDelta SSE 事件给前端。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn process_llm_stream(
|
||||
llm: &LlmClient,
|
||||
messages: &[ChatMessage],
|
||||
@ -106,6 +107,7 @@ pub async fn process_llm_stream(
|
||||
if !is_tool_call_step {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: delta,
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@ -285,6 +285,7 @@ impl StreamingToolExecutor {
|
||||
),
|
||||
is_error: result.is_error,
|
||||
metadata: result.metadata,
|
||||
skip_persist: result.skip_persist,
|
||||
}
|
||||
} else {
|
||||
result
|
||||
|
||||
@ -399,7 +399,6 @@ impl Curator {
|
||||
/// ```
|
||||
pub struct CuratorRunner {
|
||||
curator: Arc<Curator>,
|
||||
db: sqlx::SqlitePool,
|
||||
/// 当前是否已触发过（防止重复运行）
|
||||
last_run_at: Arc<Mutex<Option<DateTime<Utc>>>>,
|
||||
/// 是否暂停
|
||||
@ -414,10 +413,9 @@ pub struct CuratorRunner {
|
||||
|
||||
impl CuratorRunner {
|
||||
/// 创建运行器。
|
||||
pub fn new(curator: Curator, db: sqlx::SqlitePool) -> Self {
|
||||
pub fn new(curator: Curator, _db: sqlx::SqlitePool) -> Self {
|
||||
CuratorRunner {
|
||||
curator: Arc::new(curator),
|
||||
db,
|
||||
last_run_at: Arc::new(Mutex::new(None)),
|
||||
paused: Arc::new(AtomicBool::new(false)),
|
||||
interval: DEFAULT_CURATOR_INTERVAL,
|
||||
|
||||
@ -36,8 +36,10 @@ pub struct SubAgentRunner {
|
||||
permission_checker: Arc<PermissionChecker>,
|
||||
/// 可选的进度发送器（用于向父代理报告中间步骤）
|
||||
progress_tx: Option<UnboundedSender<AgentStreamEvent>>,
|
||||
/// 父代理的会话 ID（用于子代理消息的数据库持久化）
|
||||
/// 父代理的会话 ID（用于子代理消息 of 数据库持久化）
|
||||
parent_session_id: String,
|
||||
/// 覆盖子代理使用的模型客户端，若不指定则默认使用 app_state.llm (主推理模型)
|
||||
llm_client_override: Option<LlmClient>,
|
||||
}
|
||||
|
||||
impl SubAgentRunner {
|
||||
@ -52,6 +54,7 @@ impl SubAgentRunner {
|
||||
permission_checker: Arc::new(PermissionChecker::new()),
|
||||
progress_tx: None,
|
||||
parent_session_id: String::new(),
|
||||
llm_client_override: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -71,6 +74,7 @@ impl SubAgentRunner {
|
||||
permission_checker,
|
||||
progress_tx,
|
||||
parent_session_id: String::new(),
|
||||
llm_client_override: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -80,6 +84,12 @@ impl SubAgentRunner {
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置用于子代理的 LLM 客户端
|
||||
pub fn with_llm_client(mut self, llm: LlmClient) -> Self {
|
||||
self.llm_client_override = Some(llm);
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置是否启用 LLM 思考模式
|
||||
pub fn with_thinking(mut self, enable: bool) -> Self {
|
||||
self.config.enable_thinking = enable;
|
||||
@ -97,6 +107,7 @@ impl SubAgentRunner {
|
||||
permission_checker: Arc::new(PermissionChecker::new()),
|
||||
progress_tx: None,
|
||||
parent_session_id: String::new(),
|
||||
llm_client_override: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -116,6 +127,7 @@ impl SubAgentRunner {
|
||||
permission_checker,
|
||||
progress_tx,
|
||||
parent_session_id: String::new(),
|
||||
llm_client_override: None,
|
||||
}
|
||||
}
|
||||
|
||||
@ -261,7 +273,10 @@ impl SubAgentRunner {
|
||||
max_steps: usize,
|
||||
subagent_name: &str,
|
||||
) -> ToolOutput {
|
||||
let llm = &self.app_state.llm;
|
||||
let llm = self
|
||||
.llm_client_override
|
||||
.as_ref()
|
||||
.unwrap_or(&self.app_state.llm);
|
||||
let tool_defs = self.tool_registry.definitions();
|
||||
|
||||
// 全新上下文
|
||||
|
||||
340
src/agent/tools/astro/analyze_image.rs
Normal file
340
src/agent/tools/astro/analyze_image.rs
Normal file
@ -0,0 +1,340 @@
|
||||
// src/agent/tools/astro/analyze_image.rs
|
||||
//
|
||||
// analyze_image 工具：使用专用视觉模型流式分析图片。
|
||||
//
|
||||
// 仅在 LLM_VISION_MODEL 环境变量配置时注册。主 Agent 可能为纯文本模型，
|
||||
// 图片分析通过此工具委托给专用 vision 模型，增量结果通过 SSE 实时推送到前端。
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tokio::sync::mpsc::UnboundedSender;
|
||||
use tracing::{error, info};
|
||||
|
||||
use crate::agent::runtime::AgentStreamEvent;
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
use crate::clients::llm::LlmClient;
|
||||
|
||||
pub struct AnalyzeImageTool;
|
||||
|
||||
impl AnalyzeImageTool {
|
||||
/// 核心分析逻辑：加载图片 + 流式调用视觉模型。与 ToolContext 解耦以便测试。
|
||||
pub async fn analyze(
|
||||
image_path: &str,
|
||||
question: &str,
|
||||
vision_llm: &LlmClient,
|
||||
library_dir: &Path,
|
||||
sse_tx: Option<&UnboundedSender<AgentStreamEvent>>,
|
||||
tool_call_id: Option<String>,
|
||||
) -> ToolOutput {
|
||||
let (base64_data, mime_type) = match load_image(library_dir, image_path).await {
|
||||
Ok(data) => data,
|
||||
Err(e) => {
|
||||
error!("[analyze_image] 加载图片失败: {} ({})", e, image_path);
|
||||
return ToolOutput::error(format!("加载图片失败: {}", e));
|
||||
}
|
||||
};
|
||||
|
||||
info!(
|
||||
"[analyze_image] 流式分析图片: {} ({} bytes, type={})",
|
||||
image_path,
|
||||
base64_data.len(),
|
||||
mime_type
|
||||
);
|
||||
|
||||
let system_prompt = "你是一位专业的天体物理学家。请基于提供的图片和用户问题，给出专业、详细、清晰的中文分析。如有数学公式，使用 LaTeX 格式。";
|
||||
|
||||
// 流式输出：先发送状态头部（作为第一条 text_delta，保证持久化到工具结果中）
|
||||
let header = format!(
|
||||
"> 正在使用视觉模型（{}）分析图片: {}\n\n",
|
||||
vision_llm.model(),
|
||||
image_path
|
||||
);
|
||||
if let Some(tx) = &sse_tx {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: header.clone(),
|
||||
tool_call_id: tool_call_id.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
let sse_tx_clone = sse_tx.cloned();
|
||||
let tc_id = tool_call_id.clone();
|
||||
let result = vision_llm
|
||||
.analyze_image_stream(
|
||||
system_prompt,
|
||||
question,
|
||||
&base64_data,
|
||||
&mime_type,
|
||||
move |delta: &str| {
|
||||
if let Some(ref tx) = sse_tx_clone {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: delta.to_string(),
|
||||
tool_call_id: tc_id.clone(),
|
||||
});
|
||||
}
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(answer) => ToolOutput::success(
|
||||
format!("{}{}", header, answer),
|
||||
json!({
|
||||
"image_path": image_path,
|
||||
"question": question,
|
||||
"model": vision_llm.model(),
|
||||
}),
|
||||
),
|
||||
Err(e) => {
|
||||
error!("[analyze_image] 视觉模型调用失败: {}", e);
|
||||
ToolOutput::error(format!("图片分析失败: {}", e))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for AnalyzeImageTool {
|
||||
fn name(&self) -> &str {
|
||||
"analyze_image"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"使用专用视觉模型分析图片。可传入本地图片路径（相对于 library 目录）或 http/https URL。分析结果会实时流式输出。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"image_path": {
|
||||
"type": "string",
|
||||
"description": "图片路径。本地路径相对于 library 目录，也可以传入 http/https URL。"
|
||||
},
|
||||
"question": {
|
||||
"type": "string",
|
||||
"description": "要分析的问题，如「这张图的 x 轴和 y 轴分别代表什么？」"
|
||||
}
|
||||
},
|
||||
"required": ["image_path", "question"]
|
||||
})
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
self.execute_with_progress(args, ctx, None).await
|
||||
}
|
||||
|
||||
async fn execute_with_progress(
|
||||
&self,
|
||||
args: serde_json::Value,
|
||||
ctx: &ToolContext,
|
||||
_progress_tx: Option<&UnboundedSender<String>>,
|
||||
) -> ToolOutput {
|
||||
let image_path = match args.get("image_path").and_then(|v| v.as_str()) {
|
||||
Some(p) => p.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'image_path'"),
|
||||
};
|
||||
let question = match args.get("question").and_then(|v| v.as_str()) {
|
||||
Some(q) => q.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'question'"),
|
||||
};
|
||||
let vision_llm = match ctx.app_state.vision_llm.as_ref() {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
return ToolOutput::error(
|
||||
"图片分析功能未启用。请配置 LLM_VISION_MODEL 环境变量后重启服务。",
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
Self::analyze(
|
||||
&image_path,
|
||||
&question,
|
||||
vision_llm,
|
||||
&ctx.app_state.config.library_dir,
|
||||
ctx.sse_tx.as_ref(),
|
||||
ctx.tool_call_id.clone(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
/// 加载图片：本地路径或 HTTP(S) URL → base64 data + mime type
|
||||
async fn load_image(library_dir: &Path, path: &str) -> anyhow::Result<(String, String)> {
|
||||
if path.starts_with("http://") || path.starts_with("https://") {
|
||||
let client = reqwest::Client::new();
|
||||
let resp = client.get(path).send().await?;
|
||||
let mime = resp
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("image/png")
|
||||
.to_string();
|
||||
let bytes = resp.bytes().await?;
|
||||
let b64 = base64_encode(&bytes);
|
||||
Ok((b64, mime))
|
||||
} else {
|
||||
let abs_path = if path.starts_with('/') {
|
||||
Path::new(path).to_path_buf()
|
||||
} else {
|
||||
library_dir.join(path)
|
||||
};
|
||||
if !abs_path.exists() {
|
||||
return Err(anyhow::anyhow!("图片文件不存在: {}", abs_path.display()));
|
||||
}
|
||||
let ext = abs_path
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.unwrap_or("png")
|
||||
.to_lowercase();
|
||||
let mime = match ext.as_str() {
|
||||
"jpg" | "jpeg" => "image/jpeg",
|
||||
"gif" => "image/gif",
|
||||
"webp" => "image/webp",
|
||||
"svg" => "image/svg+xml",
|
||||
_ => "image/png",
|
||||
};
|
||||
let bytes = std::fs::read(&abs_path)?;
|
||||
let b64 = base64_encode(&bytes);
|
||||
Ok((b64, mime.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
fn base64_encode(bytes: &[u8]) -> String {
|
||||
use base64::{engine::general_purpose, Engine as _};
|
||||
general_purpose::STANDARD.encode(bytes)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// 1×1 红色像素 PNG 的 base64 编码。
|
||||
const RED_PIXEL_PNG_BASE64: &str =
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==";
|
||||
|
||||
fn base64_decode(b64: &str) -> Vec<u8> {
|
||||
use base64::{engine::general_purpose, Engine as _};
|
||||
general_purpose::STANDARD.decode(b64).unwrap()
|
||||
}
|
||||
|
||||
/// 真实视觉模型集成测试 —— 走完整的 `AnalyzeImageTool::analyze()` 路径。
|
||||
/// 需要配置 LLM_VISION_MODEL 环境变量。
|
||||
/// 运行: `LLM_VISION_MODEL=gpt-4o cargo test analyze_image --lib -- --ignored`
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn test_execute_with_progress_real() {
|
||||
let config = crate::Config::from_env();
|
||||
|
||||
if config.llm_vision_model.is_empty() {
|
||||
eprintln!("跳过: 未配置 LLM_VISION_MODEL 环境变量");
|
||||
return;
|
||||
}
|
||||
|
||||
let key = if config.llm_vision_api_key.is_empty() {
|
||||
config.llm_api_key.clone()
|
||||
} else {
|
||||
config.llm_vision_api_key.clone()
|
||||
};
|
||||
let base = if config.llm_vision_api_base.is_empty() {
|
||||
config.llm_api_base.clone()
|
||||
} else {
|
||||
config.llm_vision_api_base.clone()
|
||||
};
|
||||
let vision_llm = LlmClient::new(key, base, config.llm_vision_model.clone());
|
||||
|
||||
// 写入测试图片到临时目录
|
||||
let tmp_dir = std::env::temp_dir().join("astro_vision_test");
|
||||
std::fs::create_dir_all(&tmp_dir).unwrap();
|
||||
let img_path = tmp_dir.join("test_pixel.png");
|
||||
std::fs::write(&img_path, base64_decode(RED_PIXEL_PNG_BASE64)).unwrap();
|
||||
|
||||
// 创建 SSE 通道，收集流式事件
|
||||
let (sse_tx, mut sse_rx) = tokio::sync::mpsc::unbounded_channel::<AgentStreamEvent>();
|
||||
|
||||
// 后台收集 SSE 事件
|
||||
let collector = tokio::spawn(async move {
|
||||
let mut events: Vec<AgentStreamEvent> = Vec::new();
|
||||
while let Some(event) = sse_rx.recv().await {
|
||||
let is_done = matches!(event, AgentStreamEvent::Done);
|
||||
events.push(event);
|
||||
if is_done {
|
||||
break;
|
||||
}
|
||||
}
|
||||
events
|
||||
});
|
||||
|
||||
// 调用核心分析逻辑
|
||||
let output = AnalyzeImageTool::analyze(
|
||||
"test_pixel.png",
|
||||
"这张图片有什么内容？",
|
||||
&vision_llm,
|
||||
&tmp_dir,
|
||||
Some(&sse_tx),
|
||||
None, // tool_call_id — not needed for tests
|
||||
)
|
||||
.await;
|
||||
|
||||
// 关闭 SSE 发送端，触发 collector 完成
|
||||
let _ = sse_tx.send(AgentStreamEvent::Done);
|
||||
drop(sse_tx);
|
||||
let events = collector.await.unwrap_or_default();
|
||||
|
||||
assert!(!output.is_error, "分析应成功: {}", output.content);
|
||||
assert!(!output.content.is_empty(), "应返回非空结果");
|
||||
|
||||
// 验证流式事件
|
||||
let delta_count = events
|
||||
.iter()
|
||||
.filter(|e| matches!(e, AgentStreamEvent::TextDelta { .. }))
|
||||
.count();
|
||||
|
||||
println!("视觉模型: {}", vision_llm.model());
|
||||
println!("SSE 事件总数: {}, TextDelta: {}", events.len(), delta_count);
|
||||
println!("分析结果:\n{}", output.content);
|
||||
|
||||
assert!(
|
||||
delta_count >= 2,
|
||||
"至少应有 header + 分析结果的 TextDelta 事件"
|
||||
);
|
||||
|
||||
// 清理
|
||||
std::fs::remove_dir_all(&tmp_dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_image_from_disk() {
|
||||
let tmp_dir = std::env::temp_dir().join("astro_test");
|
||||
std::fs::create_dir_all(&tmp_dir).unwrap();
|
||||
let img_path = tmp_dir.join("test.png");
|
||||
let png_bytes = base64_decode(RED_PIXEL_PNG_BASE64);
|
||||
std::fs::write(&img_path, &png_bytes).unwrap();
|
||||
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
let (b64, mime) = rt
|
||||
.block_on(load_image(&tmp_dir, "test.png"))
|
||||
.expect("应成功加载本地图片");
|
||||
|
||||
assert_eq!(mime, "image/png");
|
||||
assert_eq!(b64, RED_PIXEL_PNG_BASE64);
|
||||
|
||||
std::fs::remove_dir_all(&tmp_dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_image_not_found() {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
let result = rt.block_on(load_image(Path::new("/tmp"), "nonexistent_image_xyz.png"));
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
@ -2,12 +2,14 @@
|
||||
//
|
||||
// 天文科研工具集 — 文献搜索/下载/解析、RAG 检索、天体目标查询、研究笔记。
|
||||
|
||||
pub mod analyze_image;
|
||||
pub mod note;
|
||||
pub mod paper;
|
||||
pub mod rag;
|
||||
pub mod search;
|
||||
pub mod target;
|
||||
|
||||
pub use analyze_image::AnalyzeImageTool;
|
||||
pub use note::SaveNoteTool;
|
||||
pub use paper::{DownloadPaperTool, GetPaperContentTool, ParsePaperTool};
|
||||
pub use rag::RagSearchTool;
|
||||
|
||||
@ -130,9 +130,9 @@ impl AgentTool for ReadFileTool {
|
||||
content.clone()
|
||||
};
|
||||
|
||||
// 截断到单次输出上限
|
||||
let truncated_content = crate::agent::tools::truncate_content(&truncated, 4000);
|
||||
let content_len = truncated_content.len();
|
||||
// read_file 不截断输出——LLM 明确要求读取完整内容，
|
||||
// 且 skip_persist 确保不会被二次持久化。
|
||||
let content_len = truncated.len();
|
||||
info!(
|
||||
"[ReadFile] 读取 {}: {} 字符 ({} 行)",
|
||||
file_path_str, content_len, total_lines
|
||||
@ -154,8 +154,8 @@ impl AgentTool for ReadFileTool {
|
||||
}
|
||||
// ── 缓存更新结束 ──
|
||||
|
||||
ToolOutput::success(
|
||||
truncated_content,
|
||||
ToolOutput::success_skip_persist(
|
||||
truncated,
|
||||
json!({"file_path": file_path_str, "total_lines": total_lines}),
|
||||
)
|
||||
}
|
||||
|
||||
@ -67,6 +67,10 @@ pub struct ToolContext {
|
||||
pub enable_thinking: bool,
|
||||
/// 附加允许目录（扩展文件沙箱范围）
|
||||
pub additional_allowed_dirs: Vec<String>,
|
||||
/// 当前工具调用 ID（工具可通过此 ID 将流式输出关联到自身）
|
||||
pub tool_call_id: Option<String>,
|
||||
/// 工具输出最大字符数（继承自 AgentConfig，工具可据此截断输出）
|
||||
pub max_output_chars: usize,
|
||||
}
|
||||
|
||||
impl ToolContext {
|
||||
@ -80,6 +84,8 @@ impl ToolContext {
|
||||
sse_tx: None,
|
||||
enable_thinking: false,
|
||||
additional_allowed_dirs: Vec::new(),
|
||||
tool_call_id: None,
|
||||
max_output_chars: 4000,
|
||||
}
|
||||
}
|
||||
|
||||
@ -96,6 +102,8 @@ impl ToolContext {
|
||||
sse_tx: None,
|
||||
enable_thinking: false,
|
||||
additional_allowed_dirs: Vec::new(),
|
||||
tool_call_id: None,
|
||||
max_output_chars: 4000,
|
||||
}
|
||||
}
|
||||
|
||||
@ -126,6 +134,18 @@ impl ToolContext {
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置工具调用 ID（用于工具的流式输出关联到自身）
|
||||
pub fn with_tool_call_id(mut self, id: String) -> Self {
|
||||
self.tool_call_id = Some(id);
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置最大输出字符数
|
||||
pub fn with_max_output_chars(mut self, max_chars: usize) -> Self {
|
||||
self.max_output_chars = max_chars;
|
||||
self
|
||||
}
|
||||
|
||||
/// 创建静默上下文（子代理使用）
|
||||
pub fn silent(app_state: Arc<AppState>) -> Self {
|
||||
ToolContext {
|
||||
@ -136,6 +156,8 @@ impl ToolContext {
|
||||
sse_tx: None,
|
||||
enable_thinking: false,
|
||||
additional_allowed_dirs: Vec::new(),
|
||||
tool_call_id: None,
|
||||
max_output_chars: 4000,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -149,6 +171,8 @@ pub struct ToolOutput {
|
||||
pub is_error: bool,
|
||||
/// 结构化元数据（给前端 Timeline 直接渲染）
|
||||
pub metadata: serde_json::Value,
|
||||
/// 跳过持久化到磁盘（用于 read_file 等已从磁盘读取内容的工具，避免级联持久化）
|
||||
pub skip_persist: bool,
|
||||
}
|
||||
|
||||
impl ToolOutput {
|
||||
@ -158,6 +182,17 @@ impl ToolOutput {
|
||||
content: content.into(),
|
||||
is_error: false,
|
||||
metadata,
|
||||
skip_persist: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建成功结果，跳过持久化
|
||||
pub fn success_skip_persist(content: impl Into<String>, metadata: serde_json::Value) -> Self {
|
||||
ToolOutput {
|
||||
content: content.into(),
|
||||
is_error: false,
|
||||
metadata,
|
||||
skip_persist: true,
|
||||
}
|
||||
}
|
||||
|
||||
@ -167,6 +202,7 @@ impl ToolOutput {
|
||||
content: msg.into(),
|
||||
is_error: true,
|
||||
metadata: json!({}),
|
||||
skip_persist: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -91,13 +91,14 @@ pub fn maybe_persist_tool_result(
|
||||
}
|
||||
|
||||
// 预览（在第一个换行处截断，避免 mid-line cut）
|
||||
// 使用 chars 迭代器保证始终落在有效 UTF-8 字符边界上
|
||||
let preview_limit = 500.min(max_chars);
|
||||
let preview: String =
|
||||
if let Some(newline_pos) = content[..preview_limit.min(content.len())].rfind('\n') {
|
||||
content[..newline_pos].to_string()
|
||||
} else {
|
||||
content.chars().take(preview_limit).collect()
|
||||
};
|
||||
let preview_chars: String = content.chars().take(preview_limit).collect();
|
||||
let preview: String = if let Some(newline_pos) = preview_chars.rfind('\n') {
|
||||
preview_chars[..newline_pos].to_string()
|
||||
} else {
|
||||
preview_chars
|
||||
};
|
||||
|
||||
let stub = format!(
|
||||
"<persisted-output>\n\
|
||||
|
||||
190
src/api/agent.rs
190
src/api/agent.rs
@ -26,12 +26,59 @@ use crate::agent::runtime::{AgentRuntime, AgentStreamEvent};
|
||||
pub struct AgentChatRequest {
|
||||
pub question: String,
|
||||
pub session_id: Option<String>,
|
||||
/// 是否启用 LLM 思考模式（默认关闭）
|
||||
/// Agent 运行模式： "default" / "deep-research" / "literature-reader"
|
||||
#[serde(default = "default_mode")]
|
||||
pub mode: String,
|
||||
/// 是否启用 LLM 思考模式。None = 由 mode 决定，Some(true/false) = 用户显式覆盖。
|
||||
#[serde(default)]
|
||||
pub thinking: bool,
|
||||
/// 是否启用协调者模式（Coordinator delegates to Workers）
|
||||
pub thinking: Option<bool>,
|
||||
/// 可选的图片附件（base64 编码 + MIME 类型）
|
||||
#[serde(default)]
|
||||
pub coordinator_mode: bool,
|
||||
pub image: Option<AttachedImage>,
|
||||
}
|
||||
|
||||
/// 用户附带的图片，用于多模态 Agent 提问。
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct AttachedImage {
|
||||
/// base64 编码的图片数据（不含 data:xxx;base64, 前缀）。与 path 互斥。
|
||||
#[serde(default)]
|
||||
pub data: String,
|
||||
/// MIME 类型，如 "image/png"、"image/jpeg"
|
||||
#[serde(default)]
|
||||
pub mime_type: String,
|
||||
/// 已有图片的相对路径（重试时复用已有文件，不再重新 base64 解码存盘）
|
||||
#[serde(default)]
|
||||
pub path: Option<String>,
|
||||
}
|
||||
|
||||
fn default_mode() -> String {
|
||||
"default".to_string()
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct AgentModeDto {
|
||||
pub id: &'static str,
|
||||
pub name: &'static str,
|
||||
pub description: &'static str,
|
||||
pub icon: &'static str,
|
||||
}
|
||||
|
||||
// ── GET /api/chat/modes ──
|
||||
// 获取可用的智能体运行模式列表
|
||||
pub async fn get_agent_modes() -> Json<Vec<AgentModeDto>> {
|
||||
use crate::agent::modes::ModeRegistry;
|
||||
let registry = ModeRegistry::builtins();
|
||||
let modes = registry
|
||||
.list()
|
||||
.iter()
|
||||
.map(|m| AgentModeDto {
|
||||
id: m.id,
|
||||
name: m.name,
|
||||
description: m.description,
|
||||
icon: m.icon,
|
||||
})
|
||||
.collect();
|
||||
Json(modes)
|
||||
}
|
||||
|
||||
pub async fn chat_agent(
|
||||
@ -39,21 +86,112 @@ pub async fn chat_agent(
|
||||
Json(req): Json<AgentChatRequest>,
|
||||
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (StatusCode, String)> {
|
||||
info!(
|
||||
"接收到智能体对话请求: question='{}', session_id={:?}",
|
||||
req.question, req.session_id
|
||||
"接收到智能体对话请求: question='{}', session_id={:?}, has_image={}",
|
||||
req.question,
|
||||
req.session_id,
|
||||
req.image.is_some()
|
||||
);
|
||||
|
||||
let runtime = AgentRuntime::new(Arc::clone(&state))
|
||||
.with_thinking(req.thinking)
|
||||
.with_coordinator_mode(req.coordinator_mode);
|
||||
// 处理图片附件：保存到磁盘，路径注入 Agent 上下文，前端和 DB 保留原始问题
|
||||
let question = req.question.clone();
|
||||
let (image_context, image_path_for_db): (Option<String>, Option<String>) = match req.image {
|
||||
Some(ref img) => {
|
||||
// 如果带有 path 字段（重试场景），复用已有文件，不重新保存
|
||||
let relative_path: String = if let Some(ref existing_path) = img.path {
|
||||
let full = state.config.library_dir.join(existing_path);
|
||||
if full.exists() {
|
||||
info!("重试复用已有图片: {}", existing_path);
|
||||
existing_path.clone()
|
||||
} else {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("图片文件不存在: {}", existing_path),
|
||||
));
|
||||
}
|
||||
} else {
|
||||
if img.data.is_empty() {
|
||||
return Err((StatusCode::BAD_REQUEST, "图片数据为空".to_string()));
|
||||
}
|
||||
if !img.mime_type.starts_with("image/") {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("不支持的图片类型: {}", img.mime_type),
|
||||
));
|
||||
}
|
||||
if state.vision_llm.is_none() {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"图片分析功能未启用。请配置 LLM_VISION_MODEL 环境变量后重试。".to_string(),
|
||||
));
|
||||
}
|
||||
let ext = img.mime_type.strip_prefix("image/").unwrap_or("png");
|
||||
let upload_dir = state
|
||||
.config
|
||||
.library_dir
|
||||
.join(".agent_images")
|
||||
.join("uploads");
|
||||
std::fs::create_dir_all(&upload_dir).map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("创建上传目录失败: {}", e),
|
||||
)
|
||||
})?;
|
||||
let filename = format!("{}.{}", uuid::Uuid::new_v4(), ext);
|
||||
let filepath = upload_dir.join(&filename);
|
||||
use base64::{engine::general_purpose, Engine as _};
|
||||
let bytes = general_purpose::STANDARD.decode(&img.data).map_err(|e| {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("图片 base64 解码失败: {}", e),
|
||||
)
|
||||
})?;
|
||||
std::fs::write(&filepath, &bytes).map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("保存图片失败: {}", e),
|
||||
)
|
||||
})?;
|
||||
let rel = filepath
|
||||
.strip_prefix(&state.config.library_dir)
|
||||
.unwrap_or(&filepath)
|
||||
.display()
|
||||
.to_string();
|
||||
info!("用户图片已保存: {}", rel);
|
||||
rel
|
||||
};
|
||||
|
||||
let ctx = format!(
|
||||
"用户上传了一张图片，已保存到: {}\n如需分析此图片，请使用 analyze_image 工具，传入 image_path=\"{}\"。",
|
||||
relative_path, relative_path
|
||||
);
|
||||
(Some(ctx), Some(relative_path))
|
||||
}
|
||||
None => (None, None),
|
||||
};
|
||||
|
||||
let mut runtime = AgentRuntime::new(Arc::clone(&state)).with_mode(&req.mode);
|
||||
// 只有 mode 未强制固定 thinking 时，用户才可以覆盖
|
||||
if runtime.mode_fixed_thinking().is_none() {
|
||||
if let Some(thinking) = req.thinking {
|
||||
runtime = runtime.with_thinking(thinking);
|
||||
}
|
||||
}
|
||||
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<AgentStreamEvent>();
|
||||
|
||||
let question = req.question.clone();
|
||||
let session_id = req.session_id.clone();
|
||||
|
||||
// 在后台 tokio 任务中执行 Agent 循环
|
||||
tokio::spawn(async move {
|
||||
match runtime.run_turn(session_id, &question, tx.clone()).await {
|
||||
match runtime
|
||||
.run_turn_with_image_context(
|
||||
session_id,
|
||||
&question,
|
||||
image_context,
|
||||
image_path_for_db,
|
||||
tx.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(sid) => {
|
||||
info!("智能体对话完成: session_id={}", sid);
|
||||
}
|
||||
@ -96,6 +234,7 @@ pub struct SessionSummary {
|
||||
pub session_id: String,
|
||||
pub title: String,
|
||||
pub model: String,
|
||||
pub mode: String,
|
||||
pub turn_count: i32,
|
||||
pub summary: Option<String>,
|
||||
pub created_at: String,
|
||||
@ -110,7 +249,7 @@ pub async fn list_sessions(
|
||||
let offset = params.offset.unwrap_or(0);
|
||||
|
||||
let rows = sqlx::query(
|
||||
"SELECT session_id, title, model, turn_count, summary, created_at, updated_at \
|
||||
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \
|
||||
FROM agent_sessions \
|
||||
WHERE deleted_at IS NULL \
|
||||
ORDER BY updated_at DESC \
|
||||
@ -133,10 +272,11 @@ pub async fn list_sessions(
|
||||
session_id: r.get(0),
|
||||
title: r.get(1),
|
||||
model: r.get(2),
|
||||
turn_count: r.get(3),
|
||||
summary: r.get(4),
|
||||
created_at: r.get(5),
|
||||
updated_at: r.get(6),
|
||||
mode: r.get(3),
|
||||
turn_count: r.get(4),
|
||||
summary: r.get(5),
|
||||
created_at: r.get(6),
|
||||
updated_at: r.get(7),
|
||||
})
|
||||
.collect();
|
||||
|
||||
@ -174,7 +314,7 @@ pub async fn get_session(
|
||||
) -> Result<Json<SessionDetail>, (StatusCode, String)> {
|
||||
// 查询会话元信息
|
||||
let session_row = sqlx::query(
|
||||
"SELECT session_id, title, model, turn_count, summary, created_at, updated_at \
|
||||
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \
|
||||
FROM agent_sessions \
|
||||
WHERE session_id = ? AND deleted_at IS NULL",
|
||||
)
|
||||
@ -193,10 +333,11 @@ pub async fn get_session(
|
||||
session_id: session_row.get(0),
|
||||
title: session_row.get(1),
|
||||
model: session_row.get(2),
|
||||
turn_count: session_row.get(3),
|
||||
summary: session_row.get(4),
|
||||
created_at: session_row.get(5),
|
||||
updated_at: session_row.get(6),
|
||||
mode: session_row.get(3),
|
||||
turn_count: session_row.get(4),
|
||||
summary: session_row.get(5),
|
||||
created_at: session_row.get(6),
|
||||
updated_at: session_row.get(7),
|
||||
};
|
||||
|
||||
// 查询消息列表（包含 lead 和 subagent 消息，前端按 agent_name/metadata 区分渲染）
|
||||
@ -587,13 +728,15 @@ pub struct RetryResponse {
|
||||
pub new_turn_index: i32,
|
||||
pub deleted_count: i64,
|
||||
pub session_id: String,
|
||||
/// 原消息附带的图片路径（如果有）
|
||||
pub image_path: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn retry_session(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Json<RetryResponse>, (StatusCode, String)> {
|
||||
let (retried_message, new_turn_index) =
|
||||
let (retried_message, new_turn_index, image_path) =
|
||||
crate::agent::runtime::session::retry_last_turn(&state.db, &session_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;
|
||||
@ -601,8 +744,9 @@ pub async fn retry_session(
|
||||
Ok(Json(RetryResponse {
|
||||
retried_message,
|
||||
new_turn_index,
|
||||
deleted_count: 0, // 数据库层不便返回，设为 0
|
||||
deleted_count: 0,
|
||||
session_id,
|
||||
image_path,
|
||||
}))
|
||||
}
|
||||
|
||||
|
||||
@ -40,8 +40,10 @@ fn prune_expired_sessions(
|
||||
let now = chrono::Utc::now();
|
||||
let expiry_duration = chrono::Duration::hours(24);
|
||||
sessions.retain(|_, last_active| {
|
||||
if let Some(dur) = now.signed_duration_since(*last_active).to_std().ok() {
|
||||
dur < expiry_duration.to_std().unwrap_or(std::time::Duration::from_secs(86400))
|
||||
if let Ok(dur) = now.signed_duration_since(*last_active).to_std() {
|
||||
dur < expiry_duration
|
||||
.to_std()
|
||||
.unwrap_or(std::time::Duration::from_secs(86400))
|
||||
} else {
|
||||
true // 尚未到达或时钟回拨，保留
|
||||
}
|
||||
@ -56,7 +58,7 @@ pub async fn login(
|
||||
let password = req.password.unwrap_or_default();
|
||||
if password == state.config.admin_password {
|
||||
let token = uuid::Uuid::new_v4().to_string();
|
||||
|
||||
|
||||
// 异常安全地记录活跃会话及其当前时间
|
||||
if let Ok(mut sessions) = state.sessions.lock() {
|
||||
sessions.insert(token.clone(), chrono::Utc::now());
|
||||
@ -94,13 +96,14 @@ pub async fn logout(
|
||||
let mut token_to_remove = None;
|
||||
|
||||
// 优先从 Authorization 提取 Token（向前兼容旧 Bearer Header）
|
||||
let auth_header = req.headers()
|
||||
let auth_header = req
|
||||
.headers()
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|val| val.to_str().ok());
|
||||
|
||||
if let Some(auth_str) = auth_header {
|
||||
if auth_str.starts_with("Bearer ") {
|
||||
token_to_remove = Some(auth_str[7..].to_string());
|
||||
if let Some(rest) = auth_str.strip_prefix("Bearer ") {
|
||||
token_to_remove = Some(rest.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
@ -143,13 +146,14 @@ pub async fn auth_middleware(
|
||||
let mut token = None;
|
||||
|
||||
// 1. 优先从 Authorization 头读取 Bearer Token（兼容）
|
||||
let auth_header = req.headers()
|
||||
let auth_header = req
|
||||
.headers()
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|val| val.to_str().ok());
|
||||
|
||||
if let Some(auth_str) = auth_header {
|
||||
if auth_str.starts_with("Bearer ") {
|
||||
token = Some(auth_str[7..].to_string());
|
||||
if let Some(rest) = auth_str.strip_prefix("Bearer ") {
|
||||
token = Some(rest.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -84,6 +84,14 @@ pub fn convert_arxiv_to_standard(doc: &ArxivPaper) -> StandardPaper {
|
||||
}
|
||||
|
||||
pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Result<()> {
|
||||
let mut conn = db.acquire().await?;
|
||||
save_paper_to_db_tx(&mut conn, p).await
|
||||
}
|
||||
|
||||
pub async fn save_paper_to_db_tx(
|
||||
conn: &mut sqlx::SqliteConnection,
|
||||
p: &StandardPaper,
|
||||
) -> anyhow::Result<()> {
|
||||
let authors_json = serde_json::to_string(&p.authors)?;
|
||||
let keywords_json = serde_json::to_string(&p.keywords)?;
|
||||
|
||||
@ -94,7 +102,7 @@ pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Res
|
||||
"SELECT bibcode, pdf_path, html_path, markdown_path, translation_path FROM papers WHERE arxiv_id = ?"
|
||||
)
|
||||
.bind(&p.arxiv_id)
|
||||
.fetch_optional(db)
|
||||
.fetch_optional(&mut *conn)
|
||||
.await?;
|
||||
|
||||
if let Some((existing_bibcode, _pdf, _html, _md, _tr)) = existing_opt {
|
||||
@ -124,7 +132,7 @@ pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Res
|
||||
.bind(p.reference_count)
|
||||
.bind(&p.doctype)
|
||||
.bind(&existing_bibcode)
|
||||
.execute(db)
|
||||
.execute(&mut *conn)
|
||||
.await?;
|
||||
|
||||
return Ok(());
|
||||
@ -168,7 +176,7 @@ pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Res
|
||||
.bind(p.citation_count)
|
||||
.bind(p.reference_count)
|
||||
.bind(&p.doctype)
|
||||
.execute(db)
|
||||
.execute(&mut *conn)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
|
||||
@ -57,6 +57,10 @@ pub struct AppState {
|
||||
pub ads: AdsClient,
|
||||
pub arxiv: ArxivClient,
|
||||
pub llm: LlmClient,
|
||||
pub medium_llm: LlmClient,
|
||||
pub fast_llm: LlmClient,
|
||||
/// 专用视觉模型（LLM_VISION_MODEL 配置时启用图片分析工具）
|
||||
pub vision_llm: Option<LlmClient>,
|
||||
pub embedding: EmbeddingClient,
|
||||
pub downloader: Downloader,
|
||||
pub harvest_status: Arc<tokio::sync::Mutex<crate::services::batch_sync::MetaSyncStatus>>,
|
||||
@ -76,7 +80,8 @@ pub struct AppState {
|
||||
/// 项目记忆管理器（跨会话持久化）
|
||||
pub memory_manager: Arc<tokio::sync::Mutex<MemoryManager>>,
|
||||
/// 活跃的登录会话 Token 及其最后活跃时间（单用户内存管理）
|
||||
pub sessions: Arc<std::sync::Mutex<std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>>>,
|
||||
pub sessions:
|
||||
Arc<std::sync::Mutex<std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>>>,
|
||||
}
|
||||
|
||||
// 统一标准化的文献格式，用于向前端传输
|
||||
@ -118,16 +123,16 @@ pub mod targets;
|
||||
pub mod handlers {
|
||||
pub use super::agent::{
|
||||
answer_question, branch_session, chat_agent, delete_session, get_agent_metrics,
|
||||
get_pending_permissions, get_pending_questions, get_session, get_session_audit,
|
||||
list_sessions, respond_permission, restore_rewound_session, retry_session, rewind_session,
|
||||
stop_agent, AgentChatRequest, AgentMetricsResponse, AuditLogEntry, BranchResponse,
|
||||
MessageRecord, RestoreResponse, RetryResponse, RewindRequest, RewindResponse,
|
||||
SessionDetail, SessionListParams, SessionSummary,
|
||||
get_agent_modes, get_pending_permissions, get_pending_questions, get_session,
|
||||
get_session_audit, list_sessions, respond_permission, restore_rewound_session,
|
||||
retry_session, rewind_session, stop_agent, AgentChatRequest, AgentMetricsResponse,
|
||||
AuditLogEntry, BranchResponse, MessageRecord, RestoreResponse, RetryResponse,
|
||||
RewindRequest, RewindResponse, SessionDetail, SessionListParams, SessionSummary,
|
||||
};
|
||||
pub use super::auth::{check_auth, login, logout};
|
||||
pub use super::helpers::{
|
||||
check_paper_paths_in_db, convert_ads_doc_to_standard, convert_arxiv_to_standard,
|
||||
get_paper_from_db, save_paper_to_db,
|
||||
get_paper_from_db, save_paper_to_db, save_paper_to_db_tx,
|
||||
};
|
||||
pub use super::notes::{
|
||||
create_note, delete_note, get_notes, CreateNoteRequest, DeleteNoteParams, GetNotesParams,
|
||||
@ -149,9 +154,8 @@ pub mod handlers {
|
||||
MetaSyncCountRequest, MetaSyncCountResponse, MetaSyncRunRequest, SavedSyncQuery,
|
||||
};
|
||||
pub use super::targets::{
|
||||
associate_target, chat_figure, chat_rag, extract_paper_targets, list_targets, query_target,
|
||||
AssociateResponse, AssociateTargetRequest, ChatFigureRequest, ChatResponse,
|
||||
ExtractTargetsRequest, ExtractTargetsResponse, RagAskRequest, TargetListParams,
|
||||
associate_target, extract_paper_targets, list_targets, query_target, AssociateResponse,
|
||||
AssociateTargetRequest, ExtractTargetsRequest, ExtractTargetsResponse, TargetListParams,
|
||||
TargetQueryParams,
|
||||
};
|
||||
pub use super::{AppState, StandardPaper};
|
||||
|
||||
@ -203,7 +203,7 @@ pub async fn translate_paper(
|
||||
let translated_markdown = crate::services::translation::translate_markdown(
|
||||
&english_markdown,
|
||||
&state.dict,
|
||||
&state.llm,
|
||||
&state.medium_llm,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
|
||||
@ -8,15 +8,8 @@ use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::AppState;
|
||||
use crate::services::rag::{ask, RagAnswer};
|
||||
use crate::services::target::{query_target_cached, TargetInfo};
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct RagAskRequest {
|
||||
pub question: String,
|
||||
pub top_k: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct TargetQueryParams {
|
||||
pub object_name: String,
|
||||
@ -39,32 +32,6 @@ pub struct AssociateResponse {
|
||||
pub target: TargetInfo,
|
||||
}
|
||||
|
||||
// ── POST /api/chat/rag ──
|
||||
pub async fn chat_rag(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(req): Json<RagAskRequest>,
|
||||
) -> Result<Json<RagAnswer>, (StatusCode, String)> {
|
||||
let top_k = req.top_k.unwrap_or(5);
|
||||
|
||||
let answer = ask(
|
||||
&state.db,
|
||||
&state.embedding,
|
||||
&state.llm,
|
||||
&req.question,
|
||||
top_k,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!("RAG 问答执行失败: {}", e);
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("RAG 问答执行失败: {}", e),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(Json(answer))
|
||||
}
|
||||
|
||||
// ── GET /api/target/query ──
|
||||
pub async fn query_target(
|
||||
State(state): State<Arc<AppState>>,
|
||||
@ -217,92 +184,3 @@ pub async fn extract_paper_targets(
|
||||
|
||||
Ok(Json(ExtractTargetsResponse { targets }))
|
||||
}
|
||||
|
||||
// ── POST /api/chat/figure ──
|
||||
#[derive(Deserialize)]
|
||||
pub struct ChatFigureRequest {
|
||||
pub bibcode: String,
|
||||
pub image_path: String,
|
||||
pub question: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct ChatResponse {
|
||||
pub answer: String,
|
||||
}
|
||||
|
||||
async fn get_image_bytes(
|
||||
library_dir: &std::path::Path,
|
||||
path: &str,
|
||||
) -> anyhow::Result<(Vec<u8>, String)> {
|
||||
if path.starts_with("http://") || path.starts_with("https://") {
|
||||
let client = reqwest::Client::new();
|
||||
let resp = client.get(path).send().await?;
|
||||
let headers = resp.headers().clone();
|
||||
let mime_type = headers
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("image/png")
|
||||
.to_string();
|
||||
let bytes = resp.bytes().await?.to_vec();
|
||||
Ok((bytes, mime_type))
|
||||
} else {
|
||||
let abs_path = library_dir.join(path);
|
||||
if !abs_path.exists() {
|
||||
return Err(anyhow::anyhow!("本地图片未找到: {:?}", abs_path));
|
||||
}
|
||||
let ext = abs_path
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.unwrap_or("png")
|
||||
.to_lowercase();
|
||||
let mime_type = match ext.as_str() {
|
||||
"jpg" | "jpeg" => "image/jpeg",
|
||||
"gif" => "image/gif",
|
||||
"webp" => "image/webp",
|
||||
"svg" => "image/svg+xml",
|
||||
_ => "image/png",
|
||||
}
|
||||
.to_string();
|
||||
let bytes = std::fs::read(&abs_path)?;
|
||||
Ok((bytes, mime_type))
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn chat_figure(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(req): Json<ChatFigureRequest>,
|
||||
) -> Result<Json<ChatResponse>, (StatusCode, String)> {
|
||||
tracing::info!(
|
||||
"接收到图表多模态提问: bibcode={}, image_path={}, question={}",
|
||||
req.bibcode,
|
||||
req.image_path,
|
||||
req.question
|
||||
);
|
||||
|
||||
let (bytes, mime_type) = get_image_bytes(&state.config.library_dir, &req.image_path)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!("获取图片失败: {}", e);
|
||||
(StatusCode::BAD_REQUEST, format!("获取图片失败: {}", e))
|
||||
})?;
|
||||
|
||||
use base64::{engine::general_purpose, Engine as _};
|
||||
let base64_data = general_purpose::STANDARD.encode(&bytes);
|
||||
|
||||
let system_prompt = "You are a professional astronomer and physicist. You are helping a researcher analyze a scientific plot, diagram, or chart extracted from a theoretical or observational astrophysics paper. Please provide an expert, detailed, and clear explanation of the figure in Chinese based on the image provided and the user's question. Format equations in standard LaTeX.";
|
||||
|
||||
let answer = state
|
||||
.llm
|
||||
.chat_completion_with_image(system_prompt, &req.question, &base64_data, &mime_type)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!("多模态大模型请求失败: {}", e);
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("大模型图表解析失败: {}", e),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(Json(ChatResponse { answer }))
|
||||
}
|
||||
|
||||
@ -31,7 +31,9 @@ pub struct FunctionCall {
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
/// 标准对话消息（OpenAI 兼容格式）
|
||||
// ── ChatMessage ─────────────────────────────────────────────────────────
|
||||
|
||||
/// 标准对话消息（OpenAI 兼容格式）。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChatMessage {
|
||||
pub role: MessageRole,
|
||||
@ -48,6 +50,11 @@ pub struct ChatMessage {
|
||||
}
|
||||
|
||||
impl ChatMessage {
|
||||
/// 获取纯文本内容。
|
||||
pub fn text(&self) -> Option<&str> {
|
||||
self.content.as_deref()
|
||||
}
|
||||
|
||||
/// 创建系统消息
|
||||
pub fn system(content: impl Into<String>) -> Self {
|
||||
ChatMessage {
|
||||
@ -302,86 +309,6 @@ impl LlmClient {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn chat_completion_with_image(
|
||||
&self,
|
||||
system_prompt: &str,
|
||||
user_prompt: &str,
|
||||
image_base64: &str,
|
||||
mime_type: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
let url = format!("{}/chat/completions", self.api_base);
|
||||
|
||||
let image_url_val = if image_base64.starts_with("data:") {
|
||||
image_base64.to_string()
|
||||
} else {
|
||||
format!("data:{};base64,{}", mime_type, image_base64)
|
||||
};
|
||||
|
||||
let payload = serde_json::json!({
|
||||
"model": self.model(),
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": system_prompt
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": user_prompt
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": image_url_val
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"temperature": 0.3
|
||||
});
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
error!("LLM 多模态接口调用失败: 状态码={}, 报错={}", status, body);
|
||||
return Err(anyhow::anyhow!("大模型多模态接口返回错误状态: {}", status));
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Message {
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Choice {
|
||||
message: Message,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LLMResponse {
|
||||
choices: Vec<Choice>,
|
||||
}
|
||||
|
||||
let res_data: LLMResponse = response.json().await?;
|
||||
if let Some(choice) = res_data.choices.first() {
|
||||
Ok(choice.message.content.clone())
|
||||
} else {
|
||||
Err(anyhow::anyhow!("大模型返回空多模态回答选项集"))
|
||||
}
|
||||
}
|
||||
|
||||
/// 多轮对话补全（含原生 Tool Calling 支持），非流式
|
||||
pub async fn chat(
|
||||
&self,
|
||||
@ -462,6 +389,87 @@ impl LlmClient {
|
||||
}
|
||||
}
|
||||
|
||||
/// 流式分析图片，通过 `on_delta` 回调推送增量文本，返回完整结果。
|
||||
pub async fn analyze_image_stream(
|
||||
&self,
|
||||
system_prompt: &str,
|
||||
question: &str,
|
||||
image_base64: &str,
|
||||
mime_type: &str,
|
||||
mut on_delta: impl FnMut(&str),
|
||||
) -> anyhow::Result<String> {
|
||||
let user_content = serde_json::json!([
|
||||
{ "type": "text", "text": question },
|
||||
{ "type": "image_url", "image_url": { "url": format!("data:{};base64,{}", mime_type, image_base64) } }
|
||||
]);
|
||||
|
||||
let payload = serde_json::json!({
|
||||
"model": self.model(),
|
||||
"messages": [
|
||||
{ "role": "system", "content": system_prompt },
|
||||
{ "role": "user", "content": user_content }
|
||||
],
|
||||
"temperature": 0.3,
|
||||
"stream": true,
|
||||
"stream_options": { "include_usage": true }
|
||||
});
|
||||
|
||||
let url = format!("{}/chat/completions", self.api_base);
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(anyhow::anyhow!(
|
||||
"视觉模型调用失败 (HTTP {}): {}",
|
||||
status,
|
||||
body
|
||||
));
|
||||
}
|
||||
|
||||
use futures_util::StreamExt;
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut accumulated = String::new();
|
||||
let mut buffer = String::new();
|
||||
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk?;
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
while let Some(pos) = buffer.find('\n') {
|
||||
let line = buffer[..pos].trim().to_string();
|
||||
buffer = buffer[pos + 1..].to_string();
|
||||
|
||||
if line.is_empty() || !line.starts_with("data: ") {
|
||||
continue;
|
||||
}
|
||||
let data = line[6..].trim().to_string();
|
||||
if data == "[DONE]" {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Ok(event) = serde_json::from_str::<serde_json::Value>(&data) {
|
||||
if let Some(delta) = event["choices"][0]["delta"]["content"].as_str() {
|
||||
accumulated.push_str(delta);
|
||||
on_delta(delta);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if accumulated.is_empty() {
|
||||
return Err(anyhow::anyhow!("视觉模型返回空结果"));
|
||||
}
|
||||
Ok(accumulated)
|
||||
}
|
||||
|
||||
/// 流式多轮对话补全（含 Tool Call 碎片累积还原），返回 SSE 事件流
|
||||
pub async fn chat_stream(
|
||||
&self,
|
||||
@ -831,27 +839,24 @@ mod tests {
|
||||
// 测试系统消息构造
|
||||
let sys = ChatMessage::system("你是一个助手");
|
||||
assert_eq!(sys.role, MessageRole::System);
|
||||
assert_eq!(sys.content.as_deref(), Some("你是一个助手"));
|
||||
assert_eq!(sys.text(), Some("你是一个助手"));
|
||||
assert!(sys.tool_calls.is_none());
|
||||
|
||||
// 测试用户消息构造
|
||||
let user = ChatMessage::user("你好");
|
||||
assert_eq!(user.role, MessageRole::User);
|
||||
assert_eq!(user.content.as_deref(), Some("你好"));
|
||||
assert_eq!(user.text(), Some("你好"));
|
||||
|
||||
// 测试助手消息构造
|
||||
let assistant = ChatMessage::assistant("你好！有什么可以帮助你的吗？");
|
||||
assert_eq!(assistant.role, MessageRole::Assistant);
|
||||
assert_eq!(
|
||||
assistant.content.as_deref(),
|
||||
Some("你好！有什么可以帮助你的吗？")
|
||||
);
|
||||
assert_eq!(assistant.text(), Some("你好！有什么可以帮助你的吗？"));
|
||||
|
||||
// 测试工具结果消息构造
|
||||
let tool_result = ChatMessage::tool_result("call_123", r#"{"result": 42}"#);
|
||||
assert_eq!(tool_result.role, MessageRole::Tool);
|
||||
assert_eq!(tool_result.tool_call_id.as_deref(), Some("call_123"));
|
||||
assert_eq!(tool_result.content.as_deref(), Some(r#"{"result": 42}"#));
|
||||
assert_eq!(tool_result.text(), Some(r#"{"result": 42}"#));
|
||||
|
||||
// 测试带工具调用的助手消息
|
||||
let tool_calls = vec![ToolCall {
|
||||
@ -864,7 +869,7 @@ mod tests {
|
||||
}];
|
||||
let assistant_tc = ChatMessage::assistant_with_tool_calls(None, tool_calls.clone());
|
||||
assert_eq!(assistant_tc.role, MessageRole::Assistant);
|
||||
assert!(assistant_tc.content.is_none());
|
||||
assert!(assistant_tc.text().is_none());
|
||||
assert_eq!(assistant_tc.tool_calls.as_ref().unwrap().len(), 1);
|
||||
assert_eq!(
|
||||
assistant_tc.tool_calls.as_ref().unwrap()[0].function.name,
|
||||
@ -924,7 +929,7 @@ mod tests {
|
||||
// 测试反序列化
|
||||
let deserialized: ChatMessage = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(deserialized.role, MessageRole::Assistant);
|
||||
assert_eq!(deserialized.content.as_deref(), Some("回答内容"));
|
||||
assert_eq!(deserialized.text(), Some("回答内容"));
|
||||
assert_eq!(
|
||||
deserialized.reasoning_content.as_deref(),
|
||||
Some("这是思考过程")
|
||||
|
||||
116
src/lib.rs
116
src/lib.rs
@ -12,6 +12,12 @@ pub struct Config {
|
||||
pub llm_model: String, // 调用的翻译大模型名称
|
||||
pub llm_fallback_model: String, // 故障转移备用模型
|
||||
pub llm_fallback_chain: Vec<String>, // 故障转移模型链（按顺序尝试）
|
||||
pub llm_medium_model: String, // 中型大模型名称
|
||||
pub llm_medium_api_key: String, // 中型大模型 API Key
|
||||
pub llm_medium_api_base: String, // 中型大模型 API 基础地址
|
||||
pub llm_fast_model: String, // 快速大模型名称
|
||||
pub llm_fast_api_key: String, // 快速大模型 API Key
|
||||
pub llm_fast_api_base: String, // 快速大模型 API 基础地址
|
||||
pub embedding_api_key: String, // 向量模型 API Key
|
||||
pub embedding_api_base: String, // 向量模型 API 基础地址
|
||||
pub embedding_model: String, // 向量模型名称
|
||||
@ -25,6 +31,12 @@ pub struct Config {
|
||||
pub skills_dir: PathBuf, // Agent Skills 目录（Markdown 知识模块）
|
||||
pub port: u16, // 后端服务监听端口
|
||||
pub admin_password: String, // 单用户管理员登录密码
|
||||
/// 专用视觉模型名称（空字符串 = 未配置，不启用图片分析能力）
|
||||
pub llm_vision_model: String,
|
||||
/// 视觉模型 API Key（空时回退到 llm_api_key）
|
||||
pub llm_vision_api_key: String,
|
||||
/// 视觉模型 API Base（空时回退到 llm_api_base）
|
||||
pub llm_vision_api_base: String,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
@ -49,6 +61,21 @@ impl Config {
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect();
|
||||
|
||||
// Tier 2: Medium Model configurations (defaulting to Tier 1)
|
||||
let llm_medium_model = env::var("LLM_MEDIUM_MODEL").unwrap_or_else(|_| llm_model.clone());
|
||||
let llm_medium_api_key =
|
||||
env::var("LLM_MEDIUM_API_KEY").unwrap_or_else(|_| llm_api_key.clone());
|
||||
let llm_medium_api_base =
|
||||
env::var("LLM_MEDIUM_API_BASE").unwrap_or_else(|_| llm_api_base.clone());
|
||||
|
||||
// Tier 3: Fast Model configurations (cascading fallback: Fast -> Medium)
|
||||
let llm_fast_model =
|
||||
env::var("LLM_FAST_MODEL").unwrap_or_else(|_| llm_medium_model.clone());
|
||||
let llm_fast_api_key =
|
||||
env::var("LLM_FAST_API_KEY").unwrap_or_else(|_| llm_medium_api_key.clone());
|
||||
let llm_fast_api_base =
|
||||
env::var("LLM_FAST_API_BASE").unwrap_or_else(|_| llm_medium_api_base.clone());
|
||||
|
||||
let embedding_api_key =
|
||||
env::var("EMBEDDING_API_KEY").unwrap_or_else(|_| llm_api_key.clone());
|
||||
let embedding_api_base =
|
||||
@ -75,8 +102,11 @@ impl Config {
|
||||
.parse::<u16>()
|
||||
.unwrap_or(8000);
|
||||
|
||||
let admin_password = env::var("ADMIN_PASSWORD")
|
||||
.unwrap_or_else(|_| "admin".to_string());
|
||||
let admin_password = env::var("ADMIN_PASSWORD").unwrap_or_else(|_| "admin".to_string());
|
||||
|
||||
let llm_vision_model = env::var("LLM_VISION_MODEL").unwrap_or_default();
|
||||
let llm_vision_api_key = env::var("LLM_VISION_API_KEY").unwrap_or_default();
|
||||
let llm_vision_api_base = env::var("LLM_VISION_API_BASE").unwrap_or_default();
|
||||
|
||||
Config {
|
||||
database_url,
|
||||
@ -86,6 +116,12 @@ impl Config {
|
||||
llm_model,
|
||||
llm_fallback_model,
|
||||
llm_fallback_chain,
|
||||
llm_medium_model,
|
||||
llm_medium_api_key,
|
||||
llm_medium_api_base,
|
||||
llm_fast_model,
|
||||
llm_fast_api_key,
|
||||
llm_fast_api_base,
|
||||
embedding_api_key,
|
||||
embedding_api_base,
|
||||
embedding_model,
|
||||
@ -99,6 +135,9 @@ impl Config {
|
||||
skills_dir,
|
||||
port,
|
||||
admin_password,
|
||||
llm_vision_model,
|
||||
llm_vision_api_key,
|
||||
llm_vision_api_base,
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -136,4 +175,77 @@ mod config_tests {
|
||||
std::env::remove_var("DATABASE_URL");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_capability_tiering_fallback() {
|
||||
// 备份并清理可能存在的环境变量
|
||||
let vars_to_test = [
|
||||
"LLM_MODEL",
|
||||
"LLM_API_KEY",
|
||||
"LLM_API_BASE",
|
||||
"LLM_MEDIUM_MODEL",
|
||||
"LLM_MEDIUM_API_KEY",
|
||||
"LLM_MEDIUM_API_BASE",
|
||||
"LLM_FAST_MODEL",
|
||||
"LLM_FAST_API_KEY",
|
||||
"LLM_FAST_API_BASE",
|
||||
];
|
||||
let backups: Vec<(&str, Option<String>)> = vars_to_test
|
||||
.iter()
|
||||
.map(|&var| (var, std::env::var(var).ok()))
|
||||
.collect();
|
||||
|
||||
for &var in &vars_to_test {
|
||||
std::env::remove_var(var);
|
||||
}
|
||||
|
||||
// 场景 1：仅设置一级模型，二级、三级应该级联回退到一级
|
||||
std::env::set_var("LLM_MODEL", "primary-gpt");
|
||||
std::env::set_var("LLM_API_KEY", "pk-123");
|
||||
std::env::set_var("LLM_API_BASE", "http://primary");
|
||||
|
||||
let config = Config::from_env();
|
||||
assert_eq!(config.llm_model, "primary-gpt");
|
||||
assert_eq!(config.llm_medium_model, "primary-gpt");
|
||||
assert_eq!(config.llm_medium_api_key, "pk-123");
|
||||
assert_eq!(config.llm_medium_api_base, "http://primary");
|
||||
assert_eq!(config.llm_fast_model, "primary-gpt");
|
||||
assert_eq!(config.llm_fast_api_key, "pk-123");
|
||||
assert_eq!(config.llm_fast_api_base, "http://primary");
|
||||
|
||||
// 场景 2：同时设置一级和二级，三级应该回退到二级
|
||||
std::env::set_var("LLM_MEDIUM_MODEL", "medium-gpt");
|
||||
std::env::set_var("LLM_MEDIUM_API_KEY", "med-456");
|
||||
std::env::set_var("LLM_MEDIUM_API_BASE", "http://medium");
|
||||
|
||||
let config = Config::from_env();
|
||||
assert_eq!(config.llm_model, "primary-gpt");
|
||||
assert_eq!(config.llm_medium_model, "medium-gpt");
|
||||
assert_eq!(config.llm_medium_api_key, "med-456");
|
||||
assert_eq!(config.llm_medium_api_base, "http://medium");
|
||||
assert_eq!(config.llm_fast_model, "medium-gpt");
|
||||
assert_eq!(config.llm_fast_api_key, "med-456");
|
||||
assert_eq!(config.llm_fast_api_base, "http://medium");
|
||||
|
||||
// 场景 3：全部设置，各司其职
|
||||
std::env::set_var("LLM_FAST_MODEL", "fast-gpt");
|
||||
std::env::set_var("LLM_FAST_API_KEY", "fast-789");
|
||||
std::env::set_var("LLM_FAST_API_BASE", "http://fast");
|
||||
|
||||
let config = Config::from_env();
|
||||
assert_eq!(config.llm_model, "primary-gpt");
|
||||
assert_eq!(config.llm_medium_model, "medium-gpt");
|
||||
assert_eq!(config.llm_fast_model, "fast-gpt");
|
||||
assert_eq!(config.llm_fast_api_key, "fast-789");
|
||||
assert_eq!(config.llm_fast_api_base, "http://fast");
|
||||
|
||||
// 恢复环境变量
|
||||
for (var, val) in backups {
|
||||
if let Some(v) = val {
|
||||
std::env::set_var(var, v);
|
||||
} else {
|
||||
std::env::remove_var(var);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
73
src/main.rs
73
src/main.rs
@ -70,8 +70,10 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Agent Skills 目录
|
||||
std::fs::create_dir_all(&config.skills_dir).unwrap_or_default();
|
||||
|
||||
// 3. 初始化本地 SQLite 数据库连接池（开启外键约束）
|
||||
// 3. 初始化本地 SQLite 数据库连接池（开启外键约束并启用 WAL 模式与繁忙等待）
|
||||
let options = SqliteConnectOptions::from_str(&config.database_url)?
|
||||
.journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
|
||||
.busy_timeout(std::time::Duration::from_secs(10))
|
||||
.foreign_keys(true)
|
||||
.create_if_missing(true);
|
||||
|
||||
@ -145,6 +147,32 @@ async fn main() -> anyhow::Result<()> {
|
||||
config.llm_api_base.clone(),
|
||||
config.llm_model.clone(),
|
||||
);
|
||||
let medium_llm = LlmClient::new(
|
||||
config.llm_medium_api_key.clone(),
|
||||
config.llm_medium_api_base.clone(),
|
||||
config.llm_medium_model.clone(),
|
||||
);
|
||||
let fast_llm = LlmClient::new(
|
||||
config.llm_fast_api_key.clone(),
|
||||
config.llm_fast_api_base.clone(),
|
||||
config.llm_fast_model.clone(),
|
||||
);
|
||||
let vision_llm = if !config.llm_vision_model.is_empty() {
|
||||
let key = if config.llm_vision_api_key.is_empty() {
|
||||
config.llm_api_key.clone()
|
||||
} else {
|
||||
config.llm_vision_api_key.clone()
|
||||
};
|
||||
let base = if config.llm_vision_api_base.is_empty() {
|
||||
config.llm_api_base.clone()
|
||||
} else {
|
||||
config.llm_vision_api_base.clone()
|
||||
};
|
||||
info!("视觉模型已启用: {}", config.llm_vision_model);
|
||||
Some(LlmClient::new(key, base, config.llm_vision_model.clone()))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let embedding = EmbeddingClient::new(
|
||||
config.embedding_api_key.clone(),
|
||||
config.embedding_api_base.clone(),
|
||||
@ -173,6 +201,9 @@ async fn main() -> anyhow::Result<()> {
|
||||
ads,
|
||||
arxiv,
|
||||
llm,
|
||||
medium_llm,
|
||||
fast_llm,
|
||||
vision_llm,
|
||||
embedding,
|
||||
downloader,
|
||||
harvest_status: Arc::new(tokio::sync::Mutex::new(
|
||||
@ -206,28 +237,32 @@ async fn main() -> anyhow::Result<()> {
|
||||
];
|
||||
|
||||
let cors_permissive = CorsLayer::new()
|
||||
.allow_origin(tower_http::cors::AllowOrigin::predicate(|_origin, _parts| {
|
||||
// 书签脚本的采集接口允许所有外域学术站点的跨域请求，并带上 Session 凭证
|
||||
true
|
||||
}))
|
||||
.allow_origin(tower_http::cors::AllowOrigin::predicate(
|
||||
|_origin, _parts| {
|
||||
// 书签脚本的采集接口允许所有外域学术站点的跨域请求，并带上 Session 凭证
|
||||
true
|
||||
},
|
||||
))
|
||||
.allow_methods(allowed_methods.clone())
|
||||
.allow_headers(tower_http::cors::AllowHeaders::mirror_request())
|
||||
.allow_credentials(true);
|
||||
|
||||
let cors_local = CorsLayer::new()
|
||||
.allow_origin(tower_http::cors::AllowOrigin::predicate(|origin, _parts| {
|
||||
if let Ok(origin_str) = origin.to_str() {
|
||||
// 仅允许本地开发或回环地址（允许任意端口，支持 Vite dev 端口 :5173）
|
||||
if origin_str.starts_with("http://localhost:")
|
||||
|| origin_str.starts_with("http://127.0.0.1:")
|
||||
|| origin_str == "http://localhost"
|
||||
|| origin_str == "http://127.0.0.1"
|
||||
{
|
||||
return true;
|
||||
.allow_origin(tower_http::cors::AllowOrigin::predicate(
|
||||
|origin, _parts| {
|
||||
if let Ok(origin_str) = origin.to_str() {
|
||||
// 仅允许本地开发或回环地址（允许任意端口，支持 Vite dev 端口 :5173）
|
||||
if origin_str.starts_with("http://localhost:")
|
||||
|| origin_str.starts_with("http://127.0.0.1:")
|
||||
|| origin_str == "http://localhost"
|
||||
|| origin_str == "http://127.0.0.1"
|
||||
{
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}))
|
||||
false
|
||||
},
|
||||
))
|
||||
.allow_methods(allowed_methods)
|
||||
.allow_headers(tower_http::cors::AllowHeaders::mirror_request())
|
||||
.allow_credentials(true);
|
||||
@ -279,14 +314,13 @@ async fn main() -> anyhow::Result<()> {
|
||||
"/sync/queries/:id",
|
||||
axum::routing::delete(handlers::delete_sync_query),
|
||||
)
|
||||
.route("/chat/rag", post(handlers::chat_rag))
|
||||
.route("/chat/figure", post(handlers::chat_figure))
|
||||
.route("/target/query", get(handlers::query_target))
|
||||
.route("/target/associate", post(handlers::associate_target))
|
||||
.route("/target/extract", post(handlers::extract_paper_targets))
|
||||
.route("/target/list", get(handlers::list_targets))
|
||||
// 智能体路由
|
||||
.route("/chat/agent", post(handlers::chat_agent))
|
||||
.route("/chat/modes", get(handlers::get_agent_modes))
|
||||
.route("/chat/metrics", get(handlers::get_agent_metrics))
|
||||
.route("/chat/sessions", get(handlers::list_sessions))
|
||||
.route(
|
||||
@ -326,6 +360,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
app_state.clone(),
|
||||
astroresearch::api::auth::auth_middleware,
|
||||
))
|
||||
.layer(axum::extract::DefaultBodyLimit::max(100 * 1024 * 1024))
|
||||
.layer(cors_local.clone());
|
||||
|
||||
let api_routes = Router::new()
|
||||
|
||||
@ -81,10 +81,11 @@ impl AssetBatch {
|
||||
status: Arc<Mutex<AssetBatchStatus>>,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
let semaphore = Arc::new(tokio::sync::Semaphore::new(3));
|
||||
let llm_client = crate::clients::llm::LlmClient::new(
|
||||
config.llm_api_key.clone(),
|
||||
config.llm_api_base.clone(),
|
||||
config.llm_model.clone(),
|
||||
config.llm_medium_api_key.clone(),
|
||||
config.llm_medium_api_base.clone(),
|
||||
config.llm_medium_model.clone(),
|
||||
);
|
||||
let embedding_client = crate::clients::llm::EmbeddingClient::new(
|
||||
config.embedding_api_key.clone(),
|
||||
@ -520,7 +521,9 @@ impl AssetBatch {
|
||||
};
|
||||
|
||||
if submitted_ok {
|
||||
let sem_clone = semaphore.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
let _permit = sem_clone.acquire().await.unwrap();
|
||||
match crate::services::parser::poll_and_extract_mineru(&batch_id, &bibcode_clone, &qiniu_clone, &config_clone).await {
|
||||
Ok(md) => {
|
||||
let paper_meta_res = sqlx::query("SELECT title, authors, pub, year, keywords FROM papers WHERE bibcode = ?")
|
||||
|
||||
@ -6,7 +6,7 @@ use tokio::sync::Mutex;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::api::handlers::{
|
||||
convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db,
|
||||
convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db_tx,
|
||||
};
|
||||
use crate::clients::ads::AdsClient;
|
||||
use crate::clients::arxiv::ArxivClient;
|
||||
@ -196,9 +196,12 @@ impl MetaSync {
|
||||
break;
|
||||
}
|
||||
let count = docs.len() as i32;
|
||||
for doc in docs {
|
||||
let paper = convert_ads_doc_to_standard(&doc);
|
||||
let _ = save_paper_to_db(&db, &paper).await;
|
||||
if let Ok(mut tx) = db.begin().await {
|
||||
for doc in docs {
|
||||
let paper = convert_ads_doc_to_standard(&doc);
|
||||
let _ = save_paper_to_db_tx(&mut tx, &paper).await;
|
||||
}
|
||||
let _ = tx.commit().await;
|
||||
}
|
||||
local_synced += count;
|
||||
start_offset += count;
|
||||
@ -264,9 +267,12 @@ impl MetaSync {
|
||||
break;
|
||||
}
|
||||
let count = papers.len() as i32;
|
||||
for p in papers {
|
||||
let paper = convert_arxiv_to_standard(&p);
|
||||
let _ = save_paper_to_db(&db, &paper).await;
|
||||
if let Ok(mut tx) = db.begin().await {
|
||||
for p in papers {
|
||||
let paper = convert_arxiv_to_standard(&p);
|
||||
let _ = save_paper_to_db_tx(&mut tx, &paper).await;
|
||||
}
|
||||
let _ = tx.commit().await;
|
||||
}
|
||||
local_synced += count;
|
||||
start_offset += count;
|
||||
|
||||
@ -63,27 +63,20 @@ pub async fn ingest_paper(
|
||||
);
|
||||
|
||||
// 先清除该文献的旧切片（重新解析场景）
|
||||
// 因为 vec0 虚拟表没有 FK，需要先查出旧的 rowid 再手动删除
|
||||
let old_rowids: Vec<(i64,)> =
|
||||
sqlx::query_as("SELECT rowid FROM paper_chunks_content WHERE bibcode = ?")
|
||||
.bind(bibcode)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
// 合并为两条原子 SQL 语句，避免循环查询和删除
|
||||
sqlx::query(
|
||||
"DELETE FROM vec_paper_chunks WHERE rowid IN (SELECT rowid FROM paper_chunks_content WHERE bibcode = ?)"
|
||||
)
|
||||
.bind(bibcode)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
if !old_rowids.is_empty() {
|
||||
info!("清除文献 {} 的 {} 条旧切片", bibcode, old_rowids.len());
|
||||
for (rid,) in &old_rowids {
|
||||
sqlx::query("DELETE FROM vec_paper_chunks WHERE rowid = ?")
|
||||
.bind(rid)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
}
|
||||
sqlx::query("DELETE FROM paper_chunks_content WHERE bibcode = ?")
|
||||
.bind(bibcode)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
}
|
||||
sqlx::query("DELETE FROM paper_chunks_content WHERE bibcode = ?")
|
||||
.bind(bibcode)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
let mut tx = pool.begin().await?;
|
||||
let mut ingested_count = 0;
|
||||
|
||||
for chunk in &chunks {
|
||||
@ -106,7 +99,7 @@ pub async fn ingest_paper(
|
||||
.bind(bibcode)
|
||||
.bind(chunk.paragraph_index as i64)
|
||||
.bind(&chunk.content)
|
||||
.execute(pool)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
let rowid = result.last_insert_rowid();
|
||||
@ -117,12 +110,14 @@ pub async fn ingest_paper(
|
||||
sqlx::query("INSERT INTO vec_paper_chunks (rowid, embedding) VALUES (?, ?)")
|
||||
.bind(rowid)
|
||||
.bind(&embedding_bytes)
|
||||
.execute(pool)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
ingested_count += 1;
|
||||
}
|
||||
|
||||
tx.commit().await?;
|
||||
|
||||
info!(
|
||||
"文献 {} 向量化完成，成功写入 {}/{} 块",
|
||||
bibcode,
|
||||
|
||||
@ -273,11 +273,70 @@ pub async fn extract_and_cache_targets(
|
||||
) -> Vec<TargetInfo> {
|
||||
let names = extract_targets(text);
|
||||
let mut results = Vec::new();
|
||||
let mut new_targets = Vec::new();
|
||||
|
||||
for name in names {
|
||||
match query_target_cached(pool, &name, Some(bibcode), client).await {
|
||||
Ok(info) => results.push(info),
|
||||
Err(e) => warn!("天体 {} 查询失败: {}", name, e),
|
||||
// 优先在 paper_targets 表中查找本地缓存
|
||||
let row = sqlx::query_as::<_, (String, Option<String>, Option<String>, Option<f64>, Option<String>, Option<f64>, Option<String>)>(
|
||||
"SELECT target_name, ra, dec, parallax, spectral_type, v_magnitude, aliases FROM paper_targets WHERE target_name = ? LIMIT 1"
|
||||
)
|
||||
.bind(&name)
|
||||
.fetch_optional(pool)
|
||||
.await;
|
||||
|
||||
if let Ok(Some((tname, ra, dec, parallax, spectral_type, v_magnitude, aliases_json))) = row
|
||||
{
|
||||
info!("天体 {} 命中本地缓存", name);
|
||||
let aliases: Vec<String> = aliases_json
|
||||
.and_then(|s| serde_json::from_str(&s).ok())
|
||||
.unwrap_or_default();
|
||||
results.push(TargetInfo {
|
||||
target_name: tname,
|
||||
ra,
|
||||
dec,
|
||||
parallax,
|
||||
spectral_type,
|
||||
v_magnitude,
|
||||
aliases,
|
||||
});
|
||||
} else {
|
||||
// 本地缓存未命中，释放连接并外部查询 Sesame
|
||||
info!("天体 {} 本地缓存未命中，查询 CDS Sesame...", name);
|
||||
// 速率节流：50ms 间隔
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
match query_sesame(&name, client).await {
|
||||
Ok(info) => {
|
||||
results.push(info.clone());
|
||||
new_targets.push(info);
|
||||
}
|
||||
Err(e) => warn!("天体 {} 查询失败: {}", name, e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 最终在事务中批量写入新天体缓存，防范高频写入锁表
|
||||
if !new_targets.is_empty() {
|
||||
if let Ok(mut tx) = pool.begin().await {
|
||||
for info in new_targets {
|
||||
let aliases_json =
|
||||
serde_json::to_string(&info.aliases).unwrap_or_else(|_| "[]".to_string());
|
||||
if let Err(e) = sqlx::query(
|
||||
"INSERT OR IGNORE INTO paper_targets (bibcode, target_name, ra, dec, parallax, spectral_type, v_magnitude, aliases) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
.bind(bibcode)
|
||||
.bind(&info.target_name)
|
||||
.bind(&info.ra)
|
||||
.bind(&info.dec)
|
||||
.bind(info.parallax)
|
||||
.bind(&info.spectral_type)
|
||||
.bind(info.v_magnitude)
|
||||
.bind(&aliases_json)
|
||||
.execute(&mut *tx)
|
||||
.await {
|
||||
error!("天体信息写入缓存失败: {}", e);
|
||||
}
|
||||
}
|
||||
let _ = tx.commit().await;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -10,6 +10,8 @@ use tracing::{info, warn};
|
||||
pub struct Dictionary {
|
||||
// 英文名词(全小写) -> 中文标准译名
|
||||
terms: HashMap<String, String>,
|
||||
// 词表首个英文单词(全小写)集合，用于在扫描文本时快速过滤，避免盲目 join/lowercase 分配内存
|
||||
first_words: HashSet<String>,
|
||||
}
|
||||
|
||||
impl Default for Dictionary {
|
||||
@ -22,9 +24,18 @@ impl Dictionary {
|
||||
pub fn new() -> Self {
|
||||
Dictionary {
|
||||
terms: HashMap::new(),
|
||||
first_words: HashSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
// 插入对照术语
|
||||
pub fn insert(&mut self, english: String, chinese: String) {
|
||||
if let Some(first_w) = english.split_whitespace().next() {
|
||||
self.first_words.insert(first_w.to_string());
|
||||
}
|
||||
self.terms.insert(english, chinese);
|
||||
}
|
||||
|
||||
// 从本地物理文本加载词表数据
|
||||
pub fn load_from_file<P: AsRef<Path>>(&mut self, path: P) -> anyhow::Result<()> {
|
||||
let path_ref = path.as_ref();
|
||||
@ -46,12 +57,13 @@ impl Dictionary {
|
||||
let chinese_raw = parts[1].trim();
|
||||
let chinese = clean_chinese_translation(chinese_raw);
|
||||
if !english.is_empty() && !chinese.is_empty() {
|
||||
self.terms.insert(english, chinese);
|
||||
self.insert(english, chinese);
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
self.terms.shrink_to_fit();
|
||||
self.first_words.shrink_to_fit();
|
||||
info!("天文词典加载成功，总计导入 {} 条专业术语对照", count);
|
||||
Ok(())
|
||||
}
|
||||
@ -77,16 +89,22 @@ impl Dictionary {
|
||||
let words: Vec<&str> = clean_text.split_whitespace().collect();
|
||||
let mut matched = HashSet::new();
|
||||
let mut results = Vec::new();
|
||||
let mut matched_indices = HashSet::new();
|
||||
|
||||
// 天文学词条跨度最大限制（一般多词短语不超过 6 个英文单词）
|
||||
let max_span = 6;
|
||||
let n = words.len();
|
||||
let mut next_valid_index = 0;
|
||||
|
||||
for i in 0..n {
|
||||
if matched_indices.contains(&i) {
|
||||
if i < next_valid_index {
|
||||
continue;
|
||||
}
|
||||
// 快速过滤：如果当前单词的 lowercase 不在 first_words 集合中，不可能匹配任何词条
|
||||
let first_word_lower = words[i].to_lowercase();
|
||||
if !self.first_words.contains(&first_word_lower) {
|
||||
continue;
|
||||
}
|
||||
|
||||
for len in (1..=max_span).rev() {
|
||||
if i + len <= n {
|
||||
let phrase_slice = &words[i..i + len];
|
||||
@ -96,14 +114,12 @@ impl Dictionary {
|
||||
// 避免重复匹配更长名词的子词 (如已匹配 'active galactic nucleus' 就不重复提取其中的 'nucleus')
|
||||
if !matched.contains(&phrase) {
|
||||
let chinese = self.terms.get(&phrase).unwrap().clone();
|
||||
let original_phrase = &words[i..i + len].join(" ");
|
||||
let original_phrase = phrase_slice.join(" ");
|
||||
|
||||
results.push((original_phrase.clone(), chinese.clone()));
|
||||
matched.insert(phrase.clone());
|
||||
}
|
||||
for k in i..i + len {
|
||||
matched_indices.insert(k);
|
||||
results.push((original_phrase, chinese));
|
||||
matched.insert(phrase);
|
||||
}
|
||||
next_valid_index = i + len;
|
||||
break; // 匹配到当前位置的最长词，标记后跳出对当前i的其他长度匹配尝试
|
||||
}
|
||||
}
|
||||
@ -233,15 +249,13 @@ mod tests {
|
||||
fn test_dictionary_match() {
|
||||
let mut dict = Dictionary::new();
|
||||
// 模拟词典数据
|
||||
dict.terms.insert(
|
||||
dict.insert(
|
||||
"active galactic nucleus".to_string(),
|
||||
"活动星系核".to_string(),
|
||||
);
|
||||
dict.terms
|
||||
.insert("galactic nucleus".to_string(), "星系核".to_string());
|
||||
dict.terms.insert("nucleus".to_string(), "核心".to_string());
|
||||
dict.terms
|
||||
.insert("black hole".to_string(), "黑洞".to_string());
|
||||
dict.insert("galactic nucleus".to_string(), "星系核".to_string());
|
||||
dict.insert("nucleus".to_string(), "核心".to_string());
|
||||
dict.insert("black hole".to_string(), "黑洞".to_string());
|
||||
|
||||
let text = "We study the active galactic nucleus and its central black hole.";
|
||||
let matched = dict.match_text(text);
|
||||
|
||||
Loading…
Reference in New Issue
Block a user