feat: ReAct Agent 自主科研引擎、LLM 流式工具调用、论文搜索服务化
核心新增 —— ReAct 智能体引擎 (src/agent/): - ReAct 主循环：自主推理→工具调用→观察→迭代，支持流式 SSE 推送、 上下文自动压缩、重复调用死循环检测、手动取消与超时控制 - 7 个内置工具：ADS/arXiv 论文搜索、元数据查询、全文读取、PDF/HTML 下载解析、RAG 向量语义检索、CDS Sesame 天体目标查询 LLM 客户端增强 (src/clients/llm.rs): - 新增流式/非流式 Function Calling (tool calling) 支持 - 原生 reasoning_content 推理链字段 (DeepSeek/QwQ) - 多模态图片+文本对话、TokenUsage 用量统计 搜索服务重构 (src/services/search.rs): - 抽取 ADS/arXiv 搜索为统一 service，跨平台去重、本地状态合并 - API handler 瘦身至薄封装层 下载器增强 (src/services/download.rs): - Obscura 无头浏览器反爬下载通道 (进程内/CLI 双模式) - arXiv/Springer 专用策略、ADS Scan PDF 兜底、CAPTCHA 绕过 解析器增强 (src/services/parser/mod.rs): - Markdown 输出自动附加 YAML 前端元数据头 (标题/作者/来源) 新增 API (src/api/agent.rs): - POST /api/chat/agent — SSE 流式智能体对话 - GET/DELETE /api/chat/sessions — 会话管理 CRUD - POST /api/chat/sessions/:id/stop — 紧急停止 数据库: agent_sessions + agent_messages 表 前端 (dashboard/): - 全新"智能科研"标签页 (ResearchAgentPanel) — 会话列表、ReAct 步骤时间线可视化、SSE 实时流式渲染、Markdown+KaTeX 答案展示 - Reader AI 助手集成 Agent 工作流，含工具调用步骤展示与引用溯源 - 侧边栏新增"智能科研"导航入口
This commit is contained in:
parent
22e7e1dcee
commit
b1fb884f21
24
Cargo.lock
generated
24
Cargo.lock
generated
@ -132,6 +132,8 @@ name = "astroresearch"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-stream",
|
||||
"async-trait",
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"chrono",
|
||||
@ -177,6 +179,28 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
|
||||
dependencies = [
|
||||
"async-stream-impl",
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream-impl"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.89"
|
||||
|
||||
@ -46,11 +46,13 @@ flate2 = "1.1.9"
|
||||
zip = "8.6.0"
|
||||
uuid = { version = "1.23.2", features = ["v4"] }
|
||||
tracing-appender = "0.2.5"
|
||||
obscura-browser = { path = "libs/obscura/crates/obscura-browser", optional = true }
|
||||
obscura-net = { path = "libs/obscura/crates/obscura-net", optional = true }
|
||||
obscura-browser = { path = "/home/fmq/program/AstroResearch/libs/obscura/crates/obscura-browser", optional = true }
|
||||
obscura-net = { path = "/home/fmq/program/AstroResearch/libs/obscura/crates/obscura-net", optional = true }
|
||||
libsqlite3-sys = { version = "0.27.0", features = ["bundled"] }
|
||||
sqlite-vec = "0.1.9"
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
async-trait = "0.1"
|
||||
async-stream = "0.3"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
|
||||
@ -8,10 +8,11 @@ import { LibraryPanel } from './features/library/LibraryPanel';
|
||||
import { ReaderPanel } from './features/reader/ReaderPanel';
|
||||
import { CitationPanel } from './features/citation/CitationPanel';
|
||||
import { SyncPanel } from './features/sync/SyncPanel';
|
||||
import { ResearchAgentPanel } from './features/agent/ResearchAgentPanel';
|
||||
import type { StandardPaper, CitationNetwork, NoteRecord } from './types';
|
||||
|
||||
export default function App() {
|
||||
const [activeTab, setActiveTab] = useState<'search' | 'library' | 'reader' | 'citation' | 'sync'>(() => {
|
||||
const [activeTab, setActiveTab] = useState<'search' | 'library' | 'reader' | 'citation' | 'sync' | 'agent'>(() => {
|
||||
const saved = localStorage.getItem('astro_active_tab');
|
||||
return (saved as any) || 'search';
|
||||
});
|
||||
@ -544,7 +545,7 @@ export default function App() {
|
||||
{/* 选项卡容器 */}
|
||||
<div className="flex-1 overflow-y-auto p-4 sm:p-6 md:p-8 relative z-10 w-full flex flex-col">
|
||||
<div className={`w-full flex-1 flex flex-col min-h-0 ${
|
||||
(activeTab === 'reader' || activeTab === 'citation') ? 'max-w-none' : 'max-w-7xl mx-auto'
|
||||
(activeTab === 'reader' || activeTab === 'citation' || activeTab === 'agent') ? 'max-w-none' : 'max-w-7xl mx-auto'
|
||||
}`}>
|
||||
{activeTab === 'search' && (
|
||||
<SearchPanel
|
||||
@ -675,6 +676,13 @@ export default function App() {
|
||||
{activeTab === 'sync' && (
|
||||
<SyncPanel />
|
||||
)}
|
||||
|
||||
{activeTab === 'agent' && (
|
||||
<ResearchAgentPanel
|
||||
showConfirm={showConfirm}
|
||||
showAlert={showAlert}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</main>
|
||||
|
||||
@ -1,11 +1,11 @@
|
||||
// dashboard/src/components/layout/Sidebar.tsx
|
||||
import { useState } from 'react';
|
||||
import { Search, BookOpen, GitFork, Library, RefreshCw, ChevronLeft } from 'lucide-react';
|
||||
import { Search, BookOpen, GitFork, Library, RefreshCw, ChevronLeft, Sparkles } from 'lucide-react';
|
||||
import type { StandardPaper } from '../../types';
|
||||
|
||||
interface SidebarProps {
|
||||
activeTab: 'search' | 'library' | 'reader' | 'citation' | 'sync';
|
||||
setActiveTab: (tab: 'search' | 'library' | 'reader' | 'citation' | 'sync') => void;
|
||||
activeTab: 'search' | 'library' | 'reader' | 'citation' | 'sync' | 'agent';
|
||||
setActiveTab: (tab: 'search' | 'library' | 'reader' | 'citation' | 'sync' | 'agent') => void;
|
||||
selectedPaper: StandardPaper | null;
|
||||
loadCitations: (bibcode: string) => void;
|
||||
}
|
||||
@ -95,6 +95,7 @@ export function Sidebar({ activeTab, setActiveTab, selectedPaper, loadCitations
|
||||
{ id: 'library', label: '馆藏管理', icon: Library },
|
||||
{ id: 'sync', label: '批量任务', icon: RefreshCw },
|
||||
{ id: 'reader', label: '双语阅读', icon: BookOpen },
|
||||
{ id: 'agent', label: '智能科研', icon: Sparkles },
|
||||
{ id: 'citation', label: '引用星系', icon: GitFork },
|
||||
].map((tab: { id: string; label: string; icon: any; disabled?: boolean }) => {
|
||||
const Icon = tab.icon;
|
||||
|
||||
1066
dashboard/src/features/agent/ResearchAgentPanel.tsx
Normal file
1066
dashboard/src/features/agent/ResearchAgentPanel.tsx
Normal file
File diff suppressed because it is too large
Load Diff
@ -7,7 +7,10 @@ import rehypeRaw from 'rehype-raw';
|
||||
import rehypeKatex from 'rehype-katex';
|
||||
import rehypeSanitize, { defaultSchema } from 'rehype-sanitize';
|
||||
import 'katex/dist/katex.min.css';
|
||||
import { Send, Loader, Sparkles, X, BookOpen, AlertCircle, Compass } from 'lucide-react';
|
||||
import {
|
||||
Send, Loader, Sparkles, X, BookOpen, AlertCircle, Compass,
|
||||
Brain, Settings, ChevronDown, ChevronUp, AlertTriangle, Square
|
||||
} from 'lucide-react';
|
||||
import axios from 'axios';
|
||||
|
||||
interface RetrievalSource {
|
||||
@ -17,11 +20,21 @@ interface RetrievalSource {
|
||||
distance: number;
|
||||
}
|
||||
|
||||
interface AgentStep {
|
||||
step: number;
|
||||
type: 'thought' | 'tool_call' | 'tool_result' | 'error';
|
||||
label: string;
|
||||
detail?: string;
|
||||
isError?: boolean;
|
||||
isFinished?: boolean;
|
||||
}
|
||||
|
||||
interface Message {
|
||||
sender: 'user' | 'ai';
|
||||
text: string;
|
||||
sources?: RetrievalSource[];
|
||||
imageUrl?: string;
|
||||
steps?: AgentStep[];
|
||||
}
|
||||
|
||||
interface AIAssistantPanelProps {
|
||||
@ -50,6 +63,19 @@ const SUGGESTED_QUESTIONS = [
|
||||
"文献库中关于双星合并前奏（Precursor）观测特征 of 论述有哪些？"
|
||||
];
|
||||
|
||||
function getToolDisplayName(name: string): string {
|
||||
switch (name) {
|
||||
case 'search_papers': return '检索 ADS/arXiv';
|
||||
case 'get_paper_metadata': return '获取文献详细元数据';
|
||||
case 'download_paper': return '下载文献全文资源';
|
||||
case 'parse_paper': return '结构化解析文献内容';
|
||||
case 'get_paper_content': return '获取文献全文内容';
|
||||
case 'rag_search': return '语义库检索 (RAG)';
|
||||
case 'query_target': return '查询天体物理参数';
|
||||
default: return name;
|
||||
}
|
||||
}
|
||||
|
||||
export function AIAssistantPanel({
|
||||
bibcode,
|
||||
onClose,
|
||||
@ -61,16 +87,50 @@ export function AIAssistantPanel({
|
||||
const [input, setInput] = useState('');
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [sessionId, setSessionId] = useState<string | null>(null);
|
||||
|
||||
// 步骤展开状态
|
||||
const [stepsExpanded, setStepsExpanded] = useState<Record<number, boolean>>({});
|
||||
|
||||
const chatEndRef = useRef<HTMLDivElement>(null);
|
||||
const scrollContainerRef = useRef<HTMLDivElement>(null);
|
||||
const [shouldAutoScroll, setShouldAutoScroll] = useState(true);
|
||||
|
||||
const scrollToBottom = () => {
|
||||
chatEndRef.current?.scrollIntoView({ behavior: 'smooth' });
|
||||
};
|
||||
|
||||
useEffect(() => {
|
||||
if (shouldAutoScroll) {
|
||||
scrollToBottom();
|
||||
}, [messages, loading]);
|
||||
}
|
||||
}, [messages, loading, shouldAutoScroll]);
|
||||
|
||||
const handleScroll = () => {
|
||||
if (!scrollContainerRef.current) return;
|
||||
const { scrollTop, scrollHeight, clientHeight } = scrollContainerRef.current;
|
||||
// 距离底部 50 像素以内视为贴紧底部
|
||||
const isAtBottom = scrollHeight - scrollTop - clientHeight < 50;
|
||||
setShouldAutoScroll(isAtBottom);
|
||||
};
|
||||
|
||||
// 重置会话（当 bibcode 切换时）
|
||||
useEffect(() => {
|
||||
setMessages([]);
|
||||
setSessionId(null);
|
||||
setError(null);
|
||||
setStepsExpanded({});
|
||||
}, [bibcode]);
|
||||
|
||||
// 手动停止智能体执行
|
||||
const handleStop = async () => {
|
||||
if (!sessionId) return;
|
||||
try {
|
||||
await axios.post(`/api/chat/sessions/${sessionId}/stop`);
|
||||
} catch (e) {
|
||||
console.error('停止智能体执行失败:', e);
|
||||
}
|
||||
};
|
||||
|
||||
const handleSend = async (questionText: string) => {
|
||||
if (!questionText.trim() || loading) return;
|
||||
@ -84,6 +144,7 @@ export function AIAssistantPanel({
|
||||
setMessages(prev => [...prev, userMsg]);
|
||||
setInput('');
|
||||
setLoading(true);
|
||||
setShouldAutoScroll(true);
|
||||
|
||||
const isFigureQuery = !!pendingFigure;
|
||||
const currentFigurePath = pendingFigure?.path;
|
||||
@ -92,61 +153,258 @@ export function AIAssistantPanel({
|
||||
onClearPendingFigure();
|
||||
}
|
||||
|
||||
// AI 初始消息占位符
|
||||
const aiMsgPlaceholder: Message = {
|
||||
sender: 'ai',
|
||||
text: '',
|
||||
steps: [],
|
||||
sources: []
|
||||
};
|
||||
setMessages(prev => [...prev, aiMsgPlaceholder]);
|
||||
|
||||
try {
|
||||
if (isFigureQuery && currentFigurePath) {
|
||||
// 图表多模态分析依然调用特定接口
|
||||
const res = await axios.post<{ answer: string }>('/api/chat/figure', {
|
||||
bibcode,
|
||||
image_path: currentFigurePath,
|
||||
question: questionText
|
||||
});
|
||||
const aiMsg: Message = {
|
||||
sender: 'ai',
|
||||
text: res.data.answer
|
||||
};
|
||||
setMessages(prev => [...prev, aiMsg]);
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
const last = next[next.length - 1];
|
||||
if (last && last.sender === 'ai') {
|
||||
last.text = res.data.answer;
|
||||
}
|
||||
return next;
|
||||
});
|
||||
setLoading(false);
|
||||
} else {
|
||||
// 调用智能体流式对话接口（带文献上下文提示，确保 Agent 优先分析当前阅读的文献）
|
||||
const contextPrefix = `[当前正在阅读文献: ${bibcode}] `;
|
||||
const requestQuery = questionText.includes(bibcode) ? questionText : `${contextPrefix}${questionText}`;
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/chat/agent', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify({
|
||||
question: requestQuery,
|
||||
session_id: sessionId,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP error! status: ${response.status}`);
|
||||
}
|
||||
|
||||
const reader = response.body?.getReader();
|
||||
if (!reader) {
|
||||
throw new Error('ReadableStream not supported');
|
||||
}
|
||||
|
||||
const decoder = new TextDecoder('utf-8');
|
||||
let buffer = '';
|
||||
|
||||
while (true) {
|
||||
const { value, done } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split('\n');
|
||||
buffer = lines.pop() || '';
|
||||
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
if (!trimmed) continue;
|
||||
|
||||
if (trimmed.startsWith('data:')) {
|
||||
const dataStr = trimmed.slice(5).trim();
|
||||
if (!dataStr) continue;
|
||||
|
||||
try {
|
||||
const event = JSON.parse(dataStr);
|
||||
|
||||
switch (event.type) {
|
||||
case 'session':
|
||||
if (!sessionId) {
|
||||
setSessionId(event.session_id);
|
||||
}
|
||||
break;
|
||||
case 'thought':
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
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 === 'thought');
|
||||
if (!stepObj) {
|
||||
stepObj = { step: event.step, type: 'thought', label: '思考推理中' };
|
||||
last.steps.push(stepObj);
|
||||
}
|
||||
stepObj.detail = event.content;
|
||||
}
|
||||
return next;
|
||||
});
|
||||
break;
|
||||
case 'tool_call':
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
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');
|
||||
if (!stepObj) {
|
||||
stepObj = { step: event.step, type: 'tool_call', label: `调用: ${getToolDisplayName(event.name)}`, isFinished: false };
|
||||
last.steps.push(stepObj);
|
||||
}
|
||||
stepObj.detail = typeof event.arguments === 'string' ? event.arguments : JSON.stringify(event.arguments);
|
||||
}
|
||||
return next;
|
||||
});
|
||||
break;
|
||||
case 'tool_result':
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
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');
|
||||
if (stepObj) {
|
||||
stepObj.label = `调用: ${getToolDisplayName(event.name)} (已获取观测数据)`;
|
||||
stepObj.isError = event.is_error;
|
||||
stepObj.isFinished = true;
|
||||
}
|
||||
|
||||
// 提取 RAG 检索到的参考来源文献，并在 sources 中进行去重和追加
|
||||
if (event.name === 'rag_search' && event.metadata && event.metadata.sources) {
|
||||
const newSources: RetrievalSource[] = event.metadata.sources.map((s: any) => ({
|
||||
bibcode: s.bibcode,
|
||||
paragraph_index: s.paragraph_index,
|
||||
content: s.preview || '',
|
||||
distance: s.distance,
|
||||
}));
|
||||
|
||||
const existing = last.sources || [];
|
||||
const combined = [...existing];
|
||||
|
||||
for (const src of newSources) {
|
||||
if (!combined.some(c => c.bibcode === src.bibcode && c.paragraph_index === src.paragraph_index)) {
|
||||
combined.push(src);
|
||||
}
|
||||
}
|
||||
last.sources = combined;
|
||||
}
|
||||
}
|
||||
return next;
|
||||
});
|
||||
break;
|
||||
case 'text_delta':
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
const last = next[next.length - 1];
|
||||
if (last && last.sender === 'ai') {
|
||||
last.text += event.content;
|
||||
}
|
||||
return next;
|
||||
});
|
||||
break;
|
||||
case 'error':
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
const last = next[next.length - 1];
|
||||
if (last && last.sender === 'ai') {
|
||||
if (!last.steps) last.steps = [];
|
||||
last.steps.push({
|
||||
step: 99,
|
||||
type: 'error',
|
||||
label: `错误返回`,
|
||||
detail: event.message
|
||||
});
|
||||
}
|
||||
return next;
|
||||
});
|
||||
break;
|
||||
}
|
||||
} catch (e) {
|
||||
console.error('解析流式 SSE 失败:', e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
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
|
||||
});
|
||||
|
||||
const aiMsg: Message = {
|
||||
sender: 'ai',
|
||||
text: res.data.answer,
|
||||
sources: res.data.sources
|
||||
};
|
||||
setMessages(prev => [...prev, aiMsg]);
|
||||
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;
|
||||
}
|
||||
return next;
|
||||
});
|
||||
setLoading(false);
|
||||
}
|
||||
}
|
||||
} catch (err: any) {
|
||||
console.error('问答请求失败:', err);
|
||||
setError(err.response?.data?.error || err.message || '网络请求错误，请稍后重试');
|
||||
} finally {
|
||||
// 如果报错，清空多余的消息占位符
|
||||
setMessages(prev => {
|
||||
const next = [...prev];
|
||||
if (next.length > 0 && next[next.length - 1].sender === 'ai' && !next[next.length - 1].text) {
|
||||
next.pop();
|
||||
}
|
||||
return next;
|
||||
});
|
||||
setLoading(false);
|
||||
}
|
||||
};
|
||||
|
||||
const toggleSteps = (msgIdx: number) => {
|
||||
setStepsExpanded(prev => ({
|
||||
...prev,
|
||||
[msgIdx]: !prev[msgIdx]
|
||||
}));
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="console-panel rounded-xl border border-slate-200 bg-slate-50 flex flex-col overflow-hidden relative shadow-sm h-full">
|
||||
{/* 头部面板 */}
|
||||
<div className="px-4 py-3.5 border-b border-slate-200 flex items-center justify-between bg-white shrink-0">
|
||||
<div className="flex items-center gap-1.5">
|
||||
<div className="flex items-center gap-1.5 select-none">
|
||||
<Sparkles className="w-4 h-4 text-sky-600 animate-pulse" />
|
||||
<span className="text-xs font-bold text-slate-800">Astro RAG 学术助手</span>
|
||||
<span className="text-xs font-bold text-slate-800">科研智能问答助手</span>
|
||||
</div>
|
||||
<button onClick={onClose} className="text-slate-400 hover:text-slate-600 transition-colors">
|
||||
<button onClick={onClose} className="text-slate-400 hover:text-slate-600 transition-colors cursor-pointer">
|
||||
<X className="w-4 h-4" />
|
||||
</button>
|
||||
</div>
|
||||
|
||||
{/* 对话消息区 */}
|
||||
<div className="flex-1 overflow-y-auto p-4 space-y-4 min-h-0">
|
||||
<div
|
||||
ref={scrollContainerRef}
|
||||
onScroll={handleScroll}
|
||||
className="flex-1 overflow-y-auto p-4 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>
|
||||
<p className="text-[11px] text-slate-500 leading-relaxed font-semibold">
|
||||
基于向量检索（SQLite-Vec）与学术大语言模型，跨越所有已解析的 Markdown 文献为您提供深度知识整合与溯源解答。
|
||||
基于 ReAct 框架的智能科研体。支持自动读取已解析的 Markdown 文献，在向量库进行 RAG 查询以及天体参数检索。
|
||||
</p>
|
||||
</div>
|
||||
|
||||
@ -158,7 +416,7 @@ export function AIAssistantPanel({
|
||||
<button
|
||||
key={idx}
|
||||
onClick={() => handleSend(q)}
|
||||
className="w-full text-left bg-white hover:bg-slate-100 border border-slate-200 rounded-lg p-3 text-xs text-slate-700 leading-relaxed font-medium transition-colors shadow-sm cursor-pointer"
|
||||
className="w-full text-left bg-white hover:bg-slate-100 border border-slate-200 rounded-lg p-3 text-xs text-slate-700 leading-relaxed font-medium transition-all shadow-sm cursor-pointer hover:border-sky-300"
|
||||
>
|
||||
{q}
|
||||
</button>
|
||||
@ -168,15 +426,15 @@ export function AIAssistantPanel({
|
||||
</div>
|
||||
) : (
|
||||
messages.map((msg, index) => (
|
||||
<div key={index} className={`flex flex-col ${msg.sender === 'user' ? 'items-end' : 'items-start'} space-y-1`}>
|
||||
<div key={index} className={`flex flex-col ${msg.sender === 'user' ? 'items-end' : 'items-start'} space-y-1.5`}>
|
||||
<span className="text-[10px] font-bold text-slate-400 px-1">
|
||||
{msg.sender === 'user' ? '我' : 'Astro AI'}
|
||||
{msg.sender === 'user' ? '我' : '智能科研助手'}
|
||||
</span>
|
||||
<div
|
||||
className={`max-w-[90%] rounded-xl px-4 py-3 text-xs leading-relaxed font-medium shadow-sm border ${
|
||||
className={`max-w-[95%] rounded-xl px-4 py-3 text-xs leading-relaxed font-medium shadow-sm border ${
|
||||
msg.sender === 'user'
|
||||
? 'bg-sky-600 text-white border-sky-600'
|
||||
: 'bg-white text-slate-800 border-slate-200'
|
||||
? 'bg-sky-600 text-white border-sky-600 select-text'
|
||||
: 'bg-white text-slate-800 border-slate-200 select-text'
|
||||
}`}
|
||||
>
|
||||
{msg.sender === 'user' ? (
|
||||
@ -187,6 +445,49 @@ export function AIAssistantPanel({
|
||||
<p className="whitespace-pre-wrap">{msg.text}</p>
|
||||
</div>
|
||||
) : (
|
||||
<div className="space-y-2.5">
|
||||
{/* Collapsible Steps list */}
|
||||
{msg.steps && msg.steps.length > 0 && (
|
||||
<div className="border border-slate-100 bg-slate-50/50 rounded-lg p-2 select-none">
|
||||
<button
|
||||
onClick={() => toggleSteps(index)}
|
||||
className="flex items-center justify-between w-full text-left text-[10px] font-bold text-slate-500 hover:text-slate-700 transition-colors cursor-pointer"
|
||||
>
|
||||
<span className="flex items-center gap-1">
|
||||
<Brain className="w-3.5 h-3.5 text-purple-500" />
|
||||
<span>推理与工具链展开 ({msg.steps.length} 步)</span>
|
||||
</span>
|
||||
{stepsExpanded[index] ? <ChevronUp className="w-3.5 h-3.5" /> : <ChevronDown className="w-3.5 h-3.5" />}
|
||||
</button>
|
||||
|
||||
{stepsExpanded[index] && (
|
||||
<div className="mt-2 pl-2.5 border-l border-slate-200 space-y-1.5 transition-all">
|
||||
{msg.steps.map((step, sIdx) => (
|
||||
<div key={sIdx} className="text-[9px] text-slate-655 flex flex-col gap-0.5">
|
||||
<span className="font-bold flex items-center gap-1">
|
||||
{step.type === 'thought' ? (
|
||||
<Brain className="w-3 h-3 text-purple-500 shrink-0" />
|
||||
) : step.type === 'error' ? (
|
||||
<AlertTriangle className="w-3 h-3 text-rose-500 shrink-0" />
|
||||
) : (
|
||||
<Settings className={`w-3 h-3 text-sky-600 shrink-0 ${step.isFinished ? '' : 'animate-spin'}`} />
|
||||
)}
|
||||
<span>{step.label}</span>
|
||||
</span>
|
||||
{step.detail && (
|
||||
<span className="font-mono text-[9px] opacity-75 whitespace-pre-wrap block bg-white px-1.5 py-0.5 rounded border border-slate-100 max-h-24 overflow-y-auto">
|
||||
{step.detail}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Final output text */}
|
||||
{msg.text ? (
|
||||
<div className="prose prose-sm max-w-none text-slate-800 leading-relaxed prose-headings:text-slate-900 prose-headings:font-bold prose-strong:text-slate-900 prose-code:text-sky-700 prose-img:rounded-lg">
|
||||
<ReactMarkdown
|
||||
remarkPlugins={[remarkMath, remarkGfm]}
|
||||
@ -195,19 +496,23 @@ export function AIAssistantPanel({
|
||||
{msg.text}
|
||||
</ReactMarkdown>
|
||||
</div>
|
||||
) : (
|
||||
!loading && <span className="text-slate-400 italic">正在构建最终结论...</span>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* 关联来源展示 */}
|
||||
{msg.sender === 'ai' && msg.sources && msg.sources.length > 0 && (
|
||||
<div className="w-full mt-2 pl-2 space-y-1.5 border-l-2 border-slate-200">
|
||||
<span className="text-[9px] font-bold text-slate-400">参考来源文献（点击跳转）:</span>
|
||||
<div className="grid grid-cols-1 gap-1.5 w-[90%]">
|
||||
<span className="text-[9px] font-bold text-slate-400 select-none">参考来源文献（点击跳转）:</span>
|
||||
<div className="grid grid-cols-1 gap-1.5 w-[95%]">
|
||||
{msg.sources.map((src, sIdx) => (
|
||||
<button
|
||||
key={sIdx}
|
||||
onClick={() => onJumpToSource(src.bibcode, src.paragraph_index)}
|
||||
className="flex items-center justify-between text-left bg-white hover:bg-slate-100 border border-slate-200 rounded-md px-2.5 py-1.5 text-[10px] text-slate-655 font-semibold transition-all shadow-sm cursor-pointer"
|
||||
className="flex items-center justify-between text-left bg-white hover:bg-slate-100 border border-slate-200 rounded-md px-2.5 py-1.5 text-[10px] text-slate-655 font-semibold transition-all shadow-sm cursor-pointer hover:border-sky-300"
|
||||
title={src.content}
|
||||
>
|
||||
<div className="flex items-center gap-1.5 truncate mr-2">
|
||||
@ -227,15 +532,15 @@ export function AIAssistantPanel({
|
||||
))
|
||||
)}
|
||||
|
||||
{loading && (
|
||||
{loading && messages.length > 0 && !messages[messages.length - 1].text && (
|
||||
<div className="flex items-center space-x-2 text-slate-500 pl-2">
|
||||
<Loader className="w-4 h-4 animate-spin text-sky-600" />
|
||||
<span className="text-[10px] font-bold">文献检索及综述生成中...</span>
|
||||
<span className="text-[10px] font-bold">文献检索及推理结论构建中...</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{error && (
|
||||
<div className="flex items-start gap-2 bg-red-50 border border-red-200 text-red-700 rounded-lg p-3 text-xs w-[90%] font-semibold">
|
||||
<div className="flex items-start gap-2 bg-red-50 border border-red-200 text-red-700 rounded-lg p-3 text-xs w-[95%] font-semibold">
|
||||
<AlertCircle className="w-4 h-4 text-red-500 shrink-0 mt-0.5" />
|
||||
<div>
|
||||
<div className="font-bold">查询发生错误</div>
|
||||
@ -249,7 +554,7 @@ export function AIAssistantPanel({
|
||||
|
||||
{/* 待提问图片预览 */}
|
||||
{pendingFigure && (
|
||||
<div className="px-3 py-2 bg-slate-100 border-t border-slate-200 flex items-center justify-between shrink-0">
|
||||
<div className="px-3 py-2 bg-slate-100 border-t border-slate-200 flex items-center justify-between shrink-0 select-none">
|
||||
<div className="flex items-center gap-2">
|
||||
<img
|
||||
src={pendingFigure.url}
|
||||
@ -288,13 +593,24 @@ export function AIAssistantPanel({
|
||||
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"
|
||||
/>
|
||||
{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() || loading}
|
||||
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"
|
||||
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>
|
||||
</form>
|
||||
</div>
|
||||
|
||||
36
migrations/20260615000000_agent_sessions.sql
Normal file
36
migrations/20260615000000_agent_sessions.sql
Normal file
@ -0,0 +1,36 @@
|
||||
-- 科研智能体会话与消息持久化表
|
||||
-- 支持 ReAct 框架的多轮对话、工具调用与思考链存储
|
||||
|
||||
CREATE TABLE IF NOT EXISTS agent_sessions (
|
||||
session_id TEXT PRIMARY KEY,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
model TEXT NOT NULL DEFAULT '',
|
||||
turn_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_error TEXT,
|
||||
summary TEXT,
|
||||
metadata TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
deleted_at DATETIME
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS agent_messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL,
|
||||
turn_index INTEGER NOT NULL DEFAULT 0,
|
||||
step_index INTEGER NOT NULL DEFAULT 0,
|
||||
role TEXT NOT NULL CHECK(role IN ('system', 'user', 'assistant', 'tool')),
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
thought TEXT,
|
||||
tool_calls TEXT,
|
||||
tool_call_id TEXT,
|
||||
token_count INTEGER NOT NULL DEFAULT 0,
|
||||
metadata TEXT,
|
||||
raw_json TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_messages_session ON agent_messages(session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_messages_turn ON agent_messages(session_id, turn_index);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_sessions_updated ON agent_sessions(updated_at) WHERE deleted_at IS NULL;
|
||||
6
src/agent/mod.rs
Normal file
6
src/agent/mod.rs
Normal file
@ -0,0 +1,6 @@
|
||||
// src/agent/mod.rs
|
||||
// 科研智能体模块
|
||||
// 基于 ReAct 框架实现 Thought -> Action -> Observation 循环
|
||||
|
||||
pub mod tools;
|
||||
pub mod runtime;
|
||||
786
src/agent/runtime.rs
Normal file
786
src/agent/runtime.rs
Normal file
@ -0,0 +1,786 @@
|
||||
// src/agent/runtime.rs
|
||||
//
|
||||
// 科研智能体运行时核心模块。
|
||||
// 实现 ReAct 循环：Thought -> Action (工具调用) -> Observation -> Thought...
|
||||
// 支持会话持久化、上下文压缩、死循环检测和 SSE 流式输出。
|
||||
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, warn, error};
|
||||
use serde::Serialize;
|
||||
use sqlx::SqlitePool;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::{
|
||||
ChatMessage, LlmClient, MessageRole, StreamEvent,
|
||||
};
|
||||
use super::tools::{ToolContext, ToolOutput, ToolRegistry};
|
||||
|
||||
/// Agent 配置参数
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AgentConfig {
|
||||
/// 最大 ReAct 迭代次数
|
||||
pub max_steps: usize,
|
||||
/// 同质调用检测阈值（连续相同调用次数）
|
||||
pub duplicate_call_threshold: usize,
|
||||
/// 工具执行超时时间（秒）
|
||||
pub tool_timeout_secs: u64,
|
||||
/// 工具输出最大字符数
|
||||
pub max_tool_output_chars: usize,
|
||||
/// 上下文 Token 估算上限（触发自动摘要压缩）
|
||||
pub context_char_limit: usize,
|
||||
}
|
||||
|
||||
impl Default for AgentConfig {
|
||||
fn default() -> Self {
|
||||
AgentConfig {
|
||||
max_steps: std::env::var("AGENT_MAX_STEPS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(8),
|
||||
duplicate_call_threshold: 3,
|
||||
tool_timeout_secs: std::env::var("AGENT_TOOL_TIMEOUT_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(120),
|
||||
max_tool_output_chars: std::env::var("AGENT_MAX_TOOL_OUTPUT_CHARS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(4000),
|
||||
context_char_limit: std::env::var("AGENT_CONTEXT_CHAR_LIMIT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(16000),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// SSE 流式事件（发送给前端）
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum AgentStreamEvent {
|
||||
/// 会话创建/恢复
|
||||
#[serde(rename = "session")]
|
||||
Session {
|
||||
session_id: String,
|
||||
title: String,
|
||||
},
|
||||
/// 智能体思考过程
|
||||
#[serde(rename = "thought")]
|
||||
Thought {
|
||||
content: String,
|
||||
step: usize,
|
||||
},
|
||||
/// 工具调用开始
|
||||
#[serde(rename = "tool_call")]
|
||||
ToolCall {
|
||||
name: String,
|
||||
arguments: serde_json::Value,
|
||||
step: usize,
|
||||
},
|
||||
/// 工具执行结果（Observation）
|
||||
#[serde(rename = "tool_result")]
|
||||
ToolResult {
|
||||
name: String,
|
||||
output: String,
|
||||
is_error: bool,
|
||||
metadata: serde_json::Value,
|
||||
step: usize,
|
||||
},
|
||||
/// 文本增量流式输出（最终回答）
|
||||
#[serde(rename = "text_delta")]
|
||||
TextDelta {
|
||||
content: String,
|
||||
},
|
||||
/// Token 使用统计
|
||||
#[serde(rename = "usage")]
|
||||
Usage {
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
total_tokens: u32,
|
||||
},
|
||||
/// 错误通知
|
||||
#[serde(rename = "error")]
|
||||
Error {
|
||||
message: String,
|
||||
},
|
||||
/// 完成标记
|
||||
#[serde(rename = "done")]
|
||||
Done,
|
||||
}
|
||||
|
||||
/// 同质调用检测器
|
||||
#[derive(Debug, Default)]
|
||||
struct DuplicateDetector {
|
||||
last_call: Option<(String, String)>, // (tool_name, arguments)
|
||||
consecutive_count: usize,
|
||||
}
|
||||
|
||||
impl DuplicateDetector {
|
||||
/// 记录一次调用，返回是否检测到死循环
|
||||
fn record(&mut self, tool_name: &str, arguments: &str, threshold: usize) -> bool {
|
||||
let key = (tool_name.to_string(), arguments.to_string());
|
||||
if self.last_call.as_ref() == Some(&key) {
|
||||
self.consecutive_count += 1;
|
||||
if self.consecutive_count >= threshold {
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
self.last_call = Some(key);
|
||||
self.consecutive_count = 1;
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// 智能体运行时
|
||||
pub struct AgentRuntime {
|
||||
app_state: Arc<AppState>,
|
||||
config: AgentConfig,
|
||||
tool_registry: ToolRegistry,
|
||||
}
|
||||
|
||||
impl AgentRuntime {
|
||||
/// 创建新的运行时实例
|
||||
pub fn new(app_state: Arc<AppState>) -> Self {
|
||||
AgentRuntime {
|
||||
app_state,
|
||||
config: AgentConfig::default(),
|
||||
tool_registry: ToolRegistry::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带自定义配置的运行时实例
|
||||
pub fn with_config(app_state: Arc<AppState>, config: AgentConfig) -> Self {
|
||||
AgentRuntime {
|
||||
app_state,
|
||||
config,
|
||||
tool_registry: ToolRegistry::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行完整的智能体对话回合（流式 SSE 输出）
|
||||
///
|
||||
/// 流程：
|
||||
/// 1. 加载或创建会话
|
||||
/// 2. 构建消息上下文
|
||||
/// 3. ReAct 循环：LLM 调用 -> 工具执行 -> 结果注入 -> 再次调用 ...
|
||||
/// 4. 最终回答流式输出
|
||||
/// 5. 持久化所有消息
|
||||
pub async fn run_turn(
|
||||
&self,
|
||||
session_id: Option<String>,
|
||||
question: &str,
|
||||
tx: mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
) -> anyhow::Result<String> {
|
||||
let db = &self.app_state.db;
|
||||
let llm = &self.app_state.llm;
|
||||
|
||||
// 1. 创建或恢复会话
|
||||
let sid = match session_id {
|
||||
Some(id) => {
|
||||
// 验证会话存在
|
||||
let exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)"
|
||||
)
|
||||
.bind(&id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if !exists {
|
||||
return Err(anyhow::anyhow!("会话 {} 不存在或已删除", id));
|
||||
}
|
||||
id
|
||||
}
|
||||
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?;
|
||||
new_id
|
||||
}
|
||||
};
|
||||
|
||||
let _ = tx.send(AgentStreamEvent::Session {
|
||||
session_id: sid.clone(),
|
||||
title: String::new(),
|
||||
});
|
||||
|
||||
// 2. 加载历史消息（过滤掉 thought 字段，仅保留纯对话上下文）
|
||||
let mut messages = self.load_history_for_llm(db, &sid).await?;
|
||||
|
||||
// 获取当前轮次号
|
||||
let turn_index: i32 = sqlx::query_scalar(
|
||||
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages WHERE session_id = ?"
|
||||
)
|
||||
.bind(&sid)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
// 3. 构建系统提示词
|
||||
if messages.is_empty() || messages[0].role != MessageRole::System {
|
||||
messages.insert(0, ChatMessage::system(self.system_prompt()));
|
||||
}
|
||||
|
||||
// 4. 添加用户消息
|
||||
messages.push(ChatMessage::user(question));
|
||||
self.save_message(db, &sid, turn_index, 0, &ChatMessage::user(question), None).await?;
|
||||
|
||||
// 5. ReAct 循环
|
||||
let tool_defs = self.tool_registry.definitions();
|
||||
let tool_ctx = ToolContext {
|
||||
app_state: Arc::clone(&self.app_state),
|
||||
};
|
||||
let mut duplicate_detector = DuplicateDetector::default();
|
||||
let mut step = 0;
|
||||
|
||||
loop {
|
||||
step += 1;
|
||||
// 检查是否被用户手动中止
|
||||
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
if cancelled.remove(&sid) {
|
||||
warn!("[AgentRuntime] 用户手动中止了会话 {} 的智能体执行", sid);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: "用户已手动中止执行。".to_string(),
|
||||
});
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 上下文安全检查
|
||||
let context_chars: usize = messages.iter()
|
||||
.filter_map(|m| m.content.as_ref())
|
||||
.map(|c| c.len())
|
||||
.sum();
|
||||
|
||||
if context_chars > self.config.context_char_limit {
|
||||
info!("[AgentRuntime] 上下文超限 ({} > {})，触发压缩", context_chars, self.config.context_char_limit);
|
||||
self.compress_context(&mut messages, llm).await;
|
||||
}
|
||||
|
||||
// 调用 LLM（使用 `chat_stream` 实时流式读取，支持思维过程/最终回答的流式发送和中止检测）
|
||||
let mut stream_rx = match llm.chat_stream(&messages, &tool_defs).await {
|
||||
Ok(rx) => rx,
|
||||
Err(e) => {
|
||||
error!("[AgentRuntime] LLM stream 调用失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("大模型流式调用失败: {}", e),
|
||||
});
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
let mut accumulated_content = String::new();
|
||||
let mut accumulated_reasoning = String::new();
|
||||
let mut accumulated_tool_calls: Option<Vec<crate::clients::llm::ToolCall>> = None;
|
||||
let mut usage: Option<crate::clients::llm::TokenUsage> = None;
|
||||
let mut is_tool_call_step = false;
|
||||
|
||||
let cancel_fut = async {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if let Ok(cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
if cancelled.contains(&sid) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
enum StreamLoopResult {
|
||||
Success,
|
||||
Error(String),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
let stream_loop_res = {
|
||||
let mut cancel_pinned = Box::pin(cancel_fut);
|
||||
let mut error_msg = None;
|
||||
let mut cancelled = false;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
event_opt = stream_rx.recv() => {
|
||||
match event_opt {
|
||||
Some(event) => {
|
||||
match event {
|
||||
StreamEvent::ReasoningDelta(delta) => {
|
||||
accumulated_reasoning.push_str(&delta);
|
||||
// 实时流式发送思考过程给前端
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: accumulated_reasoning.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated_content.push_str(&delta);
|
||||
// 如果目前还没发现是工具调用步骤，就实时流式发送文本给前端作为最终回答
|
||||
if !is_tool_call_step {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
content: delta,
|
||||
});
|
||||
}
|
||||
}
|
||||
StreamEvent::ToolCallsComplete(tool_calls) => {
|
||||
is_tool_call_step = true;
|
||||
accumulated_tool_calls = Some(tool_calls);
|
||||
}
|
||||
StreamEvent::ToolCallDelta { .. } => {}
|
||||
StreamEvent::Usage(u) => {
|
||||
usage = Some(u);
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
break;
|
||||
}
|
||||
StreamEvent::Error(e) => {
|
||||
error_msg = Some(e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
_ = &mut cancel_pinned => {
|
||||
cancelled = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if cancelled {
|
||||
StreamLoopResult::Cancelled
|
||||
} else if let Some(e) = error_msg {
|
||||
StreamLoopResult::Error(e)
|
||||
} else {
|
||||
StreamLoopResult::Success
|
||||
}
|
||||
};
|
||||
|
||||
match stream_loop_res {
|
||||
StreamLoopResult::Success => {}
|
||||
StreamLoopResult::Error(e_str) => {
|
||||
error!("[AgentRuntime] 流式读取错误: {}", e_str);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("大模型流式读取失败: {}", e_str),
|
||||
});
|
||||
break;
|
||||
}
|
||||
StreamLoopResult::Cancelled => {
|
||||
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
cancelled.remove(&sid);
|
||||
}
|
||||
warn!("[AgentRuntime] 在流式调用期间被用户手动中止，会话 ID: {}", sid);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: "用户已手动中止执行。".to_string(),
|
||||
});
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 处理 Thought（优先使用原生推理内容，否则如果属于工具调用步骤，使用 accumulated_content 存储）
|
||||
let mut thought_content = None;
|
||||
if !accumulated_reasoning.is_empty() {
|
||||
thought_content = Some(accumulated_reasoning.clone());
|
||||
}
|
||||
|
||||
if thought_content.is_none() && is_tool_call_step {
|
||||
if !accumulated_content.is_empty() {
|
||||
// 有工具调用时，content 被视为 Thought
|
||||
thought_content = Some(accumulated_content.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 如果是在工具调用步骤中产生的前言描述，而我们之前没实时以 Thought 发送过，此时统一作为 Thought 发送给前端展示
|
||||
if is_tool_call_step {
|
||||
if let Some(ref thought_text) = thought_content {
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: thought_text.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let reasoning_option = if accumulated_reasoning.is_empty() { None } else { Some(accumulated_reasoning.clone()) };
|
||||
|
||||
// 无工具调用 = 最终回答
|
||||
if accumulated_tool_calls.is_none() || accumulated_tool_calls.as_ref().unwrap().is_empty() {
|
||||
// 保存助手最终回答消息
|
||||
let assistant_msg = ChatMessage::assistant_with_reasoning(
|
||||
Some(accumulated_content.clone()),
|
||||
reasoning_option.clone(),
|
||||
None,
|
||||
);
|
||||
|
||||
self.save_message(db, &sid, turn_index, step as i32, &assistant_msg, reasoning_option.as_deref()).await?;
|
||||
messages.push(assistant_msg);
|
||||
|
||||
// 如果有原生推理内容且之前没发送过，发送给前端展示最终思维链
|
||||
if let Some(ref thought_text) = reasoning_option {
|
||||
if thought_content.is_none() {
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: thought_text.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 流式发送 Done 或 Token 消耗
|
||||
if let Some(u) = usage {
|
||||
let _ = tx.send(AgentStreamEvent::Usage {
|
||||
prompt_tokens: u.prompt_tokens,
|
||||
completion_tokens: u.completion_tokens,
|
||||
total_tokens: u.total_tokens,
|
||||
});
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
let tool_calls = accumulated_tool_calls.unwrap();
|
||||
|
||||
// 有工具调用 —— 构建 assistant 消息（含 tool_calls 和 reasoning_content）
|
||||
let assistant_msg = ChatMessage::assistant_with_reasoning(
|
||||
if accumulated_content.is_empty() { None } else { Some(accumulated_content.clone()) },
|
||||
reasoning_option.clone(),
|
||||
Some(tool_calls.clone()),
|
||||
);
|
||||
self.save_message(db, &sid, turn_index, step as i32, &assistant_msg, reasoning_option.as_deref()).await?;
|
||||
messages.push(assistant_msg);
|
||||
|
||||
// 逐个执行工具
|
||||
for tool_call in &tool_calls {
|
||||
let tool_name = &tool_call.function.name;
|
||||
let tool_args_str = &tool_call.function.arguments;
|
||||
|
||||
// 死循环检测
|
||||
if duplicate_detector.record(tool_name, tool_args_str, self.config.duplicate_call_threshold) {
|
||||
warn!("[AgentRuntime] 检测到死循环：{} 连续调用 {} 次", tool_name, self.config.duplicate_call_threshold);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("检测到工具 {} 的重复调用，已自动终止循环。", tool_name),
|
||||
});
|
||||
|
||||
// 注入错误 tool result 让 LLM 知道要停止
|
||||
let error_msg = ChatMessage::tool_result(
|
||||
&tool_call.id,
|
||||
format!("错误：工具 {} 被连续重复调用 {} 次，参数完全相同。请停止重复调用并直接给出目前收集到的答案。", tool_name, self.config.duplicate_call_threshold),
|
||||
);
|
||||
messages.push(error_msg);
|
||||
continue;
|
||||
}
|
||||
|
||||
// 解析参数
|
||||
let args: serde_json::Value = match serde_json::from_str(tool_args_str) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
let error_output = format!("工具参数 JSON 解析失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
name: tool_name.clone(),
|
||||
output: error_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
});
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &error_output);
|
||||
self.save_message(db, &sid, turn_index, step as i32, &tool_msg, None).await?;
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// 发送工具调用事件
|
||||
let _ = tx.send(AgentStreamEvent::ToolCall {
|
||||
name: tool_name.clone(),
|
||||
arguments: args.clone(),
|
||||
step,
|
||||
});
|
||||
|
||||
enum ToolResultEnum {
|
||||
Success(ToolOutput),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
// 执行工具（带超时保护和手动中止检测）
|
||||
let tool_res = match self.tool_registry.get(tool_name) {
|
||||
Some(tool) => {
|
||||
let timeout = std::time::Duration::from_secs(self.config.tool_timeout_secs);
|
||||
let tool_fut = tool.execute(args, &tool_ctx);
|
||||
let cancel_fut = async {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if let Ok(cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
if cancelled.contains(&sid) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
res = tokio::time::timeout(timeout, tool_fut) => {
|
||||
match res {
|
||||
Ok(output) => ToolResultEnum::Success(output),
|
||||
Err(_) => ToolResultEnum::Success(ToolOutput::error(format!("工具 {} 执行超时（{}秒）", tool_name, self.config.tool_timeout_secs))),
|
||||
}
|
||||
}
|
||||
_ = cancel_fut => {
|
||||
ToolResultEnum::Cancelled
|
||||
}
|
||||
}
|
||||
}
|
||||
None => ToolResultEnum::Success(ToolOutput::error(format!("未知工具: {}", tool_name))),
|
||||
};
|
||||
|
||||
let output = match tool_res {
|
||||
ToolResultEnum::Success(out) => out,
|
||||
ToolResultEnum::Cancelled => {
|
||||
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
|
||||
cancelled.remove(&sid);
|
||||
}
|
||||
warn!("[AgentRuntime] 在工具 {} 执行期间被用户手动中止，会话 ID: {}", tool_name, sid);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: "用户已手动中止执行。".to_string(),
|
||||
});
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
// 发送工具结果事件
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
name: tool_name.clone(),
|
||||
output: output.content.clone(),
|
||||
is_error: output.is_error,
|
||||
metadata: output.metadata.clone(),
|
||||
step,
|
||||
});
|
||||
|
||||
// 截断工具输出
|
||||
let truncated_content = if output.content.len() > self.config.max_tool_output_chars {
|
||||
let truncated: String = output.content.chars().take(self.config.max_tool_output_chars).collect();
|
||||
format!("{}...\n[输出已截断，原始长度: {} 字符]", truncated, output.content.len())
|
||||
} else {
|
||||
output.content.clone()
|
||||
};
|
||||
|
||||
// 构建 tool result 消息
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &truncated_content);
|
||||
self.save_message(db, &sid, turn_index, step as i32, &tool_msg, None).await?;
|
||||
messages.push(tool_msg);
|
||||
}
|
||||
}
|
||||
|
||||
// 6. 更新会话元信息
|
||||
let new_turn_count: i32 = sqlx::query_scalar(
|
||||
"SELECT COUNT(DISTINCT turn_index) FROM agent_messages WHERE session_id = ?"
|
||||
)
|
||||
.bind(&sid)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
// 首轮自动生成标题
|
||||
if new_turn_count <= 1 {
|
||||
let title = self.generate_title(question);
|
||||
sqlx::query("UPDATE agent_sessions SET title = ?, turn_count = ?, updated_at = CURRENT_TIMESTAMP WHERE session_id = ?")
|
||||
.bind(&title)
|
||||
.bind(new_turn_count)
|
||||
.bind(&sid)
|
||||
.execute(db)
|
||||
.await?;
|
||||
} else {
|
||||
sqlx::query("UPDATE agent_sessions SET turn_count = ?, updated_at = CURRENT_TIMESTAMP WHERE session_id = ?")
|
||||
.bind(new_turn_count)
|
||||
.bind(&sid)
|
||||
.execute(db)
|
||||
.await?;
|
||||
}
|
||||
|
||||
let _ = tx.send(AgentStreamEvent::Done);
|
||||
|
||||
Ok(sid)
|
||||
}
|
||||
|
||||
/// 系统提示词
|
||||
fn system_prompt(&self) -> String {
|
||||
"你是一位专业的天体物理学研究助手，具备丰富的天文学知识。你可以使用以下工具帮助用户进行科研工作：\n\
|
||||
\n\
|
||||
- search_papers: 检索天文学文献（ADS/arXiv）\n\
|
||||
- get_paper_content: 获取文献全文内容（自动下载、解析）\n\
|
||||
- read_local_file: 快速读取已解析的本地文献\n\
|
||||
- rag_search: 在已向量化的文献库中进行语义检索\n\
|
||||
- query_target: 查询天体物理属性（坐标、光谱型等）\n\
|
||||
\n\
|
||||
请遵循以下原则：\n\
|
||||
1. 先思考用户的问题需要什么信息，再决定调用哪些工具。\n\
|
||||
2. 优先使用已有的本地文献资源（read_local_file / rag_search），必要时再检索新文献。\n\
|
||||
3. 回答时引用具体文献来源，使用 ADS bibcode 标注。\n\
|
||||
4. 对于数学公式，使用标准 LaTeX 格式。\n\
|
||||
5. 用中文回答用户的问题，但保持科学术语的准确性（可附带英文原文）。\n\
|
||||
6. 如果一个工具调用失败，不要重复使用完全相同的参数重试，尝试换一种方式。".to_string()
|
||||
}
|
||||
|
||||
/// 从数据库加载历史消息（包含 thought 作为 reasoning_content，以备原生思考模型使用）
|
||||
async fn load_history_for_llm(
|
||||
&self,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<Vec<ChatMessage>> {
|
||||
let rows: Vec<(String, String, Option<String>, Option<String>, Option<String>)> = sqlx::query_as(
|
||||
"SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \
|
||||
WHERE session_id = ? ORDER BY id ASC"
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for (role_str, content, tool_calls_json, tool_call_id, _thought) in rows {
|
||||
let role = match role_str.as_str() {
|
||||
"system" => MessageRole::System,
|
||||
"user" => MessageRole::User,
|
||||
"assistant" => MessageRole::Assistant,
|
||||
"tool" => MessageRole::Tool,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let tool_calls: Option<Vec<crate::clients::llm::ToolCall>> = tool_calls_json
|
||||
.and_then(|json_str| serde_json::from_str(&json_str).ok());
|
||||
|
||||
messages.push(ChatMessage {
|
||||
role,
|
||||
content: if content.is_empty() { None } else { Some(content) },
|
||||
tool_call_id,
|
||||
tool_calls,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
/// 保存消息到数据库
|
||||
async fn save_message(
|
||||
&self,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
step_index: i32,
|
||||
msg: &ChatMessage,
|
||||
thought: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let role = match msg.role {
|
||||
MessageRole::System => "system",
|
||||
MessageRole::User => "user",
|
||||
MessageRole::Assistant => "assistant",
|
||||
MessageRole::Tool => "tool",
|
||||
};
|
||||
|
||||
let content = msg.content.as_deref().unwrap_or("");
|
||||
let tool_calls_json = msg.tool_calls.as_ref()
|
||||
.map(|tc| serde_json::to_string(tc).unwrap_or_default());
|
||||
let tool_call_id = msg.tool_call_id.as_deref();
|
||||
let token_count = content.len() as i32 / 4; // 粗略估算
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count) \
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(turn_index)
|
||||
.bind(step_index)
|
||||
.bind(role)
|
||||
.bind(content)
|
||||
.bind(thought)
|
||||
.bind(&tool_calls_json)
|
||||
.bind(tool_call_id)
|
||||
.bind(token_count)
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 上下文压缩：保留系统提示、最近的 user 消息、以及最近的 tool_calls/tool 对
|
||||
async fn compress_context(
|
||||
&self,
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
llm: &LlmClient,
|
||||
) {
|
||||
if messages.len() <= 4 {
|
||||
return;
|
||||
}
|
||||
|
||||
// 保留系统消息
|
||||
let system_msg = messages.first().cloned();
|
||||
|
||||
// 找到安全切割点：必须保证 assistant(tool_calls) 和后续 tool(result) 不被切断
|
||||
// 策略：保留最近 6 条消息 + 系统消息
|
||||
let keep_count = 6.min(messages.len() - 1);
|
||||
let to_summarize = &messages[1..messages.len() - keep_count];
|
||||
|
||||
if to_summarize.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
// 生成摘要
|
||||
let summary_content: String = to_summarize.iter()
|
||||
.filter_map(|m| {
|
||||
let role = match m.role {
|
||||
MessageRole::User => "用户",
|
||||
MessageRole::Assistant => "助手",
|
||||
MessageRole::Tool => "工具",
|
||||
_ => return None,
|
||||
};
|
||||
m.content.as_ref().map(|c| {
|
||||
let preview: String = c.chars().take(200).collect();
|
||||
format!("[{}] {}", role, preview)
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let summary_prompt = format!(
|
||||
"请用简洁的中文总结以下对话历史的要点（不超过500字）：\n\n{}",
|
||||
summary_content
|
||||
);
|
||||
|
||||
let summary = match llm.chat_completion(
|
||||
"你是一个对话摘要助手。请提取对话的关键信息和结论。",
|
||||
&summary_prompt,
|
||||
).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
warn!("[AgentRuntime] 上下文摘要生成失败: {}，回退为简单截断", e);
|
||||
format!("[历史摘要] 此前进行了 {} 轮对话交互", to_summarize.len())
|
||||
}
|
||||
};
|
||||
|
||||
// 重建消息列表
|
||||
let recent = messages[messages.len() - keep_count..].to_vec();
|
||||
messages.clear();
|
||||
if let Some(sys) = system_msg {
|
||||
messages.push(sys);
|
||||
}
|
||||
messages.push(ChatMessage::user(format!("[历史对话摘要]\n{}", summary)));
|
||||
messages.extend(recent);
|
||||
|
||||
info!("[AgentRuntime] 上下文压缩完成，消息数: {}", messages.len());
|
||||
}
|
||||
|
||||
/// 根据用户首条问题生成会话标题
|
||||
fn generate_title(&self, question: &str) -> String {
|
||||
let chars: String = question.chars().take(50).collect();
|
||||
if question.len() > 50 {
|
||||
format!("{}...", chars)
|
||||
} else {
|
||||
chars
|
||||
}
|
||||
}
|
||||
}
|
||||
659
src/agent/tools.rs
Normal file
659
src/agent/tools.rs
Normal file
@ -0,0 +1,659 @@
|
||||
// src/agent/tools.rs
|
||||
//
|
||||
// 科研智能体工具集定义与实现。
|
||||
// 每个工具遵循 AgentTool trait，向大模型声明 JSON Schema 参数定义，
|
||||
// 并在 execute 中调用已有的服务层完成实际业务操作。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, error};
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::ToolDefinition;
|
||||
|
||||
/// 工具执行上下文，封装全局共享状态
|
||||
pub struct ToolContext {
|
||||
pub app_state: Arc<AppState>,
|
||||
}
|
||||
|
||||
/// 工具执行结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolOutput {
|
||||
/// 给大模型阅读的截断文本
|
||||
pub content: String,
|
||||
/// 是否为错误
|
||||
pub is_error: bool,
|
||||
/// 结构化元数据（给前端 Timeline 直接渲染）
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
impl ToolOutput {
|
||||
/// 创建成功结果
|
||||
pub fn success(content: impl Into<String>, metadata: serde_json::Value) -> Self {
|
||||
ToolOutput {
|
||||
content: content.into(),
|
||||
is_error: false,
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建错误结果
|
||||
pub fn error(msg: impl Into<String>) -> Self {
|
||||
ToolOutput {
|
||||
content: msg.into(),
|
||||
is_error: true,
|
||||
metadata: json!({}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 智能体工具 trait
|
||||
#[async_trait]
|
||||
pub trait AgentTool: Send + Sync {
|
||||
/// 工具名称（与 LLM function calling 的 name 保持一致）
|
||||
fn name(&self) -> &str;
|
||||
/// 工具描述（告知 LLM 何时应该调用该工具）
|
||||
fn description(&self) -> &str;
|
||||
/// JSON Schema 格式的参数定义
|
||||
fn parameters(&self) -> serde_json::Value;
|
||||
/// 执行工具逻辑
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput;
|
||||
}
|
||||
|
||||
/// 工具注册表，管理所有可用工具
|
||||
pub struct ToolRegistry {
|
||||
tools: Vec<Box<dyn AgentTool>>,
|
||||
}
|
||||
impl ToolRegistry {
|
||||
/// 创建默认工具注册表（包含全部科研工具）
|
||||
pub fn new() -> Self {
|
||||
let tools: Vec<Box<dyn AgentTool>> = vec![
|
||||
Box::new(SearchPapersTool),
|
||||
Box::new(GetPaperMetadataTool),
|
||||
Box::new(DownloadPaperTool),
|
||||
Box::new(ParsePaperTool),
|
||||
Box::new(GetPaperContentTool),
|
||||
Box::new(RagSearchTool),
|
||||
Box::new(QueryTargetTool),
|
||||
];
|
||||
ToolRegistry { tools }
|
||||
}
|
||||
|
||||
/// 根据名称查找工具
|
||||
pub fn get(&self, name: &str) -> Option<&dyn AgentTool> {
|
||||
self.tools.iter().find(|t| t.name() == name).map(|t| t.as_ref())
|
||||
}
|
||||
|
||||
/// 生成所有工具的 ToolDefinition 列表（用于发送给 LLM）
|
||||
pub fn definitions(&self) -> Vec<ToolDefinition> {
|
||||
self.tools.iter().map(|t| {
|
||||
ToolDefinition::new(t.name(), t.description(), t.parameters())
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// 截断文本到指定最大字符数
|
||||
fn truncate_content(s: &str, max_chars: usize) -> String {
|
||||
if s.len() <= max_chars {
|
||||
s.to_string()
|
||||
} else {
|
||||
let truncated: String = s.chars().take(max_chars).collect();
|
||||
format!("{}\n\n[... 内容已截断，共 {} 字符 ...]", truncated, s.len())
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 1. SearchPapersTool ──────────────────────────
|
||||
|
||||
/// 文献搜索工具：调用 ADS/arXiv 进行跨库检索
|
||||
pub struct SearchPapersTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for SearchPapersTool {
|
||||
fn name(&self) -> &str { "search_papers" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"搜索天文学文献。支持 NASA ADS 和 arXiv 跨平台联合检索，结果自动合并去重，并关联本地馆藏状态与引用关系网。输入关键词或高级检索式，返回匹配的文献列表。\
|
||||
适用于：查找相关文献、了解研究领域现状、获取特定主题的论文。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "搜索关键词或高级检索式。支持语法：\
|
||||
1. 字段限定：au:\"作者\" 或 author:\"作者\"、ti:\"标题\" 或 title:\"标题\"、abs:\"摘要关键字\"；\
|
||||
2. 年份限定：year:2020（单年）或 year:2020-2025（年份区间）；\
|
||||
3. 逻辑运算：支持 AND、OR、NOT 逻辑组合及括号分组，如 '(ti:subdwarf OR ti:\"white dwarf\") AND year:2020-2025'；\
|
||||
4. 短语匹配：用双引号 \"\" 包含精确匹配短语，如 '\"Gaia BH1\"'。\
|
||||
所有的中文标点符号（如“”()，；）在后台均会自动清洗转换。"
|
||||
},
|
||||
"rows": {
|
||||
"type": "integer",
|
||||
"description": "返回结果数量，默认5，最大20",
|
||||
"default": 5
|
||||
}
|
||||
},
|
||||
"required": ["query"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let query = match args.get("query").and_then(|q| q.as_str()) {
|
||||
Some(q) => q.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'query'"),
|
||||
};
|
||||
let rows = args.get("rows").and_then(|r| r.as_i64()).unwrap_or(5).min(20) as i32;
|
||||
|
||||
info!("[SearchPapersTool] 执行文献搜索: query='{}', rows={}", query, rows);
|
||||
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::search::search_papers(state, &query, "all", 0, rows, "relevance").await {
|
||||
Ok(results) => {
|
||||
if results.is_empty() {
|
||||
return ToolOutput::success("未找到匹配的文献。请尝试调整搜索关键词。", json!({ "count": 0 }));
|
||||
}
|
||||
|
||||
// 格式化结果（保留详细信息给 LLM，但不含摘要且不作截断）
|
||||
let display_results: Vec<serde_json::Value> = results.iter().map(|p| {
|
||||
let first_author = p.authors.first().cloned().unwrap_or_else(|| "未知".to_string());
|
||||
json!({
|
||||
"bibcode": p.bibcode,
|
||||
"title": p.title,
|
||||
"first_author": first_author,
|
||||
"year": p.year,
|
||||
"citation_count": p.citation_count,
|
||||
})
|
||||
}).collect();
|
||||
|
||||
let content = display_results.iter().enumerate().map(|(i, r)| {
|
||||
format!(
|
||||
"{}. [{}] {} ({})\n 第一作者: {}\n 被引: {} 次",
|
||||
i + 1,
|
||||
r["bibcode"].as_str().unwrap_or(""),
|
||||
r["title"].as_str().unwrap_or(""),
|
||||
r["year"].as_str().unwrap_or(""),
|
||||
r["first_author"].as_str().unwrap_or("未知"),
|
||||
r["citation_count"].as_i64().unwrap_or(0)
|
||||
)
|
||||
}).collect::<Vec<_>>().join("\n\n");
|
||||
|
||||
ToolOutput::success(
|
||||
content,
|
||||
json!({
|
||||
"count": results.len(),
|
||||
"papers": display_results
|
||||
})
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
error!("[SearchPapersTool] 检索失败: {}", e);
|
||||
ToolOutput::error(format!("文献检索失败: {}", e))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 1b. GetPaperMetadataTool ──────────────────────────
|
||||
|
||||
/// 获取文献元数据工具：获取指定文献的完整元数据（包含完整标题、所有作者、出版期刊、关键字、引用数、完整摘要等）
|
||||
pub struct GetPaperMetadataTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GetPaperMetadataTool {
|
||||
fn name(&self) -> &str { "get_paper_metadata" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"获取指定文献的完整元数据信息（包括完整标题、所有作者、出版期刊、关键字、引用数、完整摘要等）。\
|
||||
适用于：需要查看某篇文献的详细信息、阅读完整摘要以评估文献相关性。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符，支持 ADS Bibcode（如 '2024ApJ...960..123A'）、DOI（如 '10.3847/1538-4357/ad0c5a'）或 arXiv ID（如 '2401.12345'）"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
|
||||
info!("[GetPaperMetadataTool] 获取文献元数据: {}", bibcode);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::api::helpers::get_paper_from_db(&state.db, &state.config.library_dir, &bibcode).await {
|
||||
Ok(paper) => {
|
||||
let content = format!(
|
||||
"文献元数据 [{}]：\n\
|
||||
标题: {}\n\
|
||||
作者: {}\n\
|
||||
年份: {}\n\
|
||||
期刊: {}\n\
|
||||
关键字: {}\n\
|
||||
DOI: {}\n\
|
||||
arXiv ID: {}\n\
|
||||
引用数: {} 次\n\
|
||||
参考文献数: {} 次\n\
|
||||
文献类型: {}\n\
|
||||
已下载: {}\n\
|
||||
已解析为 Markdown: {}\n\
|
||||
摘要:\n{}",
|
||||
paper.bibcode,
|
||||
paper.title,
|
||||
paper.authors.join(", "),
|
||||
paper.year,
|
||||
paper.pub_journal,
|
||||
paper.keywords.join(", "),
|
||||
paper.doi,
|
||||
paper.arxiv_id,
|
||||
paper.citation_count,
|
||||
paper.reference_count,
|
||||
paper.doctype,
|
||||
paper.is_downloaded,
|
||||
paper.has_markdown,
|
||||
paper.abstract_text
|
||||
);
|
||||
ToolOutput::success(content, json!(paper))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("获取文献 {} 元数据失败: {}", bibcode, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 2. GetPaperContentTool ──────────────────────────
|
||||
|
||||
/// 获取文献内容工具：仅从本地读取并获取已解析的文献 Markdown 全文内容
|
||||
pub struct GetPaperContentTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GetPaperContentTool {
|
||||
fn name(&self) -> &str { "get_paper_content" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"读取并在本地库中获取已解析的文献 Markdown 完整文本内容。\
|
||||
注意：本工具仅能读取已在数据库注册且已解析的文献，不会自动触发下载或解析。若文献未下载或未解析，本工具会返回详细指引，提示先依次调用 download_paper 和 parse_paper。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符，支持 ADS Bibcode（如 '2024ApJ...960..123A'）、DOI（如 '10.3847/1538-4357/ad0c5a'）或 arXiv ID（如 '2401.12345'）"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
|
||||
info!("[GetPaperContentTool] 获取文献内容: {}", bibcode);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
let paths = crate::api::helpers::check_paper_paths_in_db(&state.db, &state.config.library_dir, &bibcode).await;
|
||||
let md_opt = match paths {
|
||||
Ok(Some((_, _, md_opt, _))) => md_opt,
|
||||
Ok(None) => return ToolOutput::error(format!("获取文献内容失败：该文献未在本地数据库中注册，请先使用 search_papers 搜索该文献。")),
|
||||
Err(e) => return ToolOutput::error(format!("获取文献内容失败: {}", e)),
|
||||
};
|
||||
|
||||
let md_rel = match md_opt {
|
||||
Some(rel) => rel,
|
||||
None => return ToolOutput::error(format!("获取文献内容失败：该文献尚未完成结构化解析。如果未下载，请先调用 download_paper；如果已下载，请先调用 parse_paper 进行解析。")),
|
||||
};
|
||||
|
||||
let md_abs = state.config.library_dir.join(&md_rel);
|
||||
if !md_abs.exists() {
|
||||
return ToolOutput::error(format!("获取文献内容失败：文献本地 Markdown 文件已丢失，请重新调用 parse_paper 进行解析。"));
|
||||
}
|
||||
|
||||
match std::fs::read_to_string(&md_abs) {
|
||||
Ok(content) => ToolOutput::success(
|
||||
content.clone(),
|
||||
json!({ "bibcode": bibcode, "chars": content.len() })
|
||||
),
|
||||
Err(e) => ToolOutput::error(format!("获取文献内容失败，读取本地文件错误: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 2a. DownloadPaperTool ──────────────────────────
|
||||
|
||||
/// 下载文献全文资源工具：仅下载文献全文资源（PDF/HTML）至本地图书馆
|
||||
pub struct DownloadPaperTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for DownloadPaperTool {
|
||||
fn name(&self) -> &str { "download_paper" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"下载指定文献的全文资源（PDF 或 HTML）至本地图书馆，为后续的结构化解析做好准备。\
|
||||
适用于：需要阅读或分析新搜寻到的、尚未下载的文献。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符，支持 ADS Bibcode（如 '2024ApJ...960..123A'）、DOI（如 '10.3847/1538-4357/ad0c5a'）或 arXiv ID（如 '2401.12345'）"
|
||||
},
|
||||
"force": {
|
||||
"type": "boolean",
|
||||
"description": "是否强制重新下载（即使本地已下载该文献）"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);
|
||||
|
||||
info!("[DownloadPaperTool] 下载文献全文资源: {}, 强制重下: {}", bibcode, force);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match state.downloader.download_paper_service(
|
||||
&state.db,
|
||||
&state.config.library_dir,
|
||||
&bibcode,
|
||||
force,
|
||||
)
|
||||
.await {
|
||||
Ok(paper) => ToolOutput::success(
|
||||
format!("文献 {} 全文资源下载成功。格式 - PDF: {}, HTML: {}", bibcode, paper.has_pdf, paper.has_html),
|
||||
json!({ "bibcode": bibcode, "has_pdf": paper.has_pdf, "has_html": paper.has_html })
|
||||
),
|
||||
Err(e) => ToolOutput::error(format!("文献 {} 下载失败: {}", bibcode, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 2b. ParsePaperTool ──────────────────────────
|
||||
|
||||
/// 结构化解析文献内容工具：仅对已下载的物理资源进行结构化解析生成 Markdown
|
||||
pub struct ParsePaperTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for ParsePaperTool {
|
||||
fn name(&self) -> &str { "parse_paper" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"将指定文献本地已下载的 HTML 或 PDF 资源解析为结构化的 Markdown 文本，并保存至本地 Markdown 文件夹。\
|
||||
注意：调用此工具前必须确保文献已被成功下载（已执行 download_paper）。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"bibcode": {
|
||||
"type": "string",
|
||||
"description": "文献的唯一标识符，支持 ADS Bibcode（如 '2024ApJ...960..123A'）、DOI（如 '10.3847/1538-4357/ad0c5a'）或 arXiv ID（如 '2401.12345'）"
|
||||
},
|
||||
"force": {
|
||||
"type": "boolean",
|
||||
"description": "是否强制重新解析（即使本地已解析过该文献）"
|
||||
}
|
||||
},
|
||||
"required": ["bibcode"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
|
||||
Some(b) => b.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
|
||||
};
|
||||
let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);
|
||||
|
||||
info!("[ParsePaperTool] 结构化解析文献内容: {}, 强制重析: {}", bibcode, force);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::parser::parse_paper_service(
|
||||
&state.db,
|
||||
&state.config.library_dir,
|
||||
&state.qiniu,
|
||||
&state.config,
|
||||
&bibcode,
|
||||
force,
|
||||
)
|
||||
.await {
|
||||
Ok(markdown) => ToolOutput::success(
|
||||
format!("文献 {} 结构化解析成功。解析后 Markdown 字符总数: {}", bibcode, markdown.len()),
|
||||
json!({ "bibcode": bibcode, "chars": markdown.len() })
|
||||
),
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("请先下载") {
|
||||
ToolOutput::error(format!("文献 {} 解析失败：未检测到已下载的本地资源文件，请先调用 download_paper 工具进行下载。", bibcode))
|
||||
} else {
|
||||
ToolOutput::error(format!("文献 {} 解析失败: {}", bibcode, msg))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
// ────────────────────────── 4. RagSearchTool ──────────────────────────
|
||||
|
||||
/// RAG 向量检索工具：基于语义相似度检索文献切片
|
||||
pub struct RagSearchTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for RagSearchTool {
|
||||
fn name(&self) -> &str { "rag_search" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"在已向量化的文献库中进行语义检索。输入自然语言问题，返回最相关的文献片段。\
|
||||
适用于：跨多篇文献查找特定信息、回答需要综合多个来源的问题。要求文献已完成向量化（embed）。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": {
|
||||
"type": "string",
|
||||
"description": "用于语义检索的自然语言问题"
|
||||
},
|
||||
"top_k": {
|
||||
"type": "integer",
|
||||
"description": "返回最相关的片段数量，默认5",
|
||||
"default": 5
|
||||
}
|
||||
},
|
||||
"required": ["question"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let question = match args.get("question").and_then(|q| q.as_str()) {
|
||||
Some(q) => q.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'question'"),
|
||||
};
|
||||
let top_k = args.get("top_k").and_then(|k| k.as_u64()).unwrap_or(5) as usize;
|
||||
|
||||
info!("[RagSearchTool] 执行语义检索: question='{}', top_k={}", question, top_k);
|
||||
let state = &ctx.app_state;
|
||||
|
||||
match crate::services::rag::retrieve(&state.db, &state.embedding, &question, top_k).await {
|
||||
Ok(results) => {
|
||||
if results.is_empty() {
|
||||
return ToolOutput::success(
|
||||
"未找到相关的文献片段。文献库中可能尚无向量化数据，请先对目标文献执行向量化操作。",
|
||||
json!({ "count": 0 })
|
||||
);
|
||||
}
|
||||
|
||||
let content = results.iter().enumerate().map(|(i, r)| {
|
||||
format!(
|
||||
"[片段 {} | 来源: {} | 段落: {} | 相似度距离: {:.4}]\n{}",
|
||||
i + 1, r.bibcode, r.paragraph_index, r.distance, r.content
|
||||
)
|
||||
}).collect::<Vec<_>>().join("\n\n---\n\n");
|
||||
|
||||
let sources: Vec<serde_json::Value> = results.iter().map(|r| {
|
||||
json!({
|
||||
"bibcode": r.bibcode,
|
||||
"paragraph_index": r.paragraph_index,
|
||||
"distance": r.distance,
|
||||
"preview": truncate_content(&r.content, 100)
|
||||
})
|
||||
}).collect();
|
||||
|
||||
ToolOutput::success(
|
||||
truncate_content(&content, 4000),
|
||||
json!({ "count": results.len(), "sources": sources })
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("RAG 语义检索失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ────────────────────────── 5. QueryTargetTool ──────────────────────────
|
||||
|
||||
/// 天体信息查询工具：通过 CDS Sesame 查询天体物理属性
|
||||
pub struct QueryTargetTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for QueryTargetTool {
|
||||
fn name(&self) -> &str { "query_target" }
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"查询天体的基本物理属性信息。输入天体名称，返回坐标 (RA/Dec)、视星等、光谱型、视差等属性。\
|
||||
数据来源为 CDS SIMBAD/Sesame 名称解析服务。适用于：获取天体基本参数、验证天体身份。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"object_name": {
|
||||
"type": "string",
|
||||
"description": "天体名称，如 'NGC 6752', 'GD 358', 'HD 209458', 'M 31' 等"
|
||||
}
|
||||
},
|
||||
"required": ["object_name"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let object_name = match args.get("object_name").and_then(|n| n.as_str()) {
|
||||
Some(n) => n.to_string(),
|
||||
None => return ToolOutput::error("缺少必需参数 'object_name'"),
|
||||
};
|
||||
|
||||
info!("[QueryTargetTool] 查询天体信息: {}", object_name);
|
||||
let state = &ctx.app_state;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
match crate::services::target::query_target_cached(&state.db, &object_name, None, &client).await {
|
||||
Ok(info) => {
|
||||
let content = format!(
|
||||
"天体: {}\nRA: {}\nDec: {}\n视差: {}\n光谱型: {}\nV星等: {}\n别名: {}",
|
||||
info.target_name,
|
||||
info.ra.as_deref().unwrap_or("未知"),
|
||||
info.dec.as_deref().unwrap_or("未知"),
|
||||
info.parallax.map(|p| format!("{:.4} mas", p)).unwrap_or_else(|| "未知".to_string()),
|
||||
info.spectral_type.as_deref().unwrap_or("未知"),
|
||||
info.v_magnitude.map(|v| format!("{:.2}", v)).unwrap_or_else(|| "未知".to_string()),
|
||||
if info.aliases.is_empty() { "无".to_string() } else { info.aliases.join(", ") }
|
||||
);
|
||||
|
||||
ToolOutput::success(
|
||||
content,
|
||||
json!({
|
||||
"target_name": info.target_name,
|
||||
"ra": info.ra,
|
||||
"dec": info.dec,
|
||||
"parallax": info.parallax,
|
||||
"spectral_type": info.spectral_type,
|
||||
"v_magnitude": info.v_magnitude,
|
||||
"aliases": info.aliases
|
||||
})
|
||||
)
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("天体 '{}' 查询失败: {}", object_name, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_truncate_content_short() {
|
||||
let text = "Hello, world!";
|
||||
assert_eq!(truncate_content(text, 100), text);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_content_long() {
|
||||
let text = "a".repeat(5000);
|
||||
let result = truncate_content(&text, 100);
|
||||
assert!(result.contains("内容已截断"));
|
||||
assert!(result.contains("5000"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_output_success() {
|
||||
let output = ToolOutput::success("ok", json!({"key": "value"}));
|
||||
assert!(!output.is_error);
|
||||
assert_eq!(output.content, "ok");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_output_error() {
|
||||
let output = ToolOutput::error("something went wrong");
|
||||
assert!(output.is_error);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_registry_definitions() {
|
||||
let registry = ToolRegistry::new();
|
||||
let defs = registry.definitions();
|
||||
assert_eq!(defs.len(), 7);
|
||||
assert!(defs.iter().any(|d| d.function.name == "search_papers"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "get_paper_metadata"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "download_paper"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "parse_paper"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "get_paper_content"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "rag_search"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "query_target"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_registry_get() {
|
||||
let registry = ToolRegistry::new();
|
||||
assert!(registry.get("search_papers").is_some());
|
||||
assert!(registry.get("nonexistent").is_none());
|
||||
}
|
||||
}
|
||||
244
src/api/agent.rs
Normal file
244
src/api/agent.rs
Normal file
@ -0,0 +1,244 @@
|
||||
// src/api/agent.rs
|
||||
//
|
||||
// 科研智能体 API 控制器。
|
||||
// 提供 SSE 流式对话接口和会话管理 CRUD 接口。
|
||||
|
||||
use axum::{
|
||||
extract::{Path, Query, State},
|
||||
http::StatusCode,
|
||||
response::sse::{Event, Sse},
|
||||
Json,
|
||||
};
|
||||
use futures_util::stream::Stream;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::Row;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, error};
|
||||
|
||||
use super::AppState;
|
||||
use crate::agent::runtime::{AgentRuntime, AgentStreamEvent};
|
||||
|
||||
// ── POST /api/chat/agent ──
|
||||
// SSE 流式智能体对话接口
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct AgentChatRequest {
|
||||
pub question: String,
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn chat_agent(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(req): Json<AgentChatRequest>,
|
||||
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (StatusCode, String)> {
|
||||
info!("接收到智能体对话请求: question='{}', session_id={:?}", req.question, req.session_id);
|
||||
|
||||
let runtime = AgentRuntime::new(Arc::clone(&state));
|
||||
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 {
|
||||
Ok(sid) => {
|
||||
info!("智能体对话完成: session_id={}", sid);
|
||||
}
|
||||
Err(e) => {
|
||||
error!("智能体对话执行出错: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: format!("智能体执行错误: {}", e),
|
||||
});
|
||||
let _ = tx.send(AgentStreamEvent::Done);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// 将 mpsc 通道转换为 SSE 事件流
|
||||
let stream = async_stream::stream! {
|
||||
while let Some(event) = rx.recv().await {
|
||||
let data = serde_json::to_string(&event).unwrap_or_default();
|
||||
let is_done = matches!(event, AgentStreamEvent::Done);
|
||||
yield Ok(Event::default().data(data));
|
||||
if is_done {
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Sse::new(stream))
|
||||
}
|
||||
|
||||
// ── GET /api/chat/sessions ──
|
||||
// 获取会话列表
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SessionListParams {
|
||||
pub limit: Option<i64>,
|
||||
pub offset: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SessionSummary {
|
||||
pub session_id: String,
|
||||
pub title: String,
|
||||
pub model: String,
|
||||
pub turn_count: i32,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
pub async fn list_sessions(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Query(params): Query<SessionListParams>,
|
||||
) -> Result<Json<Vec<SessionSummary>>, (StatusCode, String)> {
|
||||
let limit = params.limit.unwrap_or(50);
|
||||
let offset = params.offset.unwrap_or(0);
|
||||
|
||||
let rows = sqlx::query(
|
||||
"SELECT session_id, title, model, turn_count, created_at, updated_at \
|
||||
FROM agent_sessions \
|
||||
WHERE deleted_at IS NULL \
|
||||
ORDER BY updated_at DESC \
|
||||
LIMIT ? OFFSET ?"
|
||||
)
|
||||
.bind(limit)
|
||||
.bind(offset)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询会话列表失败: {}", e)))?;
|
||||
|
||||
let sessions: Vec<SessionSummary> = rows.iter().map(|r| {
|
||||
SessionSummary {
|
||||
session_id: r.get(0),
|
||||
title: r.get(1),
|
||||
model: r.get(2),
|
||||
turn_count: r.get(3),
|
||||
created_at: r.get(4),
|
||||
updated_at: r.get(5),
|
||||
}
|
||||
}).collect();
|
||||
|
||||
Ok(Json(sessions))
|
||||
}
|
||||
|
||||
// ── GET /api/chat/sessions/:id ──
|
||||
// 获取单个会话的全部消息历史
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SessionDetail {
|
||||
pub session: SessionSummary,
|
||||
pub messages: Vec<MessageRecord>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct MessageRecord {
|
||||
pub id: i64,
|
||||
pub turn_index: i32,
|
||||
pub step_index: i32,
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
pub thought: Option<String>,
|
||||
pub tool_calls: Option<serde_json::Value>,
|
||||
pub tool_call_id: Option<String>,
|
||||
pub token_count: i32,
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
pub async fn get_session(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Json<SessionDetail>, (StatusCode, String)> {
|
||||
// 查询会话元信息
|
||||
let session_row = sqlx::query(
|
||||
"SELECT session_id, title, model, turn_count, created_at, updated_at \
|
||||
FROM agent_sessions \
|
||||
WHERE session_id = ? AND deleted_at IS NULL"
|
||||
)
|
||||
.bind(&session_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询会话失败: {}", e)))?
|
||||
.ok_or((StatusCode::NOT_FOUND, format!("会话 {} 不存在", session_id)))?;
|
||||
|
||||
let session = SessionSummary {
|
||||
session_id: session_row.get(0),
|
||||
title: session_row.get(1),
|
||||
model: session_row.get(2),
|
||||
turn_count: session_row.get(3),
|
||||
created_at: session_row.get(4),
|
||||
updated_at: session_row.get(5),
|
||||
};
|
||||
|
||||
// 查询消息列表
|
||||
let msg_rows = sqlx::query(
|
||||
"SELECT id, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, created_at \
|
||||
FROM agent_messages \
|
||||
WHERE session_id = ? \
|
||||
ORDER BY id ASC"
|
||||
)
|
||||
.bind(&session_id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询消息列表失败: {}", e)))?;
|
||||
|
||||
let messages: Vec<MessageRecord> = msg_rows.iter().map(|r| {
|
||||
let tool_calls_json: Option<String> = r.get(6);
|
||||
let metadata_json: Option<String> = r.get(9);
|
||||
|
||||
MessageRecord {
|
||||
id: r.get(0),
|
||||
turn_index: r.get(1),
|
||||
step_index: r.get(2),
|
||||
role: r.get(3),
|
||||
content: r.get(4),
|
||||
thought: r.get(5),
|
||||
tool_calls: tool_calls_json.and_then(|s| serde_json::from_str(&s).ok()),
|
||||
tool_call_id: r.get(7),
|
||||
token_count: r.get(8),
|
||||
metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
|
||||
created_at: r.get(10),
|
||||
}
|
||||
}).collect();
|
||||
|
||||
Ok(Json(SessionDetail { session, messages }))
|
||||
}
|
||||
|
||||
// ── DELETE /api/chat/sessions/:id ──
|
||||
// 软删除会话
|
||||
|
||||
pub async fn delete_session(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let result = sqlx::query(
|
||||
"UPDATE agent_sessions SET deleted_at = CURRENT_TIMESTAMP WHERE session_id = ? AND deleted_at IS NULL"
|
||||
)
|
||||
.bind(&session_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("删除会话失败: {}", e)))?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err((StatusCode::NOT_FOUND, format!("会话 {} 不存在或已删除", session_id)));
|
||||
}
|
||||
|
||||
info!("会话已软删除: {}", session_id);
|
||||
Ok(Json(serde_json::json!({ "status": "deleted", "session_id": session_id })))
|
||||
}
|
||||
|
||||
// ── POST /api/chat/sessions/:id/stop ──
|
||||
// 手动停止智能体执行接口
|
||||
pub async fn stop_agent(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
if let Ok(mut cancelled) = state.cancelled_runs.lock() {
|
||||
cancelled.insert(session_id.clone());
|
||||
}
|
||||
info!("已接收并记录手动中止请求，会话 ID: {}", session_id);
|
||||
Ok(Json(serde_json::json!({ "status": "stopping", "session_id": session_id })))
|
||||
}
|
||||
@ -165,9 +165,26 @@ pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Res
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_paper_from_db(db: &SqlitePool, library_dir: &std::path::Path, bibcode: &str) -> anyhow::Result<StandardPaper> {
|
||||
let r = sqlx::query("SELECT bibcode, title, authors, year, pub, keywords, abstract, doi, arxiv_id, citation_count, reference_count, pdf_path, html_path, markdown_path, translation_path, doctype, EXISTS(SELECT 1 FROM paper_chunks_content WHERE bibcode = papers.bibcode) FROM papers WHERE bibcode = ?")
|
||||
.bind(bibcode)
|
||||
pub async fn get_paper_from_db(db: &SqlitePool, library_dir: &std::path::Path, identifier: &str) -> anyhow::Result<StandardPaper> {
|
||||
let clean_id = identifier.trim();
|
||||
let clean_doi = clean_id
|
||||
.trim_start_matches("doi:")
|
||||
.trim_start_matches("DOI:")
|
||||
.trim_start_matches("https://doi.org/")
|
||||
.trim_start_matches("http://doi.org/")
|
||||
.trim();
|
||||
let clean_arxiv = clean_id
|
||||
.trim_start_matches("arxiv:")
|
||||
.trim_start_matches("arXiv:")
|
||||
.trim_start_matches("pdf/")
|
||||
.trim();
|
||||
|
||||
let r = sqlx::query("SELECT bibcode, title, authors, year, pub, keywords, abstract, doi, arxiv_id, citation_count, reference_count, pdf_path, html_path, markdown_path, translation_path, doctype, EXISTS(SELECT 1 FROM paper_chunks_content WHERE bibcode = papers.bibcode) FROM papers WHERE bibcode = ? OR doi = ? OR arxiv_id = ? OR LOWER(doi) = LOWER(?) OR arxiv_id = ?")
|
||||
.bind(clean_id)
|
||||
.bind(clean_doi)
|
||||
.bind(clean_arxiv)
|
||||
.bind(clean_doi)
|
||||
.bind(clean_arxiv)
|
||||
.fetch_one(db)
|
||||
.await?;
|
||||
|
||||
@ -222,10 +239,27 @@ pub async fn get_paper_from_db(db: &SqlitePool, library_dir: &std::path::Path, b
|
||||
pub async fn check_paper_paths_in_db(
|
||||
db: &SqlitePool,
|
||||
library_dir: &std::path::Path,
|
||||
bibcode: &str
|
||||
identifier: &str
|
||||
) -> anyhow::Result<Option<(Option<String>, Option<String>, Option<String>, Option<String>)>> {
|
||||
let r_opt = sqlx::query("SELECT pdf_path, html_path, markdown_path, translation_path FROM papers WHERE bibcode = ?")
|
||||
.bind(bibcode)
|
||||
let clean_id = identifier.trim();
|
||||
let clean_doi = clean_id
|
||||
.trim_start_matches("doi:")
|
||||
.trim_start_matches("DOI:")
|
||||
.trim_start_matches("https://doi.org/")
|
||||
.trim_start_matches("http://doi.org/")
|
||||
.trim();
|
||||
let clean_arxiv = clean_id
|
||||
.trim_start_matches("arxiv:")
|
||||
.trim_start_matches("arXiv:")
|
||||
.trim_start_matches("pdf/")
|
||||
.trim();
|
||||
|
||||
let r_opt = sqlx::query("SELECT pdf_path, html_path, markdown_path, translation_path FROM papers WHERE bibcode = ? OR doi = ? OR arxiv_id = ? OR LOWER(doi) = LOWER(?) OR arxiv_id = ?")
|
||||
.bind(clean_id)
|
||||
.bind(clean_doi)
|
||||
.bind(clean_arxiv)
|
||||
.bind(clean_doi)
|
||||
.bind(clean_arxiv)
|
||||
.fetch_optional(db)
|
||||
.await?;
|
||||
|
||||
@ -374,6 +408,12 @@ mod tests {
|
||||
assert_eq!(retrieved.authors, paper.authors);
|
||||
assert_eq!(retrieved.keywords, paper.keywords);
|
||||
|
||||
// 读取 by DOI / DOI 前缀
|
||||
let retrieved_by_doi = get_paper_from_db(&pool, std::path::Path::new(""), "10.1000/test.doi").await?;
|
||||
assert_eq!(retrieved_by_doi.bibcode, paper.bibcode);
|
||||
let retrieved_by_doi_prefix = get_paper_from_db(&pool, std::path::Path::new(""), "doi:10.1000/test.doi").await?;
|
||||
assert_eq!(retrieved_by_doi_prefix.bibcode, paper.bibcode);
|
||||
|
||||
// 检查路径状态（初始为 None）
|
||||
let paths = check_paper_paths_in_db(&pool, std::path::Path::new(""), "2026A&A...123..456X").await?;
|
||||
assert!(paths.is_some());
|
||||
|
||||
@ -19,12 +19,12 @@ pub struct AppState {
|
||||
pub ads: AdsClient,
|
||||
pub arxiv: ArxivClient,
|
||||
pub llm: LlmClient,
|
||||
pub multimodal_llm: LlmClient,
|
||||
pub embedding: EmbeddingClient,
|
||||
pub downloader: Downloader,
|
||||
pub harvest_status: Arc<tokio::sync::Mutex<crate::services::batch_sync::MetaSyncStatus>>,
|
||||
pub batch_status: Arc<tokio::sync::Mutex<crate::services::batch_sync::AssetBatchStatus>>,
|
||||
pub active_bibcode: Arc<tokio::sync::Mutex<Option<String>>>,
|
||||
pub cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
|
||||
}
|
||||
|
||||
// 统一标准化的文献格式，用于向前端传输
|
||||
@ -57,6 +57,7 @@ pub mod papers;
|
||||
pub mod notes;
|
||||
pub mod sync;
|
||||
pub mod targets;
|
||||
pub mod agent;
|
||||
|
||||
// 提供兼容的 handlers 命名空间，避免修改 main.rs / batch_sync.rs 里的导入
|
||||
pub mod handlers {
|
||||
@ -88,5 +89,9 @@ pub mod handlers {
|
||||
RagAskRequest, TargetQueryParams, AssociateTargetRequest, TargetListParams, AssociateResponse,
|
||||
ExtractTargetsRequest, ExtractTargetsResponse, ChatFigureRequest, ChatResponse,
|
||||
};
|
||||
pub use super::agent::{
|
||||
chat_agent, list_sessions, get_session, delete_session, stop_agent,
|
||||
AgentChatRequest, SessionSummary, SessionDetail, MessageRecord, SessionListParams,
|
||||
};
|
||||
pub use super::{AppState, StandardPaper};
|
||||
}
|
||||
|
||||
@ -7,13 +7,12 @@ use axum::{
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use std::fs;
|
||||
use tracing::{info, warn, error};
|
||||
use tracing::{info, error};
|
||||
use sqlx::Row;
|
||||
|
||||
use super::{AppState, StandardPaper};
|
||||
use super::helpers::{
|
||||
convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db,
|
||||
get_paper_from_db, check_paper_paths_in_db,
|
||||
convert_ads_doc_to_standard, save_paper_to_db, get_paper_from_db, check_paper_paths_in_db,
|
||||
};
|
||||
|
||||
// 检索请求参数
|
||||
@ -36,88 +35,11 @@ pub async fn search_papers(
|
||||
let rows = params.rows.unwrap_or(10);
|
||||
let start = params.start.unwrap_or(0);
|
||||
let sort = params.sort.as_deref().unwrap_or("relevance");
|
||||
let mut results = Vec::new();
|
||||
|
||||
// 1. 检索 NASA ADS
|
||||
if source == "all" || source == "ads" {
|
||||
if !state.config.ads_api_key.is_empty() {
|
||||
match state.ads.search(¶ms.q, start, rows, sort).await {
|
||||
Ok(docs) => {
|
||||
for doc in docs {
|
||||
let paper = convert_ads_doc_to_standard(&doc);
|
||||
|
||||
// 入库 SQLite
|
||||
if let Err(e) = save_paper_to_db(&state.db, &paper).await {
|
||||
warn!("保存 ADS 文献至数据库失败: {}", e);
|
||||
match crate::services::search::search_papers(&state, ¶ms.q, &source, start, rows, sort).await {
|
||||
Ok(results) => Ok(Json(results)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string())),
|
||||
}
|
||||
|
||||
// 保存引用/参考文献关联拓扑
|
||||
if let Some(refs) = doc.reference {
|
||||
for ref_bib in refs {
|
||||
let _ = sqlx::query("INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)")
|
||||
.bind(&paper.bibcode)
|
||||
.bind(&ref_bib)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
if let Some(cits) = doc.citation {
|
||||
for cit_bib in cits {
|
||||
let _ = sqlx::query("INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)")
|
||||
.bind(&cit_bib)
|
||||
.bind(&paper.bibcode)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
results.push(paper);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("ADS 检索执行失败: {}", e);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
warn!("ADS_API_KEY 未配置，跳过 ADS 检索。");
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 检索 arXiv
|
||||
if source == "all" || source == "arxiv" {
|
||||
match state.arxiv.search(¶ms.q, start, rows, sort).await {
|
||||
Ok(papers) => {
|
||||
for p in papers {
|
||||
let paper = convert_arxiv_to_standard(&p);
|
||||
|
||||
// 入库 SQLite (使用 arXiv ID 暂作主键以作记录)
|
||||
if let Err(e) = save_paper_to_db(&state.db, &paper).await {
|
||||
warn!("保存 arXiv 文献至数据库失败: {}", e);
|
||||
}
|
||||
|
||||
results.push(paper);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("arXiv 检索执行失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 对两端获取的数据进行去重合并，增加对相同 arxiv_id 的判断
|
||||
let mut unique_results: Vec<StandardPaper> = Vec::new();
|
||||
for r in results {
|
||||
if !unique_results.iter().any(|u| u.bibcode == r.bibcode || (!u.doi.is_empty() && u.doi == r.doi) || (!u.arxiv_id.is_empty() && u.arxiv_id == r.arxiv_id)) {
|
||||
let mut final_paper = r.clone();
|
||||
// 如果本地数据库存在该文献，直接从数据库读取标准元数据（包括 is_downloaded, has_markdown, pdf_error, html_error 等）
|
||||
if let Ok(db_paper) = get_paper_from_db(&state.db, &state.config.library_dir, &r.bibcode).await {
|
||||
final_paper = db_paper;
|
||||
}
|
||||
unique_results.push(final_paper);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Json(unique_results))
|
||||
}
|
||||
|
||||
// ── POST /api/download ──
|
||||
@ -135,76 +57,16 @@ pub async fn download_paper(
|
||||
let force = req.force.unwrap_or(false);
|
||||
info!("接收到文献下载指令，标识符: {}, 强制重下: {}", req.bibcode, force);
|
||||
|
||||
let paper = get_paper_from_db(&state.db, &state.config.library_dir, &req.bibcode)
|
||||
let paper = state.downloader.download_paper_service(
|
||||
&state.db,
|
||||
&state.config.library_dir,
|
||||
&req.bibcode,
|
||||
force,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("未找到该文献记录: {}", e)))?;
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
// force=true 时清除旧路径记录，强制重新下载
|
||||
if force {
|
||||
sqlx::query("UPDATE papers SET pdf_path = NULL, html_path = NULL WHERE bibcode = ?")
|
||||
.bind(&req.bibcode)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("重置下载状态失败: {}", e)))?;
|
||||
}
|
||||
|
||||
// 下载策略：
|
||||
// 1. 如有 arXiv ID，优先走 arXiv 直连（绕过出版商防护墙，成功率高）
|
||||
// 2. 若 arXiv 直连失败，回退走 ADS 网关多级回退（PUB_PDF → EPRINT_PDF → CrossRef）
|
||||
let (pdf_res, html_res) = if !paper.arxiv_id.is_empty() {
|
||||
info!("[下载] 优先使用 arXiv 通道: {}", paper.arxiv_id);
|
||||
let res = state.downloader.download_arxiv_direct(&paper.arxiv_id, &state.config.library_dir).await;
|
||||
if res.0.is_ok() || res.1.is_ok() {
|
||||
res
|
||||
} else {
|
||||
warn!("[下载] arXiv 通道下载失败，开始回退至 ADS/出版商通道: {}", req.bibcode);
|
||||
let doi_opt = if !paper.doi.is_empty() { Some(paper.doi.as_str()) } else { None };
|
||||
state.downloader.download_paper(&req.bibcode, doi_opt, &state.config.library_dir).await
|
||||
}
|
||||
} else {
|
||||
let doi_opt = if !paper.doi.is_empty() { Some(paper.doi.as_str()) } else { None };
|
||||
state.downloader.download_paper(&req.bibcode, doi_opt, &state.config.library_dir).await
|
||||
};
|
||||
|
||||
if pdf_res.is_err() && html_res.is_err() {
|
||||
let pdf_err = pdf_res.as_ref().err().unwrap();
|
||||
let html_err = html_res.as_ref().err().unwrap();
|
||||
error!("文献 {} PDF 和 HTML 均下载失败，无可用物理文件格式", req.bibcode);
|
||||
|
||||
let pdf_db_err = format!("error: {}", pdf_err);
|
||||
let html_db_err = format!("error: {}", html_err);
|
||||
let _ = sqlx::query("UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?")
|
||||
.bind(&pdf_db_err)
|
||||
.bind(&html_db_err)
|
||||
.bind(&req.bibcode)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
return Err((StatusCode::INTERNAL_SERVER_ERROR, format!("下载失败。PDF: {}, HTML: {}", pdf_err, html_err)));
|
||||
}
|
||||
|
||||
let pdf_rel = match pdf_res {
|
||||
Ok(p) => Some(p.strip_prefix(&state.config.library_dir).unwrap_or(&p).to_string_lossy().to_string()),
|
||||
Err(_) => None, // 只要有一方下载成功，失败的一方字段置空（NULL），避免在 path 字段中留存报错日志
|
||||
};
|
||||
let html_rel = match html_res {
|
||||
Ok(p) => Some(p.strip_prefix(&state.config.library_dir).unwrap_or(&p).to_string_lossy().to_string()),
|
||||
Err(_) => None, // 只要有一方下载成功，失败的一方字段置空（NULL），避免在 path 字段中留存报错日志
|
||||
};
|
||||
|
||||
// 回写存储路径至数据库
|
||||
sqlx::query("UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?")
|
||||
.bind(&pdf_rel)
|
||||
.bind(&html_rel)
|
||||
.bind(&req.bibcode)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("更新数据库失败: {}", e)))?;
|
||||
|
||||
let mut updated_paper = paper;
|
||||
updated_paper.is_downloaded = true;
|
||||
|
||||
Ok(Json(updated_paper))
|
||||
Ok(Json(paper))
|
||||
}
|
||||
|
||||
// ── POST /api/parse ──
|
||||
@ -226,112 +88,29 @@ pub async fn parse_paper(
|
||||
) -> Result<Json<ParseResponse>, (StatusCode, String)> {
|
||||
info!("接收到文献结构化解析指令: {} (强制重新解析: {:?})", req.bibcode, req.force);
|
||||
|
||||
let (pdf_opt, html_opt, md_opt, _) = check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("获取文献路径失败: {}", e)))?
|
||||
.ok_or((StatusCode::NOT_FOUND, "该文献未注册在数据库中".to_string()))?;
|
||||
|
||||
let force = req.force.unwrap_or(false);
|
||||
|
||||
// 如果先前已经解析成功过且非强制重新解析，直读 Markdown 文件返回
|
||||
if !force {
|
||||
if let Some(md_rel) = md_opt {
|
||||
let md_abs = state.config.library_dir.join(&md_rel);
|
||||
if md_abs.exists() {
|
||||
if let Ok(content) = fs::read_to_string(&md_abs) {
|
||||
return Ok(Json(ParseResponse { markdown: content }));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut parsed_markdown = String::new();
|
||||
let mut relative_md_path = String::new();
|
||||
|
||||
// 查询该文献的元数据以生成 Markdown YAML 头部信息
|
||||
let paper = get_paper_from_db(&state.db, &state.config.library_dir, &req.bibcode)
|
||||
let markdown = crate::services::parser::parse_paper_service(
|
||||
&state.db,
|
||||
&state.config.library_dir,
|
||||
&state.qiniu,
|
||||
&state.config,
|
||||
&req.bibcode,
|
||||
force,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("未找到该文献元数据记录: {}", e)))?;
|
||||
|
||||
// 策略 1：HTML 优先解析
|
||||
if let Some(html_rel) = html_opt {
|
||||
let html_abs = state.config.library_dir.join(&html_rel);
|
||||
if html_abs.exists() {
|
||||
match crate::services::parser::html_to_markdown(&html_abs) {
|
||||
Ok(md) => {
|
||||
let front_matter = format!(
|
||||
"---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"https://ui.adsabs.harvard.edu/abs/{}/abstract\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n",
|
||||
serde_json::to_string(&paper.title).unwrap_or_else(|_| format!("\"{}\"", paper.title)),
|
||||
paper.authors.iter().map(|a| format!("\"{}\"", a)).collect::<Vec<_>>().join(", "),
|
||||
serde_json::to_string(&paper.pub_journal).unwrap_or_else(|_| format!("\"{}\"", paper.pub_journal)),
|
||||
paper.bibcode,
|
||||
paper.year,
|
||||
paper.keywords.join(",")
|
||||
);
|
||||
parsed_markdown = format!("{}{}", front_matter, md);
|
||||
let md_filename = format!("{}.md", req.bibcode);
|
||||
let md_dest = state.config.library_dir.join("Markdown").join(&md_filename);
|
||||
fs::create_dir_all(md_dest.parent().unwrap()).unwrap_or_default();
|
||||
if fs::write(&md_dest, &parsed_markdown).is_ok() {
|
||||
relative_md_path = format!("Markdown/{}", md_filename);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("HTML 转换为 Markdown 失败 {}: {}。将自动降级使用 PDF 结构化解析。", req.bibcode, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 策略 2：回退至 PDF 远程 MinerU 解析
|
||||
if parsed_markdown.is_empty() {
|
||||
if let Some(pdf_rel) = pdf_opt {
|
||||
let pdf_abs = state.config.library_dir.join(&pdf_rel);
|
||||
if pdf_abs.exists() {
|
||||
match crate::services::parser::parse_pdf_via_mineru(&pdf_abs, &state.qiniu, &state.config).await {
|
||||
Ok(md) => {
|
||||
let front_matter = format!(
|
||||
"---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"https://ui.adsabs.harvard.edu/abs/{}/abstract\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n",
|
||||
serde_json::to_string(&paper.title).unwrap_or_else(|_| format!("\"{}\"", paper.title)),
|
||||
paper.authors.iter().map(|a| format!("\"{}\"", a)).collect::<Vec<_>>().join(", "),
|
||||
serde_json::to_string(&paper.pub_journal).unwrap_or_else(|_| format!("\"{}\"", paper.pub_journal)),
|
||||
paper.bibcode,
|
||||
paper.year,
|
||||
paper.keywords.join(",")
|
||||
);
|
||||
parsed_markdown = format!("{}{}", front_matter, md);
|
||||
let md_filename = format!("{}.md", req.bibcode);
|
||||
let md_dest = state.config.library_dir.join("Markdown").join(&md_filename);
|
||||
fs::create_dir_all(md_dest.parent().unwrap()).unwrap_or_default();
|
||||
if fs::write(&md_dest, &parsed_markdown).is_ok() {
|
||||
relative_md_path = format!("Markdown/{}", md_filename);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("PDF layout 远程 MinerU 解析失败: {}", e);
|
||||
return Err((StatusCode::INTERNAL_SERVER_ERROR, format!("PDF 结构解析失败: {}", e)));
|
||||
}
|
||||
}
|
||||
.map_err(|e| {
|
||||
let msg = e.to_string();
|
||||
let status = if msg.contains("未注册") || msg.contains("未找到") || msg.contains("丢失") || msg.contains("未在数据库中注册") {
|
||||
StatusCode::NOT_FOUND
|
||||
} else if msg.contains("请先下载") {
|
||||
StatusCode::BAD_REQUEST
|
||||
} else {
|
||||
error!("文献 {} 解析失败：本地 PDF 文件 {:?} 丢失", req.bibcode, pdf_abs);
|
||||
return Err((StatusCode::NOT_FOUND, "本地 PDF 文件未找到".to_string()));
|
||||
}
|
||||
} else {
|
||||
error!("文献 {} 解析失败：请先下载该文献的 HTML 或 PDF 文件", req.bibcode);
|
||||
return Err((StatusCode::BAD_REQUEST, "请先下载该文献的 HTML 或 PDF 文件".to_string()));
|
||||
}
|
||||
}
|
||||
StatusCode::INTERNAL_SERVER_ERROR
|
||||
};
|
||||
(status, msg)
|
||||
})?;
|
||||
|
||||
// 更新本地解析路径至 SQLite 数据库
|
||||
if !relative_md_path.is_empty() {
|
||||
let _ = sqlx::query("UPDATE papers SET markdown_path = ? WHERE bibcode = ?")
|
||||
.bind(&relative_md_path)
|
||||
.bind(&req.bibcode)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(Json(ParseResponse { markdown: parsed_markdown }))
|
||||
Ok(Json(ParseResponse { markdown }))
|
||||
}
|
||||
|
||||
// ── POST /api/translate ──
|
||||
|
||||
@ -241,7 +241,7 @@ pub async fn chat_figure(
|
||||
|
||||
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.multimodal_llm.chat_completion_with_image(
|
||||
let answer = state.llm.chat_completion_with_image(
|
||||
system_prompt,
|
||||
&req.question,
|
||||
&base64_data,
|
||||
|
||||
@ -42,7 +42,7 @@ enum Commands {
|
||||
|
||||
/// 手动关联天体到特定文献
|
||||
TargetAssociate {
|
||||
/// 文献 bibcode
|
||||
/// 文献唯一标识符，支持 Bibcode、DOI 或 arXiv ID
|
||||
bibcode: String,
|
||||
/// 天体名称
|
||||
object_name: String,
|
||||
@ -56,7 +56,7 @@ enum Commands {
|
||||
|
||||
/// 向量化导入指定文献的 Markdown 文件
|
||||
Ingest {
|
||||
/// 文献 bibcode
|
||||
/// 文献唯一标识符，支持 Bibcode、DOI 或 arXiv ID
|
||||
bibcode: String,
|
||||
/// Markdown 文件路径
|
||||
markdown_path: String,
|
||||
|
||||
@ -1,7 +1,198 @@
|
||||
// src/clients/llm.rs
|
||||
use serde::Deserialize;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use reqwest::Client;
|
||||
use tracing::error;
|
||||
use futures_util::StreamExt;
|
||||
|
||||
/// 消息角色枚举（OpenAI 兼容）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum MessageRole {
|
||||
System,
|
||||
User,
|
||||
Assistant,
|
||||
Tool,
|
||||
}
|
||||
|
||||
/// 工具调用结构体
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCall {
|
||||
pub id: String,
|
||||
#[serde(rename = "type")]
|
||||
pub call_type: String,
|
||||
pub function: FunctionCall,
|
||||
}
|
||||
|
||||
/// 函数调用详情
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct FunctionCall {
|
||||
pub name: String,
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
/// 标准对话消息（OpenAI 兼容格式）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChatMessage {
|
||||
pub role: MessageRole,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_call_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning_content: Option<String>,
|
||||
}
|
||||
|
||||
impl ChatMessage {
|
||||
/// 创建系统消息
|
||||
pub fn system(content: impl Into<String>) -> Self {
|
||||
ChatMessage {
|
||||
role: MessageRole::System,
|
||||
content: Some(content.into()),
|
||||
tool_call_id: None,
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建用户消息
|
||||
pub fn user(content: impl Into<String>) -> Self {
|
||||
ChatMessage {
|
||||
role: MessageRole::User,
|
||||
content: Some(content.into()),
|
||||
tool_call_id: None,
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建助手消息
|
||||
pub fn assistant(content: impl Into<String>) -> Self {
|
||||
ChatMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: Some(content.into()),
|
||||
tool_call_id: None,
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带工具调用的助手消息
|
||||
pub fn assistant_with_tool_calls(content: Option<String>, tool_calls: Vec<ToolCall>) -> Self {
|
||||
ChatMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content,
|
||||
tool_call_id: None,
|
||||
tool_calls: Some(tool_calls),
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带工具调用和推理内容的助手消息
|
||||
pub fn assistant_with_reasoning(content: Option<String>, reasoning_content: Option<String>, tool_calls: Option<Vec<ToolCall>>) -> Self {
|
||||
ChatMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content,
|
||||
tool_call_id: None,
|
||||
tool_calls,
|
||||
name: None,
|
||||
reasoning_content,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建工具结果消息
|
||||
pub fn tool_result(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
|
||||
ChatMessage {
|
||||
role: MessageRole::Tool,
|
||||
content: Some(content.into()),
|
||||
tool_call_id: Some(tool_call_id.into()),
|
||||
tool_calls: None,
|
||||
name: None,
|
||||
reasoning_content: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 工具定义（声明给大模型可用的工具）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolDefinition {
|
||||
#[serde(rename = "type")]
|
||||
pub tool_type: String,
|
||||
pub function: FunctionDef,
|
||||
}
|
||||
|
||||
/// 函数定义
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct FunctionDef {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub parameters: serde_json::Value,
|
||||
}
|
||||
|
||||
impl ToolDefinition {
|
||||
pub fn new(name: impl Into<String>, description: impl Into<String>, parameters: serde_json::Value) -> Self {
|
||||
ToolDefinition {
|
||||
tool_type: "function".to_string(),
|
||||
function: FunctionDef {
|
||||
name: name.into(),
|
||||
description: description.into(),
|
||||
parameters,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 非流式对话结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompletionResult {
|
||||
/// 回复内容（纯文本）
|
||||
pub content: Option<String>,
|
||||
/// 推理内容（模型原生 Thinking）
|
||||
pub reasoning_content: Option<String>,
|
||||
/// 工具调用列表
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
/// Token 使用统计
|
||||
pub usage: TokenUsage,
|
||||
}
|
||||
|
||||
/// Token 使用统计
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
pub struct TokenUsage {
|
||||
pub prompt_tokens: u32,
|
||||
pub completion_tokens: u32,
|
||||
pub total_tokens: u32,
|
||||
}
|
||||
|
||||
/// 流式 SSE 事件
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum StreamEvent {
|
||||
/// 文本增量
|
||||
TextDelta(String),
|
||||
/// 推理内容增量
|
||||
ReasoningDelta(String),
|
||||
/// 工具调用增量碎片
|
||||
ToolCallDelta {
|
||||
index: usize,
|
||||
id: Option<String>,
|
||||
name: Option<String>,
|
||||
arguments_delta: String,
|
||||
},
|
||||
/// 工具调用完成后还原的完整 ToolCall 列表
|
||||
ToolCallsComplete(Vec<ToolCall>),
|
||||
/// Token 使用统计
|
||||
Usage(TokenUsage),
|
||||
/// 消息结束
|
||||
Done,
|
||||
/// 错误
|
||||
Error(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct LlmClient {
|
||||
@ -165,6 +356,263 @@ impl LlmClient {
|
||||
Err(anyhow::anyhow!("大模型返回空多模态回答选项集"))
|
||||
}
|
||||
}
|
||||
|
||||
/// 多轮对话补全（含原生 Tool Calling 支持），非流式
|
||||
pub async fn chat(&self, messages: &[ChatMessage], tools: &[ToolDefinition]) -> anyhow::Result<CompletionResult> {
|
||||
let url = format!("{}/chat/completions", self.api_base);
|
||||
|
||||
let mut payload = serde_json::json!({
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"temperature": 0.3
|
||||
});
|
||||
|
||||
// 如果是通义千问 (DashScope) 或者是 Qwen 模型，自动开启思考模式
|
||||
if self.api_base.contains("dashscope.aliyuncs.com") || self.model.to_lowercase().contains("qwen") {
|
||||
if let Some(obj) = payload.as_object_mut() {
|
||||
obj.insert("enable_thinking".to_string(), serde_json::json!(true));
|
||||
}
|
||||
}
|
||||
|
||||
if !tools.is_empty() {
|
||||
payload["tools"] = serde_json::to_value(tools)?;
|
||||
}
|
||||
|
||||
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 chat 接口调用失败: 状态码={}, 报错={}", status, body);
|
||||
return Err(anyhow::anyhow!("大模型 chat 接口返回错误状态: {} - {}", status, body));
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ResponseMessage {
|
||||
content: Option<String>,
|
||||
tool_calls: Option<Vec<ToolCall>>,
|
||||
reasoning_content: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Choice {
|
||||
message: ResponseMessage,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChatResponse {
|
||||
choices: Vec<Choice>,
|
||||
usage: Option<TokenUsage>,
|
||||
}
|
||||
|
||||
let res_data: ChatResponse = response.json().await?;
|
||||
if let Some(choice) = res_data.choices.into_iter().next() {
|
||||
Ok(CompletionResult {
|
||||
content: choice.message.content,
|
||||
reasoning_content: choice.message.reasoning_content,
|
||||
tool_calls: choice.message.tool_calls.unwrap_or_default(),
|
||||
usage: res_data.usage.unwrap_or_default(),
|
||||
})
|
||||
} else {
|
||||
Err(anyhow::anyhow!("大模型返回空对话选项集"))
|
||||
}
|
||||
}
|
||||
|
||||
/// 流式多轮对话补全（含 Tool Call 碎片累积还原），返回 SSE 事件流
|
||||
pub async fn chat_stream(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
tools: &[ToolDefinition],
|
||||
) -> anyhow::Result<tokio::sync::mpsc::UnboundedReceiver<StreamEvent>> {
|
||||
let url = format!("{}/chat/completions", self.api_base);
|
||||
|
||||
let mut payload = serde_json::json!({
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"temperature": 0.3,
|
||||
"stream": true,
|
||||
"stream_options": { "include_usage": true }
|
||||
});
|
||||
|
||||
// 如果是通义千问 (DashScope) 或者是 Qwen 模型，自动开启思考模式
|
||||
if self.api_base.contains("dashscope.aliyuncs.com") || self.model.to_lowercase().contains("qwen") {
|
||||
if let Some(obj) = payload.as_object_mut() {
|
||||
obj.insert("enable_thinking".to_string(), serde_json::json!(true));
|
||||
}
|
||||
}
|
||||
|
||||
if !tools.is_empty() {
|
||||
payload["tools"] = serde_json::to_value(tools)?;
|
||||
}
|
||||
|
||||
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 stream 接口调用失败: 状态码={}, 报错={}", status, body);
|
||||
return Err(anyhow::anyhow!("大模型 stream 接口返回错误状态: {} - {}", status, body));
|
||||
}
|
||||
|
||||
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let mut byte_stream = response.bytes_stream();
|
||||
|
||||
tokio::spawn(async move {
|
||||
// 流式 Tool Call 碎片累积器
|
||||
let mut tool_call_accumulators: std::collections::HashMap<usize, (String, String, String)> = std::collections::HashMap::new();
|
||||
let mut buffer = String::new();
|
||||
|
||||
while let Some(chunk_result) = byte_stream.next().await {
|
||||
let chunk = match chunk_result {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
let _ = tx.send(StreamEvent::Error(format!("流式读取错误: {}", e)));
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
// 按行解析 SSE 事件
|
||||
while let Some(line_end) = buffer.find('\n') {
|
||||
let line = buffer[..line_end].trim().to_string();
|
||||
buffer = buffer[line_end + 1..].to_string();
|
||||
|
||||
if line.is_empty() || !line.starts_with("data: ") {
|
||||
continue;
|
||||
}
|
||||
|
||||
let data = &line[6..];
|
||||
if data == "[DONE]" {
|
||||
// 流结束前，将累积的 tool calls 还原并发送
|
||||
if !tool_call_accumulators.is_empty() {
|
||||
let mut indices: Vec<usize> = tool_call_accumulators.keys().cloned().collect();
|
||||
indices.sort();
|
||||
let tool_calls: Vec<ToolCall> = indices.into_iter().map(|idx| {
|
||||
let (id, name, args) = tool_call_accumulators.remove(&idx).unwrap();
|
||||
ToolCall {
|
||||
id,
|
||||
call_type: "function".to_string(),
|
||||
function: FunctionCall { name, arguments: args },
|
||||
}
|
||||
}).collect();
|
||||
let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls));
|
||||
}
|
||||
let _ = tx.send(StreamEvent::Done);
|
||||
return;
|
||||
}
|
||||
|
||||
// 解析 JSON delta
|
||||
let parsed: serde_json::Value = match serde_json::from_str(data) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
// 提取 usage（在最后一条 chunk 中 stream_options 返回）
|
||||
if let Some(usage_val) = parsed.get("usage") {
|
||||
if !usage_val.is_null() {
|
||||
if let Ok(usage) = serde_json::from_value::<TokenUsage>(usage_val.clone()) {
|
||||
let _ = tx.send(StreamEvent::Usage(usage));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(choices) = parsed.get("choices").and_then(|c| c.as_array()) {
|
||||
if let Some(choice) = choices.first() {
|
||||
let delta = &choice["delta"];
|
||||
|
||||
// 文本增量
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
if !content.is_empty() {
|
||||
let _ = tx.send(StreamEvent::TextDelta(content.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
// 推理内容增量
|
||||
if let Some(reasoning) = delta.get("reasoning_content").and_then(|c| c.as_str()) {
|
||||
if !reasoning.is_empty() {
|
||||
let _ = tx.send(StreamEvent::ReasoningDelta(reasoning.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
// 工具调用增量
|
||||
if let Some(tool_calls) = delta.get("tool_calls").and_then(|t| t.as_array()) {
|
||||
for tc in tool_calls {
|
||||
let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
|
||||
let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
|
||||
let fn_name = tc.get("function").and_then(|f| f.get("name")).and_then(|n| n.as_str()).map(|s| s.to_string());
|
||||
let fn_args = tc.get("function").and_then(|f| f.get("arguments")).and_then(|a| a.as_str()).unwrap_or("");
|
||||
|
||||
let entry = tool_call_accumulators.entry(index).or_insert_with(|| (
|
||||
String::new(), String::new(), String::new()
|
||||
));
|
||||
if let Some(ref id_str) = id {
|
||||
entry.0 = id_str.clone();
|
||||
}
|
||||
if let Some(ref name_str) = fn_name {
|
||||
entry.1 = name_str.clone();
|
||||
}
|
||||
entry.2.push_str(fn_args);
|
||||
|
||||
let _ = tx.send(StreamEvent::ToolCallDelta {
|
||||
index,
|
||||
id,
|
||||
name: fn_name,
|
||||
arguments_delta: fn_args.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 检查是否 finish_reason == "tool_calls"
|
||||
if let Some(finish_reason) = choice.get("finish_reason").and_then(|f| f.as_str()) {
|
||||
if finish_reason == "tool_calls" && !tool_call_accumulators.is_empty() {
|
||||
let mut indices: Vec<usize> = tool_call_accumulators.keys().cloned().collect();
|
||||
indices.sort();
|
||||
let tool_calls: Vec<ToolCall> = indices.into_iter().map(|idx| {
|
||||
let (id, name, args) = tool_call_accumulators.remove(&idx).unwrap();
|
||||
ToolCall {
|
||||
id,
|
||||
call_type: "function".to_string(),
|
||||
function: FunctionCall { name, arguments: args },
|
||||
}
|
||||
}).collect();
|
||||
let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 异常结束（连接断开等）
|
||||
if !tool_call_accumulators.is_empty() {
|
||||
let mut indices: Vec<usize> = tool_call_accumulators.keys().cloned().collect();
|
||||
indices.sort();
|
||||
let tool_calls: Vec<ToolCall> = indices.into_iter().map(|idx| {
|
||||
let (id, name, args) = tool_call_accumulators.remove(&idx).unwrap();
|
||||
ToolCall {
|
||||
id,
|
||||
call_type: "function".to_string(),
|
||||
function: FunctionCall { name, arguments: args },
|
||||
}
|
||||
}).collect();
|
||||
let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls));
|
||||
}
|
||||
let _ = tx.send(StreamEvent::Done);
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
@ -266,6 +714,117 @@ mod tests {
|
||||
assert_eq!(client.model(), "model");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_message_construction() {
|
||||
// 测试系统消息构造
|
||||
let sys = ChatMessage::system("你是一个助手");
|
||||
assert_eq!(sys.role, MessageRole::System);
|
||||
assert_eq!(sys.content.as_deref(), Some("你是一个助手"));
|
||||
assert!(sys.tool_calls.is_none());
|
||||
|
||||
// 测试用户消息构造
|
||||
let user = ChatMessage::user("你好");
|
||||
assert_eq!(user.role, MessageRole::User);
|
||||
assert_eq!(user.content.as_deref(), Some("你好"));
|
||||
|
||||
// 测试助手消息构造
|
||||
let assistant = ChatMessage::assistant("你好！有什么可以帮助你的吗？");
|
||||
assert_eq!(assistant.role, MessageRole::Assistant);
|
||||
assert_eq!(assistant.content.as_deref(), 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}"#));
|
||||
|
||||
// 测试带工具调用的助手消息
|
||||
let tool_calls = vec![ToolCall {
|
||||
id: "call_abc".to_string(),
|
||||
call_type: "function".to_string(),
|
||||
function: FunctionCall {
|
||||
name: "get_weather".to_string(),
|
||||
arguments: r#"{"city": "北京"}"#.to_string(),
|
||||
},
|
||||
}];
|
||||
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_eq!(assistant_tc.tool_calls.as_ref().unwrap().len(), 1);
|
||||
assert_eq!(assistant_tc.tool_calls.as_ref().unwrap()[0].function.name, "get_weather");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_definition_construction() {
|
||||
let tool = ToolDefinition::new(
|
||||
"search",
|
||||
"搜索互联网内容",
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "搜索关键词"
|
||||
}
|
||||
},
|
||||
"required": ["query"]
|
||||
}),
|
||||
);
|
||||
assert_eq!(tool.tool_type, "function");
|
||||
assert_eq!(tool.function.name, "search");
|
||||
assert_eq!(tool.function.description, "搜索互联网内容");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_message_serialization() {
|
||||
// 测试序列化时 skip_serializing_if 生效
|
||||
let msg = ChatMessage::user("测试消息");
|
||||
let json = serde_json::to_value(&msg).unwrap();
|
||||
assert_eq!(json.get("role").unwrap(), "user");
|
||||
assert_eq!(json.get("content").unwrap(), "测试消息");
|
||||
// None 字段不应出现在 JSON 中
|
||||
assert!(json.get("tool_call_id").is_none());
|
||||
assert!(json.get("tool_calls").is_none());
|
||||
assert!(json.get("name").is_none());
|
||||
assert!(json.get("reasoning_content").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chat_message_reasoning_serialization() {
|
||||
// 测试带有 reasoning_content 的消息序列化和反序列化
|
||||
let msg = ChatMessage::assistant_with_reasoning(
|
||||
Some("回答内容".to_string()),
|
||||
Some("这是思考过程".to_string()),
|
||||
None,
|
||||
);
|
||||
let json = serde_json::to_value(&msg).unwrap();
|
||||
assert_eq!(json.get("role").unwrap(), "assistant");
|
||||
assert_eq!(json.get("content").unwrap(), "回答内容");
|
||||
assert_eq!(json.get("reasoning_content").unwrap(), "这是思考过程");
|
||||
assert!(json.get("tool_calls").is_none());
|
||||
|
||||
// 测试反序列化
|
||||
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.reasoning_content.as_deref(), Some("这是思考过程"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_message_role_serialization() {
|
||||
// 测试角色枚举的 rename_all = "lowercase" 序列化
|
||||
let role_json = serde_json::to_string(&MessageRole::System).unwrap();
|
||||
assert_eq!(role_json, r#""system""#);
|
||||
let role_json = serde_json::to_string(&MessageRole::Assistant).unwrap();
|
||||
assert_eq!(role_json, r#""assistant""#);
|
||||
let role_json = serde_json::to_string(&MessageRole::Tool).unwrap();
|
||||
assert_eq!(role_json, r#""tool""#);
|
||||
|
||||
// 测试反序列化
|
||||
let role: MessageRole = serde_json::from_str(r#""user""#).unwrap();
|
||||
assert_eq!(role, MessageRole::User);
|
||||
}
|
||||
|
||||
/// 真实网络集成测试 —— 需要配置 LLM_API_KEY 和 EMBEDDING_API_KEY
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
|
||||
14
src/lib.rs
14
src/lib.rs
@ -10,9 +10,6 @@ pub struct Config {
|
||||
pub llm_api_key: String, // 大语言模型 API Key
|
||||
pub llm_api_base: String, // 大语言模型 API 基础地址
|
||||
pub llm_model: String, // 调用的翻译大模型名称
|
||||
pub multimodal_api_key: String, // 多模态模型 API Key (如不设置回退到 llm_api_key)
|
||||
pub multimodal_api_base: String,// 多模态模型 API 基础地址 (如不设置回退到 llm_api_base)
|
||||
pub multimodal_model: String, // 多模态模型名称 (如不设置回退到 llm_model)
|
||||
pub embedding_api_key: String, // 向量模型 API Key
|
||||
pub embedding_api_base: String,// 向量模型 API 基础地址
|
||||
pub embedding_model: String, // 向量模型名称
|
||||
@ -40,13 +37,6 @@ impl Config {
|
||||
let llm_model = env::var("LLM_MODEL")
|
||||
.unwrap_or_else(|_| "gpt-4o-mini".to_string());
|
||||
|
||||
let multimodal_api_key = env::var("MULTIMODAL_API_KEY")
|
||||
.unwrap_or_else(|_| llm_api_key.clone());
|
||||
let multimodal_api_base = env::var("MULTIMODAL_API_BASE")
|
||||
.unwrap_or_else(|_| llm_api_base.clone());
|
||||
let multimodal_model = env::var("MULTIMODAL_MODEL")
|
||||
.unwrap_or_else(|_| llm_model.clone());
|
||||
|
||||
let embedding_api_key = env::var("EMBEDDING_API_KEY")
|
||||
.unwrap_or_else(|_| llm_api_key.clone());
|
||||
let embedding_api_base = env::var("EMBEDDING_API_BASE")
|
||||
@ -76,9 +66,6 @@ impl Config {
|
||||
llm_api_key,
|
||||
llm_api_base,
|
||||
llm_model,
|
||||
multimodal_api_key,
|
||||
multimodal_api_base,
|
||||
multimodal_model,
|
||||
embedding_api_key,
|
||||
embedding_api_base,
|
||||
embedding_model,
|
||||
@ -97,6 +84,7 @@ impl Config {
|
||||
pub mod api;
|
||||
pub mod clients;
|
||||
pub mod services;
|
||||
pub mod agent;
|
||||
|
||||
#[cfg(test)]
|
||||
mod config_tests {
|
||||
|
||||
14
src/main.rs
14
src/main.rs
@ -122,11 +122,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
config.llm_api_base.clone(),
|
||||
config.llm_model.clone(),
|
||||
);
|
||||
let multimodal_llm = LlmClient::new(
|
||||
config.multimodal_api_key.clone(),
|
||||
config.multimodal_api_base.clone(),
|
||||
config.multimodal_model.clone(),
|
||||
);
|
||||
let embedding = EmbeddingClient::new(
|
||||
config.embedding_api_key.clone(),
|
||||
config.embedding_api_base.clone(),
|
||||
@ -141,12 +136,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
ads,
|
||||
arxiv,
|
||||
llm,
|
||||
multimodal_llm,
|
||||
embedding,
|
||||
downloader,
|
||||
harvest_status: Arc::new(tokio::sync::Mutex::new(astroresearch::services::batch_sync::MetaSyncStatus::new())),
|
||||
batch_status: Arc::new(tokio::sync::Mutex::new(astroresearch::services::batch_sync::AssetBatchStatus::new())),
|
||||
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
|
||||
cancelled_runs: Arc::new(std::sync::Mutex::new(std::collections::HashSet::new())),
|
||||
});
|
||||
|
||||
// 7. 设置 Axum 路由、CORS 头以及 React 仪表盘静态资源托管
|
||||
@ -184,7 +179,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
.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("/target/list", get(handlers::list_targets))
|
||||
// 智能体路由
|
||||
.route("/chat/agent", post(handlers::chat_agent))
|
||||
.route("/chat/sessions", get(handlers::list_sessions))
|
||||
.route("/chat/sessions/:id", get(handlers::get_session).delete(handlers::delete_session))
|
||||
.route("/chat/sessions/:id/stop", post(handlers::stop_agent));
|
||||
|
||||
// 静态文件资源代理托管（当前端打包至 dashboard/dist 后，直接挂载到主域名根路由）
|
||||
let serve_dir = ServeDir::new("dashboard/dist")
|
||||
|
||||
@ -14,8 +14,10 @@ use std::path::{Path, PathBuf};
|
||||
use reqwest::header::{HeaderMap, HeaderValue};
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use url::Url;
|
||||
use tracing::{info, warn, debug};
|
||||
use tracing::{info, warn, debug, error};
|
||||
use anyhow::{Context, Result};
|
||||
use crate::api::StandardPaper;
|
||||
use crate::api::helpers::{get_paper_from_db, check_paper_paths_in_db};
|
||||
|
||||
// ─── 浏览器伪装辅助 ────────────────────────────────────────────
|
||||
|
||||
@ -991,6 +993,93 @@ impl Downloader {
|
||||
|
||||
(pdf_res, html_res)
|
||||
}
|
||||
|
||||
/// 高层文献下载服务，处理数据库记录、多级回退下载、数据库路径写回
|
||||
pub async fn download_paper_service(
|
||||
&self,
|
||||
db: &sqlx::SqlitePool,
|
||||
library_dir: &std::path::Path,
|
||||
bibcode: &str,
|
||||
force: bool,
|
||||
) -> anyhow::Result<StandardPaper> {
|
||||
let paper = get_paper_from_db(db, library_dir, bibcode)
|
||||
.await?;
|
||||
|
||||
if force {
|
||||
sqlx::query("UPDATE papers SET pdf_path = NULL, html_path = NULL WHERE bibcode = ?")
|
||||
.bind(bibcode)
|
||||
.execute(db)
|
||||
.await?;
|
||||
} else {
|
||||
// 如果已经下载了，直接返回
|
||||
let paths = check_paper_paths_in_db(db, library_dir, bibcode).await?;
|
||||
if let Some((pdf_opt, html_opt, _, _)) = paths {
|
||||
if pdf_opt.is_some() || html_opt.is_some() {
|
||||
let mut updated_paper = paper.clone();
|
||||
updated_paper.is_downloaded = true;
|
||||
updated_paper.has_pdf = pdf_opt.is_some();
|
||||
updated_paper.has_html = html_opt.is_some();
|
||||
return Ok(updated_paper);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 调用底层的 download pipeline
|
||||
let (pdf_res, html_res) = if !paper.arxiv_id.is_empty() {
|
||||
info!("[下载] 优先使用 arXiv 通道: {}", paper.arxiv_id);
|
||||
let res = self.download_arxiv_direct(&paper.arxiv_id, library_dir).await;
|
||||
if res.0.is_ok() || res.1.is_ok() {
|
||||
res
|
||||
} else {
|
||||
warn!("[下载] arXiv 通道下载失败，开始回退至 ADS/出版商通道: {}", bibcode);
|
||||
let doi_opt = if !paper.doi.is_empty() { Some(paper.doi.as_str()) } else { None };
|
||||
self.download_paper(bibcode, doi_opt, library_dir).await
|
||||
}
|
||||
} else {
|
||||
let doi_opt = if !paper.doi.is_empty() { Some(paper.doi.as_str()) } else { None };
|
||||
self.download_paper(bibcode, doi_opt, library_dir).await
|
||||
};
|
||||
|
||||
if pdf_res.is_err() && html_res.is_err() {
|
||||
let pdf_err = pdf_res.as_ref().err().unwrap();
|
||||
let html_err = html_res.as_ref().err().unwrap();
|
||||
error!("文献 {} PDF 和 HTML 均下载失败，无可用物理文件格式", bibcode);
|
||||
|
||||
let pdf_db_err = format!("error: {}", pdf_err);
|
||||
let html_db_err = format!("error: {}", html_err);
|
||||
let _ = sqlx::query("UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?")
|
||||
.bind(&pdf_db_err)
|
||||
.bind(&html_db_err)
|
||||
.bind(bibcode)
|
||||
.execute(db)
|
||||
.await;
|
||||
|
||||
return Err(anyhow::anyhow!("下载失败。PDF: {}, HTML: {}", pdf_err, html_err));
|
||||
}
|
||||
|
||||
let pdf_rel = match pdf_res {
|
||||
Ok(p) => Some(p.strip_prefix(library_dir).unwrap_or(&p).to_string_lossy().to_string()),
|
||||
Err(_) => None,
|
||||
};
|
||||
let html_rel = match html_res {
|
||||
Ok(p) => Some(p.strip_prefix(library_dir).unwrap_or(&p).to_string_lossy().to_string()),
|
||||
Err(_) => None,
|
||||
};
|
||||
|
||||
sqlx::query("UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?")
|
||||
.bind(&pdf_rel)
|
||||
.bind(&html_rel)
|
||||
.bind(bibcode)
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
let mut updated_paper = paper;
|
||||
updated_paper.is_downloaded = true;
|
||||
updated_paper.has_pdf = pdf_rel.is_some();
|
||||
updated_paper.has_html = html_rel.is_some();
|
||||
|
||||
Ok(updated_paper)
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_arxiv_version(arxiv_id: &str) -> String {
|
||||
|
||||
@ -7,6 +7,7 @@ pub mod logging;
|
||||
pub mod chunker;
|
||||
pub mod rag;
|
||||
pub mod target;
|
||||
pub mod search;
|
||||
|
||||
pub mod batch_sync {
|
||||
pub use super::batch::*;
|
||||
|
||||
@ -7,7 +7,8 @@ pub mod common;
|
||||
pub mod pdf;
|
||||
|
||||
use std::path::Path;
|
||||
use tracing::info;
|
||||
use tracing::{info, warn, error};
|
||||
use crate::api::helpers::{check_paper_paths_in_db, get_paper_from_db};
|
||||
|
||||
use ar5iv::Ar5ivParser;
|
||||
use aanda::AandaParser;
|
||||
@ -102,6 +103,122 @@ pub fn html_to_markdown(html_path: &Path) -> anyhow::Result<String> {
|
||||
Ok(cleaned)
|
||||
}
|
||||
|
||||
/// 高层文献解析服务，支持 HTML 优先解析和 PDF 降级 MinerU 远程解析，并自动将解析结果写入本地与数据库中
|
||||
pub async fn parse_paper_service(
|
||||
db: &sqlx::SqlitePool,
|
||||
library_dir: &std::path::Path,
|
||||
qiniu: &crate::clients::qiniu::QiniuClient,
|
||||
config: &crate::Config,
|
||||
bibcode: &str,
|
||||
force: bool,
|
||||
) -> anyhow::Result<String> {
|
||||
info!("接收到文献结构化解析服务调用: {} (强制重新解析: {:?})", bibcode, force);
|
||||
|
||||
let (pdf_opt, html_opt, md_opt, _) = check_paper_paths_in_db(db, library_dir, bibcode)
|
||||
.await?
|
||||
.ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?;
|
||||
|
||||
// 如果先前已经解析成功过且非强制重新解析，直读 Markdown 文件返回
|
||||
if !force {
|
||||
if let Some(md_rel) = md_opt {
|
||||
let md_abs = library_dir.join(&md_rel);
|
||||
if md_abs.exists() {
|
||||
if let Ok(content) = std::fs::read_to_string(&md_abs) {
|
||||
return Ok(content);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut parsed_markdown = String::new();
|
||||
let mut relative_md_path = String::new();
|
||||
|
||||
// 查询该文献的元数据以生成 Markdown YAML 头部信息
|
||||
let paper = get_paper_from_db(db, library_dir, bibcode)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("未找到该文献元数据记录: {}", e))?;
|
||||
|
||||
// 策略 1：HTML 优先解析
|
||||
if let Some(html_rel) = html_opt {
|
||||
let html_abs = library_dir.join(&html_rel);
|
||||
if html_abs.exists() {
|
||||
match html_to_markdown(&html_abs) {
|
||||
Ok(md) => {
|
||||
let front_matter = format!(
|
||||
"---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"https://ui.adsabs.harvard.edu/abs/{}/abstract\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n",
|
||||
serde_json::to_string(&paper.title).unwrap_or_else(|_| format!("\"{}\"", paper.title)),
|
||||
paper.authors.iter().map(|a| format!("\"{}\"", a)).collect::<Vec<_>>().join(", "),
|
||||
serde_json::to_string(&paper.pub_journal).unwrap_or_else(|_| format!("\"{}\"", paper.pub_journal)),
|
||||
paper.bibcode,
|
||||
paper.year,
|
||||
paper.keywords.join(",")
|
||||
);
|
||||
parsed_markdown = format!("{}{}", front_matter, md);
|
||||
let md_filename = format!("{}.md", bibcode);
|
||||
let md_dest = library_dir.join("Markdown").join(&md_filename);
|
||||
std::fs::create_dir_all(md_dest.parent().unwrap()).unwrap_or_default();
|
||||
if std::fs::write(&md_dest, &parsed_markdown).is_ok() {
|
||||
relative_md_path = format!("Markdown/{}", md_filename);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("HTML 转换为 Markdown 失败 {}: {}。将自动降级使用 PDF 结构化解析。", bibcode, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 策略 2：回退至 PDF 远程 MinerU 解析
|
||||
if parsed_markdown.is_empty() {
|
||||
if let Some(pdf_rel) = pdf_opt {
|
||||
let pdf_abs = library_dir.join(&pdf_rel);
|
||||
if pdf_abs.exists() {
|
||||
match parse_pdf_via_mineru(&pdf_abs, qiniu, config).await {
|
||||
Ok(md) => {
|
||||
let front_matter = format!(
|
||||
"---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"https://ui.adsabs.harvard.edu/abs/{}/abstract\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n",
|
||||
serde_json::to_string(&paper.title).unwrap_or_else(|_| format!("\"{}\"", paper.title)),
|
||||
paper.authors.iter().map(|a| format!("\"{}\"", a)).collect::<Vec<_>>().join(", "),
|
||||
serde_json::to_string(&paper.pub_journal).unwrap_or_else(|_| format!("\"{}\"", paper.pub_journal)),
|
||||
paper.bibcode,
|
||||
paper.year,
|
||||
paper.keywords.join(",")
|
||||
);
|
||||
parsed_markdown = format!("{}{}", front_matter, md);
|
||||
let md_filename = format!("{}.md", bibcode);
|
||||
let md_dest = library_dir.join("Markdown").join(&md_filename);
|
||||
std::fs::create_dir_all(md_dest.parent().unwrap()).unwrap_or_default();
|
||||
if std::fs::write(&md_dest, &parsed_markdown).is_ok() {
|
||||
relative_md_path = format!("Markdown/{}", md_filename);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("PDF layout 远程 MinerU 解析失败: {}", e);
|
||||
return Err(anyhow::anyhow!("PDF 结构解析失败: {}", e));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
error!("文献 {} 解析失败：本地 PDF 文件 {:?} 丢失", bibcode, pdf_abs);
|
||||
return Err(anyhow::anyhow!("本地 PDF 文件未找到"));
|
||||
}
|
||||
} else {
|
||||
error!("文献 {} 解析失败：请先下载该文献的 HTML 或 PDF 文件", bibcode);
|
||||
return Err(anyhow::anyhow!("请先下载该文献的 HTML 或 PDF 文件"));
|
||||
}
|
||||
}
|
||||
|
||||
// 更新本地解析路径至 SQLite 数据库
|
||||
if !relative_md_path.is_empty() {
|
||||
sqlx::query("UPDATE papers SET markdown_path = ? WHERE bibcode = ?")
|
||||
.bind(&relative_md_path)
|
||||
.bind(bibcode)
|
||||
.execute(db)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(parsed_markdown)
|
||||
}
|
||||
|
||||
// ── 测试 ──────────────────────────────────────────────────────────
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
101
src/services/search.rs
Normal file
101
src/services/search.rs
Normal file
@ -0,0 +1,101 @@
|
||||
// src/services/search.rs
|
||||
use tracing::{warn, error};
|
||||
use crate::api::AppState;
|
||||
use crate::api::StandardPaper;
|
||||
use crate::api::helpers::{
|
||||
convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db,
|
||||
get_paper_from_db,
|
||||
};
|
||||
|
||||
/// 统一检索逻辑，合并去重 ADS 和 arXiv 数据，并自动记录引用拓扑与同步本地状态
|
||||
pub async fn search_papers(
|
||||
state: &AppState,
|
||||
q: &str,
|
||||
source: &str,
|
||||
start: i32,
|
||||
rows: i32,
|
||||
sort: &str,
|
||||
) -> anyhow::Result<Vec<StandardPaper>> {
|
||||
let mut results = Vec::new();
|
||||
|
||||
// 1. 检索 NASA ADS
|
||||
if source == "all" || source == "ads" {
|
||||
if !state.config.ads_api_key.is_empty() {
|
||||
match state.ads.search(q, start, rows, sort).await {
|
||||
Ok(docs) => {
|
||||
for doc in docs {
|
||||
let paper = convert_ads_doc_to_standard(&doc);
|
||||
|
||||
// 入库 SQLite
|
||||
if let Err(e) = save_paper_to_db(&state.db, &paper).await {
|
||||
warn!("保存 ADS 文献至数据库失败: {}", e);
|
||||
}
|
||||
|
||||
// 保存引用/参考文献关联拓扑
|
||||
if let Some(refs) = doc.reference {
|
||||
for ref_bib in refs {
|
||||
let _ = sqlx::query("INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)")
|
||||
.bind(&paper.bibcode)
|
||||
.bind(&ref_bib)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
if let Some(cits) = doc.citation {
|
||||
for cit_bib in cits {
|
||||
let _ = sqlx::query("INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)")
|
||||
.bind(&cit_bib)
|
||||
.bind(&paper.bibcode)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
results.push(paper);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("ADS 检索执行失败: {}", e);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
warn!("ADS_API_KEY 未配置，跳过 ADS 检索。");
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 检索 arXiv
|
||||
if source == "all" || source == "arxiv" {
|
||||
match state.arxiv.search(q, start, rows, sort).await {
|
||||
Ok(papers) => {
|
||||
for p in papers {
|
||||
let paper = convert_arxiv_to_standard(&p);
|
||||
|
||||
// 入库 SQLite (使用 arXiv ID 暂作主键以作记录)
|
||||
if let Err(e) = save_paper_to_db(&state.db, &paper).await {
|
||||
warn!("保存 arXiv 文献至数据库失败: {}", e);
|
||||
}
|
||||
|
||||
results.push(paper);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("arXiv 检索执行失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 对两端获取的数据进行去重合并，增加对相同 arxiv_id 的判断
|
||||
let mut unique_results: Vec<StandardPaper> = Vec::new();
|
||||
for r in results {
|
||||
if !unique_results.iter().any(|u| u.bibcode == r.bibcode || (!u.doi.is_empty() && u.doi == r.doi) || (!u.arxiv_id.is_empty() && u.arxiv_id == r.arxiv_id)) {
|
||||
let mut final_paper = r.clone();
|
||||
// 如果本地数据库存在该文献，直接从数据库读取标准元数据（包括 is_downloaded, has_markdown, pdf_error, html_error 等）
|
||||
if let Ok(db_paper) = get_paper_from_db(&state.db, &state.config.library_dir, &r.bibcode).await {
|
||||
final_paper = db_paper;
|
||||
}
|
||||
unique_results.push(final_paper);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(unique_results)
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user