refactor: 服务层抽象下沉、异步锁全栈迁移、客户端韧性加固与移动端适配

- 服务层拆分:删除 api/helpers.rs,新增 citation/note/session/pipeline/paper/vision 独立服务模块
  - Agent 工具精简:paper_content+paper_outline 合并为 paper.rs,图片分析逻辑下沉至 services/vision
  - 并发模型升级:std::sync::{Mutex,RwLock} → tokio::sync::{Mutex,RwLock},消除 async
  上下文中的阻塞风险
  - 客户端加固:HTTP 客户端统一超时配置、ADS 429 / arXiv 503 自动重试、构造函数返回 Result
  - 启动安全:全局 panic hook 日志化、空密码拒绝启动、向量表维度不匹配需显式确认
  - CLI 扩展:构建完整 AppState 复用服务层,新增 Content/Outline/Citations/Search/Process 子命令
  - 前端:移动端汉堡菜单、侧栏滑出面板、引用星系触屏手势(单指拖拽/双指缩放)
This commit is contained in:
Asfmq 2026-06-30 19:26:01 +08:00
parent c5fd5b0d66
commit f885c0a4a8
90 changed files with 5184 additions and 3916 deletions

View File

@ -44,6 +44,8 @@ LLM_MODEL=gpt-4o-mini
# EMBEDDING_MODEL=text-embedding-3-small
# 向量维度(需与所选模型输出维度一致,默认 1536
# EMBEDDING_DIM=1536
# 维度变更确认:修改 EMBEDDING_DIM 后需设置此变量为 1 才能启动(防止误改导致数据丢失)
# EMBEDDING_DIM_FORCE_REBUILD=1
# ─────────────────────────────────────────────────────────────────────────────
# Chunker 结构化切片参数(影响向量化粒度)
@ -85,8 +87,8 @@ LIBRARY_DIR=./library
SKILLS_DIR=./skills
# 后端服务监听端口
PORT=8000
# 认证密码
ADMIN_PASSWORD=fmq
# 认证密码(必填,启动时未设置将拒绝启动)
ADMIN_PASSWORD=your_strong_password_here
# ─────────────────────────────────────────────────────────────────────────────
# 5. 日志配置

View File

@ -21,7 +21,7 @@ path = "src/bin/cli.rs"
[dependencies]
tokio = { version = "1", features = ["full"] }
axum = { version = "0.7", features = ["macros", "multipart"] }
tower-http = { version = "0.5", features = ["cors", "fs", "trace"] }
tower-http = { version = "0.5", features = ["cors", "fs", "trace", "set-header"] }
sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "sqlite", "chrono", "json"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"

View File

@ -1,7 +1,7 @@
// dashboard/src/App.tsx
import { useState, useEffect, useCallback, useRef } from 'react';
import axios from 'axios';
import { Loader, BookOpen, GitFork, Lock } from 'lucide-react';
import { Loader, BookOpen, GitFork, Lock, Menu } from 'lucide-react';
import { Sidebar } from './components/layout/Sidebar';
import { SearchPanel } from './pages/SearchPanel';
import { LibraryPanel } from './pages/LibraryPanel';
@ -24,6 +24,9 @@ import { useNotes } from './hooks/useNotes';
import { useCitations } from './hooks/useCitations';
export default function App() {
// 移动端菜单显示状态
const [isMobileSidebarOpen, setIsMobileSidebarOpen] = useState(false);
// 1. 全局自定义 Dialog 弹窗管理 (Alert / Confirm)
const [dialog, setDialog] = useState<{
type: 'alert' | 'confirm';
@ -318,11 +321,41 @@ export default function App() {
selectedPaper={library.selectedPaper}
loadCitations={citations.loadCitations}
onLogout={auth.handleLogout}
isOpen={isMobileSidebarOpen}
onClose={() => setIsMobileSidebarOpen(false)}
/>
{/* 主工作区 */}
<main className="flex-1 flex flex-col overflow-hidden relative">
<div className="flex-1 overflow-y-auto p-4 sm:p-6 md:p-8 relative z-10 w-full flex flex-col">
{/* 移动端顶部 Header */}
<header className="lg:hidden flex items-center justify-between px-4 py-3 bg-white border-b border-slate-200 select-none shrink-0 z-20">
<div className="flex items-center gap-3">
<button
type="button"
onClick={() => setIsMobileSidebarOpen(true)}
className="p-1.5 rounded-lg border border-slate-200 hover:bg-slate-55 text-slate-600 transition-all cursor-pointer flex items-center justify-center"
title="打开菜单"
>
<Menu className="w-5 h-5" />
</button>
<div className="flex items-center gap-2">
<div className="w-6 h-6">
<Logo gradientId="mobileHeaderStarGrad" />
</div>
<span className="text-xs font-bold text-slate-800 tracking-wider">AstroResearch</span>
</div>
</div>
<span className="text-[10px] font-bold px-2.5 py-1 rounded bg-slate-100 text-slate-600 font-sans tracking-wide uppercase">
{activeTab === 'search' && '统一检索'}
{activeTab === 'library' && '馆藏管理'}
{activeTab === 'reader' && '双语阅读'}
{activeTab === 'citation' && '引用星系'}
{activeTab === 'sync' && '批量任务'}
{activeTab === 'agent' && '智能科研'}
</span>
</header>
<div className="flex-1 overflow-y-auto p-4 sm:p-6 md:p-8 relative w-full flex flex-col">
<div className={`w-full flex-1 flex flex-col min-h-0 ${
(activeTab === 'reader' || activeTab === 'citation' || activeTab === 'agent') ? 'max-w-none' : 'max-w-7xl mx-auto'
}`}>

View File

@ -298,6 +298,29 @@ export function CitationGalaxyCanvas({ networks, activeNetwork, nodeLimit, onNod
render();
let initialTouchDistance = 0;
let initialScale = 1.0;
const findNodeAtCoordinate = (clientX: number, clientY: number) => {
const rect = canvas.getBoundingClientRect();
const touchX = clientX - rect.left;
const touchY = clientY - rect.top;
const cx = rect.width / 2;
const cy = rect.height / 2;
const gx = (touchX - cx - offsetX) / scale + cx;
const gy = (touchY - cy - offsetY) / scale + cy;
for (const node of nodes) {
const dx = node.x - gx;
const dy = node.y - gy;
const dist = Math.sqrt(dx * dx + dy * dy);
if (dist < node.radius + 12) { // 移动端稍微加大触碰敏感区
return node;
}
}
return null;
};
const handleMouseDown = (e: MouseEvent) => {
isDragging = true;
dragStartX = e.clientX - offsetX;
@ -379,12 +402,68 @@ export function CitationGalaxyCanvas({ networks, activeNetwork, nodeLimit, onNod
offsetY = mouseY - cy - (gy - cy) * scale;
};
// 触屏手势处理
const handleTouchStart = (e: TouchEvent) => {
if (e.touches.length === 1) {
isDragging = true;
hasDragged = false;
dragStartX = e.touches[0].clientX - offsetX;
dragStartY = e.touches[0].clientY - offsetY;
// 触碰选中节点
hoveredNode = findNodeAtCoordinate(e.touches[0].clientX, e.touches[0].clientY);
} else if (e.touches.length === 2) {
isDragging = false;
const dx = e.touches[0].clientX - e.touches[1].clientX;
const dy = e.touches[0].clientY - e.touches[1].clientY;
initialTouchDistance = Math.sqrt(dx * dx + dy * dy) || 1;
initialScale = scale;
}
};
const handleTouchMove = (e: TouchEvent) => {
if (isDragging && e.touches.length === 1) {
const dx = e.touches[0].clientX - dragStartX;
const dy = e.touches[0].clientY - dragStartY;
if (Math.sqrt((dx - offsetX) ** 2 + (dy - offsetY) ** 2) > 3) {
hasDragged = true;
}
offsetX = dx;
offsetY = dy;
e.preventDefault(); // 阻止滚动
} else if (e.touches.length === 2) {
const dx = e.touches[0].clientX - e.touches[1].clientX;
const dy = e.touches[0].clientY - e.touches[1].clientY;
const dist = Math.sqrt(dx * dx + dy * dy) || 1;
// 双指开合缩放
const newScale = initialScale * (dist / initialTouchDistance);
scale = Math.max(0.15, Math.min(5.0, newScale));
e.preventDefault(); // 阻止默认双指缩放
}
};
const handleTouchEnd = () => {
isDragging = false;
if (!hasDragged && hoveredNode) {
if (hoveredNode.id !== activeNetwork.bibcode) {
onNodeClick(hoveredNode.id);
}
}
hoveredNode = null;
};
canvas.addEventListener('mousedown', handleMouseDown);
canvas.addEventListener('mousemove', handleMouseMove);
canvas.addEventListener('mouseup', handleMouseUp);
canvas.addEventListener('mouseleave', handleMouseLeave);
canvas.addEventListener('click', handleCanvasClick);
canvas.addEventListener('wheel', handleWheel, { passive: false });
// 注册触控事件
canvas.addEventListener('touchstart', handleTouchStart, { passive: false });
canvas.addEventListener('touchmove', handleTouchMove, { passive: false });
canvas.addEventListener('touchend', handleTouchEnd);
return () => {
cancelAnimationFrame(animationFrameId);
@ -394,6 +473,11 @@ export function CitationGalaxyCanvas({ networks, activeNetwork, nodeLimit, onNod
canvas.removeEventListener('mouseleave', handleMouseLeave);
canvas.removeEventListener('click', handleCanvasClick);
canvas.removeEventListener('wheel', handleWheel);
// 注销触控事件
canvas.removeEventListener('touchstart', handleTouchStart);
canvas.removeEventListener('touchmove', handleTouchMove);
canvas.removeEventListener('touchend', handleTouchEnd);
};
}, [networks, activeNetwork, onNodeClick]);

View File

@ -190,7 +190,7 @@ export function AgentMessageList({
: '探索性科研研讨'}
</h3>
<p className="text-[10px] text-slate-400 font-semibold mt-0.5">
(ReAct Thought Action Observation)
</p>
</div>
</div>

View File

@ -37,173 +37,206 @@ export function AgentSessionSidebar({
onNewSession,
onDeleteSession,
}: AgentSessionSidebarProps) {
return (
<div className={`transition-all duration-300 ease-in-out ${collapsed ? 'w-0 overflow-hidden opacity-0 border-r-0' : 'w-64 border-r border-slate-200'} bg-slate-50 flex flex-col justify-between shrink-0 select-none`}>
<div className="flex flex-col min-h-0 flex-1">
<div className="p-4 border-b border-slate-200 bg-white flex items-center justify-between shrink-0">
<span className="text-xs font-extrabold text-slate-800 tracking-wider flex items-center gap-1.5">
<Clock className="w-3.5 h-3.5 text-blueprint" />
{!collapsed && <span></span>}
</span>
<div className="flex items-center gap-1 shrink-0">
<button
onClick={onNewSession}
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-600 hover:text-slate-800 transition-all cursor-pointer shadow-2xs"
title="新建会话"
>
<Plus className="w-3.5 h-3.5" />
</button>
<button
onClick={() => onToggleCollapse(true)}
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-500 hover:text-slate-700 transition-all cursor-pointer"
title="收起侧栏"
>
<PanelLeftClose className="w-3.5 h-3.5" />
</button>
</div>
</div>
const handleSelectSession = (id: string | null) => {
setCurrentSessionId(id);
if (typeof window !== 'undefined' && window.innerWidth < 1024) {
onToggleCollapse(true);
}
};
{/* 搜索框区域 */}
<div className="px-3 py-2 border-b border-slate-200/60 bg-white flex flex-col gap-1.5 shrink-0">
<div className="relative">
<input
type="text"
value={searchQuery}
onChange={e => setSearchQuery(e.target.value)}
placeholder="搜索历史会话或消息内容..."
className="w-full pl-8 pr-7 py-1.5 rounded-md bg-slate-50 border border-slate-200 text-slate-800 placeholder-slate-400 focus:outline-none focus:border-blueprint focus:bg-white transition-all text-[11px] font-medium"
/>
<div className="absolute left-2.5 top-2.5 text-slate-400">
<Search className="w-3.5 h-3.5 text-slate-400" />
</div>
{searchQuery && (
const handleCreateNewSession = () => {
onNewSession();
if (typeof window !== 'undefined' && window.innerWidth < 1024) {
onToggleCollapse(true);
}
};
return (
<>
{/* 移动端背景遮罩层 */}
{!collapsed && (
<button
type="button"
onClick={() => onToggleCollapse(true)}
className="fixed inset-0 bg-slate-900/30 backdrop-blur-xs z-35 lg:hidden cursor-pointer w-full h-full border-none outline-none"
aria-label="关闭侧栏"
/>
)}
<div
className={`bg-slate-50 flex flex-col justify-between shrink-0 select-none transition-all duration-300 ease-in-out fixed inset-y-0 left-0 z-40 lg:relative lg:translate-x-0 ${
collapsed
? '-translate-x-full lg:translate-x-0 lg:w-0 lg:overflow-hidden lg:opacity-0 lg:border-r-0'
: 'translate-x-0 w-64 border-r border-slate-200 lg:w-64'
}`}
>
<div className="flex flex-col min-h-0 flex-1">
<div className="p-4 border-b border-slate-200 bg-white flex items-center justify-between shrink-0">
<span className="text-xs font-extrabold text-slate-800 tracking-wider flex items-center gap-1.5">
<Clock className="w-3.5 h-3.5 text-blueprint" />
<span></span>
</span>
<div className="flex items-center gap-1 shrink-0">
<button
onClick={() => setSearchQuery('')}
className="absolute right-2 top-2 p-0.5 rounded-full hover:bg-slate-200 text-slate-400 hover:text-slate-600 transition-colors cursor-pointer"
onClick={handleCreateNewSession}
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-650 hover:text-slate-800 transition-all cursor-pointer shadow-2xs"
title="新建会话"
>
<X className="w-3 h-3" />
<Plus className="w-3.5 h-3.5" />
</button>
<button
onClick={() => onToggleCollapse(true)}
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-555 hover:text-slate-700 transition-all cursor-pointer"
title="收起侧栏"
>
<PanelLeftClose className="w-3.5 h-3.5" />
</button>
</div>
</div>
{/* 搜索框区域 */}
<div className="px-3 py-2 border-b border-slate-200/60 bg-white flex flex-col gap-1.5 shrink-0">
<div className="relative">
<input
type="text"
value={searchQuery}
onChange={e => setSearchQuery(e.target.value)}
placeholder="搜索历史会话或消息内容..."
className="w-full pl-8 pr-7 py-1.5 rounded-md bg-slate-50 border border-slate-200 text-slate-800 placeholder-slate-400 focus:outline-none focus:border-blueprint focus:bg-white transition-all text-[11px] font-medium"
/>
<div className="absolute left-2.5 top-2.5 text-slate-400">
<Search className="w-3.5 h-3.5 text-slate-400" />
</div>
{searchQuery && (
<button
onClick={() => setSearchQuery('')}
className="absolute right-2 top-2 p-0.5 rounded-full hover:bg-slate-200 text-slate-400 hover:text-slate-600 transition-colors cursor-pointer"
>
<X className="w-3 h-3" />
</button>
)}
</div>
{currentSessionId && (
<label className="flex items-center gap-1.5 text-[10px] text-slate-500 cursor-pointer font-medium select-none ml-0.5">
<input
type="checkbox"
checked={searchScopeOnlyCurrent}
onChange={e => setSearchScopeOnlyCurrent(e.target.checked)}
className="rounded text-blueprint border-slate-300 focus:ring-blueprint w-3 h-3 cursor-pointer"
/>
<span></span>
</label>
)}
</div>
{currentSessionId && (
<label className="flex items-center gap-1.5 text-[10px] text-slate-500 cursor-pointer font-medium select-none ml-0.5">
<input
type="checkbox"
checked={searchScopeOnlyCurrent}
onChange={e => setSearchScopeOnlyCurrent(e.target.checked)}
className="rounded text-blueprint border-slate-300 focus:ring-blueprint w-3 h-3 cursor-pointer"
/>
<span></span>
</label>
)}
</div>
<div className="flex-1 overflow-y-auto p-2.5 space-y-1.5 scrollbar-thin">
{searchQuery.trim() ? (
loadingSearch ? (
<div className="flex items-center justify-center p-8 text-slate-400 text-xs gap-2">
<Loader className="w-3.5 h-3.5 animate-spin text-blueprint" />
<span>...</span>
</div>
) : searchResults.length === 0 ? (
<div className="text-center py-12 text-slate-400 text-[11px] italic">
</div>
) : (
searchResults.map((result, idx) => {
const isActive = result.session_id === currentSessionId;
const isMessage = result.result_type.startsWith('message/');
const role = isMessage ? result.result_type.split('/')[1] : null;
<div className="flex-1 overflow-y-auto p-2.5 space-y-1.5 scrollbar-thin">
{searchQuery.trim() ? (
loadingSearch ? (
<div className="flex items-center justify-center p-8 text-slate-400 text-xs gap-2">
<Loader className="w-3.5 h-3.5 animate-spin text-blueprint" />
<span>...</span>
</div>
) : searchResults.length === 0 ? (
<div className="text-center py-12 text-slate-400 text-[11px] italic">
</div>
) : (
searchResults.map((result, idx) => {
const isActive = result.session_id === currentSessionId;
const isMessage = result.result_type.startsWith('message/');
const role = isMessage ? result.result_type.split('/')[1] : null;
return (
<button
key={`${result.session_id}-${idx}`}
onClick={() => setCurrentSessionId(result.session_id)}
className={`w-full text-left p-2.5 rounded-md border transition-all duration-200 flex flex-col gap-1 cursor-pointer ${
isActive
? 'bg-blueprint/5 border-blueprint text-blueprint font-bold shadow-2xs'
: 'border-transparent bg-white hover:bg-slate-100 text-slate-650 shadow-3xs'
}`}
>
<div className="flex items-center justify-between w-full">
<span className="text-[10px] font-bold text-slate-400 flex items-center gap-1">
{isMessage ? (
<>
<MessageSquare className="w-3 h-3 text-slate-400" />
<span>{role === 'user' ? '提问' : '解答'}</span>
</>
) : (
<>
<BookOpen className="w-3 h-3 text-blueprint" />
<span></span>
</>
)}
</span>
{result.created_at && (
<span className="text-[8px] text-slate-400 font-mono">
{result.created_at.split(' ')[0]}
</span>
)}
</div>
<span className="text-xs leading-snug line-clamp-1 block text-left font-bold text-slate-800">
{result.title || '无标题会话'}
</span>
<p
className="text-[10px] text-slate-500 font-medium leading-normal block text-left line-clamp-3 bg-slate-50/50 p-1.5 rounded border border-slate-100/50"
dangerouslySetInnerHTML={{ __html: result.snippet }}
/>
</button>
);
})
)
) : (
// 正常的会话列表渲染
loadingSessions ? (
<div className="flex items-center justify-center p-8 text-slate-400 text-xs gap-2">
<Loader className="w-3.5 h-3.5 animate-spin text-blueprint" />
<span>...</span>
</div>
) : sessions.length === 0 ? (
<div className="text-center py-12 text-slate-400 text-[11px] italic">
</div>
) : (
sessions.map(session => {
const isActive = session.session_id === currentSessionId;
return (
<button
key={session.session_id}
onClick={() => setCurrentSessionId(session.session_id)}
className={`w-full text-left p-3 rounded-md border transition-all duration-200 group flex items-start justify-between gap-2 cursor-pointer ${
isActive
? 'bg-blueprint/5 border-blueprint text-blueprint font-bold shadow-2xs'
: 'border-transparent bg-transparent hover:bg-slate-100 text-slate-655'
}`}
>
<div className="flex-1 min-w-0 flex flex-col gap-1 text-left">
<span className="text-xs leading-snug line-clamp-2 block text-left">
{session.title || '无标题会话'}
</span>
<span className="text-[9px] text-slate-400 font-semibold block text-left">
{session.turn_count} {session.updated_at.split(' ')[0]}
</span>
</div>
return (
<button
onClick={(e) => onDeleteSession(session.session_id, e)}
className="text-slate-400 hover:text-red-655 p-1 rounded-md opacity-0 group-hover:opacity-100 hover:bg-white border border-transparent hover:border-slate-200 transition-all cursor-pointer shrink-0"
title="删除此会话"
key={`${result.session_id}-${idx}`}
onClick={() => handleSelectSession(result.session_id)}
className={`w-full text-left p-2.5 rounded-md border transition-all duration-200 flex flex-col gap-1 cursor-pointer ${
isActive
? 'bg-blueprint/5 border-blueprint text-blueprint font-bold shadow-2xs'
: 'border-transparent bg-white hover:bg-slate-100 text-slate-650 shadow-3xs'
}`}
>
<Trash2 className="w-3 h-3" />
<div className="flex items-center justify-between w-full">
<span className="text-[10px] font-bold text-slate-400 flex items-center gap-1">
{isMessage ? (
<>
<MessageSquare className="w-3 h-3 text-slate-400" />
<span>{role === 'user' ? '提问' : '解答'}</span>
</>
) : (
<>
<BookOpen className="w-3 h-3 text-blueprint" />
<span></span>
</>
)}
</span>
{result.created_at && (
<span className="text-[8px] text-slate-400 font-mono">
{result.created_at.split(' ')[0]}
</span>
)}
</div>
<span className="text-xs leading-snug line-clamp-1 block text-left font-bold text-slate-800">
{result.title || '无标题会话'}
</span>
<p
className="text-[10px] text-slate-500 font-medium leading-normal block text-left line-clamp-3 bg-slate-50/50 p-1.5 rounded border border-slate-100/50"
dangerouslySetInnerHTML={{ __html: result.snippet }}
/>
</button>
</button>
);
})
)
)}
);
})
)
) : (
// 正常的会话列表渲染
loadingSessions ? (
<div className="flex items-center justify-center p-8 text-slate-400 text-xs gap-2">
<Loader className="w-3.5 h-3.5 animate-spin text-blueprint" />
<span>...</span>
</div>
) : sessions.length === 0 ? (
<div className="text-center py-12 text-slate-400 text-[11px] italic">
</div>
) : (
sessions.map(session => {
const isActive = session.session_id === currentSessionId;
return (
<button
key={session.session_id}
onClick={() => handleSelectSession(session.session_id)}
className={`w-full text-left p-3 rounded-md border transition-all duration-200 group flex items-start justify-between gap-2 cursor-pointer ${
isActive
? 'bg-blueprint/5 border-blueprint text-blueprint font-bold shadow-2xs'
: 'border-transparent bg-transparent hover:bg-slate-100 text-slate-655'
}`}
>
<div className="flex-1 min-w-0 flex flex-col gap-1 text-left">
<span className="text-xs leading-snug line-clamp-2 block text-left">
{session.title || '无标题会话'}
</span>
<span className="text-[9px] text-slate-400 font-semibold block text-left">
{session.turn_count} {session.updated_at.split(' ')[0]}
</span>
</div>
<button
onClick={(e) => onDeleteSession(session.session_id, e)}
className="text-slate-400 hover:text-red-655 p-1 rounded-md opacity-0 group-hover:opacity-100 hover:bg-white border border-transparent hover:border-slate-200 transition-all cursor-pointer shrink-0"
title="删除此会话"
>
<Trash2 className="w-3 h-3" />
</button>
</button>
);
})
)
)}
</div>
</div>
</div>
</div>
</>
);
}

View File

@ -75,7 +75,7 @@ export function ToolCallCard({
<div className="flex items-center justify-between">
<div className="flex items-center gap-1.5 text-[10px] font-extrabold text-emerald-700 tracking-wider uppercase">
<Eye className="w-3.5 h-3.5" />
<span> (Observation)</span>
<span></span>
{result!.isError && (
<span className="px-1.5 py-0.2 rounded bg-rose-50 text-rose-700 border border-rose-100 text-[9px]">

View File

@ -12,207 +12,225 @@ interface SidebarProps {
selectedPaper: StandardPaper | null;
loadCitations: (bibcode: string) => void;
onLogout: () => void;
isOpen?: boolean;
onClose?: () => void;
}
export function Sidebar({ activeTab, setActiveTab, selectedPaper, loadCitations, onLogout }: SidebarProps) {
export function Sidebar({ activeTab, setActiveTab, selectedPaper, loadCitations, onLogout, isOpen = false, onClose }: SidebarProps) {
const [isCollapsed, setIsCollapsed] = useState(false);
const effectiveCollapsed = isCollapsed && !isOpen;
return (
<aside
className={`bg-slate-50 border-r border-slate-200 flex flex-col justify-between py-6 z-20 select-none transition-all duration-355 cubic-bezier(0.4, 0, 0.2, 1) ${
isCollapsed ? 'w-16 px-2' : 'w-64 px-4'
}`}
>
<div>
{/* 系统LOGO与折叠控制区 */}
<div className={`mb-8 flex items-center transition-all duration-300 ${
isCollapsed ? 'justify-center px-0' : 'justify-between px-3'
}`}>
{/* Logo & 标题文字 */}
<div className="flex items-center gap-3 min-w-0">
{/* Logo按钮 (只在折叠状态下可点击展开) */}
<button
type="button"
disabled={!isCollapsed}
onClick={() => setIsCollapsed(false)}
className={`flex items-center justify-center shrink-0 rounded-lg transition-all duration-300 ${
isCollapsed
? 'w-11 h-11 bg-white hover:bg-slate-50 border border-slate-200 cursor-pointer shadow-xs hover:shadow-sm'
: 'w-9 h-9 bg-transparent border border-transparent cursor-default'
}`}
title={isCollapsed ? "展开导航" : undefined}
>
<div className={`transition-all duration-300 flex items-center justify-center ${isCollapsed ? 'w-8 h-8' : 'w-9 h-9'}`}>
<Logo gradientId="sidebarStarGrad" />
</div>
</button>
{/* 系统说明文字 */}
<div
className={`flex flex-col transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 scale-95 -translate-x-2 pointer-events-none select-none overflow-hidden h-0'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0'
}`}
>
<h1 className="text-sm font-bold text-slate-800 tracking-wider whitespace-nowrap">AstroResearch</h1>
<span className="text-[11px] text-slate-500 block font-medium font-sans whitespace-nowrap"></span>
</div>
</div>
{/* 折叠控制按钮 (圆润悬浮效果) */}
<button
type="button"
onClick={() => setIsCollapsed(true)}
className={`p-1.5 rounded-full bg-slate-100 hover:bg-slate-200 border border-slate-200 text-slate-500 hover:text-slate-800 transition-all duration-300 cursor-pointer shadow-2xs shrink-0 flex items-center justify-center hover:scale-105 active:scale-95 ${
isCollapsed
? 'opacity-0 scale-75 pointer-events-none w-0 h-0 p-0 border-0 overflow-hidden'
: 'opacity-100 scale-100'
}`}
title="收起导航"
>
<ChevronLeft className="w-3.5 h-3.5" />
</button>
</div>
{/* 导航菜单列表 */}
<nav className="space-y-1">
{[
{ id: 'search' as TabId, label: '统一检索', icon: Search },
{ id: 'library' as TabId, label: '馆藏管理', icon: Library },
{ id: 'reader' as TabId, label: '双语阅读', icon: BookOpen },
{ id: 'citation' as TabId, label: '引用星系', icon: GitFork },
{ id: 'sync' as TabId, label: '批量任务', icon: RefreshCw },
{ id: 'agent' as TabId, label: '智能科研', icon: Sparkles },
].map((tab) => {
const Icon = tab.icon;
const isActive = activeTab === tab.id;
return (
<button
key={tab.id}
title={isCollapsed ? tab.label : undefined}
onClick={() => {
setActiveTab(tab.id);
if (tab.id === 'citation' && selectedPaper) {
loadCitations(selectedPaper.bibcode);
}
}}
className={`w-full flex items-center rounded-lg text-xs font-semibold tracking-wider transition-all duration-300 border ${
isCollapsed ? 'px-2 py-2.5 justify-center' : 'px-3 py-2.5'
} ${
isActive
? 'bg-slate-100 border-slate-200 text-slate-850 shadow-xs'
: 'border-transparent text-slate-650 hover:bg-slate-100 hover:text-slate-800'
}`}
>
<Icon className={`w-4 h-4 shrink-0 transition-colors duration-300 ${isActive ? 'text-slate-700' : 'text-slate-500'}`} />
<span
className={`truncate transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 pointer-events-none select-none overflow-hidden scale-90 -translate-x-2'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0 ml-3'
}`}
>
{tab.label}
</span>
</button>
);
})}
</nav>
</div>
{/* 底部当前选定文献提示 (平滑动画版本) 与 退出登录 */}
<div className="space-y-3">
{selectedPaper ? (
<div
className={`relative overflow-hidden transition-all duration-300 border rounded-lg ${
isCollapsed
? 'p-0 border-transparent bg-transparent flex justify-center'
: 'p-3.5 border-slate-200 bg-slate-100/50'
}`}
title={isCollapsed ? `当前选定文献: ${selectedPaper.title}` : undefined}
>
{/* 折叠下的微型图书图标 */}
<div
className={`transition-all duration-300 flex items-center justify-center rounded-lg border border-slate-200 bg-slate-50 text-slate-650 shrink-0 shadow-xs ${
isCollapsed
? 'w-9 h-9 opacity-100 scale-100'
: 'w-0 h-0 opacity-0 scale-75 overflow-hidden'
}`}
>
<BookOpen className="w-4 h-4" />
</div>
{/* 展开下的完整详情 */}
<div
className={`transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 max-h-0 pointer-events-none select-none overflow-hidden'
: 'opacity-100 max-w-[200px] max-h-40'
}`}
>
<span className="text-[9px] font-bold text-slate-500 tracking-widest block mb-1"></span>
<h4 className="text-xs text-slate-800 font-bold line-clamp-2 mb-2 leading-relaxed">{selectedPaper.title}</h4>
<div className="flex items-center justify-between text-[10px] font-medium text-slate-500 gap-2">
<span className="shrink-0">: {selectedPaper.year}</span>
<span className="truncate max-w-[90px] font-mono">{selectedPaper.bibcode}</span>
</div>
</div>
</div>
) : (
<div
className={`relative overflow-hidden transition-all duration-300 border rounded-lg ${
isCollapsed
? 'p-0 border-transparent bg-transparent flex justify-center'
: 'p-3 border-slate-200 bg-slate-100/30'
}`}
title={isCollapsed ? "未选定研究目标" : undefined}
>
{/* 折叠下的微型馆藏图标 */}
<div
className={`transition-all duration-300 flex items-center justify-center rounded-lg border border-slate-200 bg-slate-100/30 text-slate-400 shrink-0 ${
isCollapsed
? 'w-9 h-9 opacity-100 scale-100'
: 'w-0 h-0 opacity-0 scale-75 overflow-hidden'
}`}
>
<Library className="w-4 h-4" />
</div>
{/* 展开下的提示文字 */}
<div
className={`text-center transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 max-h-0 pointer-events-none select-none overflow-hidden'
: 'opacity-100 max-w-[200px] max-h-12'
}`}
>
<span className="text-[10px] text-slate-400 font-medium tracking-wide"></span>
</div>
</div>
)}
{/* 退出登录按钮 */}
<>
{/* 移动端背景遮罩层 */}
{isOpen && (
<button
type="button"
onClick={onLogout}
className={`w-full flex items-center rounded-lg text-xs font-semibold tracking-wider transition-all duration-300 border border-transparent text-slate-500 hover:bg-rose-50/50 hover:text-rose-600 cursor-pointer group ${
isCollapsed ? 'px-2 py-2.5 justify-center' : 'px-3 py-2.5'
}`}
title={isCollapsed ? "退出登录" : undefined}
>
<LogOut className="w-4 h-4 shrink-0 text-slate-400 group-hover:text-rose-500 transition-colors" />
<span
className={`truncate transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 pointer-events-none select-none overflow-hidden scale-90 -translate-x-2'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0 ml-3'
onClick={onClose}
className="fixed inset-0 bg-slate-900/30 backdrop-blur-xs z-35 lg:hidden cursor-pointer w-full h-full border-none outline-none"
aria-label="关闭导航"
/>
)}
<aside
className={`bg-slate-50 border-r border-slate-200 flex flex-col justify-between py-6 z-40 select-none transition-all duration-355 cubic-bezier(0.4, 0, 0.2, 1) fixed inset-y-0 left-0 lg:relative lg:translate-x-0 ${
isOpen ? 'translate-x-0' : '-translate-x-full lg:translate-x-0'
} ${
isOpen ? 'w-64 px-4' : (effectiveCollapsed ? 'w-16 px-2' : 'w-64 px-4')
}`}
>
<div>
{/* 系统LOGO与折叠控制区 */}
<div className={`mb-8 flex items-center transition-all duration-300 ${
effectiveCollapsed ? 'justify-center px-0' : 'justify-between px-3'
}`}>
{/* Logo & 标题文字 */}
<div className="flex items-center gap-3 min-w-0">
{/* Logo按钮 (只在折叠状态下可点击展开) */}
<button
type="button"
disabled={!effectiveCollapsed}
onClick={() => setIsCollapsed(false)}
className={`flex items-center justify-center shrink-0 rounded-lg transition-all duration-300 ${
effectiveCollapsed
? 'w-11 h-11 bg-white hover:bg-slate-55 border border-slate-200 cursor-pointer shadow-xs hover:shadow-sm'
: 'w-9 h-9 bg-transparent border border-transparent cursor-default'
}`}
title={effectiveCollapsed ? "展开导航" : undefined}
>
<div className={`transition-all duration-300 flex items-center justify-center ${effectiveCollapsed ? 'w-8 h-8' : 'w-9 h-9'}`}>
<Logo gradientId="sidebarStarGrad" />
</div>
</button>
{/* 系统说明文字 */}
<div
className={`flex flex-col transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 scale-95 -translate-x-2 pointer-events-none select-none overflow-hidden h-0'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0'
}`}
>
<h1 className="text-sm font-bold text-slate-800 tracking-wider whitespace-nowrap">AstroResearch</h1>
<span className="text-[11px] text-slate-500 block font-medium font-sans whitespace-nowrap"></span>
</div>
</div>
{/* 折叠控制按钮 (圆润悬浮效果,仅在桌面端显示) */}
<button
type="button"
onClick={() => setIsCollapsed(true)}
className={`p-1.5 rounded-full bg-slate-100 hover:bg-slate-200 border border-slate-200 text-slate-505 hover:text-slate-800 transition-all duration-300 cursor-pointer shadow-2xs shrink-0 items-center justify-center hover:scale-105 active:scale-95 hidden lg:flex ${
effectiveCollapsed
? 'opacity-0 scale-75 pointer-events-none w-0 h-0 p-0 border-0 overflow-hidden'
: 'opacity-100 scale-100'
}`}
title="收起导航"
>
<ChevronLeft className="w-3.5 h-3.5" />
</button>
</div>
{/* 导航菜单列表 */}
<nav className="space-y-1">
{[
{ id: 'search' as TabId, label: '统一检索', icon: Search },
{ id: 'library' as TabId, label: '馆藏管理', icon: Library },
{ id: 'reader' as TabId, label: '双语阅读', icon: BookOpen },
{ id: 'citation' as TabId, label: '引用星系', icon: GitFork },
{ id: 'sync' as TabId, label: '批量任务', icon: RefreshCw },
{ id: 'agent' as TabId, label: '智能科研', icon: Sparkles },
].map((tab) => {
const Icon = tab.icon;
const isActive = activeTab === tab.id;
return (
<button
key={tab.id}
title={effectiveCollapsed ? tab.label : undefined}
onClick={() => {
setActiveTab(tab.id);
if (tab.id === 'citation' && selectedPaper) {
loadCitations(selectedPaper.bibcode);
}
if (onClose) onClose(); // 移动端切页时自动收起
}}
className={`w-full flex items-center rounded-lg text-xs font-semibold tracking-wider transition-all duration-300 border ${
effectiveCollapsed ? 'px-2 py-2.5 justify-center' : 'px-3 py-2.5'
} ${
isActive
? 'bg-slate-100 border-slate-200 text-slate-850 shadow-xs'
: 'border-transparent text-slate-650 hover:bg-slate-100 hover:text-slate-800'
}`}
>
<Icon className={`w-4 h-4 shrink-0 transition-colors duration-300 ${isActive ? 'text-slate-700' : 'text-slate-500'}`} />
<span
className={`truncate transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 pointer-events-none select-none overflow-hidden scale-90 -translate-x-2'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0 ml-3'
}`}
>
{tab.label}
</span>
</button>
);
})}
</nav>
</div>
{/* 底部当前选定文献提示 (平滑动画版本) 与 退出登录 */}
<div className="space-y-3">
{selectedPaper ? (
<div
className={`relative overflow-hidden transition-all duration-300 border rounded-lg ${
effectiveCollapsed
? 'p-0 border-transparent bg-transparent flex justify-center'
: 'p-3.5 border-slate-200 bg-slate-100/50'
}`}
title={effectiveCollapsed ? `当前选定文献: ${selectedPaper.title}` : undefined}
>
{/* 折叠下的微型图书图标 */}
<div
className={`transition-all duration-300 flex items-center justify-center rounded-lg border border-slate-200 bg-slate-55 text-slate-650 shrink-0 shadow-xs ${
effectiveCollapsed
? 'w-9 h-9 opacity-100 scale-100'
: 'w-0 h-0 opacity-0 scale-75 overflow-hidden'
}`}
>
<BookOpen className="w-4 h-4" />
</div>
{/* 展开下的完整详情 */}
<div
className={`transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 max-h-0 pointer-events-none select-none overflow-hidden'
: 'opacity-100 max-w-[200px] max-h-40'
}`}
>
<span className="text-[9px] font-bold text-slate-500 tracking-widest block mb-1"></span>
<h4 className="text-xs text-slate-800 font-bold line-clamp-2 mb-2 leading-relaxed">{selectedPaper.title}</h4>
<div className="flex items-center justify-between text-[10px] font-medium text-slate-500 gap-2">
<span className="shrink-0">: {selectedPaper.year}</span>
<span className="truncate max-w-[90px] font-mono">{selectedPaper.bibcode}</span>
</div>
</div>
</div>
) : (
<div
className={`relative overflow-hidden transition-all duration-300 border rounded-lg ${
effectiveCollapsed
? 'p-0 border-transparent bg-transparent flex justify-center'
: 'p-3 border-slate-200 bg-slate-100/30'
}`}
title={effectiveCollapsed ? "未选定研究目标" : undefined}
>
{/* 折叠下的微型馆藏图标 */}
<div
className={`transition-all duration-300 flex items-center justify-center rounded-lg border border-slate-200 bg-slate-100/30 text-slate-400 shrink-0 ${
effectiveCollapsed
? 'w-9 h-9 opacity-100 scale-100'
: 'w-0 h-0 opacity-0 scale-75 overflow-hidden'
}`}
>
<Library className="w-4 h-4" />
</div>
{/* 展开下的提示文字 */}
<div
className={`text-center transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 max-h-0 pointer-events-none select-none overflow-hidden'
: 'opacity-100 max-w-[200px] max-h-12'
}`}
>
<span className="text-[10px] text-slate-400 font-medium tracking-wide"></span>
</div>
</div>
)}
{/* 退出登录按钮 */}
<button
type="button"
onClick={onLogout}
className={`w-full flex items-center rounded-lg text-xs font-semibold tracking-wider transition-all duration-300 border border-transparent text-slate-500 hover:bg-rose-50/50 hover:text-rose-600 cursor-pointer group ${
effectiveCollapsed ? 'px-2 py-2.5 justify-center' : 'px-3 py-2.5'
}`}
title={effectiveCollapsed ? "退出登录" : undefined}
>
退
</span>
</button>
</div>
</aside>
<LogOut className="w-4 h-4 shrink-0 text-slate-400 group-hover:text-rose-505 transition-colors" />
<span
className={`truncate transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 pointer-events-none select-none overflow-hidden scale-90 -translate-x-2'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0 ml-3'
}`}
>
退
</span>
</button>
</div>
</aside>
</>
);
}

View File

@ -1,4 +1,5 @@
// dashboard/src/components/reader/BilingualViewer.tsx
import { useState, useEffect } from 'react';
import { createPortal } from 'react-dom';
import ReactMarkdown from 'react-markdown';
import remarkMath from 'remark-math';
@ -128,6 +129,28 @@ export function BilingualViewer({
hoverCardPos,
children,
}: BilingualViewerProps) {
// 监测屏幕宽度以判断是否在移动端/平板设备上(宽度 < 1024px
const [isMobileViewport, setIsMobileViewport] = useState(() => {
if (typeof window !== 'undefined') {
return window.innerWidth < 1024;
}
return false;
});
useEffect(() => {
const handleResize = () => {
setIsMobileViewport(window.innerWidth < 1024);
};
window.addEventListener('resize', handleResize);
return () => window.removeEventListener('resize', handleResize);
}, []);
const gridTemplateColumns = isMobileViewport
? (viewMode === 'bilingual' ? '1fr 1fr' : '1fr')
: (viewMode === 'bilingual'
? (showNotesPanel ? '1fr 1fr 380px' : '1fr 1fr')
: (showNotesPanel ? '1fr 380px' : '1fr'));
// 解析英文和中文段落的 Front Matter 头部元数据
const { metadata: engMeta, pureMarkdown: engPure } = parseMarkdownFrontMatter(englishText);
const { metadata: chnMeta, pureMarkdown: chnPure } = parseMarkdownFrontMatter(chineseText);
@ -150,11 +173,7 @@ export function BilingualViewer({
return (
<div
className="flex-1 grid gap-6 overflow-hidden w-full min-h-0"
style={{
gridTemplateColumns: viewMode === 'bilingual'
? (showNotesPanel ? '1fr 1fr 380px' : '1fr 1fr')
: (showNotesPanel ? '1fr 380px' : '1fr')
}}
style={{ gridTemplateColumns }}
>
{/* 英文正文视窗 (或双语对照) */}
{(viewMode === 'english' || viewMode === 'bilingual') && (

View File

@ -69,15 +69,24 @@ export function ReaderNotesSidebar({
if (!showNotesPanel) return null;
return (
<div className="console-panel rounded-lg border border-slate-200 bg-slate-50 flex flex-col overflow-hidden relative shadow-sm h-full">
<div className="px-4 py-3 border-b border-slate-200 flex items-center justify-between bg-white shrink-0">
<span className="text-xs font-bold text-slate-800">
{sidebarTab === 'notes' ? '观测记录手札' : '文献 AI 问答'}
</span>
<button onClick={() => setShowNotesPanel(false)} className="text-slate-400 hover:text-slate-600 transition-colors cursor-pointer">
<X className="w-4 h-4" />
</button>
</div>
<>
{/* 移动端遮罩层 */}
<button
type="button"
onClick={() => setShowNotesPanel(false)}
className="fixed inset-0 bg-slate-900/30 backdrop-blur-xs z-35 lg:hidden cursor-pointer w-full h-full border-none outline-none"
aria-label="关闭侧栏"
/>
<div className="console-panel rounded-lg border border-slate-200 bg-slate-50 flex flex-col overflow-hidden shadow-sm h-full fixed top-0 right-0 h-full w-[88vw] sm:w-[380px] z-40 lg:relative lg:top-auto lg:right-auto lg:h-full lg:w-auto lg:z-10">
<div className="px-4 py-3 border-b border-slate-200 flex items-center justify-between bg-white shrink-0">
<span className="text-xs font-bold text-slate-800">
{sidebarTab === 'notes' ? '观测记录手札' : '文献 AI 问答'}
</span>
<button onClick={() => setShowNotesPanel(false)} className="text-slate-400 hover:text-slate-600 transition-colors cursor-pointer">
<X className="w-4 h-4" />
</button>
</div>
{/* 标签栏选择器 */}
<div className="flex border-b border-slate-200 bg-white select-none shrink-0">
@ -281,5 +290,6 @@ export function ReaderNotesSidebar({
</div>
)}
</div>
</>
);
}

View File

@ -57,18 +57,18 @@ export function ReaderToolbar({
showPdf,
}: ReaderToolbarProps) {
return (
<div className="flex items-center justify-between border-b border-slate-200 pb-3 shrink-0">
<div className="flex-1 min-w-0 pr-4">
<div className="flex flex-col xl:flex-row xl:items-center xl:justify-between border-b border-slate-200 pb-3 shrink-0 gap-3">
<div className="min-w-0 w-full xl:flex-1 pr-0 xl:pr-4">
<h2 className="text-sm font-bold text-slate-900 line-clamp-1 leading-snug" title={selectedPaper.title}>
{selectedPaper.title}
</h2>
<div className="flex items-center gap-2 text-xs text-slate-500 mt-1 font-semibold">
<div className="flex flex-wrap items-center gap-x-2 gap-y-1 text-xs text-slate-500 mt-1 font-semibold">
<span>: {selectedPaper.pub_journal || '未标注'}</span>
<span></span>
<span className="hidden sm:inline"></span>
<span>: {selectedPaper.bibcode}</span>
</div>
</div>
<div className="flex gap-2 items-center relative">
<div className="flex flex-wrap gap-2 items-center relative w-full xl:w-auto">
{/* 快速切换文献菜单 */}
<div className="relative shrink-0">
<button
@ -175,7 +175,7 @@ export function ReaderToolbar({
<button
type="button"
onClick={() => setViewMode('bilingual')}
className={`hidden md:block px-2.5 py-1 rounded-md text-[10px] sm:text-xs font-bold transition-all cursor-pointer ${
className={`hidden lg:block px-2.5 py-1 rounded-md text-[10px] sm:text-xs font-bold transition-all cursor-pointer ${
viewMode === 'bilingual'
? 'bg-white text-slate-800 shadow-sm'
: 'text-slate-500 hover:text-slate-700'

View File

@ -94,7 +94,12 @@ export function useResearchAgent({ showConfirm, showAlert }: UseResearchAgentPro
const fileInputRef = useRef<HTMLInputElement>(null);
const [loadingSessions, setLoadingSessions] = useState(false);
const [loadingHistory, setLoadingHistory] = useState(false);
const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
const [sidebarCollapsed, setSidebarCollapsed] = useState(() => {
if (typeof window !== 'undefined') {
return window.innerWidth < 1024;
}
return false;
});
// 全文搜索历史记录状态
const [searchQuery, setSearchQuery] = useState('');

View File

@ -71,7 +71,7 @@ export function ReaderPanel(props: ReaderPanelProps) {
// 局部管理阅读视角模式(原文/中文/对照)
const [viewMode, setViewMode] = useState<'bilingual' | 'english' | 'chinese'>(() => {
if (typeof window !== 'undefined' && window.innerWidth < 768) {
if (typeof window !== 'undefined' && window.innerWidth < 1024) {
return 'english';
}
if (!chineseText) {
@ -80,16 +80,51 @@ export function ReaderPanel(props: ReaderPanelProps) {
return 'bilingual';
});
// 无翻译时强制回退到英文模式,有翻译时自动切换到双语对照
const [userOverrode, setUserOverrode] = useState(false);
// 监听窗口尺寸变化,处理移动端与桌面端的自适应视角切换
useEffect(() => {
if (typeof window === 'undefined') return;
let prevWidth = window.innerWidth;
const handleResize = () => {
const currentWidth = window.innerWidth;
const wasMobile = prevWidth < 1024;
const isMobile = currentWidth < 1024;
if (wasMobile !== isMobile) {
if (isMobile) {
// 从桌面缩放到移动端:切换显示英文原文
setViewMode('english');
} else {
// 从移动端拉宽到桌面端:若有翻译则显示双语对照,否则显示原文
setViewMode(chineseText ? 'bilingual' : 'english');
}
}
prevWidth = currentWidth;
};
window.addEventListener('resize', handleResize);
return () => window.removeEventListener('resize', handleResize);
}, [chineseText]);
// 无翻译时强制回退到英文模式,有翻译时且在桌面端时自动切换到双语对照
useEffect(() => {
if (userOverrode) return;
const isMobile = typeof window !== 'undefined' && window.innerWidth < 1024;
if (isMobile) {
if (viewMode !== 'english') {
queueMicrotask(() => setViewMode('english'));
}
return;
}
if (!chineseText && viewMode === 'bilingual') {
queueMicrotask(() => setViewMode('english'));
} else if (chineseText && viewMode === 'english') {
queueMicrotask(() => setViewMode('bilingual'));
}
}, [chineseText]);
}, [chineseText, viewMode, userOverrode]);
const handleViewModeChange = (mode: 'bilingual' | 'english' | 'chinese') => {
setUserOverrode(true);

View File

@ -0,0 +1,72 @@
-- migrations/20260629000000_fix_foreign_key_on_update.sql
-- 修复外键级联更新约束 (ON UPDATE CASCADE)
PRAGMA foreign_keys = OFF;
-- 1. 重建 notes 表
CREATE TABLE notes_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
bibcode TEXT NOT NULL,
paragraph_index INTEGER NOT NULL,
note_text TEXT NOT NULL DEFAULT '',
highlight_color TEXT NOT NULL DEFAULT 'yellow',
selected_text TEXT NOT NULL DEFAULT '',
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (bibcode) REFERENCES papers(bibcode) ON DELETE CASCADE ON UPDATE CASCADE
);
INSERT INTO notes_new (id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at)
SELECT id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at FROM notes;
DROP TABLE notes;
ALTER TABLE notes_new RENAME TO notes;
CREATE INDEX IF NOT EXISTS idx_notes_bibcode ON notes(bibcode);
-- 2. 重建 paper_chunks_content 表
CREATE TABLE paper_chunks_content_new (
rowid INTEGER PRIMARY KEY AUTOINCREMENT,
bibcode TEXT,
paragraph_index INTEGER,
content TEXT,
headings TEXT DEFAULT 'Document',
section_index INTEGER DEFAULT 0,
FOREIGN KEY(bibcode) REFERENCES papers(bibcode) ON DELETE CASCADE ON UPDATE CASCADE
);
INSERT INTO paper_chunks_content_new (rowid, bibcode, paragraph_index, content, headings, section_index)
SELECT rowid, bibcode, paragraph_index, content, headings, section_index FROM paper_chunks_content;
DROP TABLE paper_chunks_content;
ALTER TABLE paper_chunks_content_new RENAME TO paper_chunks_content;
-- 3. 重建 paper_targets 表
CREATE TABLE paper_targets_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
bibcode TEXT,
target_name TEXT,
ra TEXT,
dec TEXT,
parallax REAL,
spectral_type TEXT,
v_magnitude REAL,
aliases TEXT,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
otype TEXT,
oname TEXT,
pm_ra REAL,
pm_de REAL,
radial_velocity REAL,
parallax_err REAL,
photometry TEXT,
FOREIGN KEY(bibcode) REFERENCES papers(bibcode) ON DELETE CASCADE ON UPDATE CASCADE,
UNIQUE(bibcode, target_name)
);
INSERT INTO paper_targets_new (id, bibcode, target_name, ra, dec, parallax, spectral_type, v_magnitude, aliases, created_at, otype, oname, pm_ra, pm_de, radial_velocity, parallax_err, photometry)
SELECT id, bibcode, target_name, ra, dec, parallax, spectral_type, v_magnitude, aliases, created_at, otype, oname, pm_ra, pm_de, radial_velocity, parallax_err, photometry FROM paper_targets;
DROP TABLE paper_targets;
ALTER TABLE paper_targets_new RENAME TO paper_targets;
CREATE INDEX IF NOT EXISTS idx_paper_targets_name ON paper_targets(target_name);
PRAGMA foreign_keys = ON;

View File

@ -0,0 +1,4 @@
-- migrations/20260630000000_add_paper_chunks_content_bibcode_index.sql
-- 为 paper_chunks_content 表添加 bibcode 索引,提升按 bibcode 查询/删除的性能
CREATE INDEX IF NOT EXISTS idx_paper_chunks_content_bibcode ON paper_chunks_content(bibcode);

View File

@ -3,7 +3,8 @@
// 内置 Hooks — CancellationHook、MetricsHook、AuditLogHook。
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use tokio::sync::Mutex;
use async_trait::async_trait;
use sqlx::SqlitePool;
@ -67,25 +68,23 @@ impl AgentHook for CancellationHook {
}
async fn pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseAction {
if let Ok(runs) = self.cancelled_runs.lock() {
if runs.contains(&ctx.session_id) {
warn!(
"[CancellationHook] 会话 {} 已被用户取消,阻止工具 {} 执行",
ctx.session_id, ctx.tool_name
);
return PreToolUseAction::Block {
reason: "用户已手动中止执行".to_string(),
};
}
let runs = self.cancelled_runs.lock().await;
if runs.contains(&ctx.session_id) {
warn!(
"[CancellationHook] 会话 {} 已被用户取消,阻止工具 {} 执行",
ctx.session_id, ctx.tool_name
);
return PreToolUseAction::Block {
reason: "用户已手动中止执行".to_string(),
};
}
PreToolUseAction::Continue
}
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
// 清理取消状态
if let Ok(mut runs) = self.cancelled_runs.lock() {
runs.remove(&ctx.session_id);
}
let mut runs = self.cancelled_runs.lock().await;
runs.remove(&ctx.session_id);
}
}
@ -112,9 +111,9 @@ impl MetricsHook {
MetricsHook { data }
}
/// 返回当前指标快照(锁异常时返回 None
/// 返回当前指标快照(锁失败时返回 None
pub fn snapshot(&self) -> Option<MetricsData> {
self.data.lock().ok().map(|d| d.clone())
self.data.try_lock().ok().map(|d| d.clone())
}
/// 获取 Arc 引用,供外部持有
@ -139,9 +138,8 @@ impl AgentHook for MetricsHook {
}
async fn on_session_start(&self, ctx: &SessionStartContext) {
if let Ok(mut data) = self.data.lock() {
data.session_id = Some(ctx.session_id.clone());
}
let mut data = self.data.lock().await;
data.session_id = Some(ctx.session_id.clone());
}
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
@ -149,15 +147,14 @@ impl AgentHook for MetricsHook {
"[Metrics] step={} tool={} is_error={} elapsed={}ms",
ctx.step, ctx.tool_name, ctx.is_error, ctx.elapsed_ms
);
if let Ok(mut data) = self.data.lock() {
*data
.tool_call_counts
.entry(ctx.tool_name.clone())
.or_insert(0) += 1;
data.total_steps = ctx.step;
if ctx.is_error {
data.total_errors += 1;
}
let mut data = self.data.lock().await;
*data
.tool_call_counts
.entry(ctx.tool_name.clone())
.or_insert(0) += 1;
data.total_steps = ctx.step;
if ctx.is_error {
data.total_errors += 1;
}
PostToolUseAction::Continue
}

View File

@ -52,10 +52,9 @@ pub use builtins::{AuditLogHook, CancellationHook, ContextDeduplicator, MetricsH
mod tests {
use super::*;
use async_trait::async_trait;
use serde_json::{json, Value};
use sqlx::SqlitePool;
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use tokio::sync::Mutex;
struct TestHook {
name: String,
@ -204,7 +203,7 @@ mod tests {
async fn test_cancellation_hook_blocks_when_cancelled() {
let mut cancelled = HashSet::new();
cancelled.insert("test_session".to_string());
let cancelled_runs = Arc::new(std::sync::Mutex::new(cancelled));
let cancelled_runs = Arc::new(tokio::sync::Mutex::new(cancelled));
let hook = CancellationHook::new(cancelled_runs);
let ctx = PreToolUseContext {
@ -220,7 +219,7 @@ mod tests {
#[tokio::test]
async fn test_cancellation_hook_allows_when_not_cancelled() {
let cancelled_runs = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
let cancelled_runs = Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new()));
let hook = CancellationHook::new(cancelled_runs);
let ctx = PreToolUseContext {
@ -281,7 +280,7 @@ mod tests {
"start_tracker"
}
async fn on_session_start(&self, ctx: &SessionStartContext) {
self.started.lock().unwrap().push(ctx.session_id.clone());
self.started.lock().await.push(ctx.session_id.clone());
}
}
@ -300,7 +299,7 @@ mod tests {
registry.run_on_session_start(&ctx).await;
// 验证 hook 确实被调用了
let calls = started.lock().unwrap();
let calls = started.lock().await;
assert_eq!(calls.len(), 1);
assert_eq!(calls[0], "test_sid");
}
@ -325,16 +324,16 @@ mod tests {
"lifecycle_tracker"
}
async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {
*self.subagent_start.lock().unwrap() = true;
*self.subagent_start.lock().await = true;
}
async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {
*self.subagent_stop.lock().unwrap() = true;
*self.subagent_stop.lock().await = true;
}
async fn on_pre_compact(&self, _ctx: &PreCompactContext) {
*self.pre_compact.lock().unwrap() = true;
*self.pre_compact.lock().await = true;
}
async fn on_post_compact(&self, _ctx: &PostCompactContext) {
*self.post_compact.lock().unwrap() = true;
*self.post_compact.lock().await = true;
}
}
@ -381,19 +380,16 @@ mod tests {
// 验证所有 4 个 lifecycle hook 都被调用了
assert!(
*subagent_start.lock().unwrap(),
*subagent_start.lock().await,
"on_subagent_start should be called"
);
assert!(
*subagent_stop.lock().unwrap(),
*subagent_stop.lock().await,
"on_subagent_stop should be called"
);
assert!(*pre_compact.lock().await, "on_pre_compact should be called");
assert!(
*pre_compact.lock().unwrap(),
"on_pre_compact should be called"
);
assert!(
*post_compact.lock().unwrap(),
*post_compact.lock().await,
"on_post_compact should be called"
);
}
@ -412,7 +408,7 @@ mod tests {
"prompt_tracker"
}
async fn on_user_prompt_submit(&self, ctx: &UserPromptSubmitContext) {
self.prompts.lock().unwrap().push(ctx.prompt.clone());
self.prompts.lock().await.push(ctx.prompt.clone());
}
}
@ -431,7 +427,7 @@ mod tests {
registry.run_on_user_prompt_submit(&ctx).await;
// 验证 hook 被调用并接收了正确的 prompt
let calls = prompts.lock().unwrap();
let calls = prompts.lock().await;
assert_eq!(calls.len(), 1);
assert_eq!(calls[0], "Hello, agent!");
}

View File

@ -52,8 +52,8 @@ impl HookRegistry {
/// 使得 AgentRuntime.get_metrics() 能查询到实际运行数据。
pub fn with_builtins(
db: SqlitePool,
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
metrics_data: Option<Arc<std::sync::Mutex<MetricsData>>>,
cancelled_runs: Arc<tokio::sync::Mutex<std::collections::HashSet<String>>>,
metrics_data: Option<Arc<tokio::sync::Mutex<MetricsData>>>,
) -> Self {
let mut registry = Self::new();
registry.add(Box::new(CancellationHook::new(cancelled_runs)));

View File

@ -388,7 +388,7 @@ mod tests {
#[test]
fn test_find_duplicate_skips_historical() {
let mut entries = vec![MemoryEntry {
let entries = vec![MemoryEntry {
slug: "historical-one".to_string(),
name: "历史记忆".to_string(),
description: "已过时".to_string(),

View File

@ -416,7 +416,6 @@ impl MemoryManager {
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
#[test]
fn test_build_empty_reminder() {
@ -484,7 +483,7 @@ mod tests {
// 归档的旧版本也会被加载reload 不排除 _v1.md 文件),所以 entries.len() >= 1
assert!(
mgr.entries().len() >= 1,
!mgr.entries().is_empty(),
"at least the active entry should exist"
);

View File

@ -245,7 +245,7 @@ mod tests {
#[test]
fn test_extract_slugs_string_array() {
let entries = vec![
let entries = [
make_entry("alpha", "Alpha", "first"),
make_entry("beta", "Beta", "second"),
make_entry("gamma", "Gamma", "third"),
@ -258,7 +258,7 @@ mod tests {
#[test]
fn test_extract_slugs_numeric_fallback() {
let entries = vec![make_entry("x", "X", "x"), make_entry("y", "Y", "y")];
let entries = [make_entry("x", "X", "x"), make_entry("y", "Y", "y")];
let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect();
let result = extract_slugs("[0]", &candidates);
@ -267,7 +267,7 @@ mod tests {
#[test]
fn test_extract_slugs_invalid_json() {
let entries = vec![make_entry("x", "X", "x")];
let entries = [make_entry("x", "X", "x")];
let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect();
assert!(extract_slugs("not json at all", &candidates).is_err());

View File

@ -65,6 +65,7 @@ impl CompactionCircuitBreaker {
/// 设置可注入的时间源(仅用于测试)。
#[cfg(test)]
#[allow(dead_code)]
fn with_time_source(mut self, f: Box<dyn Fn() -> Instant + Send>) -> Self {
self.time_source = Some(f);
self

View File

@ -512,10 +512,10 @@ mod tests {
#[test]
fn test_backoff_delay() {
let d0 = backoff_delay(0, None);
assert!(d0 >= 500 && d0 <= 700);
assert!((500..=700).contains(&d0));
let d3 = backoff_delay(3, None);
assert!(d3 >= 4000 && d3 <= 5000);
assert!((4000..=5000).contains(&d3));
let d10 = backoff_delay(10, None);
assert!(d10 <= 40_000);

View File

@ -413,37 +413,7 @@ pub async fn execute_parallel(
// 存储待处理的权限请求
{
let mut perms = match app_state.pending_permissions.lock() {
Ok(p) => p,
Err(e) => {
warn!("[Executor] 权限系统内部错误: {}", e);
let err_output =
format!("权限系统内部错误,工具 {} 被拒绝", prep.tool_name);
let _ = tx.send(AgentStreamEvent::ToolResult {
tool_call_id: prep.tool_call_id.clone(),
name: prep.tool_name.clone(),
output: err_output.clone(),
is_error: true,
metadata: serde_json::json!({}),
step,
});
let err_msg =
ChatMessage::tool_result(&prep.tool_call_id, &err_output);
save_tool_message_sync(db, &sid, turn_index, step, &err_msg);
tool_messages.push(ToolResultMessage {
chat_message: err_msg,
was_error: true,
});
// 内部错误 → 记录拒绝追踪
if let Some(dt) = denial_tracker {
if let Ok(mut tracker) = dt.lock() {
tracker.record_denial();
}
}
denied_indices.insert(i);
continue;
}
};
let mut perms = app_state.pending_permissions.lock().await;
perms.insert(
perm_id.clone(),
PendingPermission {
@ -461,9 +431,7 @@ pub async fn execute_parallel(
let perm_result = tokio::time::timeout(timeout_dur, resp_rx).await;
// 清理待处理的权限请求
if let Ok(mut perms) = app_state.pending_permissions.lock() {
perms.remove(&perm_id);
}
app_state.pending_permissions.lock().await.remove(&perm_id);
match perm_result {
Ok(Ok(response)) if response.allowed => {
@ -570,11 +538,10 @@ pub async fn execute_parallel(
let cancel_handle = tokio::spawn(async move {
loop {
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
if let Ok(locked) = app_state_ref.cancelled_runs.lock() {
if locked.contains(&sid_ref) {
cancel_flag.store(true, Ordering::SeqCst);
return;
}
let locked = app_state_ref.cancelled_runs.lock().await;
if locked.contains(&sid_ref) {
cancel_flag.store(true, Ordering::SeqCst);
return;
}
}
});

View File

@ -14,7 +14,6 @@
use lru::LruCache;
use std::num::NonZeroUsize;
use std::path::Path;
use tracing::{info, warn};
/// 缓存条目最大数量
@ -82,35 +81,27 @@ impl FileStateCache {
}
}
/// 规范化路径 key确保一致性
/// 规范化路径 key确保一致性—— 仅做字符串规范化,不做磁盘 I/O
fn normalize_key(path: &str) -> String {
// 去除尾随斜杠,规范化重复斜杠
let p = Path::new(path);
// 尝试 canonicalize跟随符号链接失败则用简单的字符串规范化
match p.canonicalize() {
Ok(canon) => canon.to_string_lossy().to_string(),
Err(_) => {
// 简单规范化:折叠重复的 /
let mut result = String::with_capacity(path.len());
let mut prev_slash = false;
for ch in path.chars() {
if ch == '/' || ch == '\\' {
if !prev_slash {
result.push('/');
prev_slash = true;
}
} else {
result.push(ch);
prev_slash = false;
}
// 简单规范化:折叠重复的 / 和 \
let mut result = String::with_capacity(path.len());
let mut prev_slash = false;
for ch in path.chars() {
if ch == '/' || ch == '\\' {
if !prev_slash {
result.push('/');
prev_slash = true;
}
// 去除尾随 /
if result.ends_with('/') && result.len() > 1 {
result.pop();
}
result
} else {
result.push(ch);
prev_slash = false;
}
}
// 去除尾随 /
if result.ends_with('/') && result.len() > 1 {
result.pop();
}
result
}
/// 获取缓存的条目,返回 Some(&FileState) 若存在

View File

@ -74,7 +74,7 @@ pub struct AgentRuntime {
/// 后台任务通知队列(支持 bg_task_run/bg_task_check
bg_notification_queue: Arc<BgNotificationQueue>,
/// 指标采集 hook 的数据引用(供 API 查询)
metrics_data: Arc<std::sync::Mutex<super::hooks::MetricsData>>,
metrics_data: Arc<tokio::sync::Mutex<super::hooks::MetricsData>>,
/// 压缩熔断器(跨 turn 共享,防止无限压缩循环)
compaction_breaker: Arc<std::sync::Mutex<circuit_breaker::CompactionCircuitBreaker>>,
/// 权限检查器
@ -111,7 +111,7 @@ impl AgentRuntime {
apply_mode_config(&mut config, mode);
let queue = Arc::new(BgNotificationQueue::new());
let metrics_data = Arc::new(std::sync::Mutex::new(super::hooks::MetricsData::default()));
let metrics_data = Arc::new(tokio::sync::Mutex::new(super::hooks::MetricsData::default()));
let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config));
let denial_tracker = Arc::new(std::sync::Mutex::new(denial_tracker::DenialTracker::new(
config.denial_max_consecutive,
@ -182,7 +182,7 @@ impl AgentRuntime {
apply_mode_config(&mut config, mode);
let queue = Arc::new(BgNotificationQueue::new());
let metrics_data = Arc::new(std::sync::Mutex::new(super::hooks::MetricsData::default()));
let metrics_data = Arc::new(tokio::sync::Mutex::new(super::hooks::MetricsData::default()));
let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config));
let denial_tracker = Arc::new(std::sync::Mutex::new(denial_tracker::DenialTracker::new(
config.denial_max_consecutive,
@ -240,8 +240,7 @@ impl AgentRuntime {
/// 返回当前运行指标快照(锁异常时返回默认值)
pub fn get_metrics(&self) -> super::hooks::MetricsData {
self.metrics_data
.lock()
.ok()
.try_lock()
.map(|m| m.clone())
.unwrap_or_default()
}
@ -493,11 +492,10 @@ impl AgentRuntime {
// 注册当前会话的权限检查器(如不存在则从全局配置初始化)
{
if let Ok(mut checkers) = self.app_state.session_permission_checkers.write() {
checkers
.entry(sid.clone())
.or_insert_with(|| (*self.permission_checker).clone());
}
let mut checkers = self.app_state.session_permission_checkers.write().await;
checkers
.entry(sid.clone())
.or_insert_with(|| (*self.permission_checker).clone());
}
let tool_defs = self.tool_registry.definitions();
@ -525,11 +523,8 @@ impl AgentRuntime {
// 检查用户取消
let is_cancelled = {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
cancelled.remove(sid)
} else {
false
}
let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(sid)
};
if is_cancelled {
@ -874,8 +869,9 @@ impl AgentRuntime {
self.app_state
.session_permission_checkers
.read()
.ok()
.and_then(|checkers| checkers.get(sid).cloned())
.await
.get(sid)
.cloned()
};
let exec_result = executor::execute_parallel(
&prepared_calls,
@ -948,9 +944,8 @@ impl AgentRuntime {
}
if exec_result.was_cancelled {
if let Ok(mut locked) = self.app_state.cancelled_runs.lock() {
locked.remove(sid);
}
let mut locked = self.app_state.cancelled_runs.lock().await;
locked.remove(sid);
warn!(
"[AgentRuntime] 工具执行期间被用户手动中止,会话 ID: {}",
sid
@ -1018,9 +1013,8 @@ impl AgentRuntime {
match output.status {
StreamStatus::Success => return Some(output),
StreamStatus::Cancelled => {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
cancelled.remove(session_id);
}
let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id);
warn!(
"[AgentRuntime] 流式调用期间被用户手动中止,会话 ID: {}",
session_id
@ -1098,14 +1092,13 @@ impl AgentRuntime {
}
// 检查用户取消
if let Ok(cancelled) = self.app_state.cancelled_runs.lock() {
if cancelled.contains(session_id) {
warn!("[AgentRuntime] 退避重试期间被用户取消");
let _ = tx.send(AgentStreamEvent::Error {
message: "用户已手动中止执行。".to_string(),
});
return None;
}
let cancelled = self.app_state.cancelled_runs.lock().await;
if cancelled.contains(session_id) {
warn!("[AgentRuntime] 退避重试期间被用户取消");
let _ = tx.send(AgentStreamEvent::Error {
message: "用户已手动中止执行。".to_string(),
});
return None;
}
// 重试 LLM 调用
@ -1127,9 +1120,8 @@ impl AgentRuntime {
return Some(retry_output);
}
StreamStatus::Cancelled => {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
cancelled.remove(session_id);
}
let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id);
return None;
}
StreamStatus::Error(_) => {
@ -1237,9 +1229,8 @@ impl AgentRuntime {
return Some(retry_output);
}
StreamStatus::Cancelled => {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
cancelled.remove(session_id);
}
let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id);
warn!("[AgentRuntime] 恢复期间被用户中止");
let _ = tx.send(AgentStreamEvent::Error {
message: "用户已手动中止执行。".to_string(),
@ -1341,7 +1332,7 @@ impl AgentRuntime {
if let Some(skills) = self
.app_state
.skill_registry
.read()
.try_read()
.ok()
.and_then(|r| r.build_reminder())
{

View File

@ -102,7 +102,8 @@ impl ToolPartitioner {
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::sync::Arc;
use tokio::sync::RwLock;
use super::*;
use crate::agent::skills::SkillRegistry;

View File

@ -1017,7 +1017,8 @@ mod tests {
#[tokio::test]
async fn test_create_new_session() {
let db = setup_db_with_mode().await;
let llm = LlmClient::new("key".into(), "http://localhost".into(), "test-model".into());
let llm =
LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()).unwrap();
let info = create_or_resume_session(&db, None, &llm, "deep-research")
.await
@ -1040,7 +1041,8 @@ mod tests {
#[tokio::test]
async fn test_resume_existing_session() {
let db = setup_db_with_mode().await;
let llm = LlmClient::new("key".into(), "http://localhost".into(), "test-model".into());
let llm =
LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()).unwrap();
// 先创建一个会话
let created = create_or_resume_session(&db, None, &llm, "literature-reader")
@ -1067,7 +1069,8 @@ mod tests {
#[tokio::test]
async fn test_resume_nonexistent_session_errors() {
let db = setup_db_with_mode().await;
let llm = LlmClient::new("key".into(), "http://localhost".into(), "test-model".into());
let llm =
LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()).unwrap();
let result =
create_or_resume_session(&db, Some("nonexistent-id".into()), &llm, "default").await;

View File

@ -43,7 +43,7 @@ pub async fn process_llm_stream(
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
step: usize,
session_id: &str,
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
cancelled_runs: Arc<tokio::sync::Mutex<std::collections::HashSet<String>>>,
enable_thinking: bool,
) -> StreamOutput {
// 1. 发起 LLM 流式调用
@ -76,10 +76,9 @@ pub async fn process_llm_stream(
let cancel_fut = async {
loop {
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
if let Ok(cancelled) = cancelled_runs.lock() {
if cancelled.contains(&sid) {
return;
}
let cancelled = cancelled_runs.lock().await;
if cancelled.contains(&sid) {
return;
}
}
};

View File

@ -549,8 +549,6 @@ mod tests {
/// 构建测试用的 ToolContext。
async fn make_test_tool_context() -> ToolContext {
use crate::agent::memory::MemoryManager;
use crate::agent::runtime::file_cache::FileStateCache;
use crate::agent::runtime::permission::PermissionChecker;
use crate::agent::skills::SkillRegistry;
use crate::api::AppState;
use crate::clients::ads::AdsClient;
@ -563,16 +561,17 @@ mod tests {
use crate::services::translation::Dictionary;
use crate::Config;
use std::path::PathBuf;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::Arc;
let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();
sqlx::migrate!("./migrations").run(&pool).await.unwrap();
let config = Config::from_env();
let llm = LlmClient::new("tk".into(), "http://localhost".into(), "m".into());
let embedding = EmbeddingClient::new("tk".into(), "http://localhost".into(), "e".into());
let ads = AdsClient::new("tk".into());
let arxiv = ArxivClient::new();
let llm = LlmClient::new("tk".into(), "http://localhost".into(), "m".into()).unwrap();
let embedding =
EmbeddingClient::new("tk".into(), "http://localhost".into(), "e".into()).unwrap();
let ads = AdsClient::new("tk".into()).unwrap();
let arxiv = ArxivClient::new().unwrap();
let qiniu = QiniuClient::new(
"ak".into(),
"sk".into(),
@ -593,20 +592,27 @@ mod tests {
vision_llm: None,
embedding,
downloader: Downloader::new().expect("downloader"),
http_client: reqwest::Client::new(),
harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())),
batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(Mutex::new(std::collections::HashSet::new())),
skill_registry: Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from("/tmp/sk")))),
pending_questions: Arc::new(Mutex::new(std::collections::HashMap::new())),
pending_permissions: Arc::new(Mutex::new(std::collections::HashMap::new())),
session_permission_checkers: Arc::new(RwLock::new(std::collections::HashMap::new())),
cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
skill_registry: Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(PathBuf::from(
"/tmp/sk",
)))),
pending_questions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
pending_permissions: Arc::new(
tokio::sync::Mutex::new(std::collections::HashMap::new()),
),
session_permission_checkers: Arc::new(tokio::sync::RwLock::new(
std::collections::HashMap::new(),
)),
sse_broadcast: None,
memory_manager: Arc::new(tokio::sync::Mutex::new(MemoryManager::new(PathBuf::from(
"/tmp/test_mem",
)))),
sessions: Arc::new(Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(Mutex::new(std::collections::HashMap::new())),
sessions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
});
ToolContext::new(app_state)
@ -614,7 +620,7 @@ mod tests {
/// 创建包含指定 Mock 工具的 StreamingToolExecutor。
async fn make_executor(tools: Vec<MockAgentTool>) -> StreamingToolExecutor {
let skill_registry = Arc::new(std::sync::RwLock::new(
let skill_registry = Arc::new(tokio::sync::RwLock::new(
crate::agent::skills::SkillRegistry::new(std::path::PathBuf::from("/tmp/sk")),
));
let mut registry = ToolRegistry::new(skill_registry);

View File

@ -38,8 +38,9 @@ pub mod types;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock};
use std::sync::Arc;
use std::time::SystemTime;
use tokio::sync::RwLock;
use tracing::{info, warn};
// ── Frontmatter ────────────────────────────────────────────────────────────
@ -267,7 +268,7 @@ impl SkillRegistry {
use notify::{RecursiveMode, Watcher};
use std::time::Duration;
let skills_dir = match self_arc.read() {
let skills_dir = match self_arc.try_read() {
Ok(r) => r.skills_dir.clone(),
Err(_) => {
warn!("[SkillRegistry] RwLock 异常,无法启动文件监视器");
@ -310,7 +311,7 @@ impl SkillRegistry {
// 等待 debounce 窗口
while rx.recv_timeout(Duration::from_millis(300)).is_ok() {}
info!("[SkillRegistry] 检测到 skill 文件变更,自动刷新");
if let Ok(mut registry) = self_arc.write() {
if let Ok(mut registry) = self_arc.try_write() {
let reg: &mut SkillRegistry = &mut registry;
reg.refresh();
}
@ -576,7 +577,6 @@ pub fn substitute_variables(body: &str, skill_dir: &Path, session_id: Option<&st
/// ```ignore
/// let creator = SkillCreator::new(skills_dir);
/// let created = creator.create_from_pattern(&pattern)?;
#[cfg(test)]
mod tests {
use super::*;

View File

@ -146,7 +146,7 @@ impl AgentTool for AskUserTool {
.unwrap_or_default();
let question_id = uuid::Uuid::new_v4().to_string();
let short_id = question_id[..8].to_string();
let short_id = question_id[..12].to_string();
let question = UserQuestion {
question_id: short_id.clone(),
@ -169,12 +169,7 @@ impl AgentTool for AskUserTool {
// 将通道发送端存储到 AppState 的待处理问题列表
{
let mut pending = match ctx.app_state.pending_questions.lock() {
Ok(p) => p,
Err(_) => {
return ToolOutput::error("待处理问题队列不可用(内部锁异常),请稍后重试");
}
};
let mut pending = ctx.app_state.pending_questions.lock().await;
pending.insert(
short_id.clone(),
crate::api::PendingQuestion {
@ -201,9 +196,11 @@ impl AgentTool for AskUserTool {
);
// 清理
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() {
pending.remove(&short_id);
}
ctx.app_state
.pending_questions
.lock()
.await
.remove(&short_id);
let free_text = answer.free_text.unwrap_or_default();
let response = if !answer.answers.is_empty() {
@ -228,16 +225,20 @@ impl AgentTool for AskUserTool {
}
Ok(Err(_)) => {
// 通道关闭(发送端被 drop
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() {
pending.remove(&short_id);
}
ctx.app_state
.pending_questions
.lock()
.await
.remove(&short_id);
ToolOutput::error("用户取消了回答")
}
Err(_) => {
// 超时
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() {
pending.remove(&short_id);
}
ctx.app_state
.pending_questions
.lock()
.await
.remove(&short_id);
ToolOutput::error(format!(
"等待用户回答超时 (5 分钟)。问题: {}",
question_text

View File

@ -10,7 +10,7 @@ use std::path::{Path, PathBuf};
use async_trait::async_trait;
use serde_json::json;
use tokio::sync::mpsc::UnboundedSender;
use tracing::{error, info};
use tracing::error;
use crate::agent::runtime::AgentStreamEvent;
use crate::agent::tools::filesystem::security::{has_path_traversal, is_path_allowed};
@ -29,57 +29,17 @@ impl AnalyzeImageTool {
sse_tx: Option<&UnboundedSender<AgentStreamEvent>>,
tool_call_id: Option<String>,
) -> ToolOutput {
let (base64_data, mime_type) = match load_image(library_dir, image_path).await {
Ok(data) => data,
Err(e) => {
error!("[analyze_image] 加载图片失败: {} ({})", e, image_path);
return ToolOutput::error(format!("加载图片失败: {}", e));
}
};
info!(
"[analyze_image] 流式分析图片: {} ({} bytes, type={})",
match crate::services::vision::analyze_image_stream_service(
image_path,
base64_data.len(),
mime_type
);
let system_prompt = "你是一位专业的天体物理学家。请基于提供的图片和用户问题,给出专业、详细、清晰的中文分析。如有数学公式,使用 LaTeX 格式。";
// 流式输出:先发送状态头部(作为第一条 text_delta保证持久化到工具结果中
let header = format!(
"> 正在使用视觉模型({})分析图片: {}\n\n",
vision_llm.model(),
image_path
);
if let Some(tx) = &sse_tx {
let _ = tx.send(AgentStreamEvent::TextDelta {
content: header.clone(),
tool_call_id: tool_call_id.clone(),
});
}
let sse_tx_clone = sse_tx.cloned();
let tc_id = tool_call_id.clone();
let result = vision_llm
.analyze_image_stream(
system_prompt,
question,
&base64_data,
&mime_type,
move |delta: &str| {
if let Some(ref tx) = sse_tx_clone {
let _ = tx.send(AgentStreamEvent::TextDelta {
content: delta.to_string(),
tool_call_id: tc_id.clone(),
});
}
},
)
.await;
match result {
Ok(answer) => ToolOutput::success(
question,
vision_llm,
library_dir,
sse_tx,
tool_call_id,
)
.await
{
Ok((header, answer)) => ToolOutput::success(
format!("{}{}", header, answer),
json!({
"image_path": image_path,
@ -88,7 +48,7 @@ impl AnalyzeImageTool {
}),
),
Err(e) => {
error!("[analyze_image] 视觉模型调用失败: {}", e);
error!("[analyze_image] 视觉分析大模型服务执行失败: {}", e);
ToolOutput::error(format!("图片分析失败: {}", e))
}
}
@ -192,72 +152,6 @@ impl AgentTool for AnalyzeImageTool {
}
}
/// 加载图片:本地路径或 HTTP(S) URL → base64 data + mime type
async fn load_image(library_dir: &Path, path: &str) -> anyhow::Result<(String, String)> {
if path.starts_with("http://") || path.starts_with("https://") {
// ── SSRF 防护 ──
let parsed =
reqwest::Url::parse(path).map_err(|e| anyhow::anyhow!("无效的图片 URL: {}", e))?;
let host = parsed
.host_str()
.ok_or_else(|| anyhow::anyhow!("URL 缺少 host"))?;
// 字面量快检 + 异步 DNS 深度校验(不阻塞 tokio worker
if crate::utils::ssrf::is_host_literal_blocked(host) {
return Err(anyhow::anyhow!(
"出于安全考虑,禁止请求内网/保留地址: {}",
host
));
}
if crate::utils::ssrf::is_host_blocked_async(host).await {
return Err(anyhow::anyhow!("DNS 解析后指向内网/保留地址: {}", host));
}
// 使用安全的重定向策略:允许跟随合法 CDN 重定向S3/Imgur 等),
// 但每次跳转都对目标 URL 做 scheme + host 字面量校验,拦截内网目标。
let client = reqwest::Client::builder()
.redirect(crate::utils::ssrf::safe_redirect_policy())
.build()?;
let resp = client.get(path).send().await?;
let mime = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.unwrap_or("image/png")
.to_string();
let bytes = resp.bytes().await?;
let b64 = base64_encode(&bytes);
Ok((b64, mime))
} else {
let abs_path = if path.starts_with('/') {
Path::new(path).to_path_buf()
} else {
library_dir.join(path)
};
if !abs_path.exists() {
return Err(anyhow::anyhow!("图片文件不存在: {}", abs_path.display()));
}
let ext = abs_path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("png")
.to_lowercase();
let mime = match ext.as_str() {
"jpg" | "jpeg" => "image/jpeg",
"gif" => "image/gif",
"webp" => "image/webp",
"svg" => "image/svg+xml",
_ => "image/png",
};
let bytes = std::fs::read(&abs_path)?;
let b64 = base64_encode(&bytes);
Ok((b64, mime.to_string()))
}
}
fn base64_encode(bytes: &[u8]) -> String {
use base64::{engine::general_purpose, Engine as _};
general_purpose::STANDARD.encode(bytes)
}
#[cfg(test)]
mod tests {
use super::*;
@ -294,7 +188,7 @@ mod tests {
} else {
config.llm_vision_api_base.clone()
};
let vision_llm = LlmClient::new(key, base, config.llm_vision_model.clone());
let vision_llm = LlmClient::new(key, base, config.llm_vision_model.clone()).unwrap();
// 写入测试图片到临时目录
let tmp_dir = std::env::temp_dir().join("astro_vision_test");
@ -366,7 +260,7 @@ mod tests {
let rt = tokio::runtime::Runtime::new().unwrap();
let (b64, mime) = rt
.block_on(load_image(&tmp_dir, "test.png"))
.block_on(crate::services::vision::load_image(&tmp_dir, "test.png"))
.expect("应成功加载本地图片");
assert_eq!(mime, "image/png");
@ -378,7 +272,10 @@ mod tests {
#[test]
fn test_load_image_not_found() {
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(load_image(Path::new("/tmp"), "nonexistent_image_xyz.png"));
let result = rt.block_on(crate::services::vision::load_image(
Path::new("/tmp"),
"nonexistent_image_xyz.png",
));
assert!(result.is_err());
}
}

View File

@ -21,8 +21,7 @@ pub use research::citation_network::GetCitationNetworkTool;
pub use research::library_search::SearchLocalLibraryTool;
pub use research::metadata::GetPaperMetadataTool;
pub use research::note::SaveNoteTool;
pub use research::paper_content::GetPaperContentTool;
pub use research::paper_outline::GetPaperOutlineTool;
pub use research::paper::{GetPaperContentTool, GetPaperOutlineTool};
pub use research::rag::RagSearchTool;
pub use research::target::QueryTargetTool;

View File

@ -7,20 +7,10 @@
use async_trait::async_trait;
use serde_json::json;
use sqlx::FromRow;
use tracing::info;
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
/// 引用条目(关联查询结果)
#[derive(Debug, FromRow)]
struct CitationEntry {
bibcode: String,
title: String,
authors: Option<String>,
year: Option<String>,
citation_count: Option<i32>,
}
use crate::services::citation::{get_citations_paginated, search_citations};
pub struct GetCitationNetworkTool;
@ -117,7 +107,7 @@ impl AgentTool for GetCitationNetworkTool {
let state = &ctx.app_state;
// 验证源文献存在
let paper = match crate::api::helpers::get_paper_from_db(
let paper = match crate::services::paper::get_paper_from_db(
&state.db,
&state.config.library_dir,
&bibcode,
@ -133,313 +123,134 @@ impl AgentTool for GetCitationNetworkTool {
}
};
// 构建基础查询
let (join_col, filter_col) = match direction {
"citations" => ("source_bibcode", "target_bibcode"),
_ => ("target_bibcode", "source_bibcode"),
};
let order_clause = match sort {
"year" => "p.year DESC",
"citation_count" => "p.citation_count DESC",
_ => "p.bibcode ASC",
};
// 有 query → 引用查找模式:在关联文献中按作者+年份搜索
if let Some(ref q) = query {
let (author_part, year_part) = parse_query(q);
let results = search_citations(
&state.db,
&paper.bibcode,
join_col,
filter_col,
&author_part,
&year_part,
)
.await;
match search_citations(&state.db, &paper.bibcode, direction, q).await {
Ok(results) => {
if results.is_empty() {
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
return ToolOutput::success(
format!(
"在 {} 的{}列表中未找到匹配 '{}' 的文献。请尝试调整搜索词(如仅用作者姓氏)。",
paper.bibcode, dir_label, q
),
json!({ "bibcode": paper.bibcode, "query": q, "direction": direction, "count": 0, "results": [] }),
);
}
if results.is_empty() {
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
return ToolOutput::success(
format!(
"在 {} 的{}列表中未找到匹配 '{}' 的文献。请尝试调整搜索词(如仅用作者姓氏)。",
paper.bibcode, dir_label, q
),
json!({ "bibcode": paper.bibcode, "query": q, "direction": direction, "count": 0, "results": [] }),
);
}
let content = results
.iter()
.enumerate()
.map(|(i, r)| {
format!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
i + 1,
r.bibcode,
r.title,
r.year.as_deref().unwrap_or(""),
r.authors.join(", "),
r.citation_count,
)
})
.collect::<Vec<_>>()
.join("\n\n");
let display: Vec<serde_json::Value> = results.iter().map(citation_to_json).collect();
let content = display
.iter()
.enumerate()
.map(|(i, r)| {
format!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
i + 1,
r["bibcode"].as_str().unwrap_or(""),
r["title"].as_str().unwrap_or(""),
r["year"].as_str().unwrap_or(""),
r["authors"].as_str().unwrap_or(""),
r["citation_count"].as_i64().unwrap_or(0),
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
ToolOutput::success(
format!(
"在 {} 的{}列表中找到 {} 篇匹配 '{}' 的文献:\n\n{}",
paper.bibcode,
dir_label,
results.len(),
q,
content
),
json!({
"bibcode": paper.bibcode,
"query": q,
"direction": direction,
"count": results.len(),
"results": results,
"mode": "search"
}),
)
})
.collect::<Vec<_>>()
.join("\n\n");
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
ToolOutput::success(
format!(
"在 {} 的{}列表中找到 {} 篇匹配 '{}' 的文献:\n\n{}",
paper.bibcode,
dir_label,
results.len(),
q,
content
),
json!({
"bibcode": paper.bibcode,
"query": q,
"direction": direction,
"count": results.len(),
"results": display,
"mode": "search"
}),
)
}
Err(e) => ToolOutput::error(format!("引文检索失败: {}", e)),
}
} else {
// 无 query → 分页浏览模式
let total: i64 = sqlx::query_scalar(&format!(
"SELECT COUNT(*) FROM citations_references WHERE {} = ?",
filter_col
))
.bind(&paper.bibcode)
.fetch_one(&state.db)
.await
.unwrap_or(0);
let query_sql = format!(
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.{} \
WHERE cr.{} = ? \
ORDER BY {} \
LIMIT ? OFFSET ?",
join_col, filter_col, order_clause
);
let rows: Vec<CitationEntry> = sqlx::query_as(&query_sql)
.bind(&paper.bibcode)
.bind(limit)
.bind(offset)
.fetch_all(&state.db)
match get_citations_paginated(&state.db, &paper.bibcode, direction, sort, offset, limit)
.await
.unwrap_or_default();
{
Ok((rows, total)) => {
let page_count = rows.len();
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
let sort_label = match sort {
"citation_count" => "按被引次数降序",
"year" => "按年份降序",
_ => "默认顺序",
};
let page_count = rows.len();
let display: Vec<serde_json::Value> = rows.iter().map(citation_to_json).collect();
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
let sort_label = match sort {
"citation_count" => "按被引次数降序",
"year" => "按年份降序",
_ => "默认顺序",
};
let content = if rows.is_empty() {
format!("文献 {} 暂无{}记录。", paper.bibcode, dir_label)
} else {
let items = display
.iter()
.enumerate()
.map(|(i, r)| {
let content = if rows.is_empty() {
format!("文献 {} 暂无{}记录。", paper.bibcode, dir_label)
} else {
let items = rows
.iter()
.enumerate()
.map(|(i, r)| {
format!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
offset as usize + i + 1,
r.bibcode,
r.title,
r.year.as_deref().unwrap_or(""),
r.authors.join(", "),
r.citation_count,
)
})
.collect::<Vec<_>>()
.join("\n\n");
format!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
offset as usize + i + 1,
r["bibcode"].as_str().unwrap_or(""),
r["title"].as_str().unwrap_or(""),
r["year"].as_str().unwrap_or(""),
r["authors"].as_str().unwrap_or(""),
r["citation_count"].as_i64().unwrap_or(0),
"文献 {} 的{}列表({},共 {} 篇,第 {}-{} 条):\n\n{}",
paper.bibcode,
dir_label,
sort_label,
total,
offset + 1,
offset + page_count as i64,
items,
)
})
.collect::<Vec<_>>()
.join("\n\n");
format!(
"文献 {} 的{}列表({},共 {} 篇,第 {}-{} 条):\n\n{}",
paper.bibcode,
dir_label,
sort_label,
total,
offset + 1,
offset + page_count as i64,
items,
)
};
};
ToolOutput::success(
content,
json!({
"bibcode": paper.bibcode,
"direction": direction,
"sort": sort,
"total": total,
"offset": offset,
"limit": limit,
"page_count": page_count,
"results": display,
"mode": "browse"
}),
)
ToolOutput::success(
content,
json!({
"bibcode": paper.bibcode,
"direction": direction,
"sort": sort,
"total": total,
"offset": offset,
"limit": limit,
"page_count": page_count,
"results": rows,
"mode": "browse"
}),
)
}
Err(e) => ToolOutput::error(format!("获取引文列表失败: {}", e)),
}
}
}
}
/// 解析查询字符串,提取作者部分和年份部分。
/// "Lei 2023" → ("Lei", Some("2023"))
/// "Zhang et al. 2020" → ("Zhang et al.", Some("2020"))
/// "Smith" → ("Smith", None)
fn parse_query(q: &str) -> (String, Option<String>) {
// 去掉 "et al." 变体中的句号,保留为空格分隔
let cleaned = q.replace("et al.", "et al");
let parts: Vec<&str> = cleaned.split_whitespace().collect();
// 找最后一个看起来像年份的 token4位数字19xx 或 20xx
let mut year: Option<String> = None;
let mut author_tokens: Vec<&str> = Vec::new();
for (i, &p) in parts.iter().enumerate().rev() {
if year.is_none()
&& p.len() == 4
&& p.chars().all(|c| c.is_ascii_digit())
&& (p.starts_with("19") || p.starts_with("20"))
{
year = Some(p.to_string());
// 年份之前的是作者,之后的不属于查询
author_tokens = parts[..i].to_vec();
break;
}
}
if year.is_none() {
// 没有找到年份,整个字符串作为作者搜索
author_tokens = parts;
}
let author_part = if author_tokens.is_empty() {
q.to_string()
} else {
author_tokens.join(" ")
};
(author_part, year)
}
/// 在引用关系中搜索匹配的文献
async fn search_citations(
db: &sqlx::SqlitePool,
bibcode: &str,
join_col: &str,
filter_col: &str,
author: &str,
year: &Option<String>,
) -> Vec<CitationEntry> {
let author_pattern = format!("%{}%", author);
if let Some(ref y) = year {
let sql = format!(
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.{} \
WHERE cr.{} = ? AND p.authors LIKE ? AND p.year = ?",
join_col, filter_col
);
sqlx::query_as::<_, CitationEntry>(&sql)
.bind(bibcode)
.bind(&author_pattern)
.bind(y)
.fetch_all(db)
.await
.unwrap_or_default()
} else {
let sql = format!(
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.{} \
WHERE cr.{} = ? AND p.authors LIKE ?",
join_col, filter_col
);
sqlx::query_as::<_, CitationEntry>(&sql)
.bind(bibcode)
.bind(&author_pattern)
.fetch_all(db)
.await
.unwrap_or_default()
}
}
fn citation_to_json(e: &CitationEntry) -> serde_json::Value {
let authors: Vec<String> = e
.authors
.as_deref()
.and_then(|a| serde_json::from_str(a).ok())
.unwrap_or_default();
let first_author = authors
.first()
.cloned()
.unwrap_or_else(|| "未知".to_string());
json!({
"bibcode": e.bibcode,
"title": e.title,
"authors": authors,
"first_author": first_author,
"year": e.year,
"citation_count": e.citation_count.unwrap_or(0),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_query_with_year() {
let (author, year) = parse_query("Lei 2023");
assert_eq!(author, "Lei");
assert_eq!(year, Some("2023".into()));
}
#[test]
fn test_parse_query_et_al() {
let (author, year) = parse_query("Zhang et al. 2020");
assert_eq!(author, "Zhang et al");
assert_eq!(year, Some("2020".into()));
}
#[test]
fn test_parse_query_author_only() {
let (author, year) = parse_query("Smith");
assert_eq!(author, "Smith");
assert_eq!(year, None);
}
#[test]
fn test_parse_query_full_name() {
let (author, year) = parse_query("Gaia Collaboration 2018");
assert_eq!(author, "Gaia Collaboration");
assert_eq!(year, Some("2018".into()));
}
}

View File

@ -54,8 +54,12 @@ impl AgentTool for GetPaperMetadataTool {
info!("[GetPaperMetadata] 获取文献元数据: {}", bibcode);
let state = &ctx.app_state;
match crate::api::helpers::get_paper_from_db(&state.db, &state.config.library_dir, &bibcode)
.await
match crate::services::paper::get_paper_from_db(
&state.db,
&state.config.library_dir,
&bibcode,
)
.await
{
Ok(paper) => {
let content = format!(

View File

@ -5,8 +5,7 @@ pub mod citation_network;
pub mod library_search;
pub mod metadata;
pub mod note;
pub mod paper_content;
pub mod paper_outline;
pub mod paper;
pub mod rag;
pub mod target;
@ -14,7 +13,6 @@ pub use citation_network::GetCitationNetworkTool;
pub use library_search::SearchLocalLibraryTool;
pub use metadata::GetPaperMetadataTool;
pub use note::SaveNoteTool;
pub use paper_content::GetPaperContentTool;
pub use paper_outline::GetPaperOutlineTool;
pub use paper::{GetPaperContentTool, GetPaperOutlineTool};
pub use rag::RagSearchTool;
pub use target::QueryTargetTool;

View File

@ -57,41 +57,14 @@ impl AgentTool for SaveNoteTool {
info!("[SaveNote] 保存笔记: {}", title);
let state = &ctx.app_state;
// 文件名安全化
let safe_title: String = title
.chars()
.map(|c| {
if c.is_alphanumeric() || c == '-' || c == '_' || c == ' ' {
c
} else {
'_'
}
})
.collect::<String>()
.replace(' ', "_");
let notes_dir = state.config.library_dir.join("notes");
if let Err(e) = std::fs::create_dir_all(&notes_dir) {
return ToolOutput::error(format!("创建笔记目录失败: {}", e));
}
let filename = format!("{}.md", safe_title);
let filepath = notes_dir.join(&filename);
let full_content = format!(
"---\ntitle: {}\ndate: {}\ngenerated_by: AstroResearch Agent\n---\n\n{}",
title,
chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
content
);
match std::fs::write(&filepath, &full_content) {
Ok(_) => ToolOutput::success(
match crate::services::note::save_note_service(&state.config.library_dir, &title, &content)
{
Ok((filename, filepath, size)) => ToolOutput::success(
format!("笔记已保存: {}", filename),
json!({
"filename": filename,
"path": filepath.to_string_lossy(),
"size": full_content.len()
"size": size
}),
),
Err(e) => ToolOutput::error(format!("保存笔记失败: {}", e)),

View File

@ -0,0 +1,197 @@
// src/agent/tools/astro/research/paper_reader_tools.rs
//
// 文献内容读取与大纲提取的智能体工具定义。
use async_trait::async_trait;
use serde_json::json;
use tracing::info;
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
// ==========================================
// 1. GetPaperOutlineTool (获取文献大纲目录)
// ==========================================
#[derive(Default)]
pub struct GetPaperOutlineTool;
#[async_trait]
impl AgentTool for GetPaperOutlineTool {
fn name(&self) -> &'static str {
"get_paper_outline"
}
fn description(&self) -> &'static str {
"获取文献的章节大纲目录,列出所有章节的序号和标题。\n\
使便 token"
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"bibcode": {
"type": "string",
"description": "文献的唯一标识符bibcode如 2006BaltA..15...69W"
}
},
"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!("[GetPaperOutline] 提取大纲: {}", bibcode);
let state = &ctx.app_state;
let mode = crate::services::paper::ReadMode::Outline { max_level: Some(3) };
match crate::services::paper::read_paper_content(
&state.db,
&state.config.library_dir,
&bibcode,
mode,
)
.await
{
Ok(res) => {
let sections = res.outline.unwrap_or_default();
if sections.is_empty() {
return ToolOutput::success(
"该文献未检测到章节标题(## 或 ### 格式)。请直接读取全文。",
json!({ "bibcode": bibcode, "sections": [] }),
);
}
let display_sections: Vec<serde_json::Value> = sections
.iter()
.map(|s| {
json!({
"index": s.index,
"level": s.level,
"heading": s.heading.clone(),
})
})
.collect();
let header = format!("文献 {} 章节大纲(共 {} 节):\n", bibcode, sections.len());
let footer = "\n\n使用 get_paper_content 并传入 section_index 或 section_name 读取指定章节内容。";
ToolOutput::success(
header + &res.content + footer,
json!({
"bibcode": bibcode,
"section_count": sections.len(),
"sections": display_sections
}),
)
}
Err(e) => ToolOutput::error(e.to_string()),
}
}
}
// ==========================================
// 2. GetPaperContentTool (获取文献具体内容)
// ==========================================
#[derive(Default)]
pub struct GetPaperContentTool;
#[async_trait]
impl AgentTool for GetPaperContentTool {
fn name(&self) -> &'static str {
"get_paper_content"
}
fn description(&self) -> &'static str {
"获取文献的具体内容(全文或指定章节内容)。\n\
section_index section_name \n\
使 get_paper_outline token "
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"bibcode": {
"type": "string",
"description": "文献的唯一标识符bibcode如 2006BaltA..15...69W"
},
"section_index": {
"type": "integer",
"description": "章节序号0-based。从 get_paper_outline 返回的 sections[index] 获取。与 section_name 二选一"
},
"section_name": {
"type": "string",
"description": "章节标题文本(支持模糊匹配,如 'introduction')。与 section_index 二选一"
}
},
"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 section_index = args
.get("section_index")
.and_then(|s| s.as_u64())
.map(|n| n as usize);
let section_name = args
.get("section_name")
.and_then(|s| s.as_str())
.map(String::from);
info!(
"[GetPaperContent] bibcode={}, section_index={:?}, section_name={:?}",
bibcode, section_index, section_name
);
let state = &ctx.app_state;
let mode = if let Some(idx) = section_index {
crate::services::paper::ReadMode::SectionIndex(idx)
} else if let Some(name) = section_name {
crate::services::paper::ReadMode::SectionName(name)
} else {
crate::services::paper::ReadMode::Full {
include_translation: false,
}
};
let label = match &mode {
crate::services::paper::ReadMode::SectionIndex(idx) => Some(format!("#{}", idx)),
crate::services::paper::ReadMode::SectionName(name) => Some(name.clone()),
_ => None,
};
let is_full = matches!(mode, crate::services::paper::ReadMode::Full { .. });
match crate::services::paper::read_paper_content(
&state.db,
&state.config.library_dir,
&bibcode,
mode,
)
.await
{
Ok(res) => ToolOutput::success(
res.content.clone(),
json!({
"bibcode": bibcode,
"chars": res.content.len(),
"section": label,
"full_text": is_full
}),
),
Err(e) => ToolOutput::error(e.to_string()),
}
}
}

View File

@ -1,166 +0,0 @@
// src/agent/tools/astro/research/paper_content.rs
// GetPaperContentTool — 文献内容读取(全文 + 按章节)
use async_trait::async_trait;
use serde_json::json;
use tracing::info;
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
pub struct GetPaperContentTool;
#[async_trait]
impl AgentTool for GetPaperContentTool {
fn name(&self) -> &str {
"get_paper_content"
}
fn description(&self) -> &str {
"读取已解析文献的 Markdown 内容。默认返回全文;可通过 section_index序号或 section_name标题名\
使 get_paper_outline token \
"
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"bibcode": {
"type": "string",
"description": "文献的唯一标识符,支持 ADS Bibcode、DOI 或 arXiv ID"
},
"section_index": {
"type": "integer",
"description": "章节序号0-based。从 get_paper_outline 返回的 sections[index] 获取。与 section_name 二选一"
},
"section_name": {
"type": "string",
"description": "章节标题名称,大小写不敏感,支持模糊匹配。如 'Introduction'、'Results'、'Discussion'。与 section_index 二选一"
}
},
"required": ["bibcode"]
})
}
fn group(&self) -> &str {
"as:research"
}
fn is_readonly(&self) -> bool {
true
}
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
true
}
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 section_index = args
.get("section_index")
.and_then(|s| s.as_u64())
.map(|n| n as usize);
let section_name = args
.get("section_name")
.and_then(|s| s.as_str())
.map(String::from);
let want_section = section_index.is_some() || section_name.is_some();
info!(
"[GetPaperContent] bibcode={}, section_index={:?}, section_name={:?}",
bibcode, section_index, section_name
);
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(
"获取文献内容失败:该文献未在本地数据库中注册,请先使用 search_papers 搜索该文献。",
),
Err(e) => return ToolOutput::error(format!("获取文献内容失败: {}", e)),
};
let md_rel = match md_opt {
Some(rel) => rel,
None => {
return ToolOutput::error(
"获取文献内容失败:该文献尚未完成结构化解析。请使用 process_paper 执行 download 和 parse 任务。",
)
}
};
let md_abs = state.config.library_dir.join(&md_rel);
if !md_abs.exists() {
return ToolOutput::error(
"获取文献内容失败:文献本地 Markdown 文件已丢失,请重新使用 process_paper 执行 parse 任务。",
);
}
let content = match std::fs::read_to_string(&md_abs) {
Ok(c) => c,
Err(e) => {
return ToolOutput::error(format!("获取文献内容失败,读取本地文件错误: {}", e))
}
};
// 按章节读取
if want_section {
let section_content = if let Some(idx) = section_index {
crate::services::section_parser::extract_section_by_index(&content, idx)
} else if let Some(ref name) = section_name {
crate::services::section_parser::extract_section_by_name(&content, name)
} else {
None
};
match section_content {
Some(sec) => {
let label = section_name
.or_else(|| section_index.map(|i| format!("#{}", i)))
.unwrap_or_default();
ToolOutput::success(
sec.clone(),
json!({
"bibcode": bibcode,
"chars": sec.len(),
"section": label,
"full_text": false
}),
)
}
None => {
let hint = if section_name.is_some() {
"请使用 get_paper_outline 查看可用的章节标题,确认章节名称拼写正确。"
} else {
"请使用 get_paper_outline 查看该文献的章节数量。"
};
ToolOutput::error(format!("获取章节内容失败:未找到匹配的章节。{}", hint))
}
}
} else {
// 返回全文(原行为)
ToolOutput::success(
content.clone(),
json!({
"bibcode": bibcode,
"chars": content.len(),
"section": null,
"full_text": true
}),
)
}
}
}

View File

@ -1,135 +0,0 @@
// src/agent/tools/astro/research/paper_outline.rs
// GetPaperOutlineTool — 提取文献章节大纲(目录)
use async_trait::async_trait;
use serde_json::json;
use tracing::info;
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
pub struct GetPaperOutlineTool;
#[async_trait]
impl AgentTool for GetPaperOutlineTool {
fn name(&self) -> &str {
"get_paper_outline"
}
fn description(&self) -> &str {
"提取已解析文献的章节大纲(目录)。返回所有章节标题、层级和序号,供后续按章节读取使用。\
get_paper_content section_index/section_name 使"
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"bibcode": {
"type": "string",
"description": "文献的唯一标识符,支持 ADS Bibcode、DOI 或 arXiv ID"
}
},
"required": ["bibcode"]
})
}
fn group(&self) -> &str {
"as:research"
}
fn is_readonly(&self) -> bool {
true
}
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
true
}
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!("[GetPaperOutline] 提取大纲: {}", bibcode);
let state = &ctx.app_state;
let paths =
match crate::api::helpers::check_paper_paths_in_db(
&state.db,
&state.config.library_dir,
&bibcode,
)
.await
{
Ok(Some((_, _, md_opt, _))) => md_opt,
Ok(None) => return ToolOutput::error(
"获取大纲失败:该文献未在本地数据库中注册。请先使用 search_papers 搜索该文献。",
),
Err(e) => return ToolOutput::error(format!("获取大纲失败: {}", e)),
};
let md_rel = match paths {
Some(rel) => rel,
None => {
return ToolOutput::error(
"获取大纲失败:该文献尚未完成结构化解析。请先使用 process_paper 执行 download 和 parse 任务。",
)
}
};
let md_abs = state.config.library_dir.join(&md_rel);
let content = match std::fs::read_to_string(&md_abs) {
Ok(c) => c,
Err(e) => return ToolOutput::error(format!("获取大纲失败,读取文件错误: {}", e)),
};
let sections = crate::services::section_parser::extract_outline(&content, 3);
if sections.is_empty() {
return ToolOutput::success(
"该文献未检测到章节标题(## 或 ### 格式)。请使用 get_paper_content 直接读取全文。",
json!({ "bibcode": bibcode, "sections": [] }),
);
}
let display_sections: Vec<serde_json::Value> = sections
.iter()
.map(|s| {
json!({
"index": s.index,
"level": s.level,
"heading": s.heading,
})
})
.collect();
let content_display = sections
.iter()
.map(|s| {
let prefix = "#".repeat(s.level);
format!(
"[{}] {} {} ({} 字符)",
s.index,
prefix,
s.heading,
s.char_end - s.char_start
)
})
.collect::<Vec<_>>()
.join("\n");
let header = format!("文献 {} 章节大纲(共 {} 节):\n", bibcode, sections.len());
let footer =
"\n\n使用 get_paper_content 并传入 section_index 或 section_name 读取指定章节内容。";
ToolOutput::success(
header + &content_display + footer,
json!({
"bibcode": bibcode,
"section_count": sections.len(),
"sections": display_sections
}),
)
}
}

View File

@ -54,9 +54,13 @@ impl AgentTool for QueryTargetTool {
info!("[QueryTarget] 查询天体: {}", 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
match crate::services::target::query_target_cached(
&state.db,
&object_name,
None,
&state.http_client,
)
.await
{
Ok(target) => {
let mut content = String::new();

View File

@ -3,7 +3,6 @@
use async_trait::async_trait;
use serde_json::json;
use tracing::info;
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
@ -87,175 +86,42 @@ impl AgentTool for ProcessPaperTool {
}
let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);
let state = &ctx.app_state;
let mut completed: Vec<String> = Vec::new();
let mut has_pdf = false;
let mut has_html = false;
let mut has_markdown = false;
let mut chunk_count = 0usize;
for task in &tasks {
match task.as_str() {
"download" => {
info!("[ProcessPaper] download: {}", bibcode);
match state
.downloader
.download_paper_service(
&state.db,
&state.config.library_dir,
&bibcode,
force,
)
.await
{
Ok(paper) => {
has_pdf = paper.has_pdf;
has_html = paper.has_html;
completed.push("download".into());
}
Err(e) => {
return ToolOutput::error(format!(
"文献 {} 下载失败(已完成: [{}]: {}",
bibcode,
completed.join(", "),
e
));
}
}
}
"parse" => {
info!("[ProcessPaper] parse: {}", bibcode);
match crate::services::parser::parse_paper_service(
&state.db,
&state.config.library_dir,
&state.qiniu,
&state.config,
&bibcode,
force,
)
.await
{
Ok(markdown) => {
has_markdown = true;
completed.push("parse".into());
info!("[ProcessPaper] parse done: {} chars", markdown.len());
}
Err(e) => {
let msg = e.to_string();
let hint = if msg.contains("请先下载") {
"未检测到已下载的本地资源文件,请确保 tasks 中 download 在 parse 之前执行。"
} else {
""
};
return ToolOutput::error(format!(
"文献 {} 解析失败(已完成: [{}]: {} {}",
bibcode,
completed.join(", "),
msg,
hint
));
}
}
}
"embed" => {
info!("[ProcessPaper] embed: {}", bibcode);
// 读取 markdown 文件
let paths = crate::api::helpers::check_paper_paths_in_db(
&state.db,
&state.config.library_dir,
&bibcode,
)
.await;
let md_rel = match paths {
Ok(Some((_, _, Some(md), _))) => md,
Ok(_) => {
return ToolOutput::error(format!(
"文献 {} 向量化失败(已完成: [{}]):未找到已解析的 Markdown 文件,请确保 tasks 中 parse 在 embed 之前执行",
bibcode,
completed.join(", ")
));
}
Err(e) => {
return ToolOutput::error(format!(
"文献 {} 向量化失败(已完成: [{}]: {}",
bibcode,
completed.join(", "),
e
));
}
};
let md_abs = state.config.library_dir.join(&md_rel);
let markdown_text = match std::fs::read_to_string(&md_abs) {
Ok(t) => t,
Err(e) => {
return ToolOutput::error(format!(
"文献 {} 向量化失败(已完成: [{}]):读取 Markdown 文件错误: {}",
bibcode,
completed.join(", "),
e
));
}
};
match crate::services::rag::ingest_paper(
&state.db,
&state.embedding,
&bibcode,
&markdown_text,
None, // 使用默认 ChunkerConfig结构化分块
)
.await
{
Ok(n) => {
chunk_count = n;
completed.push("embed".into());
info!("[ProcessPaper] embed done: {} chunks", n);
}
Err(e) => {
return ToolOutput::error(format!(
"文献 {} 向量化失败(已完成: [{}]: {}",
bibcode,
completed.join(", "),
e
));
}
}
}
other => {
return ToolOutput::error(format!(
"无效的任务类型 '{}',仅支持: download, parse, embed",
other
));
}
match crate::services::pipeline::process_paper_pipeline(state, &bibcode, &tasks, force)
.await
{
Ok(res) => {
let msg = format!(
"文献 {} 处理完成。已执行: [{}]。PDF: {}, HTML: {}, Markdown: {}, 向量块数: {}",
bibcode,
res.completed.join(""),
res.has_pdf,
res.has_html,
res.has_markdown,
res.chunk_count
);
ToolOutput::success(
msg,
json!({
"bibcode": res.bibcode,
"completed": res.completed,
"has_pdf": res.has_pdf,
"has_html": res.has_html,
"has_markdown": res.has_markdown,
"chunk_count": res.chunk_count
}),
)
}
Err(e) => {
let msg = e.to_string();
let hint = if msg.contains("请先下载") {
"(提示:未检测到已下载的本地资源文件,请确保 tasks 中 download 在 parse 之前执行)"
} else {
""
};
ToolOutput::error(format!("文献 {} 处理失败: {}{}", bibcode, msg, hint))
}
}
let msg = format!(
"文献 {} 处理完成。已执行: [{}]。PDF: {}, HTML: {}, Markdown: {}, 向量块数: {}",
bibcode,
completed.join(""),
has_pdf,
has_html,
has_markdown,
chunk_count
);
ToolOutput::success(
msg,
json!({
"bibcode": bibcode,
"completed": completed,
"has_pdf": has_pdf,
"has_html": has_html,
"has_markdown": has_markdown,
"chunk_count": chunk_count
}),
)
}
}

View File

@ -182,10 +182,11 @@ impl AgentTool for RunBashTool {
/// 安全命令白名单(只读或数据处理类命令)。
/// 这些命令的首个单词匹配时自动允许执行,不再进行黑名单检查。
/// 注意: awk 可通过 system() 执行任意命令tee 可写入任意路径,已移除。
const SAFE_COMMANDS: &[&str] = &[
"ls", "cat", "head", "tail", "find", "grep", "wc", "echo", "pwd", "sort", "uniq", "cut", "tr",
"awk", "sed", "jq", "diff", "file", "stat", "du", "df", "env", "printenv", "which", "basename",
"dirname", "realpath", "readlink", "xargs", "tee", "date", "sleep", "true", "false",
"sed", "jq", "diff", "file", "stat", "du", "df", "env", "printenv", "which", "basename",
"dirname", "realpath", "readlink", "xargs", "date", "sleep", "true", "false",
];
/// 危险命令黑名单(交互式、提权、远程访问、破坏性命令)。
@ -234,7 +235,6 @@ const DANGEROUS_COMMANDS: &[&str] = &[
"nmap",
// 登录相关
"login",
"su",
"sulogin",
];
@ -346,6 +346,37 @@ fn has_command_substitution_bypass(command: &str) -> bool {
false
}
/// 语言运行时白名单 — 只允许直接执行脚本文件,禁止 -c/-e 标志执行字符串代码。
const LANGUAGE_RUNTIMES: &[&str] = &[
"python", "python3", "ruby", "node", "perl", "php", "lua", "tcl",
];
/// 检测语言运行时是否通过 -c/-e 标志执行内联代码(可绕过沙箱)。
fn has_language_runtime_inline_code(first_word: &str, command: &str) -> bool {
if !LANGUAGE_RUNTIMES.contains(&first_word) {
return false;
}
let args: Vec<&str> = command.split_whitespace().collect();
// 跳过环境变量赋值和首词本身,检查后续参数
for arg in &args[1..] {
if arg.starts_with('-') {
// 匹配 -c, -e, -exec, --command, --eval 等
let flag = arg.trim_start_matches('-');
if flag == "c"
|| flag == "e"
|| flag.starts_with("c=")
|| flag.starts_with("e=")
|| flag == "exec"
|| flag == "command"
|| flag == "eval"
{
return true;
}
}
}
false
}
/// 检查命令是否包含危险的参数模式(用于非白名单非黑名单命令的二次检查)。
fn check_dangerous_args(command: &str) -> Option<String> {
let lower = command.to_lowercase();
@ -386,6 +417,13 @@ fn validate_bash_command(command: &str) -> Option<String> {
return Some("不允许执行空命令".to_string());
}
// 彻底禁止 Shell 逻辑操作符,防止命令拼接注入
if has_shell_operators(trimmed) {
return Some(
"命令包含 Shell 操作符(如 ;, |, &, <, > 等)。为了安全起见,如需执行复杂逻辑或管道拼接,请先将命令写入 .sh 脚本文件,然后通过 bash 运行该脚本。".to_string(),
);
}
// 检测命令替换绕过
if has_command_substitution_bypass(trimmed) {
return Some(
@ -399,6 +437,14 @@ fn validate_bash_command(command: &str) -> Option<String> {
return Some("无法解析命令名称".to_string());
}
// 拦截语言运行时的 -c/-e 内联代码执行(可绕过沙箱)
if has_language_runtime_inline_code(&first_word, trimmed) {
return Some(format!(
"不允许通过 {} 的 -c/-e 标志执行内联代码。请将代码写入脚本文件后执行。",
first_word
));
}
// 白名单优先:安全命令放行前,先做危险参数二次检查。
// 白名单命令find/grep/xargs/sed/tee 等)可携带破坏性参数绕过黑名单,
// 故此处对 -delete、xargs rm、sed -i、写入系统目录等模式做拦截。
@ -426,6 +472,17 @@ fn validate_bash_command(command: &str) -> Option<String> {
None
}
/// 检查命令是否包含危险的 Shell 操作符(拼接、重定向、命令替换等)
fn has_shell_operators(command: &str) -> bool {
let operators = [";", "&", "|", "<", ">", "$(", "`"];
for op in operators {
if command.contains(op) {
return true;
}
}
false
}
/// 检查命令是否需要用户权限确认。
/// 安全白名单中的命令不需确认,其他命令需要。
pub fn bash_needs_permission(command: &str) -> bool {
@ -487,9 +544,11 @@ mod tests {
}
#[test]
fn test_whitelist_allows_safe_commands_with_pipes() {
assert!(validate_bash_command("cat file.txt | grep foo | wc -l").is_none());
assert!(validate_bash_command("find . -name '*.rs' | xargs wc").is_none());
fn test_whitelist_blocks_shell_operators() {
assert!(validate_bash_command("cat file.txt | grep foo | wc -l").is_some());
assert!(validate_bash_command("find . -name '*.rs' | xargs wc").is_some());
assert!(validate_bash_command("ls; rm -rf /").is_some());
assert!(validate_bash_command("echo hello > file.txt").is_some());
}
// ── validate_bash_command: 白名单命令的危险参数二次检查 ──
@ -501,12 +560,7 @@ mod tests {
"find / -delete",
"find . -name x -exec rm -rf {} \\;",
"find . -execdir rm {} +",
"cat files.txt | xargs rm -rf",
"sed -i 's/old/new/g' /etc/hosts",
"echo bad > /etc/passwd",
"tee /etc/passwd",
"echo key >> ~/.ssh/authorized_keys",
"find / -name foo > /var/log/x",
// 多空格/制表符变形绕过测试
"xargs rm",
"sed\t-i\tfile",
@ -607,10 +661,10 @@ mod tests {
}
#[test]
fn test_allows_dollar_parens_in_args() {
// $() 在参数位置(非命令首词)是合法的,如命令替换结果作为参数
assert!(validate_bash_command("echo $(date)").is_none());
assert!(validate_bash_command("grep $(cat pattern.txt) file.txt").is_none());
fn test_blocks_dollar_parens_in_args() {
// 由于为了彻底防御拼接和复杂语法执行,现在全面禁用了 $() 和 backtick
assert!(validate_bash_command("echo $(date)").is_some());
assert!(validate_bash_command("grep $(cat pattern.txt) file.txt").is_some());
}
// ── validate_bash_command: 危险参数模式 ──
@ -643,13 +697,27 @@ mod tests {
#[test]
fn test_allows_common_dev_tools() {
// 开发常用命令不在白名单但也不在黑名单,应默认放行
assert!(validate_bash_command("python script.py").is_none());
// 开发常用命令不在白名单但也不在黑名单,需权限确认
assert!(validate_bash_command("git status").is_none());
assert!(validate_bash_command("cargo build").is_none());
assert!(validate_bash_command("npm test").is_none());
}
#[test]
fn test_blocks_language_runtime_inline_code() {
// 语言运行时 -c/-e 标志应被拦截
assert!(validate_bash_command("python -c 'import os; os.system(\"rm -rf /\")'").is_some());
assert!(validate_bash_command("python3 -e 'print(1)'").is_some());
assert!(validate_bash_command("ruby -e 'puts 1'").is_some());
assert!(validate_bash_command("node -c 'console.log(1)'").is_some());
assert!(validate_bash_command("perl -e 'print 1'").is_some());
assert!(validate_bash_command("lua -e 'print(1)'").is_some());
// 但直接执行脚本文件应允许
assert!(validate_bash_command("python script.py").is_none());
assert!(validate_bash_command("python3 /path/to/script.py arg1").is_none());
assert!(validate_bash_command("node index.js").is_none());
}
#[test]
fn test_network_commands_blocking() {
// 网络命令始终被阻止。

View File

@ -67,7 +67,11 @@ impl AgentTool for GlobFilesTool {
return ToolOutput::error("路径不在允许的沙箱范围内");
}
let glob_pattern = search_path.join(pattern);
let glob_pattern = if std::path::Path::new(pattern).is_absolute() {
return ToolOutput::error("glob 模式不能使用绝对路径");
} else {
search_path.join(pattern)
};
let pattern_str = glob_pattern.to_string_lossy().to_string();
match glob::glob(&pattern_str) {

View File

@ -113,6 +113,17 @@ impl AgentTool for ReadFileTool {
}
// ── 去重检查结束 ──
// 预检文件大小,避免 OOM
const MAX_READ_BYTES: u64 = 10 * 1024 * 1024; // 10MB
if let Ok(meta) = std::fs::metadata(&path) {
if meta.len() > MAX_READ_BYTES {
return ToolOutput::error(format!(
"文件过大 ({} MB),超过 10MB 限制。请使用 read_file 的 offset/limit 参数分段读取。",
meta.len() / (1024 * 1024)
));
}
}
match std::fs::read_to_string(&path) {
Ok(content) => {
let total_lines = content.lines().count();

View File

@ -43,7 +43,14 @@ pub fn is_path_allowed(path: &Path, ctx: &ToolContext) -> bool {
/// 检查路径字符串是否包含穿越尝试
pub fn has_path_traversal(path_str: &str) -> bool {
path_str.contains("..") || path_str.contains('~')
// Check for path components that indicate traversal
let parts: Vec<&str> = path_str.split('/').collect();
for part in &parts {
if *part == ".." || *part == "~" {
return true;
}
}
false
}
/// 规范化用户提供的路径

View File

@ -13,7 +13,8 @@
use async_trait::async_trait;
use serde_json::json;
use std::sync::{Arc, RwLock};
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::agent::runtime::file_cache::FileStateCache;
use crate::agent::skills::SkillRegistry;
@ -39,8 +40,7 @@ pub use astro::research::citation_network::GetCitationNetworkTool;
pub use astro::research::library_search::SearchLocalLibraryTool;
pub use astro::research::metadata::GetPaperMetadataTool;
pub use astro::research::note::SaveNoteTool;
pub use astro::research::paper_content::GetPaperContentTool;
pub use astro::research::paper_outline::GetPaperOutlineTool;
pub use astro::research::paper::{GetPaperContentTool, GetPaperOutlineTool};
pub use astro::research::rag::RagSearchTool;
pub use astro::research::target::QueryTargetTool;
pub use astro::system::process::ProcessPaperTool;
@ -668,11 +668,15 @@ impl ToolRegistry {
/// 截断文本到指定最大字符数
pub fn truncate_content(s: &str, max_chars: usize) -> String {
if s.len() <= max_chars {
let char_count = s.chars().count();
if char_count <= max_chars {
s.to_string()
} else {
let truncated: String = s.chars().take(max_chars).collect();
format!("{}\n\n[... 内容已截断,共 {} 字符 ...]", truncated, s.len())
format!(
"{}\n\n[... 内容已截断,共 {} 字符 ...]",
truncated, char_count
)
}
}

View File

@ -41,7 +41,23 @@ pub fn maybe_persist_tool_result(
);
}
let file_path = tool_results_dir.join(format!("{}.txt", tool_call_id));
// Sanitize tool_call_id to prevent path traversal
let safe_id: String = tool_call_id
.chars()
.filter(|c| c.is_alphanumeric() || *c == '_' || *c == '-')
.collect();
if safe_id.is_empty() {
let truncated: String = content.chars().take(max_chars).collect();
return (
format!(
"{}...\n[输出已截断,原始长度: {} 字符]",
truncated,
content.len()
),
None,
);
}
let file_path = tool_results_dir.join(format!("{}.txt", safe_id));
// 独占创建——只在文件不存在时写入,保证幂等性
let written = match std::fs::OpenOptions::new()
@ -161,7 +177,7 @@ mod tests {
let large2 = "b".repeat(5000);
let (_, path1) = maybe_persist_tool_result(&large1, "call_3", 100, &dir);
let (content2, path2) = maybe_persist_tool_result(&large2, "call_3", 100, &dir);
let (_content2, path2) = maybe_persist_tool_result(&large2, "call_3", 100, &dir);
assert!(path1.is_some());
assert!(path2.is_some());

View File

@ -53,10 +53,11 @@ impl AgentTool for SearchHistoryTool {
}
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
let query = match args.get("query").and_then(|v| v.as_str()) {
let raw_query = match args.get("query").and_then(|v| v.as_str()) {
Some(q) if !q.is_empty() => q,
_ => return ToolOutput::error("缺少 query 参数"),
};
let query = crate::services::search::sanitize_fts5_query(raw_query);
let scope = args.get("scope").and_then(|v| v.as_str()).unwrap_or("all");
let db = &ctx.app_state.db;
@ -70,7 +71,7 @@ impl AgentTool for SearchHistoryTool {
JOIN agent_sessions s ON s.session_id = fts.session_id \
WHERE agent_sessions_fts MATCH $1 ORDER BY rank LIMIT 10",
)
.bind(query)
.bind(&query)
.fetch_all(db)
.await
{
@ -100,7 +101,7 @@ impl AgentTool for SearchHistoryTool {
WHERE agent_messages_fts MATCH $1 AND m.active = 1 \
ORDER BY rank LIMIT 10",
)
.bind(query)
.bind(&query)
.fetch_all(db)
.await
{

View File

@ -8,8 +8,9 @@
use async_trait::async_trait;
use serde_json::json;
use std::sync::{Arc, RwLock};
use tracing::{info, warn};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::info;
use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
use crate::agent::skills::{substitute_variables, SkillRegistry};
@ -76,14 +77,8 @@ impl AgentTool for LoadSkillTool {
info!("[LoadSkill] 加载 skill: {}", skill_name);
// 从缓存注册表读取单语句guard 自动 drop不跨越 await
let skill = match self.registry.read() {
Ok(reg) => reg.get_skill(&skill_name).cloned(),
Err(e) => {
warn!("[LoadSkill] RwLock poisoned: {:?}", e);
return ToolOutput::error("技能注册表不可用(内部锁异常),请稍后重试");
}
};
// 从缓存注册表读取
let skill = self.registry.read().await.get_skill(&skill_name).cloned();
let skill = match skill {
Some(s) => s,
None => {
@ -94,29 +89,20 @@ impl AgentTool for LoadSkillTool {
}
};
// 记录使用统计单语句write guard 自动 drop
if let Ok(mut reg) = self.registry.write() {
// 记录使用统计
{
let mut reg = self.registry.write().await;
reg.record_usage(&skill_name);
} else {
warn!("[LoadSkill] 无法记录使用统计RwLock poisoned");
}
// 变量替换
let session_id = if ctx.app_state.config.database_url.contains("session") {
"current"
// 变量替换——有 session_id 时注入当前会话 ID
let session_id: Option<&str> = if ctx.session_id.is_empty() {
None
} else {
""
Some(&ctx.session_id)
};
let body = substitute_variables(
&skill.body,
&skill.skill_dir,
if session_id.is_empty() {
None
} else {
Some(session_id)
},
);
let body = substitute_variables(&skill.body, &skill.skill_dir, session_id);
// 构建 skill 目录的绝对路径(用于 "Base directory" 前缀)
let skill_dir_path = skill

View File

@ -210,7 +210,7 @@ mod tests {
use crate::services::translation::Dictionary;
use crate::Config;
use std::path::PathBuf;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::Arc;
/// 构建一个最小化的 ToolContext 用于测试(不依赖真实外部服务)
async fn make_test_tool_context() -> ToolContext {
@ -227,16 +227,18 @@ mod tests {
"test_key".into(),
"http://localhost".into(),
"test-model".into(),
);
)
.unwrap();
let medium_llm = llm.clone();
let fast_llm = llm.clone();
let embedding = EmbeddingClient::new(
"test_key".into(),
"http://localhost".into(),
"test-embed".into(),
);
let ads = AdsClient::new("test_token".into());
let arxiv = ArxivClient::new();
)
.unwrap();
let ads = AdsClient::new("test_token".into()).unwrap();
let arxiv = ArxivClient::new().unwrap();
let qiniu = QiniuClient::new(
"test_ak".into(),
"test_sk".into(),
@ -259,22 +261,27 @@ mod tests {
vision_llm: None,
embedding,
downloader,
http_client: reqwest::Client::new(),
harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())),
batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(Mutex::new(std::collections::HashSet::new())),
skill_registry: Arc::new(RwLock::new(crate::agent::skills::SkillRegistry::new(
PathBuf::from("/tmp/test_skills"),
))),
pending_questions: Arc::new(Mutex::new(std::collections::HashMap::new())),
pending_permissions: Arc::new(Mutex::new(std::collections::HashMap::new())),
session_permission_checkers: Arc::new(RwLock::new(std::collections::HashMap::new())),
cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
skill_registry: Arc::new(tokio::sync::RwLock::new(
crate::agent::skills::SkillRegistry::new(PathBuf::from("/tmp/test_skills")),
)),
pending_questions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
pending_permissions: Arc::new(
tokio::sync::Mutex::new(std::collections::HashMap::new()),
),
session_permission_checkers: Arc::new(tokio::sync::RwLock::new(
std::collections::HashMap::new(),
)),
sse_broadcast: None,
memory_manager: Arc::new(tokio::sync::Mutex::new(
crate::agent::memory::MemoryManager::new(PathBuf::from("/tmp/test_memory")),
)),
sessions: Arc::new(Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(Mutex::new(std::collections::HashMap::new())),
sessions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
});
ToolContext {

View File

@ -10,7 +10,6 @@ use axum::{
};
use futures_util::stream::Stream;
use serde::{Deserialize, Serialize};
use sqlx::Row;
use std::convert::Infallible;
use std::sync::Arc;
use tracing::{error, info};
@ -129,7 +128,8 @@ pub async fn chat_agent(
.library_dir
.join(".agent_images")
.join("uploads");
std::fs::create_dir_all(&upload_dir)
tokio::fs::create_dir_all(&upload_dir)
.await
.map_err(|e| AppError::internal(format!("创建上传目录失败: {}", e)))?;
let filename = format!("{}.{}", uuid::Uuid::new_v4(), ext);
let filepath = upload_dir.join(&filename);
@ -137,7 +137,8 @@ pub async fn chat_agent(
let bytes = general_purpose::STANDARD
.decode(&img.data)
.map_err(|e| AppError::bad_request(format!("图片 base64 解码失败: {}", e)))?;
std::fs::write(&filepath, &bytes)
tokio::fs::write(&filepath, &bytes)
.await
.map_err(|e| AppError::internal(format!("保存图片失败: {}", e)))?;
let rel = filepath
.strip_prefix(&state.config.library_dir)
@ -241,143 +242,30 @@ pub struct SessionListParams {
pub offset: Option<i64>,
}
#[derive(Debug, Serialize)]
pub struct SessionSummary {
pub session_id: String,
pub title: String,
pub model: String,
pub mode: String,
pub turn_count: i32,
pub summary: Option<String>,
pub created_at: String,
pub updated_at: String,
}
pub async fn list_sessions(
State(state): State<Arc<AppState>>,
Query(params): Query<SessionListParams>,
) -> ApiResult<Json<Vec<SessionSummary>>> {
) -> ApiResult<Json<Vec<crate::services::session::SessionSummary>>> {
let limit = params.limit.unwrap_or(50);
let offset = params.offset.unwrap_or(0);
let rows = sqlx::query(
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \
FROM agent_sessions \
WHERE deleted_at IS NULL \
ORDER BY updated_at DESC \
LIMIT ? OFFSET ?",
)
.bind(limit)
.bind(offset)
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("查询会话列表失败: {}", e)))?;
let sessions: Vec<SessionSummary> = rows
.iter()
.map(|r| SessionSummary {
session_id: r.get(0),
title: r.get(1),
model: r.get(2),
mode: r.get(3),
turn_count: r.get(4),
summary: r.get(5),
created_at: r.get(6),
updated_at: r.get(7),
})
.collect();
let sessions = crate::services::session::list_sessions_service(&state.db, limit, offset)
.await
.map_err(|e| AppError::internal(format!("查询会话列表失败: {}", e)))?;
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 agent_name: String,
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>,
) -> ApiResult<Json<SessionDetail>> {
// 查询会话元信息
let session_row = sqlx::query(
"SELECT session_id, title, model, mode, turn_count, summary, 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| AppError::internal(format!("查询会话失败: {}", e)))?
.ok_or_else(|| AppError::not_found(format!("会话 {} 不存在", session_id)))?;
) -> ApiResult<Json<crate::services::session::SessionDetail>> {
let detail = crate::services::session::get_session_detail_service(&state.db, &session_id)
.await
.map_err(|e| AppError::internal(format!("查询会话详情失败: {}", e)))?
.ok_or_else(|| AppError::not_found(format!("会话 {} 不存在", session_id)))?;
let session = SessionSummary {
session_id: session_row.get(0),
title: session_row.get(1),
model: session_row.get(2),
mode: session_row.get(3),
turn_count: session_row.get(4),
summary: session_row.get(5),
created_at: session_row.get(6),
updated_at: session_row.get(7),
};
// 查询消息列表(包含 lead 和 subagent 消息,前端按 agent_name/metadata 区分渲染)
let msg_rows = sqlx::query(
"SELECT id, agent_name, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, created_at \
FROM agent_messages \
WHERE session_id = ? AND active = 1 \
ORDER BY id ASC"
)
.bind(&session_id)
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("查询消息列表失败: {}", e)))?;
let messages: Vec<MessageRecord> = msg_rows
.iter()
.map(|r| {
let tool_calls_json: Option<String> = r.get(7);
let metadata_json: Option<String> = r.get(10);
MessageRecord {
id: r.get(0),
agent_name: r.get(1),
turn_index: r.get(2),
step_index: r.get(3),
role: r.get(4),
content: r.get(5),
thought: r.get(6),
tool_calls: tool_calls_json.and_then(|s| serde_json::from_str(&s).ok()),
tool_call_id: r.get(8),
token_count: r.get(9),
metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
created_at: r.get(11),
}
})
.collect();
Ok(Json(SessionDetail { session, messages }))
Ok(Json(detail))
}
// ── DELETE /api/chat/sessions/:id ──
@ -387,15 +275,11 @@ pub async fn delete_session(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
) -> ApiResult<Json<serde_json::Value>> {
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| AppError::internal(format!("删除会话失败: {}", e)))?;
let success = crate::services::session::delete_session_service(&state.db, &session_id)
.await
.map_err(|e| AppError::internal(format!("删除会话失败: {}", e)))?;
if result.rows_affected() == 0 {
if !success {
return Err(AppError::not_found(format!(
"会话 {} 不存在或已删除",
session_id
@ -414,9 +298,8 @@ pub async fn stop_agent(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
) -> ApiResult<Json<serde_json::Value>> {
if let Ok(mut cancelled) = state.cancelled_runs.lock() {
cancelled.insert(session_id.clone());
}
let mut cancelled = state.cancelled_runs.lock().await;
cancelled.insert(session_id.clone());
info!("已接收并记录手动中止请求,会话 ID: {}", session_id);
Ok(Json(
serde_json::json!({ "status": "stopping", "session_id": session_id }),
@ -426,130 +309,26 @@ pub async fn stop_agent(
// ── GET /api/chat/metrics ──
// 返回聚合的智能体运行指标
#[derive(Debug, Serialize)]
pub struct AgentMetricsResponse {
pub total_sessions: i64,
pub total_tool_calls: i64,
pub tool_call_breakdown: serde_json::Value,
pub avg_steps_per_session: f64,
pub error_rate: f64,
}
pub async fn get_agent_metrics(
State(state): State<Arc<AppState>>,
) -> ApiResult<Json<AgentMetricsResponse>> {
// 总会话数
let total_sessions: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions WHERE deleted_at IS NULL")
.fetch_one(&state.db)
.await
.unwrap_or(0);
) -> ApiResult<Json<crate::services::session::AgentMetrics>> {
let metrics = crate::services::session::get_agent_metrics_service(&state.db)
.await
.map_err(|e| AppError::internal(format!("获取系统指标失败: {}", e)))?;
// 工具调用统计(从审计日志聚合)
let total_tool_calls: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_audit_log WHERE status IN ('OK', 'FAIL')")
.fetch_one(&state.db)
.await
.unwrap_or(0);
// 各工具调用次数
let breakdown_rows: Vec<(String, i64)> = sqlx::query_as(
"SELECT COALESCE(tool_name, 'unknown'), COUNT(*) as cnt \
FROM agent_audit_log \
WHERE status IN ('OK', 'FAIL') \
GROUP BY tool_name \
ORDER BY cnt DESC",
)
.fetch_all(&state.db)
.await
.unwrap_or_default();
let tool_call_breakdown: serde_json::Value = breakdown_rows
.iter()
.map(|(name, cnt)| serde_json::json!({ name: cnt }))
.fold(serde_json::json!({}), |mut acc, v| {
if let serde_json::Value::Object(map) = &mut acc {
if let serde_json::Value::Object(v_map) = v {
for (k, val) in v_map {
map.insert(k.clone(), val.clone());
}
}
}
acc
});
// 平均步数
let avg_steps: f64 = sqlx::query_scalar(
"SELECT COALESCE(AVG(CAST(turn_count AS REAL)), 0.0) \
FROM agent_sessions WHERE deleted_at IS NULL",
)
.fetch_one(&state.db)
.await
.unwrap_or(0.0);
// 错误率
let total_errors: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_audit_log WHERE status = 'FAIL'")
.fetch_one(&state.db)
.await
.unwrap_or(0);
let error_rate = if total_tool_calls > 0 {
total_errors as f64 / total_tool_calls as f64
} else {
0.0
};
Ok(Json(AgentMetricsResponse {
total_sessions,
total_tool_calls,
tool_call_breakdown,
avg_steps_per_session: avg_steps,
error_rate,
}))
Ok(Json(metrics))
}
// ── GET /api/chat/sessions/:id/audit ──
// 返回指定会话的审计日志
#[derive(Debug, Serialize)]
pub struct AuditLogEntry {
pub id: i64,
pub step: i32,
pub tool_name: Option<String>,
pub status: String,
pub elapsed_ms: i32,
pub output_preview: Option<String>,
pub created_at: String,
}
pub async fn get_session_audit(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
) -> ApiResult<Json<Vec<AuditLogEntry>>> {
let rows = sqlx::query(
"SELECT id, step, tool_name, status, elapsed_ms, output_preview, created_at \
FROM agent_audit_log \
WHERE session_id = ? \
ORDER BY id ASC",
)
.bind(&session_id)
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("查询审计日志失败: {}", e)))?;
let entries: Vec<AuditLogEntry> = rows
.iter()
.map(|r| AuditLogEntry {
id: r.get(0),
step: r.get(1),
tool_name: r.get(2),
status: r.get(3),
elapsed_ms: r.get(4),
output_preview: r.get(5),
created_at: r.get(6),
})
.collect();
) -> ApiResult<Json<Vec<crate::services::session::AuditLogEntry>>> {
let entries = crate::services::session::get_session_audit_service(&state.db, &session_id)
.await
.map_err(|e| AppError::internal(format!("查询审计日志失败: {}", e)))?;
Ok(Json(entries))
}
@ -570,12 +349,7 @@ pub async fn answer_question(
) -> ApiResult<Json<serde_json::Value>> {
use crate::agent::tools::ask_user::UserAnswer;
let mut pending = match state.pending_questions.lock() {
Ok(p) => p,
Err(_) => {
return Err(AppError::internal("服务器内部状态异常,请稍后重试"));
}
};
let mut pending = state.pending_questions.lock().await;
let question_id = req.question_id.clone();
match pending.remove(&question_id) {
@ -608,10 +382,7 @@ pub async fn answer_question(
pub async fn get_pending_questions(
State(state): State<Arc<AppState>>,
) -> Json<Vec<serde_json::Value>> {
let pending = match state.pending_questions.lock() {
Ok(p) => p,
Err(_) => return Json(Vec::new()),
};
let pending = state.pending_questions.lock().await;
let questions: Vec<serde_json::Value> = pending
.iter()
.map(|(id, pq)| {
@ -630,10 +401,7 @@ pub async fn respond_permission(
Path(session_id): Path<String>,
Json(req): Json<super::PermissionResponse>,
) -> ApiResult<Json<serde_json::Value>> {
let mut perms = match state.pending_permissions.lock() {
Ok(p) => p,
Err(_) => return Err(AppError::internal("内部状态异常")),
};
let mut perms = state.pending_permissions.lock().await;
// 按 tool_call_id 查找匹配的权限请求
let perm_id = perms
@ -667,10 +435,7 @@ pub async fn respond_permission(
pub async fn get_pending_permissions(
State(state): State<Arc<AppState>>,
) -> Json<Vec<serde_json::Value>> {
let perms = match state.pending_permissions.lock() {
Ok(p) => p,
Err(_) => return Json(Vec::new()),
};
let perms = state.pending_permissions.lock().await;
let result: Vec<serde_json::Value> = perms
.iter()
.map(|(id, p)| {

View File

@ -25,9 +25,32 @@ pub struct LoginResponse {
// 辅助函数:从 Cookie 字符串中解析出 session_id 的值
fn parse_session_cookie(cookie_str: &str) -> Option<String> {
for cookie in cookie_str.split(';') {
let parts: Vec<&str> = cookie.split('=').map(|s| s.trim()).collect();
if parts.len() == 2 && parts[0] == "session_id" {
return Some(parts[1].to_string());
let mut parts = cookie.splitn(2, '=');
if let (Some(name), Some(value)) = (parts.next(), parts.next()) {
if name.trim() == "session_id" {
return Some(value.trim().to_string());
}
}
}
None
}
/// 从请求头中提取认证 Token。
/// 优先从 Authorization: Bearer 头读取,其次从 Cookie 中的 session_id 读取。
fn extract_auth_token(headers: &HeaderMap) -> Option<String> {
// 1. Authorization: Bearer <token>
if let Some(auth_str) = headers
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
{
if let Some(rest) = auth_str.strip_prefix("Bearer ") {
return Some(rest.to_string());
}
}
// 2. Cookie: session_id=<token>
if let Some(cookie_header) = headers.get(header::COOKIE) {
if let Ok(cookie_str) = cookie_header.to_str() {
return parse_session_cookie(cookie_str);
}
}
None
@ -53,25 +76,6 @@ fn verify_password(input: &str, expected: &str) -> bool {
input_bytes.ct_eq(expected_bytes).into()
}
#[cfg(test)]
mod tests {
use super::verify_password;
#[test]
fn test_verify_password_correct() {
assert!(verify_password("secret123", "secret123"));
assert!(verify_password("", ""));
}
#[test]
fn test_verify_password_wrong() {
assert!(!verify_password("secret123", "secret456"));
assert!(!verify_password("secret", "secret123"));
assert!(!verify_password("secret123", "secret"));
assert!(!verify_password("wrong", ""));
}
}
// 辅助函数清除过期的会话记录24小时未活跃即过期
fn prune_expired_sessions(
sessions: &mut std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>,
@ -89,53 +93,63 @@ fn prune_expired_sessions(
});
}
/// 最大会话数上限,超过后拒绝新登录以防止内存耗尽
const MAX_SESSIONS: usize = 1000;
// 登录接口:验证密码成功后,写入 HttpOnly Cookie
pub async fn login(
State(state): State<Arc<AppState>>,
headers: axum::http::header::HeaderMap,
_headers: axum::http::header::HeaderMap,
axum::extract::ConnectInfo(addr): axum::extract::ConnectInfo<std::net::SocketAddr>,
Json(body): Json<LoginRequest>,
) -> Result<impl IntoResponse, AppError> {
// 提取客户端 IP用于速率限制
let client_ip = headers
.get("x-forwarded-for")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.split(',').next())
.unwrap_or("unknown")
.trim()
.to_string();
// 使用真实连接 IP 做限速X-Forwarded-For 可伪造,不应用于安全判断)
let client_ip = addr.ip().to_string();
// 检查速率限制5 次失败 / 5 分钟窗口
const MAX_FAILURES: u32 = 5;
const WINDOW_SECS: u64 = 300;
{
if let Ok(mut limiter) = state.login_rate_limiter.lock() {
if let Some((count, first_fail)) = limiter.get(&client_ip) {
let elapsed = first_fail.elapsed().as_secs();
if elapsed < WINDOW_SECS && *count >= MAX_FAILURES {
let remaining = WINDOW_SECS - elapsed;
return Err(AppError::unauthorized(format!(
"登录尝试过于频繁,请在 {} 秒后重试",
remaining
)));
}
if elapsed >= WINDOW_SECS {
limiter.remove(&client_ip);
}
let mut limiter = state.login_rate_limiter.lock().await;
// 批量清理过期条目,防止内存无限增长
limiter.retain(|_, (_, first_fail)| first_fail.elapsed().as_secs() < WINDOW_SECS);
if let Some((count, first_fail)) = limiter.get(&client_ip) {
let elapsed = first_fail.elapsed().as_secs();
if elapsed < WINDOW_SECS && *count >= MAX_FAILURES {
let remaining = WINDOW_SECS - elapsed;
return Err(AppError::unauthorized(format!(
"登录尝试过于频繁,请在 {} 秒后重试",
remaining
)));
}
}
}
let password = body.password.unwrap_or_default();
let password = body
.password
.ok_or_else(|| AppError::bad_request("缺少密码字段"))?;
if password.is_empty() {
return Err(AppError::bad_request("密码不能为空"));
}
if verify_password(&password, &state.config.admin_password) {
// 成功:清除该 IP 的失败记录
if let Ok(mut limiter) = state.login_rate_limiter.lock() {
limiter.remove(&client_ip);
}
state.login_rate_limiter.lock().await.remove(&client_ip);
let token = uuid::Uuid::new_v4().to_string();
// 异常安全地记录活跃会话及其当前时间
if let Ok(mut sessions) = state.sessions.lock() {
{
let mut sessions = state.sessions.lock().await;
// 先清理过期会话
prune_expired_sessions(&mut sessions);
// 检查会话上限
if sessions.len() >= MAX_SESSIONS {
return Err(AppError::unauthorized(format!(
"活跃会话数已达上限 ({}), 请稍后重试",
MAX_SESSIONS
)));
}
sessions.insert(token.clone(), chrono::Utc::now());
}
@ -159,13 +173,11 @@ pub async fn login(
))
} else {
// 失败:记录该 IP 的失败次数
if let Ok(mut limiter) = state.login_rate_limiter.lock() {
let entry = limiter
.entry(client_ip)
.or_insert((0, std::time::Instant::now()));
entry.0 += 1;
// 重置窗口起点(从首次失败开始计算)
}
let mut limiter = state.login_rate_limiter.lock().await;
let entry = limiter
.entry(client_ip)
.or_insert((0, std::time::Instant::now()));
entry.0 += 1;
Err(AppError::unauthorized("密码错误"))
}
}
@ -175,33 +187,10 @@ pub async fn logout(
State(state): State<Arc<AppState>>,
req: Request<axum::body::Body>,
) -> ApiResult<impl IntoResponse> {
let mut token_to_remove = None;
// 优先从 Authorization 提取 Token向前兼容旧 Bearer Header
let auth_header = req
.headers()
.get(axum::http::header::AUTHORIZATION)
.and_then(|val| val.to_str().ok());
if let Some(auth_str) = auth_header {
if let Some(rest) = auth_str.strip_prefix("Bearer ") {
token_to_remove = Some(rest.to_string());
}
}
// 其次从 Cookie 中读取
if token_to_remove.is_none() {
if let Some(cookie_header) = req.headers().get(axum::http::header::COOKIE) {
if let Ok(cookie_str) = cookie_header.to_str() {
token_to_remove = parse_session_cookie(cookie_str);
}
}
}
let token_to_remove = extract_auth_token(req.headers());
if let Some(ref token) = token_to_remove {
if let Ok(mut sessions) = state.sessions.lock() {
sessions.remove(token);
}
state.sessions.lock().await.remove(token);
}
// 通过 Max-Age=0 清除浏览器端的 Cookie
@ -214,7 +203,10 @@ pub async fn logout(
Ok((headers, Json(serde_json::json!({ "status": "ok" }))))
}
// 鉴权状态检查接口
/// 鉴权状态检查接口。
///
/// 注:此端点由 `auth_middleware` 保护——只有携带有效会话 Token 的请求才能到达此处,
/// 因此始终返回 `authenticated: true`。如果中间件被移除,此端点将无条件返回已认证。
pub async fn check_auth() -> Result<impl IntoResponse, StatusCode> {
Ok(Json(serde_json::json!({ "authenticated": true })))
}
@ -225,40 +217,18 @@ pub async fn auth_middleware(
req: axum::extract::Request,
next: Next,
) -> Response {
let mut token = None;
// 1. 优先从 Authorization 头读取 Bearer Token兼容
let auth_header = req
.headers()
.get(axum::http::header::AUTHORIZATION)
.and_then(|val| val.to_str().ok());
if let Some(auth_str) = auth_header {
if let Some(rest) = auth_str.strip_prefix("Bearer ") {
token = Some(rest.to_string());
}
}
// 2. 其次从 Cookie 头读取 session_id
if token.is_none() {
if let Some(cookie_header) = req.headers().get(axum::http::header::COOKIE) {
if let Ok(cookie_str) = cookie_header.to_str() {
token = parse_session_cookie(cookie_str);
}
}
}
let token = extract_auth_token(req.headers());
let mut is_valid = false;
if let Some(tok) = token {
if let Ok(mut sessions) = state.sessions.lock() {
// 顺便执行过期会话自动清理垃圾收集
prune_expired_sessions(&mut sessions);
let mut sessions = state.sessions.lock().await;
// 顺便执行过期会话自动清理垃圾收集
prune_expired_sessions(&mut sessions);
if let Some(last_active) = sessions.get_mut(&tok) {
// 会话存在且未过期,更新最后活跃时间
*last_active = chrono::Utc::now();
is_valid = true;
}
if let Some(last_active) = sessions.get_mut(&tok) {
// 会话存在且未过期,更新最后活跃时间
*last_active = chrono::Utc::now();
is_valid = true;
}
}
@ -268,3 +238,22 @@ pub async fn auth_middleware(
StatusCode::UNAUTHORIZED.into_response()
}
}
#[cfg(test)]
mod tests {
use super::verify_password;
#[test]
fn test_verify_password_correct() {
assert!(verify_password("secret123", "secret123"));
assert!(verify_password("", ""));
}
#[test]
fn test_verify_password_wrong() {
assert!(!verify_password("secret123", "secret456"));
assert!(!verify_password("secret", "secret123"));
assert!(!verify_password("secret123", "secret"));
assert!(!verify_password("wrong", ""));
}
}

View File

@ -1,496 +0,0 @@
// src/api/helpers.rs
use super::StandardPaper;
use crate::clients::ads::AdsPaperDoc;
use crate::clients::arxiv::ArxivPaper;
use sqlx::{Row, SqlitePool};
use tracing::info;
/// SQLite 单条 SQL 的参数绑定上限SQLITE_MAX_VARIABLE_NUMBER 默认 999留余量取 900
/// 批量 IN 查询时需按此分批,避免超出绑定数量。
pub const SQLITE_PARAM_LIMIT: usize = 900;
pub fn convert_ads_doc_to_standard(doc: &AdsPaperDoc) -> StandardPaper {
let title = doc
.title
.as_ref()
.and_then(|v: &Vec<String>| v.first())
.cloned()
.unwrap_or_else(|| doc.bibcode.clone());
let authors = doc.author.clone().unwrap_or_default();
let keywords = doc.keyword.clone().unwrap_or_default();
let doi = doc
.doi
.as_ref()
.and_then(|v: &Vec<String>| v.first())
.cloned()
.unwrap_or_default();
let mut arxiv_id = String::new();
if let Some(identifiers) = &doc.identifier {
for id in identifiers {
if id.starts_with("arXiv:") {
arxiv_id = id.replace("arXiv:", "").trim().to_string();
break;
}
}
}
if arxiv_id.is_empty() && doc.bibcode.starts_with("arXiv") {
arxiv_id = doc.bibcode.replace("arXiv", "").trim().to_string();
}
StandardPaper {
bibcode: doc.bibcode.clone(),
title,
authors,
year: doc.year.clone().unwrap_or_default(),
pub_journal: doc.pub_journal.clone().unwrap_or_default(),
keywords,
abstract_text: doc.abstract_text.clone().unwrap_or_default(),
doi,
arxiv_id,
citation_count: doc.citation_count.unwrap_or(0),
reference_count: doc.reference_count.unwrap_or(0),
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: doc.doctype.clone().unwrap_or_else(|| "article".to_string()),
pdf_error: None,
html_error: None,
}
}
pub fn convert_arxiv_to_standard(doc: &ArxivPaper) -> StandardPaper {
StandardPaper {
bibcode: doc.id.clone(),
title: doc.title.clone(),
authors: doc.authors.clone(),
year: doc.year.clone(),
pub_journal: "arXiv Preprint".to_string(),
keywords: Vec::new(),
abstract_text: doc.abstract_text.clone(),
doi: doc.doi.clone().unwrap_or_default(),
arxiv_id: doc.id.clone(),
citation_count: 0,
reference_count: 0,
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: "eprint".to_string(),
pdf_error: None,
html_error: None,
}
}
pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Result<()> {
let mut conn = db.acquire().await?;
save_paper_to_db_tx(&mut conn, p).await
}
pub async fn save_paper_to_db_tx(
conn: &mut sqlx::SqliteConnection,
p: &StandardPaper,
) -> anyhow::Result<()> {
let authors_json = serde_json::to_string(&p.authors)?;
let keywords_json = serde_json::to_string(&p.keywords)?;
// 1. 如果存在 arxiv_id检查是否有已存在的相同 arxiv_id 记录以防 duplicate
if !p.arxiv_id.is_empty() {
#[allow(clippy::type_complexity)]
let existing_opt: Option<(String, Option<String>, Option<String>, Option<String>, Option<String>)> = sqlx::query_as(
"SELECT bibcode, pdf_path, html_path, markdown_path, translation_path FROM papers WHERE arxiv_id = ?"
)
.bind(&p.arxiv_id)
.fetch_optional(&mut *conn)
.await?;
if let Some((existing_bibcode, _pdf, _html, _md, _tr)) = existing_opt {
if existing_bibcode != p.bibcode {
// 发现不同 bibcode 标识的同一篇文献记录,需要进行合并
// 如果已存在的记录使用的是临时 arXiv ID 作为 bibcode且新记录使用的是正式 ADS bibcode我们升级 bibcode 主键
let is_existing_temp = existing_bibcode == p.arxiv_id;
let is_new_formal = p.bibcode != p.arxiv_id;
if is_existing_temp && is_new_formal {
info!(
"发现相同 arXiv ID 的文献,将临时主键 {} 升级为正式 ADS Bibcode: {}",
existing_bibcode, p.bibcode
);
sqlx::query(
"UPDATE papers SET bibcode = ?, title = ?, authors = ?, year = ?, pub = ?, keywords = ?, abstract = ?, doi = ?, citation_count = ?, reference_count = ?, doctype = ? WHERE bibcode = ?"
)
.bind(&p.bibcode)
.bind(&p.title)
.bind(&authors_json)
.bind(&p.year)
.bind(&p.pub_journal)
.bind(&keywords_json)
.bind(&p.abstract_text)
.bind(&p.doi)
.bind(p.citation_count)
.bind(p.reference_count)
.bind(&p.doctype)
.bind(&existing_bibcode)
.execute(&mut *conn)
.await?;
return Ok(());
} else {
// 如果已存在的是正式 ADS bibcode而新插入的是临时 arXiv ID直接忽略或更新元数据而不更改主键
info!(
"发现相同 arXiv ID 的文献 {} 已存在正式记录,忽略临时 arXiv 插入",
existing_bibcode
);
return Ok(());
}
}
}
}
// 2. 正常插入/冲突更新
sqlx::query(
"INSERT INTO papers (bibcode, title, authors, year, pub, keywords, abstract, doi, arxiv_id, citation_count, reference_count, doctype) \
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \
ON CONFLICT(bibcode) DO UPDATE SET \
title=excluded.title, \
authors=excluded.authors, \
pub=excluded.pub, \
keywords=excluded.keywords, \
abstract=excluded.abstract, \
doi=excluded.doi, \
arxiv_id=excluded.arxiv_id, \
citation_count=excluded.citation_count, \
reference_count=excluded.reference_count, \
doctype=excluded.doctype"
)
.bind(&p.bibcode)
.bind(&p.title)
.bind(authors_json)
.bind(&p.year)
.bind(&p.pub_journal)
.bind(keywords_json)
.bind(&p.abstract_text)
.bind(&p.doi)
.bind(&p.arxiv_id)
.bind(p.citation_count)
.bind(p.reference_count)
.bind(&p.doctype)
.execute(&mut *conn)
.await?;
Ok(())
}
/// 从数据库行解析出 StandardPaper消除 get_paper_from_db 和 get_library 之间的重复代码)
///
/// 期望行的列顺序bibcode(0), title(1), authors(2), year(3), pub(4), keywords(5),
/// abstract(6), doi(7), arxiv_id(8), citation_count(9), reference_count(10),
/// pdf_path(11), html_path(12), markdown_path(13), translation_path(14),
/// doctype(15), has_vector(16)
pub fn parse_paper_row(
r: &sqlx::sqlite::SqliteRow,
library_dir: &std::path::Path,
) -> StandardPaper {
let pdf_path: Option<String> = r.get(11);
let html_path: Option<String> = r.get(12);
let markdown_path: Option<String> = r.get(13);
let translation_path: Option<String> = r.get(14);
let doctype_val: Option<String> = r.get(15);
let has_vector: bool = r.get(16);
let authors_str: Option<String> = r.get(2);
let authors: Vec<String> = authors_str
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let keywords_str: Option<String> = r.get(5);
let keywords: Vec<String> = keywords_str
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let is_pdf_exist = pdf_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_html_exist = html_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_md_exist = markdown_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_tr_exist = translation_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let pdf_error = pdf_path
.as_ref()
.filter(|p| p.starts_with("error:"))
.map(|p| p["error:".len()..].trim().to_string());
let html_error = html_path
.as_ref()
.filter(|p| p.starts_with("error:"))
.map(|p| p["error:".len()..].trim().to_string());
StandardPaper {
bibcode: r.get(0),
title: r.get(1),
authors,
year: r.get(3),
pub_journal: r.get(4),
keywords,
abstract_text: r.get(6),
doi: r.get(7),
arxiv_id: r.get(8),
citation_count: r.get(9),
reference_count: r.get(10),
is_downloaded: is_pdf_exist || is_html_exist,
has_pdf: is_pdf_exist,
has_html: is_html_exist,
has_markdown: is_md_exist,
has_translation: is_tr_exist,
has_vector,
doctype: doctype_val.unwrap_or_else(|| "article".to_string()),
pdf_error,
html_error,
}
}
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?;
Ok(parse_paper_row(&r, library_dir))
}
pub async fn check_paper_paths_in_db(
db: &SqlitePool,
library_dir: &std::path::Path,
identifier: &str,
) -> anyhow::Result<
Option<(
Option<String>,
Option<String>,
Option<String>,
Option<String>,
)>,
> {
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?;
if let Some(r) = r_opt {
let pdf: Option<String> = r.get(0);
let html: Option<String> = r.get(1);
let md: Option<String> = r.get(2);
let tr: Option<String> = r.get(3);
let pdf_res = pdf.filter(|p| library_dir.join(p).exists());
let html_res = html.filter(|p| library_dir.join(p).exists());
let md_res = md.filter(|p| library_dir.join(p).exists());
let tr_res = tr.filter(|p| library_dir.join(p).exists());
Ok(Some((pdf_res, html_res, md_res, tr_res)))
} else {
Ok(None)
}
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
#[test]
fn test_convert_ads_doc_to_standard() {
let doc = AdsPaperDoc {
bibcode: "2026A&A...123..456X".to_string(),
title: Some(vec!["A Test Title".to_string()]),
author: Some(vec!["Author A".to_string(), "Author B".to_string()]),
year: Some("2026".to_string()),
pub_journal: Some("Astronomy & Astrophysics".to_string()),
keyword: Some(vec!["Keyword 1".to_string()]),
abstract_text: Some("This is abstract".to_string()),
doi: Some(vec!["10.1000/test.doi".to_string()]),
citation_count: Some(5),
reference_count: Some(10),
reference: None,
citation: None,
identifier: None,
doctype: Some("article".to_string()),
};
let paper = convert_ads_doc_to_standard(&doc);
assert_eq!(paper.bibcode, "2026A&A...123..456X");
assert_eq!(paper.title, "A Test Title");
assert_eq!(paper.authors, vec!["Author A", "Author B"]);
assert_eq!(paper.year, "2026");
assert_eq!(paper.pub_journal, "Astronomy & Astrophysics");
assert_eq!(paper.keywords, vec!["Keyword 1"]);
assert_eq!(paper.abstract_text, "This is abstract");
assert_eq!(paper.doi, "10.1000/test.doi");
assert_eq!(paper.arxiv_id, "");
assert_eq!(paper.citation_count, 5);
assert_eq!(paper.reference_count, 10);
}
#[test]
fn test_convert_ads_doc_to_standard_with_arxiv_identifier() {
let doc = AdsPaperDoc {
bibcode: "2026MNRAS.530.1234A".to_string(),
title: Some(vec!["Another Test Title".to_string()]),
author: Some(vec!["Author A".to_string()]),
year: Some("2026".to_string()),
pub_journal: Some("MNRAS".to_string()),
keyword: None,
abstract_text: None,
doi: None,
citation_count: None,
reference_count: None,
reference: None,
citation: None,
identifier: Some(vec![
"2026MNRAS.530.1234A".to_string(),
"arXiv:2606.12345".to_string(),
]),
doctype: Some("article".to_string()),
};
let paper = convert_ads_doc_to_standard(&doc);
assert_eq!(paper.bibcode, "2026MNRAS.530.1234A");
assert_eq!(paper.arxiv_id, "2606.12345");
}
#[test]
fn test_convert_arxiv_to_standard() {
let doc = ArxivPaper {
id: "2606.12345".to_string(),
title: "Arxiv Title".to_string(),
authors: vec!["Author C".to_string()],
year: "2026".to_string(),
abstract_text: "Arxiv abstract".to_string(),
doi: Some("10.1000/arxiv.doi".to_string()),
pdf_url: "https://arxiv.org/pdf/2606.12345.pdf".to_string(),
};
let paper = convert_arxiv_to_standard(&doc);
assert_eq!(paper.bibcode, "2606.12345");
assert_eq!(paper.title, "Arxiv Title");
assert_eq!(paper.authors, vec!["Author C"]);
assert_eq!(paper.year, "2026");
assert_eq!(paper.pub_journal, "arXiv Preprint");
assert_eq!(paper.doi, "10.1000/arxiv.doi");
assert_eq!(paper.arxiv_id, "2606.12345");
}
#[tokio::test]
async fn test_db_operations() -> anyhow::Result<()> {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await?;
// 运行迁移
sqlx::migrate!("./migrations").run(&pool).await?;
let paper = StandardPaper {
bibcode: "2026A&A...123..456X".to_string(),
title: "A Test Title".to_string(),
authors: vec!["Author A".to_string()],
year: "2026".to_string(),
pub_journal: "Astronomy & Astrophysics".to_string(),
keywords: vec!["Keyword 1".to_string()],
abstract_text: "This is abstract".to_string(),
doi: "10.1000/test.doi".to_string(),
arxiv_id: "".to_string(),
citation_count: 5,
reference_count: 10,
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: "article".to_string(),
pdf_error: None,
html_error: None,
};
// 保存
save_paper_to_db(&pool, &paper).await?;
// 读取
let retrieved =
get_paper_from_db(&pool, std::path::Path::new(""), "2026A&A...123..456X").await?;
assert_eq!(retrieved.title, paper.title);
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());
let (pdf, html, md, tr) = paths.unwrap();
assert!(pdf.is_none());
assert!(html.is_none());
assert!(md.is_none());
assert!(tr.is_none());
Ok(())
}
}

View File

@ -11,8 +11,8 @@ use crate::Config;
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use tokio::sync::{broadcast, oneshot};
use std::sync::Arc;
use tokio::sync::{broadcast, oneshot, Mutex, RwLock};
/// 提供给前端的 SSE 事件
#[derive(Debug, Clone, Serialize)]
@ -63,6 +63,8 @@ pub struct AppState {
pub vision_llm: Option<LlmClient>,
pub embedding: EmbeddingClient,
pub downloader: Downloader,
/// 共享 HTTP 客户端(带超时),避免每次请求新建连接池
pub http_client: reqwest::Client,
pub harvest_status: Arc<tokio::sync::Mutex<crate::services::batch::MetaSyncStatus>>,
pub batch_status: Arc<tokio::sync::Mutex<crate::services::batch::AssetBatchStatus>>,
pub active_bibcode: Arc<tokio::sync::Mutex<Option<String>>>,
@ -83,42 +85,18 @@ pub struct AppState {
/// 项目记忆管理器(跨会话持久化)
pub memory_manager: Arc<tokio::sync::Mutex<MemoryManager>>,
/// 活跃的登录会话 Token 及其最后活跃时间(单用户内存管理)
pub sessions:
Arc<std::sync::Mutex<std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>>>,
pub sessions: Arc<Mutex<std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>>>,
/// 登录速率限制器IP -> 最近失败次数 + 首次失败时间)
pub login_rate_limiter:
Arc<std::sync::Mutex<std::collections::HashMap<String, (u32, std::time::Instant)>>>,
Arc<Mutex<std::collections::HashMap<String, (u32, std::time::Instant)>>>,
}
// 统一标准化的文献格式,用于向前端传输
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct StandardPaper {
pub bibcode: String,
pub title: String,
pub authors: Vec<String>,
pub year: String,
pub pub_journal: String,
pub keywords: Vec<String>,
pub abstract_text: String,
pub doi: String,
pub arxiv_id: String,
pub citation_count: i32,
pub reference_count: i32,
pub is_downloaded: bool,
pub has_pdf: bool,
pub has_html: bool,
pub has_markdown: bool,
pub has_translation: bool,
pub has_vector: bool,
pub doctype: String,
pub pdf_error: Option<String>,
pub html_error: Option<String>,
}
pub use crate::services::paper::StandardPaper;
pub mod agent;
pub mod auth;
pub mod error;
pub mod helpers;
pub mod notes;
pub mod papers;
pub mod permissions;
@ -132,15 +110,10 @@ pub mod handlers {
answer_question, branch_session, chat_agent, delete_session, get_agent_metrics,
get_agent_modes, get_pending_permissions, get_pending_questions, get_session,
get_session_audit, list_sessions, respond_permission, restore_rewound_session,
retry_session, rewind_session, stop_agent, AgentChatRequest, AgentMetricsResponse,
AuditLogEntry, BranchResponse, MessageRecord, RestoreResponse, RetryResponse,
RewindRequest, RewindResponse, SessionDetail, SessionListParams, SessionSummary,
retry_session, rewind_session, stop_agent, AgentChatRequest, BranchResponse,
RestoreResponse, RetryResponse, RewindRequest, RewindResponse, SessionListParams,
};
pub use super::auth::{check_auth, login, logout};
pub use super::helpers::{
check_paper_paths_in_db, convert_ads_doc_to_standard, convert_arxiv_to_standard,
get_paper_from_db, save_paper_to_db, save_paper_to_db_tx,
};
pub use super::notes::{
create_note, delete_note, get_notes, CreateNoteRequest, DeleteNoteParams, GetNotesParams,
NoteRecord,
@ -148,10 +121,9 @@ pub mod handlers {
pub use super::papers::{
download_paper, embed_paper, export_citations, get_active_bibcode, get_citation_network,
get_library, get_paper_detail, mark_no_resource, parse_paper, search_papers,
set_active_bibcode, translate_paper, upload_paper_file, CitationsResponse, DownloadRequest,
EmbedRequest, EmbedResponse, ExportRequest, ExportResponse, MarkNoResourceRequest,
PaperDetailResponse, ParseRequest, ParseResponse, SearchParams, TranslateRequest,
TranslateResponse,
set_active_bibcode, translate_paper, upload_paper_file, DownloadRequest, EmbedRequest,
EmbedResponse, ExportRequest, ExportResponse, MarkNoResourceRequest, PaperDetailResponse,
ParseRequest, ParseResponse, SearchParams, TranslateRequest, TranslateResponse,
};
pub use super::permissions::{
list_permission_rules, update_permission_mode, update_permission_rules,
@ -160,7 +132,7 @@ pub mod handlers {
pub use super::sync::{
delete_sync_query, get_asset_batch_status, get_meta_sync_count, get_meta_sync_status,
get_sync_queries, run_asset_batch, run_meta_sync, stop_asset_batch, AssetBatchRunRequest,
MetaSyncCountRequest, MetaSyncCountResponse, MetaSyncRunRequest, SavedSyncQuery,
MetaSyncCountRequest, MetaSyncCountResponse, MetaSyncRunRequest,
};
pub use super::targets::{
associate_target, extract_paper_targets, list_targets, query_target, AssociateResponse,
@ -168,4 +140,13 @@ pub mod handlers {
TargetQueryParams,
};
pub use super::{AppState, StandardPaper};
pub use crate::services::batch::SavedSyncQuery;
pub use crate::services::paper::{
check_paper_paths_in_db, convert_ads_doc_to_standard, convert_arxiv_to_standard,
get_paper_from_db, save_paper_to_db, save_paper_to_db_tx,
};
pub use crate::services::session::{
AgentMetrics as AgentMetricsResponse, AuditLogEntry, MessageRecord, SessionDetail,
SessionSummary,
};
}

View File

@ -4,23 +4,12 @@ use axum::{
http::StatusCode,
Json,
};
use serde::{Deserialize, Serialize};
use sqlx::Row;
use serde::Deserialize;
use std::sync::Arc;
use super::error::ApiResult;
use super::AppState;
#[derive(Debug, Serialize, Deserialize)]
pub struct NoteRecord {
pub id: i64,
pub bibcode: String,
pub paragraph_index: i64,
pub note_text: String,
pub highlight_color: String,
pub selected_text: String,
pub created_at: String,
}
pub use crate::services::note::NoteRecord;
#[derive(Deserialize)]
pub struct CreateNoteRequest {
@ -50,27 +39,18 @@ pub async fn create_note(
let highlight_color = req.highlight_color.unwrap_or_else(|| "yellow".to_string());
let selected_text = req.selected_text.unwrap_or_default();
let row = sqlx::query(
"INSERT INTO notes (bibcode, paragraph_index, note_text, highlight_color, selected_text) VALUES (?, ?, ?, ?, ?) RETURNING id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at"
let record = crate::services::note::create_highlight_note(
&state.db,
&req.bibcode,
req.paragraph_index,
&note_text,
&highlight_color,
&selected_text,
)
.bind(&req.bibcode)
.bind(req.paragraph_index)
.bind(&note_text)
.bind(&highlight_color)
.bind(&selected_text)
.fetch_one(&state.db)
.await
.map_err(|e| super::error::AppError::internal(format!("保存笔记失败: {}", e)))?;
Ok(Json(NoteRecord {
id: row.get(0),
bibcode: row.get(1),
paragraph_index: row.get(2),
note_text: row.get(3),
highlight_color: row.get(4),
selected_text: row.get(5),
created_at: row.get(6),
}))
Ok(Json(record))
}
// 查询某篇文献的全部笔记
@ -78,26 +58,9 @@ pub async fn get_notes(
State(state): State<Arc<AppState>>,
Query(params): Query<GetNotesParams>,
) -> ApiResult<Json<Vec<NoteRecord>>> {
let rows = sqlx::query(
"SELECT id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at FROM notes WHERE bibcode = ? ORDER BY paragraph_index, created_at"
)
.bind(&params.bibcode)
.fetch_all(&state.db)
.await
.map_err(|e| super::error::AppError::internal(format!("查询笔记失败: {}", e)))?;
let notes: Vec<NoteRecord> = rows
.iter()
.map(|r| NoteRecord {
id: r.get(0),
bibcode: r.get(1),
paragraph_index: r.get(2),
note_text: r.get(3),
highlight_color: r.get(4),
selected_text: r.get(5),
created_at: r.get(6),
})
.collect();
let notes = crate::services::note::list_highlight_notes(&state.db, &params.bibcode)
.await
.map_err(|e| super::error::AppError::internal(format!("查询笔记失败: {}", e)))?;
Ok(Json(notes))
}
@ -107,11 +70,16 @@ pub async fn delete_note(
State(state): State<Arc<AppState>>,
Query(params): Query<DeleteNoteParams>,
) -> ApiResult<StatusCode> {
sqlx::query("DELETE FROM notes WHERE id = ?")
.bind(params.id)
.execute(&state.db)
let success = crate::services::note::delete_highlight_note(&state.db, params.id)
.await
.map_err(|e| super::error::AppError::internal(format!("删除笔记失败: {}", e)))?;
if !success {
return Err(super::error::AppError::not_found(format!(
"未找到 ID 为 {} 的笔记",
params.id
)));
}
Ok(StatusCode::NO_CONTENT)
}

View File

@ -5,17 +5,12 @@ use axum::{
Json,
};
use serde::{Deserialize, Serialize};
use sqlx::Row;
use std::sync::Arc;
use tokio::fs;
use tracing::{error, info, warn};
use tracing::info;
use super::error::{ApiResult, AppError};
use super::helpers::{
check_paper_paths_in_db, convert_ads_doc_to_standard, get_paper_from_db, parse_paper_row,
save_paper_to_db, SQLITE_PARAM_LIMIT,
};
use super::{AppState, StandardPaper};
use crate::services::paper::get_paper_from_db;
// 检索请求参数
#[derive(Debug, Deserialize)]
@ -34,8 +29,8 @@ pub async fn search_papers(
Query(params): Query<SearchParams>,
) -> ApiResult<Json<Vec<StandardPaper>>> {
let source = params.source.unwrap_or_else(|| "all".to_string());
let rows = params.rows.unwrap_or(10);
let start = params.start.unwrap_or(0);
let rows = params.rows.unwrap_or(10).clamp(1, 200);
let start = params.start.unwrap_or(0).max(0);
let sort = params.sort.as_deref().unwrap_or("relevance");
let results =
@ -106,17 +101,17 @@ pub async fn parse_paper(
)
.await
.map_err(|e| {
// 根据错误信息推断合适的错误类型(保留原有状态码语义)
let msg = e.to_string();
if msg.contains("未注册")
|| msg.contains("未找到")
|| msg.contains("丢失")
|| msg.contains("未在数据库中注册")
{
AppError::not_found(msg)
AppError::not_found("该文献未找到,请确认标识符是否正确")
} else if msg.contains("请先下载") {
AppError::bad_request(msg)
AppError::bad_request("请先下载该文献后再尝试解析")
} else {
tracing::error!("文献解析内部错误: {}", msg);
AppError::internal(msg)
}
})?;
@ -147,229 +142,58 @@ pub async fn translate_paper(
req.bibcode, force
);
let (_, _, md_opt, tr_opt) =
check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode)
.await
.map_err(|e| AppError::not_found(format!("查询文献路径失败: {}", e)))?
.ok_or_else(|| AppError::not_found("该文献未注册在数据库中"))?;
// 若本地已存在翻译物理文件且未指明强制重译,直读本地缓存返回
if !force {
if let Some(tr_rel) = tr_opt {
let tr_abs = state.config.library_dir.join(&tr_rel);
if tr_abs.exists() {
if let Ok(content) = fs::read_to_string(&tr_abs).await {
return Ok(Json(TranslateResponse {
translation: content,
}));
}
}
}
}
// 检查英文解析文件是否存在
let md_rel = match md_opt {
Some(rel) => rel,
None => {
error!(
"文献 {} 翻译失败:文献未完成解析,缺少英文 Markdown 路径",
req.bibcode
);
return Err(AppError::bad_request("文献必须先完成解析方可翻译"));
}
};
let md_abs = state.config.library_dir.join(&md_rel);
if !md_abs.exists() {
error!(
"文献 {} 翻译失败:解析的英文 Markdown 文件 {:?} 不存在",
req.bibcode, md_abs
);
return Err(AppError::bad_request("解析 Markdown 文件丢失"));
}
let english_markdown = fs::read_to_string(&md_abs).await.map_err(|e| {
error!("文献 {} 翻译失败:读取解析内容失败: {}", req.bibcode, e);
AppError::internal(format!("读取解析内容失败: {}", e))
})?;
// 调用 LLM 翻译服务并注入对照词表
let translated_markdown = crate::services::translation::translate_markdown(
&english_markdown,
let translated_markdown = crate::services::translation::translate_paper_cached(
&state.db,
&state.config.library_dir,
&state.dict,
&state.medium_llm,
&req.bibcode,
force,
)
.await
.map_err(|e| {
error!(
"文献 {} 翻译失败:调用 LLM 翻译发生错误: {}",
req.bibcode, e
);
AppError::internal(format!("调用 LLM 翻译失败: {}", e))
let err_msg = e.to_string();
if err_msg.contains("该文献未注册") || err_msg.contains("丢失") {
AppError::not_found("该文献未在本地库中找到")
} else if err_msg.contains("必须先完成解析") {
AppError::bad_request("请先完成文献解析后再翻译")
} else {
tracing::error!("文献翻译内部错误: {}", err_msg);
AppError::internal(err_msg)
}
})?;
// 翻译结果持久化C2抽取为共享函数与批量任务复用
crate::services::translation::persist_translation(
&state.db,
&state.config.library_dir,
&req.bibcode,
&translated_markdown,
)
.await
.map_err(|e| AppError::internal(format!("翻译结果持久化失败: {}", e)))?;
Ok(Json(TranslateResponse {
translation: translated_markdown,
}))
}
#[derive(Debug, Serialize)]
pub struct CitationsResponse {
pub bibcode: String,
pub title: String,
pub citation_count: i32,
pub reference_count: i32,
pub references: Vec<String>, // 该文献参考文献 bibcode 数组
pub citations: Vec<String>, // 引用该文献的 bibcode 数组
pub citation_counts: std::collections::HashMap<String, i32>, // 相关文献与被引数映射
}
// 从 SQLite 查询引用关联,生成引用星系关系树
pub async fn get_citation_network(
State(state): State<Arc<AppState>>,
Query(params): Query<DownloadRequest>,
) -> ApiResult<Json<CitationsResponse>> {
let paper = match get_paper_from_db(&state.db, &state.config.library_dir, &params.bibcode).await
{
Ok(p) => p,
Err(_) => {
// 如果本地数据库查不到,尝试从 ADS 在线 API 动态获取
if !state.config.ads_api_key.is_empty() {
match state
.ads
.search(&format!("bibcode:{}", params.bibcode), 0, 1, "relevance")
.await
{
Ok(docs) => {
if let Some(doc) = docs.first() {
let standard_paper = convert_ads_doc_to_standard(doc);
// 保存至数据库缓存,并保存引用关联
if let Err(e) = save_paper_to_db(&state.db, &standard_paper).await {
warn!("保存引用文献至数据库失败: {}", e);
}
if let Some(refs) = &doc.reference {
for ref_bib in refs {
if let Err(e) = sqlx::query("INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)")
.bind(&standard_paper.bibcode)
.bind(ref_bib)
.execute(&state.db)
.await
{
warn!("保存引用关系失败 ({} -> {}): {}", standard_paper.bibcode, ref_bib, e);
}
}
}
if let Some(cits) = &doc.citation {
for cit_bib in cits {
if let Err(e) = sqlx::query("INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)")
.bind(cit_bib)
.bind(&standard_paper.bibcode)
.execute(&state.db)
.await
{
warn!("保存被引关系失败 ({} -> {}): {}", cit_bib, standard_paper.bibcode, e);
}
}
}
standard_paper
} else {
return Err(AppError::not_found(format!(
"在本地库及 ADS 中均未找到该文献: {}",
params.bibcode
)));
}
}
Err(e) => {
return Err(AppError::internal(format!("在线检索文献元数据失败: {}", e)));
}
}
} else {
return Err(AppError::not_found(format!(
"本地数据库未收录该文献,且未配置 ADS_API_KEY无法在线加载: {}",
params.bibcode
)));
}
) -> ApiResult<Json<crate::services::citation::CitationNetworkResult>> {
let res = crate::services::citation::get_citation_network_with_fallback(
&state.db,
&state.config.library_dir,
&state.config.ads_api_key,
&state.ads,
&params.bibcode,
)
.await
.map_err(|e| {
let msg = e.to_string();
if msg.contains("未找到该文献") {
AppError::not_found("未找到该文献的引用关系数据")
} else if msg.contains("未配置 ADS_API_KEY") {
AppError::bad_request("未配置 ADS_API_KEY无法在线加载引用数据")
} else {
tracing::error!("引用网络查询内部错误: {}", msg);
AppError::internal(msg)
}
};
})?;
// 加载引用的文献
let refs_rows =
sqlx::query("SELECT target_bibcode FROM citations_references WHERE source_bibcode = ?")
.bind(&params.bibcode)
.fetch_all(&state.db)
.await
.unwrap_or_else(|e| {
warn!("查询引用关系失败: {}", e);
Vec::new()
});
let references: Vec<String> = refs_rows.iter().map(|row| row.get(0)).collect();
// 加载被引用的文献
let cits_rows =
sqlx::query("SELECT source_bibcode FROM citations_references WHERE target_bibcode = ?")
.bind(&params.bibcode)
.fetch_all(&state.db)
.await
.unwrap_or_else(|e| {
warn!("查询被引关系失败: {}", e);
Vec::new()
});
let citations: Vec<String> = cits_rows.iter().map(|row| row.get(0)).collect();
// 加载关联文献的被引数量 (从 SQLite papers 表获取)
// 优化:用 WHERE IN 批量查询替代逐条 query_scalar避免 N+1。
let mut citation_counts = std::collections::HashMap::new();
let mut all_related = references.clone();
all_related.extend(citations.clone());
// 去重,避免重复绑定与查询
let mut unique_bibs: Vec<String> = all_related.clone();
unique_bibs.sort();
unique_bibs.dedup();
// SQLite 单条 SQL 参数上限 999分批查询
for chunk in unique_bibs.chunks(SQLITE_PARAM_LIMIT) {
if chunk.is_empty() {
continue;
}
let placeholders = vec!["?"; chunk.len()].join(", ");
let sql = format!(
"SELECT bibcode, citation_count FROM papers WHERE bibcode IN ({})",
placeholders
);
let mut query = sqlx::query(&sql);
for b in chunk {
query = query.bind(b);
}
let rows = query.fetch_all(&state.db).await.unwrap_or_else(|e| {
warn!("批量查询文献引用数失败: {}", e);
Vec::new()
});
for row in rows {
let bib: String = row.get(0);
let count: i32 = row.get(1);
citation_counts.insert(bib, count);
}
}
Ok(Json(CitationsResponse {
bibcode: paper.bibcode,
title: paper.title,
citation_count: paper.citation_count,
reference_count: paper.reference_count,
references,
citations,
citation_counts,
}))
Ok(Json(res))
}
// ── GET /api/paper ──
@ -389,23 +213,20 @@ pub async fn get_paper_detail(
.await
.map_err(|e| AppError::not_found(format!("未找到该文献数据: {}", e)))?;
let (_, _, md_opt, tr_opt) =
check_paper_paths_in_db(&state.db, &state.config.library_dir, &params.bibcode)
.await
.map_err(|e| AppError::internal(e.to_string()))?
.unwrap_or_default();
let english_content = match md_opt {
Some(rel) => fs::read_to_string(state.config.library_dir.join(rel))
.await
.ok(),
None => None,
let mode = crate::services::paper::ReadMode::Full {
include_translation: true,
};
let translation_content = match tr_opt {
Some(rel) => fs::read_to_string(state.config.library_dir.join(rel))
.await
.ok(),
None => None,
let (english_content, translation_content) = match crate::services::paper::read_paper_content(
&state.db,
&state.config.library_dir,
&params.bibcode,
mode,
)
.await
{
Ok(res) => (Some(res.content), res.translation_content),
Err(_) => (None, None),
};
Ok(Json(PaperDetailResponse {
@ -416,19 +237,25 @@ pub async fn get_paper_detail(
}
// ── GET /api/library ──
#[derive(Deserialize)]
pub struct LibraryParams {
pub limit: Option<i64>,
pub offset: Option<i64>,
}
// 获取本地图书馆文献列表
pub async fn get_library(
State(state): State<Arc<AppState>>,
Query(params): Query<LibraryParams>,
) -> ApiResult<Json<Vec<StandardPaper>>> {
let rows = 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 ORDER BY created_at DESC")
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("访问本地数据库失败: {}", e)))?;
let mut list = Vec::new();
for r in rows {
list.push(parse_paper_row(&r, &state.config.library_dir));
}
let list = crate::services::paper::get_library_list(
&state.db,
&state.config.library_dir,
params.limit,
params.offset,
)
.await
.map_err(|e| AppError::internal(format!("访问本地数据库失败: {}", e)))?;
Ok(Json(list))
}
@ -508,128 +335,26 @@ pub async fn upload_paper_file(
return Err(AppError::bad_request("bibcode 包含非法字符"));
}
if file_bytes.is_empty() {
return Err(AppError::bad_request("上传文件为空或读取失败"));
}
// 尝试将可能的 DOI 或 arXiv ID 解析为真实的 bibcode
let mut resolved_bibcode = bibcode.clone();
let exists_as_bibcode = sqlx::query("SELECT bibcode FROM papers WHERE bibcode = ?")
.bind(&resolved_bibcode)
.fetch_optional(&state.db)
.await
.unwrap_or(None)
.is_some();
if !exists_as_bibcode {
// 尝试匹配 DOI
let clean_doi = resolved_bibcode
.trim_start_matches("doi:")
.trim_start_matches("DOI:")
.trim_start_matches("https://doi.org/")
.trim_start_matches("http://doi.org/")
.trim();
if let Some(row) = sqlx::query(
"SELECT bibcode FROM papers WHERE doi = ? OR doi = ? OR LOWER(doi) = LOWER(?)",
)
.bind(clean_doi)
.bind(&resolved_bibcode)
.bind(clean_doi)
.fetch_optional(&state.db)
.await
.unwrap_or(None)
{
let found: String = row.get(0);
info!(
"上传接口:通过 DOI 匹配成功,将 '{}' 解析为 bibcode '{}'",
bibcode, found
);
resolved_bibcode = found;
let updated_paper = crate::services::paper::upload_paper_file_service(
&state.db,
&state.config.library_dir,
&bibcode,
&file_type,
&file_name,
&file_bytes,
)
.await
.map_err(|e| {
let err_msg = e.to_string();
if err_msg.contains("未找到该文献记录") {
AppError::not_found("未找到该文献记录,请先在库中检索该文献")
} else if err_msg.contains("校验失败") || err_msg.contains("不是有效的 UTF-8") {
AppError::bad_request("上传文件校验失败,请确认文件格式正确")
} else {
// 尝试匹配 arXiv ID
let clean_arxiv = resolved_bibcode
.trim_start_matches("arxiv:")
.trim_start_matches("arXiv:")
.trim_start_matches("ARXIV:")
.trim();
// 移除可能存在的版本号后缀(如 2303.12345v1 -> 2303.12345
let clean_arxiv_no_version = if let Some(pos) = clean_arxiv.find('v') {
if clean_arxiv[pos + 1..].chars().all(|c| c.is_ascii_digit()) {
&clean_arxiv[..pos]
} else {
clean_arxiv
}
} else {
clean_arxiv
};
if let Some(row) = sqlx::query(
"SELECT bibcode FROM papers WHERE arxiv_id = ? OR arxiv_id = ? OR arxiv_id LIKE ? OR arxiv_id LIKE ?"
)
.bind(clean_arxiv)
.bind(clean_arxiv_no_version)
.bind(format!("{}%", clean_arxiv_no_version))
.bind(format!("arXiv:{}%", clean_arxiv_no_version))
.fetch_optional(&state.db)
.await
.unwrap_or(None) {
let found: String = row.get(0);
info!("上传接口:通过 arXiv ID 匹配成功,将 '{}' 解析为 bibcode '{}'", bibcode, found);
resolved_bibcode = found;
}
tracing::error!("文件上传内部错误: {}", err_msg);
AppError::internal(err_msg)
}
}
let bibcode = resolved_bibcode;
// 从数据库读取该文献元数据
let _paper = get_paper_from_db(&state.db, &state.config.library_dir, &bibcode)
.await
.map_err(|e| AppError::not_found(format!("未找到该文献记录: {}", e)))?;
// 校验并保存文件
let is_pdf = file_type == "pdf" || file_name.to_lowercase().ends_with(".pdf");
let relative_path = if is_pdf {
crate::services::download::validate_pdf_content(&file_bytes)
.map_err(|e| AppError::bad_request(format!("PDF 文件内容校验失败: {}", e)))?;
let pdf_filename = format!("{}.pdf", bibcode);
let pdf_dest = state.config.library_dir.join("PDF").join(&pdf_filename);
if let Some(parent) = pdf_dest.parent() {
std::fs::create_dir_all(parent).unwrap_or_default();
}
std::fs::write(&pdf_dest, &file_bytes)
.map_err(|e| AppError::internal(format!("无法写入 PDF 文件: {}", e)))?;
format!("PDF/{}", pdf_filename)
} else {
let text_content = String::from_utf8(file_bytes)
.map_err(|_| AppError::bad_request("上传的 HTML 文件不是有效的 UTF-8 文本"))?;
crate::services::download::validate_html_content_lenient(&text_content)
.map_err(|e| AppError::bad_request(format!("HTML 文件内容校验失败: {}", e)))?;
let html_filename = format!("{}.html", bibcode);
let html_dest = state.config.library_dir.join("HTML").join(&html_filename);
if let Some(parent) = html_dest.parent() {
std::fs::create_dir_all(parent).unwrap_or_default();
}
std::fs::write(&html_dest, &text_content)
.map_err(|e| AppError::internal(format!("无法写入 HTML 文件: {}", e)))?;
format!("HTML/{}", html_filename)
};
// 更新数据库路径状态
let path_field = if is_pdf { "pdf_path" } else { "html_path" };
let sql = format!("UPDATE papers SET {} = ? WHERE bibcode = ?", path_field);
sqlx::query(&sql)
.bind(&relative_path)
.bind(&bibcode)
.execute(&state.db)
.await
.map_err(|e| AppError::internal(format!("更新数据库状态失败: {}", e)))?;
// 重新获取最新的文献信息以更新前端界面
let updated_paper = get_paper_from_db(&state.db, &state.config.library_dir, &bibcode)
.await
.map_err(|e| AppError::internal(format!("重读文献数据失败: {}", e)))?;
})?;
Ok(Json(updated_paper))
}
@ -646,27 +371,19 @@ pub async fn mark_no_resource(
Json(req): Json<MarkNoResourceRequest>,
) -> ApiResult<Json<StandardPaper>> {
let clear_flag = req.clear.unwrap_or(false);
info!(
"接收到文献无资源标记指令,标识符: {}, clear: {}",
req.bibcode, clear_flag
);
if clear_flag {
info!("接收到清除文献无资源标记指令,标识符: {}", req.bibcode);
sqlx::query("UPDATE papers SET pdf_path = NULL, html_path = NULL WHERE bibcode = ?")
.bind(&req.bibcode)
.execute(&state.db)
.await
.map_err(|e| AppError::internal(format!("清除无资源标记失败: {}", e)))?;
} else {
info!("接收到文献无资源标记指令,标识符: {}", req.bibcode);
sqlx::query("UPDATE papers SET pdf_path = 'error:no_resource', html_path = 'error:no_resource' WHERE bibcode = ?")
.bind(&req.bibcode)
.execute(&state.db)
.await
.map_err(|e| AppError::internal(format!("更新数据库无资源标记失败: {}", e)))?;
}
// 重新获取最新的文献信息以更新前端界面
let updated_paper = get_paper_from_db(&state.db, &state.config.library_dir, &req.bibcode)
.await
.map_err(|e| AppError::internal(format!("重读文献数据失败: {}", e)))?;
let updated_paper = crate::services::paper::mark_no_resource_service(
&state.db,
&state.config.library_dir,
&req.bibcode,
clear_flag,
)
.await
.map_err(|e| AppError::internal(format!("更新文献无资源标记失败: {}", e)))?;
Ok(Json(updated_paper))
}
@ -717,32 +434,23 @@ pub async fn embed_paper(
) -> ApiResult<Json<EmbedResponse>> {
info!("接收到文献向量化分块入库指令: {}", req.bibcode);
let (_, _, md_opt, _) =
check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode)
.await
.map_err(|e| AppError::not_found(format!("获取文献路径失败: {}", e)))?
.ok_or_else(|| AppError::not_found("该文献未注册在数据库中"))?;
let md_rel =
md_opt.ok_or_else(|| AppError::bad_request("文献尚未解析为 Markdown请先执行解析"))?;
let md_abs = state.config.library_dir.join(&md_rel);
if !md_abs.exists() {
return Err(AppError::not_found("文献 Markdown 文件未找到,请重新解析"));
}
let markdown_content = fs::read_to_string(&md_abs)
.await
.map_err(|e| AppError::internal(format!("读取 Markdown 文件失败: {}", e)))?;
let chunk_count = crate::services::rag::ingest_paper(
let chunk_count = crate::services::rag::embed_paper_service(
&state.db,
&state.embedding,
&state.config.library_dir,
&req.bibcode,
&markdown_content,
None,
&state.embedding,
)
.await
.map_err(|e| AppError::internal(format!("文献向量化失败: {}", e)))?;
.map_err(|e| {
let err_msg = e.to_string();
if err_msg.contains("未注册") || err_msg.contains("未找到") {
AppError::not_found(err_msg)
} else if err_msg.contains("尚未解析") || err_msg.contains("未找到") {
AppError::bad_request(err_msg)
} else {
AppError::internal(err_msg)
}
})?;
Ok(Json(EmbedResponse { chunk_count }))
}

View File

@ -52,6 +52,26 @@ pub struct PermissionActionResponse {
// ── Handlers ──
/// 验证会话是否存在(防止对不存在的会话进行权限操作)
async fn verify_session_exists(
db: &sqlx::SqlitePool,
session_id: &str,
) -> Result<(), Json<PermissionActionResponse>> {
let exists: bool =
sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ?)")
.bind(session_id)
.fetch_one(db)
.await
.unwrap_or(false);
if !exists {
return Err(Json(PermissionActionResponse {
success: false,
message: format!("会话 {} 不存在", session_id),
}));
}
Ok(())
}
/// 动态添加或移除权限规则。
pub async fn update_permission_rules(
Path(session_id): Path<String>,
@ -63,6 +83,11 @@ pub async fn update_permission_rules(
session_id, req.action, req.kind, req.rule
);
// 验证会话存在
if let Err(resp) = verify_session_exists(&state.db, &session_id).await {
return resp;
}
if !["deny", "allow", "ask"].contains(&req.kind.as_str()) {
return Json(PermissionActionResponse {
success: false,
@ -105,43 +130,31 @@ pub async fn update_permission_rules(
},
};
if let Ok(mut checkers) = state.session_permission_checkers.write() {
let checker = checkers
.entry(session_id.clone())
.or_insert_with(PermissionChecker::new);
checker.add_rule_dynamic(session_rule);
Json(PermissionActionResponse {
success: true,
message: format!("已添加 {}:{} 规则", req.kind, req.rule),
})
} else {
Json(PermissionActionResponse {
success: false,
message: "权限检查器锁定失败".to_string(),
})
}
let mut checkers = state.session_permission_checkers.write().await;
let checker = checkers
.entry(session_id.clone())
.or_insert_with(PermissionChecker::new);
checker.add_rule_dynamic(session_rule);
Json(PermissionActionResponse {
success: true,
message: format!("已添加 {}:{} 规则", req.kind, req.rule),
})
}
"remove" => {
if let Ok(mut checkers) = state.session_permission_checkers.write() {
let removed = if let Some(checker) = checkers.get_mut(&session_id) {
checker.remove_rule_dynamic(&req.rule, &req.kind)
} else {
0
};
Json(PermissionActionResponse {
success: removed > 0,
message: if removed > 0 {
format!("已移除 {}{}/{} 规则", removed, req.kind, req.rule)
} else {
format!("未找到匹配的 {}/{} 规则", req.kind, req.rule)
},
})
let mut checkers = state.session_permission_checkers.write().await;
let removed = if let Some(checker) = checkers.get_mut(&session_id) {
checker.remove_rule_dynamic(&req.rule, &req.kind)
} else {
Json(PermissionActionResponse {
success: false,
message: "权限检查器锁定失败".to_string(),
})
}
0
};
Json(PermissionActionResponse {
success: removed > 0,
message: if removed > 0 {
format!("已移除 {}{}/{} 规则", removed, req.kind, req.rule)
} else {
format!("未找到匹配的 {}/{} 规则", req.kind, req.rule)
},
})
}
other => Json(PermissionActionResponse {
success: false,
@ -156,6 +169,11 @@ pub async fn update_permission_mode(
State(state): State<Arc<AppState>>,
Json(req): Json<PermissionModeRequest>,
) -> Json<PermissionActionResponse> {
// 验证会话存在
if let Err(resp) = verify_session_exists(&state.db, &session_id).await {
return resp;
}
let mode = PermissionMode::from_str(&req.mode);
let mode_str = format!("{:?}", mode);
@ -164,23 +182,17 @@ pub async fn update_permission_mode(
session_id, mode_str
);
if let Ok(mut checkers) = state.session_permission_checkers.write() {
if let Some(checker) = checkers.get_mut(&session_id) {
checker.set_mode(mode);
Json(PermissionActionResponse {
success: true,
message: format!("权限模式已切换为: {}", mode_str),
})
} else {
Json(PermissionActionResponse {
success: false,
message: format!("会话 {} 不存在或无活跃 Agent", session_id),
})
}
let mut checkers = state.session_permission_checkers.write().await;
if let Some(checker) = checkers.get_mut(&session_id) {
checker.set_mode(mode);
Json(PermissionActionResponse {
success: true,
message: format!("权限模式已切换为: {}", mode_str),
})
} else {
Json(PermissionActionResponse {
success: false,
message: "权限检查器锁定失败".to_string(),
message: format!("会话 {} 不存在或无活跃 Agent", session_id),
})
}
}
@ -190,17 +202,7 @@ pub async fn list_permission_rules(
Path(session_id): Path<String>,
State(state): State<Arc<AppState>>,
) -> Json<PermissionRulesResponse> {
let checkers = match state.session_permission_checkers.read() {
Ok(c) => c,
Err(_) => {
return Json(PermissionRulesResponse {
deny_rules: Vec::new(),
allow_rules: Vec::new(),
ask_rules: Vec::new(),
mode: "default".to_string(),
});
}
};
let checkers = state.session_permission_checkers.read().await;
if let Some(checker) = checkers.get(&session_id) {
let mut deny_rules = Vec::new();

View File

@ -7,11 +7,12 @@ use axum::{
extract::{Query, State},
Json,
};
use serde::{Deserialize, Serialize};
use serde::Deserialize;
use std::sync::Arc;
use super::error::{ApiResult, AppError};
use super::AppState;
pub use crate::services::search::SearchResult;
#[derive(Debug, Deserialize)]
pub struct SearchParams {
@ -34,102 +35,22 @@ fn default_limit() -> i64 {
20
}
#[derive(Debug, Serialize)]
pub struct SearchResult {
pub result_type: String,
pub session_id: String,
pub title: Option<String>,
pub snippet: String,
pub created_at: Option<String>,
}
/// GET /api/search
pub async fn search(
State(state): State<Arc<AppState>>,
Query(params): Query<SearchParams>,
) -> ApiResult<Json<Vec<SearchResult>>> {
let scope = &params.scope;
let limit = params.limit.min(100);
let mut results = Vec::new();
let db = &state.db;
// 搜索会话
if scope == "all" || scope == "sessions" {
let rows = sqlx::query_as::<_, (String, String, String, Option<String>)>(
"SELECT s.session_id, s.title, \
snippet(agent_sessions_fts, 1, '<mark>', '</mark>', '...', 40) as snippet, \
s.created_at \
FROM agent_sessions_fts fts \
JOIN agent_sessions s ON s.session_id = fts.session_id \
WHERE agent_sessions_fts MATCH $1 \
ORDER BY rank LIMIT $2",
)
.bind(&params.q)
.bind(limit)
.fetch_all(db)
.await
.map_err(|e| AppError::internal(format!("搜索会话失败: {}", e)))?;
for (sid, title, snippet, created_at) in rows {
results.push(SearchResult {
result_type: "session".into(),
session_id: sid,
title: Some(title),
snippet,
created_at,
});
}
}
// 搜索消息
if scope == "all" || scope == "messages" {
let msg_rows = if let Some(ref sid) = params.session_id {
sqlx::query_as::<_, (String, String, String, String, Option<String>)>(
"SELECT fts.session_id, s.title, \
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 60) as snippet, \
fts.role, m.created_at \
FROM agent_messages_fts fts \
JOIN agent_sessions s ON s.session_id = fts.session_id \
JOIN agent_messages m ON m.rowid = fts.rowid \
WHERE agent_messages_fts MATCH $1 AND fts.session_id = $3 \
AND m.active = 1 \
ORDER BY rank LIMIT $2",
)
.bind(&params.q)
.bind(limit)
.bind(sid)
.fetch_all(db)
.await
.map_err(|e| AppError::internal(format!("搜索消息失败: {}", e)))?
} else {
sqlx::query_as::<_, (String, String, String, String, Option<String>)>(
"SELECT fts.session_id, s.title, \
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 60) as snippet, \
fts.role, m.created_at \
FROM agent_messages_fts fts \
JOIN agent_sessions s ON s.session_id = fts.session_id \
JOIN agent_messages m ON m.rowid = fts.rowid \
WHERE agent_messages_fts MATCH $1 AND m.active = 1 \
ORDER BY rank LIMIT $2",
)
.bind(&params.q)
.bind(limit)
.fetch_all(db)
.await
.map_err(|e| AppError::internal(format!("搜索消息失败: {}", e)))?
};
for (sid, title, snippet, role, created_at) in msg_rows {
results.push(SearchResult {
result_type: format!("message/{}", role),
session_id: sid,
title: Some(title),
snippet,
created_at,
});
}
}
let results = crate::services::search::search_agent_history(
&state.db,
&params.q,
&params.scope,
limit,
params.session_id.as_deref(),
)
.await
.map_err(|e| AppError::internal(format!("全文检索失败: {}", e)))?;
Ok(Json(results))
}

View File

@ -5,7 +5,6 @@ use axum::{
Json,
};
use serde::{Deserialize, Serialize};
use sqlx::Row;
use std::sync::Arc;
use tracing::error;
@ -24,7 +23,7 @@ pub async fn run_meta_sync(
State(state): State<Arc<AppState>>,
Json(req): Json<MetaSyncRunRequest>,
) -> ApiResult<StatusCode> {
// 检查并同步初始化任务状态,防止前端轮询竞态条件
// 检查并同步初始化任务状态——原子地检查+设置,消除 TOCTOU 竞态
{
let mut status = state.harvest_status.lock().await;
if status.active {
@ -100,30 +99,19 @@ pub struct AssetBatchRunRequest {
pub skip_preceding_uncompleted: Option<bool>, // 跳过前置未完成
}
struct PaperRecord {
bibcode: String,
pdf_path: Option<String>,
html_path: Option<String>,
markdown_path: Option<String>,
translation_path: Option<String>,
year: String,
created_at: String,
has_vector: bool,
has_target: bool,
}
pub async fn run_asset_batch(
State(state): State<Arc<AppState>>,
Json(req): Json<AssetBatchRunRequest>,
) -> ApiResult<StatusCode> {
// 检查是否已经在进行批量处理任务
// 原子地检查+设置 active 标志,消除 TOCTOU 竞态
{
let status = state.batch_status.lock().await;
let mut status = state.batch_status.lock().await;
if status.active {
return Err(AppError::conflict(
"当前已有文献批量任务在后台运行中,请勿重复启动",
));
}
status.active = true;
}
let target_phase = req.target_phase.clone();
@ -136,154 +124,23 @@ pub async fn run_asset_batch(
_ => return Err(AppError::bad_request("不支持的 target_phase 参数值")),
};
let rows = sqlx::query("SELECT bibcode, pdf_path, html_path, markdown_path, translation_path, year, datetime(created_at, 'localtime'), EXISTS(SELECT 1 FROM paper_chunks_content WHERE bibcode = papers.bibcode), EXISTS(SELECT 1 FROM paper_targets WHERE bibcode = papers.bibcode) FROM papers")
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("读取数据库失败: {}", e)))?;
let mut records = Vec::new();
for r in rows {
records.push(PaperRecord {
bibcode: r.get(0),
pdf_path: r.get(1),
html_path: r.get(2),
markdown_path: r.get(3),
translation_path: r.get(4),
year: r.get(5),
created_at: r.get(6),
has_vector: r.get(7),
has_target: r.get(8),
});
}
// 排序
let sort_order = req
.sort_order
.clone()
.unwrap_or_else(|| "default".to_string());
if sort_order == "pub_year_desc" {
records.sort_by(|a, b| b.year.cmp(&a.year));
} else if sort_order == "created_at_desc" {
records.sort_by(|a, b| b.created_at.cmp(&a.created_at));
}
let skip_completed = req.skip_completed.unwrap_or(false);
let skip_failed = req.skip_failed.unwrap_or(false);
let skip_preceding_failed = req.skip_preceding_failed.unwrap_or(false);
let skip_preceding_uncompleted = req.skip_preceding_uncompleted.unwrap_or(false);
let is_completed = |path: &Option<String>| -> bool {
if let Some(p) = path {
!p.is_empty() && !p.starts_with("error:") && !p.starts_with("mineru_batch:")
} else {
false
}
};
let is_failed = |path: &Option<String>| -> bool {
if let Some(p) = path {
p.starts_with("error:")
} else {
false
}
};
let is_no_resource = |path: &Option<String>| -> bool {
if let Some(p) = path {
p.starts_with("error:no_resource")
|| p.starts_with("error:无资源")
|| p.starts_with("error:无有效全文")
} else {
false
}
};
let mut target_bibcodes = Vec::new();
for rec in records {
let is_no_resource_paper = is_no_resource(&rec.pdf_path) || is_no_resource(&rec.html_path);
if is_no_resource_paper {
continue;
}
let download_completed = is_completed(&rec.pdf_path) || is_completed(&rec.html_path);
let download_failed = is_failed(&rec.pdf_path) || is_failed(&rec.html_path);
let download_uncompleted = !download_completed && !download_failed;
let parse_completed = is_completed(&rec.markdown_path);
let parse_failed = is_failed(&rec.markdown_path);
let parse_uncompleted = !parse_completed && !parse_failed;
let translate_completed = is_completed(&rec.translation_path);
let translate_failed = is_failed(&rec.translation_path);
match target_phase.as_str() {
"download" => {
if skip_completed && download_completed {
continue;
}
if skip_failed && download_failed {
continue;
}
}
"parse" => {
if skip_completed && parse_completed {
continue;
}
if skip_failed && parse_failed {
continue;
}
if skip_preceding_failed && download_failed {
continue;
}
if skip_preceding_uncompleted && download_uncompleted {
continue;
}
}
"translate" => {
if skip_completed && translate_completed {
continue;
}
if skip_failed && translate_failed {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
"embed" => {
if skip_completed && rec.has_vector {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
"target" => {
if skip_completed && rec.has_target {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
_ => {}
}
target_bibcodes.push(rec.bibcode);
}
let limit = req.limit_count.unwrap_or(100) as usize;
if target_bibcodes.len() > limit {
target_bibcodes.truncate(limit);
}
let target_bibcodes = crate::services::batch::filter_papers_for_batch(
&state.db,
&target_phase,
req.limit_count,
req.sort_order.clone(),
skip_completed,
skip_failed,
skip_preceding_failed,
skip_preceding_uncompleted,
)
.await
.map_err(|e| AppError::internal(format!("筛选待处理文献失败: {}", e)))?;
if target_bibcodes.is_empty() {
// 没有需要处理的文献 —— 返回 OK软成功前端据 status 判断)
@ -316,37 +173,16 @@ pub async fn stop_asset_batch(State(state): State<Arc<AppState>>) -> StatusCode
}
// ── GET /api/sync/queries ──
#[derive(Debug, Serialize, Deserialize)]
pub struct SavedSyncQuery {
pub id: i64,
pub query: String,
pub source: String,
pub limit_count: i32,
pub last_run: String,
}
pub async fn get_sync_queries(
State(state): State<Arc<AppState>>,
) -> ApiResult<Json<Vec<SavedSyncQuery>>> {
let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') FROM sync_queries ORDER BY last_run DESC")
.fetch_all(&state.db)
) -> ApiResult<Json<Vec<crate::services::batch::SavedSyncQuery>>> {
let list = crate::services::batch::MetaSync::list_queries(&state.db)
.await
.map_err(|e| {
error!("获取已存同步检索配置失败: {}", e);
AppError::internal(format!("获取已存同步检索配置失败: {}", e))
})?;
let mut list = Vec::new();
for r in rows {
list.push(SavedSyncQuery {
id: r.get(0),
query: r.get(1),
source: r.get(2),
limit_count: r.get(3),
last_run: r.get(4),
});
}
Ok(Json(list))
}
@ -355,15 +191,20 @@ pub async fn delete_sync_query(
State(state): State<Arc<AppState>>,
axum::extract::Path(id): axum::extract::Path<i64>,
) -> ApiResult<StatusCode> {
sqlx::query("DELETE FROM sync_queries WHERE id = ?")
.bind(id)
.execute(&state.db)
let success = crate::services::batch::MetaSync::delete_query(&state.db, id)
.await
.map_err(|e| {
error!("删除同步检索配置失败: {}", e);
AppError::internal(format!("删除同步检索配置失败: {}", e))
})?;
if !success {
return Err(AppError::not_found(format!(
"未找到 ID 为 {} 的检索配置",
id
)));
}
Ok(StatusCode::OK)
}

View File

@ -37,8 +37,7 @@ pub async fn query_target(
State(state): State<Arc<AppState>>,
Query(params): Query<TargetQueryParams>,
) -> ApiResult<Json<TargetInfo>> {
let client = reqwest::Client::new();
let info = query_target_cached(&state.db, &params.object_name, None, &client)
let info = query_target_cached(&state.db, &params.object_name, None, &state.http_client)
.await
.map_err(|e| {
tracing::error!("查询天体信息失败 ({}): {}", params.object_name, e);
@ -53,20 +52,23 @@ pub async fn associate_target(
State(state): State<Arc<AppState>>,
Json(req): Json<AssociateTargetRequest>,
) -> ApiResult<Json<AssociateResponse>> {
let client = reqwest::Client::new();
// 查询并将结果缓存/关联到对应文献
let info = query_target_cached(&state.db, &req.object_name, Some(&req.bibcode), &client)
.await
.map_err(|e| {
tracing::error!(
"手动关联天体失败 ({} -> {}): {}",
req.object_name,
req.bibcode,
e
);
AppError::internal(format!("手动关联天体失败: {}", e))
})?;
let info = query_target_cached(
&state.db,
&req.object_name,
Some(&req.bibcode),
&state.http_client,
)
.await
.map_err(|e| {
tracing::error!(
"手动关联天体失败 ({} -> {}): {}",
req.object_name,
req.bibcode,
e
);
AppError::internal(format!("手动关联天体失败: {}", e))
})?;
Ok(Json(AssociateResponse {
status: "success".to_string(),
@ -79,80 +81,12 @@ pub async fn list_targets(
State(state): State<Arc<AppState>>,
Query(params): Query<TargetListParams>,
) -> ApiResult<Json<Vec<TargetInfo>>> {
let rows = sqlx::query_as::<
_,
(
String,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<String>,
Option<f64>,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<f64>,
Option<String>,
Option<String>,
),
>(
"SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \
otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \
FROM paper_targets WHERE bibcode = ? ORDER BY target_name",
)
.bind(&params.bibcode)
.fetch_all(&state.db)
.await
.map_err(|e| {
tracing::error!("获取文献天体关联列表失败 ({}): {}", params.bibcode, e);
AppError::internal(format!("获取文献天体关联列表失败: {}", e))
})?;
let targets: Vec<TargetInfo> = rows
.into_iter()
.map(
|(
name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry_json,
aliases_json,
)| {
let aliases: Vec<String> = aliases_json
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
photometry_json.and_then(|s| serde_json::from_str(&s).ok());
TargetInfo {
target_name: name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry,
aliases,
}
},
)
.collect();
let targets = crate::services::target::list_targets_for_paper(&state.db, &params.bibcode)
.await
.map_err(|e| {
tracing::error!("获取文献天体关联列表失败 ({}): {}", params.bibcode, e);
AppError::internal(format!("获取文献天体关联列表失败: {}", e))
})?;
Ok(Json(targets))
}
@ -174,43 +108,24 @@ pub async fn extract_paper_targets(
) -> ApiResult<Json<ExtractTargetsResponse>> {
tracing::info!("接收到文献天体提取与识别指令: {}", req.bibcode);
let (_, _, md_opt, _) = crate::api::helpers::check_paper_paths_in_db(
let targets = crate::services::target::extract_and_refresh_targets(
&state.db,
&state.config.library_dir,
&req.bibcode,
&state.http_client,
)
.await
.map_err(|e| AppError::not_found(format!("获取文献路径失败: {}", e)))?
.ok_or_else(|| AppError::not_found("该文献未注册在数据库中"))?;
let md_rel =
md_opt.ok_or_else(|| AppError::bad_request("文献尚未解析为 Markdown请先执行解析"))?;
let md_abs = state.config.library_dir.join(&md_rel);
if !md_abs.exists() {
return Err(AppError::not_found("文献 Markdown 文件未找到,请重新解析"));
}
let markdown_content = tokio::fs::read_to_string(&md_abs)
.await
.map_err(|e| AppError::internal(format!("读取 Markdown 文件失败: {}", e)))?;
// 重新识别前先清除该文献已有的天体关联记录,确保陈旧和错误绑定的天体得到重置与刷新
if let Err(e) = sqlx::query("DELETE FROM paper_targets WHERE bibcode = ?")
.bind(&req.bibcode)
.execute(&state.db)
.await
{
tracing::error!("清除文献 {} 的旧天体关联失败: {}", req.bibcode, e);
}
let client = reqwest::Client::new();
let targets = crate::services::target::extract_and_cache_targets(
&state.db,
&markdown_content,
&req.bibcode,
&client,
)
.await;
.map_err(|e| {
let err_msg = e.to_string();
if err_msg.contains("该文献未注册") || err_msg.contains("未找到") {
AppError::not_found("未在本地库中找到该文献")
} else if err_msg.contains("尚未解析") {
AppError::bad_request("请先完成文献解析后再提取天体目标")
} else {
tracing::error!("天体目标提取内部错误: {}", err_msg);
AppError::internal(err_msg)
}
})?;
Ok(Json(ExtractTargetsResponse { targets }))
}

View File

@ -3,12 +3,20 @@
// AstroResearch CLI Skills Agent — 向外部 Agent如 Claude Code提供
// 标准的命令行工具接口,用于 RAG 问答、天体查询和天体关联操作。
use anyhow::Context;
use clap::{Parser, Subcommand};
use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::Arc;
use tracing_subscriber::FmtSubscriber;
use astroresearch::api::AppState;
use astroresearch::clients::ads::AdsClient;
use astroresearch::clients::arxiv::ArxivClient;
use astroresearch::clients::llm::{EmbeddingClient, LlmClient};
use astroresearch::clients::qiniu::QiniuClient;
use astroresearch::services::download::Downloader;
use astroresearch::Config;
#[derive(Parser)]
@ -61,6 +69,89 @@ enum Commands {
/// Markdown 文件路径
markdown_path: String,
},
/// 读取已解析文献的 Markdown 全文或指定章节内容
Content {
/// 文献唯一标识符,支持 Bibcode、DOI 或 arXiv ID
bibcode: String,
/// 章节序号0-based可选
#[arg(short, long)]
section_index: Option<usize>,
/// 章节名称(支持模糊匹配,可选)
#[arg(short = 'n', long)]
section_name: Option<String>,
},
/// 获取文献的章节大纲目录
Outline {
/// 文献唯一标识符,支持 Bibcode、DOI 或 arXiv ID
bibcode: String,
},
/// 浏览和检索文献的引用网络关系(参考文献或被引列表)
Citations {
/// 文献唯一标识符,支持 Bibcode、DOI 或 arXiv ID
bibcode: String,
/// 可选的模糊检索词,如 'Lei 2023'、'Zhang'
#[arg(short, long)]
query: Option<String>,
/// 查询方向references参考文献默认或 citations被引
#[arg(short, long, default_value = "references")]
direction: String,
/// 排序方式citation_count默认、year 或 default
#[arg(short, long, default_value = "citation_count")]
sort: String,
/// 分页偏移量
#[arg(long, default_value = "0")]
offset: i64,
/// 每页返回数量上限
#[arg(long, default_value = "20")]
limit: i64,
},
/// 在本地已下载的文献库中进行全文检索
LocalSearch {
/// 检索词
query: String,
/// 返回结果数量上限
#[arg(short, long, default_value = "10")]
limit: usize,
},
/// 在云端NASA ADS & arXiv联合检索天文学文献
Search {
/// 检索词或高级检索式
query: String,
/// 返回结果数量上限
#[arg(short, long, default_value = "5")]
limit: i32,
},
/// 对文献执行下载、解析 (PDF -> MD) 和向量化嵌入一站式流水线处理
Process {
/// 文献唯一标识符,支持 Bibcode、DOI 或 arXiv ID
bibcode: String,
/// 是否执行下载任务
#[arg(long)]
download: bool,
/// 是否执行 PDF 解析任务
#[arg(long)]
parse: bool,
/// 是否执行向量化嵌入任务
#[arg(long)]
embed: bool,
/// 是否强制重新执行(跳过已有结果检查)
#[arg(short, long)]
force: bool,
},
/// 保存一篇研究笔记到本地 Markdown 文件中
NoteSave {
/// 笔记标题(会被安全化作为文件名)
title: String,
/// 笔记正文内容
content: String,
},
}
#[tokio::main]
@ -71,10 +162,7 @@ async fn main() -> anyhow::Result<()> {
.finish();
tracing_subscriber::util::SubscriberInitExt::init(subscriber);
// SAFETY: 注册静态 sqlite-vec 初始化函数。transmute 是安全的,
// 因为 sqlite3_vec_init 符合 SQLite C API 的自动扩展回调签名
// (sqlite3*, char**, const sqlite3_api_routines*)。
// 该注册必须在开启任何数据库连接前执行。
// SAFETY: 注册静态 sqlite-vec 初始化函数。
unsafe {
libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute::<
*const (),
@ -102,29 +190,92 @@ async fn main() -> anyhow::Result<()> {
// 执行迁移
sqlx::migrate!("./migrations").run(&pool).await?;
// 构建极简的 AppState用于复用所有的服务层代码
let app_state = Arc::new(AppState {
config: config.clone(),
db: pool.clone(),
dict: astroresearch::services::translation::Dictionary::new(),
qiniu: QiniuClient::new(
config.qiniu_ak.clone(),
config.qiniu_sk.clone(),
config.qiniu_bucket.clone(),
config.qiniu_domain.clone(),
),
ads: AdsClient::new(config.ads_api_key.clone()).context("构建 ADS 客户端失败")?,
arxiv: ArxivClient::new().context("构建 arXiv 客户端失败")?,
llm: LlmClient::new(
config.llm_api_key.clone(),
config.llm_api_base.clone(),
config.llm_model.clone(),
)
.context("构建 LLM 客户端失败")?,
medium_llm: LlmClient::new(
config.llm_medium_api_key.clone(),
config.llm_medium_api_base.clone(),
config.llm_medium_model.clone(),
)
.context("构建 Medium LLM 客户端失败")?,
fast_llm: LlmClient::new(
config.llm_fast_api_key.clone(),
config.llm_fast_api_base.clone(),
config.llm_fast_model.clone(),
)
.context("构建 Fast LLM 客户端失败")?,
vision_llm: None,
embedding: EmbeddingClient::new(
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
)
.context("构建 Embedding 客户端失败")?,
downloader: Downloader::new()?,
http_client: reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to create shared HTTP client")?,
harvest_status: Arc::new(tokio::sync::Mutex::new(
astroresearch::services::batch::MetaSyncStatus::new(),
)),
batch_status: Arc::new(tokio::sync::Mutex::new(
astroresearch::services::batch::AssetBatchStatus::new(),
)),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
skill_registry: Arc::new(tokio::sync::RwLock::new(
astroresearch::agent::skills::SkillRegistry::new(config.skills_dir.clone()),
)),
pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
session_permission_checkers: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
sse_broadcast: {
let (tx, _) = tokio::sync::broadcast::channel(256);
Some(tx)
},
memory_manager: Arc::new(tokio::sync::Mutex::new(
astroresearch::agent::memory::MemoryManager::new(config.library_dir.clone()),
)),
sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
login_rate_limiter: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
});
let cli = Cli::parse();
match cli.command {
Commands::Rag { question, top_k } => {
let embedding = EmbeddingClient::new(
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
);
let llm = LlmClient::new(
config.llm_api_key.clone(),
config.llm_api_base.clone(),
config.llm_model.clone(),
);
let sources = astroresearch::services::rag::retrieve_hybrid(
&pool,
&app_state.embedding,
&question,
top_k,
)
.await?;
let result =
astroresearch::services::rag::ask(&pool, &embedding, &llm, &question, top_k)
.await?;
println!("\n📖 回答:\n{}\n", result.answer);
if !result.sources.is_empty() {
println!("📚 参考来源:");
for (i, src) in result.sources.iter().enumerate() {
if sources.is_empty() {
println!("\n未找到匹配的文献段落。\n");
} else {
println!("\n🔍 检索到 {} 条结果:\n", sources.len());
for (i, src) in sources.iter().enumerate() {
let heading_info = if src.headings.is_empty() || src.headings == "Document" {
String::new()
} else {
@ -138,6 +289,17 @@ async fn main() -> anyhow::Result<()> {
heading_info,
src.distance
);
// 输出片段内容的前 200 字符作为预览
let preview: String = src.content.chars().take(200).collect();
println!(
" {}{}\n",
preview,
if src.content.chars().count() > 200 {
"..."
} else {
""
}
);
}
}
}
@ -206,19 +368,259 @@ async fn main() -> anyhow::Result<()> {
markdown_path,
} => {
let content = std::fs::read_to_string(&markdown_path)?;
let embedding = EmbeddingClient::new(
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
);
let count = astroresearch::services::rag::ingest_paper(
&pool, &embedding, &bibcode, &content, None,
&pool,
&app_state.embedding,
&bibcode,
&content,
None,
)
.await?;
println!("✅ 文献 {} 向量化完成,写入 {} 个切片", bibcode, count);
}
Commands::Content {
bibcode,
section_index,
section_name,
} => {
let mode = if let Some(idx) = section_index {
astroresearch::services::paper::ReadMode::SectionIndex(idx)
} else if let Some(name) = section_name {
astroresearch::services::paper::ReadMode::SectionName(name)
} else {
astroresearch::services::paper::ReadMode::Full {
include_translation: false,
}
};
let res = astroresearch::services::paper::read_paper_content(
&pool,
&config.library_dir,
&bibcode,
mode,
)
.await?;
println!("{}", res.content);
}
Commands::Outline { bibcode } => {
let mode = astroresearch::services::paper::ReadMode::Outline { max_level: Some(3) };
let res = astroresearch::services::paper::read_paper_content(
&pool,
&config.library_dir,
&bibcode,
mode,
)
.await?;
if let Some(sections) = res.outline {
println!("📖 文献 {} 章节大纲:", res.bibcode);
for s in &sections {
let indent = " ".repeat(s.level.saturating_sub(1));
println!("{}[{}] {}", indent, s.index, s.heading);
}
} else {
println!("该文献未检测到明显的章节标题。");
}
}
Commands::Citations {
bibcode,
query,
direction,
sort,
offset,
limit,
} => {
let paper = astroresearch::services::paper::get_paper_from_db(
&pool,
&config.library_dir,
&bibcode,
)
.await?;
if let Some(ref q) = query {
let results = astroresearch::services::citation::search_citations(
&pool,
&paper.bibcode,
&direction,
q,
)
.await?;
if results.is_empty() {
println!("未找到匹配 '{}' 的引文。", q);
} else {
println!("🎯 找到 {} 篇匹配 '{}' 的引文:", results.len(), q);
for (i, r) in results.iter().enumerate() {
println!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
i + 1,
r.bibcode,
r.title,
r.year.as_deref().unwrap_or(""),
r.authors.join(", "),
r.citation_count
);
}
}
} else {
let (rows, total) = astroresearch::services::citation::get_citations_paginated(
&pool,
&paper.bibcode,
&direction,
&sort,
offset,
limit,
)
.await?;
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
println!(
"📚 文献 {} 的{}列表 (共 {} 篇, 当前展示 {}-{} 条):",
paper.bibcode,
dir_label,
total,
offset + 1,
offset + rows.len() as i64
);
for (i, r) in rows.iter().enumerate() {
println!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
offset + i as i64 + 1,
r.bibcode,
r.title,
r.year.as_deref().unwrap_or(""),
r.authors.join(", "),
r.citation_count
);
}
}
}
Commands::LocalSearch { query, limit } => {
let results =
astroresearch::services::search::search_local_library(&pool, &query, limit).await?;
if results.is_empty() {
println!("本地文献库中未找到匹配的文献。");
} else {
println!("🔎 本地检索到 {} 篇匹配文献:", results.len());
for (i, p) in results.iter().enumerate() {
println!(
"{}. [{}] {} ({})\n 第一作者: {} | 期刊: {}\n 被引: {} 次 | 已解析: {}",
i + 1,
p.bibcode,
p.title,
p.year,
p.authors
.first()
.cloned()
.unwrap_or_else(|| "未知".to_string()),
p.pub_journal,
p.citation_count,
if p.has_markdown { "" } else { "" }
);
}
}
}
Commands::Search { query, limit } => {
let results = astroresearch::services::search::search_papers(
&app_state,
&query,
"all",
0,
limit,
"relevance",
)
.await?;
if results.is_empty() {
println!("未找到任何匹配的云端文献。");
} else {
println!("☁️ 云端联合检索到 {} 篇文献:", results.len());
for (i, p) in results.iter().enumerate() {
println!(
"{}. [{}] {} ({})\n 第一作者: {}\n 被引: {} 次",
i + 1,
p.bibcode,
p.title,
p.year,
p.authors
.first()
.cloned()
.unwrap_or_else(|| "未知".to_string()),
p.citation_count
);
}
}
}
Commands::Process {
bibcode,
download,
parse,
embed,
force,
} => {
println!("⚙️ 开始处理文献 {} ...", bibcode);
let mut tasks = Vec::new();
if download {
tasks.push("download".to_string());
}
if parse {
tasks.push("parse".to_string());
}
if embed {
tasks.push("embed".to_string());
}
if tasks.is_empty() {
println!(
"⚠️ 未指定任何处理任务。请使用 --download、--parse 或 --embed 指定任务步骤。"
);
return Ok(());
}
let res = astroresearch::services::pipeline::process_paper_pipeline(
&app_state, &bibcode, &tasks, force,
)
.await?;
println!(
"🎉 文献 {} 处理成功!\n 已完成任务: {:?}\n PDF: {}\n HTML: {}\n Markdown: {}\n 写入向量分块数: {}",
res.bibcode,
res.completed,
if res.has_pdf { "" } else { "" },
if res.has_html { "" } else { "" },
if res.has_markdown { "" } else { "" },
res.chunk_count
);
}
Commands::NoteSave { title, content } => {
let (filename, filepath, size) = astroresearch::services::note::save_note_service(
&config.library_dir,
&title,
&content,
)?;
println!(
"✅ 笔记已成功保存!\n 文件名: {}\n 物理路径: {}\n 总字符数: {}",
filename,
filepath.display(),
size
);
}
}
Ok(())

View File

@ -1,7 +1,8 @@
// src/ads.rs
use anyhow::Context;
use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
use serde::{Deserialize, Serialize};
use tracing::{error, info};
use tracing::{error, info, warn};
// 原始 ADS API 返回的数据文档结构
#[derive(Debug, Clone, Serialize, Deserialize)]
@ -46,11 +47,16 @@ pub struct AdsClient {
}
impl AdsClient {
pub fn new(api_key: String) -> Self {
AdsClient {
pub fn new(api_key: String) -> anyhow::Result<Self> {
Ok(AdsClient {
api_key,
client: reqwest::Client::new(),
}
client: reqwest::Client::builder()
.redirect(crate::utils::ssrf::safe_redirect_policy())
.timeout(std::time::Duration::from_secs(30))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to create ADS HTTP client")?,
})
}
// 拼装鉴权 Header
@ -95,19 +101,43 @@ impl AdsClient {
let start_str = start.to_string();
let rows_str = rows.to_string();
let response = self
.client
.get(url)
.headers(self.headers())
.query(&[
("q", translated.as_str()),
("start", start_str.as_str()),
("rows", rows_str.as_str()),
("fl", fl),
("sort", ads_sort),
])
.send()
.await?;
// 带重试的请求发送(处理 429 速率限制)
const MAX_RETRIES: u32 = 3;
let mut response = None;
for attempt in 0..MAX_RETRIES {
let resp = self
.client
.get(url)
.headers(self.headers())
.query(&[
("q", translated.as_str()),
("start", start_str.as_str()),
("rows", rows_str.as_str()),
("fl", fl),
("sort", ads_sort),
])
.send()
.await?;
if resp.status().as_u16() == 429 && attempt < MAX_RETRIES - 1 {
let retry_after = resp
.headers()
.get("retry-after")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(5);
warn!(
"ADS 速率限制 (429){} 秒后重试 (第 {} 次)",
retry_after,
attempt + 1
);
tokio::time::sleep(std::time::Duration::from_secs(retry_after)).await;
continue;
}
response = Some(resp);
break;
}
let response = response.unwrap();
if !response.status().is_success() {
let status = response.status();
@ -258,8 +288,8 @@ mod tests {
println!("警告: 未在环境配置中检测到 ADS_API_KEY跳过 ADS 集成测试。");
}
let ads = AdsClient::new(config.ads_api_key.clone());
let arxiv = ArxivClient::new();
let ads = AdsClient::new(config.ads_api_key.clone()).unwrap();
let arxiv = ArxivClient::new().unwrap();
println!("================= 开始真实检索逻辑集成测试 =================");

View File

@ -1,7 +1,8 @@
// src/arxiv.rs
use anyhow::Context;
use regex::Regex;
use serde::{Deserialize, Serialize};
use tracing::{error, info};
use tracing::{error, info, warn};
// 统一的 arXiv 文献临时结构
#[derive(Debug, Clone, Serialize, Deserialize)]
@ -23,15 +24,26 @@ pub struct ArxivClient {
impl Default for ArxivClient {
fn default() -> Self {
Self::new()
Self::new().unwrap_or_else(|e| {
tracing::warn!("ArxivClient::new() failed in Default impl: {}", e);
// Return a client with a dummy reqwest client — searches will fail gracefully
ArxivClient {
client: reqwest::Client::new(),
}
})
}
}
impl ArxivClient {
pub fn new() -> Self {
ArxivClient {
client: reqwest::Client::new(),
}
pub fn new() -> anyhow::Result<Self> {
Ok(ArxivClient {
client: reqwest::Client::builder()
.redirect(crate::utils::ssrf::safe_redirect_policy())
.timeout(std::time::Duration::from_secs(30))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to create arXiv HTTP client")?,
})
}
// 请求 arXiv 官方的 Export 检索接口并解析返回内容,支持分页与排序
@ -66,18 +78,42 @@ impl ArxivClient {
let start_str = start.to_string();
let max_results_str = max_results.to_string();
let response = self
.client
.get(url)
.query(&[
("search_query", final_query.as_str()),
("start", start_str.as_str()),
("max_results", max_results_str.as_str()),
("sortBy", sort_by),
("sortOrder", sort_order),
])
.send()
.await?;
// 带重试的请求发送(处理 503 临时不可用)
const MAX_RETRIES: u32 = 3;
let mut response = None;
for attempt in 0..MAX_RETRIES {
let resp = self
.client
.get(url)
.query(&[
("search_query", final_query.as_str()),
("start", start_str.as_str()),
("max_results", max_results_str.as_str()),
("sortBy", sort_by),
("sortOrder", sort_order),
])
.send()
.await?;
if resp.status().as_u16() == 503 && attempt < MAX_RETRIES - 1 {
let retry_after = resp
.headers()
.get("retry-after")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(5);
warn!(
"arXiv 服务暂时不可用 (503){} 秒后重试 (第 {} 次)",
retry_after,
attempt + 1
);
tokio::time::sleep(std::time::Duration::from_secs(retry_after)).await;
continue;
}
response = Some(resp);
break;
}
let response = response.unwrap();
if !response.status().is_success() {
let status = response.status();
@ -120,9 +156,12 @@ impl ArxivClient {
}
let xml_content = response.text().await?;
let total_re =
Regex::new(r"<opensearch:totalResults[^>]*>(\d+)</opensearch:totalResults>").unwrap();
if let Some(caps) = total_re.captures(&xml_content) {
// Use the same cached regex from parse_arxiv_xml
use std::sync::LazyLock;
static TOTAL_RE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"<opensearch:totalResults[^>]*>(\d+)</opensearch:totalResults>").unwrap()
});
if let Some(caps) = TOTAL_RE.captures(&xml_content) {
if let Ok(count) = caps[1].parse::<i32>() {
return Ok(count);
}
@ -133,27 +172,36 @@ impl ArxivClient {
// 使用正则表达式手动提取 XML 内容,避免由于命名空间前缀不同造成的反序列化问题
fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
use std::sync::LazyLock;
static ENTRY_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<entry>(.*?)</entry>").unwrap());
static ID_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<id>http://arxiv.org/abs/(.*?)(?:v\d+)?</id>").unwrap());
static TITLE_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<title>(.*?)</title>").unwrap());
static SUMMARY_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<summary>(.*?)</summary>").unwrap());
static PUBLISHED_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<published>(\d{4})-\d{2}-\d{2}").unwrap());
static AUTHOR_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<author>\s*<name>(.*?)</name>").unwrap());
static DOI_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<arxiv:doi[^>]*>(.*?)</arxiv:doi>").unwrap());
static FALLBACK_ID_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<id>(.*?)</id>").unwrap());
let mut papers = Vec::new();
let entry_re = Regex::new(r"(?s)<entry>(.*?)</entry>").unwrap();
let id_re = Regex::new(r"<id>http://arxiv.org/abs/(.*?)(?:v\d+)?</id>").unwrap();
let title_re = Regex::new(r"(?s)<title>(.*?)</title>").unwrap();
let summary_re = Regex::new(r"(?s)<summary>(.*?)</summary>").unwrap();
let published_re = Regex::new(r"<published>(\d{4})-\d{2}-\d{2}").unwrap();
let author_re = Regex::new(r"(?s)<author>\s*<name>(.*?)</name>").unwrap();
let doi_re = Regex::new(r"<arxiv:doi[^>]*>(.*?)</arxiv:doi>").unwrap();
let fallback_id_re = Regex::new(r"<id>(.*?)</id>").unwrap();
for cap in entry_re.captures_iter(xml) {
for cap in ENTRY_RE.captures_iter(xml) {
let entry_content = &cap[1];
// 提取并清洗 ID
let id = id_re
let id = ID_RE
.captures(entry_content)
.map(|c| c[1].trim().to_string())
.unwrap_or_else(|| {
fallback_id_re
FALLBACK_ID_RE
.captures(entry_content)
.map(|c| c[1].trim().to_string())
.unwrap_or_default()
@ -164,7 +212,7 @@ fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
}
// 提取标题,清理换行与连续空格
let mut title = title_re
let mut title = TITLE_RE
.captures(entry_content)
.map(|c| c[1].to_string())
.unwrap_or_default();
@ -175,7 +223,7 @@ fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
.to_string();
// 提取摘要
let mut abstract_text = summary_re
let mut abstract_text = SUMMARY_RE
.captures(entry_content)
.map(|c| c[1].to_string())
.unwrap_or_default();
@ -186,14 +234,14 @@ fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
.to_string();
// 提取发布年份
let year = published_re
let year = PUBLISHED_RE
.captures(entry_content)
.map(|c| c[1].to_string())
.unwrap_or_else(|| "未知".to_string());
// 提取作者列表
let mut authors = Vec::new();
for auth_cap in author_re.captures_iter(entry_content) {
for auth_cap in AUTHOR_RE.captures_iter(entry_content) {
let author_name = auth_cap[1].trim().to_string();
if !author_name.is_empty() {
authors.push(author_name);
@ -201,7 +249,7 @@ fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
}
// 提取关联 DOI
let doi = doi_re
let doi = DOI_RE
.captures(entry_content)
.map(|c| c[1].trim().to_string());

View File

@ -2,6 +2,7 @@
//
// ChatCompleter trait + LlmClient 实现(对话补全、流式对话、图片分析)。
use anyhow::Context;
use async_trait::async_trait;
use futures_util::StreamExt;
use reqwest::Client;
@ -35,13 +36,19 @@ impl ChatCompleter for LlmClient {
}
impl LlmClient {
pub fn new(api_key: String, api_base: String, model: String) -> Self {
LlmClient {
pub fn new(api_key: String, api_base: String, model: String) -> anyhow::Result<Self> {
let client = Client::builder()
.timeout(std::time::Duration::from_secs(120))
.connect_timeout(std::time::Duration::from_secs(10))
.redirect(crate::utils::ssrf::safe_redirect_policy())
.build()
.context("Failed to create LLM HTTP client")?;
Ok(LlmClient {
api_key,
api_base,
model: Arc::new(std::sync::RwLock::new(model)),
client: Client::new(),
}
client,
})
}
pub fn model(&self) -> String {
@ -97,8 +104,9 @@ impl LlmClient {
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
error!("LLM 接口调用失败: 状态码={}, 报错={}", status, body);
let _body = response.text().await.unwrap_or_default();
// 仅记录 status code不记录 body可能包含 API key 等敏感信息)
error!("LLM 接口调用失败: 状态码={}", status);
return Err(anyhow::anyhow!("大模型接口返回错误状态: {}", status));
}
@ -165,13 +173,9 @@ impl LlmClient {
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
));
let _body = response.text().await.unwrap_or_default();
error!("LLM chat 接口调用失败: 状态码={}", status);
return Err(anyhow::anyhow!("大模型 chat 接口返回错误状态: {}", status));
}
#[derive(Deserialize)]
@ -242,12 +246,9 @@ impl LlmClient {
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(anyhow::anyhow!(
"视觉模型调用失败 (HTTP {}): {}",
status,
body
));
let _body = response.text().await.unwrap_or_default();
// 仅记录 status code不记录 body可能包含 API key 等敏感信息)
return Err(anyhow::anyhow!("视觉模型调用失败 (HTTP {})", status,));
}
use futures_util::StreamExt;
@ -334,12 +335,12 @@ impl LlmClient {
.get("retry-after")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<u64>().ok());
let body = response.text().await.unwrap_or_default();
error!("LLM stream 接口调用失败: 状态码={}, 报错={}", status, body);
let _body = response.text().await.unwrap_or_default();
// 仅记录 status code不记录 body可能包含 API key 等敏感信息)
error!("LLM stream 接口调用失败: 状态码={}", status);
return Err(anyhow::anyhow!(
"HTTP {}: {} | retry_after={:?}",
"HTTP {} | retry_after={:?}",
status_code,
body,
retry_after
));
}
@ -565,7 +566,8 @@ mod tests {
#[test]
fn test_llm_client_initialization() {
let client = LlmClient::new("key".to_string(), "base".to_string(), "model".to_string());
let client =
LlmClient::new("key".to_string(), "base".to_string(), "model".to_string()).unwrap();
assert_eq!(client.api_key(), "key");
assert_eq!(client.api_base(), "base");
assert_eq!(client.model(), "model");
@ -708,7 +710,8 @@ mod tests {
config.llm_api_key.clone(),
config.llm_api_base.clone(),
config.llm_model.clone(),
);
)
.unwrap();
match llm
.chat_completion("You are a helpful assistant.", "Say Hello!")
.await
@ -735,7 +738,8 @@ mod tests {
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
);
)
.unwrap();
let test_text = "active galactic nucleus";
match embedding_client.create_embedding(test_text).await {
Ok(vector) => {

View File

@ -2,6 +2,7 @@
//
// Embedder trait + EmbeddingClient 实现(文本向量化)。
use anyhow::Context;
use async_trait::async_trait;
use reqwest::Client;
use serde::Deserialize;
@ -67,49 +68,75 @@ impl Embedder for EmbeddingClient {
let mut all_embeddings: Vec<Vec<f32>> = Vec::with_capacity(texts.len());
// 分片处理:每批 BATCH_SIZE 条,避免单请求过大触发 413/429
for chunk in texts.chunks(BATCH_SIZE) {
// 并发分片处理:每批 BATCH_SIZE 条semaphore(5) 控制并发度
let sem = std::sync::Arc::new(tokio::sync::Semaphore::new(5));
let mut tasks = tokio::task::JoinSet::new();
for (batch_idx, chunk) in texts.chunks(BATCH_SIZE).enumerate() {
let payload = serde_json::json!({
"model": self.model(),
"input": chunk,
});
let client = self.client.clone();
let api_key = self.api_key.clone();
let url = url.clone();
let sem = sem.clone();
let batch_size = chunk.len();
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&payload)
.send()
.await?;
tasks.spawn(async move {
let _permit = sem.acquire().await;
let response = client
.post(&url)
.header("Authorization", format!("Bearer {}", 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!(
"Embedding 批量接口调用失败: 状态码={}, 报错={}",
status, body
);
return Err(anyhow::anyhow!("向量接口返回错误状态: {}", status));
(batch_idx, batch_size, response)
});
}
let num_batches = texts.chunks(BATCH_SIZE).len();
let mut batch_results: Vec<Option<(usize, Vec<Vec<f32>>)>> = vec![None; num_batches];
while let Some(result) = tasks.join_next().await {
match result {
Ok((batch_idx, batch_size, Ok(response))) => {
if !response.status().is_success() {
let status = response.status();
return Err(anyhow::anyhow!("向量接口返回错误状态: {}", status));
}
let res_data: EmbeddingResponse = response.json().await?;
// 按 index 排序(每批内 0-based
let mut indexed: Vec<(usize, Vec<f32>)> = res_data
.data
.into_iter()
.map(|d| (d.index, d.embedding))
.collect();
indexed.sort_by_key(|(i, _)| *i);
if indexed.len() != batch_size {
return Err(anyhow::anyhow!(
"向量批量接口返回数量({})与输入({})不一致",
indexed.len(),
batch_size
));
}
batch_results[batch_idx] =
Some((batch_idx, indexed.into_iter().map(|(_, v)| v).collect()));
}
Ok((_batch_idx, _batch_size, Err(e))) => {
return Err(e.into());
}
Err(join_err) => {
return Err(anyhow::anyhow!("并发向量任务 panicked: {}", join_err));
}
}
}
let res_data: EmbeddingResponse = response.json().await?;
// OpenAI 不保证 data 数组顺序,按 index 排序index 在每批内 0-based
let mut indexed: Vec<(usize, Vec<f32>)> = res_data
.data
.into_iter()
.map(|d| (d.index, d.embedding))
.collect();
indexed.sort_by_key(|(i, _)| *i);
if indexed.len() != chunk.len() {
return Err(anyhow::anyhow!(
"向量批量接口返回数量({})与输入({})不一致",
indexed.len(),
chunk.len()
));
}
all_embeddings.extend(indexed.into_iter().map(|(_, v)| v));
// 按批次顺序重组结果
for (_, embeddings) in batch_results.into_iter().flatten() {
all_embeddings.extend(embeddings);
}
Ok(all_embeddings)
@ -117,13 +144,19 @@ impl Embedder for EmbeddingClient {
}
impl EmbeddingClient {
pub fn new(api_key: String, api_base: String, model: String) -> Self {
EmbeddingClient {
pub fn new(api_key: String, api_base: String, model: String) -> anyhow::Result<Self> {
let client = Client::builder()
.timeout(std::time::Duration::from_secs(60))
.connect_timeout(std::time::Duration::from_secs(10))
.redirect(crate::utils::ssrf::safe_redirect_policy())
.build()
.context("Failed to create Embedding HTTP client")?;
Ok(EmbeddingClient {
api_key,
api_base,
model,
client: Client::new(),
}
client,
})
}
pub fn model(&self) -> &str {
@ -157,8 +190,9 @@ impl EmbeddingClient {
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
error!("Embedding 接口调用失败: 状态码={}, 报错={}", status, body);
let _body = response.text().await.unwrap_or_default();
// 仅记录 status code不记录 body可能包含 API key 等敏感信息)
error!("Embedding 接口调用失败: 状态码={}", status);
return Err(anyhow::anyhow!("向量接口返回错误状态: {}", status));
}
@ -188,7 +222,8 @@ mod tests {
#[test]
fn test_embedding_client_initialization() {
let client =
EmbeddingClient::new("key".to_string(), "base".to_string(), "model".to_string());
EmbeddingClient::new("key".to_string(), "base".to_string(), "model".to_string())
.unwrap();
assert_eq!(client.api_key(), "key");
assert_eq!(client.api_base(), "base");
assert_eq!(client.model(), "model");

View File

@ -103,7 +103,7 @@ impl Config {
8000
});
let admin_password = env::var("ADMIN_PASSWORD").unwrap_or_else(|_| "admin".to_string());
let admin_password = env::var("ADMIN_PASSWORD").unwrap_or_default();
let llm_vision_model = env::var("LLM_VISION_MODEL").unwrap_or_default();
let llm_vision_api_key = env::var("LLM_VISION_API_KEY").unwrap_or_default();

View File

@ -2,6 +2,7 @@
use anyhow::Context;
use axum::{
http::HeaderValue,
routing::{get, post},
Router,
};
@ -9,9 +10,10 @@ use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
use std::collections::HashMap;
use std::net::SocketAddr;
use std::str::FromStr;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::Arc;
use tower_http::cors::CorsLayer;
use tower_http::services::ServeDir;
use tower_http::set_header::SetResponseHeaderLayer;
use tracing::{error, info, warn};
use astroresearch::agent::skills::SkillRegistry;
@ -29,6 +31,24 @@ async fn main() -> anyhow::Result<()> {
// 1. 初始化日志记录器并保留异步写保护 Guard
let _logging_guards = astroresearch::services::logging::init_logging()?;
// 1.2 设置全局 panic hook — 记录 panic 信息到日志而非 stderr
// 防止 Agent 工具执行中的 panic 导致 tokio 任务静默失败
std::panic::set_hook(Box::new(|panic| {
let payload = panic.payload();
let msg = if let Some(s) = payload.downcast_ref::<&str>() {
s.to_string()
} else if let Some(s) = payload.downcast_ref::<String>() {
s.clone()
} else {
"Box<dyn Any>".to_string()
};
let location = panic
.location()
.map(|l| format!(" at {}:{}", l.file(), l.line()))
.unwrap_or_default();
tracing::error!("PANIC{}: {}", location, msg);
}));
info!("正在启动 AstroResearch 天文学文献辅助系统后端服务...");
// 1.5 静态注册 sqlite-vec 自动扩展
@ -57,7 +77,7 @@ async fn main() -> anyhow::Result<()> {
"系统配置成功载入。本地 SQLite 连接串: {}",
config.database_url
);
// 弱密码告警
// 密码安全检查
let pwd = &config.admin_password;
let default_pwds = [
"admin",
@ -68,7 +88,9 @@ async fn main() -> anyhow::Result<()> {
"admin123",
"astroresearch",
];
if pwd.len() < 6 {
if pwd.is_empty() {
anyhow::bail!("启动失败: ADMIN_PASSWORD 未设置。请在 .env 文件中设置一个强密码后重试。");
} else if pwd.len() < 6 {
warn!(
"⚠️ 安全警告: ADMIN_PASSWORD 长度仅 {} 个字符,属弱密码。请设至少 8 位的强密码。",
pwd.len()
@ -121,9 +143,24 @@ async fn main() -> anyhow::Result<()> {
if let Some((sql,)) = existing_sql {
let expected_pattern = format!("float[{}]", embedding_dim);
if !sql.contains(&expected_pattern) {
info!(
"检测到已存在的向量表维度不匹配,正在重建以适配当前维度: {}...",
embedding_dim
let force_rebuild = std::env::var("EMBEDDING_DIM_FORCE_REBUILD")
.map(|v| v == "1" || v == "true")
.unwrap_or(false);
if !force_rebuild {
anyhow::bail!(
"向量表维度不匹配 (当前: {}, 期望: {})。\n\
\n\
EMBEDDING_DIM_FORCE_REBUILD=1 ",
sql,
expected_pattern
);
}
warn!(
"⚠️ EMBEDDING_DIM_FORCE_REBUILD=1 已设置,开始重建向量表。\
: {}, : {}",
sql, expected_pattern
);
sqlx::query("DROP TABLE IF EXISTS vec_paper_chunks")
.execute(&pool)
@ -131,6 +168,10 @@ async fn main() -> anyhow::Result<()> {
sqlx::query("DELETE FROM paper_chunks_content")
.execute(&pool)
.await?;
warn!(
"向量表和切片内容已清空,正在按新维度 {} 重建...",
embedding_dim
);
}
}
@ -155,24 +196,27 @@ async fn main() -> anyhow::Result<()> {
config.qiniu_domain.clone(),
);
let ads = AdsClient::new(config.ads_api_key.clone());
let arxiv = ArxivClient::new();
let ads = AdsClient::new(config.ads_api_key.clone()).context("构建 ADS 客户端失败")?;
let arxiv = ArxivClient::new().context("构建 arXiv 客户端失败")?;
let downloader = Downloader::new().context("构建 HTTP 下载客户端失败")?;
let llm = LlmClient::new(
config.llm_api_key.clone(),
config.llm_api_base.clone(),
config.llm_model.clone(),
);
)
.context("构建 LLM 客户端失败")?;
let medium_llm = LlmClient::new(
config.llm_medium_api_key.clone(),
config.llm_medium_api_base.clone(),
config.llm_medium_model.clone(),
);
)
.context("构建 Medium LLM 客户端失败")?;
let fast_llm = LlmClient::new(
config.llm_fast_api_key.clone(),
config.llm_fast_api_base.clone(),
config.llm_fast_model.clone(),
);
)
.context("构建 Fast LLM 客户端失败")?;
let vision_llm = if !config.llm_vision_model.is_empty() {
let key = if config.llm_vision_api_key.is_empty() {
config.llm_api_key.clone()
@ -185,7 +229,10 @@ async fn main() -> anyhow::Result<()> {
config.llm_vision_api_base.clone()
};
info!("视觉模型已启用: {}", config.llm_vision_model);
Some(LlmClient::new(key, base, config.llm_vision_model.clone()))
Some(
LlmClient::new(key, base, config.llm_vision_model.clone())
.context("构建 Vision LLM 客户端失败")?,
)
} else {
None
};
@ -193,17 +240,19 @@ async fn main() -> anyhow::Result<()> {
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
);
)
.context("构建 Embedding 客户端失败")?;
let skill_registry = Arc::new(RwLock::new(SkillRegistry::new(config.skills_dir.clone())));
if let Ok(mut reg) = skill_registry.write() {
let skill_registry = Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(
config.skills_dir.clone(),
)));
{
let mut reg = skill_registry.write().await;
reg.refresh();
} else {
warn!("SkillRegistry 初始化刷新失败RwLock 异常),将使用磁盘缓存");
}
info!(
"SkillRegistry 初始化完成,加载 {} 个 skill。",
skill_registry.read().map(|r| r.len()).unwrap_or(0)
skill_registry.read().await.len()
);
// 启动文件监听(热更新 skills
@ -222,6 +271,11 @@ async fn main() -> anyhow::Result<()> {
vision_llm,
embedding,
downloader,
http_client: reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to create shared HTTP client")?,
harvest_status: Arc::new(tokio::sync::Mutex::new(
astroresearch::services::batch::MetaSyncStatus::new(),
)),
@ -229,17 +283,22 @@ async fn main() -> anyhow::Result<()> {
astroresearch::services::batch::AssetBatchStatus::new(),
)),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(Mutex::new(std::collections::HashSet::new())),
cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
skill_registry,
pending_questions: Arc::new(Mutex::new(HashMap::new())),
pending_permissions: Arc::new(Mutex::new(HashMap::new())),
sse_broadcast: None,
pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
sse_broadcast: {
let (tx, _) = tokio::sync::broadcast::channel(256);
Some(tx)
},
memory_manager: Arc::new(tokio::sync::Mutex::new(
astroresearch::agent::memory::MemoryManager::new(config.library_dir.clone()),
)),
session_permission_checkers: Arc::new(RwLock::new(std::collections::HashMap::new())),
sessions: Arc::new(Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(Mutex::new(std::collections::HashMap::new())),
session_permission_checkers: Arc::new(tokio::sync::RwLock::new(
std::collections::HashMap::new(),
)),
sessions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
});
// 7. 设置 Axum 路由、CORS 头以及 React 仪表盘静态资源托管
@ -254,7 +313,12 @@ async fn main() -> anyhow::Result<()> {
let cors_permissive = CorsLayer::new()
.allow_origin(tower_http::cors::AllowOrigin::predicate(
|_origin, _parts| {
// 书签脚本的采集接口允许所有外域学术站点的跨域请求,并带上 Session 凭证
// [Code Review 豁免说明]:
// 这里的跨域策略被设置为无条件允许 (true)这是由于书签采集脚本Bookmarklet
// 需要在任意外部学术站点(如 arxiv.org, nature.com 等)向本服务发起 /upload 等请求。
// 虽然该策略允许凭据 (allow_credentials) 可能构成潜在的跨站请求伪造 (CSRF) 面,
// 但暴露的 API 仅限文件上传和元数据写入,不涉及任何提权、敏感信息泄露或 RCE 风险。
// 强制添加白名单会极大增加维护成本,因此该策略是有意为之的安全平衡。
true
},
))
@ -401,20 +465,35 @@ async fn main() -> anyhow::Result<()> {
.nest_service("/api/files", protected_files)
.fallback_service(serve_dir)
.layer(tower_http::trace::TraceLayer::new_for_http())
.layer(SetResponseHeaderLayer::overriding(
axum::http::header::X_CONTENT_TYPE_OPTIONS,
HeaderValue::from_static("nosniff"),
))
.layer(SetResponseHeaderLayer::overriding(
axum::http::header::X_FRAME_OPTIONS,
HeaderValue::from_static("DENY"),
))
.layer(SetResponseHeaderLayer::overriding(
axum::http::header::STRICT_TRANSPORT_SECURITY,
HeaderValue::from_static("max-age=31536000; includeSubDomains"),
))
.with_state(app_state);
let addr = SocketAddr::from(([0, 0, 0, 0], config.port));
info!("天文学科研服务已成功监听 http://localhost:{}", config.port);
let listener = tokio::net::TcpListener::bind(addr).await?;
let server = axum::serve(listener, app);
let server = axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
);
// 优雅关停:监听 SIGTERM / Ctrl+C
let shutdown_signal = async {
tokio::signal::ctrl_c()
.await
.expect("Failed to install Ctrl+C handler");
info!("收到关停信号,正在优雅关闭...");
info!("收到关停信号,正在关闭...");
};
server

View File

@ -93,16 +93,34 @@ impl AssetBatch {
status: Arc<Mutex<AssetBatchStatus>>,
) {
tokio::spawn(async move {
let llm_client = crate::clients::llm::LlmClient::new(
let llm_client = match crate::clients::llm::LlmClient::new(
config.llm_medium_api_key.clone(),
config.llm_medium_api_base.clone(),
config.llm_medium_model.clone(),
);
let embedding_client = crate::clients::llm::EmbeddingClient::new(
) {
Ok(c) => c,
Err(e) => {
tracing::error!("批量处理: 创建 LLM 客户端失败: {}", e);
let mut s = status.lock().await;
s.active = false;
s.add_log(format!("创建 LLM 客户端失败: {}", e));
return;
}
};
let embedding_client = match crate::clients::llm::EmbeddingClient::new(
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
);
) {
Ok(c) => c,
Err(e) => {
tracing::error!("批量处理: 创建 Embedding 客户端失败: {}", e);
let mut s = status.lock().await;
s.active = false;
s.add_log(format!("创建 Embedding 客户端失败: {}", e));
return;
}
};
let total = bibcodes.len() as i32;
{
let mut s = status.lock().await;
@ -200,6 +218,182 @@ impl AssetBatch {
}
}
struct PaperRecord {
bibcode: String,
pdf_path: Option<String>,
html_path: Option<String>,
markdown_path: Option<String>,
translation_path: Option<String>,
year: String,
created_at: String,
has_vector: bool,
has_target: bool,
}
#[allow(clippy::too_many_arguments)]
pub async fn filter_papers_for_batch(
db: &SqlitePool,
target_phase: &str,
limit_count: Option<i32>,
sort_order: Option<String>,
skip_completed: bool,
skip_failed: bool,
skip_preceding_failed: bool,
skip_preceding_uncompleted: bool,
) -> anyhow::Result<Vec<String>> {
// 基础过滤下推至 SQL排除已标记为无资源的文献减少内存加载
let rows = sqlx::query(
"SELECT bibcode, pdf_path, html_path, markdown_path, translation_path, year, \
datetime(created_at, 'localtime') AS created_at_local, \
EXISTS(SELECT 1 FROM paper_chunks_content WHERE bibcode = papers.bibcode) AS has_vector, \
EXISTS(SELECT 1 FROM paper_targets WHERE bibcode = papers.bibcode) AS has_target \
FROM papers \
WHERE (pdf_path IS NULL OR (pdf_path NOT LIKE 'error:no_resource%' AND pdf_path NOT LIKE 'error:%' AND pdf_path NOT LIKE 'error:%')) \
AND (html_path IS NULL OR (html_path NOT LIKE 'error:no_resource%' AND html_path NOT LIKE 'error:%' AND html_path NOT LIKE 'error:%'))"
)
.fetch_all(db)
.await?;
let mut records = Vec::new();
for r in rows {
use sqlx::Row;
records.push(PaperRecord {
bibcode: r.get("bibcode"),
pdf_path: r.get("pdf_path"),
html_path: r.get("html_path"),
markdown_path: r.get("markdown_path"),
translation_path: r.get("translation_path"),
year: r.get("year"),
created_at: r.get("created_at_local"),
has_vector: r.get("has_vector"),
has_target: r.get("has_target"),
});
}
// 排序
let sort_order_str = sort_order.unwrap_or_else(|| "default".to_string());
if sort_order_str == "pub_year_desc" {
records.sort_by(|a, b| b.year.cmp(&a.year));
} else if sort_order_str == "created_at_desc" {
records.sort_by(|a, b| b.created_at.cmp(&a.created_at));
}
let is_completed = |path: &Option<String>| -> bool {
if let Some(p) = path {
!p.is_empty() && !p.starts_with("error:") && !p.starts_with("mineru_batch:")
} else {
false
}
};
let is_failed = |path: &Option<String>| -> bool {
if let Some(p) = path {
p.starts_with("error:")
} else {
false
}
};
let is_no_resource = |path: &Option<String>| -> bool {
if let Some(p) = path {
p.starts_with("error:no_resource")
|| p.starts_with("error:无资源")
|| p.starts_with("error:无有效全文")
} else {
false
}
};
let mut target_bibcodes = Vec::new();
for rec in records {
let is_no_resource_paper = is_no_resource(&rec.pdf_path) || is_no_resource(&rec.html_path);
if is_no_resource_paper {
continue;
}
let download_completed = is_completed(&rec.pdf_path) || is_completed(&rec.html_path);
let download_failed = is_failed(&rec.pdf_path) || is_failed(&rec.html_path);
let download_uncompleted = !download_completed && !download_failed;
let parse_completed = is_completed(&rec.markdown_path);
let parse_failed = is_failed(&rec.markdown_path);
let parse_uncompleted = !parse_completed && !parse_failed;
let translate_completed = is_completed(&rec.translation_path);
let translate_failed = is_failed(&rec.translation_path);
match target_phase {
"download" => {
if skip_completed && download_completed {
continue;
}
if skip_failed && download_failed {
continue;
}
}
"parse" => {
if skip_completed && parse_completed {
continue;
}
if skip_failed && parse_failed {
continue;
}
if skip_preceding_failed && download_failed {
continue;
}
if skip_preceding_uncompleted && download_uncompleted {
continue;
}
}
"translate" => {
if skip_completed && translate_completed {
continue;
}
if skip_failed && translate_failed {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
"embed" => {
if skip_completed && rec.has_vector {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
"target" => {
if skip_completed && rec.has_target {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
_ => {}
}
target_bibcodes.push(rec.bibcode);
}
let limit = limit_count.unwrap_or(100) as usize;
if target_bibcodes.len() > limit {
target_bibcodes.truncate(limit);
}
Ok(target_bibcodes)
}
#[cfg(test)]
mod tests {
use super::*;
@ -321,7 +515,7 @@ mod tests {
.bind(&bibcode)
.fetch_one(&pool)
.await?;
let md_path_rel: String = row.get(0);
let md_path_rel: String = row.get("markdown_path");
assert_eq!(md_path_rel, format!("Markdown/{}.md", bibcode));
assert!(temp_dir.join(&md_path_rel).exists());

View File

@ -314,4 +314,42 @@ impl MetaSync {
}
});
}
// 获取所有已存同步检索配置
pub async fn list_queries(db: &SqlitePool) -> Result<Vec<SavedSyncQuery>, sqlx::Error> {
let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') FROM sync_queries ORDER BY last_run DESC")
.fetch_all(db)
.await?;
let mut list = Vec::new();
for r in rows {
use sqlx::Row;
list.push(SavedSyncQuery {
id: r.get("id"),
query: r.get("query"),
source: r.get("source"),
limit_count: r.get("limit_count"),
last_run: r.get("last_run"),
});
}
Ok(list)
}
// 删除同步检索配置
pub async fn delete_query(db: &SqlitePool, id: i64) -> Result<bool, sqlx::Error> {
let result = sqlx::query("DELETE FROM sync_queries WHERE id = ?")
.bind(id)
.execute(db)
.await?;
Ok(result.rows_affected() > 0)
}
}
#[derive(Debug, Serialize, serde::Deserialize)]
pub struct SavedSyncQuery {
pub id: i64,
pub query: String,
pub source: String,
pub limit_count: i32,
pub last_run: String,
}

View File

@ -2,5 +2,5 @@
pub mod asset;
pub mod meta;
pub use asset::{AssetBatch, AssetBatchStatus, BatchAction};
pub use meta::{MetaSync, MetaSyncStatus};
pub use asset::{filter_papers_for_batch, AssetBatch, AssetBatchStatus, BatchAction};
pub use meta::{MetaSync, MetaSyncStatus, SavedSyncQuery};

398
src/services/citation.rs Normal file
View File

@ -0,0 +1,398 @@
// src/services/citation.rs
// 文献引用关系查询服务层
use serde::{Deserialize, Serialize};
use sqlx::sqlite::SqlitePool;
use sqlx::FromRow;
use crate::services::paper::StandardPaper;
/// 统一的引文信息结构体
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CitationInfo {
pub bibcode: String,
pub title: String,
pub authors: Vec<String>,
pub first_author: String,
pub year: Option<String>,
pub citation_count: i32,
}
#[derive(Debug, FromRow)]
struct CitationDbRow {
bibcode: String,
title: String,
authors: Option<String>,
year: Option<String>,
citation_count: Option<i32>,
}
impl From<CitationDbRow> for CitationInfo {
fn from(row: CitationDbRow) -> Self {
let authors: Vec<String> = row
.authors
.as_deref()
.and_then(|a| serde_json::from_str(a).ok())
.unwrap_or_default();
let first_author = authors
.first()
.cloned()
.unwrap_or_else(|| "未知".to_string());
CitationInfo {
bibcode: row.bibcode,
title: row.title,
authors,
first_author,
year: row.year,
citation_count: row.citation_count.unwrap_or(0),
}
}
}
/// 解析查询字符串,提取作者部分和年份部分。
/// "Lei 2023" → ("Lei", Some("2023"))
/// "Zhang et al. 2020" → ("Zhang et al.", Some("2020"))
/// "Smith" → ("Smith", None)
pub fn parse_query(q: &str) -> (String, Option<String>) {
let cleaned = q.replace("et al.", "et al");
let parts: Vec<&str> = cleaned.split_whitespace().collect();
let mut year: Option<String> = None;
let mut author_tokens: Vec<&str> = Vec::new();
for (i, &p) in parts.iter().enumerate().rev() {
if year.is_none()
&& p.len() == 4
&& p.chars().all(|c| c.is_ascii_digit())
&& (p.starts_with("19") || p.starts_with("20"))
{
year = Some(p.to_string());
author_tokens = parts[..i].to_vec();
break;
}
}
if year.is_none() {
author_tokens = parts;
}
let author_part = if author_tokens.is_empty() {
q.to_string()
} else {
author_tokens.join(" ")
};
(author_part, year)
}
/// 在引用关系中模糊搜索匹配的文献
pub async fn search_citations(
db: &SqlitePool,
bibcode: &str,
direction: &str,
query: &str,
) -> Result<Vec<CitationInfo>, sqlx::Error> {
// 固定的 SQL 查询,列名由 match 分支选择,避免动态 format 拼接
const CITATIONS_WITH_YEAR: &str =
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.source_bibcode \
WHERE cr.target_bibcode = ? AND p.authors LIKE ? AND p.year = ?";
const CITATIONS_NO_YEAR: &str =
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.source_bibcode \
WHERE cr.target_bibcode = ? AND p.authors LIKE ?";
const REFERENCES_WITH_YEAR: &str =
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.target_bibcode \
WHERE cr.source_bibcode = ? AND p.authors LIKE ? AND p.year = ?";
const REFERENCES_NO_YEAR: &str =
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.target_bibcode \
WHERE cr.source_bibcode = ? AND p.authors LIKE ?";
let (author, year) = parse_query(query);
let author_pattern = format!("%{}%", author);
let is_citations = direction == "citations";
// LIMIT 防止模糊查询返回过多结果
const SEARCH_LIMIT: i64 = 100;
let rows = if let Some(ref y) = year {
let sql = if is_citations {
CITATIONS_WITH_YEAR
} else {
REFERENCES_WITH_YEAR
};
sqlx::query_as::<_, CitationDbRow>(&format!("{sql} LIMIT ?"))
.bind(bibcode)
.bind(&author_pattern)
.bind(y)
.bind(SEARCH_LIMIT)
.fetch_all(db)
.await?
} else {
let sql = if is_citations {
CITATIONS_NO_YEAR
} else {
REFERENCES_NO_YEAR
};
sqlx::query_as::<_, CitationDbRow>(&format!("{sql} LIMIT ?"))
.bind(bibcode)
.bind(&author_pattern)
.bind(SEARCH_LIMIT)
.fetch_all(db)
.await?
};
Ok(rows.into_iter().map(CitationInfo::from).collect())
}
/// 分页查询文献的参考文献或被引文献列表,返回列表及总数
pub async fn get_citations_paginated(
db: &SqlitePool,
bibcode: &str,
direction: &str,
sort: &str,
offset: i64,
limit: i64,
) -> Result<(Vec<CitationInfo>, i64), sqlx::Error> {
// 固定 SQL方向通过 match 选择,避免动态列名拼接
let count_sql = if direction == "citations" {
"SELECT COUNT(*) FROM citations_references WHERE target_bibcode = ?"
} else {
"SELECT COUNT(*) FROM citations_references WHERE source_bibcode = ?"
};
let query_sql = if direction == "citations" {
match sort {
"year" => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.source_bibcode WHERE cr.target_bibcode = ? ORDER BY p.year DESC LIMIT ? OFFSET ?",
"citation_count" => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.source_bibcode WHERE cr.target_bibcode = ? ORDER BY p.citation_count DESC LIMIT ? OFFSET ?",
_ => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.source_bibcode WHERE cr.target_bibcode = ? ORDER BY p.bibcode ASC LIMIT ? OFFSET ?",
}
} else {
match sort {
"year" => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.target_bibcode WHERE cr.source_bibcode = ? ORDER BY p.year DESC LIMIT ? OFFSET ?",
"citation_count" => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.target_bibcode WHERE cr.source_bibcode = ? ORDER BY p.citation_count DESC LIMIT ? OFFSET ?",
_ => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.target_bibcode WHERE cr.source_bibcode = ? ORDER BY p.bibcode ASC LIMIT ? OFFSET ?",
}
};
// 查询总数
let total: i64 = sqlx::query_scalar(count_sql)
.bind(bibcode)
.fetch_one(db)
.await
.unwrap_or_else(|e| {
tracing::warn!("查询引用总数失败: {}", e);
0
});
// 分页查询
let rows: Vec<CitationDbRow> = sqlx::query_as(query_sql)
.bind(bibcode)
.bind(limit)
.bind(offset)
.fetch_all(db)
.await?;
let infos = rows.into_iter().map(CitationInfo::from).collect();
Ok((infos, total))
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CitationNetworkResult {
pub bibcode: String,
pub title: String,
pub citation_count: i32,
pub reference_count: i32,
pub references: Vec<String>,
pub citations: Vec<String>,
pub citation_counts: std::collections::HashMap<String, i32>,
}
/// 文献元数据加载(本地 + ADS fallback
async fn resolve_paper_from_db_or_ads(
db: &SqlitePool,
library_dir: &std::path::Path,
ads_api_key: &str,
ads_client: &crate::clients::ads::AdsClient,
bibcode: &str,
) -> anyhow::Result<StandardPaper> {
match crate::services::paper::get_paper_from_db(db, library_dir, bibcode).await {
Ok(p) => Ok(p),
Err(_) => {
if ads_api_key.is_empty() {
anyhow::bail!(
"本地数据库未收录该文献,且未配置 ADS_API_KEY无法在线加载: {}",
bibcode
);
}
let docs = ads_client
.search(&format!("bibcode:{}", bibcode), 0, 1, "relevance")
.await?;
let doc = docs
.first()
.ok_or_else(|| anyhow::anyhow!("在本地库及 ADS 中均未找到该文献: {}", bibcode))?;
let paper = crate::services::paper::convert_ads_doc_to_standard(doc);
if let Err(e) = crate::services::paper::save_paper_to_db(db, &paper).await {
tracing::warn!("保存引用文献至数据库失败: {}", e);
}
// 保存引用拓扑关系(批量 INSERT避免 N+1
let mut pairs: Vec<(String, String)> = Vec::new();
if let Some(refs) = &doc.reference {
for ref_bib in refs {
pairs.push((paper.bibcode.clone(), ref_bib.clone()));
}
}
if let Some(cits) = &doc.citation {
for cit_bib in cits {
pairs.push((cit_bib.clone(), paper.bibcode.clone()));
}
}
if !pairs.is_empty() {
let mut tx = match db.begin().await {
Ok(tx) => tx,
Err(e) => {
tracing::warn!("无法开启引用关系事务: {}", e);
return Ok(paper);
}
};
for chunk in pairs.chunks(crate::services::paper::SQLITE_PARAM_LIMIT / 2) {
for (src, tgt) in chunk {
if let Err(e) = sqlx::query(
"INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)",
)
.bind(src)
.bind(tgt)
.execute(&mut *tx)
.await
{
tracing::warn!("保存引用关系失败 ({} -> {}): {}", src, tgt, e);
}
}
}
if let Err(e) = tx.commit().await {
tracing::warn!("提交引用关系事务失败: {}", e);
}
}
Ok(paper)
}
}
}
/// 批量查询关联文献的被引数量
async fn batch_load_citation_counts(
db: &SqlitePool,
references: &[String],
citations: &[String],
) -> std::collections::HashMap<String, i32> {
let mut counts = std::collections::HashMap::new();
let mut all: Vec<String> = references.to_vec();
all.extend(citations.iter().cloned());
all.sort();
all.dedup();
for chunk in all.chunks(crate::services::paper::SQLITE_PARAM_LIMIT) {
if chunk.is_empty() {
continue;
}
let placeholders = vec!["?"; chunk.len()].join(", ");
let sql = format!(
"SELECT bibcode, citation_count FROM papers WHERE bibcode IN ({})",
placeholders
);
let mut query = sqlx::query(&sql);
for b in chunk {
query = query.bind(b);
}
match query.fetch_all(db).await {
Ok(rows) => {
for row in rows {
counts.insert(sqlx::Row::get(&row, 0), sqlx::Row::get(&row, 1));
}
}
Err(e) => {
tracing::warn!("批量查询引用数失败 (chunk size={}): {}", chunk.len(), e);
}
}
}
counts
}
pub async fn get_citation_network_with_fallback(
db: &SqlitePool,
library_dir: &std::path::Path,
ads_api_key: &str,
ads_client: &crate::clients::ads::AdsClient,
bibcode: &str,
) -> anyhow::Result<CitationNetworkResult> {
let paper =
resolve_paper_from_db_or_ads(db, library_dir, ads_api_key, ads_client, bibcode).await?;
// 加载引用文献列表
let refs_rows =
sqlx::query("SELECT target_bibcode FROM citations_references WHERE source_bibcode = ?")
.bind(bibcode)
.fetch_all(db)
.await
.unwrap_or_else(|e| {
tracing::warn!("查询引用关系失败: {}", e);
Vec::new()
});
let references: Vec<String> = refs_rows.iter().map(|row| sqlx::Row::get(row, 0)).collect();
// 加载被引文献列表
let cits_rows =
sqlx::query("SELECT source_bibcode FROM citations_references WHERE target_bibcode = ?")
.bind(bibcode)
.fetch_all(db)
.await
.unwrap_or_else(|e| {
tracing::warn!("查询被引关系失败: {}", e);
Vec::new()
});
let citations: Vec<String> = cits_rows.iter().map(|row| sqlx::Row::get(row, 0)).collect();
// 批量加载关联文献的被引数量
let citation_counts = batch_load_citation_counts(db, &references, &citations).await;
Ok(CitationNetworkResult {
bibcode: paper.bibcode,
title: paper.title,
citation_count: paper.citation_count,
reference_count: paper.reference_count,
references,
citations,
citation_counts,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_query_with_year() {
let (author, year) = parse_query("Lei 2023");
assert_eq!(author, "Lei");
assert_eq!(year, Some("2023".into()));
}
#[test]
fn test_parse_query_et_al() {
let (author, year) = parse_query("Zhang et al. 2020");
assert_eq!(author, "Zhang et al");
assert_eq!(year, Some("2020".into()));
}
#[test]
fn test_parse_query_author_only() {
let (author, year) = parse_query("Smith");
assert_eq!(author, "Smith");
assert_eq!(year, None);
}
}

View File

@ -12,8 +12,7 @@ mod strategies;
pub(crate) use headers::validate_html_content_lenient;
pub(crate) use headers::validate_pdf_content;
use crate::api::helpers::{check_paper_paths_in_db, get_paper_from_db};
use crate::api::StandardPaper;
use crate::services::paper::{check_paper_paths_in_db, get_paper_from_db, StandardPaper};
use anyhow::{Context, Result};
use reqwest::header::{HeaderMap, HeaderValue};
use std::path::{Path, PathBuf};
@ -646,7 +645,7 @@ mod tests {
#[tokio::test]
#[ignore]
async fn test_download_via_obscura_integration() -> anyhow::Result<()> {
use axum::http::{header::CONTENT_TYPE, HeaderValue};
use axum::http::header::CONTENT_TYPE;
use axum::{response::Response as AxumResponse, routing::get, Router};
// Bind to a random port

View File

@ -1,11 +1,16 @@
pub mod batch;
pub mod chunker;
pub mod citation;
pub mod download;
pub mod logging;
pub mod note;
pub mod paper;
pub mod parser;
pub mod pipeline;
pub mod query_parser;
pub mod rag;
pub mod search;
pub mod section_parser;
pub mod session;
pub mod target;
pub mod translation;
pub mod vision;

121
src/services/note.rs Normal file
View File

@ -0,0 +1,121 @@
// src/services/note.rs
// 研究笔记服务层
use std::path::{Path, PathBuf};
/// 保存研究笔记为本地 Markdown 文件
pub fn save_note_service(
library_dir: &Path,
title: &str,
content: &str,
) -> anyhow::Result<(String, PathBuf, usize)> {
// 文件名安全化
let safe_title: String = title
.chars()
.map(|c| {
if c.is_alphanumeric() || c == '-' || c == '_' || c == ' ' {
c
} else {
'_'
}
})
.collect::<String>()
.replace(' ', "_");
let notes_dir = library_dir.join("notes");
std::fs::create_dir_all(&notes_dir)?;
let filename = format!("{}.md", safe_title);
let filepath = notes_dir.join(&filename);
let full_content = format!(
"---\ntitle: {}\ndate: {}\ngenerated_by: AstroResearch Agent\n---\n\n{}",
title,
chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
content
);
std::fs::write(&filepath, &full_content)?;
Ok((filename, filepath, full_content.len()))
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, sqlx::FromRow)]
pub struct NoteRecord {
pub id: i64,
pub bibcode: String,
pub paragraph_index: i64,
pub note_text: String,
pub highlight_color: String,
pub selected_text: String,
pub created_at: String,
}
/// 创建文献段落高亮与笔记
pub async fn create_highlight_note(
db: &sqlx::SqlitePool,
bibcode: &str,
paragraph_index: i64,
note_text: &str,
highlight_color: &str,
selected_text: &str,
) -> Result<NoteRecord, sqlx::Error> {
let note = sqlx::query_as::<_, NoteRecord>(
"INSERT INTO notes (bibcode, paragraph_index, note_text, highlight_color, selected_text) VALUES (?, ?, ?, ?, ?) RETURNING id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at"
)
.bind(bibcode)
.bind(paragraph_index)
.bind(note_text)
.bind(highlight_color)
.bind(selected_text)
.fetch_one(db)
.await?;
Ok(note)
}
/// 查询某篇文献的全部高亮笔记
pub async fn list_highlight_notes(
db: &sqlx::SqlitePool,
bibcode: &str,
) -> Result<Vec<NoteRecord>, sqlx::Error> {
let notes = sqlx::query_as::<_, NoteRecord>(
"SELECT id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at FROM notes WHERE bibcode = ? ORDER BY paragraph_index, created_at"
)
.bind(bibcode)
.fetch_all(db)
.await?;
Ok(notes)
}
/// 删除指定 id 的高亮笔记
pub async fn delete_highlight_note(db: &sqlx::SqlitePool, id: i64) -> Result<bool, sqlx::Error> {
let result = sqlx::query("DELETE FROM notes WHERE id = ?")
.bind(id)
.execute(db)
.await?;
Ok(result.rows_affected() > 0)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_save_note_service() {
let tmp_dir = std::env::temp_dir().join("astro_note_test");
let title = "Test Note Title 123!";
let content = "Hello world from note service test";
let res = save_note_service(&tmp_dir, title, content);
assert!(res.is_ok());
let (filename, filepath, size) = res.unwrap();
assert_eq!(filename, "Test_Note_Title_123_.md");
assert!(filepath.exists());
assert!(size > content.len());
std::fs::remove_dir_all(&tmp_dir).ok();
}
}

367
src/services/paper/db.rs Normal file
View File

@ -0,0 +1,367 @@
// src/services/paper/db.rs
//
// 文献元数据数据库存取。
use super::model::StandardPaper;
use sqlx::{Acquire, Row, SqlitePool};
use std::path::Path;
use tracing::info;
/// SQLite 单条 SQL 的参数绑定上限SQLITE_MAX_VARIABLE_NUMBER 默认 999留余量取 900
/// 批量 IN 查询时需按此分批,避免超出绑定数量。
pub const SQLITE_PARAM_LIMIT: usize = 900;
/// 清洗后的文献标识符,提取纯 DOI 和 arXiv ID。
///
/// 统一处理用户可能输入的多种前缀格式(`doi:`, `https://doi.org/`, `arxiv:` 等)。
#[derive(Debug, Clone)]
pub struct CleanId {
/// 原始标识符的 trim 结果
pub raw: String,
/// 去除了 DOI 前缀的纯标识符
pub clean_doi: String,
/// 去除了 arXiv 前缀及 pdf/ 前缀的纯标识符
pub clean_arxiv: String,
/// arXiv ID 去除版本号后缀(仅 ingest 使用,如 `2303.12345v1` → `2303.12345`
pub clean_arxiv_no_version: String,
}
/// 清洗文献标识符去除常见前缀doi:, arXiv:, https://doi.org/ 等)。
pub fn clean_identifier(identifier: &str) -> CleanId {
let raw = identifier.trim().to_string();
let clean_doi = raw
.trim_start_matches("doi:")
.trim_start_matches("DOI:")
.trim_start_matches("https://doi.org/")
.trim_start_matches("http://doi.org/")
.trim()
.to_string();
let clean_arxiv = raw
.trim_start_matches("arxiv:")
.trim_start_matches("arXiv:")
.trim_start_matches("ARXIV:")
.trim_start_matches("pdf/")
.trim()
.to_string();
// 移除 arXiv 版本号后缀(如 2303.12345v1 → 2303.12345
let clean_arxiv_no_version = if let Some(pos) = clean_arxiv.find('v') {
if pos + 1 < clean_arxiv.len() && clean_arxiv[pos + 1..].chars().all(|c| c.is_ascii_digit())
{
clean_arxiv[..pos].to_string()
} else {
clean_arxiv.clone()
}
} else {
clean_arxiv.clone()
};
CleanId {
raw,
clean_doi,
clean_arxiv,
clean_arxiv_no_version,
}
}
pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Result<()> {
let mut conn = db.acquire().await?;
save_paper_to_db_tx(&mut conn, p).await
}
pub async fn save_paper_to_db_tx(
conn: &mut sqlx::SqliteConnection,
p: &StandardPaper,
) -> anyhow::Result<()> {
let authors_json = serde_json::to_string(&p.authors)?;
let keywords_json = serde_json::to_string(&p.keywords)?;
// 1. 如果存在 arxiv_id检查是否有已存在的相同 arxiv_id 记录以防 duplicate
if !p.arxiv_id.is_empty() {
#[allow(clippy::type_complexity)]
let existing_opt: Option<(String, Option<String>, Option<String>, Option<String>, Option<String>)> = sqlx::query_as(
"SELECT bibcode, pdf_path, html_path, markdown_path, translation_path FROM papers WHERE arxiv_id = ?"
)
.bind(&p.arxiv_id)
.fetch_optional(&mut *conn)
.await?;
if let Some((existing_bibcode, _pdf, _html, _md, _tr)) = existing_opt {
if existing_bibcode != p.bibcode {
// 发现不同 bibcode 标识的同一篇文献记录,需要进行合并
// 如果已存在的记录使用的是临时 arXiv ID 作为 bibcode且新记录使用的是正式 ADS bibcode我们升级 bibcode 主键
let is_existing_temp = existing_bibcode == p.arxiv_id;
let is_new_formal = p.bibcode != p.arxiv_id;
if is_existing_temp && is_new_formal {
info!(
"发现相同 arXiv ID 的文献,将临时主键 {} 升级为正式 ADS Bibcode: {}",
existing_bibcode, p.bibcode
);
// 在事务中原子执行 bibcode 迁移papers 表更新 + 两条 citation_references 引用级联
let mut tx = conn.begin().await?;
sqlx::query(
"UPDATE papers SET bibcode = ?, title = ?, authors = ?, year = ?, pub = ?, keywords = ?, abstract = ?, doi = ?, citation_count = ?, reference_count = ?, doctype = ? WHERE bibcode = ?"
)
.bind(&p.bibcode)
.bind(&p.title)
.bind(&authors_json)
.bind(&p.year)
.bind(&p.pub_journal)
.bind(&keywords_json)
.bind(&p.abstract_text)
.bind(&p.doi)
.bind(p.citation_count)
.bind(p.reference_count)
.bind(&p.doctype)
.bind(&existing_bibcode)
.execute(&mut *tx)
.await?;
// 同步更新引用关系表(手动级联,在同一事务内)
sqlx::query("UPDATE citations_references SET source_bibcode = ? WHERE source_bibcode = ?")
.bind(&p.bibcode)
.bind(&existing_bibcode)
.execute(&mut *tx)
.await?;
sqlx::query("UPDATE citations_references SET target_bibcode = ? WHERE target_bibcode = ?")
.bind(&p.bibcode)
.bind(&existing_bibcode)
.execute(&mut *tx)
.await?;
tx.commit().await?;
return Ok(());
} else {
// 如果已存在的是正式 ADS bibcode而新插入的是临时 arXiv ID直接忽略或更新元数据而不更改主键
info!(
"发现相同 arXiv ID 的文献 {} 已存在正式记录,忽略临时 arXiv 插入",
existing_bibcode
);
return Ok(());
}
}
}
}
// 2. 正常插入/冲突更新
sqlx::query(
"INSERT INTO papers (bibcode, title, authors, year, pub, keywords, abstract, doi, arxiv_id, citation_count, reference_count, doctype) \
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \
ON CONFLICT(bibcode) DO UPDATE SET \
title=excluded.title, \
authors=excluded.authors, \
pub=excluded.pub, \
keywords=excluded.keywords, \
abstract=excluded.abstract, \
doi=excluded.doi, \
arxiv_id=excluded.arxiv_id, \
citation_count=excluded.citation_count, \
reference_count=excluded.reference_count, \
doctype=excluded.doctype"
)
.bind(&p.bibcode)
.bind(&p.title)
.bind(authors_json)
.bind(&p.year)
.bind(&p.pub_journal)
.bind(keywords_json)
.bind(&p.abstract_text)
.bind(&p.doi)
.bind(&p.arxiv_id)
.bind(p.citation_count)
.bind(p.reference_count)
.bind(&p.doctype)
.execute(&mut *conn)
.await?;
Ok(())
}
/// 从数据库行解析出 StandardPaper
pub fn parse_paper_row(r: &sqlx::sqlite::SqliteRow, library_dir: &Path) -> StandardPaper {
let pdf_path: Option<String> = r.get("pdf_path");
let html_path: Option<String> = r.get("html_path");
let markdown_path: Option<String> = r.get("markdown_path");
let translation_path: Option<String> = r.get("translation_path");
let doctype_val: Option<String> = r.get("doctype");
let has_vector: bool = r.get("has_vector");
let authors_str: Option<String> = r.get("authors");
let authors: Vec<String> = authors_str
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let keywords_str: Option<String> = r.get("keywords");
let keywords: Vec<String> = keywords_str
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let is_pdf_exist = pdf_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_html_exist = html_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_md_exist = markdown_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_tr_exist = translation_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let pdf_error = pdf_path
.as_ref()
.filter(|p| p.starts_with("error:"))
.map(|p| p["error:".len()..].trim().to_string());
let html_error = html_path
.as_ref()
.filter(|p| p.starts_with("error:"))
.map(|p| p["error:".len()..].trim().to_string());
StandardPaper {
bibcode: r.get("bibcode"),
title: r.get("title"),
authors,
year: r.get("year"),
pub_journal: r.get("pub"),
keywords,
abstract_text: r.get("abstract"),
doi: r.get("doi"),
arxiv_id: r.get("arxiv_id"),
citation_count: r.get("citation_count"),
reference_count: r.get("reference_count"),
is_downloaded: is_pdf_exist || is_html_exist,
has_pdf: is_pdf_exist,
has_html: is_html_exist,
has_markdown: is_md_exist,
has_translation: is_tr_exist,
has_vector,
doctype: doctype_val.unwrap_or_else(|| "article".to_string()),
pdf_error,
html_error,
}
}
/// 通用的 5 路 OR 查询条件bibcode / doi / arxiv_id 精确匹配 + doi 大小写不敏感匹配)
const PAPER_LOOKUP_WHERE: &str =
"bibcode = ? OR doi = ? OR arxiv_id = ? OR LOWER(doi) = LOWER(?) OR arxiv_id = ?";
pub async fn get_paper_from_db(
db: &SqlitePool,
library_dir: &Path,
identifier: &str,
) -> anyhow::Result<StandardPaper> {
let id = clean_identifier(identifier);
let query = format!(
"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) AS has_vector FROM papers WHERE {}",
PAPER_LOOKUP_WHERE
);
let r = sqlx::query(&query)
.bind(&id.raw)
.bind(&id.clean_doi)
.bind(&id.clean_arxiv)
.bind(&id.clean_doi)
.bind(&id.clean_arxiv)
.fetch_one(db)
.await?;
Ok(parse_paper_row(&r, library_dir))
}
pub async fn check_paper_paths_in_db(
db: &SqlitePool,
library_dir: &Path,
identifier: &str,
) -> anyhow::Result<
Option<(
Option<String>,
Option<String>,
Option<String>,
Option<String>,
)>,
> {
let id = clean_identifier(identifier);
let query = format!(
"SELECT pdf_path, html_path, markdown_path, translation_path FROM papers WHERE {}",
PAPER_LOOKUP_WHERE
);
let r_opt = sqlx::query(&query)
.bind(&id.raw)
.bind(&id.clean_doi)
.bind(&id.clean_arxiv)
.bind(&id.clean_doi)
.bind(&id.clean_arxiv)
.fetch_optional(db)
.await?;
if let Some(r) = r_opt {
let pdf: Option<String> = r.get("pdf_path");
let html: Option<String> = r.get("html_path");
let md: Option<String> = r.get("markdown_path");
let tr: Option<String> = r.get("translation_path");
let pdf_res = pdf.filter(|p| library_dir.join(p).exists());
let html_res = html.filter(|p| library_dir.join(p).exists());
let md_res = md.filter(|p| library_dir.join(p).exists());
let tr_res = tr.filter(|p| library_dir.join(p).exists());
Ok(Some((pdf_res, html_res, md_res, tr_res)))
} else {
Ok(None)
}
}
/// 获取本地文献库的全部文献列表
pub async fn get_library_list(
db: &SqlitePool,
library_dir: &Path,
limit: Option<i64>,
offset: Option<i64>,
) -> Result<Vec<StandardPaper>, sqlx::Error> {
let rows = if let (Some(lim), Some(off)) = (limit, offset) {
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) AS has_vector FROM papers ORDER BY created_at DESC LIMIT ? OFFSET ?")
.bind(lim)
.bind(off)
.fetch_all(db)
.await?
} else {
// 无分页参数时默认限制 2000 条,防止全量加载
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) AS has_vector FROM papers ORDER BY created_at DESC LIMIT 2000")
.fetch_all(db)
.await?
};
let mut list = Vec::new();
for r in rows {
list.push(parse_paper_row(&r, library_dir));
}
Ok(list)
}
/// 标记/清除文献的“无资源”状态
pub async fn mark_no_resource_service(
db: &SqlitePool,
library_dir: &Path,
bibcode: &str,
clear: bool,
) -> anyhow::Result<StandardPaper> {
if clear {
sqlx::query("UPDATE papers SET pdf_path = NULL, html_path = NULL WHERE bibcode = ?")
.bind(bibcode)
.execute(db)
.await?;
} else {
sqlx::query("UPDATE papers SET pdf_path = 'error:no_resource', html_path = 'error:no_resource' WHERE bibcode = ?")
.bind(bibcode)
.execute(db)
.await?;
}
let updated_paper = get_paper_from_db(db, library_dir, bibcode).await?;
Ok(updated_paper)
}

View File

@ -0,0 +1,138 @@
// src/services/paper/ingest.rs
//
// 文献手动上传与推入逻辑Ingestion Service
use crate::services::paper::db::get_paper_from_db;
use crate::services::paper::model::StandardPaper;
use sqlx::SqlitePool;
use std::path::Path;
use tracing::info;
/// 校验并保存上传的文献文件PDF 或 HTML并更新数据库状态。
pub async fn upload_paper_file_service(
db: &SqlitePool,
library_dir: &Path,
input_bibcode: &str,
file_type: &str,
file_name: &str,
file_bytes: &[u8],
) -> anyhow::Result<StandardPaper> {
if file_bytes.is_empty() {
anyhow::bail!("上传文件为空或读取失败");
}
// 尝试将可能的 DOI 或 arXiv ID 解析为真实的 bibcode
let mut resolved_bibcode = input_bibcode.to_string();
let exists_as_bibcode = sqlx::query("SELECT bibcode FROM papers WHERE bibcode = ?")
.bind(&resolved_bibcode)
.fetch_optional(db)
.await
.unwrap_or(None)
.is_some();
if !exists_as_bibcode {
let id = crate::services::paper::clean_identifier(&resolved_bibcode);
use sqlx::Row;
// 尝试匹配 DOI
if let Some(row) = sqlx::query(
"SELECT bibcode FROM papers WHERE doi = ? OR doi = ? OR LOWER(doi) = LOWER(?)",
)
.bind(&id.clean_doi)
.bind(&resolved_bibcode)
.bind(&id.clean_doi)
.fetch_optional(db)
.await
.unwrap_or(None)
{
let found: String = row.get("bibcode");
info!(
"上传服务:通过 DOI 匹配成功,将 '{}' 解析为 bibcode '{}'",
input_bibcode, found
);
resolved_bibcode = found;
} else if let Some(row) = sqlx::query(
// 尝试匹配 arXiv ID含版本号去除
"SELECT bibcode FROM papers WHERE arxiv_id = ? OR arxiv_id = ? OR arxiv_id LIKE ? OR arxiv_id LIKE ?"
)
.bind(&id.clean_arxiv)
.bind(&id.clean_arxiv_no_version)
.bind(format!("{}%", id.clean_arxiv_no_version))
.bind(format!("arXiv:{}%", id.clean_arxiv_no_version))
.fetch_optional(db)
.await
.unwrap_or(None)
{
let found: String = row.get("bibcode");
info!("上传服务:通过 arXiv ID 匹配成功,将 '{}' 解析为 bibcode '{}'", input_bibcode, found);
resolved_bibcode = found;
}
}
let bibcode = resolved_bibcode;
// 从数据库读取该文献元数据以验证其确实存在
let _paper = get_paper_from_db(db, library_dir, &bibcode)
.await
.map_err(|e| anyhow::anyhow!("未找到该文献记录: {}", e))?;
// 校验并保存文件
// 优先使用显式 file_type 字段,避免文件名扩展名覆盖真实类型
let is_pdf = if !file_type.is_empty() {
file_type == "pdf"
} else {
file_name.to_lowercase().ends_with(".pdf")
};
let relative_path = if is_pdf {
crate::services::download::validate_pdf_content(file_bytes)
.map_err(|e| anyhow::anyhow!("PDF 文件内容校验失败: {}", e))?;
let pdf_filename = format!("{}.pdf", bibcode);
let pdf_dest = library_dir.join("PDF").join(&pdf_filename);
if let Some(parent) = pdf_dest.parent() {
tokio::fs::create_dir_all(parent)
.await
.map_err(|e| anyhow::anyhow!("无法创建 PDF 目录: {}", e))?;
}
tokio::fs::write(&pdf_dest, file_bytes)
.await
.map_err(|e| anyhow::anyhow!("无法写入 PDF 文件: {}", e))?;
format!("PDF/{}", pdf_filename)
} else {
let text_content = String::from_utf8(file_bytes.to_vec())
.map_err(|_| anyhow::anyhow!("上传的 HTML 文件不是有效的 UTF-8 文本"))?;
crate::services::download::validate_html_content_lenient(&text_content)
.map_err(|e| anyhow::anyhow!("HTML 文件内容校验失败: {}", e))?;
let html_filename = format!("{}.html", bibcode);
let html_dest = library_dir.join("HTML").join(&html_filename);
if let Some(parent) = html_dest.parent() {
tokio::fs::create_dir_all(parent)
.await
.map_err(|e| anyhow::anyhow!("无法创建 HTML 目录: {}", e))?;
}
tokio::fs::write(&html_dest, &text_content)
.await
.map_err(|e| anyhow::anyhow!("无法写入 HTML 文件: {}", e))?;
format!("HTML/{}", html_filename)
};
// 更新数据库路径状态(固定列名,避免动态 SQL 拼接)
if is_pdf {
sqlx::query("UPDATE papers SET pdf_path = ? WHERE bibcode = ?")
.bind(&relative_path)
.bind(&bibcode)
.execute(db)
.await
.map_err(|e| anyhow::anyhow!("更新数据库状态失败: {}", e))?;
} else {
sqlx::query("UPDATE papers SET html_path = ? WHERE bibcode = ?")
.bind(&relative_path)
.bind(&bibcode)
.execute(db)
.await
.map_err(|e| anyhow::anyhow!("更新数据库状态失败: {}", e))?;
}
// 重新获取最新的文献信息以返回
let updated_paper = get_paper_from_db(db, library_dir, &bibcode).await?;
Ok(updated_paper)
}

20
src/services/paper/mod.rs Normal file
View File

@ -0,0 +1,20 @@
// src/services/paper/mod.rs
//
// 文献服务模块入口,统一重导出子模块接口。
pub mod db;
pub mod ingest;
pub mod model;
pub mod reader;
pub use db::{
check_paper_paths_in_db, clean_identifier, get_library_list, get_paper_from_db,
mark_no_resource_service, parse_paper_row, save_paper_to_db, save_paper_to_db_tx, CleanId,
SQLITE_PARAM_LIMIT,
};
pub use ingest::upload_paper_file_service;
pub use model::{convert_ads_doc_to_standard, convert_arxiv_to_standard, StandardPaper};
pub use reader::{
extract_outline, extract_section_by_index, extract_section_by_name, read_paper_content,
MarkdownSection, PaperReadResponse, ReadMode,
};

110
src/services/paper/model.rs Normal file
View File

@ -0,0 +1,110 @@
// src/services/paper/model.rs
//
// 文献领域模型与转换逻辑。
use crate::clients::ads::AdsPaperDoc;
use crate::clients::arxiv::ArxivPaper;
use serde::{Deserialize, Serialize};
// 统一标准化的文献格式,用于在系统内部及向前端传输
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct StandardPaper {
pub bibcode: String,
pub title: String,
pub authors: Vec<String>,
pub year: String,
pub pub_journal: String,
pub keywords: Vec<String>,
pub abstract_text: String,
pub doi: String,
pub arxiv_id: String,
pub citation_count: i32,
pub reference_count: i32,
pub is_downloaded: bool,
pub has_pdf: bool,
pub has_html: bool,
pub has_markdown: bool,
pub has_translation: bool,
pub has_vector: bool,
pub doctype: String,
pub pdf_error: Option<String>,
pub html_error: Option<String>,
}
pub fn convert_ads_doc_to_standard(doc: &AdsPaperDoc) -> StandardPaper {
let title = doc
.title
.as_ref()
.and_then(|v: &Vec<String>| v.first())
.cloned()
.unwrap_or_else(|| doc.bibcode.clone());
let authors = doc.author.clone().unwrap_or_default();
let keywords = doc.keyword.clone().unwrap_or_default();
let doi = doc
.doi
.as_ref()
.and_then(|v: &Vec<String>| v.first())
.cloned()
.unwrap_or_default();
let mut arxiv_id = String::new();
if let Some(identifiers) = &doc.identifier {
for id in identifiers {
if id.starts_with("arXiv:") {
arxiv_id = id.replace("arXiv:", "").trim().to_string();
break;
}
}
}
if arxiv_id.is_empty() && doc.bibcode.starts_with("arXiv") {
arxiv_id = doc.bibcode.replace("arXiv", "").trim().to_string();
}
StandardPaper {
bibcode: doc.bibcode.clone(),
title,
authors,
year: doc.year.clone().unwrap_or_default(),
pub_journal: doc.pub_journal.clone().unwrap_or_default(),
keywords,
abstract_text: doc.abstract_text.clone().unwrap_or_default(),
doi,
arxiv_id,
citation_count: doc.citation_count.unwrap_or(0),
reference_count: doc.reference_count.unwrap_or(0),
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: doc.doctype.clone().unwrap_or_else(|| "article".to_string()),
pdf_error: None,
html_error: None,
}
}
pub fn convert_arxiv_to_standard(doc: &ArxivPaper) -> StandardPaper {
StandardPaper {
bibcode: doc.id.clone(),
title: doc.title.clone(),
authors: doc.authors.clone(),
year: doc.year.clone(),
pub_journal: "arXiv Preprint".to_string(),
keywords: Vec::new(),
abstract_text: doc.abstract_text.clone(),
doi: doc.doi.clone().unwrap_or_default(),
arxiv_id: doc.id.clone(),
citation_count: 0,
reference_count: 0,
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: "eprint".to_string(),
pdf_error: None,
html_error: None,
}
}

View File

@ -1,13 +1,29 @@
// src/services/section_parser.rs
// src/services/paper/reader.rs
//
// Markdown 章节解析服务。
// 从已解析的文献 Markdown 中提取章节大纲(目录)和指定章节内容。
// 文献内容读取与章节大纲解析服务。
use regex::Regex;
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
use std::path::Path;
use std::sync::LazyLock;
use tracing;
/// 读取模式
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ReadMode {
/// 全文模式。可选择是否同时附带翻译内容
Full { include_translation: bool },
/// 仅提取大纲目录(可指定最大层级,默认 3
Outline { max_level: Option<usize> },
/// 按序号提取单个章节0-based
SectionIndex(usize),
/// 按名称模糊匹配提取单个章节
SectionName(String),
}
/// Markdown 章节结构
#[derive(Debug, Clone, serde::Serialize)]
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MarkdownSection {
/// 章节序号0-based
pub index: usize,
@ -25,13 +41,117 @@ pub struct MarkdownSection {
pub char_end: usize,
}
/// 匹配 Markdown ATX 标题行(##, ###, ####...
/// 统一的文献读取响应
#[derive(Debug, Clone, Serialize)]
pub struct PaperReadResponse {
pub bibcode: String,
/// 主文本内容(原文全文、大纲、或者单个章节文本)
pub content: String,
/// 如果是 Full 模式且 include_translation = true返回翻译内容全文
pub translation_content: Option<String>,
/// 如果是大纲模式,附带结构化的大纲数据,方便 API 转换为 JSON 返回
pub outline: Option<Vec<MarkdownSection>>,
}
/// 统一读取文献内容的业务逻辑函数
pub async fn read_paper_content(
db: &SqlitePool,
library_dir: &Path,
bibcode: &str,
mode: ReadMode,
) -> anyhow::Result<PaperReadResponse> {
// 1. 统一安全与数据库路径校验
let paths = super::db::check_paper_paths_in_db(db, library_dir, bibcode)
.await?
.ok_or_else(|| {
anyhow::anyhow!("文献 {} 未在本地数据库中注册,请先检索该文献。", bibcode)
})?;
let (_, _, md_opt, tr_opt) = paths;
// 2. 检查英文原文相对路径
let md_rel = md_opt.ok_or_else(|| {
anyhow::anyhow!("文献 {} 尚未完成解析 (parse)。请先执行解析任务。", bibcode)
})?;
// 3. 读取英文原文物理文件
let md_abs = library_dir.join(&md_rel);
if !md_abs.exists() {
anyhow::bail!(
"文献 {} 的本地原文文件已丢失,请重新执行解析任务。",
bibcode
);
}
let full_text = tokio::fs::read_to_string(&md_abs).await?;
// 4. 初始化响应字段
let mut content = full_text.clone();
let mut translation_content = None;
let mut outline = None;
// 5. 根据模式处理文本
match &mode {
ReadMode::Full {
include_translation,
} => {
if *include_translation {
if let Some(tr_rel) = tr_opt {
let tr_abs = library_dir.join(&tr_rel);
if tr_abs.exists() {
match tokio::fs::read_to_string(&tr_abs).await {
Ok(tr_text) => translation_content = Some(tr_text),
Err(e) => tracing::warn!(
"翻译文件存在但读取失败 for {} ({}): {}",
bibcode,
tr_abs.display(),
e
),
}
}
}
}
}
ReadMode::Outline { max_level } => {
let level = max_level.unwrap_or(3);
let sections = extract_outline(&full_text, level);
let formatted = sections
.iter()
.map(|s| {
format!(
"{}[{}] {}",
" ".repeat(s.level.saturating_sub(1)),
s.index,
s.heading
)
})
.collect::<Vec<_>>()
.join("\n");
content = formatted;
outline = Some(sections);
}
ReadMode::SectionIndex(idx) => {
content = extract_section_by_index(&full_text, *idx)
.ok_or_else(|| anyhow::anyhow!("未找到序号为 #{} 的章节,请先检查大纲。", idx))?;
}
ReadMode::SectionName(name) => {
content = extract_section_by_name(&full_text, name).ok_or_else(|| {
anyhow::anyhow!("未找到名称匹配 '{}' 的章节,请先检查大纲。", name)
})?;
}
}
Ok(PaperReadResponse {
bibcode: bibcode.to_string(),
content,
translation_content,
outline,
})
}
static HEADING_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^(#{2,})\s+(.+)$").unwrap());
/// 提取文献章节大纲。
///
/// `max_level` 控制最大标题级别2 = 仅 ## 章标题3 = ## + ###,以此类推。
/// 跳过 `#`(一级标题,通常是论文标题本身)。
pub fn extract_outline(content: &str, max_level: usize) -> Vec<MarkdownSection> {
let mut sections = Vec::new();
let mut index = 0usize;
@ -48,25 +168,23 @@ pub fn extract_outline(content: &str, max_level: usize) -> Vec<MarkdownSection>
index,
level,
heading,
start_line: line_num + 1, // 正文从下一行开始
end_line: 0, // 暂时标记,后续计算
char_start: 0, // 暂时标记,后续计算
start_line: line_num + 1,
end_line: 0,
char_start: 0,
char_end: 0,
});
index += 1;
}
}
// 计算每个章节的正文范围
for i in 0..sections.len() {
let next_start_line = if i + 1 < sections.len() {
sections[i + 1].start_line - 1 // 下一个标题行
sections[i + 1].start_line - 1
} else {
content.lines().count()
};
sections[i].end_line = next_start_line;
// 计算字节偏移量
let (cs, ce) = find_line_byte_range(content, sections[i].start_line, sections[i].end_line);
sections[i].char_start = cs;
sections[i].char_end = ce;
@ -75,11 +193,8 @@ pub fn extract_outline(content: &str, max_level: usize) -> Vec<MarkdownSection>
sections
}
/// 按章节名称(标题)提取内容。大小写不敏感,支持部分匹配。
pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<String> {
let lower = section_name.to_lowercase();
// 先获取大纲,再匹配
// 直接用 HEADING_RE 扫描,找到第一个匹配的标题
let mut found_start: Option<usize> = None;
let mut found_level: Option<usize> = None;
let mut found_line: Option<usize> = None;
@ -91,7 +206,7 @@ pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<Stri
let heading = caps.get(2).unwrap().as_str().trim();
if heading.to_lowercase().contains(&lower) {
found_start = Some(line_num + 1); // 正文从下一行开始
found_start = Some(line_num + 1);
found_level = Some(level);
found_line = Some(line_num);
break;
@ -103,7 +218,6 @@ pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<Stri
let section_level = found_level?;
let section_line = found_line?;
// 找到下一个同级或更高级标题作为结束
let total_lines = content.lines().count();
let mut end_line = total_lines;
for (line_num, line) in content.lines().enumerate() {
@ -125,20 +239,18 @@ pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<Stri
Some(section_content.trim().to_string())
}
/// 按序号提取章节内容0-based
pub fn extract_section_by_index(content: &str, index: usize) -> Option<String> {
let sections = extract_outline(content, 10); // max_level 足够大,捕获所有层级标题
let sections = extract_outline(content, 10);
let section = sections.get(index)?;
let section_content = &content[section.char_start..section.char_end];
Some(section_content.trim().to_string())
}
/// 根据行号范围计算字节偏移量
fn find_line_byte_range(content: &str, start_line: usize, end_line: usize) -> (usize, usize) {
let char_start = content
.lines()
.take(start_line)
.map(|l| l.len() + 1) // +1 for newline
.map(|l| l.len() + 1)
.sum::<usize>();
let char_end = if end_line >= content.lines().count() {
@ -158,6 +270,16 @@ fn find_line_byte_range(content: &str, start_line: usize, end_line: usize) -> (u
mod tests {
use super::*;
#[test]
fn test_read_mode_serialize() {
let mode = ReadMode::Full {
include_translation: true,
};
let serialized = serde_json::to_string(&mode).unwrap();
assert!(serialized.contains("Full"));
assert!(serialized.contains("include_translation"));
}
#[test]
fn test_extract_outline_basic() {
let md = "\
@ -283,7 +405,6 @@ More text.
## Results
Results text.
";
// "Methods" section should include both subsections but stop at "Results"
let content = extract_section_by_name(md, "Methods").unwrap();
assert!(content.contains("Subsection text"));
assert!(content.contains("More text"));

View File

@ -6,7 +6,7 @@ mod generic;
mod iop;
pub mod pdf;
use crate::api::helpers::{check_paper_paths_in_db, get_paper_from_db};
use crate::services::paper::{check_paper_paths_in_db, get_paper_from_db};
use std::path::Path;
use tracing::{error, info, warn};
@ -51,6 +51,24 @@ fn detect_parser(html: &str) -> Box<dyn JournalParser> {
// 将 PDF/MinerU 函数重新导出到 parser 模块层级
pub use pdf::{parse_pdf_via_mineru, poll_and_extract_mineru, submit_pdf_to_mineru};
/// 为解析后的 Markdown 构建 YAML 前置元数据front matter
fn build_front_matter(paper: &crate::services::paper::StandardPaper) -> String {
let title_json =
serde_json::to_string(&paper.title).unwrap_or_else(|_| format!("\"{}\"", paper.title));
let author_list = paper
.authors
.iter()
.map(|a| format!("\"{}\"", a))
.collect::<Vec<_>>()
.join(", ");
let pub_json = serde_json::to_string(&paper.pub_journal)
.unwrap_or_else(|_| format!("\"{}\"", paper.pub_journal));
format!(
"---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"https://ui.adsabs.harvard.edu/abs/{}/abstract\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n",
title_json, author_list, pub_json, paper.bibcode, paper.year, paper.keywords.join(",")
)
}
/// HTML → Markdown 核心解析管线
pub fn html_to_markdown(html_path: &Path) -> anyhow::Result<String> {
info!("正在解析 HTML → Markdown: {:?}", html_path);
@ -161,22 +179,14 @@ pub async fn parse_paper_service(
tokio::task::spawn_blocking(move || html_to_markdown(&html_abs_for_blocking)).await;
match md_result {
Ok(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(",")
);
let front_matter = build_front_matter(&paper);
parsed_markdown = format!("{}{}", front_matter, md);
let md_filename = format!("{}.md", bibcode);
let md_dest = library_dir.join("Markdown").join(&md_filename);
if let Some(parent) = md_dest.parent() {
std::fs::create_dir_all(parent).unwrap_or_default();
tokio::fs::create_dir_all(parent).await.unwrap_or_default();
}
if std::fs::write(&md_dest, &parsed_markdown).is_ok() {
if tokio::fs::write(&md_dest, &parsed_markdown).await.is_ok() {
relative_md_path = format!("Markdown/{}", md_filename);
}
}
@ -203,15 +213,7 @@ pub async fn parse_paper_service(
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(",")
);
let front_matter = build_front_matter(&paper);
parsed_markdown = format!("{}{}", front_matter, md);
let md_filename = format!("{}.md", bibcode);
let md_dest = library_dir.join("Markdown").join(&md_filename);
@ -258,9 +260,8 @@ pub async fn parse_paper_service(
// ── 测试 ──────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::common::{convert_html_tables_to_markdown, postprocess_markdown};
use super::common::postprocess_markdown;
use super::html_to_markdown;
use std::io::Write;
#[test]
fn test_convert_html_tables_to_markdown() {

114
src/services/pipeline.rs Normal file
View File

@ -0,0 +1,114 @@
// src/services/pipeline.rs
//
// 文献处理流水线服务:下载 → 解析 → 向量化。
use crate::api::AppState;
use tracing::info;
#[derive(Debug, Clone)]
pub struct PipelineResult {
pub bibcode: String,
pub completed: Vec<String>,
pub has_pdf: bool,
pub has_html: bool,
pub has_markdown: bool,
pub chunk_count: usize,
}
/// 执行文献处理流水线。按序执行指定的任务步骤("download", "parse", "embed")。
pub async fn process_paper_pipeline(
state: &AppState,
bibcode: &str,
tasks: &[String],
force: bool,
) -> anyhow::Result<PipelineResult> {
let mut completed = Vec::new();
let mut has_pdf = false;
let mut has_html = false;
let mut has_markdown = false;
let mut chunk_count = 0;
for task in tasks {
match task.as_str() {
"download" => {
info!("[Pipeline] download: {}", bibcode);
let paper = state
.downloader
.download_paper_service(&state.db, &state.config.library_dir, bibcode, force)
.await?;
has_pdf = paper.has_pdf;
has_html = paper.has_html;
completed.push("download".to_string());
}
"parse" => {
info!("[Pipeline] parse: {}", bibcode);
let markdown = crate::services::parser::parse_paper_service(
&state.db,
&state.config.library_dir,
&state.qiniu,
&state.config,
bibcode,
force,
)
.await?;
let markdown_len = markdown.len();
drop(markdown); // 释放大字符串内存,仅保留长度
has_markdown = true;
completed.push("parse".to_string());
info!("[Pipeline] parse done: {} chars", markdown_len);
}
"embed" => {
info!("[Pipeline] embed: {}", bibcode);
let paths = crate::services::paper::check_paper_paths_in_db(
&state.db,
&state.config.library_dir,
bibcode,
)
.await?;
let md_rel = match paths {
Some((_, _, Some(md), _)) => md,
_ => {
return Err(anyhow::anyhow!(
"未找到已解析的 Markdown 文件,请确保 tasks 中 parse 在 embed 之前执行"
))
}
};
let md_abs = state.config.library_dir.join(&md_rel);
let markdown_text = tokio::fs::read_to_string(&md_abs).await?;
let count = crate::services::rag::ingest_paper(
&state.db,
&state.embedding,
bibcode,
&markdown_text,
None, // 使用默认 ChunkerConfig
)
.await?;
chunk_count = count;
completed.push("embed".to_string());
info!("[Pipeline] embed done: {} chunks", count);
}
other => {
return Err(anyhow::anyhow!(
"无效的任务类型 '{}',仅支持: download, parse, embed",
other
));
}
}
}
Ok(PipelineResult {
bibcode: bibcode.to_string(),
completed,
has_pdf,
has_html,
has_markdown,
chunk_count,
})
}

View File

@ -1,15 +1,15 @@
// src/services/rag.rs
//
// RAG (Retrieval-Augmented Generation) 问答编排模块。
// RAG (Retrieval-Augmented Generation) 检索模块。
// 职责:
// 1. Ingest: 将文献 Markdown 切片 -> 向量化 -> 写入 SQLite vec0 + content 双表
// 2. Retrieve: 向量化用户提问 -> 相似度检索 -> 返回 Top-K 片段
// 3. Complete: 组装检索上下文 -> 调用 LLM 生成最终回答
// 3. Hybrid: 稠密向量 + 稀疏 BM25 双路融合检索 (RRF)
use sqlx::SqlitePool;
use tracing::{error, info, warn};
use crate::clients::llm::{ChatCompleter, Embedder};
use crate::clients::llm::Embedder;
use crate::services::chunker::{chunk_markdown, ChunkerConfig, StructuralChunk};
/// RAG 检索结果条目
@ -29,17 +29,8 @@ pub struct RetrievalResult {
pub distance: f64,
}
/// RAG 问答最终结果
#[derive(Debug, Clone, serde::Serialize)]
pub struct RagAnswer {
/// LLM 生成的回答
pub answer: String,
/// 检索到的参考来源
pub sources: Vec<RetrievalResult>,
}
// ─────────────────────── Ingest ───────────────────────
/// RAG 检索结果条目
///
/// 将一篇文献的 Markdown 内容切片、向量化并写入数据库
///
/// 流程:
@ -68,27 +59,33 @@ pub async fn ingest_paper(
chunks.len()
);
// 先清除该文献的旧切片(重新解析场景)
// 合并为两条原子 SQL 语句,避免循环查询和删除
sqlx::query(
"DELETE FROM vec_paper_chunks WHERE rowid IN (SELECT rowid FROM paper_chunks_content WHERE bibcode = ?)"
)
.bind(bibcode)
.execute(pool)
.await?;
sqlx::query("DELETE FROM paper_chunks_content WHERE bibcode = ?")
// 先清除该文献的旧切片(重新解析场景)。
// 在事务内原子执行两条 DELETE防止并发读取观测到 vec0 和 content 表的不一致状态。
{
let mut tx = pool.begin().await?;
sqlx::query(
"DELETE FROM vec_paper_chunks WHERE rowid IN (SELECT rowid FROM paper_chunks_content WHERE bibcode = ?)"
)
.bind(bibcode)
.execute(pool)
.execute(&mut *tx)
.await?;
sqlx::query("DELETE FROM paper_chunks_content WHERE bibcode = ?")
.bind(bibcode)
.execute(&mut *tx)
.await?;
tx.commit().await?;
}
// ── 事务外批量计算所有 embedding ──
// 避免长事务跨越数十次网络往返SQLite 写锁阻塞)。
let contents: Vec<String> = chunks.iter().map(|c| c.content.clone()).collect();
let embeddings: Vec<Option<Vec<f32>>> = match embedding_client.embed_batch(&contents).await {
Ok(all) => all.into_iter().map(Some).collect(),
Err(e) => {
// 批量失败时降级为逐条计算,保留"单块失败跳过、其余继续"的容错语义
// 批量失败时降级为逐条计算,保留"单块失败跳过、其余继续"的容错语义。
// TODO: 后续将 Embedder trait 改为 Arc<dyn Embedder> 以支持并发降级semaphore=3
warn!("文献 {} 批量向量化失败,降级逐条计算: {}", bibcode, e);
let mut out = Vec::with_capacity(chunks.len());
for chunk in &chunks {
@ -210,8 +207,8 @@ async fn retrieve_sparse(
query: &str,
top_k: usize,
) -> anyhow::Result<Vec<(String, i64, String, usize, String, i64)>> {
// 清洗查询:转义 FTS5 特殊字符,构建 MATCH 字符串
let sanitized = query.replace('"', "\"\"").replace('\'', "''");
// 清洗查询:使用统一的 FTS5 净化函数,防止操作符注入
let sanitized = crate::services::search::sanitize_fts5_query(query);
let fts_query = format!("\"{}\"", sanitized);
let rows: Vec<(String, i64, String, String, i64)> = sqlx::query_as(
@ -359,66 +356,6 @@ pub async fn retrieve_hybrid(
Ok(scored.into_iter().take(top_k).map(|(r, _)| r).collect())
}
// ─────────────────────── Complete ───────────────────────
/// 执行完整的 RAG 问答流程:检索 + LLM 生成
pub async fn ask(
pool: &SqlitePool,
embedding_client: &dyn Embedder,
llm_client: &dyn ChatCompleter,
question: &str,
top_k: usize,
) -> anyhow::Result<RagAnswer> {
// 1. 检索相关片段
let sources = retrieve(pool, embedding_client, question, top_k).await?;
if sources.is_empty() {
return Ok(RagAnswer {
answer: "未找到与该问题相关的文献内容。请先导入相关文献的 Markdown 版本。".to_string(),
sources: vec![],
});
}
// 2. 组装上下文(含章节路径信息)
let context = sources
.iter()
.enumerate()
.map(|(i, r)| {
let heading_info = if r.headings.is_empty() || r.headings == "Document" {
String::new()
} else {
format!(" (章节: {})", r.headings)
};
format!(
"[来源 {}: {} §{}{}]\n{}",
i + 1,
r.bibcode,
r.paragraph_index,
heading_info,
r.content
)
})
.collect::<Vec<_>>()
.join("\n\n---\n\n");
// 3. 构建 System Prompt
let system_prompt = "你是一位专业的天体物理学研究助手。\
\
使 [ N] 便\
\
";
let user_content = format!(
"## 检索到的参考片段\n\n{}\n\n---\n\n## 用户问题\n\n{}",
context, question
);
// 4. 调用 LLM
let answer = llm_client.complete(system_prompt, &user_content).await?;
Ok(RagAnswer { answer, sources })
}
// ─────────────────────── Helpers ───────────────────────
/// 将 f32 向量转换为字节切片(小端序),用于 sqlite-vec 的二进制格式
@ -430,6 +367,29 @@ fn embedding_to_bytes(embedding: &[f32]) -> Vec<u8> {
bytes
}
/// 向量化单篇文献服务:读取 Markdown 物理文件、判断其存在,并向量化切片入库
pub async fn embed_paper_service(
db: &SqlitePool,
library_dir: &std::path::Path,
bibcode: &str,
embedding_client: &dyn Embedder,
) -> anyhow::Result<usize> {
let (_, _, md_opt, _) =
crate::services::paper::check_paper_paths_in_db(db, library_dir, bibcode)
.await?
.ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?;
let md_rel = md_opt.ok_or_else(|| anyhow::anyhow!("文献尚未解析为 Markdown请先执行解析"))?;
let md_abs = library_dir.join(&md_rel);
if !md_abs.exists() {
anyhow::bail!("文献 Markdown 文件未找到,请重新解析");
}
let markdown_content = tokio::fs::read_to_string(&md_abs).await?;
let count = ingest_paper(db, embedding_client, bibcode, &markdown_content, None).await?;
Ok(count)
}
#[cfg(test)]
mod tests {
use super::*;

View File

@ -1,13 +1,11 @@
// src/services/search.rs
use std::collections::HashSet;
use crate::api::helpers::{
convert_ads_doc_to_standard, convert_arxiv_to_standard, get_paper_from_db, save_paper_to_db,
use crate::api::AppState;
use crate::services::paper::{
convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db, StandardPaper,
SQLITE_PARAM_LIMIT,
};
use crate::api::AppState;
use crate::api::StandardPaper;
use sqlx::Row;
use tracing::{error, warn};
/// 统一检索逻辑,合并去重 ADS 和 arXiv 数据,并自动记录引用拓扑与同步本地状态。
@ -96,18 +94,14 @@ pub async fn search_papers(
}
// 5. 批量回填本地库的标准元数据is_downloaded/has_markdown/pdf_error 等),避免 N+1。
// get_paper_from_db 内部还做了 library_dir 文件存在性检查,故仍逐条调用,
// 但先用一次 IN 查询过滤出"数据库中确实存在"的 bibcode仅对这些调用 get_paper_from_db。
// 用一次批量 IN 查询直接返回完整 StandardPaper无需逐条调用 get_paper_from_db。
let all_bibs: Vec<String> = unique_results.iter().map(|p| p.bibcode.clone()).collect();
let existing_bibs: HashSet<String> = batch_fetch_existing_bibcodes(&state.db, &all_bibs).await;
let existing_papers =
batch_fetch_existing_papers(&state.db, &state.config.library_dir, &all_bibs).await;
for paper in unique_results.iter_mut() {
if existing_bibs.contains(&paper.bibcode) {
if let Ok(db_paper) =
get_paper_from_db(&state.db, &state.config.library_dir, &paper.bibcode).await
{
*paper = db_paper;
}
if let Some(db_paper) = existing_papers.get(&paper.bibcode) {
*paper = db_paper.clone();
}
}
@ -137,41 +131,52 @@ async fn save_citation_topology(
return;
}
// 单事务批量 INSERT分批避免超参数上限
// 单一事务批量 INSERT分批仅用于避免 SQLite 参数上限
let mut tx = match db.begin().await {
Ok(tx) => tx,
Err(e) => {
warn!("无法开启引用关系事务: {}", e);
return;
}
};
for chunk in pairs.chunks(SQLITE_PARAM_LIMIT / 2) {
if let Ok(mut tx) = db.begin().await {
for (src, tgt) in chunk {
if let Err(e) = sqlx::query(
"INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)",
)
.bind(src)
.bind(tgt)
.execute(&mut *tx)
.await
{
warn!("保存引用关系失败 ({} -> {}): {}", src, tgt, e);
}
}
if let Err(e) = tx.commit().await {
warn!("提交引用关系事务失败: {}", e);
for (src, tgt) in chunk {
if let Err(e) = sqlx::query(
"INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)",
)
.bind(src)
.bind(tgt)
.execute(&mut *tx)
.await
{
warn!("保存引用关系失败 ({} -> {}): {}", src, tgt, e);
}
}
}
if let Err(e) = tx.commit().await {
warn!("提交引用关系事务失败: {}", e);
}
}
/// 批量查询数据库中已存在的 bibcode 集合WHERE IN 分批),避免 N+1。
async fn batch_fetch_existing_bibcodes(
/// 批量查询数据库中已存在的文献完整元数据WHERE IN 分批),避免 N+1。
async fn batch_fetch_existing_papers(
db: &sqlx::SqlitePool,
library_dir: &std::path::Path,
bibcodes: &[String],
) -> HashSet<String> {
let mut existing = HashSet::new();
) -> std::collections::HashMap<String, StandardPaper> {
use crate::services::paper::parse_paper_row;
let mut papers = std::collections::HashMap::new();
for chunk in bibcodes.chunks(SQLITE_PARAM_LIMIT) {
if chunk.is_empty() {
continue;
}
let placeholders = vec!["?"; chunk.len()].join(", ");
let sql = format!(
"SELECT bibcode FROM papers WHERE bibcode IN ({})",
"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) AS has_vector \
FROM papers WHERE bibcode IN ({})",
placeholders
);
let mut query = sqlx::query(&sql);
@ -179,13 +184,28 @@ async fn batch_fetch_existing_bibcodes(
query = query.bind(b);
}
if let Ok(rows) = query.fetch_all(db).await {
for row in rows {
let b: String = row.get(0);
existing.insert(b);
for row in &rows {
let paper = parse_paper_row(row, library_dir);
papers.insert(paper.bibcode.clone(), paper);
}
}
}
existing
papers
}
/// FTS5 查询净化:转义特殊字符,防止 FTS5 操作符注入。
///
/// FTS5 特殊字符(`"`、`'`、`*`)和结构操作符(`:`、`AND`、`OR`、`NOT`
/// 会被替换为空格,确保查询作为纯文本搜索而非布尔表达式执行。
pub fn sanitize_fts5_query(query: &str) -> String {
query
.replace('"', "\"\"")
.replace('\'', "''")
.replace(['*', ':'], " ")
.replace(" AND ", " ")
.replace(" OR ", " ")
.replace(" NOT ", " ")
.replace(" NEAR ", " ")
}
/// 本地文献库 FTS5 全文检索。
@ -197,8 +217,8 @@ pub async fn search_local_library(
query: &str,
limit: usize,
) -> anyhow::Result<Vec<StandardPaper>> {
// FTS5 查询需要转义特殊字符
let sanitized = query.replace('"', "\"\"").replace('\'', "''");
// FTS5 查询净化,防止操作符注入
let sanitized = sanitize_fts5_query(query);
let fts_query = format!("\"{}\"", sanitized);
let rows = sqlx::query_as::<_, PaperRow>(
@ -266,3 +286,100 @@ struct PaperRow {
markdown_path: Option<String>,
translation_path: Option<String>,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SearchResult {
pub result_type: String,
pub session_id: String,
pub title: Option<String>,
pub snippet: String,
pub created_at: Option<String>,
}
/// 跨会话/消息的智能体历史全文检索。
pub async fn search_agent_history(
db: &sqlx::SqlitePool,
q: &str,
scope: &str,
limit: i64,
session_id: Option<&str>,
) -> Result<Vec<SearchResult>, sqlx::Error> {
let mut results = Vec::new();
let sanitized = sanitize_fts5_query(q);
// 搜索会话
if scope == "all" || scope == "sessions" {
let rows = sqlx::query_as::<_, (String, String, String, Option<String>)>(
"SELECT s.session_id, s.title, \
snippet(agent_sessions_fts, 1, '<mark>', '</mark>', '...', 40) as snippet, \
s.created_at \
FROM agent_sessions_fts fts \
JOIN agent_sessions s ON s.session_id = fts.session_id \
WHERE agent_sessions_fts MATCH $1 \
ORDER BY rank LIMIT $2",
)
.bind(&sanitized)
.bind(limit)
.fetch_all(db)
.await?;
for (sid, title, snippet, created_at) in rows {
results.push(SearchResult {
result_type: "session".into(),
session_id: sid,
title: Some(title),
snippet,
created_at,
});
}
}
// 搜索消息
if scope == "all" || scope == "messages" {
let msg_rows = if let Some(sid) = session_id {
sqlx::query_as::<_, (String, String, String, String, Option<String>)>(
"SELECT fts.session_id, s.title, \
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 60) as snippet, \
fts.role, m.created_at \
FROM agent_messages_fts fts \
JOIN agent_sessions s ON s.session_id = fts.session_id \
JOIN agent_messages m ON m.rowid = fts.rowid \
WHERE agent_messages_fts MATCH $1 AND fts.session_id = $3 \
AND m.active = 1 \
ORDER BY rank LIMIT $2",
)
.bind(&sanitized)
.bind(limit)
.bind(sid)
.fetch_all(db)
.await?
} else {
sqlx::query_as::<_, (String, String, String, String, Option<String>)>(
"SELECT fts.session_id, s.title, \
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 60) as snippet, \
fts.role, m.created_at \
FROM agent_messages_fts fts \
JOIN agent_sessions s ON s.session_id = fts.session_id \
JOIN agent_messages m ON m.rowid = fts.rowid \
WHERE agent_messages_fts MATCH $1 AND m.active = 1 \
ORDER BY rank LIMIT $2",
)
.bind(&sanitized)
.bind(limit)
.fetch_all(db)
.await?
};
for (sid, title, snippet, role, created_at) in msg_rows {
results.push(SearchResult {
result_type: format!("message/{}", role),
session_id: sid,
title: Some(title),
snippet,
created_at,
});
}
}
Ok(results)
}

290
src/services/session.rs Normal file
View File

@ -0,0 +1,290 @@
// src/services/session.rs
//
// 智能体运行会话管理与指标统计服务层
use serde::{Deserialize, Serialize};
use sqlx::{Row, SqlitePool};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionSummary {
pub session_id: String,
pub title: String,
pub model: String,
pub mode: String,
pub turn_count: i32,
pub summary: Option<String>,
pub created_at: String,
pub updated_at: String,
}
/// 查询会话列表,支持分页
pub async fn list_sessions_service(
db: &SqlitePool,
limit: i64,
offset: i64,
) -> Result<Vec<SessionSummary>, sqlx::Error> {
let rows = sqlx::query(
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \
FROM agent_sessions \
WHERE deleted_at IS NULL \
ORDER BY updated_at DESC \
LIMIT ? OFFSET ?",
)
.bind(limit)
.bind(offset)
.fetch_all(db)
.await?;
let sessions: Vec<SessionSummary> = rows
.iter()
.map(|r| SessionSummary {
session_id: r.get("session_id"),
title: r.get("title"),
model: r.get("model"),
mode: r.get("mode"),
turn_count: r.get("turn_count"),
summary: r.get("summary"),
created_at: r.get("created_at"),
updated_at: r.get("updated_at"),
})
.collect();
Ok(sessions)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageRecord {
pub id: i64,
pub agent_name: String,
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,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionDetail {
pub session: SessionSummary,
pub messages: Vec<MessageRecord>,
}
/// 获取单个会话的详细信息(会话信息与激活的消息历史)
pub async fn get_session_detail_service(
db: &SqlitePool,
session_id: &str,
) -> Result<Option<SessionDetail>, sqlx::Error> {
// 查询会话元信息
let session_row = match sqlx::query(
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \
FROM agent_sessions \
WHERE session_id = ? AND deleted_at IS NULL",
)
.bind(session_id)
.fetch_optional(db)
.await?
{
Some(row) => row,
None => return Ok(None),
};
let session = SessionSummary {
session_id: session_row.get("session_id"),
title: session_row.get("title"),
model: session_row.get("model"),
mode: session_row.get("mode"),
turn_count: session_row.get("turn_count"),
summary: session_row.get("summary"),
created_at: session_row.get("created_at"),
updated_at: session_row.get("updated_at"),
};
// 查询消息列表(限制 500 条,防止会话消息过多导致内存膨胀)
const MAX_MESSAGES: i64 = 500;
let msg_rows = sqlx::query(
"SELECT id, agent_name, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, created_at \
FROM agent_messages \
WHERE session_id = ? AND active = 1 \
ORDER BY id ASC \
LIMIT ?"
)
.bind(session_id)
.bind(MAX_MESSAGES)
.fetch_all(db)
.await?;
let messages: Vec<MessageRecord> = msg_rows
.iter()
.map(|r| {
let tool_calls_json: Option<String> = r.get("tool_calls");
let metadata_json: Option<String> = r.get("metadata");
MessageRecord {
id: r.get("id"),
agent_name: r.get("agent_name"),
turn_index: r.get("turn_index"),
step_index: r.get("step_index"),
role: r.get("role"),
content: r.get("content"),
thought: r.get("thought"),
tool_calls: tool_calls_json.and_then(|s| serde_json::from_str(&s).ok()),
tool_call_id: r.get("tool_call_id"),
token_count: r.get("token_count"),
metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
created_at: r.get("created_at"),
}
})
.collect();
Ok(Some(SessionDetail { session, messages }))
}
/// 软删除指定会话
pub async fn delete_session_service(
db: &SqlitePool,
session_id: &str,
) -> Result<bool, sqlx::Error> {
let result = sqlx::query(
"UPDATE agent_sessions SET deleted_at = CURRENT_TIMESTAMP WHERE session_id = ? AND deleted_at IS NULL"
)
.bind(session_id)
.execute(db)
.await?;
Ok(result.rows_affected() > 0)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentMetrics {
pub total_sessions: i64,
pub total_tool_calls: i64,
pub tool_call_breakdown: serde_json::Value,
pub avg_steps_per_session: f64,
pub error_rate: f64,
}
/// 获取全局运行监控与性能分析指标。
/// 使用 `try_join!` 并行执行独立查询,减少串行等待。
pub async fn get_agent_metrics_service(db: &SqlitePool) -> Result<AgentMetrics, sqlx::Error> {
let sessions_fut = async {
let total: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions WHERE deleted_at IS NULL")
.fetch_one(db)
.await
.unwrap_or(0);
let avg: f64 = sqlx::query_scalar(
"SELECT COALESCE(AVG(CAST(turn_count AS REAL)), 0.0) \
FROM agent_sessions WHERE deleted_at IS NULL",
)
.fetch_one(db)
.await
.unwrap_or(0.0);
(total, avg)
};
let audit_fut = async {
let total: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM agent_audit_log WHERE status IN ('OK', 'FAIL')",
)
.fetch_one(db)
.await
.unwrap_or(0);
let errors: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_audit_log WHERE status = 'FAIL'")
.fetch_one(db)
.await
.unwrap_or(0);
(total, errors)
};
let breakdown_fut = async {
sqlx::query_as::<_, (String, i64)>(
"SELECT COALESCE(tool_name, 'unknown'), COUNT(*) as cnt \
FROM agent_audit_log \
WHERE status IN ('OK', 'FAIL') \
GROUP BY tool_name \
ORDER BY cnt DESC",
)
.fetch_all(db)
.await
.unwrap_or_default()
};
let ((total_sessions, avg_steps), (total_tool_calls, total_errors), breakdown_rows) =
tokio::join!(sessions_fut, audit_fut, breakdown_fut);
let tool_call_breakdown: serde_json::Value = breakdown_rows
.iter()
.map(|(name, cnt): &(String, i64)| serde_json::json!({ name: cnt }))
.fold(serde_json::json!({}), |mut acc, v| {
if let serde_json::Value::Object(map) = &mut acc {
if let serde_json::Value::Object(v_map) = v {
for (k, val) in v_map {
map.insert(k.clone(), val.clone());
}
}
}
acc
});
let error_rate = if total_tool_calls > 0 {
total_errors as f64 / total_tool_calls as f64
} else {
0.0
};
Ok(AgentMetrics {
total_sessions,
total_tool_calls,
tool_call_breakdown,
avg_steps_per_session: avg_steps,
error_rate,
})
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditLogEntry {
pub id: i64,
pub step: i32,
pub tool_name: Option<String>,
pub status: String,
pub elapsed_ms: i32,
pub output_preview: Option<String>,
pub created_at: String,
}
/// 查询单个会话的全部操作审计日志
pub async fn get_session_audit_service(
db: &SqlitePool,
session_id: &str,
) -> Result<Vec<AuditLogEntry>, sqlx::Error> {
let rows = sqlx::query(
"SELECT id, step, tool_name, status, elapsed_ms, output_preview, created_at \
FROM agent_audit_log \
WHERE session_id = ? \
ORDER BY id ASC",
)
.bind(session_id)
.fetch_all(db)
.await?;
let entries: Vec<AuditLogEntry> = rows
.iter()
.map(|r| AuditLogEntry {
id: r.get("id"),
step: r.get("step"),
tool_name: r.get("tool_name"),
status: r.get("status"),
elapsed_ms: r.get("elapsed_ms"),
output_preview: r.get("output_preview"),
created_at: r.get("created_at"),
})
.collect();
Ok(entries)
}

View File

@ -31,6 +31,52 @@ pub struct TargetInfo {
pub aliases: Vec<String>,
}
/// paper_targets 表的查询行类型,避免 14 字段元组重复定义
#[derive(Debug, sqlx::FromRow)]
struct TargetDbRow {
target_name: String,
ra: Option<String>,
dec: Option<String>,
parallax: Option<f64>,
parallax_err: Option<f64>,
spectral_type: Option<String>,
v_magnitude: Option<f64>,
otype: Option<String>,
oname: Option<String>,
pm_ra: Option<f64>,
pm_de: Option<f64>,
radial_velocity: Option<f64>,
photometry: Option<String>,
aliases: Option<String>,
}
impl From<TargetDbRow> for TargetInfo {
fn from(r: TargetDbRow) -> Self {
let aliases: Vec<String> = r
.aliases
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
r.photometry.and_then(|s| serde_json::from_str(&s).ok());
TargetInfo {
target_name: r.target_name,
ra: r.ra,
dec: r.dec,
parallax: r.parallax,
parallax_err: r.parallax_err,
spectral_type: r.spectral_type,
v_magnitude: r.v_magnitude,
otype: r.otype,
oname: r.oname,
pm_ra: r.pm_ra,
pm_de: r.pm_de,
radial_velocity: r.radial_velocity,
photometry,
aliases,
}
}
}
/// 从文本中提取可能的天体标识符
///
/// 支持的星表命名规范(基于 IAU 标准):
@ -202,8 +248,9 @@ fn parse_sesame_xml(xml: &str, original_name: &str) -> anyhow::Result<TargetInfo
// 提取别名列表
let mut aliases = Vec::new();
let alias_re = Regex::new(r"<alias>([^<]+)</alias>").unwrap();
for cap in alias_re.captures_iter(xml) {
static ALIAS_RE: std::sync::LazyLock<Regex> =
std::sync::LazyLock::new(|| Regex::new(r"<alias>([^<]+)</alias>").unwrap());
for cap in ALIAS_RE.captures_iter(xml) {
if let Some(alias) = cap.get(1) {
aliases.push(alias.as_str().trim().to_string());
}
@ -211,7 +258,10 @@ fn parse_sesame_xml(xml: &str, original_name: &str) -> anyhow::Result<TargetInfo
// 如果连坐标都查不到,说明 Sesame 无法识别这个天体
if ra.is_none() && dec.is_none() && aliases.is_empty() {
warn!("Sesame 未能识别天体: {}", original_name);
return Err(anyhow::anyhow!(
"CDS Sesame 无法识别天体: {}",
original_name
));
}
Ok(TargetInfo {
@ -290,25 +340,7 @@ pub async fn query_target_cached(
client: &Client,
) -> anyhow::Result<TargetInfo> {
// 1. 查本地缓存
let row = sqlx::query_as::<
_,
(
String,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<String>,
Option<f64>,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<f64>,
Option<String>,
Option<String>,
),
>(
let row = sqlx::query_as::<_, TargetDbRow>(
"SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \
otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \
FROM paper_targets WHERE target_name = ? LIMIT 1",
@ -317,45 +349,9 @@ pub async fn query_target_cached(
.fetch_optional(pool)
.await?;
if let Some((
name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry_json,
aliases_json,
)) = row
{
if let Some(db_row) = row {
info!("天体 {} 命中本地缓存", target_name);
let aliases: Vec<String> = aliases_json
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
photometry_json.and_then(|s| serde_json::from_str(&s).ok());
return Ok(TargetInfo {
target_name: name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry,
aliases,
});
return Ok(TargetInfo::from(db_row));
}
// 2. 缓存 Miss调用 Sesame
@ -418,25 +414,7 @@ pub async fn extract_and_cache_targets(
for name in names {
// 优先在 paper_targets 表中查找本地缓存
let row = sqlx::query_as::<
_,
(
String,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<String>,
Option<f64>,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<f64>,
Option<String>,
Option<String>,
),
>(
let row = sqlx::query_as::<_, TargetDbRow>(
"SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \
otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \
FROM paper_targets WHERE target_name = ? LIMIT 1",
@ -445,45 +423,9 @@ pub async fn extract_and_cache_targets(
.fetch_optional(pool)
.await;
if let Ok(Some((
tname,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry_json,
aliases_json,
))) = row
{
if let Ok(Some(db_row)) = row {
info!("天体 {} 命中本地缓存", name);
let aliases: Vec<String> = aliases_json
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
photometry_json.and_then(|s| serde_json::from_str(&s).ok());
results.push(TargetInfo {
target_name: tname,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry,
aliases,
});
results.push(TargetInfo::from(db_row));
} else {
// 本地缓存未命中,释放连接并外部查询 Sesame
info!("天体 {} 本地缓存未命中,查询 CDS Sesame...", name);
@ -544,6 +486,121 @@ pub async fn extract_and_cache_targets(
results
}
/// 获取指定文献关联的所有天体信息
pub async fn list_targets_for_paper(
db: &SqlitePool,
bibcode: &str,
) -> Result<Vec<TargetInfo>, sqlx::Error> {
let rows = sqlx::query_as::<
_,
(
String,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<String>,
Option<f64>,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<f64>,
Option<String>,
Option<String>,
),
>(
"SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \
otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \
FROM paper_targets WHERE bibcode = ? ORDER BY target_name",
)
.bind(bibcode)
.fetch_all(db)
.await?;
let targets: Vec<TargetInfo> = rows
.into_iter()
.map(
|(
name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry_json,
aliases_json,
)| {
let aliases: Vec<String> = aliases_json
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
photometry_json.and_then(|s| serde_json::from_str(&s).ok());
TargetInfo {
target_name: name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry,
aliases,
}
},
)
.collect();
Ok(targets)
}
/// 提取文献中的所有天体,并刷新关联记录缓存
pub async fn extract_and_refresh_targets(
db: &SqlitePool,
library_dir: &std::path::Path,
bibcode: &str,
client: &reqwest::Client,
) -> anyhow::Result<Vec<TargetInfo>> {
let (_, _, md_opt, _) =
crate::services::paper::check_paper_paths_in_db(db, library_dir, bibcode)
.await?
.ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?;
let md_rel = md_opt.ok_or_else(|| anyhow::anyhow!("文献尚未解析为 Markdown请先执行解析"))?;
let md_abs = library_dir.join(&md_rel);
if !md_abs.exists() {
anyhow::bail!("文献 Markdown 文件未找到,请重新解析");
}
let markdown_content = tokio::fs::read_to_string(&md_abs).await?;
// 重新识别前先清除该文献已有的天体关联记录。
// 使用显式事务确保 DELETE 原子执行,避免并发竞态导致的数据不一致。
{
let mut tx = db.begin().await?;
sqlx::query("DELETE FROM paper_targets WHERE bibcode = ?")
.bind(bibcode)
.execute(&mut *tx)
.await?;
tx.commit().await?;
}
let targets = extract_and_cache_targets(db, &markdown_content, bibcode, client).await;
Ok(targets)
}
#[cfg(test)]
mod tests {
use super::*;

View File

@ -198,9 +198,9 @@ pub async fn persist_translation(
let tr_filename = format!("{}.md", bibcode);
let tr_dest = library_dir.join("Translation").join(&tr_filename);
if let Some(parent) = tr_dest.parent() {
std::fs::create_dir_all(parent)?;
tokio::fs::create_dir_all(parent).await?;
}
std::fs::write(&tr_dest, translated_markdown)?;
tokio::fs::write(&tr_dest, translated_markdown).await?;
let relative_tr_path = format!("Translation/{}", tr_filename);
sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?")
@ -270,6 +270,55 @@ pub async fn translate_markdown(
}
}
/// 中英双栏对比翻译,带物理文件缓存。若缓存存在且不强制重译,则直读本地;否则调用大语言模型进行翻译并写入缓存。
pub async fn translate_paper_cached(
db: &sqlx::SqlitePool,
library_dir: &Path,
dict: &Dictionary,
llm_client: &crate::clients::llm::LlmClient,
bibcode: &str,
force: bool,
) -> anyhow::Result<String> {
let (_, _, md_opt, tr_opt) =
crate::services::paper::check_paper_paths_in_db(db, library_dir, bibcode)
.await?
.ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?;
// 若本地已存在翻译物理文件且未指明强制重译,直读本地缓存返回
if !force {
if let Some(tr_rel) = tr_opt {
let tr_abs = library_dir.join(&tr_rel);
if tr_abs.exists() {
if let Ok(content) = tokio::fs::read_to_string(&tr_abs).await {
return Ok(content);
}
}
}
}
// 检查英文解析文件是否存在
let md_rel = match md_opt {
Some(rel) => rel,
None => {
anyhow::bail!("文献必须先完成解析方可翻译");
}
};
let md_abs = library_dir.join(&md_rel);
if !md_abs.exists() {
anyhow::bail!("解析 Markdown 文件丢失");
}
let english_markdown = tokio::fs::read_to_string(&md_abs).await?;
// 调用 LLM 翻译服务并注入对照词表
let translated_markdown = translate_markdown(&english_markdown, dict, llm_client).await?;
// 翻译结果持久化
persist_translation(db, library_dir, bibcode, &translated_markdown).await?;
Ok(translated_markdown)
}
#[cfg(test)]
mod tests {
use super::*;

129
src/services/vision.rs Normal file
View File

@ -0,0 +1,129 @@
// src/services/vision.rs
//
// 图像加载、SSRF 防护与视觉大模型流式分析服务层
use crate::agent::runtime::AgentStreamEvent;
use crate::clients::llm::LlmClient;
use std::path::Path;
use tracing::info;
/// 图像分析服务:加载图片并通过大语言模型进行流式分析。
pub async fn analyze_image_stream_service(
image_path: &str,
question: &str,
vision_llm: &LlmClient,
library_dir: &Path,
sse_tx: Option<&tokio::sync::mpsc::UnboundedSender<AgentStreamEvent>>,
tool_call_id: Option<String>,
) -> anyhow::Result<(String, String)> {
let (base64_data, mime_type) = load_image(library_dir, image_path).await?;
info!(
"[vision_service] 流式分析图片: {} ({} bytes, type={})",
image_path,
base64_data.len(),
mime_type
);
let system_prompt = "你是一位专业的天体物理学家。请基于提供的图片和用户问题,给出专业、详细、清晰的中文分析。如有数学公式,使用 LaTeX 格式。";
// 流式输出:先发送状态头部(作为第一条 text_delta保证持久化到工具结果中
let header = format!(
"> 正在使用视觉模型({})分析图片: {}\n\n",
vision_llm.model(),
image_path
);
if let Some(tx) = sse_tx {
let _ = tx.send(AgentStreamEvent::TextDelta {
content: header.clone(),
tool_call_id: tool_call_id.clone(),
});
}
let sse_tx_clone = sse_tx.cloned();
let tc_id = tool_call_id.clone();
let answer = vision_llm
.analyze_image_stream(
system_prompt,
question,
&base64_data,
&mime_type,
move |delta: &str| {
if let Some(ref tx) = sse_tx_clone {
let _ = tx.send(AgentStreamEvent::TextDelta {
content: delta.to_string(),
tool_call_id: tc_id.clone(),
});
}
},
)
.await?;
Ok((header, answer))
}
/// 加载图片:本地路径或 HTTP(S) URL → base64 data + mime type
pub async fn load_image(library_dir: &Path, path: &str) -> anyhow::Result<(String, String)> {
if path.starts_with("http://") || path.starts_with("https://") {
// ── SSRF 防护 ──
let parsed =
reqwest::Url::parse(path).map_err(|e| anyhow::anyhow!("无效的图片 URL: {}", e))?;
let host = parsed
.host_str()
.ok_or_else(|| anyhow::anyhow!("URL 缺少 host"))?;
// 字面量快检 + 异步 DNS 深度校验(不阻塞 tokio worker
if crate::utils::ssrf::is_host_literal_blocked(host) {
return Err(anyhow::anyhow!(
"出于安全考虑,禁止请求内网/保留地址: {}",
host
));
}
if crate::utils::ssrf::is_host_blocked_async(host).await {
return Err(anyhow::anyhow!("DNS 解析后指向内网/保留地址: {}", host));
}
// 使用安全的重定向策略:允许跟随合法 CDN 重定向S3/Imgur 等),
// 但每次跳转都对目标 URL 做 scheme + host 字面量校验,拦截内网目标。
let client = reqwest::Client::builder()
.redirect(crate::utils::ssrf::safe_redirect_policy())
.build()?;
let resp = client.get(path).send().await?;
let mime = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.unwrap_or("image/png")
.to_string();
let bytes = resp.bytes().await?;
let b64 = base64_encode(&bytes);
Ok((b64, mime))
} else {
let abs_path = if path.starts_with('/') {
Path::new(path).to_path_buf()
} else {
library_dir.join(path)
};
if !abs_path.exists() {
return Err(anyhow::anyhow!("图片文件不存在: {}", abs_path.display()));
}
let ext = abs_path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("png")
.to_lowercase();
let mime = match ext.as_str() {
"jpg" | "jpeg" => "image/jpeg",
"gif" => "image/gif",
"webp" => "image/webp",
"svg" => "image/svg+xml",
_ => "image/png",
};
let bytes = tokio::fs::read(&abs_path).await?;
let b64 = base64_encode(&bytes);
Ok((b64, mime.to_string()))
}
}
fn base64_encode(bytes: &[u8]) -> String {
use base64::{engine::general_purpose, Engine as _};
general_purpose::STANDARD.encode(bytes)
}

View File

@ -97,16 +97,34 @@ pub fn is_host_literal_blocked(host: &str) -> bool {
false
}
/// 校验 URL 的 scheme 与 host 字面量是否安全(同步、无 DNS
/// 同步 DNS 解析并校验主机 IP 是否属于内网/保留网段。
/// 用于 reqwest 重定向等同步上下文。
pub fn is_host_blocked_sync(host: &str) -> bool {
if let Ok(ip) = host.parse::<IpAddr>() {
return ip_is_blocked(&ip);
}
// 使用标准库进行同步 DNS 解析
use std::net::ToSocketAddrs;
let lookup = format!("{}:80", host);
if let Ok(addrs) = lookup.to_socket_addrs() {
addrs.into_iter().any(|sa| ip_is_blocked(&sa.ip()))
} else {
// 解析失败按危险处理
true
}
}
/// 校验 URL 的 scheme、host 字面量与 DNS 解析是否安全(同步)。
///
/// 用于重定向回调中对每个跳转目标的快速拦截。
/// 用于重定向回调中对每个跳转目标的深度拦截。
pub fn is_url_safe(url: &reqwest::Url) -> bool {
let scheme = url.scheme();
if scheme != "http" && scheme != "https" {
return false;
}
match url.host_str() {
Some(host) => !is_host_literal_blocked(host),
// 同时进行字面量快速过滤与同步 DNS 深度过滤
Some(host) => !is_host_literal_blocked(host) && !is_host_blocked_sync(host),
None => false,
}
}
@ -196,4 +214,15 @@ mod tests {
async fn test_async_allowed_public() {
assert!(!is_host_blocked_async("example.com").await);
}
#[test]
fn test_sync_blocked_loopback() {
assert!(is_host_blocked_sync("localhost"));
assert!(is_host_blocked_sync("127.0.0.1"));
}
#[test]
fn test_sync_allowed_public() {
assert!(!is_host_blocked_sync("example.com"));
}
}