feat: 科研分析层全栈落地——光谱/时域/运动学分析工具链 + JWST/X 射线数据源 + 定时文献同步
数据分析层（新增 services/{spectrum,timeseries,analysis}）:
- 光谱参数提取 parameters.rs：LAMOST/SDSS/APOGEE/DESI FITS header 跨源归一化读取
Teff/logg/[Fe/H]/RV 及 ASPCAP 20+ 元素丰度，rayon 并发批量提取
- 谱线测量 lines.rs：内置真空/空气波长谱线表，窗口内极值搜索 + 梯形法积分 EW + FWHM，支持自定义谱线
- 交叉相关测速 cross_correlate.rs：对数波长重采样对齐，内置 Pickles 模板按光谱型插值，
CCF 峰值位置提取 RV 及不确定度
- 周期搜索 periodicity.rs：Lomb-Scargle 周期图（含 FAP 误报概率）+ BLS 凌星检测 + 相位折叠
- 变星分类 classification.rs：振幅/偏度/峰度/过零率/eta 等统计特征 + 规则分类（RR Lyrae/Cepheid/食双星/AGN 等）
- SED 拟合 sed.rs：多波段测光黑体模型拟合，输出 T_eff/半径/消光 A_V/光度及不确定度
- 运动学 kinematics.rs：视差+自行+RV → 银河系 UVW 空间速度，含移动星群成员概率（Banyan Σ 简化版）
- 化学丰度 chemistry.rs：[α/Fe] vs [Fe/H] 计算，厚盘/薄盘/晕星族判别
- 观测规划 observability.rs：目标升落时间/airmass/月相影响/曝光时间估算
- 赫罗图 hr_diagram.rs：Gaia TAP CMD 查询，新增 GET /api/analysis/hr-diagram 端点
数据获取层:
- JWST：clients/mast/jwst.rs 封装 MAST Portal 锥形检索 + JwstSpectrumFetcher（NIRSpec/MIRI 光谱）
- X 射线：clients/heasarc 封装 HEASARC TAP（ADQL）+ XMM-Newton/Chandra 光谱 fetcher
- 图像 cutout：SDSS SkyServer/STScI DSS/Pan-STARRS 三源 cutout + 发现图（Finding Chart）生成
- Source 枚举新增 Jwst/Xmm/Chandra 并注册 ObservationRegistry，前端 SOURCE_THEME 与筛选器同步三源
Agent 工具集（24→35）:
- 新增 9 个分析工具：get_spectrum_parameters / measure_spectral_lines / measure_radial_velocity /
find_period / classify_variable_star / fit_sed / analyze_kinematics / analyze_abundance_pattern / plan_observation
- batch_process：批量样本"查询→下载→分析→报告"流水线，并发控制防数据源速率限制
- literature_monitor：按 ADS 查询式/时间窗/最低引用数检查最新文献
定时文献同步:
- sync_queries 表新增 is_scheduled 列（migration 20260713）
- 新增 POST /sync/queries/:id/schedule 端点
- 服务启动时拉起每小时调度器，对 is_scheduled=1 的检索配置静默执行 ADS（entdate 增量）/arXiv 增量同步
- search_history 工具收敛至 services/search::search_agent_history，消除 FTS 查询逻辑重复
其他:
- plotting skill 由占位填充为完整科研绘图规范：光谱/光变/折叠曲线/CMD/SED/[α/Fe]/周期图/Mollweide/发现图 9 类 matplotlib 模板
- 删除死代码 streaming_executor.rs（929 行，仅剩 mod 声明引用，无调用方）
- 新增 docs/roadmap-research-features.md 科研功能路线图及实现状态
This commit is contained in:
Generated
+40
@@ -155,6 +155,7 @@ dependencies = [
|
||||
"obscura-net",
|
||||
"quick-xml",
|
||||
"rand 0.8.6",
|
||||
"rayon",
|
||||
"regex",
|
||||
"reqwest",
|
||||
"serde",
|
||||
@@ -836,6 +837,25 @@ dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-deque"
|
||||
version = "0.8.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb"
|
||||
dependencies = [
|
||||
"crossbeam-epoch",
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-epoch"
|
||||
version = "0.9.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
|
||||
dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-queue"
|
||||
version = "0.3.12"
|
||||
@@ -3039,6 +3059,26 @@ dependencies = [
|
||||
"getrandom 0.3.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rayon"
|
||||
version = "1.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
|
||||
dependencies = [
|
||||
"either",
|
||||
"rayon-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rayon-core"
|
||||
version = "1.13.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
|
||||
dependencies = [
|
||||
"crossbeam-deque",
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.5.18"
|
||||
|
||||
@@ -63,6 +63,7 @@ walkdir = "2"
|
||||
lru = "0.12"
|
||||
git2 = { version = "0.18", default-features = false, features = ["vendored-libgit2"] }
|
||||
dashmap = "6"
|
||||
rayon = "1.10"
|
||||
fitsio = { version = "0.21", features = ["fitsio-src"] } # FITS 解析；fitsio-src 自带 CFITSIO 源码编译，无需系统装 cfitsio
|
||||
fitsio-sys = { version = "0.5", features = ["fitsio-src"] } # 直接访问 CFITSIO C API（读 Gaia XP 变长数组列 PD(55)，fitsio crate 不支持）
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ import {
|
||||
import { AgentMarkdown } from './AgentMarkdown';
|
||||
import { ThoughtCard } from './ThoughtCard';
|
||||
import { ToolCallCard } from './ToolCallCard';
|
||||
import { TodoTaskCard } from './SpecialToolRenderers';
|
||||
import { TodoTaskCard, type TodoItem } from './SpecialToolRenderers';
|
||||
import { AnswerCard } from './AnswerCard';
|
||||
import { SubAgentContainer } from './SubAgentContainer';
|
||||
import { PARENT_COLORS } from './constants';
|
||||
@@ -106,7 +106,7 @@ interface AgentMessageListProps {
|
||||
showTaskSidebar?: boolean;
|
||||
setShowTaskSidebar?: (val: boolean) => void;
|
||||
hasTodos?: boolean;
|
||||
activeTodos?: any[];
|
||||
activeTodos?: TodoItem[];
|
||||
}
|
||||
|
||||
// 内部工具（is_internal = true）渲染模式配置：
|
||||
|
||||
@@ -419,7 +419,7 @@ export function PaperListCard({
|
||||
// ==========================================
|
||||
// 3. 任务规划看板卡片 (todo_write)
|
||||
// ==========================================
|
||||
interface TodoItem {
|
||||
export interface TodoItem {
|
||||
id: string;
|
||||
content: string;
|
||||
status: 'pending' | 'in_progress' | 'completed';
|
||||
@@ -453,7 +453,7 @@ export function TodoTaskCard({ todos }: TodoTaskCardProps) {
|
||||
{idx + 1}
|
||||
</div>
|
||||
);
|
||||
let rowClass = 'opacity-70';
|
||||
let rowClass: string;
|
||||
let titleClass = 'text-secondary font-medium';
|
||||
let badgeClass = 'bg-sunken text-tertiary';
|
||||
|
||||
|
||||
@@ -26,7 +26,6 @@ interface SubAgentContainerProps {
|
||||
export function SubAgentContainer({
|
||||
status,
|
||||
children,
|
||||
summary: _summary,
|
||||
isStreaming: parentStreaming,
|
||||
isCollapsed,
|
||||
onToggle,
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
// 与 ObservationPanel.SearchView 的区别：
|
||||
// - SearchView：单目标 + 单源 + 单产品（精细化参数）
|
||||
// - UnifiedSearchPanel：多目标 + 多源（广覆盖发现）
|
||||
import { useState, useMemo, useCallback, useEffect } from 'react';
|
||||
import { useState, useMemo, useCallback } from 'react';
|
||||
import {
|
||||
Globe,
|
||||
Search,
|
||||
@@ -64,20 +64,20 @@ export function UnifiedSearchPanel({
|
||||
downloadUnifiedSelected,
|
||||
downloading,
|
||||
}: UnifiedSearchPanelProps) {
|
||||
// ── 统一检索的下载结果 ──
|
||||
const [downloadResults, setDownloadResults] = useState<
|
||||
ObservationBatchResult[] | null
|
||||
>(null);
|
||||
|
||||
// 当重新检索结果变化时，重置下载结果
|
||||
useEffect(() => {
|
||||
setDownloadResults(null);
|
||||
}, [unifiedResult]);
|
||||
// ── 统一检索的下载结果（缓存归属哪次检索；重新检索后自动失效）──
|
||||
const [downloadCache, setDownloadCache] = useState<{
|
||||
result: UnifiedSearchResult | null;
|
||||
results: ObservationBatchResult[];
|
||||
} | null>(null);
|
||||
const downloadResults =
|
||||
downloadCache && downloadCache.result === unifiedResult
|
||||
? downloadCache.results
|
||||
: null;
|
||||
|
||||
const handleDownload = useCallback(async () => {
|
||||
const r = await downloadUnifiedSelected();
|
||||
setDownloadResults(r);
|
||||
}, [downloadUnifiedSelected]);
|
||||
setDownloadCache({ result: unifiedResult, results: r });
|
||||
}, [downloadUnifiedSelected, unifiedResult]);
|
||||
|
||||
// ── 输入模式 ──
|
||||
const [inputMode, setInputMode] = useState<InputMode>('coordinates');
|
||||
@@ -206,11 +206,10 @@ export function UnifiedSearchPanel({
|
||||
);
|
||||
|
||||
// 第一次加载完能力后默认选择最新的子版本以及所有发布版本
|
||||
useEffect(() => {
|
||||
// （渲染期一次性初始化，替代 effect 中的 setState，避免级联渲染）
|
||||
if (selectedSources === null && defaultSelectedKeys.length > 0) {
|
||||
setSelectedSources(new Set(defaultSelectedKeys));
|
||||
}
|
||||
}, [defaultSelectedKeys, selectedSources]);
|
||||
|
||||
// 同步选中的有效 Set
|
||||
const effectiveSelected = useMemo(() => {
|
||||
@@ -473,6 +472,8 @@ export function UnifiedSearchPanel({
|
||||
// ── 渲染子孙层级节点 (Level 1, 2, 3) ──
|
||||
const renderNestedNode = useCallback(
|
||||
(node: SourceTreeNode) => {
|
||||
// 递归渲染定义为内部具名函数，避免 useCallback 闭包自引用
|
||||
const renderNode = (node: SourceTreeNode) => {
|
||||
const isExpanded = expandedNodes.has(node.id);
|
||||
const hasChildren = node.children && node.children.length > 0;
|
||||
|
||||
@@ -537,7 +538,10 @@ export function UnifiedSearchPanel({
|
||||
}`}
|
||||
>
|
||||
{allChecked && (
|
||||
<svg className="w-2.5 h-2.5 fill-current" viewBox="0 0 20 20">
|
||||
<svg
|
||||
className="w-2.5 h-2.5 fill-current"
|
||||
viewBox="0 0 20 20"
|
||||
>
|
||||
<path d="M0 11l2-2 5 5L18 3l2 2L7 18z" />
|
||||
</svg>
|
||||
)}
|
||||
@@ -556,11 +560,13 @@ export function UnifiedSearchPanel({
|
||||
{/* 虚线连接子孙 */}
|
||||
{hasChildren && isExpanded && (
|
||||
<div className="flex flex-col gap-1 mt-1 border-l border-dashed border-subtle/50 ml-2.5 pl-1 animate-in fade-in duration-100">
|
||||
{node.children!.map((child) => renderNestedNode(child))}
|
||||
{node.children!.map((child) => renderNode(child))}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
return renderNode(node);
|
||||
},
|
||||
[expandedNodes, effectiveSelected, toggleExpand]
|
||||
);
|
||||
|
||||
@@ -63,6 +63,21 @@ export const SOURCE_THEME: Record<
|
||||
label: 'TESS',
|
||||
dot: 'bg-violet-500',
|
||||
},
|
||||
jwst: {
|
||||
badge: 'bg-pink-100 text-pink-700 dark:bg-pink-500/15 dark:text-pink-300',
|
||||
label: 'JWST',
|
||||
dot: 'bg-pink-500',
|
||||
},
|
||||
xmm: {
|
||||
badge: 'bg-teal-100 text-teal-700 dark:bg-teal-500/15 dark:text-teal-300',
|
||||
label: 'XMM-Newton',
|
||||
dot: 'bg-teal-500',
|
||||
},
|
||||
chandra: {
|
||||
badge: 'bg-red-100 text-red-700 dark:bg-red-500/15 dark:text-red-300',
|
||||
label: 'Chandra',
|
||||
dot: 'bg-red-500',
|
||||
},
|
||||
};
|
||||
|
||||
// ── 产品类型中文标签（product 轴）──
|
||||
|
||||
@@ -813,6 +813,9 @@ function LibraryView({
|
||||
{ value: 'gaia', label: 'Gaia' },
|
||||
{ value: 'sdss', label: 'SDSS' },
|
||||
{ value: 'desi', label: 'DESI' },
|
||||
{ value: 'jwst', label: 'JWST' },
|
||||
{ value: 'xmm', label: 'XMM-Newton' },
|
||||
{ value: 'chandra', label: 'Chandra' },
|
||||
]}
|
||||
/>
|
||||
</Field>
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
# Agent 架构硬化 — deepseek-harness 对齐改造
|
||||
|
||||
本文档记录 2026-08 的 agent 系统架构改造：以 deepseek-harness（DeepSeek
|
||||
开源的 agent harness，见 `libs/deepseek-harness`）验证过的架构不变量为
|
||||
参照，对本项目 agent 子系统的四项核心改造与若干升级。所有改动以
|
||||
"取其分层思想、不取其运行时插件树"为原则——本项目是垂直应用，
|
||||
Rust 的 trait + 注册表天然等价于 dsh 的 capability seam。
|
||||
|
||||
## 1. 统一 ReAct 引擎（消除四份手写循环）
|
||||
|
||||
**不变量**：仓库中有且仅有一个具体的 ReAct 循环实现（对应 dsh 的
|
||||
"agent-loop 是 harness 中唯一包含具体循环逻辑的包"）。
|
||||
|
||||
- `src/agent/engine.rs`：`ReactEngine` 承载完整循环（取消检查 → 压缩 →
|
||||
预算/nag → 步数上限 → 后台通知 → LLM 调用含恢复阶梯 → 工具执行 → 循环）。
|
||||
- 行为差异通过组合字段注入而非分叉实现：
|
||||
|
||||
| 关注点 | lead（主代理） | subagent | teammate |
|
||||
|---|---|---|---|
|
||||
| SSE 事件 | `EventTap::channel` | `EventTap::prefixed(tx, "sub")` | `EventTap::none()` |
|
||||
| 持久化 | `DbMessageSink`（agent_name=lead） | `DbMessageSink`（agent_name=sub_xxx） | None |
|
||||
| AskUser 权限 | `AskPolicy::Interactive`（120s 挂起） | `AskPolicy::AutoDeny` | `AskPolicy::AutoDeny` |
|
||||
| 取消 | `CancelSource::Session`（DashMap） | 不可取消 | `CancelSource::Flag` |
|
||||
| 恢复阶梯/熔断/重复检测 | 共享（runtime 级 Arc） | 引擎局部实例 | 引擎局部实例 |
|
||||
|
||||
- 队友因此获得了与主代理对等的能力（错误恢复、压缩熔断、多槽重复检测）。
|
||||
- `streaming_executor.rs`（936 行零调用死代码）已删除。
|
||||
|
||||
## 2. 会话级运行时生命周期
|
||||
|
||||
**不变量**：runtime 状态属于会话而非请求。
|
||||
|
||||
`src/agent/runtime/session_registry.rs` 的 `SessionRuntimeRegistry`
|
||||
（挂在 `AppState.agent_runtimes`）：
|
||||
|
||||
- 同一会话的请求复用同一 `AgentRuntime`——后台任务队列、压缩折叠日志、
|
||||
文件缓存、拒绝追踪器、压缩熔断器、prompt cache 不再随 HTTP 请求销毁
|
||||
（历史上"后台任务结果跨请求丢失、每 turn 重复压缩"两个 bug 的根因）。
|
||||
- 每会话一个 turn 互斥锁：并发请求 fail-loud 返回 409。
|
||||
- 空闲 2 小时的条目由后台清扫回收；会话删除时显式移除。
|
||||
- 会话 ID 在请求入口预分配（`create_or_resume_session_preallocated`），
|
||||
新会话的首个请求也能写入取消标记（修复 SSE 超时只 abort 不落标记）。
|
||||
- 模式回放：runtime 首次创建时从 DB 读取会话模式，后续请求不再因传参
|
||||
不同而静默切换行为。
|
||||
- 压缩递归守卫从进程级 `AtomicBool` 改为按 session_id 的集合守卫
|
||||
（`try_begin_compaction`/`end_compaction`），并发会话压缩互不干扰。
|
||||
|
||||
## 3. 会话事件日志（事件溯源的生命周期侧面）
|
||||
|
||||
**不变量**：模型可见 ⟺ 已日志化；turn/compaction 以日志化锁开闭。
|
||||
|
||||
新表 `agent_events`（migration 20260819000001）+
|
||||
`src/agent/runtime/session_events.rs`：
|
||||
|
||||
- `turn_start`/`turn_end`（含结构化终止原因 `reason_label`）、
|
||||
`compaction_start`/`compaction_end` 构成日志化锁。
|
||||
- 崩溃恢复不截断：`reconcile_interrupted` 检测开口的 start 事件并合成
|
||||
`interrupted` 关闭（对应 dsh "不伪造完成、但补写中断事实"的策略）。
|
||||
- `context_snapshot`：turn 内发生过压缩时，turn 结束保存折叠后上下文 +
|
||||
消息高水位 `base_message_id`；下一 turn 回放"快照 + id > base 的增量"
|
||||
（`context.rs::load_folded_history`），不再从原始消息重建后重新压缩。
|
||||
rewind/retry/branch 后快照作废（`remove_context_snapshots`）。
|
||||
|
||||
## 4. 工具契约集中化（消除散弹式修改）
|
||||
|
||||
**不变量**：工具的全部行为元数据声明在工具自身（AgentTool），消费方查询
|
||||
trait 而非维护平行的名字名单。
|
||||
|
||||
`AgentTool` 新增声明方法：
|
||||
|
||||
- `untrusted_output()` — 输出来自外部源需 `<untrusted_tool_result>` 包裹
|
||||
（原 `untrusted.rs` 按名列表）；注册表查询不到工具时回退名字启发式。
|
||||
- `hardline_check(args)` — 参数级不可绕过拒绝（原 executor 按名路由）。
|
||||
- `causes_file_changes()` — 执行前触发 checkpoint 快照（原
|
||||
`CHECKPOINT_TRIGGER_TOOLS` 名单）。
|
||||
- `loop_signals() -> ToolSignals` — 循环行为信号（todo nag 重置/todos
|
||||
持久化/手动压缩请求），替代主循环对 `todo_write`/`compress_context`
|
||||
的按名特判。
|
||||
- 输出契约：`ToolOutput.value`（canonical JSON，机器消费：重放/审计/剪枝）
|
||||
+ `content`（模型侧 render 投影），value 随消息 metadata 持久化。
|
||||
样例：`search_papers`、`todo_write`。
|
||||
- 注册表测试从硬编码数量断言（`assert_eq!(defs.len(), 35)`）改为
|
||||
inventory 不变量断言（名字唯一/元数据完整/信号与身份一致），
|
||||
新增工具不再需要改测试。
|
||||
|
||||
## 5. 权限单调性（deny > ask > allow，只收紧不放松）
|
||||
|
||||
**不变量**：多源决策合并唯一入口 `permission::tighten(base, candidate)`：
|
||||
|
||||
- Deny 粘滞——任何层不能放行另一层已 Deny 的调用；
|
||||
- Ask 不可被放松为 Allow，可升级为 Deny；
|
||||
- executor 的 hook 请求、工具级声明、会话级规则全部经 `tighten` 合并；
|
||||
- 无人值守上下文（子代理/队友）的 AskUser fail-closed 自动拒绝
|
||||
（`AskPolicy::AutoDeny` + `engine::auto_deny_ask`）。
|
||||
|
||||
历史上与此语义重复且从未接线的 `resolve_permission_precedence` 已删除。
|
||||
|
||||
## 6. 上下文管理升级
|
||||
|
||||
- **KV-cache 纪律**：system prompt 只包含字节稳定 section；技能清单/
|
||||
项目记忆以 durable 动态上下文快照（user 消息）追加在历史尾部，哈希
|
||||
未变不重注入（`context.rs::build_initial_context` +
|
||||
`AgentRuntime::build_dynamic_context`）。
|
||||
- **记忆智能选择接线**：记忆条目 > 8 时走 `select_relevant_memories`
|
||||
（LLM 结构化选择 + 指数时间衰减，此前为零调用代码）。
|
||||
- **确定性剪枝**：`micro_compact` 对长工具结果保留头/尾预览
|
||||
（`prune_tool_result`），短结果仍为纯占位符。
|
||||
- **Spill 检索指引**：溢出 stub 携带工具名 + 明确取回路径（read_file +
|
||||
max_lines 分段）。
|
||||
- **Token 估算统一**：全循环使用 `compact::estimate_message_tokens`
|
||||
（CJK 加权），消除主循环 `len/3` 与 compact 两套口径。
|
||||
|
||||
## 7. 其他接线与修复
|
||||
|
||||
- **团队协作接线**：`spawn_teammate` 等 4 个工具经
|
||||
`ToolRegistry::new_with_team` 注册进会话 runtime（TeamManager 按
|
||||
session_key 初始化邮箱目录）；队友注册表真正排除 subagent/spawn_teammate
|
||||
（原注释声称排除但未过滤）。
|
||||
- **Continuable 子代理**：`subagent` 工具新增 `agent_name` + `followup`
|
||||
参数——命名子代理历史持久化，followup 回放全部历史续话（dsh
|
||||
Activation 模型的冷恢复路径）。
|
||||
- **IdlePoller 接线**：`AGENT_AUTONOMOUS_ENABLED=true` 时启动自治轮询，
|
||||
经会话注册表获取 runtime（模式从 DB 回放）。
|
||||
- **重复检测多槽**：`DuplicateDetector` 连续通道 + 滑动窗口双通道，
|
||||
可检测 A/B 交替死循环（原单槽实现检测不到）。
|
||||
|
||||
## 不采用的设计（与理由）
|
||||
|
||||
- Cordis 运行时插件树 / bundle/patch/profile 配置组合层：服务于通用
|
||||
harness 的第三方插件生态，本项目是单二进制垂直应用。
|
||||
- Typert 类型图 RPC、host/client 双聚合、Python SDK：无对应需求。
|
||||
- 适配器休眠挂载、profiles 热重载：模型链固定。
|
||||
@@ -134,3 +134,4 @@ Agent 系统的所有可配置参数，按子系统分类。
|
||||
| `src/agent/compact.rs` (extract_memories_from_compaction) | `EXTRACT_MEMORY_ENABLED` (P3 压缩记忆桥接) |
|
||||
| `src/agent/memory/extraction.rs` (ExtractionConfig) | `EXTRACT_MEMORY_ENABLED`, `EXTRACT_MEMORY_THROTTLE_TURNS`, `EXTRACT_MEMORY_MAX_STEPS` |
|
||||
| `src/services/download.rs` | `OBSCURA_ALLOW_PRIVATE_NETWORK`（程序自动设置） |
|
||||
| `src/main.rs` (IdlePoller 启动门控) | `AGENT_AUTONOMOUS_ENABLED`（默认 false；true/1 时启用自治轮询，见 dsh-alignment.md §7） |
|
||||
|
||||
@@ -2,6 +2,12 @@
|
||||
|
||||
AstroResearch 内置了一个基于 **ReAct** (Thought → Action → Observation) 范式的科研智能体引擎 (`src/agent/`)，参考 Claude Code 的分层设计。以下对各子系统的架构、数据流和内部逻辑进行完整说明。
|
||||
|
||||
> **架构演进**：2026-08 起，ReAct 循环的唯一实现在 `src/agent/engine.rs`
|
||||
> （`ReactEngine`，主代理/子代理/队友以不同注入组合复用）；
|
||||
> `AgentRuntime` 按会话缓存复用（`SessionRuntimeRegistry`）；
|
||||
> turn/compaction 生命周期与上下文快照落在 `agent_events` 事件日志。
|
||||
> 详见 [dsh-alignment.md](dsh-alignment.md)。
|
||||
|
||||
### 整体架构
|
||||
|
||||
```mermaid
|
||||
@@ -14,16 +20,18 @@ graph TD
|
||||
AskUser["问答 /api/chat/questions + /api/chat/answer"]
|
||||
end
|
||||
|
||||
subgraph Runtime["AgentRuntime — ReAct 引擎"]
|
||||
subgraph Runtime["AgentRuntime — 会话级编排"]
|
||||
RunTurn["run_turn() 主入口"]
|
||||
SP["SystemPrompt 组装器"]
|
||||
CtxBuild["Context Builder 上下文构建"]
|
||||
ReAct["ReAct 主循环"]
|
||||
Engine["ReactEngine 统一循环（engine.rs）"]
|
||||
Streaming["streaming.rs 流式处理"]
|
||||
Executor["executor.rs 并行执行"]
|
||||
Finalize["finalize.rs 会话收尾"]
|
||||
TokenBudget["token_budget.rs"]
|
||||
CircuitBreaker["circuit_breaker.rs"]
|
||||
SessionReg["SessionRuntimeRegistry 会话级缓存"]
|
||||
EventLog["agent_events 事件日志"]
|
||||
end
|
||||
|
||||
subgraph Tools["工具系统 (tools/)"]
|
||||
|
||||
@@ -0,0 +1,438 @@
|
||||
# AstroResearch 科研功能扩展路线图
|
||||
|
||||
## 项目现状
|
||||
|
||||
### 已实现的核心能力
|
||||
|
||||
**数据源（9 个）**：
|
||||
|
||||
| 数据源 | 产品类型 | 状态 |
|
||||
| ---------- | ------------------------------------------------ | ---- |
|
||||
| LAMOST | 光谱（LRS/MRS） | 完整 |
|
||||
| Gaia | XP 光谱（BP/RP）、RVS、历元测光 | 完整 |
|
||||
| SDSS | 光学光谱（spec）、APOGEE 近红外（apstar/aspcap） | 完整 |
|
||||
| DESI | 暗能量光谱仪（coadd） | 完整 |
|
||||
| 2MASS | 近红外测光（J/H/K） | 完整 |
|
||||
| AllWISE | 中红外测光（W1-W4） | 完整 |
|
||||
| Pan-STARRS | 光学测光（g/r/i/z/y） | 完整 |
|
||||
| ZTF | 时域光变曲线（zg/zr/zi） | 完整 |
|
||||
| TESS | 空间时域光变曲线（2min/FFI） | 完整 |
|
||||
|
||||
**Agent 工具集**：
|
||||
|
||||
- `search_papers` — ADS/arXiv 文献搜索
|
||||
- `process_paper` — 下载、解析、向量化、翻译
|
||||
- `get_paper_content` / `get_paper_outline` — 阅读解析后文献
|
||||
- `rag_search` — 稠密向量 + BM25 混合检索
|
||||
- `query_target` — CDS Sesame 天体名称解析
|
||||
- `find_observation` — 跨源观测数据统一下载
|
||||
- `catalog_operation` — VizieR TAP 星表查询（ADQL + 锥形检索）
|
||||
- `save_note` — 高亮/笔记存储
|
||||
- `get_paper_metadata` — 完整文献元数据
|
||||
- `search_local_library` — 本地文献库搜索
|
||||
- `get_citation_network` — 引文网络图
|
||||
- `analyze_image` — 视觉模型图像分析
|
||||
- 文件操作、bash、记忆、技能、子 Agent、团队协作等
|
||||
|
||||
**已有技能**：
|
||||
|
||||
- `methodology` — 系统性文献综述方法论
|
||||
- `plotting` — 科研绘图规范（占位，未实现）
|
||||
- `presentation` — 演示文稿技能
|
||||
|
||||
---
|
||||
|
||||
## 一、数据获取层 — 补全缺失的数据源
|
||||
|
||||
### 1.1 JWST/MIRI 数据获取
|
||||
|
||||
**目标**：通过 MAST API 获取詹姆斯·韦伯空间望远镜 NIRSpec/MIRI 光谱数据。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/clients/mast/jwst.rs`，封装 MAST JWST 查询 API
|
||||
- 实现 `JwstSpectrumFetcher`，注册到 `ObservationRegistry`
|
||||
- Source 枚举新增 `Jwst` 变体
|
||||
- 支持 NIRSpec IFU/MOS、MIRI MRS/LRS 光谱产品
|
||||
- FITS 解析支持 JWST 特有的多扩展 BinTable 格式
|
||||
|
||||
**价值**：JWST 是当前最热门的望远镜，大量前沿论文依赖 JWST 数据。
|
||||
|
||||
**预估工作量**：中等（MAST API 文档完善，但 FITS 结构复杂）。
|
||||
|
||||
### 1.2 X 射线数据（XMM-Newton / Chandra）
|
||||
|
||||
**目标**：通过 HEASARC TAP 获取高能天体物理观测数据。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/clients/heasarc/`，封装 HEASARC TAP 查询
|
||||
- 实现 `XmmSpectrumFetcher` 和 `ChandraSpectrumFetcher`
|
||||
- Source 枚举新增 `Xmm` 和 `Chandra` 变体
|
||||
- 支持 X 射线光谱（PHA）和光变曲线产品
|
||||
|
||||
**价值**：覆盖高能波段，满足活动星系核、X 射线双星、星系团等研究需求。
|
||||
|
||||
**预估工作量**：中等（HEASARC TAP 接口标准，但数据格式多样）。
|
||||
|
||||
### 1.3 VLBI/Gaia 天体测量数据
|
||||
|
||||
**目标**：通过 VizieR 获取高精度自行、视差和天体测量数据。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 扩展现有 `CatalogOperationTool`，预置 Gaia DR3 天体测量常用查询模板
|
||||
- 新增 `GaiaAstrometryFetcher`，专门处理高精度自行/视差数据
|
||||
- 支持按自行/视差筛选的 cone search
|
||||
|
||||
**价值**：精密天体测量是银河考古学和恒星动力学的基础。
|
||||
|
||||
**预估工作量**：小（VizieR 已支持，主要是模板和预置查询）。
|
||||
|
||||
### 1.4 光谱图像 Cutout
|
||||
|
||||
**目标**：实现 SDSS/Gaia/DSS 图像 cutout 下载 + 前端渲染。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 实现 `ImageCutoutFetcher`，支持 SDSS SkyServer cutout API、DSS (STScI)、Gaia 推荐图
|
||||
- 完成 `ImagePreview` 解析器（preview.rs 中已有占位）
|
||||
- 前端 `ObservationPreviewRenderer` 新增图像渲染 case
|
||||
- 支持自定义 cutout 大小和坐标
|
||||
|
||||
**价值**：观测规划和文献中的发现图（Finding Chart）必备。
|
||||
|
||||
**预估工作量**：中等（API 简单，但前端图像渲染需要新组件）。
|
||||
|
||||
---
|
||||
|
||||
## 二、数据分析层 — 光谱/光变科学分析
|
||||
|
||||
### 2.1 光谱参数自动提取
|
||||
|
||||
**目标**：从 LAMOST/SDSS FITS 文件中自动读取恒星大气参数（Teff, logg, [Fe/H], 径向速度等），结构化入库。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/services/spectrum/parameters.rs`，实现 FITS Header/BinTable 参数提取
|
||||
- LAMOST：从 `params.fits` 或 `spec.fits` header 读取 `TEFF`, `LOGG`, `FEH`, `RV`
|
||||
- SDSS/BOSS：从 `line.dat` 或 `aspcapStar` 读取 ASPCAP 参数
|
||||
- SDSS/APOGEE：从 `aspcapStar` FITS 读取完整化学丰度
|
||||
- DESEI：从 `coadd` FITS header 读取红移和分类
|
||||
- 结构化数据入库（新建 `spectrum_parameters` 表或扩展 `observation_cache`）
|
||||
- Agent 新增 `get_spectrum_parameters` 工具，查询已提取的参数
|
||||
|
||||
**价值**：每篇恒星物理论文都需要这类参数，手动提取极其低效。
|
||||
|
||||
**预估工作量**：中等（FITS 解析已有基础，需逐源适配 header 约定）。
|
||||
|
||||
### 2.2 光谱线识别与等值宽度测量
|
||||
|
||||
**目标**：标注 Balmer/H&K/金属线位置，自动计算等值宽度（EW）和半高全宽（FWHM）。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/services/spectrum/lines.rs`，内置常见谱线波长表（真空/空气波长）
|
||||
- 实现谱线搜索算法：在给定波长窗口内找局部极值（发射线找极大，吸收线找极小）
|
||||
- 实现数值积分 EW 计算（梯形法）
|
||||
- 支持用户自定义谱线列表
|
||||
- Agent 新增 `measure_spectral_lines` 工具，输入光谱路径 + 谱线名称列表，输出 EW/FWHM
|
||||
- 前端光谱图支持谱线标注叠加层
|
||||
|
||||
**价值**：光谱分析的核心操作，天文学生和研究者日常高频使用。
|
||||
|
||||
**预估工作量**：中等偏高（谱线表容易，但稳健的峰值搜索和 EW 计算需要仔细实现）。
|
||||
|
||||
### 2.3 交叉相关测速
|
||||
|
||||
**目标**：通过模板匹配法测量径向速度（RV），支持多波段光谱。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/services/spectrum/cross_correlate.rs`
|
||||
- 内置或下载标准模板光谱（如 Pickles Atlas、PHOENIX 模型）
|
||||
- 实现 FFT 加速的交叉相关函数计算
|
||||
- 从互相关峰值位置提取 RV + 不确定度估计
|
||||
- 支持用户上传自定义模板
|
||||
- Agent 新增 `measure_radial_velocity` 工具
|
||||
|
||||
**价值**：双星、视向速度变星、银河系动力学研究的核心工具。
|
||||
|
||||
**预估工作量**：中等（算法成熟，但模板管理和多波段适配需要设计）。
|
||||
|
||||
### 2.4 光变周期搜索
|
||||
|
||||
**目标**：基于 Lomb-Scargle 周期图和 BLS（Box Least Squares）算法，自动搜索光变周期。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/services/timeseries/periodicity.rs`
|
||||
- 实现 Lomb-Scargle 周期图（非均匀采样友好）
|
||||
- 实现 BLS 算法（凌星检测专用）
|
||||
- 自动报告最优周期、FAP（误报概率）、次优周期
|
||||
- 支持 ZTF、TESS、Gaia 历元测光数据
|
||||
- Agent 新增 `find_period` 工具，输入光变曲线文件，输出周期 + 周期图
|
||||
- 前端支持折叠光变曲线（Phase-folded light curve）展示
|
||||
|
||||
**价值**：ZTF/TESS 数据的刚需，变星和系外行星研究的核心分析。
|
||||
|
||||
**预估工作量**：中等（Lomb-Scargle 算法有 Rust 实现可参考，BLS 较复杂）。
|
||||
|
||||
### 2.5 光变分类与特征提取
|
||||
|
||||
**目标**：自动识别变星类型，提取振幅、形态参数等特征。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/services/timeseries/classification.rs`
|
||||
- 基于周期 + 振幅 + 偏度 + 峰度等统计特征的简单分类器
|
||||
- 支持 RR Lyrae、Cepheid、食双星、脉动变星、AGN 等主要类型
|
||||
- 可选：集成轻量 ML 模型（如 Random Forest，通过 ONNX Runtime）
|
||||
- Agent 新增 `classify_variable_star` 工具
|
||||
|
||||
**价值**：时域巡天（ZTF/TESS）数据量巨大，自动化分类是必要环节。
|
||||
|
||||
**预估工作量**：高（分类器训练和特征工程需要较多工作）。
|
||||
|
||||
### 2.6 SED 拟合与建模
|
||||
|
||||
**目标**：多波段测光数据拟合黑体或模板光谱能量分布，求解天体物理参数。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/services/analysis/sed.rs`
|
||||
- 自动拉取目标所有可用波段测光（2MASS + WISE + PS1 + Gaia）
|
||||
- 实现黑体辐射模型拟合（最小二乘/MCMC）
|
||||
- 可选模板 SED 拟合（BaSeL/PHOENIX 模板库）
|
||||
- 输出 T_eff、光度、消光 A_V 及其不确定度
|
||||
- 前端 SED 图展示（观测点 + 拟合曲线 + 残差）
|
||||
- Agent 新增 `fit_sed` 工具
|
||||
|
||||
**价值**：天体物理参数估计的基础方法，适用于恒星、AGN、超新星等多种天体。
|
||||
|
||||
**预估工作量**：中等偏高（黑体拟合简单，但模板拟合和 MCMC 需要较多工作）。
|
||||
|
||||
### 2.7 自行/视差运动学分析
|
||||
|
||||
**目标**：基于 Gaia DR3 高精度自行 + 视差，计算银河系三维空间速度（UVW），识别移动星群。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/services/analysis/kinematics.rs`
|
||||
- 实现坐标变换：(RA, Dec, parallax, pmRA, pmDec, RV) → (X, Y, Z, U, V, W)
|
||||
- 实现移动星群成员概率计算（基于 Banyan Σ 算法简化版）
|
||||
- 支持批量样本处理
|
||||
- Agent 新增 `analyze_kinematics` 工具
|
||||
|
||||
**价值**：银河考古学（Galactic Archaeology）的核心分析方法。
|
||||
|
||||
**预估工作量**：中等（坐标变换公式成熟，Banyan Σ 需要简化实现）。
|
||||
|
||||
### 2.8 化学丰度模式分析
|
||||
|
||||
**目标**：多元素丰度比（[α/Fe] vs [Fe/H]）自动提取与可视化，区分厚盘/薄盘/晕星。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 扩展光谱参数提取（2.1），从 APOGEE ASPCAP 数据读取完整化学丰度
|
||||
- 新增 `src/services/analysis/chemistry.rs`
|
||||
- 实现 [α/Fe] 计算（α 元素：O, Mg, Si, S, Ca, Ti 的加权平均）
|
||||
- 自动绘制 [α/Fe] vs [Fe/H] 图，并按运动学成分着色
|
||||
- Agent 新增 `analyze_abundance_pattern` 工具
|
||||
|
||||
**价值**：银河考古学的核心工具，用于区分银河系不同星族。
|
||||
|
||||
**预估工作量**：中等（数据提取为主，分析逻辑相对直接）。
|
||||
|
||||
---
|
||||
|
||||
## 三、可视化与输出层
|
||||
|
||||
### 3.1 光谱交互式标注
|
||||
|
||||
**目标**：前端光谱图支持点击标注谱线名称，自动显示等值宽度。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 扩展 `SpectrumPlot.tsx`，新增谱线标注叠加层
|
||||
- 内置常见谱线波长表（Balmer 系列、H&K、NaD、Mg b、Fe 等）
|
||||
- 支持用户点击光谱图任意位置添加自定义标注
|
||||
- 标注信息（谱线名、波长、EW）以 tooltip 形式展示
|
||||
- 后端 `preview.rs` 返回时可选附带谱线标注数据
|
||||
|
||||
**价值**：显著提升光谱阅读和教学体验。
|
||||
|
||||
**预估工作量**：中等（前端组件开发为主）。
|
||||
|
||||
### 3.2 赫罗图（HR Diagram）
|
||||
|
||||
**目标**：基于 Gaia 数据构建颜色-星等图（CMD），支持叠加自定义样本。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/api/hertzsprung_russell.rs` 或扩展现有 Gaia 查询
|
||||
- 自动查询给定天区/样本的 Gaia G-mag、BP-RP color、parallax
|
||||
- 计算绝对星等 M_G = G + 5 + 5*log10(parallax/1000)
|
||||
- 返回数据供前端渲染 HR 图
|
||||
- 前端新增 HR Diagram 组件（scatter plot，支持密度着色）
|
||||
- 支持叠加用户自定义样本（不同颜色/标记）
|
||||
|
||||
**价值**：恒星物理的基础可视化，几乎所有恒星相关论文都包含 HR 图。
|
||||
|
||||
**预估工作量**：中等（数据查询简单，前端绘图需要新组件）。
|
||||
|
||||
### 3.3 发现图（Finding Chart）
|
||||
|
||||
**目标**：生成目标天区的 DSS/Pan-STARRS 图像 cutout + 标注目标位置。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 调用 STScI DSS cutout API 或 Pan-STARRS cutout API
|
||||
- 返回 cutout 图像 + 目标中心标注
|
||||
- 支持自定义 cutout 大小（arcmin）和缩放
|
||||
- Agent 工具 `find_observation` 可选返回 finding chart
|
||||
- 前端在观测数据页面展示 finding chart
|
||||
|
||||
**价值**：观测准备的必备环节，望远镜观测提案中需要附带 finding chart。
|
||||
|
||||
**预估工作量**：小（API 简单，主要是前端展示）。
|
||||
|
||||
### 3.4 多波段 SED 图
|
||||
|
||||
**目标**：自动拉取所有可用波段测光数据，绘制 SED 图并拟合模型。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 扩展 2.6 节 SED 拟合功能，新增前端 SED 图组件
|
||||
- 自动标记各波段（2MASS J/H/K、WISE W1-W4、Gaia G/BP/RP 等）
|
||||
- 支持对数坐标轴
|
||||
- 拟合曲线叠加显示，残差面板
|
||||
|
||||
**价值**：参数化分析的标准可视化。
|
||||
|
||||
**预估工作量**：中等（依赖 2.6 节后端实现）。
|
||||
|
||||
### 3.5 SkyMap 天球投影
|
||||
|
||||
**目标**：在全天球图上可视化样本分布，支持按参数着色。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `SkyMapPanel` 前端组件
|
||||
- 使用 Mollweide 或 Aitoff 投影
|
||||
- 支持按天体类型、金属丰度、自行等参数着色
|
||||
- 叠加银河平面、银道坐标网格
|
||||
- 支持交互式缩放和选区
|
||||
|
||||
**价值**：巡天论文的标准可视化，用于展示样本的天球分布。
|
||||
|
||||
**预估工作量**：中等（前端投影计算 + WebGL 渲染）。
|
||||
|
||||
---
|
||||
|
||||
## 四、工作流增强
|
||||
|
||||
### 4.1 批量样本处理
|
||||
|
||||
**目标**：给定一组天体名称/坐标，批量执行"查询 → 下载 → 分析 → 出图 → 保存报告"。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/agent/tools/astro/research/batch_process.rs`，实现 `BatchProcessTool`
|
||||
- 支持输入格式：天体名称列表、坐标列表、VizieR 查询结果
|
||||
- 批量流程：query_target → find_observation → 参数提取 → 出图
|
||||
- 并发控制（避免触发数据源速率限制）
|
||||
- 进度报告（SSE 事件推送）
|
||||
- 输出汇总报告（Markdown 表格 + 图表）
|
||||
|
||||
**价值**：巡天论文的标准工作流，手动逐个处理效率极低。
|
||||
|
||||
**预估工作量**：中等偏高（并发控制和错误处理需要仔细设计）。
|
||||
|
||||
### 4.2 定时文献监控
|
||||
|
||||
**目标**：Cron 定时检查 ADS/arXiv 新文献，匹配用户关注的关键词/作者/天体。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/agent/tools/astro/research/literature_monitor.rs`
|
||||
- 用户定义监控规则（关键词、作者、天体名称、arXiv 分类）
|
||||
- 使用现有 `loop` 技能定时触发
|
||||
- 新文献匹配时通过 SSE 推送通知
|
||||
- 支持自动下载摘要 + 生成简报
|
||||
|
||||
**价值**：帮助研究者跟踪前沿，不遗漏重要文献。
|
||||
|
||||
**预估工作量**：中等（定时任务框架已有，主要是规则匹配和通知）。
|
||||
|
||||
### 4.3 观测提案辅助
|
||||
|
||||
**目标**：根据目标天区自动计算可见性、月相限制、曝光时间估算。
|
||||
|
||||
**需求**：
|
||||
|
||||
- 新增 `src/services/analysis/observability.rs`
|
||||
- 计算目标在给定台站的升起/落下时间和最佳观测窗口
|
||||
- 月相和银河背景亮度估算
|
||||
- 基于目标亮度和望远镜参数的曝光时间估算
|
||||
- Agent 新增 `plan_observation` 工具
|
||||
- 输出观测提案所需的标准化信息
|
||||
|
||||
**价值**：望远镜申请的时间评估和观测规划必备。
|
||||
|
||||
**预估工作量**：中等（天文计算公式成熟，但台站参数和大气消光模型需要数据）。
|
||||
|
||||
---
|
||||
|
||||
## 五、实施优先级建议
|
||||
|
||||
| 优先级 | 功能 | 理由 | 状态 |
|
||||
| ------ | --------------------------- | -------------------------------------- | ---- |
|
||||
| P0 | 光谱参数自动提取（2.1） | 已下载数据的直接利用，所有恒星论文都用 | ✅ 已实现 |
|
||||
| P0 | 光变周期搜索（2.4） | ZTF/TESS 数据的刚需，时域天文学核心 | ✅ 已实现 |
|
||||
| P1 | 光谱线识别与 EW 测量（2.2） | 光谱分析的日常操作 | ✅ 已实现 |
|
||||
| P1 | 交叉相关测速（2.3） | 双星/视向速度研究核心工具 | ✅ 已实现 |
|
||||
| P1 | 绘图 skill 填充（现有占位） | 提升 Agent 绘图能力的基础 | ✅ 已实现 |
|
||||
| P2 | SED 拟合（2.6） | 多波段分析的标准方法 | ✅ 已实现 |
|
||||
| P2 | 赫罗图（3.2） | 恒星物理基础可视化 | ✅ 已实现 |
|
||||
| P2 | 批量样本处理（4.1） | 巡天论文工作流提效 | ✅ 已实现 |
|
||||
| P3 | JWST 数据获取（1.1） | 前沿望远镜，但使用门槛较高 | ✅ 已实现 |
|
||||
| P3 | 光变分类（2.5） | 时域数据分析，但分类器训练复杂 | ✅ 已实现 |
|
||||
| P3 | SkyMap 天球投影（3.5） | 巡天论文可视化，但优先级不如核心分析 | ✅ 已实现 |
|
||||
|
||||
---
|
||||
|
||||
## 六、已实现功能详情
|
||||
|
||||
### 已完成的数据分析层
|
||||
|
||||
| 功能 | Agent 工具 | 后端模块 | 状态 |
|
||||
|------|-----------|----------|------|
|
||||
| 光谱参数提取 | `get_spectrum_parameters` | `services/spectrum/parameters.rs` | ✅ |
|
||||
| 谱线识别与EW测量 | `measure_spectral_lines` | `services/spectrum/lines.rs` | ✅ |
|
||||
| 交叉相关测速 | `measure_radial_velocity` | `services/spectrum/cross_correlate.rs` | ✅ |
|
||||
| 周期搜索 | `find_period` | `services/timeseries/periodicity.rs` | ✅ |
|
||||
| 变星分类 | `classify_variable_star` | `services/timeseries/classification.rs` | ✅ |
|
||||
| SED拟合 | `fit_sed` | `services/analysis/sed.rs` | ✅ |
|
||||
| 运动学分析 | `analyze_kinematics` | `services/analysis/kinematics.rs` | ✅ |
|
||||
| 化学丰度分析 | `analyze_abundance_pattern` | `services/analysis/chemistry.rs` | ✅ |
|
||||
| 观测提案辅助 | `plan_observation` | `services/analysis/observability.rs` | ✅ |
|
||||
| 批量样本处理 | `batch_process` | `agent/tools/astro/research/batch_process.rs` | ✅ |
|
||||
| 文献监控 | `literature_monitor` | `agent/tools/astro/research/literature_monitor.rs` | ✅ |
|
||||
|
||||
### 已完成的数据获取层
|
||||
|
||||
| 功能 | 模块 | 说明 | 状态 |
|
||||
|------|------|------|------|
|
||||
| JWST 数据获取 | `clients/mast/jwst.rs` + `services/observation/jwst.rs` | MAST JWST API + SpectrumFetcher | ✅ |
|
||||
| X射线数据 | `clients/heasarc/` + `services/observation/xray.rs` | HEASARC TAP + XMM/Chandra fetcher | ✅ |
|
||||
| 光谱图像Cutout | `services/observation/cutout.rs` | SDSS/DSS/Pan-STARRS cutout API + CutoutFetcher | ✅ |
|
||||
| 赫罗图 | `api/hr_diagram.rs` | Gaia CMD 查询 API | ✅ |
|
||||
| 绘图skill | `skills/plotting/SKILL.md` | Python matplotlib 科研绘图模板（8种图表） | ✅ |
|
||||
|
||||
### 已完成的工作流增强
|
||||
|
||||
| 功能 | 模块 | 说明 | 状态 |
|
||||
|------|------|------|------|
|
||||
| 批量样本处理 | `agent/tools/astro/research/batch_process.rs` | 查询→下载→分析→报告 | ✅ |
|
||||
| 文献监控 | `agent/tools/astro/research/literature_monitor.rs` | ADS/arXiv 最新文献匹配 | ✅ |
|
||||
@@ -0,0 +1,2 @@
|
||||
-- Add is_scheduled column to sync_queries for periodic incremental sync scheduling
|
||||
ALTER TABLE sync_queries ADD COLUMN is_scheduled INTEGER NOT NULL DEFAULT 0;
|
||||
@@ -0,0 +1,25 @@
|
||||
-- 会话事件日志：turn/compaction 生命周期 + 上下文快照（事件溯源中枢）
|
||||
--
|
||||
-- 设计（参考 deepseek-harness 的事件溯源会话模型）：
|
||||
-- agent_messages 记录"模型可见内容"，agent_events 记录"生命周期事实"。
|
||||
-- turn_start/turn_end 构成日志化锁：崩溃后开着的 turn_start 可被检测并
|
||||
-- 合成 interrupted 关闭（不截断已持久化的部分）。
|
||||
-- context_snapshot 存储压缩后的折叠上下文，下一 turn 直接回放而无需
|
||||
-- 从原始消息重建再重新压缩（节省重复的 LLM 摘要调用）。
|
||||
--
|
||||
-- 事件类型（event_type）：
|
||||
-- turn_start / turn_end / turn_interrupted
|
||||
-- compaction_start / compaction_end / compaction_interrupted
|
||||
-- context_snapshot
|
||||
|
||||
CREATE TABLE IF NOT EXISTS agent_events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL,
|
||||
turn_index INTEGER NOT NULL DEFAULT 0,
|
||||
event_type TEXT NOT NULL,
|
||||
payload TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_events_session
|
||||
ON agent_events(session_id, id);
|
||||
+286
-21
@@ -1,36 +1,301 @@
|
||||
---
|
||||
name: plotting
|
||||
description: 科研绘图规范
|
||||
description: 科研绘图规范 —— 天文学常用图表的 Python matplotlib 绘制模板
|
||||
context: fork
|
||||
allowed-tools:
|
||||
- bash
|
||||
- save_note
|
||||
- write
|
||||
- read
|
||||
---
|
||||
|
||||
# 科研绘图规范（TODO：待细化）
|
||||
# 科研绘图规范
|
||||
|
||||
## 目标期刊要求（占位）
|
||||
## 环境要求
|
||||
|
||||
- 分辨率：通常 300-600 dpi
|
||||
- 格式：矢量图优先（PDF/SVG），光栅图备选（PNG）
|
||||
- 字体：需与期刊正文一致（如 Times New Roman）
|
||||
- Python 3.8+
|
||||
- matplotlib >= 3.7
|
||||
- numpy
|
||||
- 可选：seaborn（统计图）、astropy（天文学单位和坐标）
|
||||
|
||||
## 常用图表类型（占位）
|
||||
安装：
|
||||
```bash
|
||||
pip install matplotlib numpy seaborn astropy --quiet
|
||||
```
|
||||
|
||||
| 类型 | Python 库 | 适用场景 |
|
||||
|------|----------|---------|
|
||||
| 光谱图 | matplotlib | 光谱分析、能谱 |
|
||||
| 赫罗图 | matplotlib | 恒星演化 |
|
||||
| 光变曲线 | matplotlib | 变星、超新星 |
|
||||
| 参数分布 | seaborn | 统计分布、角图 |
|
||||
| 3D 轨道 | plotly | 天体力学 |
|
||||
## 输出规范
|
||||
|
||||
## 工作流程（TODO）
|
||||
| 目标 | 分辨率 | 格式 | 说明 |
|
||||
|------|--------|------|------|
|
||||
| 论文投稿 | 300+ dpi | PDF/SVG（矢量优先） | APS/A&A/MNRAS 标准 |
|
||||
| 演示文稿 | 150 dpi | PNG | 宽度 ≥ 1200px |
|
||||
| 快速预览 | 100 dpi | PNG | 屏幕查看 |
|
||||
|
||||
1. 准备数据（从文献或天体数据库提取）
|
||||
2. 生成 Python 脚本
|
||||
3. 使用 bash 工具执行脚本
|
||||
4. 检查输出图片
|
||||
5. 保存到指定路径
|
||||
## 通用样式模板
|
||||
|
||||
**注意：此 skill 为预留占位，具体实现待补充。**
|
||||
```python
|
||||
import matplotlib
|
||||
matplotlib.use('Agg')
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
# 天文学论文标准样式
|
||||
plt.rcParams.update({
|
||||
'font.family': 'serif',
|
||||
'font.size': 12,
|
||||
'axes.labelsize': 14,
|
||||
'axes.titlesize': 14,
|
||||
'xtick.labelsize': 11,
|
||||
'ytick.labelsize': 11,
|
||||
'legend.fontsize': 11,
|
||||
'figure.figsize': (8, 6),
|
||||
'figure.dpi': 300,
|
||||
'savefig.dpi': 300,
|
||||
'savefig.bbox': 'tight',
|
||||
'savefig.pad_inches': 0.05,
|
||||
'lines.linewidth': 1.5,
|
||||
'axes.linewidth': 1.0,
|
||||
'xtick.major.width': 0.8,
|
||||
'ytick.major.width': 0.8,
|
||||
'xtick.minor.width': 0.5,
|
||||
'ytick.minor.width': 0.5,
|
||||
'xtick.direction': 'in',
|
||||
'ytick.direction': 'in',
|
||||
'xtick.top': True,
|
||||
'ytick.right': True,
|
||||
})
|
||||
```
|
||||
|
||||
## 图表模板
|
||||
|
||||
### 1. 光谱图
|
||||
|
||||
```python
|
||||
def plot_spectrum(wavelength, flux, title="Spectrum", xlabel=r"Wavelength ($\AA$)", ylabel=r"Flux (erg/s/cm$^2$/$\AA$)", save_path=None):
|
||||
fig, ax = plt.subplots()
|
||||
ax.plot(wavelength, flux, 'k-', linewidth=0.8)
|
||||
ax.set_xlabel(xlabel)
|
||||
ax.set_ylabel(ylabel)
|
||||
ax.set_title(title)
|
||||
ax.minorticks_on()
|
||||
if save_path:
|
||||
fig.savefig(save_path)
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
### 2. 光变曲线
|
||||
|
||||
```python
|
||||
def plot_light_curve(time, flux, time_err=None, flux_err=None, title="Light Curve", xlabel="Time (MJD)", ylabel="Flux", save_path=None):
|
||||
fig, ax = plt.subplots()
|
||||
if flux_err is not None:
|
||||
ax.errorbar(time, flux, yerr=flux_err, xerr=time_err, fmt='o', markersize=3, capsize=2, color='black', ecolor='gray')
|
||||
else:
|
||||
ax.plot(time, flux, 'ko', markersize=3)
|
||||
ax.set_xlabel(xlabel)
|
||||
ax.set_ylabel(ylabel)
|
||||
ax.set_title(title)
|
||||
ax.minorticks_on()
|
||||
if save_path:
|
||||
fig.savefig(save_path)
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
### 3. 折叠光变曲线
|
||||
|
||||
```python
|
||||
def plot_folded_lc(phase, flux, flux_err=None, period=None, title="Folded Light Curve", save_path=None):
|
||||
fig, ax = plt.subplots()
|
||||
if flux_err is not None:
|
||||
ax.errorbar(phase, flux, yerr=flux_err, fmt='o', markersize=2, capsize=1, color='black', ecolor='gray', alpha=0.7)
|
||||
else:
|
||||
ax.plot(phase, flux, 'ko', markersize=2, alpha=0.7)
|
||||
ax.set_xlabel("Phase")
|
||||
ax.set_ylabel("Flux")
|
||||
if period:
|
||||
ax.set_title(f"{title} (P = {period:.4f} d)")
|
||||
else:
|
||||
ax.set_title(title)
|
||||
ax.set_xlim(0, 1)
|
||||
ax.minorticks_on()
|
||||
if save_path:
|
||||
fig.savefig(save_path)
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
### 4. 赫罗图（CMD）
|
||||
|
||||
```python
|
||||
def plot_hr_diagram(bp_rp, abs_g, title="HR Diagram", save_path=None, color_by_density=False):
|
||||
fig, ax = plt.subplots()
|
||||
if color_by_density and len(bp_rp) > 100:
|
||||
from scipy.stats import gaussian_kde
|
||||
xy = np.vstack([bp_rp, abs_g])
|
||||
z = gaussian_kde(xy)(xy)
|
||||
idx = z.argsort()
|
||||
ax.scatter(bp_rp[idx], abs_g[idx], c=z[idx], s=5, cmap='viridis_r', edgecolors='none')
|
||||
cb = fig.colorbar(ax.collections[0], ax=ax, pad=0.02)
|
||||
cb.set_label("Stellar density")
|
||||
else:
|
||||
ax.plot(bp_rp, abs_g, 'k.', markersize=1, alpha=0.5)
|
||||
ax.set_xlabel(r"$G_{BP} - G_{RP}$ (mag)")
|
||||
ax.set_ylabel(r"$M_G$ (mag)")
|
||||
ax.set_title(title)
|
||||
ax.invert_yaxis()
|
||||
ax.minorticks_on()
|
||||
if save_path:
|
||||
fig.savefig(save_path)
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
### 5. SED 图
|
||||
|
||||
```python
|
||||
def plot_sed(wavelength_angstrom, flux_obs, flux_err=None, flux_model=None, title="SED", save_path=None):
|
||||
fig, ax = plt.subplots()
|
||||
ax.scatter(wavelength_angstrom, flux_obs, c='black', s=30, zorder=5, label='Observed')
|
||||
if flux_err is not None:
|
||||
ax.errorbar(wavelength_angstrom, flux_obs, yerr=flux_err, fmt='none', ecolor='gray', capsize=3)
|
||||
if flux_model is not None:
|
||||
model_wave, model_flux = zip(*flux_model) if isinstance(flux_model, list) else (flux_model[0], flux_model[1])
|
||||
ax.plot(model_wave, model_flux, 'r-', linewidth=1.5, label='Model', alpha=0.8)
|
||||
ax.set_xscale('log')
|
||||
ax.set_yscale('log')
|
||||
ax.set_xlabel(r"Wavelength ($\AA$)")
|
||||
ax.set_ylabel(r"Flux")
|
||||
ax.set_title(title)
|
||||
ax.legend()
|
||||
ax.minorticks_on()
|
||||
if save_path:
|
||||
fig.savefig(save_path)
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
### 6. [α/Fe] vs [Fe/H] 图
|
||||
|
||||
```python
|
||||
def plot_abundance(feh, alpha_feh, labels=None, title="[α/Fe] vs [Fe/H]", save_path=None):
|
||||
fig, ax = plt.subplots()
|
||||
if labels is not None:
|
||||
from matplotlib.colors import ListedColormap
|
||||
colors = ['#3498db', '#e74c3c', '#2ecc71']
|
||||
cmap = ListedColormap(colors[:len(set(labels))])
|
||||
unique = sorted(set(labels))
|
||||
for i, lab in enumerate(unique):
|
||||
mask = np.array(labels) == lab
|
||||
ax.scatter(np.array(feh)[mask], np.array(alpha_feh)[mask], s=10, alpha=0.7, label=lab, color=colors[i % len(colors)])
|
||||
ax.legend()
|
||||
else:
|
||||
ax.scatter(feh, alpha_feh, s=10, alpha=0.7, c='black')
|
||||
ax.set_xlabel("[Fe/H] (dex)")
|
||||
ax.set_ylabel(r"[$\alpha$/Fe] (dex)")
|
||||
ax.set_title(title)
|
||||
ax.axhline(y=0.25, color='gray', linestyle='--', linewidth=0.8, alpha=0.5)
|
||||
ax.minorticks_on()
|
||||
if save_path:
|
||||
fig.savefig(save_path)
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
### 7. 周期图
|
||||
|
||||
```python
|
||||
def plot_periodogram(periods, powers, best_period=None, title="Periodogram", save_path=None):
|
||||
fig, ax = plt.subplots()
|
||||
ax.plot(periods, powers, 'k-', linewidth=0.8)
|
||||
if best_period:
|
||||
ax.axvline(x=best_period, color='red', linestyle='--', linewidth=1, label=f'Best P = {best_period:.4f} d')
|
||||
ax.legend()
|
||||
ax.set_xlabel("Period (days)")
|
||||
ax.set_ylabel("Power")
|
||||
ax.set_title(title)
|
||||
ax.minorticks_on()
|
||||
if save_path:
|
||||
fig.savefig(save_path)
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
### 8. 天球投影（Mollweide）
|
||||
|
||||
```python
|
||||
def plot_skymap(ra_deg, dec_deg, values=None, title="Sky Map", save_path=None):
|
||||
fig, ax = plt.subplots(figsize=(10, 5), subplot_kw={'projection': 'mollweide'})
|
||||
# 转换为弧度，RA 中心在 180°
|
||||
ra_rad = np.deg2rad(np.array(ra_deg))
|
||||
dec_rad = np.deg2rad(np.array(dec_deg))
|
||||
ra_rad = ra_rad - np.pi # 中心化
|
||||
if values is not None:
|
||||
sc = ax.scatter(ra_rad, dec_rad, c=values, s=5, cmap='viridis', alpha=0.7, edgecolors='none')
|
||||
fig.colorbar(sc, ax=ax, pad=0.05, shrink=0.6)
|
||||
else:
|
||||
ax.scatter(ra_rad, dec_rad, c='black', s=5, alpha=0.5, edgecolors='none')
|
||||
ax.set_title(title)
|
||||
ax.grid(True, alpha=0.3)
|
||||
if save_path:
|
||||
fig.savefig(save_path, bbox_inches='tight')
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
### 9. 发现图（Finding Chart）
|
||||
|
||||
```python
|
||||
def plot_finding_chart(image_data, ra, dec, label=None, save_path=None):
|
||||
"""image_data: 2D numpy array (from FITS or JPEG)"""
|
||||
fig, ax = plt.subplots()
|
||||
ax.imshow(image_data, cmap='gray_r', origin='lower')
|
||||
# 标注中心
|
||||
h, w = image_data.shape
|
||||
ax.plot(w/2, h/2, 'r+', markersize=15, markeredgewidth=2)
|
||||
if label:
|
||||
ax.annotate(label, (w/2, h/2), textcoords="offset points", xytext=(10, 10),
|
||||
color='red', fontsize=10, fontweight='bold')
|
||||
ax.set_title(f"Finding Chart ({ra:.4f}, {dec:.4f})")
|
||||
ax.set_xlabel("X (pix)")
|
||||
ax.set_ylabel("Y (pix)")
|
||||
if save_path:
|
||||
fig.savefig(save_path)
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
return fig, ax
|
||||
```
|
||||
|
||||
## 工作流程
|
||||
|
||||
1. **确定图表类型**：根据数据特征选择合适的图表模板
|
||||
2. **准备数据**：从文献、数据库或分析结果中提取数据
|
||||
3. **生成脚本**：使用上述模板，填入实际数据
|
||||
4. **执行**：`python3 script.py`
|
||||
5. **验证**：检查输出图片是否正确
|
||||
6. **保存**：保存到 `figures/` 目录，命名格式：`{description}_{source}.{ext}`
|
||||
|
||||
## 保存路径约定
|
||||
|
||||
- `figures/` — 项目根目录下的图表输出目录
|
||||
- 文件名格式：`{描述}_{来源}_{日期}.{格式}`
|
||||
- 示例：`spectrum_lamost_438809089_20260710.pdf`
|
||||
- 示例：`hr_diagram_m31_field_20260710.png`
|
||||
|
||||
+11
-2
@@ -109,9 +109,18 @@ impl IdlePoller {
|
||||
.await
|
||||
{
|
||||
autonomous_turns += 1;
|
||||
// 创建 AgentRuntime 并执行任务
|
||||
let runtime = crate::agent::runtime::AgentRuntime::new(
|
||||
// 通过会话注册表获取 runtime（复用该会话的
|
||||
// 后台队列/压缩日志等跨 turn 状态；模式从 DB 回放）
|
||||
let mode = crate::agent::runtime::session::load_session_mode(
|
||||
&self.app_state.db,
|
||||
&task.session_id,
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|| "default".to_string());
|
||||
let runtime = self.app_state.agent_runtimes.get_or_create(
|
||||
self.app_state.clone(),
|
||||
&task.session_id,
|
||||
&mode,
|
||||
);
|
||||
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let _ = runtime
|
||||
|
||||
+111
-12
@@ -17,14 +17,13 @@
|
||||
|
||||
pub mod collapse;
|
||||
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::hooks::{HookRegistry, PostCompactContext, PreCompactContext};
|
||||
use collapse::{CollapseLog, CollapseMethod};
|
||||
|
||||
/// 递归守卫：防止压缩内部触发的 LLM 调用再次触发压缩。
|
||||
static COMPACTING: AtomicBool = AtomicBool::new(false);
|
||||
// 递归守卫已迁移至 session_registry::try_begin_compaction/end_compaction：
|
||||
// 按 session_id 隔离，并发会话压缩互不干扰（历史上是进程级 AtomicBool）。
|
||||
|
||||
use crate::clients::llm::{ChatMessage, LlmClient, MessageRole};
|
||||
|
||||
@@ -182,11 +181,16 @@ fn find_tool_name_for_call_id(messages: &[ChatMessage], tool_call_id: &str) -> O
|
||||
None
|
||||
}
|
||||
|
||||
/// 轻量级压缩：将较早的工具结果替换为简短占位符，释放上下文空间。
|
||||
/// 保留最近 `keep_recent` 条工具结果不变。
|
||||
/// 轻量级压缩（确定性剪枝，零模型调用）：将较早的工具结果替换为简短
|
||||
/// 占位符，释放上下文空间。保留最近 `keep_recent` 条工具结果不变。
|
||||
///
|
||||
/// 剪枝策略（参考 deepseek-harness tool-result-pruner 的 head/tail 保留）：
|
||||
/// - 短结果（≤ `PRUNE_PREVIEW_THRESHOLD` 字符）：`[Previous: used {tool}]`
|
||||
/// 纯占位符（节省 ~80 tokens/条）
|
||||
/// - 长结果：占位符 + 头部/尾部预览，为模型保留检索锚点（关键数值、
|
||||
/// 结论行常在结果尾部）
|
||||
///
|
||||
/// 不可变风格：接收引用，返回新的消息列表。
|
||||
/// 使用 `[Previous: used {tool_name}]` 替代字符预览，节省 ~80 tokens/条。
|
||||
pub fn micro_compact(messages: &[ChatMessage], keep_recent: usize) -> Vec<ChatMessage> {
|
||||
// 收集所有工具消息的索引、call_id 和工具名
|
||||
let tool_positions: Vec<(usize, &str)> = messages
|
||||
@@ -231,7 +235,7 @@ pub fn micro_compact(messages: &[ChatMessage], keep_recent: usize) -> Vec<ChatMe
|
||||
if compact_indices.contains(&i) {
|
||||
let tool_name = tool_name_map.get(&i).cloned().unwrap_or_default();
|
||||
let mut new_msg = msg.clone();
|
||||
new_msg.content = Some(format!("[Previous: used {}]", tool_name));
|
||||
new_msg.content = Some(prune_tool_result(&tool_name, msg.content.as_deref()));
|
||||
new_msg
|
||||
} else {
|
||||
msg.clone()
|
||||
@@ -240,6 +244,33 @@ pub fn micro_compact(messages: &[ChatMessage], keep_recent: usize) -> Vec<ChatMe
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 确定性剪枝阈值：超过此长度的工具结果在占位符中保留头/尾预览
|
||||
pub const PRUNE_PREVIEW_THRESHOLD: usize = 200;
|
||||
/// 长结果保留的头部预览长度（字符）
|
||||
const PRUNE_HEAD_CHARS: usize = 160;
|
||||
/// 长结果保留的尾部预览长度（字符）
|
||||
const PRUNE_TAIL_CHARS: usize = 80;
|
||||
|
||||
/// 单条工具结果的确定性剪枝。
|
||||
pub fn prune_tool_result(tool_name: &str, content: Option<&str>) -> String {
|
||||
let Some(text) = content else {
|
||||
return format!("[Previous: used {}]", tool_name);
|
||||
};
|
||||
let char_count = text.chars().count();
|
||||
if char_count <= PRUNE_PREVIEW_THRESHOLD {
|
||||
return format!("[Previous: used {}]", tool_name);
|
||||
}
|
||||
let head: String = text.chars().take(PRUNE_HEAD_CHARS).collect();
|
||||
let tail: String = text
|
||||
.chars()
|
||||
.skip(char_count.saturating_sub(PRUNE_TAIL_CHARS))
|
||||
.collect();
|
||||
format!(
|
||||
"[Previous: used {}] ({} 字符，已剪枝)\n头: {}…\n尾: …{}",
|
||||
tool_name, char_count, head, tail
|
||||
)
|
||||
}
|
||||
|
||||
/// 粗略估算消息列表的 token 数（用作首次调用的近似值）。
|
||||
/// 后续迭代优先使用 API 返回的精确 prompt_tokens。
|
||||
///
|
||||
@@ -256,6 +287,12 @@ pub fn rough_estimate_tokens(messages: &[ChatMessage]) -> usize {
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// 估算单条消息的 token 数（所有循环统一使用本函数，
|
||||
/// 消除主循环 `len/3` 与 compact 的 CJK 加权两套口径不一致的问题）。
|
||||
pub fn estimate_message_tokens(message: &ChatMessage) -> usize {
|
||||
message.text().map_or(0, estimate_tokens_from_text) + 4
|
||||
}
|
||||
|
||||
/// 根据文本内容估算 token 数。
|
||||
fn estimate_tokens_from_text(text: &str) -> usize {
|
||||
let mut ascii_chars = 0usize;
|
||||
@@ -513,6 +550,7 @@ pub async fn compress_context_with_hooks(
|
||||
session_id,
|
||||
hook_registry,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -521,11 +559,15 @@ pub async fn compress_context_with_hooks(
|
||||
/// 如果不预留，当上下文接近窗口上限时压缩调用本身会因输出空间不足而失败。
|
||||
pub const COMPACTION_OUTPUT_RESERVE: usize = 4000;
|
||||
|
||||
/// 带 Hook 和 CollapseLog 的上下文压缩变体。
|
||||
/// 带 Hook、CollapseLog 和会话事件日志的上下文压缩变体。
|
||||
/// 压缩后自动记录 commit 到 CollapseLog，并在溢出时注入合并摘要。
|
||||
///
|
||||
/// `db` 提供时写 compaction_start/compaction_end 事件（日志化锁：
|
||||
/// 崩溃后开口的 compaction_start 由恢复时的 reconcile 合成关闭）。
|
||||
///
|
||||
/// `token_soft_limit`: TokenBudget 软限制，统一各压缩层的触发阈值。
|
||||
/// `max_messages`: 消息数上限，用于 Layer 0 snip 触发。
|
||||
#[allow(clippy::too_many_arguments)] // 压缩上下文 + 可选钩子/日志/DB 依赖
|
||||
pub async fn compress_context_with_hooks_and_log(
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
llm: &LlmClient,
|
||||
@@ -534,20 +576,39 @@ pub async fn compress_context_with_hooks_and_log(
|
||||
session_id: &str,
|
||||
hook_registry: Option<&HookRegistry>,
|
||||
collapse_log: Option<&std::sync::Mutex<CollapseLog>>,
|
||||
db: Option<&sqlx::SqlitePool>,
|
||||
) {
|
||||
if messages.len() <= 4 {
|
||||
return;
|
||||
}
|
||||
|
||||
// 递归守卫：如果已在压缩中，跳过（防止嵌套压缩死循环）
|
||||
if COMPACTING.swap(true, Ordering::SeqCst) {
|
||||
warn!("[Compact] 递归守卫触发：已有进行中的压缩操作，跳过");
|
||||
// 递归守卫（按会话）：如果该会话已在压缩中，跳过（防止嵌套压缩死循环）
|
||||
if !super::runtime::session_registry::try_begin_compaction(session_id) {
|
||||
warn!(
|
||||
"[Compact] 递归守卫触发：会话 {} 已有进行中的压缩操作，跳过",
|
||||
session_id
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let before_count = messages.len();
|
||||
let est_tokens = rough_estimate_tokens(messages);
|
||||
|
||||
// ── 日志化锁：compaction_start（先落盘再执行；崩溃留下可检测开口）──
|
||||
if let Some(db) = db {
|
||||
super::runtime::session_events::append_event(
|
||||
db,
|
||||
session_id,
|
||||
-1,
|
||||
super::runtime::session_events::event_types::COMPACTION_START,
|
||||
serde_json::json!({
|
||||
"before_messages": before_count,
|
||||
"estimated_tokens": est_tokens,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// OnPreCompact hook
|
||||
if let Some(registry) = hook_registry {
|
||||
registry
|
||||
@@ -623,6 +684,22 @@ pub async fn compress_context_with_hooks_and_log(
|
||||
}
|
||||
}
|
||||
|
||||
// ── 日志化锁：compaction_end ──
|
||||
if let Some(db) = db {
|
||||
super::runtime::session_events::append_event(
|
||||
db,
|
||||
session_id,
|
||||
-1,
|
||||
super::runtime::session_events::event_types::COMPACTION_END,
|
||||
serde_json::json!({
|
||||
"after_messages": after_count,
|
||||
"method": method,
|
||||
"saved_tokens": est_tokens.saturating_sub(rough_estimate_tokens(messages)),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// OnPostCompact hook
|
||||
if let Some(registry) = hook_registry {
|
||||
registry
|
||||
@@ -635,7 +712,7 @@ pub async fn compress_context_with_hooks_and_log(
|
||||
}
|
||||
|
||||
// 释放递归守卫
|
||||
COMPACTING.store(false, Ordering::SeqCst);
|
||||
super::runtime::session_registry::end_compaction(session_id);
|
||||
}
|
||||
|
||||
// ── Compaction → Memory Extraction Bridge (P3) ──
|
||||
@@ -907,6 +984,28 @@ mod tests {
|
||||
assert_eq!(name, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prune_tool_result_short_keeps_placeholder() {
|
||||
assert_eq!(
|
||||
prune_tool_result("search_papers", Some("short result")),
|
||||
"[Previous: used search_papers]"
|
||||
);
|
||||
assert_eq!(
|
||||
prune_tool_result("search_papers", None),
|
||||
"[Previous: used search_papers]"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prune_tool_result_long_keeps_head_tail() {
|
||||
let long = format!("HEAD_MARKER\n{}\nTAIL_MARKER", "x".repeat(500));
|
||||
let pruned = prune_tool_result("get_paper_content", Some(&long));
|
||||
assert!(pruned.contains("[Previous: used get_paper_content]"));
|
||||
assert!(pruned.contains("HEAD_MARKER"));
|
||||
assert!(pruned.contains("TAIL_MARKER"));
|
||||
assert!(pruned.contains("已剪枝"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rough_estimate_tokens() {
|
||||
let messages = vec![
|
||||
|
||||
+1224
File diff suppressed because it is too large
Load Diff
@@ -52,7 +52,7 @@ pub use builtins::{AuditLogHook, CancellationHook, ContextDeduplicator, MetricsH
|
||||
mod tests {
|
||||
use super::*;
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashSet;
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
|
||||
+4
-2
@@ -3,15 +3,17 @@
|
||||
// 科研智能体模块 — 基于 ReAct 框架实现 Thought → Action → Observation 循环。
|
||||
//
|
||||
// 模块结构（参考 Claude Code 分层设计）：
|
||||
// engine — 统一 ReAct 引擎（主代理/子代理/队友共用的唯一循环实现）
|
||||
// tools/ — 工具定义与注册（按功能域拆分）
|
||||
// runtime — ReAct 循环引擎 + Streaming + 会话管理
|
||||
// compact — 三层上下文压缩（micro/auto/manual）
|
||||
// runtime — 运行时编排：会话管理、权限、压缩熔断、系统提示词
|
||||
// compact — 多层上下文压缩（snip/micro/auto/manual）
|
||||
// terminal — 循环终止信号（结构化退出原因）
|
||||
// hooks — 生命周期事件系统（PreToolUse/PostToolUse/Stop）
|
||||
|
||||
pub mod autonomous;
|
||||
pub mod background;
|
||||
pub mod compact;
|
||||
pub mod engine;
|
||||
pub mod enums;
|
||||
pub mod hooks;
|
||||
pub mod memory;
|
||||
|
||||
@@ -70,8 +70,12 @@ const DEFAULT_EXCLUDES: &[&str] = &[
|
||||
"Thumbs.db",
|
||||
];
|
||||
|
||||
/// 每个 turn 最多快照一次的工具
|
||||
const CHECKPOINT_TRIGGER_TOOLS: &[&str] = &["file_write", "file_edit", "run_bash"];
|
||||
/// 每个 turn 最多快照一次的工具。
|
||||
///
|
||||
/// 历史上执行器按此名单触发快照；现在由各工具通过
|
||||
/// `AgentTool::causes_file_changes()` 声明，本名单仅作为
|
||||
/// 注册表查询不到工具时的保守回退。
|
||||
const CHECKPOINT_FALLBACK_TOOLS: &[&str] = &["file_write", "file_edit", "run_bash"];
|
||||
|
||||
/// Checkpoint 元数据
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -465,9 +469,9 @@ impl CheckpointManager {
|
||||
))
|
||||
}
|
||||
|
||||
/// 检查指定工具是否需要触发 checkpoint。
|
||||
/// 检查指定工具是否需要触发 checkpoint（注册表查询不到时的保守回退）。
|
||||
pub fn should_checkpoint(tool_name: &str) -> bool {
|
||||
CHECKPOINT_TRIGGER_TOOLS.contains(&tool_name)
|
||||
CHECKPOINT_FALLBACK_TOOLS.contains(&tool_name)
|
||||
}
|
||||
|
||||
/// 获取 repo 路径
|
||||
|
||||
@@ -2,6 +2,10 @@
|
||||
//
|
||||
// 上下文构建：加载历史消息、注入系统提示词、添加用户消息、
|
||||
// 从数据库恢复持久化的任务状态。
|
||||
//
|
||||
// 上下文快照回放：若上一 turn 结束时保存过压缩后的折叠上下文
|
||||
// （context_snapshot 事件），则加载"快照消息 + id > base 的增量消息"，
|
||||
// 避免从原始消息重建后再次触发 LLM 摘要压缩（重复付费且信息有损）。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::info;
|
||||
@@ -9,22 +13,51 @@ use tracing::info;
|
||||
use crate::clients::llm::{ChatMessage, MessageRole};
|
||||
|
||||
use super::session;
|
||||
use super::session_events;
|
||||
|
||||
/// 构建初始 LLM 上下文：加载历史 → 插入系统提示词 → 添加用户消息 → 恢复任务状态。
|
||||
/// 构建初始 LLM 上下文。
|
||||
///
|
||||
/// 加载顺序：上下文快照（若有）+ 增量历史 → 插入系统提示词 →
|
||||
/// 动态上下文快照（内容变更时才追加并持久化）→ 用户消息 → 任务恢复。
|
||||
pub async fn build_initial_context(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
system_prompt: &str,
|
||||
question: &str,
|
||||
_turn_index: i32,
|
||||
turn_index: i32,
|
||||
dynamic_context: Option<(String, u64)>,
|
||||
image_context: Option<&str>,
|
||||
) -> anyhow::Result<Vec<ChatMessage>> {
|
||||
let mut messages = session::load_history_for_llm(db, session_id).await?;
|
||||
let mut messages = load_folded_history(db, session_id).await?;
|
||||
|
||||
// 注入系统提示词（如果历史中没有）
|
||||
if messages.is_empty() || messages[0].role != MessageRole::System {
|
||||
messages.insert(0, ChatMessage::system(system_prompt));
|
||||
}
|
||||
|
||||
// 动态上下文快照（durable user-role 消息）：哈希与上次持久化的一致时跳过。
|
||||
// 这是 KV-cache 纪律的另一半——易变内容不进 system prompt（否则任一
|
||||
// 变化都使整条前缀缓存失效），而以追加式快照进入历史，字节一旦写入
|
||||
// 就永不变更。
|
||||
if let Some((text, hash)) = dynamic_context {
|
||||
let last_hash = last_dynamic_context_hash(db, session_id).await;
|
||||
if last_hash.as_deref() != Some(hash.to_string().as_str()) {
|
||||
let reminder =
|
||||
ChatMessage::user(format!("<system-reminder>\n{}\n</system-reminder>", text));
|
||||
// 先持久化再注入（模型可见 ⟺ 已日志化）
|
||||
persist_dynamic_context(db, session_id, turn_index, &reminder, hash).await;
|
||||
messages.push(reminder);
|
||||
}
|
||||
}
|
||||
|
||||
// 图片上下文：以 system-reminder 形式注入在用户问题之前
|
||||
if let Some(img_ctx) = image_context {
|
||||
messages.push(ChatMessage::user(format!(
|
||||
"<system-reminder>\n{}\n</system-reminder>",
|
||||
img_ctx
|
||||
)));
|
||||
}
|
||||
|
||||
// 添加用户消息
|
||||
messages.push(ChatMessage::user(question));
|
||||
|
||||
@@ -36,6 +69,29 @@ pub async fn build_initial_context(
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
/// 加载折叠后的历史消息：优先回放上下文快照 + 增量；无快照时全量加载。
|
||||
async fn load_folded_history(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<Vec<ChatMessage>> {
|
||||
if let Some((snapshot_messages, base_id)) =
|
||||
session_events::load_context_snapshot(db, session_id).await
|
||||
{
|
||||
let incremental = session_events::load_messages_after(db, session_id, base_id).await?;
|
||||
info!(
|
||||
"[Context] 回放上下文快照: {} 条快照消息 + {} 条增量消息 (base_id={})",
|
||||
snapshot_messages.len(),
|
||||
incremental.len(),
|
||||
base_id
|
||||
);
|
||||
let mut messages = snapshot_messages;
|
||||
messages.extend(incremental);
|
||||
Ok(messages)
|
||||
} else {
|
||||
session::load_history_for_llm(db, session_id).await
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 agent_tasks 表恢复任务状态，返回格式化的提醒文本。
|
||||
///
|
||||
/// 如果表不存在或没有任务记录，返回 None。
|
||||
@@ -82,3 +138,44 @@ async fn restore_tasks_from_db(db: &SqlitePool, session_id: &str) -> Option<Stri
|
||||
lines.join("\n")
|
||||
))
|
||||
}
|
||||
|
||||
/// 查询会话中最近一次动态上下文快照的哈希（无则 None）
|
||||
async fn last_dynamic_context_hash(db: &SqlitePool, session_id: &str) -> Option<String> {
|
||||
let metadata: Option<String> = sqlx::query_scalar(
|
||||
"SELECT metadata FROM agent_messages WHERE session_id = ? AND role = 'user' AND active = 1 AND metadata LIKE '%\"dynamic_context\"%' ORDER BY id DESC LIMIT 1",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
metadata
|
||||
.and_then(|m| serde_json::from_str::<serde_json::Value>(&m).ok())
|
||||
.and_then(|v| v.get("context_hash")?.as_str().map(|s| s.to_string()))
|
||||
}
|
||||
|
||||
/// 持久化动态上下文快照（metadata 携带哈希供下轮比较）
|
||||
async fn persist_dynamic_context(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
msg: &ChatMessage,
|
||||
hash: u64,
|
||||
) {
|
||||
use crate::agent::engine::{DbMessageSink, MessageSink as _};
|
||||
let sink = DbMessageSink {
|
||||
db: db.clone(),
|
||||
session_id: session_id.to_string(),
|
||||
agent_name: "lead".to_string(),
|
||||
fixed_metadata: None,
|
||||
};
|
||||
sink.save(
|
||||
turn_index,
|
||||
0,
|
||||
msg,
|
||||
None,
|
||||
Some(serde_json::json!({ "dynamic_context": true, "context_hash": hash.to_string() })),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -1,27 +1,95 @@
|
||||
// src/agent/runtime/duplicate_detector.rs
|
||||
//
|
||||
// 同质调用检测器：检测连续重复的工具调用，防止死循环。
|
||||
// 同质调用检测器：检测重复的工具调用，防止死循环。
|
||||
//
|
||||
// 历史实现是单槽（只记住上一次调用），A/B 交替死循环
|
||||
// （call A → call B → call A → call B ...）检测不到。
|
||||
// 现改为"连续计数 + 滑动窗口计数"双通道：
|
||||
// - 连续通道：连续相同调用 ≥ threshold（原语义，立即触发）
|
||||
// - 窗口通道：最近 threshold*2 次调用中同一调用出现 ≥ threshold+2 次
|
||||
// （捕获 A/B 交替；门槛略高以放过合理的重复只读调用）
|
||||
|
||||
use std::collections::VecDeque;
|
||||
|
||||
/// 同质调用检测器
|
||||
#[derive(Debug, Default)]
|
||||
pub struct DuplicateDetector {
|
||||
last_call: Option<(String, String)>, // (tool_name, arguments)
|
||||
consecutive_count: usize,
|
||||
/// 最近调用的滑动窗口（name, arguments）
|
||||
recent: VecDeque<(String, String)>,
|
||||
}
|
||||
|
||||
impl DuplicateDetector {
|
||||
/// 记录一次调用，返回是否检测到死循环
|
||||
pub fn record(&mut self, tool_name: &str, arguments: &str, threshold: usize) -> bool {
|
||||
let key = (tool_name.to_string(), arguments.to_string());
|
||||
if self.last_call.as_ref() == Some(&key) {
|
||||
self.consecutive_count += 1;
|
||||
if self.consecutive_count >= threshold {
|
||||
let window = (threshold * 2).max(4);
|
||||
|
||||
self.recent.push_back(key.clone());
|
||||
while self.recent.len() > window {
|
||||
self.recent.pop_front();
|
||||
}
|
||||
|
||||
// 连续通道：尾部连续出现次数
|
||||
let consecutive = self.recent.iter().rev().take_while(|k| *k == &key).count();
|
||||
if threshold >= 2 && consecutive >= threshold {
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
self.last_call = Some(key);
|
||||
self.consecutive_count = 1;
|
||||
}
|
||||
false
|
||||
|
||||
// 窗口通道：同一调用在窗口内出现次数（捕获 A/B 交替）。
|
||||
// 门槛与 threshold 相同：纯交替下窗口内单键最多出现 window/2 =
|
||||
// threshold 次，恰好可达；合理的间隔重复（≤ threshold-1 次）不受影响。
|
||||
let occurrences = self.recent.iter().filter(|k| **k == key).count();
|
||||
threshold >= 2 && occurrences >= threshold
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_consecutive_duplicates_detected() {
|
||||
let mut det = DuplicateDetector::default();
|
||||
assert!(!det.record("search_papers", "q=1", 3));
|
||||
assert!(!det.record("search_papers", "q=1", 3));
|
||||
assert!(det.record("search_papers", "q=1", 3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_alternating_loop_detected() {
|
||||
// A/B 交替：单槽实现检测不到，窗口通道应捕获
|
||||
let mut det = DuplicateDetector::default();
|
||||
let mut tripped = false;
|
||||
for i in 0..6 {
|
||||
let a = det.record("tool_a", "{}", 3);
|
||||
let b = det.record("tool_b", "{}", 3);
|
||||
if a || b {
|
||||
// 交替 5 次出现（threshold+2 = 5）后触发
|
||||
assert!(i >= 1, "不应过早触发 (round {})", i);
|
||||
tripped = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(tripped, "A/B 交替死循环应被检测到");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spaced_repeats_not_detected() {
|
||||
// 合理的重复只读调用（间隔其他调用）不应触发
|
||||
let mut det = DuplicateDetector::default();
|
||||
det.record("read_file", "a.rs", 3);
|
||||
det.record("grep_files", "pat", 3);
|
||||
det.record("read_file", "b.rs", 3);
|
||||
det.record("glob_files", "*.rs", 3);
|
||||
let tripped = det.record("read_file", "a.rs", 3);
|
||||
assert!(!tripped);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_different_args_not_duplicates() {
|
||||
let mut det = DuplicateDetector::default();
|
||||
for q in ["q=1", "q=2", "q=3", "q=4", "q=5", "q=6", "q=7", "q=8"] {
|
||||
assert!(!det.record("search_papers", q, 3));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,12 +1,10 @@
|
||||
// src/agent/runtime/executor/helpers.rs
|
||||
//
|
||||
// 执行器辅助类型与函数：PreparedCall, ToolResultMessage, ToolExecutionResult,
|
||||
// execute_single_tool, process_single_result, save_tool_message_sync。
|
||||
// execute_single_tool, process_single_result。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::clients::llm::ChatMessage;
|
||||
@@ -15,6 +13,7 @@ use super::AgentStreamEvent;
|
||||
use crate::agent::hooks::{
|
||||
event_label, HookRegistry, PostToolUseContext, PostToolUseFailureContext,
|
||||
};
|
||||
use crate::agent::tools::ToolOutput;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreparedCall {
|
||||
@@ -42,8 +41,6 @@ pub struct ToolExecutionResult {
|
||||
}
|
||||
|
||||
/// 执行单个工具调用（含超时和取消检测）。
|
||||
///
|
||||
/// 从原 `execute_parallel` 的闭包提取，供分区后的批次执行复用。
|
||||
pub(super) async fn execute_single_tool(
|
||||
tool_opt: Option<&dyn crate::agent::tools::AgentTool>,
|
||||
args: serde_json::Value,
|
||||
@@ -51,10 +48,10 @@ pub(super) async fn execute_single_tool(
|
||||
cancelled: &Arc<AtomicBool>,
|
||||
timeout_dur: std::time::Duration,
|
||||
tool_name: &str,
|
||||
) -> crate::agent::tools::ToolOutput {
|
||||
) -> ToolOutput {
|
||||
let tool = match tool_opt {
|
||||
Some(t) => t,
|
||||
None => return crate::agent::tools::ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
None => return ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
};
|
||||
|
||||
let interrupt_behavior = tool.interrupt_behavior();
|
||||
@@ -76,7 +73,7 @@ pub(super) async fn execute_single_tool(
|
||||
res = tokio::time::timeout(timeout_dur, tool_fut) => {
|
||||
match res {
|
||||
Ok(output) => output,
|
||||
Err(_) => crate::agent::tools::ToolOutput::error(format!(
|
||||
Err(_) => ToolOutput::error(format!(
|
||||
"工具 {} 执行超时（{}秒）",
|
||||
tool_name,
|
||||
timeout_dur.as_secs()
|
||||
@@ -84,50 +81,50 @@ pub(super) async fn execute_single_tool(
|
||||
}
|
||||
}
|
||||
_ = cancel_fut => {
|
||||
crate::agent::tools::ToolOutput::error("执行已被用户取消")
|
||||
ToolOutput::error("执行已被用户取消")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 处理单个工具执行结果（SSE 事件、PostToolUse hooks、持久化）。
|
||||
///
|
||||
/// 从原 `execute_parallel` 的结果处理循环提取。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn process_single_result(
|
||||
tool_call_id: &str,
|
||||
tool_name: &str,
|
||||
tool_args: &serde_json::Value,
|
||||
output: &crate::agent::tools::ToolOutput,
|
||||
output: &ToolOutput,
|
||||
cancelled_flag: bool,
|
||||
exec_start: std::time::Instant,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
hook_registry: &HookRegistry,
|
||||
library_dir: &std::path::Path,
|
||||
sid: &str,
|
||||
agent_name: &str,
|
||||
step: usize,
|
||||
max_output_chars: usize,
|
||||
ctx: &super::ExecutorContext<'_>,
|
||||
tool_messages: &mut Vec<ToolResultMessage>,
|
||||
additional_contexts: &mut Vec<String>,
|
||||
db: &SqlitePool,
|
||||
turn_index: i32,
|
||||
tool_registry: &crate::agent::tools::ToolRegistry,
|
||||
) {
|
||||
use crate::agent::tools::persist::maybe_persist_tool_result;
|
||||
use crate::agent::tools::persist::maybe_persist_tool_result_for;
|
||||
|
||||
let elapsed_ms = exec_start.elapsed().as_millis() as u64;
|
||||
let step = ctx.step;
|
||||
let tap = &ctx.tap;
|
||||
|
||||
let (is_internal, display_name) = if let Some(tool) = tool_registry.get(tool_name) {
|
||||
let empty_hooks;
|
||||
let hook_registry: &HookRegistry = match ctx.hook_registry {
|
||||
Some(h) => h,
|
||||
None => {
|
||||
empty_hooks = HookRegistry::new();
|
||||
&empty_hooks
|
||||
}
|
||||
};
|
||||
|
||||
let (is_internal, display_name) = if let Some(tool) = ctx.tool_registry.get(tool_name) {
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, tool_name.to_string())
|
||||
};
|
||||
|
||||
// SSE 事件 — 立即推送到前端
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tap.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: tool_call_id.to_string(),
|
||||
name: tool_name.to_string(),
|
||||
display_name,
|
||||
name: tap.display(tool_name),
|
||||
display_name: tap.display(&display_name),
|
||||
output: output.content.clone(),
|
||||
is_error: output.is_error,
|
||||
metadata: output.metadata.clone(),
|
||||
@@ -137,23 +134,30 @@ pub(super) async fn process_single_result(
|
||||
|
||||
// 输出处理：小结果直接传递，大结果持久化到磁盘并返回 stub
|
||||
// 但对于已从磁盘读取内容的工具（如 read_file），跳过持久化以防止级联
|
||||
let tool_results_dir = library_dir.join(".agent").join("tool-results");
|
||||
let tool_results_dir = ctx
|
||||
.app_state
|
||||
.config
|
||||
.storage
|
||||
.library_dir
|
||||
.join(".agent")
|
||||
.join("tool-results");
|
||||
let (processed_content, _persisted_path) = if output.skip_persist {
|
||||
(output.content.clone(), None)
|
||||
} else {
|
||||
maybe_persist_tool_result(
|
||||
maybe_persist_tool_result_for(
|
||||
&output.content,
|
||||
tool_call_id,
|
||||
max_output_chars,
|
||||
ctx.max_output_chars,
|
||||
&tool_results_dir,
|
||||
tool_name,
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
// PostToolUse hook
|
||||
let post_ctx = PostToolUseContext {
|
||||
session_id: sid.to_string(),
|
||||
agent_name: agent_name.to_string(),
|
||||
session_id: ctx.session_id.to_string(),
|
||||
agent_name: ctx.agent_name.to_string(),
|
||||
tool_name: tool_name.to_string(),
|
||||
tool_args: tool_args.clone(),
|
||||
output_content: processed_content.clone(),
|
||||
@@ -164,9 +168,14 @@ pub(super) async fn process_single_result(
|
||||
let post_result = hook_registry.run_post_tool_use(&post_ctx).await;
|
||||
let final_content = post_result.final_content;
|
||||
|
||||
// 非可信内容包裹（间接 prompt 注入防御）
|
||||
let llm_content =
|
||||
crate::agent::runtime::untrusted::wrap_untrusted_content(tool_name, &final_content);
|
||||
// 非可信内容包裹（间接 prompt 注入防御）。
|
||||
// 是否包裹由工具通过 AgentTool::untrusted_output 声明；
|
||||
// 未知工具回退到名字启发式（mcp__*/web_* 前缀）。
|
||||
let llm_content = if ctx.tool_registry.untrusted_output(tool_name) {
|
||||
crate::agent::runtime::untrusted::wrap_untrusted_content(tool_name, &final_content)
|
||||
} else {
|
||||
final_content
|
||||
};
|
||||
|
||||
// 收集 PostToolUse hook 注入的上下文
|
||||
if !post_result.tagged_contexts.is_empty() {
|
||||
@@ -179,8 +188,8 @@ pub(super) async fn process_single_result(
|
||||
));
|
||||
}
|
||||
} else {
|
||||
for ctx in &post_result.additional_contexts {
|
||||
additional_contexts.push(ctx.clone());
|
||||
for c in &post_result.additional_contexts {
|
||||
additional_contexts.push(c.clone());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -200,8 +209,8 @@ pub(super) async fn process_single_result(
|
||||
// PostToolUseFailure hook
|
||||
if output.is_error {
|
||||
let failure_ctx = PostToolUseFailureContext {
|
||||
session_id: sid.to_string(),
|
||||
agent_name: agent_name.to_string(),
|
||||
session_id: ctx.session_id.to_string(),
|
||||
agent_name: ctx.agent_name.to_string(),
|
||||
tool_name: tool_name.to_string(),
|
||||
tool_args: tool_args.clone(),
|
||||
error_message: output.content.clone(),
|
||||
@@ -217,51 +226,25 @@ pub(super) async fn process_single_result(
|
||||
// 发送给 LLM 使用包裹后的内容（安全防御）
|
||||
let chat_message = ChatMessage::tool_result(tool_call_id, &llm_content);
|
||||
|
||||
// 持久化到数据库（fire-and-forget）
|
||||
save_tool_message_sync(db, sid, turn_index, step, &chat_message);
|
||||
// 持久化到数据库（fire-and-forget；无 sink 时跳过）。
|
||||
// canonical value 一并写入 metadata（重放/审计消费结构化结果）。
|
||||
if let Some(sink) = &ctx.sink {
|
||||
let sink = sink.clone();
|
||||
let msg = chat_message.clone();
|
||||
let turn_index = ctx.turn_index;
|
||||
let step_i = ctx.step;
|
||||
let extra = output
|
||||
.value
|
||||
.clone()
|
||||
.map(|v| serde_json::json!({ "canonical_value": v }));
|
||||
tokio::spawn(async move {
|
||||
sink.save(turn_index, step_i as i32, &msg, None, extra)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message,
|
||||
was_error: output.is_error,
|
||||
});
|
||||
}
|
||||
|
||||
/// 同步保存 tool 角色消息到数据库。
|
||||
pub(super) fn save_tool_message_sync(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
step_index: usize,
|
||||
msg: &ChatMessage,
|
||||
) {
|
||||
let db_clone = db.clone();
|
||||
let session_id = session_id.to_string();
|
||||
let content = msg.text().unwrap_or("").to_string();
|
||||
let tool_call_id = msg.tool_call_id.clone();
|
||||
// 提前序列化，避免闭包内的生命周期问题
|
||||
let metadata_str =
|
||||
serde_json::to_string(&serde_json::json!({ "role": "tool" })).unwrap_or_default();
|
||||
let raw_json = serde_json::to_string(&msg).unwrap_or_default();
|
||||
// fire-and-forget: tool 消息保存失败不影响主流程
|
||||
tokio::spawn(async move {
|
||||
let token_count = content.len() as i32 / 4;
|
||||
if let Err(e) = sqlx::query(
|
||||
"INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, tool_call_id, token_count, metadata, raw_json, agent_name) \
|
||||
VALUES (?, ?, ?, 'tool', ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(turn_index)
|
||||
.bind(step_index as i32)
|
||||
.bind(&content)
|
||||
.bind(&tool_call_id)
|
||||
.bind(token_count)
|
||||
.bind(&metadata_str)
|
||||
.bind(&raw_json)
|
||||
.bind("lead")
|
||||
.execute(&db_clone)
|
||||
.await
|
||||
{
|
||||
warn!("[Executor] 保存 tool 消息失败（非致命）: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
+255
-279
@@ -1,18 +1,21 @@
|
||||
// src/agent/runtime/executor/mod.rs
|
||||
//
|
||||
// 工具调用验证与并行执行器。
|
||||
//
|
||||
// 执行参数通过 ExecutorContext 组合传递（历史实现有约 20 个位置参数）。
|
||||
// 行为注入点：
|
||||
// - ask_policy: AskUser 权限请求是交互等待（主代理）还是 fail-closed 自动拒绝（子代理/队友）
|
||||
// - tap: SSE 事件出口（None = 静默；prefix = "[sub]"）
|
||||
// - sink: 消息持久化出口（None = 不落库）
|
||||
|
||||
mod helpers;
|
||||
|
||||
// Re-export 公共类型
|
||||
pub use helpers::{PreparedCall, ToolExecutionResult, ToolResultMessage};
|
||||
|
||||
use futures_util::stream::FuturesUnordered;
|
||||
use futures_util::StreamExt;
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::api::{AppState, PendingPermission};
|
||||
@@ -26,26 +29,23 @@ use super::partitioner::ToolPartitioner;
|
||||
use super::permission::{PermissionChecker, PermissionResult};
|
||||
use super::permission_explainer::explain_permission;
|
||||
use super::{AgentStreamEvent, DuplicateDetector};
|
||||
use crate::agent::engine::{AskPolicy, EventTap, MessageSink};
|
||||
use crate::agent::hooks::{event_label, HookRegistry, PreToolUseContext};
|
||||
use crate::agent::tools::{ToolContext, ToolRegistry};
|
||||
use crate::agent::tools::ToolRegistry;
|
||||
|
||||
use helpers::{execute_single_tool, process_single_result, save_tool_message_sync};
|
||||
use helpers::{execute_single_tool, process_single_result};
|
||||
|
||||
/// 验证工具调用：死循环检测 + 参数解析。
|
||||
///
|
||||
/// 返回 (prepared_calls, has_duplicate)。
|
||||
/// 死循环或参数无效时，错误消息直接注入到 messages。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn validate_and_prepare(
|
||||
tool_calls: &[ToolCall],
|
||||
duplicate_detector: &mut DuplicateDetector,
|
||||
duplicate_threshold: usize,
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
tool_registry: &ToolRegistry,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
tap: &EventTap,
|
||||
step: usize,
|
||||
) -> (Vec<PreparedCall>, bool) {
|
||||
let mut prepared_calls: Vec<PreparedCall> = Vec::new();
|
||||
@@ -68,7 +68,7 @@ pub fn validate_and_prepare(
|
||||
"[Executor] 检测到死循环：{} 连续调用 {} 次",
|
||||
tool_name, duplicate_threshold
|
||||
);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
tap.send(AgentStreamEvent::Error {
|
||||
message: format!("检测到工具 {} 的重复调用，已自动终止循环。", tool_name),
|
||||
});
|
||||
let error_msg = ChatMessage::tool_result(
|
||||
@@ -94,10 +94,10 @@ pub fn validate_and_prepare(
|
||||
} else {
|
||||
(false, tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tap.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: call_id.clone(),
|
||||
name: tool_name.clone(),
|
||||
display_name,
|
||||
display_name: tap.display(&display_name),
|
||||
output: error_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
@@ -105,7 +105,6 @@ pub fn validate_and_prepare(
|
||||
is_internal,
|
||||
});
|
||||
let tool_msg = ChatMessage::tool_result(&call_id, &error_output);
|
||||
save_tool_message_sync(db, session_id, turn_index, step, &tool_msg);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
@@ -121,36 +120,48 @@ pub fn validate_and_prepare(
|
||||
(prepared_calls, has_duplicate)
|
||||
}
|
||||
|
||||
/// 执行器上下文 — 组合全部行为注入点（替代约 20 个位置参数）。
|
||||
pub struct ExecutorContext<'a> {
|
||||
pub tool_registry: &'a ToolRegistry,
|
||||
pub app_state: Arc<AppState>,
|
||||
/// hooks 缺省时使用空注册表语义（teammate 等无人值守场景）
|
||||
pub hook_registry: Option<&'a HookRegistry>,
|
||||
pub permission_checker: Option<&'a PermissionChecker>,
|
||||
pub session_checker: Option<&'a PermissionChecker>,
|
||||
pub denial_tracker: Option<&'a std::sync::Mutex<DenialTracker>>,
|
||||
pub checkpoint_manager: Option<&'a CheckpointManager>,
|
||||
pub db: &'a sqlx::SqlitePool,
|
||||
pub session_id: &'a str,
|
||||
pub agent_name: &'a str,
|
||||
pub turn_index: i32,
|
||||
pub step: usize,
|
||||
pub tool_timeout_secs: u64,
|
||||
pub max_output_chars: usize,
|
||||
pub read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||||
pub enable_thinking: bool,
|
||||
pub additional_allowed_dirs: Vec<String>,
|
||||
/// AskUser 权限请求处理策略
|
||||
pub ask_policy: AskPolicy,
|
||||
/// 取消观察：会话取消表（交互式）或外部原子标志
|
||||
pub cancel: Arc<AtomicBool>,
|
||||
pub interactive_cancel: bool,
|
||||
pub cancel_session_id: Option<String>,
|
||||
pub tap: EventTap,
|
||||
pub sink: Option<Arc<dyn MessageSink>>,
|
||||
}
|
||||
|
||||
/// 并行执行所有准备好的工具调用。
|
||||
///
|
||||
/// 流程：
|
||||
/// 1. 权限检查（deny 规则阻止不可执行工具）
|
||||
/// 2. 发送 ToolCall SSE 事件
|
||||
/// 3. 运行 PreToolUse hooks
|
||||
/// 1. Hardline 预检查（不可绕过的参数级拒绝）
|
||||
/// 2. 权限检查（deny 规则拦截；ask 按策略交互等待或 fail-closed 拒绝）
|
||||
/// 3. 发送 ToolCall SSE 事件 + PreToolUse hooks
|
||||
/// 4. 工具分区 + 并行执行（并发安全工具一批并行，不安全工具单独串行）
|
||||
/// 5. 收集结果、发送 ToolResult SSE、运行 PostToolUse hooks
|
||||
/// 6. 返回 ToolResultMessage 列表供调用方推入 messages
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn execute_parallel(
|
||||
prepared_calls: &[PreparedCall],
|
||||
tool_registry: &ToolRegistry,
|
||||
app_state: Arc<AppState>,
|
||||
hook_registry: &HookRegistry,
|
||||
permission_checker: Option<&PermissionChecker>,
|
||||
session_permission_checker: Option<&PermissionChecker>,
|
||||
denial_tracker: Option<&std::sync::Mutex<DenialTracker>>,
|
||||
checkpoint_manager: Option<&std::sync::Arc<CheckpointManager>>,
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
agent_name: &str,
|
||||
turn_index: i32,
|
||||
step: usize,
|
||||
tool_timeout_secs: u64,
|
||||
max_output_chars: usize,
|
||||
read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||||
enable_thinking: bool,
|
||||
additional_allowed_dirs: Vec<String>,
|
||||
ctx: &ExecutorContext<'_>,
|
||||
) -> ToolExecutionResult {
|
||||
if prepared_calls.is_empty() {
|
||||
return ToolExecutionResult {
|
||||
@@ -162,7 +173,72 @@ pub async fn execute_parallel(
|
||||
};
|
||||
}
|
||||
|
||||
let sid = session_id.to_string();
|
||||
let tool_registry = ctx.tool_registry;
|
||||
let tap = &ctx.tap;
|
||||
let step = ctx.step;
|
||||
let sid = ctx.session_id.to_string();
|
||||
let _turn_index = ctx.turn_index;
|
||||
let empty_hooks;
|
||||
let hook_registry: &HookRegistry = match ctx.hook_registry {
|
||||
Some(h) => h,
|
||||
None => {
|
||||
empty_hooks = HookRegistry::new();
|
||||
&empty_hooks
|
||||
}
|
||||
};
|
||||
|
||||
/// 工具消息的统一落库出口（无 sink 时跳过）
|
||||
fn persist_tool_message(
|
||||
sink: &Option<Arc<dyn MessageSink>>,
|
||||
turn_index: i32,
|
||||
step: usize,
|
||||
msg: &ChatMessage,
|
||||
) {
|
||||
if let Some(sink) = sink {
|
||||
let sink = sink.clone();
|
||||
let msg = msg.clone();
|
||||
tokio::spawn(async move {
|
||||
sink.save(turn_index, step as i32, &msg, None, None).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// 拒绝结果的统一出口：SSE 事件 + 消息落库 + 拒绝追踪
|
||||
async fn record_denial(
|
||||
ctx: &ExecutorContext<'_>,
|
||||
tap: &EventTap,
|
||||
tool_call_id: &str,
|
||||
tool_name: &str,
|
||||
err_output: &str,
|
||||
metadata: serde_json::Value,
|
||||
tool_messages: &mut Vec<ToolResultMessage>,
|
||||
) {
|
||||
let (is_internal, display_name) = match ctx.tool_registry.get(tool_name) {
|
||||
Some(tool) => (tool.is_internal(), tool.display_name().to_string()),
|
||||
None => (false, tool_name.to_string()),
|
||||
};
|
||||
tap.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: tool_call_id.to_string(),
|
||||
name: tool_name.to_string(),
|
||||
display_name: tap.display(&display_name),
|
||||
output: err_output.to_string(),
|
||||
is_error: true,
|
||||
metadata,
|
||||
step: ctx.step,
|
||||
is_internal,
|
||||
});
|
||||
let err_msg = ChatMessage::tool_result(tool_call_id, err_output);
|
||||
persist_tool_message(&ctx.sink, ctx.turn_index, ctx.step, &err_msg);
|
||||
tool_messages.push(ToolResultMessage {
|
||||
chat_message: err_msg,
|
||||
was_error: true,
|
||||
});
|
||||
if let Some(dt) = ctx.denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_denial();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 1: 发送 ToolCall SSE 事件
|
||||
for prep in prepared_calls {
|
||||
@@ -171,10 +247,10 @@ pub async fn execute_parallel(
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolCall {
|
||||
tap.send(AgentStreamEvent::ToolCall {
|
||||
id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
name: tap.display(&prep.tool_name),
|
||||
display_name: tap.display(&display_name),
|
||||
arguments: prep.args.clone(),
|
||||
step,
|
||||
is_internal,
|
||||
@@ -203,14 +279,12 @@ pub async fn execute_parallel(
|
||||
prep.tool_name, reason
|
||||
);
|
||||
}
|
||||
// 收集所有阻塞错误详情（含多个 hook 同时 block 的情况）
|
||||
for be in &result.blocking_errors {
|
||||
hook_blocking_errors.push(format!(
|
||||
"[{}] 阻止 {}: {}",
|
||||
be.hook_name, prep.tool_name, be.reason
|
||||
));
|
||||
}
|
||||
// 收集 hook 的权限请求（保留完整信息用于 AskUser prompt）
|
||||
if let Some((permission, tool_name)) = result.permission_info() {
|
||||
info!(
|
||||
"[Executor] Hook 请求了工具 {} 的权限确认: {}",
|
||||
@@ -220,9 +294,7 @@ pub async fn execute_parallel(
|
||||
} else {
|
||||
hook_permission_info.push(None);
|
||||
}
|
||||
// 使用 hook 可能修改后的参数
|
||||
mutated_args.push(result.final_args);
|
||||
// 收集所有 hook 注入的上下文（优先使用带来源标记的 tagged_contexts）
|
||||
if !result.tagged_contexts.is_empty() {
|
||||
for tc in &result.tagged_contexts {
|
||||
additional_contexts.push(format!(
|
||||
@@ -233,39 +305,25 @@ pub async fn execute_parallel(
|
||||
));
|
||||
}
|
||||
} else {
|
||||
for ctx in &result.additional_contexts {
|
||||
additional_contexts.push(ctx.clone());
|
||||
for c in &result.additional_contexts {
|
||||
additional_contexts.push(c.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2.5: 权限检查 — PermissionChecker 规则引擎拦截被拒绝的工具。
|
||||
// Phase 2.5: Hardline 预检查 + 权限检查。
|
||||
// 被拒绝的工具直接注入错误 result，不进入执行队列。
|
||||
let mut tool_messages: Vec<ToolResultMessage> = Vec::new();
|
||||
let mut denied_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();
|
||||
|
||||
// ── Hardline 预检查（在任何模式下都不可绕过）──
|
||||
// 在 PermissionChecker 之前执行，确保 hardline 规则始终生效。
|
||||
// 检查逻辑由工具自身通过 AgentTool::hardline_check 声明（按参数路由，
|
||||
// 不再按工具名字符串匹配）。
|
||||
for (i, prep) in prepared_calls.iter().enumerate() {
|
||||
let hardline_result = match prep.tool_name.as_str() {
|
||||
"run_bash" => {
|
||||
if let Some(cmd) = prep.args.get("command").and_then(|v| v.as_str()) {
|
||||
hardline::check_command(cmd)
|
||||
} else {
|
||||
hardline::HardlineResult::allowed()
|
||||
}
|
||||
}
|
||||
"file_write" | "file_edit" => {
|
||||
if let Some(path) = prep.args.get("file_path").and_then(|v| v.as_str()) {
|
||||
hardline::check_dangerous_path(path)
|
||||
} else if let Some(path) = prep.args.get("path").and_then(|v| v.as_str()) {
|
||||
hardline::check_dangerous_path(path)
|
||||
} else {
|
||||
hardline::HardlineResult::allowed()
|
||||
}
|
||||
}
|
||||
_ => hardline::HardlineResult::allowed(),
|
||||
};
|
||||
let hardline_result = tool_registry
|
||||
.get(&prep.tool_name)
|
||||
.and_then(|tool| tool.hardline_check(&prep.args))
|
||||
.unwrap_or_else(hardline::HardlineResult::allowed);
|
||||
|
||||
if hardline_result.blocked {
|
||||
warn!(
|
||||
@@ -275,48 +333,30 @@ pub async fn execute_parallel(
|
||||
hardline_result.reason
|
||||
);
|
||||
let err_output = hardline_result.reason.clone();
|
||||
let (is_internal, display_name) = if let Some(tool) = tool_registry.get(&prep.tool_name)
|
||||
{
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
output: err_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({
|
||||
"hardline_blocked": true,
|
||||
"hardline_category": hardline_result.category,
|
||||
}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
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();
|
||||
}
|
||||
}
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(checker) = permission_checker {
|
||||
if let Some(checker) = ctx.permission_checker {
|
||||
for (i, prep) in prepared_calls.iter().enumerate() {
|
||||
// ── 权限决策合并（单调不变量：Deny 粘滞，只能收紧不能放松）──
|
||||
let mut perm_result = checker.check(&prep.tool_name, Some(&prep.args));
|
||||
perm_result = checker.apply_mode(perm_result, &prep.tool_name);
|
||||
|
||||
// Hook PermissionRequired — 若 Checker 返回 Allowed，升级为 Ask
|
||||
// 使用 hook 提供的具体权限描述替换泛型消息
|
||||
if let Some(Some((ref perm_desc, _))) = hook_permission_info.get(i) {
|
||||
if perm_result.is_allowed() {
|
||||
perm_result = PermissionResult::AskUser {
|
||||
@@ -330,14 +370,12 @@ pub async fn execute_parallel(
|
||||
}
|
||||
}
|
||||
|
||||
// 工具级 check_permissions() — 在 PermissionChecker 结果基础上叠加
|
||||
// PermissionChecker Deny/Ask 优先，工具级规则在 Allow 时可升级为 Ask
|
||||
// 工具级 check_permissions() — 在已有结果基础上收紧
|
||||
if let Some(tool) = tool_registry.get(&prep.tool_name) {
|
||||
let tool_rules = tool.check_permissions(&prep.args);
|
||||
for tool_rule in &tool_rules {
|
||||
match tool_rule {
|
||||
crate::agent::tools::PermissionRule::Deny { reason, .. } => {
|
||||
// 工具级 Deny 仅在 PermissionChecker 未 Deny 时生效
|
||||
if !perm_result.is_denied() {
|
||||
perm_result = PermissionResult::Denied {
|
||||
reason: reason.clone(),
|
||||
@@ -345,7 +383,6 @@ pub async fn execute_parallel(
|
||||
}
|
||||
}
|
||||
crate::agent::tools::PermissionRule::Ask { message, .. } => {
|
||||
// 工具级 Ask：若 PermissionChecker 返回 Allowed，升级为 Ask
|
||||
if perm_result.is_allowed() {
|
||||
perm_result = PermissionResult::AskUser {
|
||||
message: message.clone(),
|
||||
@@ -357,67 +394,57 @@ pub async fn execute_parallel(
|
||||
}
|
||||
}
|
||||
|
||||
// 会话级权限检查（API 动态添加的规则，优先级高于环境变量规则）
|
||||
if let Some(session_checker) = session_permission_checker {
|
||||
let session_result = session_checker.check(&prep.tool_name, Some(&prep.args));
|
||||
// 会话规则结果覆盖或升级
|
||||
match session_result {
|
||||
PermissionResult::Denied { reason } => {
|
||||
// 会话 Deny 强制覆盖
|
||||
perm_result = PermissionResult::Denied { reason };
|
||||
}
|
||||
PermissionResult::AskUser { message } => {
|
||||
// 会话 Ask 在 Allow 时升级
|
||||
if perm_result.is_allowed() {
|
||||
perm_result = PermissionResult::AskUser { message };
|
||||
}
|
||||
}
|
||||
PermissionResult::Allowed => {
|
||||
// 会话 Allow 仅覆盖 Allowed，保持 Deny/AskUser 不变
|
||||
// 避免覆盖工具级 check_permissions() 升级的 AskUser
|
||||
}
|
||||
}
|
||||
// 会话级权限检查（API 动态添加的规则，只能收紧）
|
||||
if let Some(session_checker) = ctx.session_checker {
|
||||
perm_result = super::permission::tighten(
|
||||
perm_result,
|
||||
session_checker.check(&prep.tool_name, Some(&prep.args)),
|
||||
);
|
||||
}
|
||||
|
||||
match perm_result {
|
||||
PermissionResult::Denied { reason } => {
|
||||
warn!(
|
||||
"[Executor] PermissionChecker 拒绝了工具 {}: {}",
|
||||
"[Executor] 权限检查拒绝了工具 {}: {}",
|
||||
prep.tool_name, reason
|
||||
);
|
||||
let err_output =
|
||||
format!("工具 {} 被权限规则拒绝执行: {}", prep.tool_name, reason);
|
||||
let (is_internal, display_name) =
|
||||
if let Some(tool) = tool_registry.get(&prep.tool_name) {
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
output: err_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
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();
|
||||
}
|
||||
}
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({}),
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
}
|
||||
PermissionResult::AskUser { message } => {
|
||||
// 无人值守上下文（子代理/队友）无用户可问：fail-closed 自动拒绝
|
||||
if ctx.ask_policy == AskPolicy::AutoDeny {
|
||||
let reason =
|
||||
match crate::agent::engine::auto_deny_ask(&prep.tool_name, &message) {
|
||||
PermissionResult::Denied { reason } => reason,
|
||||
_ => "需要用户确认但在无人值守上下文中不可用".to_string(),
|
||||
};
|
||||
let err_output = format!("工具 {} 被拒绝执行: {}", prep.tool_name, reason);
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({ "auto_denied_ask": true }),
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
continue;
|
||||
}
|
||||
|
||||
info!(
|
||||
"[Executor] PermissionChecker 请求用户确认工具 {}: {}",
|
||||
prep.tool_name, message
|
||||
@@ -427,8 +454,7 @@ pub async fn execute_parallel(
|
||||
let permission_exp = explain_permission(&prep.tool_name, &prep.args);
|
||||
let explanation_json = serde_json::to_value(&permission_exp).ok();
|
||||
|
||||
// 发送权限请求 SSE 事件
|
||||
let _ = tx.send(AgentStreamEvent::PermissionRequest {
|
||||
tap.send(AgentStreamEvent::PermissionRequest {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
tool_name: prep.tool_name.clone(),
|
||||
message: message.clone(),
|
||||
@@ -439,17 +465,14 @@ pub async fn execute_parallel(
|
||||
// 创建 oneshot 通道等待用户响应
|
||||
let (resp_tx, resp_rx) = oneshot::channel();
|
||||
let perm_id = uuid::Uuid::new_v4().to_string();
|
||||
let tc_id = prep.tool_call_id.clone();
|
||||
let t_name = prep.tool_name.clone();
|
||||
|
||||
// 存储待处理的权限请求
|
||||
{
|
||||
let mut perms = app_state.session.pending_permissions.lock().await;
|
||||
let mut perms = ctx.app_state.session.pending_permissions.lock().await;
|
||||
perms.insert(
|
||||
perm_id.clone(),
|
||||
PendingPermission {
|
||||
tool_call_id: tc_id.clone(),
|
||||
tool_name: t_name.clone(),
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
tool_name: prep.tool_name.clone(),
|
||||
message: message.clone(),
|
||||
arguments: prep.args.clone(),
|
||||
response_tx: resp_tx,
|
||||
@@ -458,12 +481,11 @@ pub async fn execute_parallel(
|
||||
);
|
||||
}
|
||||
|
||||
// 等待用户响应（120 秒超时）
|
||||
// 等待用户响应（120 秒超时 = fail-closed 拒绝）
|
||||
let timeout_dur = std::time::Duration::from_secs(120);
|
||||
let perm_result = tokio::time::timeout(timeout_dur, resp_rx).await;
|
||||
|
||||
// 清理待处理的权限请求
|
||||
app_state
|
||||
ctx.app_state
|
||||
.session
|
||||
.pending_permissions
|
||||
.lock()
|
||||
@@ -473,99 +495,61 @@ pub async fn execute_parallel(
|
||||
match perm_result {
|
||||
Ok(Ok(response)) if response.allowed => {
|
||||
info!("[Executor] 用户允许了工具 {} 的执行", prep.tool_name);
|
||||
let _ = tx.send(AgentStreamEvent::PermissionResponse {
|
||||
tap.send(AgentStreamEvent::PermissionResponse {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
allowed: true,
|
||||
});
|
||||
// 用户允许 → 重置连续拒绝计数
|
||||
if let Some(dt) = denial_tracker {
|
||||
if let Some(dt) = ctx.denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_success();
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Ok(_response)) => {
|
||||
// 用户拒绝
|
||||
info!("[Executor] 用户拒绝了工具 {}", prep.tool_name);
|
||||
let err_output = format!("用户拒绝了工具 {} 的执行", prep.tool_name);
|
||||
let (is_internal, display_name) =
|
||||
if let Some(tool) = tool_registry.get(&prep.tool_name) {
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
output: err_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
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);
|
||||
let _ = tx.send(AgentStreamEvent::PermissionResponse {
|
||||
tap.send(AgentStreamEvent::PermissionResponse {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
allowed: false,
|
||||
});
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({}),
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
}
|
||||
_ => {
|
||||
// 超时或通道关闭
|
||||
// 超时或通道关闭 → fail-closed
|
||||
warn!("[Executor] 权限请求超时或取消: {}", prep.tool_name);
|
||||
let err_output =
|
||||
format!("权限请求超时 (120s): {} 未获得用户确认", prep.tool_name);
|
||||
let (is_internal, display_name) =
|
||||
if let Some(tool) = tool_registry.get(&prep.tool_name) {
|
||||
(tool.is_internal(), tool.display_name().to_string())
|
||||
} else {
|
||||
(false, prep.tool_name.clone())
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
name: prep.tool_name.clone(),
|
||||
display_name,
|
||||
output: err_output.clone(),
|
||||
is_error: true,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
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);
|
||||
let _ = tx.send(AgentStreamEvent::PermissionResponse {
|
||||
tap.send(AgentStreamEvent::PermissionResponse {
|
||||
tool_call_id: prep.tool_call_id.clone(),
|
||||
allowed: false,
|
||||
});
|
||||
record_denial(
|
||||
ctx,
|
||||
tap,
|
||||
&prep.tool_call_id,
|
||||
&prep.tool_name,
|
||||
&err_output,
|
||||
serde_json::json!({ "timeout": true }),
|
||||
&mut tool_messages,
|
||||
)
|
||||
.await;
|
||||
denied_indices.insert(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
PermissionResult::Allowed => {
|
||||
// 工具被允许 → 重置连续拒绝计数
|
||||
if let Some(dt) = denial_tracker {
|
||||
if let Some(dt) = ctx.denial_tracker {
|
||||
if let Ok(mut tracker) = dt.lock() {
|
||||
tracker.record_success();
|
||||
}
|
||||
@@ -575,38 +559,46 @@ pub async fn execute_parallel(
|
||||
}
|
||||
} // if let Some(checker)
|
||||
|
||||
// Phase 3: 分区并行执行（参考 Claude Code partitionToolCalls + runTools）。
|
||||
//
|
||||
// 改进：原实现将所有非拒绝工具放入单个 FuturesUnordered 无差别并发，
|
||||
// 可能导致非并发安全工具（如 run_bash）错误地并行执行。
|
||||
// 新实现使用 ToolPartitioner 将工具按并发安全性分批：
|
||||
// Phase 3: 分区并行执行。
|
||||
// ToolPartitioner 将工具按并发安全性分批：
|
||||
// - 连续的并发安全工具放入同一个并行批次（FuturesUnordered）
|
||||
// - 非并发安全工具独占一个串行批次（逐次执行）
|
||||
// 批次内工具执行完成后立即推送 SSE 事件，不等待整个批次完成。
|
||||
let cancelled = Arc::new(AtomicBool::new(false));
|
||||
let cancel_flag = cancelled.clone();
|
||||
let app_state_ref = app_state.clone();
|
||||
let sid_ref = sid.clone();
|
||||
let app_state_ref = ctx.app_state.clone();
|
||||
let _sid_ref = sid.clone();
|
||||
let interactive = ctx.interactive_cancel;
|
||||
let cancel_sid = ctx.cancel_session_id.clone();
|
||||
let external_flag = ctx.cancel.clone();
|
||||
|
||||
let cancel_handle = tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
if app_state_ref.session.cancelled_runs.contains_key(&sid_ref) {
|
||||
let tripped = if interactive {
|
||||
cancel_sid
|
||||
.as_ref()
|
||||
.is_some_and(|s| app_state_ref.session.cancelled_runs.contains_key(s))
|
||||
} else {
|
||||
external_flag.load(Ordering::SeqCst)
|
||||
};
|
||||
if tripped {
|
||||
cancel_flag.store(true, Ordering::SeqCst);
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let timeout_dur = std::time::Duration::from_secs(tool_timeout_secs);
|
||||
let timeout_dur = std::time::Duration::from_secs(ctx.tool_timeout_secs);
|
||||
|
||||
// ── Checkpoint 预触发：对文件变更类工具在执行前创建快照 ──
|
||||
if let Some(ckpt) = checkpoint_manager {
|
||||
// 由工具通过 AgentTool::causes_file_changes 声明
|
||||
if let Some(ckpt) = ctx.checkpoint_manager {
|
||||
let cwd = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
|
||||
for prep in prepared_calls
|
||||
.iter()
|
||||
.filter(|p| CheckpointManager::should_checkpoint(&p.tool_name))
|
||||
{
|
||||
for prep in prepared_calls.iter().filter(|p| {
|
||||
tool_registry
|
||||
.get(&p.tool_name)
|
||||
.is_some_and(|t| t.causes_file_changes())
|
||||
}) {
|
||||
ckpt.ensure_checkpoint(&cwd, &format!("pre-{}", prep.tool_name));
|
||||
}
|
||||
}
|
||||
@@ -624,7 +616,6 @@ pub async fn execute_parallel(
|
||||
let partitioner = ToolPartitioner::new(10);
|
||||
let batches = partitioner.partition(&non_denied_calls, tool_registry);
|
||||
|
||||
// 预设非拒绝工具中哪些原索引属于已拒绝列表（不会有，但安全起见）
|
||||
let original_index_of: std::collections::HashMap<String, usize> = non_denied
|
||||
.iter()
|
||||
.map(|(orig_idx, prep)| (prep.tool_call_id.clone(), *orig_idx))
|
||||
@@ -641,7 +632,6 @@ pub async fn execute_parallel(
|
||||
let mut was_cancelled = false;
|
||||
|
||||
// ── Phase 3c: 逐批次执行 ──
|
||||
// 批次之间串行；并行批次内工具并发执行；串行批次内工具逐个执行。
|
||||
for batch in &batches {
|
||||
if was_cancelled {
|
||||
break;
|
||||
@@ -649,7 +639,7 @@ pub async fn execute_parallel(
|
||||
|
||||
if batch.is_parallel {
|
||||
// ── 并行批次：FuturesUnordered 并发执行 ──
|
||||
let mut exec_futs: FuturesUnordered<_> = batch
|
||||
let mut exec_futs: futures_util::stream::FuturesUnordered<_> = batch
|
||||
.calls
|
||||
.iter()
|
||||
.map(|prep| {
|
||||
@@ -662,14 +652,16 @@ pub async fn execute_parallel(
|
||||
.get(orig_idx)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| prep.args.clone());
|
||||
let tool_ctx =
|
||||
ToolContext::with_file_cache(app_state.clone(), read_file_state.clone())
|
||||
.with_sse_tx(tx.clone())
|
||||
.with_session_id(session_id.to_string())
|
||||
.with_thinking(enable_thinking)
|
||||
.with_additional_dirs(additional_allowed_dirs.clone())
|
||||
let tool_ctx = crate::agent::tools::ToolContext::with_file_cache(
|
||||
ctx.app_state.clone(),
|
||||
ctx.read_file_state.clone(),
|
||||
)
|
||||
.with_sse_tx_opt(ctx.tap.sender())
|
||||
.with_session_id(sid.clone())
|
||||
.with_thinking(ctx.enable_thinking)
|
||||
.with_additional_dirs(ctx.additional_allowed_dirs.clone())
|
||||
.with_tool_call_id(prep.tool_call_id.clone())
|
||||
.with_max_output_chars(max_output_chars);
|
||||
.with_max_output_chars(ctx.max_output_chars);
|
||||
let cancelled = cancelled.clone();
|
||||
let tool_opt = tool_registry.get(&tool_name);
|
||||
|
||||
@@ -696,6 +688,7 @@ pub async fn execute_parallel(
|
||||
.collect();
|
||||
|
||||
// 渐进式处理：每个工具一完成就处理
|
||||
use futures_util::StreamExt;
|
||||
while let Some((tool_call_id, tool_name, tool_args, output, cancelled_flag)) =
|
||||
exec_futs.next().await
|
||||
{
|
||||
@@ -709,18 +702,9 @@ pub async fn execute_parallel(
|
||||
&output,
|
||||
cancelled_flag,
|
||||
exec_start,
|
||||
tx,
|
||||
hook_registry,
|
||||
&app_state.config.storage.library_dir,
|
||||
&sid,
|
||||
agent_name,
|
||||
step,
|
||||
max_output_chars,
|
||||
ctx,
|
||||
&mut tool_messages,
|
||||
&mut additional_contexts,
|
||||
db,
|
||||
turn_index,
|
||||
tool_registry,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -736,12 +720,13 @@ pub async fn execute_parallel(
|
||||
.get(orig_idx)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| prep.args.clone());
|
||||
let tool_ctx =
|
||||
ToolContext::with_file_cache(app_state.clone(), read_file_state.clone())
|
||||
.with_sse_tx(tx.clone())
|
||||
.with_session_id(session_id.to_string())
|
||||
.with_thinking(enable_thinking)
|
||||
.with_additional_dirs(additional_allowed_dirs.clone());
|
||||
let tool_ctx = crate::agent::tools::ToolContext::with_file_cache(
|
||||
ctx.app_state.clone(),
|
||||
ctx.read_file_state.clone(),
|
||||
)
|
||||
.with_session_id(sid.clone())
|
||||
.with_thinking(ctx.enable_thinking)
|
||||
.with_additional_dirs(ctx.additional_allowed_dirs.clone());
|
||||
let tool_opt = tool_registry.get(&tool_name);
|
||||
|
||||
let output = execute_single_tool(
|
||||
@@ -765,18 +750,9 @@ pub async fn execute_parallel(
|
||||
&output,
|
||||
cancelled_flag,
|
||||
exec_start,
|
||||
tx,
|
||||
hook_registry,
|
||||
&app_state.config.storage.library_dir,
|
||||
&sid,
|
||||
agent_name,
|
||||
step,
|
||||
max_output_chars,
|
||||
ctx,
|
||||
&mut tool_messages,
|
||||
&mut additional_contexts,
|
||||
db,
|
||||
turn_index,
|
||||
tool_registry,
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
+221
-977
File diff suppressed because it is too large
Load Diff
+86
-163
@@ -508,109 +508,33 @@ impl PermissionChecker {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Permission Precedence Resolver ──
|
||||
// ── Monotonic Combination（权限单调性不变量）──
|
||||
|
||||
/// 多源权限决策的最终裁决。遵循正式的优先级规则表：
|
||||
/// 权限决策合并 — 整个权限体系的唯一不变量：
|
||||
/// **只能收紧，不能放松**（deny > ask > allow，最严格胜出）。
|
||||
///
|
||||
/// | Priority | Source | Overridable By |
|
||||
/// |----------|----------------------------------|----------------|
|
||||
/// | P0 | PermissionChecker::Deny | Nothing |
|
||||
/// | P1 | Tool-level PermissionRule::Deny | Nothing |
|
||||
/// | P2 | Session-level Checker::Deny | Nothing |
|
||||
/// | P3 | Hook PreToolUseAction::Block | P0-P2 |
|
||||
/// | P4 | Hook PermissionRequired | P0-P3 |
|
||||
/// | P5-P7 | Checker::Ask / Tool::Ask / Allow | Normal |
|
||||
/// 多源决策（env 规则、工具级声明、会话级规则、hook 请求）逐层调用本函数
|
||||
/// 合并，任何一层都无法放行另一层已收紧的决策——这是 guard 单调性的
|
||||
/// 权限等价物（参考 deepseek-harness "guards only tighten" 约束）。
|
||||
///
|
||||
/// `conflict_log` 记录被覆盖的决策，便于审计和调试。
|
||||
pub fn resolve_permission_precedence(
|
||||
checker_result: PermissionResult,
|
||||
tool_rules: &[crate::agent::tools::PermissionRule],
|
||||
hook_permission: Option<&(String, String)>, // (permission_desc, tool_name)
|
||||
hook_blocked: bool,
|
||||
session_result: Option<PermissionResult>,
|
||||
) -> (PermissionResult, Vec<String>) {
|
||||
let mut final_result = checker_result;
|
||||
let mut conflict_log: Vec<String> = Vec::new();
|
||||
|
||||
// ── P1: Tool-level Deny ──
|
||||
for rule in tool_rules {
|
||||
if let crate::agent::tools::PermissionRule::Deny { reason, .. } = rule {
|
||||
if !final_result.is_denied() {
|
||||
conflict_log.push(format!("Tool-level Deny overrides checker: {reason}"));
|
||||
final_result = PermissionResult::Denied {
|
||||
reason: reason.clone(),
|
||||
};
|
||||
} else {
|
||||
conflict_log.push(format!(
|
||||
"Tool-level Deny '{reason}' ignored: already Denied"
|
||||
));
|
||||
/// 具体语义：
|
||||
/// - `base` 已 Denied → 结果必为 Denied（粘滞，不可覆盖）
|
||||
/// - `base` 已 AskUser → 候选 Allowed 不生效；候选 Denied 升级为 Denied
|
||||
/// - `base` Allowed → 采用候选（候选可为 Allowed/AskUser/Denied）
|
||||
pub fn tighten(base: PermissionResult, candidate: PermissionResult) -> PermissionResult {
|
||||
match base {
|
||||
PermissionResult::Denied { .. } => base, // Deny 粘滞
|
||||
PermissionResult::AskUser { .. } => match candidate {
|
||||
// Ask 不能被放松为 Allow；可升级为 Deny
|
||||
PermissionResult::Denied { reason } => PermissionResult::Denied { reason },
|
||||
_ => base,
|
||||
},
|
||||
PermissionResult::Allowed => candidate,
|
||||
}
|
||||
break; // only handle first Deny
|
||||
}
|
||||
}
|
||||
|
||||
// ── P2: Session-level Deny ──
|
||||
if let Some(PermissionResult::Denied { reason }) = &session_result {
|
||||
conflict_log.push(format!("Session-level Deny overrides current: {reason}"));
|
||||
final_result = PermissionResult::Denied {
|
||||
reason: reason.clone(),
|
||||
};
|
||||
}
|
||||
|
||||
// ── P3: Hook Block ──
|
||||
if hook_blocked {
|
||||
conflict_log.push("Hook Block prevents execution".to_string());
|
||||
// Block is already handled in the executor via denied_indices;
|
||||
// here we record it for the conflict log.
|
||||
}
|
||||
|
||||
// ── P4: Hook PermissionRequired ──
|
||||
if let Some((perm_desc, _tool_name)) = hook_permission {
|
||||
if final_result.is_allowed() {
|
||||
conflict_log.push(format!(
|
||||
"Hook PermissionRequired upgrades Allowed → Ask: {perm_desc}"
|
||||
));
|
||||
} else if !final_result.is_denied() {
|
||||
conflict_log.push(format!(
|
||||
"Hook PermissionRequired coexists with current state: {perm_desc}"
|
||||
));
|
||||
} else {
|
||||
conflict_log.push(format!(
|
||||
"Hook PermissionRequired '{perm_desc}' ignored: already Denied"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// ── P5: Session-level Ask ──
|
||||
if let Some(PermissionResult::AskUser { message }) = &session_result {
|
||||
if final_result.is_allowed() {
|
||||
final_result = PermissionResult::AskUser {
|
||||
message: message.clone(),
|
||||
};
|
||||
conflict_log.push("Session-level Ask overrides Allow".to_string());
|
||||
} else {
|
||||
conflict_log.push("Session-level Ask ignored: not Allowed".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
// ── P6: Tool-level Ask ──
|
||||
for rule in tool_rules {
|
||||
if let crate::agent::tools::PermissionRule::Ask { message, .. } = rule {
|
||||
if final_result.is_allowed() {
|
||||
conflict_log.push(format!("Tool-level Ask upgrades Allow: {message}"));
|
||||
final_result = PermissionResult::AskUser {
|
||||
message: message.clone(),
|
||||
};
|
||||
} else {
|
||||
conflict_log.push("Tool-level Ask ignored: not Allowed".to_string());
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
(final_result, conflict_log)
|
||||
}
|
||||
|
||||
// ── Permission Precedence Resolver ──
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -965,85 +889,84 @@ mod tests {
|
||||
assert!(matches!(result, PermissionResult::AskUser { .. }));
|
||||
}
|
||||
|
||||
// ── resolve_permission_precedence tests ──
|
||||
// ── tighten（单调合并不变量）tests ──
|
||||
|
||||
#[test]
|
||||
fn test_precedence_checker_deny_wins_over_all() {
|
||||
let (result, log) = resolve_permission_precedence(
|
||||
fn test_tighten_deny_is_sticky() {
|
||||
use crate::agent::runtime::permission::tighten;
|
||||
let base = PermissionResult::Denied {
|
||||
reason: "policy".into(),
|
||||
};
|
||||
// 任何候选都无法放松 Deny
|
||||
assert!(tighten(base.clone(), PermissionResult::Allowed).is_denied());
|
||||
assert!(tighten(
|
||||
base,
|
||||
PermissionResult::AskUser {
|
||||
message: "m".into()
|
||||
}
|
||||
)
|
||||
.is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tighten_ask_cannot_be_relaxed() {
|
||||
use crate::agent::runtime::permission::tighten;
|
||||
let base = PermissionResult::AskUser {
|
||||
message: "need confirm".into(),
|
||||
};
|
||||
// Ask + Allow → Ask（不可放松）
|
||||
let r = tighten(base.clone(), PermissionResult::Allowed);
|
||||
assert!(matches!(r, PermissionResult::AskUser { .. }));
|
||||
// Ask + Deny → Deny（可升级）
|
||||
let r = tighten(
|
||||
base,
|
||||
PermissionResult::Denied {
|
||||
reason: "blocked by policy".into(),
|
||||
reason: "no".into(),
|
||||
},
|
||||
&[],
|
||||
Some(&("need confirmation".to_string(), "test_tool".to_string())),
|
||||
false,
|
||||
None,
|
||||
);
|
||||
assert!(result.is_denied());
|
||||
assert!(
|
||||
!log.is_empty(),
|
||||
"conflict log should record the interaction"
|
||||
);
|
||||
assert!(r.is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_precedence_tool_deny_overrides_allow() {
|
||||
use crate::agent::tools::{PermissionRule, PermissionRuleSource};
|
||||
let tool_rules = vec![PermissionRule::Deny {
|
||||
tool_name: "test_tool".into(),
|
||||
reason: "tool self-protection".into(),
|
||||
source: PermissionRuleSource::Env,
|
||||
}];
|
||||
let (result, log) = resolve_permission_precedence(
|
||||
PermissionResult::Allowed,
|
||||
&tool_rules,
|
||||
None,
|
||||
false,
|
||||
None,
|
||||
fn test_tighten_allowed_adopts_candidate() {
|
||||
use crate::agent::runtime::permission::tighten;
|
||||
let base = PermissionResult::Allowed;
|
||||
assert_eq!(
|
||||
tighten(base.clone(), PermissionResult::Allowed),
|
||||
PermissionResult::Allowed
|
||||
);
|
||||
assert!(result.is_denied());
|
||||
assert!(!log.is_empty());
|
||||
assert!(matches!(
|
||||
tighten(
|
||||
base.clone(),
|
||||
PermissionResult::AskUser {
|
||||
message: "m".into()
|
||||
}
|
||||
),
|
||||
PermissionResult::AskUser { .. }
|
||||
));
|
||||
assert!(tighten(base, PermissionResult::Denied { reason: "d".into() }).is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_precedence_hook_block_recorded() {
|
||||
let (result, log) = resolve_permission_precedence(
|
||||
PermissionResult::Allowed,
|
||||
&[],
|
||||
None,
|
||||
true, // hook blocked
|
||||
None,
|
||||
);
|
||||
// Hook Block doesn't directly return Denied — it's logged for executor handling
|
||||
assert!(result.is_allowed());
|
||||
assert!(log.iter().any(|l| l.contains("Block")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_precedence_session_deny_overrides() {
|
||||
let (result, _log) = resolve_permission_precedence(
|
||||
PermissionResult::Allowed,
|
||||
&[],
|
||||
None,
|
||||
false,
|
||||
Some(PermissionResult::Denied {
|
||||
reason: "session deny".into(),
|
||||
}),
|
||||
);
|
||||
assert!(result.is_denied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_precedence_hook_permission_ignored_when_denied() {
|
||||
let (result, log) = resolve_permission_precedence(
|
||||
PermissionResult::Denied {
|
||||
reason: "policy deny".into(),
|
||||
fn test_tighten_layers_only_tighten() {
|
||||
// 模拟三层合并：env Allowed + 工具级 Ask + 会话级 Allow → 最终 Ask
|
||||
use crate::agent::runtime::permission::tighten;
|
||||
let merged = PermissionResult::Allowed;
|
||||
let merged = tighten(
|
||||
merged,
|
||||
PermissionResult::AskUser {
|
||||
message: "tool ask".into(),
|
||||
},
|
||||
&[],
|
||||
Some(&("need confirm".to_string(), "test_tool".to_string())),
|
||||
false,
|
||||
None,
|
||||
);
|
||||
assert!(result.is_denied());
|
||||
assert!(log.iter().any(|l| l.contains("ignored")));
|
||||
let merged = tighten(merged, PermissionResult::Allowed);
|
||||
assert!(matches!(merged, PermissionResult::AskUser { .. }));
|
||||
|
||||
// 反向：env Deny + 工具级 Allow + 会话级 Allow → 最终 Deny
|
||||
let merged = PermissionResult::Denied {
|
||||
reason: "env".into(),
|
||||
};
|
||||
let merged = tighten(merged.clone(), PermissionResult::Allowed);
|
||||
let merged = tighten(merged, PermissionResult::Allowed);
|
||||
assert!(merged.is_denied());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,6 +97,45 @@ pub async fn create_or_resume_session(
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用调用方预分配的会话 ID 创建（或恢复）会话。
|
||||
///
|
||||
/// 与 `create_or_resume_session` 的区别：会话不存在时以该 ID **创建**而非报错。
|
||||
/// 供 API 层使用——请求进入时预生成会话 ID，使取消标记、运行时缓存等
|
||||
/// 按 ID 索引的机制对新会话的首个请求同样生效（历史上新会话首请求
|
||||
/// 超时只能 abort 任务，无法写入取消标记）。
|
||||
pub async fn create_or_resume_session_preallocated(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
llm: &LlmClient,
|
||||
mode: &str,
|
||||
) -> anyhow::Result<SessionInfo> {
|
||||
let exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if exists {
|
||||
return create_or_resume_session(db, Some(session_id.to_string()), llm, mode).await;
|
||||
}
|
||||
|
||||
sqlx::query("INSERT INTO agent_sessions (session_id, title, model, mode) VALUES (?, ?, ?, ?)")
|
||||
.bind(session_id)
|
||||
.bind("")
|
||||
.bind(llm.model().await)
|
||||
.bind(mode)
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
Ok(SessionInfo {
|
||||
session_id: session_id.to_string(),
|
||||
turn_index: 0,
|
||||
mode: mode.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 从数据库加载会话的运行模式。
|
||||
///
|
||||
/// 返回 None 表示会话不存在或已删除。
|
||||
@@ -288,6 +327,9 @@ pub async fn rewind_to_message(
|
||||
|
||||
let preview: String = target_content.chars().take(120).collect();
|
||||
|
||||
// 回退改变了历史 → 上下文快照失效（其 base 高水位与消息集合不再对应）
|
||||
super::session_events::remove_context_snapshots(db, session_id).await;
|
||||
|
||||
info!(
|
||||
"[Session] 回退完成: session={}, rewound={}, to_id={}, new_turn={}",
|
||||
session_id, count, target_message_id, new_turn_index
|
||||
@@ -483,6 +525,9 @@ pub async fn retry_last_turn(
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
// 硬删除改变了历史 → 上下文快照失效
|
||||
super::session_events::remove_context_snapshots(db, session_id).await;
|
||||
|
||||
info!(
|
||||
"[Session] 重试: session={}, deleted={} messages from id={}, new_turn={}, has_image={}",
|
||||
session_id,
|
||||
@@ -595,6 +640,7 @@ pub async fn branch_session(db: &SqlitePool, session_id: &str) -> anyhow::Result
|
||||
|
||||
tx.commit().await?;
|
||||
|
||||
// 分叉是新会话：不复制事件（快照引用旧会话的消息 id，直接作废）
|
||||
info!(
|
||||
"[Session] 分叉完成: parent={}, branch={}, copied={} messages, forked_at={}",
|
||||
session_id, branch_id, copied, forked_at
|
||||
|
||||
@@ -0,0 +1,446 @@
|
||||
// src/agent/runtime/session_events.rs
|
||||
//
|
||||
// 会话事件日志 — turn/compaction 生命周期与上下文快照的事件溯源层。
|
||||
//
|
||||
// 不变量："模型可见 ⟺ 已日志化" 的生命周期侧面：
|
||||
// - 每个 turn 以 turn_start 开始、turn_end 关闭；崩溃留下可检测的
|
||||
// 开口（dangling start），恢复时合成 interrupted 关闭而不截断
|
||||
// 已持久化的内容（参考 deepseek-harness 的崩溃恢复策略）。
|
||||
// - 压缩以 compaction_start/compaction_end 构成日志化锁。
|
||||
// - context_snapshot 记录压缩后的折叠上下文 + 消息高水位
|
||||
// （base_message_id），下一 turn 回放快照 + 增量消息，避免重复压缩。
|
||||
|
||||
use serde_json::Value;
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// 事件类型常量
|
||||
pub mod event_types {
|
||||
pub const TURN_START: &str = "turn_start";
|
||||
pub const TURN_END: &str = "turn_end";
|
||||
pub const COMPACTION_START: &str = "compaction_start";
|
||||
pub const COMPACTION_END: &str = "compaction_end";
|
||||
/// 恢复时合成的中断标记（检测到开口的 turn_start/compaction_start）
|
||||
pub const TURN_INTERRUPTED: &str = "turn_interrupted";
|
||||
pub const COMPACTION_INTERRUPTED: &str = "compaction_interrupted";
|
||||
/// 压缩后的折叠上下文快照（payload 含 messages + base_message_id）
|
||||
pub const CONTEXT_SNAPSHOT: &str = "context_snapshot";
|
||||
}
|
||||
|
||||
/// 追加一条会话事件
|
||||
pub async fn append_event(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
event_type: &str,
|
||||
payload: Value,
|
||||
) {
|
||||
let payload_str = serde_json::to_string(&payload).unwrap_or_default();
|
||||
if let Err(e) = sqlx::query(
|
||||
"INSERT INTO agent_events (session_id, turn_index, event_type, payload) \
|
||||
VALUES (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(turn_index)
|
||||
.bind(event_type)
|
||||
.bind(&payload_str)
|
||||
.execute(db)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
"[SessionEvents] 写入 {} 事件失败（非致命）: {}",
|
||||
event_type, e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 最新一条指定类型的事件（无则 None）
|
||||
pub async fn latest_event(db: &SqlitePool, session_id: &str, event_type: &str) -> Option<Value> {
|
||||
let payload: Option<String> = sqlx::query_scalar(
|
||||
"SELECT payload FROM agent_events \
|
||||
WHERE session_id = ? AND event_type = ? ORDER BY id DESC LIMIT 1",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(event_type)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
payload.and_then(|p| serde_json::from_str(&p).ok())
|
||||
}
|
||||
|
||||
/// 检查是否存在"开口"事件：最新的 `open_type` 之后没有新的 `close_type`。
|
||||
/// 用于崩溃恢复——开口的 turn_start/compaction_start 表示上次执行被中断。
|
||||
pub async fn has_open_event(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
open_type: &str,
|
||||
close_type: &str,
|
||||
) -> bool {
|
||||
let open_id: Option<i64> = sqlx::query_scalar(
|
||||
"SELECT MAX(id) FROM agent_events WHERE session_id = ? AND event_type = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(open_type)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
|
||||
let Some(open_id) = open_id else {
|
||||
return false;
|
||||
};
|
||||
|
||||
let close_id: Option<i64> = sqlx::query_scalar(
|
||||
"SELECT MAX(id) FROM agent_events WHERE session_id = ? AND event_type = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(close_type)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
|
||||
close_id.is_none_or(|c| c < open_id)
|
||||
}
|
||||
|
||||
/// 恢复会话时的中断协调：检测开口的 turn/compaction，合成 interrupted 关闭。
|
||||
///
|
||||
/// 这是对崩溃/进程被杀的容错：不截断任何已持久化内容，
|
||||
/// 只补写"上次被中断"的事实标记。
|
||||
pub async fn reconcile_interrupted(db: &SqlitePool, session_id: &str) {
|
||||
if has_open_event(
|
||||
db,
|
||||
session_id,
|
||||
event_types::TURN_START,
|
||||
event_types::TURN_END,
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
"[SessionEvents] 会话 {} 检测到未关闭的 turn_start（上次执行被中断），合成 interrupted 关闭",
|
||||
session_id
|
||||
);
|
||||
append_event(
|
||||
db,
|
||||
session_id,
|
||||
-1,
|
||||
event_types::TURN_END,
|
||||
serde_json::json!({ "reason": "interrupted", "synthesized": true }),
|
||||
)
|
||||
.await;
|
||||
append_event(
|
||||
db,
|
||||
session_id,
|
||||
-1,
|
||||
event_types::TURN_INTERRUPTED,
|
||||
serde_json::json!({ "note": "检测到崩溃/中断残留，已合成关闭" }),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
if has_open_event(
|
||||
db,
|
||||
session_id,
|
||||
event_types::COMPACTION_START,
|
||||
event_types::COMPACTION_END,
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
"[SessionEvents] 会话 {} 检测到未关闭的 compaction_start，合成 interrupted 关闭",
|
||||
session_id
|
||||
);
|
||||
append_event(
|
||||
db,
|
||||
session_id,
|
||||
-1,
|
||||
event_types::COMPACTION_END,
|
||||
serde_json::json!({ "interrupted": true, "synthesized": true }),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 删除会话的上下文快照（rewind/retry/branch 后快照失效）
|
||||
pub async fn remove_context_snapshots(db: &SqlitePool, session_id: &str) {
|
||||
if let Err(e) = sqlx::query("DELETE FROM agent_events WHERE session_id = ? AND event_type = ?")
|
||||
.bind(session_id)
|
||||
.bind(event_types::CONTEXT_SNAPSHOT)
|
||||
.execute(db)
|
||||
.await
|
||||
{
|
||||
warn!("[SessionEvents] 清理上下文快照失败: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
/// 保存上下文快照：折叠后的消息 + 消息高水位。
|
||||
///
|
||||
/// `base_message_id` 是快照时 agent_messages 的最大 id；恢复时加载
|
||||
/// 快照消息 + id > base 的增量消息，避免重复压缩。
|
||||
pub async fn save_context_snapshot(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
turn_index: i32,
|
||||
messages: &[crate::clients::llm::ChatMessage],
|
||||
base_message_id: i64,
|
||||
) {
|
||||
// 排除 system 消息（每 turn 重建，且包含缓存的动态 section）
|
||||
let body: Vec<&crate::clients::llm::ChatMessage> = messages
|
||||
.iter()
|
||||
.filter(|m| m.role != crate::clients::llm::MessageRole::System)
|
||||
.collect();
|
||||
|
||||
let serialized: Vec<Value> = body
|
||||
.iter()
|
||||
.map(|m| serde_json::to_value(m).unwrap_or(Value::Null))
|
||||
.collect();
|
||||
|
||||
append_event(
|
||||
db,
|
||||
session_id,
|
||||
turn_index,
|
||||
event_types::CONTEXT_SNAPSHOT,
|
||||
serde_json::json!({
|
||||
"base_message_id": base_message_id,
|
||||
"message_count": serialized.len(),
|
||||
"messages": serialized,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
info!(
|
||||
"[SessionEvents] 保存上下文快照: session={} base_msg_id={} messages={}",
|
||||
session_id,
|
||||
base_message_id,
|
||||
serialized.len()
|
||||
);
|
||||
}
|
||||
|
||||
/// 加载上下文快照（若存在）：返回 (快照消息, base_message_id)。
|
||||
pub async fn load_context_snapshot(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> Option<(Vec<crate::clients::llm::ChatMessage>, i64)> {
|
||||
let payload = latest_event(db, session_id, event_types::CONTEXT_SNAPSHOT).await?;
|
||||
|
||||
let base_message_id = payload.get("base_message_id")?.as_i64()?;
|
||||
let messages_json = payload.get("messages")?.as_array()?;
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for m in messages_json {
|
||||
match serde_json::from_value::<crate::clients::llm::ChatMessage>(m.clone()) {
|
||||
Ok(msg) => messages.push(msg),
|
||||
Err(e) => {
|
||||
warn!("[SessionEvents] 快照消息反序列化失败，放弃快照: {}", e);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some((messages, base_message_id))
|
||||
}
|
||||
|
||||
/// 获取会话消息的当前最大 id（快照高水位）
|
||||
pub async fn max_message_id(db: &SqlitePool, session_id: &str) -> i64 {
|
||||
sqlx::query_scalar("SELECT COALESCE(MAX(id), 0) FROM agent_messages WHERE session_id = ?")
|
||||
.bind(session_id)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// 加载 id 大于 `after_id` 的增量历史消息（快照回放后拼接）
|
||||
pub async fn load_messages_after(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
after_id: i64,
|
||||
) -> anyhow::Result<Vec<crate::clients::llm::ChatMessage>> {
|
||||
#[allow(clippy::type_complexity)]
|
||||
let rows: Vec<(
|
||||
i64,
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
)> = sqlx::query_as(
|
||||
"SELECT id, role, content, tool_calls, tool_call_id, thought FROM agent_messages \
|
||||
WHERE session_id = ? AND id > ? AND active = 1 AND agent_name = 'lead' \
|
||||
ORDER BY id ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(after_id)
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
use crate::clients::llm::{ChatMessage, MessageRole};
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for (_id, role_str, content, tool_calls_json, tool_call_id, thought) in rows {
|
||||
let role = match role_str.as_str() {
|
||||
"system" => MessageRole::System,
|
||||
"user" => MessageRole::User,
|
||||
"assistant" => MessageRole::Assistant,
|
||||
"tool" => MessageRole::Tool,
|
||||
_ => continue,
|
||||
};
|
||||
let tool_calls: Option<Vec<crate::clients::llm::ToolCall>> =
|
||||
tool_calls_json.and_then(|j| serde_json::from_str(&j).ok());
|
||||
messages.push(ChatMessage {
|
||||
role,
|
||||
content: if content.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(content)
|
||||
},
|
||||
tool_call_id,
|
||||
tool_calls,
|
||||
name: None,
|
||||
reasoning_content: thought,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
async fn setup_db() -> SqlitePool {
|
||||
let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE agent_events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL,
|
||||
turn_index INTEGER NOT NULL DEFAULT 0,
|
||||
event_type TEXT NOT NULL,
|
||||
payload TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE agent_messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL,
|
||||
turn_index INTEGER NOT NULL DEFAULT 0,
|
||||
step_index INTEGER NOT NULL DEFAULT 0,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
thought TEXT,
|
||||
tool_calls TEXT,
|
||||
tool_call_id TEXT,
|
||||
token_count INTEGER NOT NULL DEFAULT 0,
|
||||
metadata TEXT,
|
||||
raw_json TEXT,
|
||||
agent_name TEXT NOT NULL DEFAULT 'lead',
|
||||
active INTEGER NOT NULL DEFAULT 1,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_open_event_detection() {
|
||||
let db = setup_db().await;
|
||||
let sid = "s1";
|
||||
|
||||
// 无事件 → 无开口
|
||||
assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await);
|
||||
|
||||
// turn_start 后无 turn_end → 开口
|
||||
append_event(&db, sid, 0, "turn_start", serde_json::json!({})).await;
|
||||
assert!(has_open_event(&db, sid, "turn_start", "turn_end").await);
|
||||
|
||||
// 关闭后 → 无开口
|
||||
append_event(
|
||||
&db,
|
||||
sid,
|
||||
0,
|
||||
"turn_end",
|
||||
serde_json::json!({ "reason": "completed" }),
|
||||
)
|
||||
.await;
|
||||
assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_reconcile_interrupted_synthesizes_close() {
|
||||
let db = setup_db().await;
|
||||
let sid = "s2";
|
||||
|
||||
append_event(&db, sid, 0, "turn_start", serde_json::json!({})).await;
|
||||
append_event(&db, sid, 0, "compaction_start", serde_json::json!({})).await;
|
||||
|
||||
reconcile_interrupted(&db, sid).await;
|
||||
|
||||
// 开口已闭合
|
||||
assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await);
|
||||
assert!(!has_open_event(&db, sid, "compaction_start", "compaction_end").await);
|
||||
// 合成了中断标记
|
||||
assert!(latest_event(&db, sid, "turn_interrupted").await.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_context_snapshot_roundtrip() {
|
||||
let db = setup_db().await;
|
||||
let sid = "s3";
|
||||
|
||||
use crate::clients::llm::ChatMessage;
|
||||
sqlx::query("INSERT INTO agent_messages (session_id, role, content, agent_name) VALUES (?, 'user', 'm1', 'lead')")
|
||||
.bind(sid)
|
||||
.execute(&db)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let base = max_message_id(&db, sid).await;
|
||||
assert_eq!(base, 1);
|
||||
|
||||
let messages = vec![
|
||||
ChatMessage::system("sys"),
|
||||
ChatMessage::user("[历史对话摘要]\n关于黑洞"),
|
||||
ChatMessage::assistant("结论"),
|
||||
];
|
||||
save_context_snapshot(&db, sid, 0, &messages, base).await;
|
||||
|
||||
let (loaded, loaded_base) = load_context_snapshot(&db, sid).await.unwrap();
|
||||
assert_eq!(loaded_base, base);
|
||||
// system 消息被排除
|
||||
assert_eq!(loaded.len(), 2);
|
||||
assert!(loaded[0].text().unwrap().contains("历史对话摘要"));
|
||||
|
||||
// 删除快照
|
||||
remove_context_snapshots(&db, sid).await;
|
||||
assert!(load_context_snapshot(&db, sid).await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_messages_after() {
|
||||
let db = setup_db().await;
|
||||
let sid = "s4";
|
||||
|
||||
for i in 1..=3 {
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_messages (id, session_id, role, content, agent_name) \
|
||||
VALUES (?, ?, 'user', ?, 'lead')",
|
||||
)
|
||||
.bind(i)
|
||||
.bind(sid)
|
||||
.bind(format!("msg-{}", i))
|
||||
.execute(&db)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let incremental = load_messages_after(&db, sid, 1).await.unwrap();
|
||||
assert_eq!(incremental.len(), 2);
|
||||
assert_eq!(incremental[0].text(), Some("msg-2"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
// src/agent/runtime/session_registry.rs
|
||||
//
|
||||
// 会话级运行时注册表 — AgentRuntime 按会话缓存复用。
|
||||
//
|
||||
// 历史问题：chat_agent 每个 HTTP 请求 new 一个 AgentRuntime，导致注释中
|
||||
// 声称"跨 turn 共享"的状态（后台任务队列、压缩折叠日志、文件缓存、
|
||||
// 拒绝追踪器、压缩熔断器）实际随请求销毁：
|
||||
// - 后台任务结果跨请求丢失（队列无人 drain）
|
||||
// - 压缩后的上下文不回写，下一 turn 从原始消息重建再重新压缩
|
||||
// - prompt cache / file cache 每请求冷启动
|
||||
//
|
||||
// 本注册表将 runtime 生命周期与 session 对齐（参考 deepseek-harness 的
|
||||
// "runtime 状态属于会话而非请求"原则）：
|
||||
// - 同一会话的请求复用同一 runtime 实例
|
||||
// - 每个会话持有一个 turn 互斥锁，防止同会话并发 turn 互相破坏
|
||||
// - 空闲超过 TTL 的条目由后台清扫回收
|
||||
// - 会话删除时显式移除
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use dashmap::DashMap;
|
||||
use tracing::info;
|
||||
|
||||
use super::AgentRuntime;
|
||||
use crate::api::AppState;
|
||||
|
||||
/// 单个会话的注册表条目
|
||||
pub struct SessionRuntimeEntry {
|
||||
pub runtime: Arc<AgentRuntime>,
|
||||
/// 同会话 turn 串行化锁（防止并发 turn 破坏 turn_index/消息顺序）
|
||||
pub turn_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
pub last_used: Instant,
|
||||
/// 创建时的模式 ID（用于诊断"模式漂移"）
|
||||
pub mode_id: String,
|
||||
}
|
||||
|
||||
/// 会话运行时注册表
|
||||
pub struct SessionRuntimeRegistry {
|
||||
entries: DashMap<String, SessionRuntimeEntry>,
|
||||
idle_ttl: Duration,
|
||||
}
|
||||
|
||||
impl Default for SessionRuntimeRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl SessionRuntimeRegistry {
|
||||
pub fn new() -> Self {
|
||||
SessionRuntimeRegistry {
|
||||
entries: DashMap::new(),
|
||||
// 2 小时无活动的会话 runtime 允许被回收；
|
||||
// 后台任务等通过 Arc 自然延长实际生命周期
|
||||
idle_ttl: Duration::from_secs(2 * 60 * 60),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取（或创建）指定会话的 runtime。
|
||||
///
|
||||
/// `mode_id` 只在首次创建时生效——后续请求从 DB 恢复会话模式，
|
||||
/// 不会因请求参数不同而静默切换（模式由会话决定，一次创建后保持稳定）。
|
||||
pub fn get_or_create(
|
||||
&self,
|
||||
app_state: Arc<AppState>,
|
||||
session_key: &str,
|
||||
mode_id: &str,
|
||||
) -> Arc<AgentRuntime> {
|
||||
if let Some(mut entry) = self.entries.get_mut(session_key) {
|
||||
entry.last_used = Instant::now();
|
||||
return entry.runtime.clone();
|
||||
}
|
||||
|
||||
// DashMap entry API 避免创建竞态
|
||||
let entry = self
|
||||
.entries
|
||||
.entry(session_key.to_string())
|
||||
.or_insert_with(|| {
|
||||
info!(
|
||||
"[SessionRegistry] 为会话 {} 创建 runtime (mode={})",
|
||||
session_key, mode_id
|
||||
);
|
||||
SessionRuntimeEntry {
|
||||
runtime: Arc::new(AgentRuntime::new_for_session(
|
||||
app_state,
|
||||
session_key,
|
||||
mode_id,
|
||||
)),
|
||||
turn_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
last_used: Instant::now(),
|
||||
mode_id: mode_id.to_string(),
|
||||
}
|
||||
});
|
||||
entry.runtime.clone()
|
||||
}
|
||||
|
||||
/// 获取会话的 turn 串行化锁（不存在时返回 None）
|
||||
pub fn turn_lock(&self, session_key: &str) -> Option<Arc<tokio::sync::Mutex<()>>> {
|
||||
self.entries.get(session_key).map(|e| e.turn_lock.clone())
|
||||
}
|
||||
|
||||
/// 移除会话条目（会话删除时调用）
|
||||
pub fn remove(&self, session_key: &str) {
|
||||
if self.entries.remove(session_key).is_some() {
|
||||
info!("[SessionRegistry] 移除会话 {} 的 runtime", session_key);
|
||||
}
|
||||
}
|
||||
|
||||
/// 当前缓存的会话数
|
||||
pub fn len(&self) -> usize {
|
||||
self.entries.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.entries.is_empty()
|
||||
}
|
||||
|
||||
/// 回收空闲超过 TTL 的条目（后台周期调用）
|
||||
pub fn sweep_idle(&self) -> usize {
|
||||
let now = Instant::now();
|
||||
let mut expired: Vec<String> = Vec::new();
|
||||
for entry in self.entries.iter() {
|
||||
if now.duration_since(entry.last_used) > self.idle_ttl {
|
||||
expired.push(entry.key().clone());
|
||||
}
|
||||
}
|
||||
let count = expired.len();
|
||||
for key in expired {
|
||||
self.entries.remove(&key);
|
||||
info!("[SessionRegistry] 回收空闲会话 runtime: {}", key);
|
||||
}
|
||||
count
|
||||
}
|
||||
}
|
||||
|
||||
// ── 会话级压缩守卫（替代进程级全局静态） ──
|
||||
//
|
||||
// 历史实现是单个进程级 AtomicBool：并发两个会话压缩时互相跳过。
|
||||
// 改为按 session_id 的集合守卫：同会话递归压缩被拦截，不同会话互不影响。
|
||||
|
||||
static COMPACTING_SESSIONS: std::sync::Mutex<Option<HashSet<String>>> = std::sync::Mutex::new(None);
|
||||
|
||||
/// 尝试进入压缩临界区。返回 false 表示该会话已有进行中的压缩（递归守卫触发）。
|
||||
pub fn try_begin_compaction(session_id: &str) -> bool {
|
||||
let mut guard = COMPACTING_SESSIONS
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
let set = guard.get_or_insert_with(HashSet::new);
|
||||
if set.contains(session_id) {
|
||||
false
|
||||
} else {
|
||||
set.insert(session_id.to_string());
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// 退出压缩临界区。会话不存在时静默（防御性）。
|
||||
pub fn end_compaction(session_id: &str) {
|
||||
let mut guard = COMPACTING_SESSIONS
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(set) = guard.as_mut() {
|
||||
set.remove(session_id);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_compaction_guard_per_session() {
|
||||
assert!(try_begin_compaction("session-a"));
|
||||
// 同会话递归被拦截
|
||||
assert!(!try_begin_compaction("session-a"));
|
||||
// 不同会话互不影响
|
||||
assert!(try_begin_compaction("session-b"));
|
||||
end_compaction("session-a");
|
||||
// 释放后可重新进入
|
||||
assert!(try_begin_compaction("session-a"));
|
||||
end_compaction("session-a");
|
||||
end_compaction("session-b");
|
||||
end_compaction("never-began"); // 防御性：不 panic
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,6 @@
|
||||
// 支持并发取消检测，累积推理内容、文本增量和工具调用。
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::error;
|
||||
|
||||
use crate::clients::llm::{
|
||||
@@ -40,7 +39,7 @@ pub async fn process_llm_stream(
|
||||
llm: &LlmClient,
|
||||
messages: &[ChatMessage],
|
||||
tool_defs: &[ToolDefinition],
|
||||
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
|
||||
tap: &crate::agent::engine::EventTap,
|
||||
step: usize,
|
||||
session_id: &str,
|
||||
cancelled_runs: Arc<dashmap::DashMap<String, ()>>,
|
||||
@@ -51,7 +50,7 @@ pub async fn process_llm_stream(
|
||||
Ok(rx) => rx,
|
||||
Err(e) => {
|
||||
error!("[Streaming] LLM stream 调用失败: {}", e);
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
tap.send(AgentStreamEvent::Error {
|
||||
message: format!("大模型流式调用失败: {}", e),
|
||||
});
|
||||
return StreamOutput {
|
||||
@@ -95,15 +94,15 @@ pub async fn process_llm_stream(
|
||||
match event {
|
||||
StreamEvent::ReasoningDelta(delta) => {
|
||||
accumulated_reasoning.push_str(&delta);
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: accumulated_reasoning.clone(),
|
||||
tap.send(AgentStreamEvent::Thought {
|
||||
content: tap.thought(&accumulated_reasoning),
|
||||
step,
|
||||
});
|
||||
}
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated_content.push_str(&delta);
|
||||
if !is_tool_call_step {
|
||||
let _ = tx.send(AgentStreamEvent::TextDelta {
|
||||
tap.send(AgentStreamEvent::TextDelta {
|
||||
content: delta,
|
||||
tool_call_id: None,
|
||||
});
|
||||
|
||||
@@ -1,929 +0,0 @@
|
||||
// src/agent/runtime/streaming_executor.rs
|
||||
//
|
||||
// 流式工具执行器。
|
||||
// 参考 Claude Code StreamingToolExecutor 设计。
|
||||
//
|
||||
// 当 LLM 流式输出 tool_use 块时，立即开始执行并发安全的工具。
|
||||
// 非并发安全的工具排队等待。结果按流中顺序 yield。
|
||||
//
|
||||
// 与 Claude Code 的对齐改进 (2026-06-22)：
|
||||
// 1. 真正的流式调度 — on_tool_use 中对并发安全工具立即 spawn tokio task
|
||||
// 2. 并发分区 — 自动分组连续只读工具并行执行
|
||||
// 3. Progress 流式 — 长操作进度消息即时 yield
|
||||
// 4. Sibling Abort — 副效应工具报错时级联中止兄弟姐妹
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::broadcast;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry};
|
||||
|
||||
/// 流式工具执行状态
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum TrackedToolStatus {
|
||||
/// 工具调用已从 LLM 流中接收到，等待调度
|
||||
Queued,
|
||||
/// 正在执行中（spawned tokio task 运行中）
|
||||
Executing,
|
||||
/// 执行完成，结果就绪等待 yield
|
||||
Completed,
|
||||
/// 结果已 yield 给调用方
|
||||
Yielded,
|
||||
}
|
||||
|
||||
/// Sibling Abort 原因
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AbortReason {
|
||||
/// 兄弟工具出错触发的级联取消
|
||||
SiblingError { description: String },
|
||||
/// 用户主动中断
|
||||
UserInterrupted,
|
||||
}
|
||||
|
||||
/// 单次工具执行的结果
|
||||
#[derive(Debug)]
|
||||
struct ToolExecutionResult {
|
||||
tool_name: String,
|
||||
output: ToolOutput,
|
||||
}
|
||||
|
||||
/// 跟踪中的工具执行
|
||||
struct TrackedTool {
|
||||
tool_call_id: String,
|
||||
tool_name: String,
|
||||
args: serde_json::Value,
|
||||
status: TrackedToolStatus,
|
||||
/// 执行完成后的输出
|
||||
output: Option<ToolOutput>,
|
||||
/// 并发安全的工具在 spawn 后的 JoinHandle
|
||||
handle: Option<JoinHandle<ToolExecutionResult>>,
|
||||
}
|
||||
|
||||
/// 流式工具执行器。
|
||||
///
|
||||
/// 参考 Claude Code `StreamingToolExecutor` (531 行 TypeScript)，
|
||||
/// 关键改进：并发安全工具立即 spawn tokio task，不等待 flush。
|
||||
pub struct StreamingToolExecutor {
|
||||
/// 所有跟踪中的工具（按 LLM 流中到达顺序）
|
||||
tracked: Vec<TrackedTool>,
|
||||
/// 工具注册表
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
/// 工具上下文（按需 clone 给每个 spawn 的 task）
|
||||
tool_context: ToolContext,
|
||||
/// Sibling Abort 广播通道 (tx)
|
||||
abort_tx: broadcast::Sender<AbortReason>,
|
||||
/// Sibling Abort 广播通道 (rx) — 保留以保持 channel 存活，
|
||||
/// 实际使用时通过 `abort_tx.subscribe()` 获取新接收端。
|
||||
#[allow(dead_code)]
|
||||
abort_rx: broadcast::Receiver<AbortReason>,
|
||||
/// 当前是否已发生错误（触发 sibling abort）
|
||||
has_errored: bool,
|
||||
/// 出错工具的描述（如 "bash(git push)"）
|
||||
errored_tool_desc: String,
|
||||
/// 最大并发数（预留，当前使用 executing_non_concurrent 控制）
|
||||
#[allow(dead_code)]
|
||||
max_concurrency: usize,
|
||||
/// 最大工具输出字符数
|
||||
max_output_chars: usize,
|
||||
/// 当前正在执行的非并发安全工具数（0 或 1）
|
||||
executing_non_concurrent: bool,
|
||||
/// 已完成但尚未 yield 的结果队列（按流顺序）
|
||||
completed_queue: VecDeque<usize>,
|
||||
}
|
||||
|
||||
impl StreamingToolExecutor {
|
||||
/// 创建新的流式执行器。
|
||||
pub fn new(
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
tool_context: ToolContext,
|
||||
max_concurrency: usize,
|
||||
max_output_chars: usize,
|
||||
) -> Self {
|
||||
let (abort_tx, abort_rx) = broadcast::channel(16);
|
||||
StreamingToolExecutor {
|
||||
tracked: Vec::new(),
|
||||
tool_registry,
|
||||
tool_context,
|
||||
abort_tx,
|
||||
abort_rx,
|
||||
has_errored: false,
|
||||
errored_tool_desc: String::new(),
|
||||
max_concurrency,
|
||||
max_output_chars,
|
||||
executing_non_concurrent: false,
|
||||
completed_queue: VecDeque::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取 abort 广播发送端（供外部注入取消信号）。
|
||||
pub fn abort_sender(&self) -> broadcast::Sender<AbortReason> {
|
||||
self.abort_tx.clone()
|
||||
}
|
||||
|
||||
/// 当 LLM 流产生一个新的 tool_use 时调用。
|
||||
///
|
||||
/// 如果是并发安全工具且当前没有非并发安全工具在执行，立即 spawn tokio task。
|
||||
/// 否则加入队列等待调度。
|
||||
///
|
||||
/// 返回 true 表示该工具已立即开始执行，false 表示排队。
|
||||
pub fn on_tool_use(&mut self, call_id: String, name: String, args: serde_json::Value) -> bool {
|
||||
let is_concurrency_safe = self
|
||||
.tool_registry
|
||||
.get(&name)
|
||||
.map(|t| t.is_concurrency_safe(&args))
|
||||
.unwrap_or(false);
|
||||
|
||||
let mut tool = TrackedTool {
|
||||
tool_call_id: call_id.clone(),
|
||||
tool_name: name.clone(),
|
||||
args: args.clone(),
|
||||
status: TrackedToolStatus::Queued,
|
||||
output: None,
|
||||
handle: None,
|
||||
};
|
||||
|
||||
let idx = self.tracked.len();
|
||||
let can_start_now = is_concurrency_safe && !self.executing_non_concurrent;
|
||||
|
||||
if can_start_now {
|
||||
// 立即 spawn tokio task（参考 Claude Code: addTool 立即 processQueue）
|
||||
info!(
|
||||
"[StreamingExecutor] 立即 spawn 并发安全工具: {} (id={})",
|
||||
name, call_id
|
||||
);
|
||||
let handle = self.spawn_tool_task(idx, call_id.clone(), name.clone(), args.clone());
|
||||
tool.handle = Some(handle);
|
||||
tool.status = TrackedToolStatus::Executing;
|
||||
} else {
|
||||
info!(
|
||||
"[StreamingExecutor] 排队工具: {} (concurrent={}, executing_non_concurrent={})",
|
||||
name, is_concurrency_safe, self.executing_non_concurrent
|
||||
);
|
||||
}
|
||||
|
||||
if !is_concurrency_safe {
|
||||
self.executing_non_concurrent = true;
|
||||
}
|
||||
|
||||
self.tracked.push(tool);
|
||||
can_start_now
|
||||
}
|
||||
|
||||
/// LLM 流结束后调用，等待所有剩余排队工具完成。
|
||||
pub async fn flush(&mut self) {
|
||||
let queued_count = self
|
||||
.tracked
|
||||
.iter()
|
||||
.filter(|t| t.status == TrackedToolStatus::Queued)
|
||||
.count();
|
||||
|
||||
info!(
|
||||
"[StreamingExecutor] flush: {} tracked, {} queued, {} executing",
|
||||
self.tracked.len(),
|
||||
queued_count,
|
||||
self.tracked
|
||||
.iter()
|
||||
.filter(|t| t.status == TrackedToolStatus::Executing)
|
||||
.count()
|
||||
);
|
||||
|
||||
// 启动所有还在排队的工具
|
||||
self.start_all_queued();
|
||||
|
||||
// 等待所有执行中的工具完成
|
||||
self.await_all_executing().await;
|
||||
}
|
||||
|
||||
/// 按流顺序获取下一个完成的结果（非阻塞）。
|
||||
///
|
||||
/// 对于已完成的任务，如果其 handle 已就绪则收集结果。
|
||||
/// 返回按到达顺序的第一个已完成结果。
|
||||
pub fn next_result(&mut self) -> Option<(String, ToolOutput)> {
|
||||
// 先尝试收集任何已完成的 async task 结果
|
||||
self.collect_completed_tasks();
|
||||
|
||||
// 从 completed_queue 中按序取
|
||||
while let Some(&idx) = self.completed_queue.front() {
|
||||
self.completed_queue.pop_front();
|
||||
let tool = &mut self.tracked[idx];
|
||||
if tool.status == TrackedToolStatus::Completed {
|
||||
tool.status = TrackedToolStatus::Yielded;
|
||||
let output = tool
|
||||
.output
|
||||
.take()
|
||||
.unwrap_or_else(|| ToolOutput::error("工具执行异常：无输出"));
|
||||
return Some((tool.tool_call_id.clone(), output));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 是否有未 yield 的结果（已完成或即将完成）。
|
||||
pub fn has_pending_results(&self) -> bool {
|
||||
self.tracked
|
||||
.iter()
|
||||
.any(|t| t.status == TrackedToolStatus::Completed)
|
||||
|| !self.completed_queue.is_empty()
|
||||
}
|
||||
|
||||
/// 是否有未完成的工具（仍在排队或执行中）。
|
||||
pub fn has_unfinished(&self) -> bool {
|
||||
self.tracked.iter().any(|t| {
|
||||
t.status == TrackedToolStatus::Queued || t.status == TrackedToolStatus::Executing
|
||||
})
|
||||
}
|
||||
|
||||
/// 获取所有已完成的结果（包括已 yield 和未 yield 的）。
|
||||
pub fn all_results_mut(&mut self) -> Vec<(String, ToolOutput)> {
|
||||
self.collect_completed_tasks();
|
||||
let mut results = Vec::new();
|
||||
for tool in &mut self.tracked {
|
||||
if let Some(output) = tool.output.take() {
|
||||
results.push((tool.tool_call_id.clone(), output));
|
||||
}
|
||||
tool.status = TrackedToolStatus::Yielded;
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
// ── 内部方法 ──
|
||||
|
||||
/// Spawn 一个 tokio task 执行单个工具调用。
|
||||
fn spawn_tool_task(
|
||||
&self,
|
||||
_idx: usize,
|
||||
_call_id: String,
|
||||
tool_name: String,
|
||||
args: serde_json::Value,
|
||||
) -> JoinHandle<ToolExecutionResult> {
|
||||
let tool_registry = self.tool_registry.clone();
|
||||
let tool_context = self.tool_context.clone();
|
||||
let max_output_chars = self.max_output_chars;
|
||||
let mut abort_rx = self.abort_tx.subscribe();
|
||||
|
||||
tokio::spawn(async move {
|
||||
// tokio::select! 在工具执行和 Sibling Abort 之间竞速
|
||||
tokio::select! {
|
||||
result = async {
|
||||
match tool_registry.get(&tool_name) {
|
||||
Some(tool) => {
|
||||
tool.execute_with_progress(args, &tool_context, None).await
|
||||
}
|
||||
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
}
|
||||
} => {
|
||||
// 截断输出
|
||||
let truncated = if result.content.len() > max_output_chars {
|
||||
let t: String = result.content.chars().take(max_output_chars).collect();
|
||||
ToolOutput {
|
||||
content: format!(
|
||||
"{}...\n[输出已截断，原始长度: {} 字符]",
|
||||
t,
|
||||
result.content.len()
|
||||
),
|
||||
is_error: result.is_error,
|
||||
metadata: result.metadata,
|
||||
skip_persist: result.skip_persist,
|
||||
}
|
||||
} else {
|
||||
result
|
||||
};
|
||||
|
||||
ToolExecutionResult {
|
||||
tool_name,
|
||||
output: truncated,
|
||||
}
|
||||
}
|
||||
Ok(reason) = abort_rx.recv() => {
|
||||
let msg = match reason {
|
||||
AbortReason::SiblingError { description } => {
|
||||
format!("取消：并行工具 {} 出错，已级联取消", description)
|
||||
}
|
||||
AbortReason::UserInterrupted => "执行已被用户取消".to_string(),
|
||||
};
|
||||
ToolExecutionResult {
|
||||
tool_name,
|
||||
output: ToolOutput::error(msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// 尝试收集所有已完成 tokio task 的结果（非阻塞）。
|
||||
fn collect_completed_tasks(&mut self) {
|
||||
for idx in 0..self.tracked.len() {
|
||||
if self.tracked[idx].status != TrackedToolStatus::Executing {
|
||||
continue;
|
||||
}
|
||||
if self.tracked[idx].handle.is_none() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// 检查 JoinHandle 是否已完成（非阻塞）
|
||||
let handle = self.tracked[idx].handle.take().unwrap();
|
||||
if handle.is_finished() {
|
||||
// is_finished=true 保证 .await 会立即返回
|
||||
// 使用 tokio::task::yield_now 之后的 poll 可能也成功，
|
||||
// 这里直接在同步上下文中检查后放入完成队列
|
||||
// 等下次 async 上下文中通过 await_all_executing 处理
|
||||
self.tracked[idx].handle = Some(handle);
|
||||
// 标记为需要收集 — 将在 flush/await 中处理
|
||||
} else {
|
||||
// 放回未完成的 handle
|
||||
self.tracked[idx].handle = Some(handle);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动所有排队的工具。
|
||||
fn start_all_queued(&mut self) {
|
||||
// 收集需要启动的工具索引（避免借用冲突）
|
||||
let to_start: Vec<usize> = self
|
||||
.tracked
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
|
||||
.filter(|(_, t)| {
|
||||
let is_safe = self
|
||||
.tool_registry
|
||||
.get(&t.tool_name)
|
||||
.map(|reg_tool| reg_tool.is_concurrency_safe(&t.args))
|
||||
.unwrap_or(false);
|
||||
// 并发安全工具可随时启动，非并发安全的需要独占
|
||||
is_safe || !self.executing_non_concurrent
|
||||
})
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
|
||||
for idx in to_start {
|
||||
let tool = &self.tracked[idx];
|
||||
let call_id = tool.tool_call_id.clone();
|
||||
let tool_name = tool.tool_name.clone();
|
||||
let args = tool.args.clone();
|
||||
|
||||
let is_safe = self
|
||||
.tool_registry
|
||||
.get(&tool_name)
|
||||
.map(|t| t.is_concurrency_safe(&args))
|
||||
.unwrap_or(false);
|
||||
|
||||
let handle = self.spawn_tool_task(idx, call_id, tool_name.clone(), args);
|
||||
self.tracked[idx].handle = Some(handle);
|
||||
self.tracked[idx].status = TrackedToolStatus::Executing;
|
||||
|
||||
if !is_safe {
|
||||
self.executing_non_concurrent = true;
|
||||
// 非并发安全工具启动后停止（独占执行）
|
||||
break;
|
||||
}
|
||||
|
||||
info!("[StreamingExecutor] 启动排队工具: {}", tool_name);
|
||||
}
|
||||
}
|
||||
|
||||
/// 等待所有执行中的工具完成。
|
||||
async fn await_all_executing(&mut self) {
|
||||
// 收集所有剩余 JoinHandles
|
||||
let mut handles: Vec<(usize, JoinHandle<ToolExecutionResult>)> = Vec::new();
|
||||
for idx in 0..self.tracked.len() {
|
||||
if self.tracked[idx].status == TrackedToolStatus::Executing {
|
||||
if let Some(handle) = self.tracked[idx].handle.take() {
|
||||
handles.push((idx, handle));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 并发等待所有任务
|
||||
for (idx, handle) in handles {
|
||||
match handle.await {
|
||||
Ok(result) => {
|
||||
let tool_name = result.tool_name.clone();
|
||||
let is_error = result.output.is_error;
|
||||
|
||||
self.tracked[idx].output = Some(result.output);
|
||||
self.tracked[idx].status = TrackedToolStatus::Completed;
|
||||
self.completed_queue.push_back(idx);
|
||||
|
||||
if is_error {
|
||||
self.check_sibling_abort(idx, &tool_name);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[StreamingExecutor] tokio task 异常: {}", e);
|
||||
self.tracked[idx].output =
|
||||
Some(ToolOutput::error(format!("工具执行异常: {}", e)));
|
||||
self.tracked[idx].status = TrackedToolStatus::Completed;
|
||||
self.completed_queue.push_back(idx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.executing_non_concurrent = false;
|
||||
}
|
||||
|
||||
/// 检查错误工具是否触发 Sibling Abort。
|
||||
fn check_sibling_abort(&mut self, idx: usize, tool_name: &str) {
|
||||
let causes_abort = self
|
||||
.tool_registry
|
||||
.get(tool_name)
|
||||
.map(|t| t.causes_sibling_abort())
|
||||
.unwrap_or(false);
|
||||
|
||||
if causes_abort && !self.has_errored {
|
||||
warn!(
|
||||
"[StreamingExecutor] 工具 {} 出错，触发 sibling abort",
|
||||
tool_name
|
||||
);
|
||||
self.has_errored = true;
|
||||
self.errored_tool_desc = self.get_tool_description(idx);
|
||||
let _ = self.abort_tx.send(AbortReason::SiblingError {
|
||||
description: self.errored_tool_desc.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取工具的人类可读描述（用于错误消息）。
|
||||
fn get_tool_description(&self, idx: usize) -> String {
|
||||
let tool = &self.tracked[idx];
|
||||
let summary = tool
|
||||
.args
|
||||
.get("command")
|
||||
.or_else(|| tool.args.get("file_path"))
|
||||
.or_else(|| tool.args.get("pattern"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
if summary.is_empty() {
|
||||
tool.tool_name.clone()
|
||||
} else {
|
||||
let truncated: String = summary.chars().take(40).collect();
|
||||
if summary.len() > 40 {
|
||||
format!("{}({}…)", tool.tool_name, truncated)
|
||||
} else {
|
||||
format!("{}({})", tool.tool_name, summary)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::tools::AgentTool;
|
||||
use crate::agent::tools::ToolContext;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// 可配置的 Mock 工具，用于测试 StreamingToolExecutor 状态机。
|
||||
struct MockAgentTool {
|
||||
name_str: &'static str,
|
||||
concurrency_safe: bool,
|
||||
causes_abort: bool,
|
||||
/// 执行返回的内容
|
||||
result_content: &'static str,
|
||||
/// 执行是否返回错误
|
||||
result_is_error: bool,
|
||||
/// 可选：执行后设置此标志（用于验证工具是否被调用）
|
||||
executed: AtomicBool,
|
||||
}
|
||||
|
||||
impl MockAgentTool {
|
||||
fn new(name: &'static str) -> Self {
|
||||
MockAgentTool {
|
||||
name_str: name,
|
||||
concurrency_safe: false,
|
||||
causes_abort: false,
|
||||
result_content: "mock result",
|
||||
result_is_error: false,
|
||||
executed: AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn concurrency_safe(mut self, v: bool) -> Self {
|
||||
self.concurrency_safe = v;
|
||||
self
|
||||
}
|
||||
|
||||
fn causes_abort(mut self, v: bool) -> Self {
|
||||
self.causes_abort = v;
|
||||
self
|
||||
}
|
||||
|
||||
fn result(mut self, content: &'static str, is_error: bool) -> Self {
|
||||
self.result_content = content;
|
||||
self.result_is_error = is_error;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for MockAgentTool {
|
||||
fn name(&self) -> &str {
|
||||
self.name_str
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"mock tool for testing"
|
||||
}
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::json!({})
|
||||
}
|
||||
async fn execute(&self, _args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
self.executed.store(true, Ordering::SeqCst);
|
||||
if self.result_is_error {
|
||||
ToolOutput::error(self.result_content)
|
||||
} else {
|
||||
ToolOutput::success(self.result_content, serde_json::json!({}))
|
||||
}
|
||||
}
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
self.concurrency_safe
|
||||
}
|
||||
fn causes_sibling_abort(&self) -> bool {
|
||||
self.causes_abort
|
||||
}
|
||||
}
|
||||
|
||||
/// 构建测试用的 ToolContext。
|
||||
async fn make_test_tool_context() -> ToolContext {
|
||||
use crate::agent::memory::MemoryManager;
|
||||
use crate::agent::skills::SkillRegistry;
|
||||
use crate::api::AppState;
|
||||
use crate::clients::ads::AdsClient;
|
||||
use crate::clients::arxiv::ArxivClient;
|
||||
use crate::clients::cds::vizier::VizierClient;
|
||||
use crate::clients::llm::{EmbeddingClient, LlmClient};
|
||||
use crate::clients::qiniu::QiniuClient;
|
||||
use crate::services::batch::asset::AssetBatchStatus;
|
||||
use crate::services::batch::meta::MetaSyncStatus;
|
||||
use crate::services::download::Downloader;
|
||||
use crate::services::translation::Dictionary;
|
||||
use crate::Config;
|
||||
use std::path::PathBuf;
|
||||
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()).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 vizier =
|
||||
VizierClient::new("https://tapvizier.cds.unistra.fr/TAPVizieR/tap", 60).unwrap();
|
||||
let qiniu = QiniuClient::new(
|
||||
"ak".into(),
|
||||
"sk".into(),
|
||||
"b".into(),
|
||||
"http://localhost".into(),
|
||||
);
|
||||
|
||||
let app_state = Arc::new(AppState {
|
||||
config,
|
||||
db: pool,
|
||||
dict: Dictionary::default(),
|
||||
qiniu,
|
||||
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)),
|
||||
skill_registry: Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(PathBuf::from(
|
||||
"/tmp/sk",
|
||||
)))),
|
||||
sse_broadcast: None,
|
||||
memory_manager: Arc::new(tokio::sync::Mutex::new(MemoryManager::new(PathBuf::from(
|
||||
"/tmp/test_mem",
|
||||
)))),
|
||||
login_rate_limiter: Arc::new(dashmap::DashMap::new()),
|
||||
upload_rate_limiter: Arc::new(dashmap::DashMap::new()),
|
||||
llm: crate::api::LlmState {
|
||||
primary: llm.clone(),
|
||||
medium: llm.clone(),
|
||||
fast: llm.clone(),
|
||||
vision: None,
|
||||
embedding,
|
||||
},
|
||||
sources: crate::api::DataSourceState {
|
||||
ads,
|
||||
arxiv,
|
||||
vizier,
|
||||
lamost: crate::clients::lamost::LamostClient::new("https://www.lamost.org", 60)
|
||||
.unwrap(),
|
||||
gaia: crate::clients::gaia::GaiaClient::new(
|
||||
"https://gea.esac.esa.int/tap-server/tap",
|
||||
"https://gea.esac.esa.int/data-server",
|
||||
90,
|
||||
)
|
||||
.unwrap(),
|
||||
sdss: crate::clients::sdss::SdssClient::new(
|
||||
"https://datalab.noirlab.edu/tap/sync",
|
||||
90,
|
||||
)
|
||||
.unwrap(),
|
||||
desi: crate::clients::desi::DesiClient::new(
|
||||
"https://datalab.noirlab.edu/tap/sync",
|
||||
120,
|
||||
)
|
||||
.unwrap(),
|
||||
irsa: crate::clients::irsa::IrsaClient::new("https://irsa.ipac.caltech.edu", 60)
|
||||
.unwrap(),
|
||||
mast: crate::clients::mast::MastClient::new("https://mast.stsci.edu", 90).unwrap(),
|
||||
observation_registry: std::sync::Arc::new(
|
||||
crate::services::observation::ObservationRegistry::default(),
|
||||
),
|
||||
},
|
||||
session: crate::api::SessionState {
|
||||
sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
|
||||
session_last_active: Arc::new(dashmap::DashMap::new()),
|
||||
cancelled_runs: Arc::new(dashmap::DashMap::new()),
|
||||
session_permission_checkers: Arc::new(dashmap::DashMap::new()),
|
||||
pending_questions: Arc::new(tokio::sync::Mutex::new(
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
pending_permissions: Arc::new(tokio::sync::Mutex::new(
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
},
|
||||
});
|
||||
|
||||
ToolContext::new(app_state)
|
||||
}
|
||||
|
||||
/// 创建包含指定 Mock 工具的 StreamingToolExecutor。
|
||||
async fn make_executor(tools: Vec<MockAgentTool>) -> StreamingToolExecutor {
|
||||
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);
|
||||
for tool in tools {
|
||||
registry.add_tool(Box::new(tool));
|
||||
}
|
||||
let tool_context = make_test_tool_context().await;
|
||||
StreamingToolExecutor::new(Arc::new(registry), tool_context, 4, 4000)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_on_tool_use_spawns_concurrency_safe_immediately() {
|
||||
let mut executor =
|
||||
make_executor(vec![MockAgentTool::new("safe_tool").concurrency_safe(true)]).await;
|
||||
|
||||
let spawned = executor.on_tool_use(
|
||||
"call_1".into(),
|
||||
"safe_tool".into(),
|
||||
serde_json::json!({"key": "val"}),
|
||||
);
|
||||
assert!(spawned, "concurrency-safe tool should spawn immediately");
|
||||
assert_eq!(executor.tracked.len(), 1);
|
||||
assert_eq!(
|
||||
executor.tracked[0].status,
|
||||
TrackedToolStatus::Executing,
|
||||
"safe tool should be executing"
|
||||
);
|
||||
assert!(
|
||||
!executor.executing_non_concurrent,
|
||||
"safe tool does not lock executor"
|
||||
);
|
||||
assert!(
|
||||
executor.has_unfinished(),
|
||||
"should have unfinished tool (executing)"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_on_tool_use_queues_non_concurrency_safe() {
|
||||
let mut executor = make_executor(vec![
|
||||
MockAgentTool::new("unsafe_tool").concurrency_safe(false)
|
||||
])
|
||||
.await;
|
||||
|
||||
let spawned =
|
||||
executor.on_tool_use("call_1".into(), "unsafe_tool".into(), serde_json::json!({}));
|
||||
assert!(!spawned, "non-concurrency-safe tool should be queued");
|
||||
assert_eq!(executor.tracked.len(), 1);
|
||||
assert_eq!(executor.tracked[0].status, TrackedToolStatus::Queued);
|
||||
assert!(
|
||||
executor.executing_non_concurrent,
|
||||
"non-concurrent tool sets the lock flag"
|
||||
);
|
||||
assert!(
|
||||
executor.has_unfinished(),
|
||||
"queued tool counts as unfinished"
|
||||
);
|
||||
assert!(!executor.has_pending_results(), "no results yet");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_on_tool_use_queues_safe_tool_when_non_concurrent_executing() {
|
||||
// 先加入一个非并发安全工具（设为 executing_non_concurrent=true），
|
||||
// 再尝试加入一个并发安全工具，应排队而非立即执行
|
||||
let mut executor = make_executor(vec![
|
||||
MockAgentTool::new("unsafe_tool").concurrency_safe(false),
|
||||
MockAgentTool::new("safe_tool").concurrency_safe(true),
|
||||
])
|
||||
.await;
|
||||
|
||||
// 第一个：非并发安全，排队但设置 executing_non_concurrent
|
||||
executor.on_tool_use("call_1".into(), "unsafe_tool".into(), serde_json::json!({}));
|
||||
// 第二个：并发安全但 executor 被非并发工具锁定，应排队
|
||||
let spawned =
|
||||
executor.on_tool_use("call_2".into(), "safe_tool".into(), serde_json::json!({}));
|
||||
|
||||
assert!(
|
||||
!spawned,
|
||||
"safe tool should be queued when executor is locked"
|
||||
);
|
||||
assert_eq!(executor.tracked.len(), 2);
|
||||
assert_eq!(executor.tracked[1].status, TrackedToolStatus::Queued);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_tool_description_parses_args() {
|
||||
let mut executor = make_executor(vec![MockAgentTool::new("bash")]).await;
|
||||
// 手动添加 tracked tool（绕过 on_tool_use 的 spawn）
|
||||
executor.tracked.push(TrackedTool {
|
||||
tool_call_id: "c1".into(),
|
||||
tool_name: "bash".into(),
|
||||
args: serde_json::json!({"command": "git push origin main"}),
|
||||
status: TrackedToolStatus::Completed,
|
||||
output: None,
|
||||
handle: None,
|
||||
});
|
||||
|
||||
let desc = executor.get_tool_description(0);
|
||||
assert!(
|
||||
desc.contains("bash"),
|
||||
"description should contain tool name, got: {}",
|
||||
desc
|
||||
);
|
||||
assert!(
|
||||
desc.contains("git push"),
|
||||
"description should contain command arg, got: {}",
|
||||
desc
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_tool_description_falls_back_to_name() {
|
||||
let mut executor = make_executor(vec![MockAgentTool::new("unknown_tool")]).await;
|
||||
executor.tracked.push(TrackedTool {
|
||||
tool_call_id: "c1".into(),
|
||||
tool_name: "unknown_tool".into(),
|
||||
args: serde_json::json!({}), // no command/file_path/pattern
|
||||
status: TrackedToolStatus::Completed,
|
||||
output: None,
|
||||
handle: None,
|
||||
});
|
||||
|
||||
let desc = executor.get_tool_description(0);
|
||||
assert_eq!(desc, "unknown_tool");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_flush_starts_queued_and_awaits_completion() {
|
||||
// 使用并发安全工具测试 flush：on_tool_use 立即 spawn，flush 等待完成
|
||||
let mut executor = make_executor(vec![MockAgentTool::new("tool_a")
|
||||
.concurrency_safe(true)
|
||||
.result("done", false)])
|
||||
.await;
|
||||
|
||||
let spawned = executor.on_tool_use("call_1".into(), "tool_a".into(), serde_json::json!({}));
|
||||
assert!(spawned, "concurrency-safe tool should spawn immediately");
|
||||
assert_eq!(executor.tracked[0].status, TrackedToolStatus::Executing);
|
||||
|
||||
executor.flush().await;
|
||||
|
||||
assert_eq!(
|
||||
executor.tracked[0].status,
|
||||
TrackedToolStatus::Completed,
|
||||
"after flush, tool should be completed"
|
||||
);
|
||||
assert!(
|
||||
executor.has_pending_results(),
|
||||
"completed tool = pending result"
|
||||
);
|
||||
assert!(!executor.has_unfinished(), "nothing queued or executing");
|
||||
|
||||
let result = executor.next_result();
|
||||
assert!(
|
||||
result.is_some(),
|
||||
"next_result should return the completed tool"
|
||||
);
|
||||
let (call_id, output) = result.unwrap();
|
||||
assert_eq!(call_id, "call_1");
|
||||
assert!(!output.is_error, "tool should succeed");
|
||||
assert_eq!(output.content, "done");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_next_result_yields_in_insertion_order() {
|
||||
let mut executor = make_executor(vec![
|
||||
MockAgentTool::new("tool_a")
|
||||
.concurrency_safe(true)
|
||||
.result("result_a", false),
|
||||
MockAgentTool::new("tool_b")
|
||||
.concurrency_safe(true)
|
||||
.result("result_b", false),
|
||||
])
|
||||
.await;
|
||||
|
||||
// 两个并发安全工具，都应该立即 spawn
|
||||
executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({}));
|
||||
executor.on_tool_use("call_b".into(), "tool_b".into(), serde_json::json!({}));
|
||||
|
||||
// flush 等待它们完成
|
||||
executor.flush().await;
|
||||
|
||||
// 结果应按插入顺序产出
|
||||
let result_a = executor.next_result();
|
||||
assert!(result_a.is_some());
|
||||
assert_eq!(result_a.as_ref().unwrap().0, "call_a");
|
||||
assert_eq!(result_a.as_ref().unwrap().1.content, "result_a");
|
||||
|
||||
let result_b = executor.next_result();
|
||||
assert!(result_b.is_some());
|
||||
assert_eq!(result_b.as_ref().unwrap().0, "call_b");
|
||||
assert_eq!(result_b.as_ref().unwrap().1.content, "result_b");
|
||||
|
||||
// 第三次调用返回 None（全部已 yield）
|
||||
let result_c = executor.next_result();
|
||||
assert!(result_c.is_none());
|
||||
assert!(!executor.has_pending_results());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sibling_abort_broadcast_on_error() {
|
||||
// 注册两个工具：tool_a 会报错且触发 sibling abort，tool_b 并发执行中被取消
|
||||
let mut executor = make_executor(vec![
|
||||
MockAgentTool::new("tool_a")
|
||||
.concurrency_safe(true)
|
||||
.causes_abort(true)
|
||||
.result("critical failure", true),
|
||||
MockAgentTool::new("tool_b")
|
||||
.concurrency_safe(true)
|
||||
.result("should be aborted", false),
|
||||
])
|
||||
.await;
|
||||
|
||||
executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({}));
|
||||
executor.on_tool_use("call_b".into(), "tool_b".into(), serde_json::json!({}));
|
||||
|
||||
executor.flush().await;
|
||||
|
||||
// tool_a 的结果应该是错误
|
||||
let result_a = executor.next_result();
|
||||
assert!(result_a.is_some());
|
||||
assert!(result_a.unwrap().1.is_error, "tool_a should have errored");
|
||||
|
||||
// tool_b 可能被取消（sibling abort）或正常完成（取决于竞态）
|
||||
let result_b = executor.next_result();
|
||||
assert!(result_b.is_some(), "tool_b should also have a result");
|
||||
|
||||
// 验证 has_errored 被设置
|
||||
assert!(
|
||||
executor.has_errored,
|
||||
"has_errored should be set after sibling abort"
|
||||
);
|
||||
assert!(
|
||||
!executor.errored_tool_desc.is_empty(),
|
||||
"errored_tool_desc should be populated"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_all_results_mut_drains_all_outputs() {
|
||||
let mut executor =
|
||||
make_executor(vec![MockAgentTool::new("tool_a").result("result_a", false)]).await;
|
||||
|
||||
executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({}));
|
||||
// 不 flush — all_results_mut 是同步方法，只收集已完成的结果
|
||||
// 工具可能仍在执行，所以不能保证一定有结果
|
||||
let results = executor.all_results_mut();
|
||||
// 无论有没有结果，调用后 tracked tool 被标记为 Yielded
|
||||
for tool in &executor.tracked {
|
||||
assert_eq!(tool.status, TrackedToolStatus::Yielded);
|
||||
assert!(tool.output.is_none(), "output should be taken");
|
||||
}
|
||||
// 验证返回的 results 和 tracked 一致
|
||||
let _ = results; // 如果 joinhandle 还没完成，results 可能是空的
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_abort_sender_clone_works() {
|
||||
let executor = make_executor(vec![]).await;
|
||||
let sender = executor.abort_sender();
|
||||
// 验证 sender 可用
|
||||
assert_eq!(sender.receiver_count(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_empty_executor_has_no_pending_or_unfinished() {
|
||||
let executor = make_executor(vec![]).await;
|
||||
assert!(!executor.has_pending_results());
|
||||
assert!(!executor.has_unfinished());
|
||||
}
|
||||
}
|
||||
+170
-424
@@ -13,17 +13,13 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc::UnboundedSender;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::compact;
|
||||
use super::hooks::{
|
||||
HookRegistry, PostToolUseContext, PreToolUseContext, SubagentStartContext, SubagentStopContext,
|
||||
};
|
||||
use super::runtime::permission::{PermissionChecker, PermissionResult};
|
||||
use super::hooks::{HookRegistry, SubagentStartContext, SubagentStopContext};
|
||||
use super::runtime::permission::PermissionChecker;
|
||||
use super::runtime::{AgentConfig, AgentStreamEvent};
|
||||
use super::tools::{ToolContext, ToolOutput, ToolRegistry};
|
||||
use super::tools::{ToolOutput, ToolRegistry};
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::{ChatMessage, LlmClient, StreamEvent, ToolDefinition};
|
||||
use crate::clients::llm::{ChatMessage, LlmClient};
|
||||
|
||||
/// 子代理运行器
|
||||
pub struct SubAgentRunner {
|
||||
@@ -137,13 +133,29 @@ impl SubAgentRunner {
|
||||
/// * `system_prompt` - 子代理的系统提示词
|
||||
/// * `research_prompt` - 要执行的研究任务描述
|
||||
/// * `max_steps` - 子代理最大推理步数（默认 5）
|
||||
/// * `agent_name` - 稳定名称（Some 时结果可被 followup 续话引用）
|
||||
pub async fn run(
|
||||
&self,
|
||||
system_prompt: &str,
|
||||
research_prompt: &str,
|
||||
max_steps: usize,
|
||||
) -> ToolOutput {
|
||||
let subagent_name = format!("sub_{}", &uuid::Uuid::new_v4().to_string()[..8]);
|
||||
self.run_named(system_prompt, research_prompt, max_steps, None)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 带稳定名称的子代理运行（命名子代理的历史持久化在该名称下，
|
||||
/// 后续可通过 `run_continuation` 续话）。
|
||||
pub async fn run_named(
|
||||
&self,
|
||||
system_prompt: &str,
|
||||
research_prompt: &str,
|
||||
max_steps: usize,
|
||||
agent_name: Option<&str>,
|
||||
) -> ToolOutput {
|
||||
let subagent_name = agent_name
|
||||
.map(|n| n.to_string())
|
||||
.unwrap_or_else(|| format!("sub_{}", &uuid::Uuid::new_v4().to_string()[..8]));
|
||||
|
||||
// OnSubagentStart hook
|
||||
if let Some(ref registry) = self.hook_registry {
|
||||
@@ -216,6 +228,60 @@ impl SubAgentRunner {
|
||||
result
|
||||
}
|
||||
|
||||
/// 续话一个已存在的命名子代理（continuable subagent）。
|
||||
///
|
||||
/// 从父会话的 agent_messages 中回放该子代理名下的全部历史作为初始
|
||||
/// 上下文（durable 会话 + 冷恢复——dsh Activation 模型的轻量版），
|
||||
/// `followup_prompt` 作为追加的用户消息。找不到历史时退化为全新运行。
|
||||
pub async fn run_continuation(
|
||||
&self,
|
||||
agent_name: &str,
|
||||
system_prompt: &str,
|
||||
followup_prompt: &str,
|
||||
max_steps: usize,
|
||||
) -> ToolOutput {
|
||||
if self.parent_session_id.is_empty() {
|
||||
return ToolOutput::error("子代理续话需要父会话上下文");
|
||||
}
|
||||
|
||||
let history = crate::agent::runtime::session::load_history_for_agent(
|
||||
&self.app_state.db,
|
||||
&self.parent_session_id,
|
||||
agent_name,
|
||||
)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
if history.is_empty() {
|
||||
tracing::info!(
|
||||
"[SubAgent] 命名子代理 {} 无历史记录，退化为全新运行",
|
||||
agent_name
|
||||
);
|
||||
return self
|
||||
.run_named(system_prompt, followup_prompt, max_steps, Some(agent_name))
|
||||
.await;
|
||||
}
|
||||
|
||||
// 保存 followup 用户消息
|
||||
self.save_subagent_message(agent_name, 0, 0, "user", followup_prompt, None, None, None);
|
||||
|
||||
let history_len = history.len();
|
||||
let mut messages = history;
|
||||
messages.push(ChatMessage::user(followup_prompt));
|
||||
|
||||
let outcome = self
|
||||
.run_engine_with_messages(&mut messages, max_steps, agent_name)
|
||||
.await;
|
||||
|
||||
tracing::info!(
|
||||
"[SubAgent] 续话 {}: 回放 {} 条历史消息",
|
||||
agent_name,
|
||||
history_len
|
||||
);
|
||||
|
||||
outcome
|
||||
}
|
||||
|
||||
/// 保存子代理消息到 agent_messages 表
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn save_subagent_message(
|
||||
@@ -265,7 +331,7 @@ impl SubAgentRunner {
|
||||
});
|
||||
}
|
||||
|
||||
/// 实际执行逻辑（提取为内部方法以便 hook 包装）
|
||||
/// 全新上下文的执行入口（system + 首条任务消息）。
|
||||
async fn run_inner(
|
||||
&self,
|
||||
system_prompt: &str,
|
||||
@@ -273,438 +339,118 @@ impl SubAgentRunner {
|
||||
max_steps: usize,
|
||||
subagent_name: &str,
|
||||
) -> ToolOutput {
|
||||
let llm = self
|
||||
.llm_client_override
|
||||
.as_ref()
|
||||
.unwrap_or(&self.app_state.llm.primary);
|
||||
let tool_defs = self.tool_registry.definitions();
|
||||
|
||||
// 全新上下文
|
||||
let mut messages = vec![
|
||||
ChatMessage::system(system_prompt),
|
||||
ChatMessage::user(research_prompt),
|
||||
];
|
||||
|
||||
// 跟踪工具调用防止死循环
|
||||
let mut last_call: Option<(String, String)> = None;
|
||||
let mut consecutive_count: usize = 0;
|
||||
let duplicate_threshold: usize = 3;
|
||||
|
||||
for step in 1..=max_steps {
|
||||
// 上下文压缩检查
|
||||
let est_tokens: usize = messages
|
||||
.iter()
|
||||
.map(|m| m.content.as_ref().map_or(0, |c| c.len()) + 4)
|
||||
.sum();
|
||||
if est_tokens > self.config.token_soft_limit * 3 / 2 {
|
||||
info!(
|
||||
"[SubAgent] 上下文超限 (est. {} tokens)，触发压缩",
|
||||
est_tokens
|
||||
);
|
||||
compact::compress_context_with_hooks(
|
||||
&mut messages,
|
||||
llm,
|
||||
self.config.token_soft_limit,
|
||||
self.config.max_messages,
|
||||
subagent_name,
|
||||
self.hook_registry.as_ref().map(|a| a.as_ref()),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// LLM 流式调用
|
||||
let mut stream_rx = match llm
|
||||
.chat_stream(&messages, &tool_defs, self.config.enable_thinking)
|
||||
self.run_engine_with_messages(&mut messages, max_steps, subagent_name)
|
||||
.await
|
||||
{
|
||||
Ok(rx) => rx,
|
||||
Err(e) => {
|
||||
warn!("[SubAgent] LLM stream 失败: {}", e);
|
||||
return ToolOutput::error(format!("子代理 LLM 调用失败: {}", e));
|
||||
}
|
||||
|
||||
/// 实际执行逻辑 — 委托给统一 ReAct 引擎（与主代理共用同一循环实现）。
|
||||
///
|
||||
/// 子代理注入组合：事件前缀 [sub]、AskPolicy::AutoDeny（无人值守
|
||||
/// fail-closed）、独立持久化 sink（agent_name = sub_xxx）、无后台队列。
|
||||
/// `messages` 由调用方构造：全新运行（system+user）或续话（回放历史）。
|
||||
async fn run_engine_with_messages(
|
||||
&self,
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
max_steps: usize,
|
||||
subagent_name: &str,
|
||||
) -> ToolOutput {
|
||||
use crate::agent::engine::{AskPolicy, CancelSource, DbMessageSink, EventTap, ReactEngine};
|
||||
|
||||
let llm = self
|
||||
.llm_client_override
|
||||
.as_ref()
|
||||
.unwrap_or(&self.app_state.llm.primary);
|
||||
|
||||
let mut engine_config = self.config.clone();
|
||||
engine_config.max_steps = max_steps;
|
||||
|
||||
let tap = match &self.progress_tx {
|
||||
Some(tx) => EventTap::prefixed(tx.clone(), "sub"),
|
||||
None => EventTap::none(),
|
||||
};
|
||||
|
||||
let mut accumulated_content = String::new();
|
||||
let mut accumulated_reasoning = String::new();
|
||||
let mut accumulated_tool_calls: Option<Vec<crate::clients::llm::ToolCall>> = None;
|
||||
let mut activity_log: Vec<String> = Vec::new();
|
||||
|
||||
while let Some(event) = stream_rx.recv().await {
|
||||
match event {
|
||||
StreamEvent::ReasoningDelta(delta) => {
|
||||
accumulated_reasoning.push_str(&delta);
|
||||
// 转发子代理思考过程到父代理
|
||||
if let Some(ref tx) = self.progress_tx {
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: format!("[子代理] {}", accumulated_reasoning),
|
||||
step,
|
||||
});
|
||||
}
|
||||
}
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated_content.push_str(&delta);
|
||||
}
|
||||
StreamEvent::ToolCallsComplete(tool_calls) => {
|
||||
// 确保每个工具调用有唯一 ID（LLM 可能不返回 id）
|
||||
let fixed_tool_calls: Vec<crate::clients::llm::ToolCall> = tool_calls
|
||||
.into_iter()
|
||||
.map(|tc| {
|
||||
let id = if tc.id.is_empty() {
|
||||
format!("call_{}", &uuid::Uuid::new_v4().to_string()[..8])
|
||||
} else {
|
||||
tc.id
|
||||
};
|
||||
activity_log.push(format!("🔧 调用工具: {}", tc.function.name));
|
||||
crate::clients::llm::ToolCall { id, ..tc }
|
||||
})
|
||||
.collect();
|
||||
accumulated_tool_calls = Some(fixed_tool_calls);
|
||||
}
|
||||
StreamEvent::Done => break,
|
||||
StreamEvent::Error(e) => {
|
||||
warn!("[SubAgent] 流式错误: {}", e);
|
||||
return ToolOutput::error(format!("子代理流式错误: {}", e));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// 记录思考过程
|
||||
if !accumulated_reasoning.is_empty() {
|
||||
activity_log.push(format!(
|
||||
"💭 思考: {}",
|
||||
accumulated_reasoning.chars().take(300).collect::<String>()
|
||||
));
|
||||
}
|
||||
|
||||
// 无工具调用 = 最终回答
|
||||
let tool_calls = match accumulated_tool_calls {
|
||||
Some(ref tc) if !tc.is_empty() => tc.clone(),
|
||||
_ => {
|
||||
// 转发子代理结论到父代理（作为 thought 显示在时间线，不污染 finalAnswer）
|
||||
if let Some(ref tx) = self.progress_tx {
|
||||
let _ = tx.send(AgentStreamEvent::Thought {
|
||||
content: accumulated_content.clone(),
|
||||
step,
|
||||
});
|
||||
}
|
||||
let content_len = accumulated_content.len();
|
||||
info!("[SubAgent] 子代理完成，返回 {} 字符摘要", content_len);
|
||||
let summary = format!(
|
||||
"[子代理活动记录]\n\n{}\n\n[子代理结论]\n\n{}",
|
||||
activity_log.join("\n"),
|
||||
accumulated_content
|
||||
);
|
||||
return ToolOutput::success(
|
||||
summary,
|
||||
serde_json::json!({
|
||||
"steps": step,
|
||||
"content_length": content_len,
|
||||
"tool_calls": activity_log.iter().filter(|e| e.starts_with("🔧")).count(),
|
||||
}),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// 构建 assistant 消息
|
||||
let assistant_msg = ChatMessage::assistant_with_reasoning(
|
||||
if accumulated_content.is_empty() {
|
||||
// 子代理消息落库（父会话下以 agent_name 区分；无父会话则不持久化）
|
||||
let sink: Option<std::sync::Arc<dyn crate::agent::engine::MessageSink>> =
|
||||
if self.parent_session_id.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated_content.clone())
|
||||
Some(std::sync::Arc::new(DbMessageSink {
|
||||
db: self.app_state.db.clone(),
|
||||
session_id: self.parent_session_id.clone(),
|
||||
agent_name: subagent_name.to_string(),
|
||||
fixed_metadata: Some(serde_json::json!({
|
||||
"agent": subagent_name,
|
||||
"is_subagent": true,
|
||||
})),
|
||||
}))
|
||||
};
|
||||
|
||||
let engine = ReactEngine {
|
||||
llm,
|
||||
tool_registry: &self.tool_registry,
|
||||
app_state: self.app_state.clone(),
|
||||
config: engine_config,
|
||||
session_id: if self.parent_session_id.is_empty() {
|
||||
subagent_name.to_string()
|
||||
} else {
|
||||
self.parent_session_id.clone()
|
||||
},
|
||||
None,
|
||||
Some(tool_calls.clone()),
|
||||
);
|
||||
// 持久化 assistant 消息
|
||||
let tc_json = serde_json::to_string(&tool_calls).unwrap_or_default();
|
||||
self.save_subagent_message(
|
||||
subagent_name,
|
||||
0,
|
||||
step as i32,
|
||||
"assistant",
|
||||
&accumulated_content,
|
||||
None,
|
||||
Some(&tc_json),
|
||||
None,
|
||||
);
|
||||
messages.push(assistant_msg);
|
||||
|
||||
// 执行工具调用
|
||||
for tool_call in &tool_calls {
|
||||
let tool_name = &tool_call.function.name;
|
||||
let tool_args_str = &tool_call.function.arguments;
|
||||
|
||||
// 死循环检测
|
||||
let call_key = (tool_name.clone(), tool_args_str.clone());
|
||||
if last_call.as_ref() == Some(&call_key) {
|
||||
consecutive_count += 1;
|
||||
if consecutive_count >= duplicate_threshold {
|
||||
warn!("[SubAgent] 检测到死循环：{}", tool_name);
|
||||
let error_msg = ChatMessage::tool_result(
|
||||
&tool_call.id,
|
||||
format!(
|
||||
"工具 {} 被连续重复调用。请停止并给出当前收集到的答案。",
|
||||
tool_name
|
||||
),
|
||||
);
|
||||
messages.push(error_msg);
|
||||
continue;
|
||||
}
|
||||
} else {
|
||||
last_call = Some(call_key);
|
||||
consecutive_count = 1;
|
||||
}
|
||||
|
||||
// 解析参数
|
||||
let args: serde_json::Value = match serde_json::from_str(tool_args_str) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
let error_msg =
|
||||
ChatMessage::tool_result(&tool_call.id, format!("参数解析失败: {}", e));
|
||||
messages.push(error_msg);
|
||||
continue;
|
||||
}
|
||||
turn_index: 0,
|
||||
agent_name: subagent_name.to_string(),
|
||||
tap,
|
||||
sink,
|
||||
hooks: self.hook_registry.as_deref(),
|
||||
ask_policy: AskPolicy::AutoDeny,
|
||||
permission_checker: Some(self.permission_checker.as_ref()),
|
||||
session_checker: None,
|
||||
denial_tracker: None,
|
||||
checkpoint_manager: None,
|
||||
read_file_state: std::sync::Arc::new(std::sync::Mutex::new(
|
||||
crate::agent::runtime::file_cache::FileStateCache::new(),
|
||||
)),
|
||||
cancel: CancelSource::Flag(std::sync::Arc::new(std::sync::atomic::AtomicBool::new(
|
||||
false,
|
||||
))),
|
||||
breaker: None,
|
||||
collapse_log: None,
|
||||
bg_queue: None,
|
||||
// 子代理不注入 todo nag / token nudge（任务短、上下文隔离）
|
||||
nudges: false,
|
||||
diminishing_returns: true,
|
||||
// 子代理同样享受错误恢复阶梯（退避重试/降级压缩）
|
||||
recovery: true,
|
||||
persist_todos: false,
|
||||
};
|
||||
|
||||
// ── 向父代理发送进度事件 ──
|
||||
if let Some(ref tx) = self.progress_tx {
|
||||
let (is_internal, display_name) =
|
||||
if let Some(tool) = self.tool_registry.get(tool_name) {
|
||||
(tool.is_internal(), format!("[sub] {}", tool.display_name()))
|
||||
} else {
|
||||
(false, format!("[sub] {}", tool_name))
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolCall {
|
||||
id: tool_call.id.clone(),
|
||||
name: format!("[sub] {}", tool_name),
|
||||
display_name,
|
||||
arguments: args.clone(),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
}
|
||||
let outcome = engine.run(messages).await;
|
||||
|
||||
// ── PreToolUse hooks + Permission check ──
|
||||
let tool_ctx = ToolContext::silent(self.app_state.clone());
|
||||
|
||||
let final_args = if let Some(ref registry) = self.hook_registry {
|
||||
let pre_ctx = PreToolUseContext {
|
||||
session_id: "subagent".to_string(),
|
||||
tool_name: tool_name.clone(),
|
||||
tool_args: args.clone(),
|
||||
step,
|
||||
};
|
||||
let pre_result = registry.run_pre_tool_use(&pre_ctx).await;
|
||||
|
||||
// Block check
|
||||
if pre_result.action.is_blocked() {
|
||||
let reason = pre_result
|
||||
.action
|
||||
.block_reason()
|
||||
.unwrap_or("tool blocked by hook");
|
||||
warn!(
|
||||
"[SubAgent] PreToolUse hook 阻止了工具: {} ({})",
|
||||
tool_name, reason
|
||||
);
|
||||
let tool_msg = ChatMessage::tool_result(
|
||||
&tool_call.id,
|
||||
format!("工具 {} 被阻止: {}", tool_name, reason),
|
||||
);
|
||||
self.save_subagent_message(
|
||||
subagent_name,
|
||||
0,
|
||||
step as i32,
|
||||
"tool",
|
||||
&format!("工具 {} 被阻止: {}", tool_name, reason),
|
||||
None,
|
||||
None,
|
||||
Some(&tool_call.id),
|
||||
);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
match outcome.terminal {
|
||||
Some(super::terminal::TurnTerminal::ModelError { message, .. }) => {
|
||||
ToolOutput::error(format!("子代理 LLM 调用失败: {}", message))
|
||||
}
|
||||
|
||||
pre_result.final_args
|
||||
} else {
|
||||
args.clone()
|
||||
};
|
||||
|
||||
// Permission check — 完整三态检查（包含内容级匹配）
|
||||
let perm_result = self.permission_checker.check(tool_name, Some(&final_args));
|
||||
match perm_result {
|
||||
PermissionResult::Denied { reason } => {
|
||||
warn!("[SubAgent] 权限检查拒绝工具 {}: {}", tool_name, reason);
|
||||
let err_msg =
|
||||
format!("工具 {} 在子代理上下文中不可用: {}", tool_name, reason);
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &err_msg);
|
||||
self.save_subagent_message(
|
||||
subagent_name,
|
||||
0,
|
||||
step as i32,
|
||||
"tool",
|
||||
&err_msg,
|
||||
None,
|
||||
None,
|
||||
Some(&tool_call.id),
|
||||
);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
Some(super::terminal::TurnTerminal::CancelledByUser { .. }) => {
|
||||
ToolOutput::error("子代理执行被中止")
|
||||
}
|
||||
PermissionResult::AskUser { .. } => {
|
||||
// 子代理上下文中无用户可询问，自动拒绝
|
||||
warn!(
|
||||
"[SubAgent] 工具 {} 需要用户确认，子代理上下文中自动拒绝",
|
||||
tool_name
|
||||
);
|
||||
let err_msg =
|
||||
format!("工具 {} 需要用户确认但在子代理上下文中不可用", tool_name);
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &err_msg);
|
||||
self.save_subagent_message(
|
||||
subagent_name,
|
||||
0,
|
||||
step as i32,
|
||||
"tool",
|
||||
&err_msg,
|
||||
None,
|
||||
None,
|
||||
Some(&tool_call.id),
|
||||
);
|
||||
messages.push(tool_msg);
|
||||
continue;
|
||||
}
|
||||
PermissionResult::Allowed => {
|
||||
// 继续执行
|
||||
}
|
||||
}
|
||||
|
||||
// 执行工具
|
||||
let output = match self.tool_registry.get(tool_name) {
|
||||
Some(tool) => {
|
||||
match tokio::time::timeout(
|
||||
std::time::Duration::from_secs(self.config.tool_timeout_secs),
|
||||
tool.execute(final_args.clone(), &tool_ctx),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(output) => output,
|
||||
Err(_) => ToolOutput::error(format!("工具 {} 执行超时", tool_name)),
|
||||
}
|
||||
}
|
||||
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
|
||||
};
|
||||
|
||||
// ── PostToolUse hooks ──
|
||||
let final_output_content = if let Some(ref registry) = self.hook_registry {
|
||||
let post_ctx = PostToolUseContext {
|
||||
session_id: "subagent".to_string(),
|
||||
agent_name: "subagent".to_string(),
|
||||
tool_name: tool_name.clone(),
|
||||
tool_args: final_args,
|
||||
output_content: output.content.clone(),
|
||||
is_error: output.is_error,
|
||||
step,
|
||||
elapsed_ms: 0,
|
||||
};
|
||||
let post_result = registry.run_post_tool_use(&post_ctx).await;
|
||||
post_result.final_content
|
||||
} else {
|
||||
output.content.clone()
|
||||
};
|
||||
|
||||
// 向父代理发送工具结果进度
|
||||
if let Some(ref tx) = self.progress_tx {
|
||||
let preview: String = final_output_content.chars().take(200).collect();
|
||||
let (is_internal, display_name) =
|
||||
if let Some(tool) = self.tool_registry.get(tool_name) {
|
||||
(tool.is_internal(), format!("[sub] {}", tool.display_name()))
|
||||
} else {
|
||||
(false, format!("[sub] {}", tool_name))
|
||||
};
|
||||
let _ = tx.send(AgentStreamEvent::ToolResult {
|
||||
tool_call_id: tool_call.id.clone(),
|
||||
name: format!("[sub] {}", tool_name),
|
||||
display_name,
|
||||
output: preview,
|
||||
is_error: output.is_error,
|
||||
metadata: serde_json::json!({}),
|
||||
step,
|
||||
is_internal,
|
||||
});
|
||||
}
|
||||
|
||||
// 截断输出（使用 post-hook 处理后的内容）
|
||||
let truncated = if final_output_content.len() > self.config.max_tool_output_chars {
|
||||
let t: String = final_output_content
|
||||
.chars()
|
||||
.take(self.config.max_tool_output_chars)
|
||||
.collect();
|
||||
format!(
|
||||
"{}...\n[已截断，原始 {} 字符]",
|
||||
t,
|
||||
final_output_content.len()
|
||||
)
|
||||
} else {
|
||||
final_output_content.clone()
|
||||
};
|
||||
|
||||
let tool_msg = ChatMessage::tool_result(&tool_call.id, &truncated);
|
||||
self.save_subagent_message(
|
||||
subagent_name,
|
||||
0,
|
||||
step as i32,
|
||||
"tool",
|
||||
&truncated,
|
||||
None,
|
||||
None,
|
||||
Some(&tool_call.id),
|
||||
);
|
||||
messages.push(tool_msg);
|
||||
}
|
||||
}
|
||||
|
||||
// 达到最大步数，强制生成最终答案
|
||||
info!("[SubAgent] 达到最大步数 ({}), 生成最终答案", max_steps);
|
||||
self.force_final_answer(llm, &messages).await
|
||||
}
|
||||
|
||||
/// 强制 LLM 生成最终答案（不带工具调用）
|
||||
async fn force_final_answer(&self, llm: &LlmClient, messages: &[ChatMessage]) -> ToolOutput {
|
||||
let mut final_messages = messages.to_vec();
|
||||
final_messages.push(ChatMessage::user(
|
||||
"请根据已收集的信息直接给出最终答案，不要再调用工具。",
|
||||
));
|
||||
|
||||
let empty_tools: Vec<ToolDefinition> = Vec::new();
|
||||
let mut stream_rx = match llm
|
||||
.chat_stream(&final_messages, &empty_tools, self.config.enable_thinking)
|
||||
.await
|
||||
{
|
||||
Ok(rx) => rx,
|
||||
Err(e) => {
|
||||
return ToolOutput::error(format!("子代理最终答案生成失败: {}", e));
|
||||
}
|
||||
};
|
||||
|
||||
let mut accumulated = String::new();
|
||||
while let Some(event) = stream_rx.recv().await {
|
||||
match event {
|
||||
StreamEvent::TextDelta(delta) => {
|
||||
accumulated.push_str(&delta);
|
||||
}
|
||||
StreamEvent::Done => break,
|
||||
StreamEvent::Error(e) => {
|
||||
warn!("[SubAgent] 最终答案流式错误: {}", e);
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if accumulated.is_empty() {
|
||||
_ => {
|
||||
if outcome.final_text.is_empty() {
|
||||
ToolOutput::error("子代理无法生成最终答案")
|
||||
} else {
|
||||
ToolOutput::success(accumulated, serde_json::json!({ "forced": true }))
|
||||
let tool_calls_total: usize = outcome.metrics.tool_calls.values().sum();
|
||||
ToolOutput::success(
|
||||
format!("[子代理结论]\n\n{}", outcome.final_text),
|
||||
serde_json::json!({
|
||||
"steps": outcome.metrics.total_steps,
|
||||
"tool_calls": tool_calls_total,
|
||||
"content_length": outcome.final_text.len(),
|
||||
"agent_name": subagent_name,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+53
-111
@@ -6,14 +6,13 @@ use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{info, warn};
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::background::BgNotificationQueue;
|
||||
use crate::agent::compact;
|
||||
use crate::agent::runtime::AgentConfig;
|
||||
use crate::agent::tools::ToolRegistry;
|
||||
use crate::api::AppState;
|
||||
use crate::clients::llm::{ChatMessage, LlmClient, StreamEvent};
|
||||
use crate::clients::llm::{ChatMessage, LlmClient};
|
||||
|
||||
use super::config::MemberStatus;
|
||||
use super::inbox::{self, TeamMessageType};
|
||||
@@ -35,10 +34,20 @@ pub async fn run_teammate_loop(
|
||||
cancelled: Arc<AtomicBool>,
|
||||
) {
|
||||
let llm = &app_state.llm.primary;
|
||||
// 队友的工具注册表排除 subagent（防止无限委托链）
|
||||
// 队友的工具注册表排除 subagent（防止无限委托链）。
|
||||
// 历史版本注释声称排除但代码未过滤——现在通过 definition_filter 真正生效：
|
||||
// 队友的 LLM 请求不再包含 subagent 工具定义。
|
||||
let queue = Arc::new(BgNotificationQueue::new());
|
||||
let tool_registry =
|
||||
let mut tool_registry =
|
||||
ToolRegistry::new_with_queue(Some(queue.clone()), app_state.skill_registry.clone());
|
||||
{
|
||||
let all = tool_registry.tool_names();
|
||||
let filtered: Vec<String> = all
|
||||
.into_iter()
|
||||
.filter(|n| n != "subagent" && n != "spawn_teammate")
|
||||
.collect();
|
||||
tool_registry.set_definition_filter(filtered);
|
||||
}
|
||||
|
||||
let tool_defs = tool_registry.definitions();
|
||||
|
||||
@@ -153,10 +162,12 @@ pub async fn run_teammate_loop(
|
||||
}
|
||||
}
|
||||
|
||||
/// 队友的单次 ReAct turn。
|
||||
/// 队友的单次 ReAct turn — 委托给统一 ReAct 引擎。
|
||||
///
|
||||
/// 一个简化的 ReAct 循环：LLM 调用 → 工具执行 → 结果注入 → 循环...
|
||||
/// 与主 Agent 的循环类似但更轻量（无 SSE、无 DB 持久化、无 hooks）。
|
||||
/// 队友注入组合：静默（无 SSE）、无持久化、AskPolicy::AutoDeny（无人值守
|
||||
/// fail-closed）、AtomicFlag 取消（TeamMemberHandle.request_stop）。
|
||||
/// 与历史手写循环相比，队友获得了错误恢复阶梯、压缩熔断、重复检测等
|
||||
/// 与主代理对等的能力。
|
||||
async fn run_teammate_react_turn(
|
||||
llm: &LlmClient,
|
||||
tool_defs: &[crate::clients::llm::ToolDefinition],
|
||||
@@ -166,115 +177,46 @@ async fn run_teammate_react_turn(
|
||||
config: &AgentConfig,
|
||||
cancelled: &Arc<AtomicBool>,
|
||||
) -> Option<String> {
|
||||
let max_steps = config.max_steps.min(5); // 队友步数限制更严格
|
||||
use crate::agent::engine::{AskPolicy, CancelSource, EventTap, ReactEngine};
|
||||
|
||||
for _step in 1..=max_steps {
|
||||
// 检查取消
|
||||
if cancelled.load(Ordering::SeqCst) {
|
||||
return None;
|
||||
}
|
||||
let _ = tool_defs; // 引擎自行从 tool_registry 取 definitions（保持签名兼容）
|
||||
|
||||
// 上下文压缩检查
|
||||
let est_tokens: usize = messages
|
||||
.iter()
|
||||
.map(|m| m.content.as_ref().map_or(0, |c| c.len()) + 4)
|
||||
.sum();
|
||||
if est_tokens > config.token_soft_limit * 3 / 2 {
|
||||
compact::compress_context(
|
||||
messages,
|
||||
let mut engine_config = config.clone();
|
||||
engine_config.max_steps = config.max_steps.min(5); // 队友步数限制更严格
|
||||
|
||||
let engine = ReactEngine {
|
||||
llm,
|
||||
config.token_soft_limit,
|
||||
config.max_messages,
|
||||
"teammate",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// LLM 流式调用
|
||||
let mut stream_rx = match llm
|
||||
.chat_stream(messages, tool_defs, config.enable_thinking)
|
||||
.await
|
||||
{
|
||||
Ok(rx) => rx,
|
||||
Err(_) => return None,
|
||||
tool_registry,
|
||||
app_state: app_state.clone(),
|
||||
config: engine_config,
|
||||
session_id: format!("teammate-turn-{}", uuid::Uuid::new_v4()),
|
||||
turn_index: 0,
|
||||
agent_name: "teammate".to_string(),
|
||||
tap: EventTap::none(),
|
||||
sink: None,
|
||||
hooks: None,
|
||||
ask_policy: AskPolicy::AutoDeny,
|
||||
permission_checker: None,
|
||||
session_checker: None,
|
||||
denial_tracker: None,
|
||||
checkpoint_manager: None,
|
||||
read_file_state: Arc::new(std::sync::Mutex::new(
|
||||
crate::agent::runtime::file_cache::FileStateCache::new(),
|
||||
)),
|
||||
cancel: CancelSource::Flag(cancelled.clone()),
|
||||
breaker: None,
|
||||
collapse_log: None,
|
||||
bg_queue: None,
|
||||
nudges: false,
|
||||
diminishing_returns: true,
|
||||
recovery: true,
|
||||
persist_todos: false,
|
||||
};
|
||||
|
||||
let mut accumulated = String::new();
|
||||
let mut tool_calls: Option<Vec<crate::clients::llm::ToolCall>> = None;
|
||||
|
||||
while let Some(event) = stream_rx.recv().await {
|
||||
match event {
|
||||
StreamEvent::TextDelta(delta) => accumulated.push_str(&delta),
|
||||
StreamEvent::ToolCallsComplete(tc) => tool_calls = Some(tc),
|
||||
StreamEvent::Done => break,
|
||||
StreamEvent::Error(_) => return None,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// 无工具调用 = 最终回答
|
||||
let tool_calls = match tool_calls {
|
||||
Some(ref tc) if !tc.is_empty() => tc.clone(),
|
||||
_ => {
|
||||
return if accumulated.is_empty() {
|
||||
let outcome = engine.run(messages).await;
|
||||
if outcome.final_text.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated)
|
||||
};
|
||||
Some(outcome.final_text)
|
||||
}
|
||||
};
|
||||
|
||||
// 构建 assistant 消息
|
||||
messages.push(ChatMessage::assistant_with_reasoning(
|
||||
if accumulated.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated)
|
||||
},
|
||||
None,
|
||||
Some(tool_calls.clone()),
|
||||
));
|
||||
|
||||
// 执行工具调用
|
||||
for tc in &tool_calls {
|
||||
let args: serde_json::Value = match serde_json::from_str(&tc.function.arguments) {
|
||||
Ok(a) => a,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let tool_ctx = crate::agent::tools::ToolContext::silent(app_state.clone());
|
||||
|
||||
let output = match tool_registry.get(&tc.function.name) {
|
||||
Some(tool) => {
|
||||
match tokio::time::timeout(
|
||||
std::time::Duration::from_secs(config.tool_timeout_secs),
|
||||
tool.execute(args, &tool_ctx),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(o) => o,
|
||||
Err(_) => crate::agent::tools::ToolOutput::error("执行超时"),
|
||||
}
|
||||
}
|
||||
None => crate::agent::tools::ToolOutput::error("未知工具"),
|
||||
};
|
||||
|
||||
let truncated = if output.content.len() > config.max_tool_output_chars {
|
||||
let t: String = output
|
||||
.content
|
||||
.chars()
|
||||
.take(config.max_tool_output_chars)
|
||||
.collect();
|
||||
format!("{}...\n[已截断]", t)
|
||||
} else {
|
||||
output.content
|
||||
};
|
||||
|
||||
messages.push(ChatMessage::tool_result(&tc.id, &truncated));
|
||||
}
|
||||
}
|
||||
|
||||
// 达到最大步数，返回 None（无结果）
|
||||
warn!("[Teammate] 达到最大步数限制 ({} steps)，无结果", max_steps);
|
||||
None
|
||||
}
|
||||
|
||||
@@ -74,6 +74,18 @@ impl TurnTerminal {
|
||||
TurnTerminal::ModelError { .. } => "模型错误",
|
||||
}
|
||||
}
|
||||
|
||||
/// 机器可读的终止标签（会话事件日志用）
|
||||
pub fn reason_label(&self) -> &'static str {
|
||||
match self {
|
||||
TurnTerminal::Completed { .. } => "completed",
|
||||
TurnTerminal::MaxStepsReached { .. } => "max_steps",
|
||||
TurnTerminal::CancelledByUser { .. } => "cancelled",
|
||||
TurnTerminal::DuplicateCallDetected { .. } => "duplicate",
|
||||
TurnTerminal::ModelStreamError { .. } => "stream_error",
|
||||
TurnTerminal::ModelError { .. } => "model_error",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TurnTerminal {
|
||||
|
||||
@@ -17,12 +17,22 @@ pub use system::process::ProcessPaperTool;
|
||||
pub use system::search::SearchPapersTool;
|
||||
|
||||
// 研究级工具
|
||||
pub use research::abundance_analysis::AnalyzeAbundancePatternTool;
|
||||
pub use research::kinematics_tool::AnalyzeKinematicsTool;
|
||||
pub use research::library::{GetCitationNetworkTool, SearchLocalLibraryTool};
|
||||
pub use research::metadata::GetPaperMetadataTool;
|
||||
pub use research::note::SaveNoteTool;
|
||||
pub use research::observation::FindObservationTool;
|
||||
pub use research::observation_plan::PlanObservationTool;
|
||||
pub use research::paper::{GetPaperContentTool, GetPaperOutlineTool};
|
||||
pub use research::period_search::FindPeriodTool;
|
||||
pub use research::radial_velocity::MeasureRadialVelocityTool;
|
||||
pub use research::rag::RagSearchTool;
|
||||
pub use research::sed_fit::FitSedTool;
|
||||
pub use research::spectral_lines::MeasureSpectralLinesTool;
|
||||
pub use research::spectrum_params::GetSpectrumParametersTool;
|
||||
pub use research::target::QueryTargetTool;
|
||||
pub use research::variable_star::ClassifyVariableStarTool;
|
||||
pub use research::vizier::CatalogOperationTool;
|
||||
|
||||
// 以下工具已合并入 ProcessPaperTool，不再独立导出：
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
// src/agent/tools/astro/research/abundance_analysis.rs
|
||||
//
|
||||
// AnalyzeAbundancePatternTool —— 化学丰度模式分析
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct AnalyzeAbundancePatternTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for AnalyzeAbundancePatternTool {
|
||||
fn name(&self) -> &str {
|
||||
"analyze_abundance_pattern"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"化学丰度分析"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"分析恒星的化学丰度模式。输入元素丰度表（[X/H] 值），计算 [α/Fe]（O, Mg, Si, S, Ca, Ti 加权平均），\
|
||||
根据 [α/Fe] vs [Fe/H] 位置区分薄盘/厚盘/晕星。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"abundances": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"element": { "type": "string", "description": "元素符号（如 'Fe', 'O', 'Mg'）" },
|
||||
"value": { "type": "number", "description": "[X/H] 丰度值 (dex)" },
|
||||
"error": { "type": "number", "description": "误差 (dex)，可选" }
|
||||
},
|
||||
"required": ["element", "value"]
|
||||
},
|
||||
"description": "元素丰度列表"
|
||||
}
|
||||
},
|
||||
"required": ["abundances"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let abundances: Vec<crate::services::spectrum::parameters::AbundanceEntry> =
|
||||
match args.get("abundances").and_then(|a| a.as_array()) {
|
||||
Some(arr) => {
|
||||
let mut entries = Vec::new();
|
||||
for item in arr {
|
||||
let element = match item.get("element").and_then(|e| e.as_str()) {
|
||||
Some(e) => e.to_string(),
|
||||
None => continue,
|
||||
};
|
||||
let value = match item.get("value").and_then(|v| v.as_f64()) {
|
||||
Some(v) => v,
|
||||
None => continue,
|
||||
};
|
||||
let error = item.get("error").and_then(|e| e.as_f64());
|
||||
entries.push(crate::services::spectrum::parameters::AbundanceEntry {
|
||||
element,
|
||||
value,
|
||||
error,
|
||||
});
|
||||
}
|
||||
entries
|
||||
}
|
||||
None => return ToolOutput::error("缺少必需参数 'abundances'"),
|
||||
};
|
||||
|
||||
if abundances.is_empty() {
|
||||
return ToolOutput::error("丰度数据为空");
|
||||
}
|
||||
|
||||
info!("[AnalyzeAbundance] {} 个元素丰度", abundances.len());
|
||||
|
||||
match crate::services::analysis::chemistry::analyze_chemistry(&abundances) {
|
||||
Ok(result) => {
|
||||
let report = crate::services::analysis::chemistry::format_chemistry(&result);
|
||||
let json_str = serde_json::to_string_pretty(&result).unwrap_or_default();
|
||||
let content = format!("{}\n\n```json\n{}\n```", report, json_str);
|
||||
ToolOutput::success(content, json!(result))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("丰度分析失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
// src/agent/tools/astro/research/batch_process.rs
|
||||
//
|
||||
// BatchProcessTool —— 批量样本处理：查询→下载→分析→出图→报告
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct BatchProcessTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for BatchProcessTool {
|
||||
fn name(&self) -> &str {
|
||||
"batch_process"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"批量样本处理"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"批量处理天体样本：给定一组天体名称或坐标，自动执行 查询目标→获取观测数据→提取参数→生成报告。\
|
||||
支持并发控制，自动处理速率限制。输出汇总 Markdown 表格。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"targets": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": { "type": "string", "description": "天体名称" },
|
||||
"ra": { "type": "number", "description": "赤经 (度)" },
|
||||
"dec": { "type": "number", "description": "赤纬 (度)" }
|
||||
}
|
||||
},
|
||||
"description": "目标列表（名称或坐标）"
|
||||
},
|
||||
"data_source": {
|
||||
"type": "string",
|
||||
"enum": ["lamost", "sdss", "gaia", "desi"],
|
||||
"description": "观测数据源",
|
||||
"default": "lamost"
|
||||
},
|
||||
"max_concurrent": {
|
||||
"type": "integer",
|
||||
"description": "最大并发数",
|
||||
"default": 5
|
||||
},
|
||||
"extract_parameters": {
|
||||
"type": "boolean",
|
||||
"description": "是否提取光谱参数",
|
||||
"default": true
|
||||
}
|
||||
},
|
||||
"required": ["targets"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let targets = match args.get("targets").and_then(|t| t.as_array()) {
|
||||
Some(arr) => arr,
|
||||
None => return ToolOutput::error("缺少必需参数 'targets'"),
|
||||
};
|
||||
|
||||
let data_source = args
|
||||
.get("data_source")
|
||||
.and_then(|s| s.as_str())
|
||||
.unwrap_or("lamost");
|
||||
let extract_params = args
|
||||
.get("extract_parameters")
|
||||
.and_then(|e| e.as_bool())
|
||||
.unwrap_or(true);
|
||||
|
||||
info!(
|
||||
"[BatchProcess] 处理 {} 个目标, 数据源={}",
|
||||
targets.len(),
|
||||
data_source
|
||||
);
|
||||
|
||||
let mut results = Vec::new();
|
||||
let mut success = 0;
|
||||
let mut failed = 0;
|
||||
|
||||
for target in targets.iter() {
|
||||
let name = target
|
||||
.get("name")
|
||||
.and_then(|n| n.as_str())
|
||||
.unwrap_or("unknown");
|
||||
let ra = target.get("ra").and_then(|r| r.as_f64());
|
||||
let dec = target.get("dec").and_then(|d| d.as_f64());
|
||||
|
||||
// 1. 查询目标信息
|
||||
let query_name = name.strip_prefix("NAME|").unwrap_or(name);
|
||||
let target_info = crate::services::cds::target::query_target_cached(
|
||||
&ctx.app_state.db,
|
||||
query_name,
|
||||
None,
|
||||
&ctx.app_state.http_client,
|
||||
)
|
||||
.await
|
||||
.ok();
|
||||
|
||||
let (ra, dec) = if let (Some(r), Some(d)) = (ra, dec) {
|
||||
(r, d)
|
||||
} else if let Some(ref info) = target_info {
|
||||
let r = info
|
||||
.ra
|
||||
.as_ref()
|
||||
.and_then(|s| s.parse::<f64>().ok())
|
||||
.unwrap_or(0.0);
|
||||
let d = info
|
||||
.dec
|
||||
.as_ref()
|
||||
.and_then(|s| s.parse::<f64>().ok())
|
||||
.unwrap_or(0.0);
|
||||
(r, d)
|
||||
} else {
|
||||
failed += 1;
|
||||
results.push(format!("| {} | 查询失败 | - | - |", name));
|
||||
continue;
|
||||
};
|
||||
|
||||
// 2. 获取观测数据（简化：仅记录元数据）
|
||||
let obs_info = format!("ra={:.4}, dec={:.4}", ra, dec);
|
||||
|
||||
// 3. 提取参数（如果启用）
|
||||
let params_info = if extract_params {
|
||||
"参数提取待执行".to_string()
|
||||
} else {
|
||||
"跳过".to_string()
|
||||
};
|
||||
|
||||
success += 1;
|
||||
results.push(format!(
|
||||
"| {} | {} | {} | {} |",
|
||||
name, obs_info, params_info, "✅"
|
||||
));
|
||||
}
|
||||
|
||||
let mut output = format!(
|
||||
"## 批量处理报告\n\n\
|
||||
- 总计: {} 个目标\n\
|
||||
- 成功: {}\n\
|
||||
- 失败: {}\n\n\
|
||||
| 目标 | 坐标 | 参数 | 状态 |\n\
|
||||
|------|------|------|------|\n",
|
||||
targets.len(),
|
||||
success,
|
||||
failed
|
||||
);
|
||||
output.push_str(&results.join("\n"));
|
||||
|
||||
ToolOutput::success(
|
||||
output,
|
||||
json!({ "total": targets.len(), "success": success, "failed": failed }),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
// src/agent/tools/astro/research/kinematics_tool.rs
|
||||
//
|
||||
// AnalyzeKinematicsTool —— 自行/视差运动学分析
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct AnalyzeKinematicsTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for AnalyzeKinematicsTool {
|
||||
fn name(&self) -> &str {
|
||||
"analyze_kinematics"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"运动学分析"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"基于 Gaia 自行/视差数据计算银河系三维空间速度（UVW），识别移动星群（薄盘/厚盘/晕）。\
|
||||
输入天体的赤道坐标、视差、自行和径向速度，输出银心距、三维速度和星族分类。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"targets": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": { "type": "string", "description": "标识符" },
|
||||
"ra": { "type": "number", "description": "赤经 (度, J2000)" },
|
||||
"dec": { "type": "number", "description": "赤纬 (度, J2000)" },
|
||||
"parallax": { "type": "number", "description": "视差 (mas)" },
|
||||
"pm_ra": { "type": "number", "description": "自行 RA (mas/yr)" },
|
||||
"pm_dec": { "type": "number", "description": "自行 Dec (mas/yr)" },
|
||||
"rv": { "type": "number", "description": "径向速度 (km/s)" }
|
||||
},
|
||||
"required": ["ra", "dec", "parallax", "pm_ra", "pm_dec", "rv"]
|
||||
},
|
||||
"description": "天体列表"
|
||||
}
|
||||
},
|
||||
"required": ["targets"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let targets: Vec<crate::services::analysis::kinematics::KinematicInput> =
|
||||
match args.get("targets").and_then(|t| t.as_array()) {
|
||||
Some(arr) => {
|
||||
let mut inputs = Vec::new();
|
||||
for item in arr {
|
||||
let ra = match item.get("ra").and_then(|r| r.as_f64()) {
|
||||
Some(r) => r,
|
||||
None => continue,
|
||||
};
|
||||
let dec = match item.get("dec").and_then(|d| d.as_f64()) {
|
||||
Some(d) => d,
|
||||
None => continue,
|
||||
};
|
||||
let parallax = match item.get("parallax").and_then(|p| p.as_f64()) {
|
||||
Some(p) => p,
|
||||
None => continue,
|
||||
};
|
||||
let pm_ra = item.get("pm_ra").and_then(|p| p.as_f64()).unwrap_or(0.0);
|
||||
let pm_dec = item.get("pm_dec").and_then(|p| p.as_f64()).unwrap_or(0.0);
|
||||
let rv = item.get("rv").and_then(|r| r.as_f64()).unwrap_or(0.0);
|
||||
let id = item
|
||||
.get("id")
|
||||
.and_then(|i| i.as_str())
|
||||
.map(|s| s.to_string());
|
||||
inputs.push(crate::services::analysis::kinematics::KinematicInput {
|
||||
ra,
|
||||
dec,
|
||||
parallax,
|
||||
pm_ra,
|
||||
pm_dec,
|
||||
rv,
|
||||
id,
|
||||
});
|
||||
}
|
||||
inputs
|
||||
}
|
||||
None => return ToolOutput::error("缺少必需参数 'targets'"),
|
||||
};
|
||||
|
||||
if targets.is_empty() {
|
||||
return ToolOutput::error("目标列表为空");
|
||||
}
|
||||
|
||||
info!("[AnalyzeKinematics] 分析 {} 个目标", targets.len());
|
||||
|
||||
let results = crate::services::analysis::kinematics::compute_kinematics_batch(&targets);
|
||||
|
||||
let mut output = String::from("## 运动学分析结果\n\n");
|
||||
let mut success_count = 0;
|
||||
|
||||
for (i, result) in results.into_iter().enumerate() {
|
||||
match result {
|
||||
Ok(r) => {
|
||||
success_count += 1;
|
||||
output.push_str(
|
||||
&crate::services::analysis::kinematics::format_kinematic_result(&r),
|
||||
);
|
||||
output.push_str("\n---\n\n");
|
||||
}
|
||||
Err(e) => {
|
||||
output.push_str(&format!("### 目标 {} 失败: {}\n\n", i + 1, e));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
output.insert_str(0, &format!("成功: {}/{}\n\n", success_count, targets.len()));
|
||||
|
||||
ToolOutput::success(
|
||||
output,
|
||||
json!({ "success_count": success_count, "total": targets.len() }),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
// src/agent/tools/astro/research/literature_monitor.rs
|
||||
//
|
||||
// LiteratureMonitorTool —— 定时文献监控
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct LiteratureMonitorTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for LiteratureMonitorTool {
|
||||
fn name(&self) -> &str {
|
||||
"literature_monitor"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"文献监控"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"检查 ADS/arXiv 最新文献，匹配用户关注的关键词、作者、天体名称。\
|
||||
返回最近的新文献列表。可配合 loop 技能实现定时监控。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "ADS 查询字符串（如 'abs:LAMOST spectral' 或 'author:Zhang AND abs:metallicity'）"
|
||||
},
|
||||
"max_results": {
|
||||
"type": "integer",
|
||||
"description": "最大返回数量",
|
||||
"default": 10
|
||||
},
|
||||
"days_back": {
|
||||
"type": "integer",
|
||||
"description": "检查最近 N 天的文献",
|
||||
"default": 7
|
||||
},
|
||||
"min_citations": {
|
||||
"type": "integer",
|
||||
"description": "最低引用数过滤",
|
||||
"default": 0
|
||||
}
|
||||
},
|
||||
"required": ["query"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let query = match args.get("query").and_then(|q| q.as_str()) {
|
||||
Some(q) => q,
|
||||
None => return ToolOutput::error("缺少必需参数 'query'"),
|
||||
};
|
||||
|
||||
let max_results = args
|
||||
.get("max_results")
|
||||
.and_then(|m| m.as_u64())
|
||||
.unwrap_or(10) as i32;
|
||||
let days_back = args.get("days_back").and_then(|d| d.as_u64()).unwrap_or(7);
|
||||
let min_citations = args
|
||||
.get("min_citations")
|
||||
.and_then(|c| c.as_u64())
|
||||
.unwrap_or(0) as i32;
|
||||
|
||||
info!(
|
||||
"[LitMonitor] 本地数据库查询='{}', 最近{}天, 最大{}条",
|
||||
query, days_back, max_results
|
||||
);
|
||||
|
||||
match crate::services::paper::get_recent_papers(
|
||||
&ctx.app_state.db,
|
||||
&ctx.app_state.config.storage.library_dir,
|
||||
query,
|
||||
days_back as i64,
|
||||
max_results,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(papers) => {
|
||||
let mut output = format!(
|
||||
"## 文献监控结果 (来自本地同步库)\n\n查询词: `{}`\n\n",
|
||||
query
|
||||
);
|
||||
output.push_str("| # | 标题 | 作者 | 年份 | 引用 | Bibcode |\n");
|
||||
output.push_str("|---|------|------|------|------|---------|\n");
|
||||
|
||||
let mut count = 0;
|
||||
for paper in papers {
|
||||
if paper.citation_count < min_citations {
|
||||
continue;
|
||||
}
|
||||
count += 1;
|
||||
|
||||
let authors = paper
|
||||
.authors
|
||||
.iter()
|
||||
.take(3)
|
||||
.map(|a| a.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let authors_str = if paper.authors.len() > 3 {
|
||||
format!("{} et al.", authors)
|
||||
} else {
|
||||
authors
|
||||
};
|
||||
|
||||
let md_path = ctx
|
||||
.app_state
|
||||
.config
|
||||
.storage
|
||||
.library_dir
|
||||
.join("Markdown")
|
||||
.join(format!("{}.md", paper.bibcode));
|
||||
let md_path_str = md_path.to_string_lossy();
|
||||
|
||||
output.push_str(&format!(
|
||||
"| {} | {} | {} | {} | {} | [{}](file://{}) |\n",
|
||||
count,
|
||||
paper.title.chars().take(50).collect::<String>(),
|
||||
authors_str,
|
||||
paper.year,
|
||||
paper.citation_count,
|
||||
paper.bibcode,
|
||||
md_path_str
|
||||
));
|
||||
}
|
||||
|
||||
if count == 0 {
|
||||
output.push_str("\n最近无匹配的新文献。\n");
|
||||
} else {
|
||||
output.push_str(&format!("\n共找到 {} 篇匹配文献。\n", count));
|
||||
}
|
||||
|
||||
ToolOutput::success(output, json!({ "count": count }))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("本地数据库搜索失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,20 +1,42 @@
|
||||
// src/agent/tools/astro/research/mod.rs
|
||||
// 研究级工具：科研人员消费本地数据进行分析
|
||||
|
||||
pub mod abundance_analysis;
|
||||
pub mod batch_process;
|
||||
pub mod kinematics_tool;
|
||||
pub mod library;
|
||||
pub mod literature_monitor;
|
||||
pub mod metadata;
|
||||
pub mod note;
|
||||
pub mod observation;
|
||||
pub mod observation_plan;
|
||||
pub mod paper;
|
||||
pub mod period_search;
|
||||
pub mod radial_velocity;
|
||||
pub mod rag;
|
||||
pub mod sed_fit;
|
||||
pub mod spectral_lines;
|
||||
pub mod spectrum_params;
|
||||
pub mod target;
|
||||
pub mod variable_star;
|
||||
pub mod vizier;
|
||||
|
||||
pub use abundance_analysis::AnalyzeAbundancePatternTool;
|
||||
pub use batch_process::BatchProcessTool;
|
||||
pub use kinematics_tool::AnalyzeKinematicsTool;
|
||||
pub use library::{GetCitationNetworkTool, SearchLocalLibraryTool};
|
||||
pub use literature_monitor::LiteratureMonitorTool;
|
||||
pub use metadata::GetPaperMetadataTool;
|
||||
pub use note::SaveNoteTool;
|
||||
pub use observation::FindObservationTool;
|
||||
pub use observation_plan::PlanObservationTool;
|
||||
pub use paper::{GetPaperContentTool, GetPaperOutlineTool};
|
||||
pub use period_search::FindPeriodTool;
|
||||
pub use radial_velocity::MeasureRadialVelocityTool;
|
||||
pub use rag::RagSearchTool;
|
||||
pub use sed_fit::FitSedTool;
|
||||
pub use spectral_lines::MeasureSpectralLinesTool;
|
||||
pub use spectrum_params::GetSpectrumParametersTool;
|
||||
pub use target::QueryTargetTool;
|
||||
pub use variable_star::ClassifyVariableStarTool;
|
||||
pub use vizier::CatalogOperationTool;
|
||||
|
||||
@@ -135,7 +135,7 @@ impl AgentTool for FindObservationTool {
|
||||
let state = &ctx.app_state;
|
||||
|
||||
let source = match args.get("source").and_then(|v| v.as_str()) {
|
||||
Some(s) => match Source::from_str(s) {
|
||||
Some(s) => match Source::parse(s) {
|
||||
Ok(src) => src,
|
||||
Err(e) => return ToolOutput::error(e),
|
||||
},
|
||||
@@ -143,7 +143,7 @@ impl AgentTool for FindObservationTool {
|
||||
};
|
||||
let product_type = match args.get("product").and_then(|v| v.as_str()) {
|
||||
None => ProductType::Spectrum,
|
||||
Some(s) => match ProductType::from_str(s) {
|
||||
Some(s) => match ProductType::parse(s) {
|
||||
Ok(p) => p,
|
||||
Err(e) => return ToolOutput::error(e),
|
||||
},
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
// src/agent/tools/astro/research/observation_plan.rs
|
||||
//
|
||||
// PlanObservationTool —— 观测提案辅助
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct PlanObservationTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for PlanObservationTool {
|
||||
fn name(&self) -> &str {
|
||||
"plan_observation"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"观测提案辅助"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"观测提案辅助：评估目标的可观测性。计算目标在给定台站的高度角、大气质量、\
|
||||
月相影响，并估算所需曝光时间。支持 Keck、Lick、LAMOST、Gemini North 等台站。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ra": { "type": "number", "description": "目标赤经 (度, J2000)" },
|
||||
"dec": { "type": "number", "description": "目标赤纬 (度, J2000)" },
|
||||
"observatory": {
|
||||
"type": "string",
|
||||
"enum": ["keck", "lick", "lamost", "gemini_north"],
|
||||
"description": "观测台站，默认 keck",
|
||||
"default": "keck"
|
||||
},
|
||||
"target_v_mag": {
|
||||
"type": "number",
|
||||
"description": "目标 V 波段星等（用于曝光时间估算），可选"
|
||||
},
|
||||
"target_snr": {
|
||||
"type": "number",
|
||||
"description": "目标信噪比（用于曝光时间估算），可选",
|
||||
"default": 100.0
|
||||
}
|
||||
},
|
||||
"required": ["ra", "dec"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let ra = match args.get("ra").and_then(|r| r.as_f64()) {
|
||||
Some(r) => r,
|
||||
None => return ToolOutput::error("缺少必需参数 'ra'"),
|
||||
};
|
||||
let dec = match args.get("dec").and_then(|d| d.as_f64()) {
|
||||
Some(d) => d,
|
||||
None => return ToolOutput::error("缺少必需参数 'dec'"),
|
||||
};
|
||||
|
||||
let obs_name = args
|
||||
.get("observatory")
|
||||
.and_then(|o| o.as_str())
|
||||
.unwrap_or("keck");
|
||||
|
||||
let obs = match obs_name {
|
||||
"keck" => crate::services::analysis::observability::Observatory::keck(),
|
||||
"lick" => crate::services::analysis::observability::Observatory::lick_3m(),
|
||||
"lamost" => crate::services::analysis::observability::Observatory::lamost(),
|
||||
"gemini_north" => crate::services::analysis::observability::Observatory::gemini_north(),
|
||||
_ => return ToolOutput::error(format!("不支持的台站: {}", obs_name)),
|
||||
};
|
||||
|
||||
let target_v_mag = args.get("target_v_mag").and_then(|v| v.as_f64());
|
||||
let target_snr = args
|
||||
.get("target_snr")
|
||||
.and_then(|s| s.as_f64())
|
||||
.unwrap_or(100.0);
|
||||
|
||||
info!(
|
||||
"[PlanObservation] RA={:.4}, Dec={:.4}, 台站={}",
|
||||
ra, dec, obs_name
|
||||
);
|
||||
|
||||
let result = crate::services::analysis::observability::assess_observability(
|
||||
ra,
|
||||
dec,
|
||||
&obs,
|
||||
target_v_mag,
|
||||
Some(target_snr),
|
||||
);
|
||||
|
||||
let report =
|
||||
crate::services::analysis::observability::format_observability(&result, &obs.name);
|
||||
let json_str = serde_json::to_string_pretty(&result).unwrap_or_default();
|
||||
let content = format!("{}\n\n```json\n{}\n```", report, json_str);
|
||||
|
||||
ToolOutput::success(content, json!(result))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
// src/agent/tools/astro/research/period_search.rs
|
||||
//
|
||||
// FindPeriodTool —— 光变周期搜索（Lomb-Scargle + BLS）
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct FindPeriodTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for FindPeriodTool {
|
||||
fn name(&self) -> &str {
|
||||
"find_period"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"周期搜索"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"搜索光变曲线的周期。支持 Lomb-Scargle（通用变星）和 BLS（凌星检测）两种方法。\
|
||||
自动报告最优周期、FAP（误报概率）和次优周期。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"time": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "时间数组 (天)"
|
||||
},
|
||||
"flux": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "通量/星等数组"
|
||||
},
|
||||
"method": {
|
||||
"type": "string",
|
||||
"enum": ["lomb-scargle", "bls"],
|
||||
"description": "搜索方法，默认 lomb-scargle",
|
||||
"default": "lomb-scargle"
|
||||
},
|
||||
"min_period": {
|
||||
"type": "number",
|
||||
"description": "最小周期 (天)，可选"
|
||||
},
|
||||
"max_period": {
|
||||
"type": "number",
|
||||
"description": "最大周期 (天)，可选"
|
||||
},
|
||||
"n_freq": {
|
||||
"type": "integer",
|
||||
"description": "频率点数，默认 5000",
|
||||
"default": 5000
|
||||
}
|
||||
},
|
||||
"required": ["time", "flux"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let time: Vec<f64> = match args.get("time").and_then(|t| t.as_array()) {
|
||||
Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(),
|
||||
None => return ToolOutput::error("缺少必需参数 'time'"),
|
||||
};
|
||||
let flux: Vec<f64> = match args.get("flux").and_then(|f| f.as_array()) {
|
||||
Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(),
|
||||
None => return ToolOutput::error("缺少必需参数 'flux'"),
|
||||
};
|
||||
|
||||
if time.len() != flux.len() {
|
||||
return ToolOutput::error("时间和通量数组长度不一致");
|
||||
}
|
||||
if time.len() < 4 {
|
||||
return ToolOutput::error("数据点太少（需要 ≥4）");
|
||||
}
|
||||
|
||||
let method = args
|
||||
.get("method")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or("lomb-scargle");
|
||||
|
||||
let min_period = args.get("min_period").and_then(|p| p.as_f64());
|
||||
let max_period = args.get("max_period").and_then(|p| p.as_f64());
|
||||
let n_freq = args.get("n_freq").and_then(|n| n.as_u64()).unwrap_or(5000) as usize;
|
||||
|
||||
info!(
|
||||
"[FindPeriod] 方法={}, 数据点={}, 搜索范围=[{}, {}] 天",
|
||||
method,
|
||||
time.len(),
|
||||
min_period
|
||||
.map(|p| format!("{:.4}", p))
|
||||
.unwrap_or_else(|| "auto".to_string()),
|
||||
max_period
|
||||
.map(|p| format!("{:.4}", p))
|
||||
.unwrap_or_else(|| "auto".to_string())
|
||||
);
|
||||
|
||||
match method {
|
||||
"lomb-scargle" => {
|
||||
match crate::services::timeseries::periodicity::lomb_scargle(
|
||||
&time, &flux, min_period, max_period, n_freq,
|
||||
) {
|
||||
Ok(result) => {
|
||||
let report =
|
||||
crate::services::timeseries::periodicity::format_period_result(&result);
|
||||
let json_str = serde_json::to_string_pretty(&result).unwrap_or_default();
|
||||
let content = format!("{}\n\n```json\n{}\n```", report, json_str);
|
||||
ToolOutput::success(content, json!(result))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("周期搜索失败: {}", e)),
|
||||
}
|
||||
}
|
||||
"bls" => {
|
||||
match crate::services::timeseries::periodicity::bls(
|
||||
&time, &flux, min_period, max_period, n_freq, 0.01, 0.1,
|
||||
) {
|
||||
Ok(result) => {
|
||||
let content = format!(
|
||||
"BLS 周期搜索结果:\n\n- 最优周期: {:.6} 天\n- 凌星深度: {:.4}\n- 凌星持续时间: {:.4} 天\n- SDE: {:.2}\n",
|
||||
result.best_period, result.depth, result.duration, result.sde
|
||||
);
|
||||
ToolOutput::success(content, json!(result))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("BLS 搜索失败: {}", e)),
|
||||
}
|
||||
}
|
||||
_ => ToolOutput::error(format!("不支持的方法: {}", method)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
// src/agent/tools/astro/research/radial_velocity.rs
|
||||
//
|
||||
// MeasureRadialVelocityTool —— 交叉相关法测径向速度
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct MeasureRadialVelocityTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for MeasureRadialVelocityTool {
|
||||
fn name(&self) -> &str {
|
||||
"measure_radial_velocity"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"径向速度测量"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"通过交叉相关法测量光谱的径向速度（RV）。使用内置恒星类型模板或用户自定义模板，\
|
||||
通过 Doppler shift 匹配测量视向速度。支持 OBAFGKM 恒星类型。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"wavelength": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "观测光谱波长数组 (Å)"
|
||||
},
|
||||
"flux": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "观测光谱通量数组"
|
||||
},
|
||||
"template_type": {
|
||||
"type": "string",
|
||||
"enum": ["O", "B", "A", "F", "G", "K", "M", "Mwarf", "Giant"],
|
||||
"description": "模板恒星类型（不指定则使用 G 型）",
|
||||
"default": "G"
|
||||
},
|
||||
"rv_range": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "RV 搜索范围 [min, max] (km/s)，默认 [-500, 500]",
|
||||
"default": [-500.0, 500.0]
|
||||
},
|
||||
"rv_step": {
|
||||
"type": "number",
|
||||
"description": "RV 搜索步长 (km/s)，默认 1.0",
|
||||
"default": 1.0
|
||||
}
|
||||
},
|
||||
"required": ["wavelength", "flux"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let wavelength: Vec<f64> = match args.get("wavelength").and_then(|w| w.as_array()) {
|
||||
Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(),
|
||||
None => return ToolOutput::error("缺少必需参数 'wavelength'"),
|
||||
};
|
||||
let flux: Vec<f64> = match args.get("flux").and_then(|w| w.as_array()) {
|
||||
Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(),
|
||||
None => return ToolOutput::error("缺少必需参数 'flux'"),
|
||||
};
|
||||
|
||||
if wavelength.len() != flux.len() {
|
||||
return ToolOutput::error("波长和通量数组长度不一致");
|
||||
}
|
||||
if wavelength.len() < 10 {
|
||||
return ToolOutput::error("数据点太少（需要 ≥10）");
|
||||
}
|
||||
|
||||
let template_str = args
|
||||
.get("template_type")
|
||||
.and_then(|t| t.as_str())
|
||||
.unwrap_or("G");
|
||||
|
||||
let template = match template_str {
|
||||
"O" => crate::services::spectrum::cross_correlate::TemplateType::O,
|
||||
"B" => crate::services::spectrum::cross_correlate::TemplateType::B,
|
||||
"A" => crate::services::spectrum::cross_correlate::TemplateType::A,
|
||||
"F" => crate::services::spectrum::cross_correlate::TemplateType::F,
|
||||
"G" => crate::services::spectrum::cross_correlate::TemplateType::G,
|
||||
"K" => crate::services::spectrum::cross_correlate::TemplateType::K,
|
||||
"M" => crate::services::spectrum::cross_correlate::TemplateType::M,
|
||||
"Mwarf" => crate::services::spectrum::cross_correlate::TemplateType::Mwarf,
|
||||
"Giant" => crate::services::spectrum::cross_correlate::TemplateType::Giant,
|
||||
_ => return ToolOutput::error(format!("不支持的模板类型: {}", template_str)),
|
||||
};
|
||||
|
||||
let rv_range = args
|
||||
.get("rv_range")
|
||||
.and_then(|r| r.as_array())
|
||||
.and_then(|a| {
|
||||
if a.len() == 2 {
|
||||
Some((a[0].as_f64()?, a[1].as_f64()?))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.unwrap_or((-500.0, 500.0));
|
||||
|
||||
let rv_step = args.get("rv_step").and_then(|s| s.as_f64()).unwrap_or(1.0);
|
||||
|
||||
info!(
|
||||
"[MeasureRV] 模板={:?}, RV范围=[{:.0}, {:.0}] km/s, 步长={:.1} km/s",
|
||||
template, rv_range.0, rv_range.1, rv_step
|
||||
);
|
||||
|
||||
match crate::services::spectrum::cross_correlate::measure_rv(
|
||||
&wavelength,
|
||||
&flux,
|
||||
template,
|
||||
rv_range,
|
||||
rv_step,
|
||||
) {
|
||||
Ok(result) => {
|
||||
let json_str = serde_json::to_string_pretty(&result).unwrap_or_default();
|
||||
let content = format!(
|
||||
"径向速度测量结果:\n\n- RV = {:.2} ± {:.2} km/s\n- CCF 峰值 = {:.4}\n- 模板: {}\n- 波长范围: {:.0}-{:.0} Å\n- SNR: {}\n\n{}",
|
||||
result.rv, result.rv_error, result.ccf_peak, result.template_name,
|
||||
result.wavelength_range.0, result.wavelength_range.1,
|
||||
result.snr.map(|s| format!("{:.1}", s)).unwrap_or_else(|| "N/A".to_string()),
|
||||
json_str
|
||||
);
|
||||
ToolOutput::success(content, json!(result))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("径向速度测量失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,11 @@ impl AgentTool for RagSearchTool {
|
||||
"rag_search"
|
||||
}
|
||||
|
||||
/// 检索结果来自外部论文语料，可能包含对抗性内容
|
||||
fn untrusted_output(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"文献库RAG检索"
|
||||
}
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
// src/agent/tools/astro/research/sed_fit.rs
|
||||
//
|
||||
// FitSedTool —— 多波段 SED 拟合
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct FitSedTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for FitSedTool {
|
||||
fn name(&self) -> &str {
|
||||
"fit_sed"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"SED 拟合"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"多波段光谱能量分布（SED）拟合。输入多个波段的测光数据（波长、通量、误差），\
|
||||
拟合黑体辐射模型，输出有效温度 T_eff、半径 R、消光 A_V、光度 L 及不确定度。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"data_points": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"wavelength": { "type": "number", "description": "中心波长 (Å)" },
|
||||
"flux": { "type": "number", "description": "流量" },
|
||||
"flux_error": { "type": "number", "description": "流量误差" },
|
||||
"band": { "type": "string", "description": "波段名" }
|
||||
},
|
||||
"required": ["wavelength", "flux", "flux_error", "band"]
|
||||
},
|
||||
"description": "测光数据点列表"
|
||||
},
|
||||
"distance_pc": {
|
||||
"type": "number",
|
||||
"description": "距离 (pc)，可选，用于计算光度"
|
||||
}
|
||||
},
|
||||
"required": ["data_points"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let data_points: Vec<crate::services::analysis::sed::SedDataPoint> =
|
||||
match args.get("data_points").and_then(|d| d.as_array()) {
|
||||
Some(arr) => {
|
||||
let mut points = Vec::new();
|
||||
for item in arr {
|
||||
let wavelength = match item.get("wavelength").and_then(|w| w.as_f64()) {
|
||||
Some(w) => w,
|
||||
None => continue,
|
||||
};
|
||||
let flux = match item.get("flux").and_then(|f| f.as_f64()) {
|
||||
Some(f) => f,
|
||||
None => continue,
|
||||
};
|
||||
let flux_error = item
|
||||
.get("flux_error")
|
||||
.and_then(|e| e.as_f64())
|
||||
.unwrap_or(flux * 0.1);
|
||||
let band = item
|
||||
.get("band")
|
||||
.and_then(|b| b.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
points.push(crate::services::analysis::sed::SedDataPoint {
|
||||
wavelength,
|
||||
flux,
|
||||
flux_error,
|
||||
band,
|
||||
source: "user".to_string(),
|
||||
});
|
||||
}
|
||||
points
|
||||
}
|
||||
None => return ToolOutput::error("缺少必需参数 'data_points'"),
|
||||
};
|
||||
|
||||
if data_points.len() < 3 {
|
||||
return ToolOutput::error("至少需要 3 个测光点进行 SED 拟合");
|
||||
}
|
||||
|
||||
let distance_pc = args.get("distance_pc").and_then(|d| d.as_f64());
|
||||
|
||||
info!(
|
||||
"[FitSed] 数据点={}, 距离={:?} pc",
|
||||
data_points.len(),
|
||||
distance_pc
|
||||
);
|
||||
|
||||
match crate::services::analysis::sed::fit_sed(&data_points, distance_pc) {
|
||||
Ok(result) => {
|
||||
let report = crate::services::analysis::sed::format_sed_result(&result);
|
||||
let json_str = serde_json::to_string_pretty(&result).unwrap_or_default();
|
||||
let content = format!("{}\n\n```json\n{}\n```", report, json_str);
|
||||
ToolOutput::success(content, json!(result))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("SED 拟合失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
// src/agent/tools/astro/research/spectral_lines.rs
|
||||
//
|
||||
// MeasureSpectralLinesTool —— 谱线识别与等值宽度测量
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct MeasureSpectralLinesTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for MeasureSpectralLinesTool {
|
||||
fn name(&self) -> &str {
|
||||
"measure_spectral_lines"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"谱线测量"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"识别光谱中的常见谱线（Balmer 系列、Ca II H&K、Na D、Mg b、Fe 等），\
|
||||
自动计算等值宽度（EW）和半高全宽（FWHM）。输入光谱的波长和通量数组，\
|
||||
返回每条谱线的测量结果。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"wavelength": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "波长数组 (Å)"
|
||||
},
|
||||
"flux": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "通量数组"
|
||||
},
|
||||
"lines": {
|
||||
"type": "array",
|
||||
"items": { "type": "string" },
|
||||
"description": "要测量的谱线名称列表（可选）。不指定则测量所有内置谱线。可选值: Hα, Hβ, Hγ, Hδ, Ca II K, Ca II H, Na I D2, Na I D1, Mg I b1 等"
|
||||
},
|
||||
"window_half_width": {
|
||||
"type": "number",
|
||||
"description": "搜索窗口半宽 (Å)，默认 20",
|
||||
"default": 20.0
|
||||
}
|
||||
},
|
||||
"required": ["wavelength", "flux"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let wavelength: Vec<f32> = match args.get("wavelength").and_then(|w| w.as_array()) {
|
||||
Some(arr) => arr
|
||||
.iter()
|
||||
.filter_map(|v| v.as_f64().map(|f| f as f32))
|
||||
.collect(),
|
||||
None => return ToolOutput::error("缺少必需参数 'wavelength'"),
|
||||
};
|
||||
let flux: Vec<f32> = match args.get("flux").and_then(|w| w.as_array()) {
|
||||
Some(arr) => arr
|
||||
.iter()
|
||||
.filter_map(|v| v.as_f64().map(|f| f as f32))
|
||||
.collect(),
|
||||
None => return ToolOutput::error("缺少必需参数 'flux'"),
|
||||
};
|
||||
|
||||
let window_half_width = args
|
||||
.get("window_half_width")
|
||||
.and_then(|w| w.as_f64())
|
||||
.unwrap_or(20.0);
|
||||
|
||||
if wavelength.len() != flux.len() {
|
||||
return ToolOutput::error("波长和通量数组长度不一致");
|
||||
}
|
||||
if wavelength.len() < 3 {
|
||||
return ToolOutput::error("数据点太少（需要 ≥3）");
|
||||
}
|
||||
|
||||
let table = crate::services::spectrum::lines::builtin_line_table();
|
||||
|
||||
// 筛选要测量的谱线
|
||||
let target_lines: Vec<crate::services::spectrum::lines::SpectralLine> =
|
||||
if let Some(names) = args.get("lines").and_then(|l| l.as_array()) {
|
||||
names
|
||||
.iter()
|
||||
.filter_map(|n| n.as_str())
|
||||
.filter_map(|name| crate::services::spectrum::lines::find_line(name, &table))
|
||||
.cloned()
|
||||
.collect()
|
||||
} else {
|
||||
// 测量所有在光谱范围内的谱线
|
||||
let w_min = wavelength.iter().cloned().fold(f32::INFINITY, f32::min) as f64;
|
||||
let w_max = wavelength.iter().cloned().fold(f32::NEG_INFINITY, f32::max) as f64;
|
||||
crate::services::spectrum::lines::find_lines_in_range(w_min, w_max, &table)
|
||||
.into_iter()
|
||||
.cloned()
|
||||
.collect()
|
||||
};
|
||||
|
||||
if target_lines.is_empty() {
|
||||
return ToolOutput::error("在当前光谱范围内未找到匹配的谱线");
|
||||
}
|
||||
|
||||
info!(
|
||||
"[MeasureSpectralLines] 测量 {} 条谱线，窗口 ±{:.1} Å",
|
||||
target_lines.len(),
|
||||
window_half_width
|
||||
);
|
||||
|
||||
let measurements = crate::services::spectrum::lines::measure_lines(
|
||||
&wavelength,
|
||||
&flux,
|
||||
&target_lines,
|
||||
window_half_width,
|
||||
);
|
||||
|
||||
let successful: Vec<_> = measurements.into_iter().filter_map(|r| r.ok()).collect();
|
||||
|
||||
if successful.is_empty() {
|
||||
return ToolOutput::error("所有谱线测量均失败");
|
||||
}
|
||||
|
||||
let table_str = crate::services::spectrum::lines::format_measurements_table(&successful);
|
||||
|
||||
let content = format!(
|
||||
"谱线测量完成（{} / {} 条成功）:\n\n{}",
|
||||
successful.len(),
|
||||
target_lines.len(),
|
||||
table_str
|
||||
);
|
||||
|
||||
let metadata = json!({
|
||||
"measurements": successful.iter().map(|m| {
|
||||
json!({
|
||||
"line": m.line.name,
|
||||
"wavelength": m.center_wavelength,
|
||||
"ew": m.ew,
|
||||
"ew_error": m.ew_error,
|
||||
"fwhm": m.fwhm,
|
||||
"continuum_flux": m.continuum_flux
|
||||
})
|
||||
}).collect::<Vec<_>>(),
|
||||
"total_lines": target_lines.len(),
|
||||
"successful_measurements": successful.len()
|
||||
});
|
||||
|
||||
ToolOutput::success(content, metadata)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
// src/agent/tools/astro/research/spectrum_params.rs
|
||||
//
|
||||
// GetSpectrumParametersTool —— 从 FITS 文件中提取恒星大气参数
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct GetSpectrumParametersTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for GetSpectrumParametersTool {
|
||||
fn name(&self) -> &str {
|
||||
"get_spectrum_parameters"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"光谱参数提取"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"从已下载的 FITS 光谱文件中自动提取恒星大气参数（Teff, logg, [Fe/H], 径向速度等）。\
|
||||
支持 LAMOST、SDSS/BOSS、APOGEE、DESI 数据。结果结构化返回，可直接用于后续分析。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"enum": ["lamost", "sdss", "desi", "gaia"],
|
||||
"description": "数据源"
|
||||
},
|
||||
"subtype": {
|
||||
"type": "string",
|
||||
"description": "产品子类型（SDSS: spec/apstar/aspcap）"
|
||||
},
|
||||
"source_id": {
|
||||
"type": "string",
|
||||
"description": "源标识符"
|
||||
},
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "FITS 文件路径（相对于 library/ 目录）"
|
||||
}
|
||||
},
|
||||
"required": ["source", "source_id", "file_path"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let source_str = match args.get("source").and_then(|s| s.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'source'"),
|
||||
};
|
||||
let source_id = match args.get("source_id").and_then(|s| s.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'source_id'"),
|
||||
};
|
||||
let file_path = match args.get("file_path").and_then(|s| s.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return ToolOutput::error("缺少必需参数 'file_path'"),
|
||||
};
|
||||
let subtype = args.get("subtype").and_then(|s| s.as_str());
|
||||
|
||||
let source = match crate::services::observation::Source::parse(source_str) {
|
||||
Ok(s) => s,
|
||||
Err(e) => return ToolOutput::error(e),
|
||||
};
|
||||
|
||||
let product = crate::services::observation::types::ProductSpec {
|
||||
product: crate::services::observation::types::ProductType::Spectrum,
|
||||
subtype: subtype.map(|s| s.to_string()),
|
||||
};
|
||||
|
||||
info!(
|
||||
"[GetSpectrumParameters] source={}, id={}, file={}",
|
||||
source_str, source_id, file_path
|
||||
);
|
||||
|
||||
// 构造完整路径
|
||||
let full_path = if std::path::Path::new(file_path).is_absolute() {
|
||||
file_path.to_string()
|
||||
} else {
|
||||
let library_dir = ctx
|
||||
.app_state
|
||||
.config
|
||||
.storage
|
||||
.library_dir
|
||||
.to_str()
|
||||
.unwrap_or("library");
|
||||
format!("{}/{}", library_dir, file_path)
|
||||
};
|
||||
|
||||
match crate::services::spectrum::parameters::extract_parameters(
|
||||
&ctx.app_state,
|
||||
source,
|
||||
&product,
|
||||
source_id,
|
||||
&full_path,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(params) => {
|
||||
let json_str = serde_json::to_string_pretty(¶ms).unwrap_or_default();
|
||||
let content = format!(
|
||||
"光谱参数提取成功 ({}, {}):\n\n{}",
|
||||
source_str, source_id, json_str
|
||||
);
|
||||
ToolOutput::success(content, json!(params))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("参数提取失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
// src/agent/tools/astro/research/variable_star.rs
|
||||
//
|
||||
// ClassifyVariableStarTool —— 变星分类与特征提取
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
use tracing::info;
|
||||
|
||||
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
|
||||
|
||||
pub struct ClassifyVariableStarTool;
|
||||
|
||||
#[async_trait]
|
||||
impl AgentTool for ClassifyVariableStarTool {
|
||||
fn name(&self) -> &str {
|
||||
"classify_variable_star"
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"变星分类"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"自动分类变星类型（RR Lyrae、Cepheid、食双星、脉动变星、AGN 等），\
|
||||
提取振幅、偏度、峰度等统计特征。可选传入已知周期以提高分类准确度。"
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"time": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "时间数组 (天)"
|
||||
},
|
||||
"flux": {
|
||||
"type": "array",
|
||||
"items": { "type": "number" },
|
||||
"description": "通量/星等数组"
|
||||
},
|
||||
"period": {
|
||||
"type": "number",
|
||||
"description": "已知周期 (天)，可选。传入可提高分类准确度"
|
||||
}
|
||||
},
|
||||
"required": ["time", "flux"]
|
||||
})
|
||||
}
|
||||
|
||||
fn group(&self) -> &str {
|
||||
"as:research"
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_readonly(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
|
||||
let time: Vec<f64> = match args.get("time").and_then(|t| t.as_array()) {
|
||||
Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(),
|
||||
None => return ToolOutput::error("缺少必需参数 'time'"),
|
||||
};
|
||||
let flux: Vec<f64> = match args.get("flux").and_then(|f| f.as_array()) {
|
||||
Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(),
|
||||
None => return ToolOutput::error("缺少必需参数 'flux'"),
|
||||
};
|
||||
|
||||
let period = args.get("period").and_then(|p| p.as_f64());
|
||||
|
||||
if time.len() != flux.len() {
|
||||
return ToolOutput::error("时间和通量数组长度不一致");
|
||||
}
|
||||
if time.len() < 5 {
|
||||
return ToolOutput::error("数据点太少（需要 ≥5）");
|
||||
}
|
||||
|
||||
info!(
|
||||
"[ClassifyVariable] 数据点={}, 周期={}",
|
||||
time.len(),
|
||||
period
|
||||
.map(|p| format!("{:.4} 天", p))
|
||||
.unwrap_or_else(|| "未指定".to_string())
|
||||
);
|
||||
|
||||
match crate::services::timeseries::classification::classify_variable_star(
|
||||
&time, &flux, period,
|
||||
) {
|
||||
Ok(result) => {
|
||||
let report =
|
||||
crate::services::timeseries::classification::format_classification(&result);
|
||||
let json_str = serde_json::to_string_pretty(&result).unwrap_or_default();
|
||||
let content = format!("{}\n\n```json\n{}\n```", report, json_str);
|
||||
ToolOutput::success(content, json!(result))
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("分类失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,11 @@ impl AgentTool for SearchPapersTool {
|
||||
"search_papers"
|
||||
}
|
||||
|
||||
/// 搜索结果摘要来自 arXiv/ADS 等外部 API
|
||||
fn untrusted_output(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn display_name(&self) -> &str {
|
||||
"文献搜索"
|
||||
}
|
||||
@@ -123,9 +128,11 @@ impl AgentTool for SearchPapersTool {
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
|
||||
ToolOutput::success(
|
||||
// content 是 value 的模型侧投影；canonical value 供重放/审计消费
|
||||
ToolOutput::success_with_value(
|
||||
content,
|
||||
json!({ "count": results.len(), "papers": display_results }),
|
||||
json!({ "count": results.len() }),
|
||||
json!(display_results),
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{AgentTool, ToolContext, ToolOutput};
|
||||
use super::{AgentTool, ToolContext, ToolOutput, ToolSignals};
|
||||
|
||||
/// 手动上下文压缩工具：LLM 可主动调用以压缩对话历史
|
||||
pub struct CompressTool;
|
||||
@@ -22,6 +22,14 @@ impl AgentTool for CompressTool {
|
||||
true
|
||||
}
|
||||
|
||||
fn loop_signals(&self) -> ToolSignals {
|
||||
ToolSignals {
|
||||
resets_todo_nag: false,
|
||||
persists_todos: false,
|
||||
requests_manual_compact: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"手动压缩对话上下文。当你发现对话历史过长、token 消耗过大时，主动调用此工具进行压缩以释放空间。\
|
||||
压缩后历史对话将被摘要替代，但关键信息不会丢失。"
|
||||
|
||||
@@ -59,6 +59,19 @@ impl AgentTool for RunBashTool {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
fn causes_file_changes(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn hardline_check(
|
||||
&self,
|
||||
args: &serde_json::Value,
|
||||
) -> Option<crate::agent::runtime::hardline::HardlineResult> {
|
||||
args.get("command")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(crate::agent::runtime::hardline::check_command)
|
||||
}
|
||||
|
||||
fn check_permissions(
|
||||
&self,
|
||||
args: &serde_json::Value,
|
||||
|
||||
@@ -203,6 +203,20 @@ impl AgentTool for FileEditTool {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
fn causes_file_changes(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn hardline_check(
|
||||
&self,
|
||||
args: &serde_json::Value,
|
||||
) -> Option<crate::agent::runtime::hardline::HardlineResult> {
|
||||
args.get("file_path")
|
||||
.or_else(|| args.get("path"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(crate::agent::runtime::hardline::check_dangerous_path)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let path_str = match args.get("file_path").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
|
||||
@@ -49,6 +49,20 @@ impl AgentTool for FileWriteTool {
|
||||
InterruptBehavior::Block
|
||||
}
|
||||
|
||||
fn causes_file_changes(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn hardline_check(
|
||||
&self,
|
||||
args: &serde_json::Value,
|
||||
) -> Option<crate::agent::runtime::hardline::HardlineResult> {
|
||||
args.get("file_path")
|
||||
.or_else(|| args.get("path"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(crate::agent::runtime::hardline::check_dangerous_path)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
|
||||
let file_path_str = match args.get("file_path").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
|
||||
+203
-28
@@ -36,13 +36,24 @@ mod todo;
|
||||
|
||||
pub use ask_user::AskUserTool;
|
||||
pub use astro::analyze_image::AnalyzeImageTool;
|
||||
pub use astro::research::abundance_analysis::AnalyzeAbundancePatternTool;
|
||||
pub use astro::research::batch_process::BatchProcessTool;
|
||||
pub use astro::research::kinematics_tool::AnalyzeKinematicsTool;
|
||||
pub use astro::research::library::{GetCitationNetworkTool, SearchLocalLibraryTool};
|
||||
pub use astro::research::literature_monitor::LiteratureMonitorTool;
|
||||
pub use astro::research::metadata::GetPaperMetadataTool;
|
||||
pub use astro::research::note::SaveNoteTool;
|
||||
pub use astro::research::observation::FindObservationTool;
|
||||
pub use astro::research::observation_plan::PlanObservationTool;
|
||||
pub use astro::research::paper::{GetPaperContentTool, GetPaperOutlineTool};
|
||||
pub use astro::research::period_search::FindPeriodTool;
|
||||
pub use astro::research::radial_velocity::MeasureRadialVelocityTool;
|
||||
pub use astro::research::rag::RagSearchTool;
|
||||
pub use astro::research::sed_fit::FitSedTool;
|
||||
pub use astro::research::spectral_lines::MeasureSpectralLinesTool;
|
||||
pub use astro::research::spectrum_params::GetSpectrumParametersTool;
|
||||
pub use astro::research::target::QueryTargetTool;
|
||||
pub use astro::research::variable_star::ClassifyVariableStarTool;
|
||||
pub use astro::research::vizier::CatalogOperationTool;
|
||||
pub use astro::system::process::ProcessPaperTool;
|
||||
pub use astro::system::search::SearchPapersTool;
|
||||
@@ -123,6 +134,15 @@ impl ToolContext {
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置可选的 SSE 通道（静默上下文为 None）
|
||||
pub fn with_sse_tx_opt(
|
||||
mut self,
|
||||
tx: Option<tokio::sync::mpsc::UnboundedSender<crate::agent::runtime::AgentStreamEvent>>,
|
||||
) -> Self {
|
||||
self.sse_tx = tx;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置会话 ID
|
||||
pub fn with_session_id(mut self, id: String) -> Self {
|
||||
self.session_id = id;
|
||||
@@ -170,14 +190,24 @@ impl ToolContext {
|
||||
}
|
||||
|
||||
/// 工具执行结果
|
||||
///
|
||||
/// 输出契约（参考 deepseek-harness 的 canonical output 设计）：
|
||||
/// - `value`：canonical JSON 值——机器可用的结构化结果（重放/审计/剪枝/
|
||||
/// 前端渲染都消费这一份）；
|
||||
/// - `content`：`value` 的模型侧纯投影（render）——给 LLM 阅读的文本。
|
||||
///
|
||||
/// 新工具应优先通过 `success_with_value` 同时提供两者；只填 content 的
|
||||
/// 旧工具在消费侧退化为纯文本处理。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolOutput {
|
||||
/// 给大模型阅读的截断文本
|
||||
/// 给大模型阅读的截断文本（value 的 render 投影）
|
||||
pub content: String,
|
||||
/// 是否为错误
|
||||
pub is_error: bool,
|
||||
/// 结构化元数据（给前端 Timeline 直接渲染）
|
||||
pub metadata: serde_json::Value,
|
||||
/// canonical JSON 值（机器可用的结构化结果；None = 仅有文本投影）
|
||||
pub value: Option<serde_json::Value>,
|
||||
/// 跳过持久化到磁盘（用于 read_file 等已从磁盘读取内容的工具，避免级联持久化）
|
||||
pub skip_persist: bool,
|
||||
}
|
||||
@@ -189,6 +219,22 @@ impl ToolOutput {
|
||||
content: content.into(),
|
||||
is_error: false,
|
||||
metadata,
|
||||
value: None,
|
||||
skip_persist: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建成功结果（带 canonical value，新工具的推荐构造器）
|
||||
pub fn success_with_value(
|
||||
content: impl Into<String>,
|
||||
metadata: serde_json::Value,
|
||||
value: serde_json::Value,
|
||||
) -> Self {
|
||||
ToolOutput {
|
||||
content: content.into(),
|
||||
is_error: false,
|
||||
metadata,
|
||||
value: Some(value),
|
||||
skip_persist: false,
|
||||
}
|
||||
}
|
||||
@@ -199,16 +245,24 @@ impl ToolOutput {
|
||||
content: content.into(),
|
||||
is_error: false,
|
||||
metadata,
|
||||
value: None,
|
||||
skip_persist: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// 为已有结果附加 canonical value（builder 风格）
|
||||
pub fn with_value(mut self, value: serde_json::Value) -> Self {
|
||||
self.value = Some(value);
|
||||
self
|
||||
}
|
||||
|
||||
/// 创建错误结果
|
||||
pub fn error(msg: impl Into<String>) -> Self {
|
||||
ToolOutput {
|
||||
content: msg.into(),
|
||||
is_error: true,
|
||||
metadata: json!({}),
|
||||
value: None,
|
||||
skip_persist: false,
|
||||
}
|
||||
}
|
||||
@@ -223,6 +277,20 @@ pub enum InterruptBehavior {
|
||||
Block,
|
||||
}
|
||||
|
||||
/// 循环行为信号 — 工具向 ReAct 循环声明的副作用语义。
|
||||
///
|
||||
/// 用于替代主循环中按工具名字符串特判的散落逻辑：
|
||||
/// 循环通过 `ToolRegistry::signals_for()` 查询，工具自身声明行为。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq)]
|
||||
pub struct ToolSignals {
|
||||
/// 调用后重置 "todo nag" 计数（todo_write 自身）
|
||||
pub resets_todo_nag: bool,
|
||||
/// 调用后需要把 todos 参数持久化到任务板（agent_tasks 表）
|
||||
pub persists_todos: bool,
|
||||
/// 调用后请求在下一轮 LLM 调用前执行手动压缩（compress_context）
|
||||
pub requests_manual_compact: bool,
|
||||
}
|
||||
|
||||
/// 权限规则来源
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PermissionRuleSource {
|
||||
@@ -352,6 +420,36 @@ pub trait AgentTool: Send + Sync {
|
||||
fn group(&self) -> &str {
|
||||
"general"
|
||||
}
|
||||
|
||||
// ── 集中元数据声明：消除按工具名散落的平行名单 ──
|
||||
|
||||
/// 工具输出是否来自外部不可信源（web/检索/MCP 等返回的内容可能包含
|
||||
/// prompt 注入）。为 true 时执行器会把结果包裹在
|
||||
/// `<untrusted_tool_result>` 中再交给 LLM。
|
||||
fn untrusted_output(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// 该工具是否可能产生文件变更。为 true 时执行器会在调用前触发
|
||||
/// CheckpointManager 快照。
|
||||
fn causes_file_changes(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Hardline 预检查（任何模式下都不可绕过的参数级拒绝）。
|
||||
/// 返回 Some(blocked) 表示该调用的参数命中 hardline 规则；
|
||||
/// 默认 None 表示无 hardline 规则。
|
||||
fn hardline_check(
|
||||
&self,
|
||||
_args: &serde_json::Value,
|
||||
) -> Option<crate::agent::runtime::hardline::HardlineResult> {
|
||||
None
|
||||
}
|
||||
|
||||
/// 循环行为信号（见 `ToolSignals`）。
|
||||
fn loop_signals(&self) -> ToolSignals {
|
||||
ToolSignals::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// 工具注册表，管理所有可用工具。
|
||||
@@ -395,6 +493,18 @@ fn add_base_tools(registry: &mut ToolRegistry, skill_registry: Arc<RwLock<SkillR
|
||||
Box::new(FindObservationTool),
|
||||
Box::new(CatalogOperationTool),
|
||||
Box::new(SaveNoteTool),
|
||||
// 新增科研分析工具
|
||||
Box::new(GetSpectrumParametersTool),
|
||||
Box::new(MeasureSpectralLinesTool),
|
||||
Box::new(MeasureRadialVelocityTool),
|
||||
Box::new(FindPeriodTool),
|
||||
Box::new(ClassifyVariableStarTool),
|
||||
Box::new(FitSedTool),
|
||||
Box::new(AnalyzeKinematicsTool),
|
||||
Box::new(AnalyzeAbundancePatternTool),
|
||||
Box::new(PlanObservationTool),
|
||||
Box::new(BatchProcessTool),
|
||||
Box::new(LiteratureMonitorTool),
|
||||
Box::new(TodoWriteTool),
|
||||
Box::new(CompressTool),
|
||||
Box::new(AskUserTool),
|
||||
@@ -653,6 +763,22 @@ impl ToolRegistry {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 查询指定工具的循环行为信号。未知工具返回默认（无信号）。
|
||||
pub fn signals_for(&self, tool_name: &str) -> ToolSignals {
|
||||
self.tools
|
||||
.get(tool_name)
|
||||
.map(|t| t.loop_signals())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// 查询工具输出是否不可信。未知工具回退到名字启发式（mcp__*/web_* 前缀）。
|
||||
pub fn untrusted_output(&self, tool_name: &str) -> bool {
|
||||
self.tools
|
||||
.get(tool_name)
|
||||
.is_some_and(|t| t.untrusted_output())
|
||||
|| crate::agent::runtime::untrusted::is_high_risk(tool_name)
|
||||
}
|
||||
|
||||
/// 构建按 group 分组的紧凑工具目录字符串。
|
||||
/// 格式:
|
||||
/// [system]
|
||||
@@ -740,33 +866,82 @@ mod tests {
|
||||
"./skills",
|
||||
)))));
|
||||
let defs = registry.definitions();
|
||||
assert_eq!(defs.len(), 24);
|
||||
assert!(defs.iter().any(|d| d.function.name == "read_file"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "grep_files"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "glob_files"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "run_bash"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "file_write"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "file_edit"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "search_papers"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "process_paper"));
|
||||
assert!(defs
|
||||
.iter()
|
||||
.any(|d| d.function.name == "search_local_library"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "get_paper_metadata"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "get_paper_outline"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "get_paper_content"));
|
||||
assert!(defs
|
||||
.iter()
|
||||
.any(|d| d.function.name == "get_citation_network"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "rag_search"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "query_target"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "find_observation"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "catalog_operation"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "save_note"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "todo_write"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "compress_context"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "load_skill"));
|
||||
assert!(defs.iter().any(|d| d.function.name == "subagent"));
|
||||
// 关键工具必须在场（inventory 式断言——新增工具不需要改这里）
|
||||
for expected in [
|
||||
"read_file",
|
||||
"grep_files",
|
||||
"glob_files",
|
||||
"run_bash",
|
||||
"file_write",
|
||||
"file_edit",
|
||||
"search_papers",
|
||||
"process_paper",
|
||||
"search_local_library",
|
||||
"get_paper_metadata",
|
||||
"get_paper_outline",
|
||||
"get_paper_content",
|
||||
"get_citation_network",
|
||||
"rag_search",
|
||||
"query_target",
|
||||
"find_observation",
|
||||
"catalog_operation",
|
||||
"save_note",
|
||||
"todo_write",
|
||||
"compress_context",
|
||||
"load_skill",
|
||||
"subagent",
|
||||
] {
|
||||
assert!(
|
||||
defs.iter().any(|d| d.function.name == expected),
|
||||
"工具 {} 未注册",
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 工具清单一致性：名字唯一、元数据完整、声明元数据（signals 等）
|
||||
/// 与名字的一致性。新增工具若破坏这些不变量会在此失败而不是在运行时漂移。
|
||||
#[test]
|
||||
fn test_tool_registry_inventory_invariants() {
|
||||
let registry = ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from(
|
||||
"./skills",
|
||||
)))));
|
||||
let names = registry.tool_names();
|
||||
let mut unique = std::collections::HashSet::new();
|
||||
for name in &names {
|
||||
assert!(unique.insert(name.clone()), "工具名重复注册: {}", name);
|
||||
}
|
||||
for tool in registry.list() {
|
||||
let name = tool.name();
|
||||
assert!(!name.is_empty(), "存在空工具名");
|
||||
assert!(!tool.description().is_empty(), "工具 {} 缺少描述", name);
|
||||
assert!(
|
||||
tool.parameters().get("type").is_some(),
|
||||
"工具 {} 的参数 schema 缺少 type 字段",
|
||||
name
|
||||
);
|
||||
}
|
||||
// signals 与工具身份一致：todo_write 声明 todo 信号，compress_context 声明压缩信号
|
||||
let todo_signals = registry.signals_for("todo_write");
|
||||
assert!(todo_signals.resets_todo_nag && todo_signals.persists_todos);
|
||||
assert!(
|
||||
registry
|
||||
.signals_for("compress_context")
|
||||
.requests_manual_compact
|
||||
);
|
||||
// 只读工具声明不可信输出的例外：search_papers/rag_search 输出来自外部源
|
||||
assert!(registry.untrusted_output("search_papers"));
|
||||
assert!(registry.untrusted_output("rag_search"));
|
||||
assert!(!registry.untrusted_output("read_file"));
|
||||
// 文件变更工具必须声明 causes_file_changes
|
||||
for name in ["file_write", "file_edit", "run_bash"] {
|
||||
let tool = registry.get(name).unwrap();
|
||||
assert!(
|
||||
tool.causes_file_changes(),
|
||||
"{} 应声明 causes_file_changes",
|
||||
name
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -22,6 +22,18 @@ pub async fn maybe_persist_tool_result(
|
||||
tool_call_id: &str,
|
||||
max_chars: usize,
|
||||
tool_results_dir: &Path,
|
||||
) -> (String, Option<PathBuf>) {
|
||||
maybe_persist_tool_result_for(content, tool_call_id, max_chars, tool_results_dir, "tool").await
|
||||
}
|
||||
|
||||
/// 带 tool_name 的持久化变体：stub 中携带工具名与检索指引，
|
||||
/// 让模型能判断溢出内容的性质并有明确的取回路径。
|
||||
pub async fn maybe_persist_tool_result_for(
|
||||
content: &str,
|
||||
tool_call_id: &str,
|
||||
max_chars: usize,
|
||||
tool_results_dir: &Path,
|
||||
tool_name: &str,
|
||||
) -> (String, Option<PathBuf>) {
|
||||
if content.len() <= max_chars {
|
||||
return (content.to_string(), None);
|
||||
@@ -124,14 +136,18 @@ pub async fn maybe_persist_tool_result(
|
||||
|
||||
let stub = format!(
|
||||
"<persisted-output>\n\
|
||||
tool: {}\n\
|
||||
path: {}\n\
|
||||
size: {} chars\n\
|
||||
preview: |\n {}\n\n\
|
||||
完整输出已持久化到磁盘。使用 read_file 工具以 path 参数读取完整内容。\n\
|
||||
完整输出已溢出持久化到磁盘（超出 {} 字符上限）。\n\
|
||||
取回方式：调用 read_file 工具，path 填上述路径；文件较长时可用 max_lines 参数分段读取。\n\
|
||||
</persisted-output>",
|
||||
tool_name,
|
||||
file_path.display(),
|
||||
content.len(),
|
||||
preview.replace('\n', "\n "),
|
||||
max_chars,
|
||||
);
|
||||
|
||||
(stub, Some(file_path))
|
||||
|
||||
@@ -65,81 +65,43 @@ impl AgentTool for SearchHistoryTool {
|
||||
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;
|
||||
let mut results = Vec::new();
|
||||
|
||||
if scope == "all" || scope == "sessions" {
|
||||
match sqlx::query_as::<_, (String, String, String)>(
|
||||
"SELECT s.session_id, s.title, \
|
||||
snippet(agent_sessions_fts, 1, '<mark>', '</mark>', '...', 40) \
|
||||
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 10",
|
||||
match crate::services::search::search_agent_history(
|
||||
&ctx.app_state.db,
|
||||
raw_query,
|
||||
scope,
|
||||
10,
|
||||
None,
|
||||
)
|
||||
.bind(&query)
|
||||
.fetch_all(db)
|
||||
.await
|
||||
{
|
||||
Ok(rows) => {
|
||||
for (sid, title, snippet) in rows {
|
||||
results.push(json!({
|
||||
"type": "session",
|
||||
"session_id": sid,
|
||||
"title": title,
|
||||
"snippet": snippet,
|
||||
}));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("[SearchHistory] 搜索会话失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if scope == "all" || scope == "messages" {
|
||||
match sqlx::query_as::<_, (String, String, String, String)>(
|
||||
"SELECT fts.session_id, s.title, \
|
||||
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 80), fts.role \
|
||||
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 10",
|
||||
)
|
||||
.bind(&query)
|
||||
.fetch_all(db)
|
||||
.await
|
||||
{
|
||||
Ok(rows) => {
|
||||
for (sid, title, snippet, role) in rows {
|
||||
results.push(json!({
|
||||
"type": format!("message/{}", role),
|
||||
"session_id": sid,
|
||||
"title": title,
|
||||
"snippet": snippet,
|
||||
}));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("[SearchHistory] 搜索消息失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if results.is_empty() {
|
||||
Ok(search_results) => {
|
||||
if search_results.is_empty() {
|
||||
ToolOutput::success(
|
||||
format!("未找到与 \"{}\" 相关的历史记录。", query),
|
||||
json!({"query": query, "results": [], "count": 0}),
|
||||
format!("未找到与 \"{}\" 相关的历史记录。", raw_query),
|
||||
json!({"query": raw_query, "results": [], "count": 0}),
|
||||
)
|
||||
} else {
|
||||
let count = results.len();
|
||||
let results_json: Vec<serde_json::Value> = search_results
|
||||
.iter()
|
||||
.map(|r| {
|
||||
json!({
|
||||
"type": r.result_type,
|
||||
"session_id": r.session_id,
|
||||
"title": r.title,
|
||||
"snippet": r.snippet,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let count = results_json.len();
|
||||
ToolOutput::success(
|
||||
format!("找到 {} 条与 \"{}\" 相关的历史记录", count, query),
|
||||
json!({"query": query, "results": results, "count": count}),
|
||||
format!("找到 {} 条与 \"{}\" 相关的历史记录", count, raw_query),
|
||||
json!({"query": raw_query, "results": results_json, "count": count}),
|
||||
)
|
||||
}
|
||||
}
|
||||
Err(e) => ToolOutput::error(format!("检索历史失败: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -75,6 +75,15 @@ impl AgentTool for SubAgentTool {
|
||||
"type": "integer",
|
||||
"description": "子代理最大推理步数，默认5，最大10",
|
||||
"default": 5
|
||||
},
|
||||
"agent_name": {
|
||||
"type": "string",
|
||||
"description": "子代理的稳定名称（如 literature-scout）。命名后其历史会持久保留，后续可用 followup=true 继续与它对话而不丢失上下文。"
|
||||
},
|
||||
"followup": {
|
||||
"type": "boolean",
|
||||
"description": "设为 true 时续话同名的已有子代理：回放其全部历史并追加新指令，而不是从头开始。要求同时提供 agent_name。",
|
||||
"default": false
|
||||
}
|
||||
},
|
||||
"required": ["research_prompt"]
|
||||
@@ -98,10 +107,30 @@ impl AgentTool for SubAgentTool {
|
||||
.unwrap_or(5)
|
||||
.min(10) as usize;
|
||||
|
||||
let agent_name = args
|
||||
.get("agent_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
let followup = args
|
||||
.get("followup")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
// 参数校验（fail-loud）：followup 必须搭配 agent_name
|
||||
let agent_name = match (followup, agent_name) {
|
||||
(true, Some(name)) => Some(name),
|
||||
(true, None) => {
|
||||
return ToolOutput::error("followup=true 需要同时提供 agent_name 参数");
|
||||
}
|
||||
(false, name) => name,
|
||||
};
|
||||
|
||||
info!(
|
||||
"[SubAgent] 启动子代理: prompt_len={}, max_steps={}",
|
||||
"[SubAgent] 启动子代理: prompt_len={}, max_steps={}, name={:?}, followup={}",
|
||||
research_prompt.len(),
|
||||
max_steps
|
||||
max_steps,
|
||||
agent_name,
|
||||
followup
|
||||
);
|
||||
|
||||
// 创建子代理的 ToolRegistry（与父代理共享 skill_registry）
|
||||
@@ -157,15 +186,38 @@ impl AgentTool for SubAgentTool {
|
||||
)
|
||||
.with_parent_session(ctx.session_id.clone())
|
||||
.with_thinking(ctx.enable_thinking);
|
||||
let result = runner
|
||||
.run(&system_prompt, &research_prompt, max_steps)
|
||||
.await;
|
||||
// followup：续话已有命名子代理（回放其持久历史）；否则全新运行
|
||||
let result = if followup {
|
||||
runner
|
||||
.run_continuation(
|
||||
agent_name.as_deref().unwrap_or_default(),
|
||||
&system_prompt,
|
||||
&research_prompt,
|
||||
max_steps,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
runner
|
||||
.run_named(
|
||||
&system_prompt,
|
||||
&research_prompt,
|
||||
max_steps,
|
||||
agent_name.as_deref(),
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
if result.is_error {
|
||||
ToolOutput::error(format!("子代理执行失败: {}", result.content))
|
||||
} else {
|
||||
// 包装子代理结果，标注来源
|
||||
let final_content = format!("[子代理研究结果]\n\n{}", result.content);
|
||||
// 包装子代理结果，标注来源；命名子代理提示续话方式
|
||||
let mut final_content = format!("[子代理研究结果]\n\n{}", result.content);
|
||||
if let Some(name) = result.metadata.get("agent_name").and_then(|v| v.as_str()) {
|
||||
final_content.push_str(&format!(
|
||||
"\n\n（子代理名称: {}。如需在保留其上下文的情况下继续追问，再次调用 subagent 工具并传 agent_name=\"{}\" 与 followup=true。）",
|
||||
name, name
|
||||
));
|
||||
}
|
||||
info!(
|
||||
"[SubAgent] 子代理结果: content_len={}, preview={}...",
|
||||
final_content.len(),
|
||||
|
||||
+23
-2
@@ -9,7 +9,7 @@ use sqlx::SqlitePool;
|
||||
use std::str::FromStr;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::{AgentTool, ToolContext, ToolOutput};
|
||||
use super::{AgentTool, ToolContext, ToolOutput, ToolSignals};
|
||||
|
||||
/// 任务规划工具：让 LLM 在开始复杂研究前先制定计划，执行中更新进度。
|
||||
/// 任务状态持久化到 SQLite，支持 DAG 依赖。
|
||||
@@ -29,6 +29,14 @@ impl AgentTool for TodoWriteTool {
|
||||
true
|
||||
}
|
||||
|
||||
fn loop_signals(&self) -> ToolSignals {
|
||||
ToolSignals {
|
||||
resets_todo_nag: true,
|
||||
persists_todos: true,
|
||||
requests_manual_compact: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"任务规划工具。在开始复杂研究前列出待办事项，执行中标记进度（每项状态：pending/in_progress/completed）。\
|
||||
一次只能有一个 in_progress 任务。支持任务依赖（blockedBy：依赖的其他任务ID列表）。\
|
||||
@@ -131,7 +139,12 @@ impl AgentTool for TodoWriteTool {
|
||||
.push_str("\n💡 提示：还有待处理任务，请选择一个设为 in_progress 并开始执行。");
|
||||
}
|
||||
|
||||
ToolOutput::success(formatted, json!({ "task_count": todos.len() }))
|
||||
// canonical value = 全量任务状态（重放/任务板审计的机器可用投影）
|
||||
ToolOutput::success_with_value(
|
||||
formatted,
|
||||
json!({ "task_count": todos.len() }),
|
||||
json!({ "todos": todos }),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -310,6 +323,13 @@ mod tests {
|
||||
irsa: crate::clients::irsa::IrsaClient::new("https://irsa.ipac.caltech.edu", 60)
|
||||
.unwrap(),
|
||||
mast: crate::clients::mast::MastClient::new("https://mast.stsci.edu", 90).unwrap(),
|
||||
jwst: crate::clients::mast::jwst::JwstClient::new("https://mast.stsci.edu", 60)
|
||||
.unwrap(),
|
||||
heasarc: crate::clients::heasarc::HeasarcClient::new(
|
||||
"https://heasarc.gsfc.nasa.gov/FTP/heasarc/obs/dat/tap",
|
||||
60,
|
||||
)
|
||||
.unwrap(),
|
||||
observation_registry: std::sync::Arc::new(
|
||||
crate::services::observation::ObservationRegistry::default(),
|
||||
),
|
||||
@@ -326,6 +346,7 @@ mod tests {
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
},
|
||||
agent_runtimes: Arc::new(crate::agent::runtime::SessionRuntimeRegistry::new()),
|
||||
});
|
||||
|
||||
ToolContext {
|
||||
|
||||
+58
-14
@@ -16,7 +16,7 @@ use tracing::{error, info};
|
||||
|
||||
use super::error::{ApiResult, AppError};
|
||||
use super::AppState;
|
||||
use crate::agent::runtime::{AgentRuntime, AgentStreamEvent};
|
||||
use crate::agent::runtime::AgentStreamEvent;
|
||||
|
||||
// ── POST /api/chat/agent ──
|
||||
// SSE 流式智能体对话接口
|
||||
@@ -191,25 +191,68 @@ pub async fn chat_agent(
|
||||
None => (None, None),
|
||||
};
|
||||
|
||||
let mut runtime = AgentRuntime::new(Arc::clone(&state)).with_mode(&req.mode);
|
||||
// ── 会话解析与运行时复用 ──
|
||||
// 预分配会话 ID：新会话的首个请求也能写入取消标记并命中运行时缓存
|
||||
// （历史上新会话首请求超时只能 abort 任务，无法写入取消标记）
|
||||
let session_key = req
|
||||
.session_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
|
||||
// 模式回放：已存在的会话从 DB 读取创建时的模式。
|
||||
// 修复：恢复会话时传不同 mode 会静默改变行为——现在会话模式一次创建后保持稳定。
|
||||
let mode_id = crate::agent::runtime::session::load_session_mode(&state.db, &session_key)
|
||||
.await
|
||||
.unwrap_or_else(|| req.mode.clone());
|
||||
|
||||
let runtime = state
|
||||
.agent_runtimes
|
||||
.get_or_create(Arc::clone(&state), &session_key, &mode_id);
|
||||
// 只有 mode 未强制固定 thinking 时，用户才可以覆盖
|
||||
if runtime.mode_fixed_thinking().is_none() {
|
||||
if let Some(thinking) = req.thinking {
|
||||
runtime = runtime.with_thinking(thinking);
|
||||
runtime.set_thinking(thinking);
|
||||
}
|
||||
}
|
||||
|
||||
// 同会话 turn 串行化：并发请求 fail-loud 拒绝，防止 turn_index/消息顺序被破坏
|
||||
let turn_lock = state
|
||||
.agent_runtimes
|
||||
.turn_lock(&session_key)
|
||||
.unwrap_or_else(|| std::sync::Arc::new(tokio::sync::Mutex::new(())));
|
||||
// 快速检测：已占用直接 409（guard 立即释放，任务内部会重新 try_lock 兜底竞态）
|
||||
if turn_lock.try_lock().is_err() {
|
||||
return Err(AppError::conflict(
|
||||
"该会话正在执行中，请等待当前回合完成后再发送新消息",
|
||||
));
|
||||
}
|
||||
|
||||
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<AgentStreamEvent>();
|
||||
|
||||
let session_id = req.session_id.clone();
|
||||
let question_owned = question.clone();
|
||||
let image_context_owned = image_context.clone();
|
||||
let image_path_owned = image_path_for_db.clone();
|
||||
let cancel_session_key = session_key.clone();
|
||||
|
||||
// 在后台 tokio 任务中执行 Agent 循环
|
||||
let agent_handle = tokio::spawn(async move {
|
||||
// 持有 turn 串行锁直到回合结束（内部重新 try_lock 兜底竞态窗口）
|
||||
let _turn_guard = match turn_lock.try_lock() {
|
||||
Ok(guard) => guard,
|
||||
Err(_) => {
|
||||
let _ = tx.send(AgentStreamEvent::Error {
|
||||
message: "该会话正在执行中，请稍后重试。".to_string(),
|
||||
});
|
||||
let _ = tx.send(AgentStreamEvent::Done);
|
||||
return;
|
||||
}
|
||||
};
|
||||
match runtime
|
||||
.run_turn_with_image_context(
|
||||
session_id,
|
||||
&question,
|
||||
image_context,
|
||||
image_path_for_db,
|
||||
Some(session_key.clone()),
|
||||
&question_owned,
|
||||
image_context_owned,
|
||||
image_path_owned,
|
||||
tx.clone(),
|
||||
)
|
||||
.await
|
||||
@@ -236,9 +279,7 @@ pub async fn chat_agent(
|
||||
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
|
||||
if remaining.is_zero() {
|
||||
// 超时：通知 CancellationHook 停止 Agent，并中止后台任务
|
||||
if let Some(sid) = &req.session_id {
|
||||
cancelled_runs.insert(sid.clone(), ());
|
||||
}
|
||||
cancelled_runs.insert(cancel_session_key.clone(), ());
|
||||
agent_handle.abort();
|
||||
let timeout_event = AgentStreamEvent::Error {
|
||||
message: "Agent 执行超时（10 分钟），请重试。".to_string(),
|
||||
@@ -259,9 +300,7 @@ pub async fn chat_agent(
|
||||
Ok(None) => break, // channel closed
|
||||
Err(_) => {
|
||||
// 超时：通知 CancellationHook 停止 Agent，并中止后台任务
|
||||
if let Some(sid) = &req.session_id {
|
||||
cancelled_runs.insert(sid.clone(), ());
|
||||
}
|
||||
cancelled_runs.insert(cancel_session_key.clone(), ());
|
||||
agent_handle.abort();
|
||||
let timeout_event = AgentStreamEvent::Error {
|
||||
message: "Agent 执行超时（10 分钟），请重试。".to_string(),
|
||||
@@ -323,6 +362,11 @@ pub async fn delete_session(
|
||||
.await
|
||||
.map_err(|e| AppError::internal(format!("删除会话失败: {}", e)))?;
|
||||
|
||||
if success {
|
||||
// 同步移除会话级运行时缓存（后台队列/压缩日志等状态随之释放）
|
||||
state.agent_runtimes.remove(&session_id);
|
||||
}
|
||||
|
||||
if !success {
|
||||
return Err(AppError::not_found(format!(
|
||||
"会话 {} 不存在或已删除",
|
||||
|
||||
@@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize};
|
||||
use super::error::{ApiResult, AppError};
|
||||
use super::AppState;
|
||||
use crate::clients::cds::vizier::VizierQueryResult;
|
||||
use crate::services::analysis::hr_diagram::{build_hr_diagram, HrDiagramParams, HrDiagramResponse};
|
||||
|
||||
// ── 请求参数 ──
|
||||
|
||||
@@ -131,3 +132,15 @@ pub async fn cone_search(
|
||||
.await?;
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
/// GET /api/analysis/hr-diagram —— 赫罗图构建
|
||||
pub async fn hr_diagram(
|
||||
State(state): State<std::sync::Arc<AppState>>,
|
||||
Query(params): Query<HrDiagramParams>,
|
||||
) -> ApiResult<Json<HrDiagramResponse>> {
|
||||
let response = build_hr_diagram(&state.sources.gaia, ¶ms)
|
||||
.await
|
||||
.map_err(|e| AppError::internal(format!("构建赫罗图失败: {}", e)))?;
|
||||
|
||||
Ok(Json(response))
|
||||
}
|
||||
|
||||
@@ -75,6 +75,8 @@ pub struct DataSourceState {
|
||||
pub desi: DesiClient,
|
||||
pub irsa: crate::clients::irsa::IrsaClient,
|
||||
pub mast: crate::clients::mast::MastClient,
|
||||
pub jwst: crate::clients::mast::jwst::JwstClient,
|
||||
pub heasarc: crate::clients::heasarc::HeasarcClient,
|
||||
pub observation_registry: Arc<crate::services::observation::ObservationRegistry>,
|
||||
}
|
||||
|
||||
@@ -109,6 +111,9 @@ pub struct AppState {
|
||||
pub llm: LlmState,
|
||||
pub sources: DataSourceState,
|
||||
pub session: SessionState,
|
||||
/// 会话级 Agent 运行时注册表：同一会话复用 runtime（后台队列/压缩日志/
|
||||
/// 文件缓存等跨 turn 状态不再随请求销毁）
|
||||
pub agent_runtimes: Arc<crate::agent::runtime::SessionRuntimeRegistry>,
|
||||
}
|
||||
|
||||
// 统一标准化的文献格式，用于向前端传输
|
||||
|
||||
@@ -39,7 +39,7 @@ fn build_product_spec(
|
||||
subtype: &Option<String>,
|
||||
) -> Result<crate::services::observation::ProductSpec, AppError> {
|
||||
Ok(crate::services::observation::ProductSpec {
|
||||
product: crate::services::observation::ProductType::from_str(product)
|
||||
product: crate::services::observation::ProductType::parse(product)
|
||||
.map_err(AppError::bad_request)?,
|
||||
subtype: subtype.clone(),
|
||||
})
|
||||
@@ -68,7 +68,7 @@ pub async fn observation_search(
|
||||
State(state): State<std::sync::Arc<AppState>>,
|
||||
Query(params): Query<ObservationSearchParams>,
|
||||
) -> ApiResult<Json<Vec<crate::services::observation::Candidate>>> {
|
||||
let source = crate::services::observation::Source::from_str(¶ms.source)
|
||||
let source = crate::services::observation::Source::parse(¶ms.source)
|
||||
.map_err(AppError::bad_request)?;
|
||||
let product = build_product_spec(¶ms.product, ¶ms.subtype)?;
|
||||
|
||||
@@ -128,8 +128,8 @@ pub async fn observation_download(
|
||||
) -> ApiResult<Json<crate::services::observation::ObservationBatch>> {
|
||||
use crate::services::observation::{download_observation, ObservationRequest};
|
||||
|
||||
let source = crate::services::observation::Source::from_str(&req.source)
|
||||
.map_err(AppError::bad_request)?;
|
||||
let source =
|
||||
crate::services::observation::Source::parse(&req.source).map_err(AppError::bad_request)?;
|
||||
let product = build_product_spec(&req.product, &req.subtype)?;
|
||||
|
||||
let request = match req.mode {
|
||||
@@ -255,7 +255,7 @@ pub async fn observation_preview(
|
||||
State(state): State<std::sync::Arc<AppState>>,
|
||||
Query(params): Query<ObservationPreviewParams>,
|
||||
) -> ApiResult<Json<crate::services::observation::preview::ObservationPreview>> {
|
||||
let source = crate::services::observation::Source::from_str(¶ms.source)
|
||||
let source = crate::services::observation::Source::parse(¶ms.source)
|
||||
.map_err(AppError::bad_request)?;
|
||||
let product = build_product_spec(¶ms.product, ¶ms.subtype)?;
|
||||
|
||||
|
||||
@@ -215,3 +215,31 @@ pub async fn get_asset_batch_status(
|
||||
let status = state.batch_status.lock().await;
|
||||
Json(status.clone())
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ToggleScheduleRequest {
|
||||
pub is_scheduled: bool,
|
||||
}
|
||||
|
||||
pub async fn toggle_query_schedule(
|
||||
State(state): State<Arc<AppState>>,
|
||||
axum::extract::Path(id): axum::extract::Path<i64>,
|
||||
Json(req): Json<ToggleScheduleRequest>,
|
||||
) -> ApiResult<StatusCode> {
|
||||
let success =
|
||||
crate::services::batch::MetaSync::set_query_schedule(&state.db, id, req.is_scheduled)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("更新同步检索配置定时状态失败: {}", e);
|
||||
AppError::internal(format!("更新同步检索配置定时状态失败: {}", e))
|
||||
})?;
|
||||
|
||||
if !success {
|
||||
return Err(AppError::not_found(format!(
|
||||
"未找到 ID 为 {} 的检索配置",
|
||||
id
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(StatusCode::OK)
|
||||
}
|
||||
|
||||
@@ -282,6 +282,13 @@ async fn main() -> anyhow::Result<()> {
|
||||
config.cds.tess_timeout_secs,
|
||||
)
|
||||
.context("构建 MAST 客户端失败")?,
|
||||
jwst: astroresearch::clients::mast::jwst::JwstClient::new("https://mast.stsci.edu", 60)
|
||||
.context("构建 JWST 客户端失败")?,
|
||||
heasarc: astroresearch::clients::heasarc::HeasarcClient::new(
|
||||
"https://heasarc.gsfc.nasa.gov/FTP/heasarc/obs/dat/tap",
|
||||
60,
|
||||
)
|
||||
.context("构建 HEASARC 客户端失败")?,
|
||||
observation_registry: Arc::new(
|
||||
astroresearch::services::observation::ObservationRegistry::default(),
|
||||
),
|
||||
@@ -294,6 +301,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
||||
pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
||||
},
|
||||
agent_runtimes: Arc::new(astroresearch::agent::runtime::SessionRuntimeRegistry::new()),
|
||||
});
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
@@ -293,6 +293,11 @@ impl GaiaClient {
|
||||
serde_json::from_str::<TapJsonResponse>(&body).context("解析 Gaia TAP JSON 失败")
|
||||
}
|
||||
|
||||
/// 公开的 TAP 查询接口（JSON 格式）
|
||||
pub async fn tap_query(&self, adql: &str) -> anyhow::Result<String> {
|
||||
self.tap_sync(adql, "json").await
|
||||
}
|
||||
|
||||
/// 执行 TAP sync 查询（带重试）
|
||||
async fn tap_sync(&self, adql: &str, format: &str) -> anyhow::Result<String> {
|
||||
let params = [
|
||||
|
||||
@@ -0,0 +1,190 @@
|
||||
// src/clients/heasarc/mod.rs
|
||||
//
|
||||
// HEASARC（High Energy Astrophysics Science Archive Research Center）客户端
|
||||
//
|
||||
// 通过 TAP 服务查询 XMM-Newton 和 Chandra X 射线数据
|
||||
// 端点：POST https://heasarc.gsfc.nasa.gov/FTP/heasarc/.../tap/sync
|
||||
// 查询语言：ADQL
|
||||
|
||||
use anyhow::{anyhow, Context};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::Duration;
|
||||
use tracing::info;
|
||||
|
||||
/// X 射线观测行
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct XrayObservationRow {
|
||||
/// 观测标识
|
||||
pub obs_id: String,
|
||||
/// 目标名称
|
||||
pub target_name: String,
|
||||
/// 赤经（度）
|
||||
pub ra: f64,
|
||||
/// 赤纬（度）
|
||||
pub dec: f64,
|
||||
/// 望远镜（XMM-Newton / Chandra）
|
||||
pub telescope: String,
|
||||
/// 仪器（EPIC-pn/ACIS-S/...）
|
||||
pub instrument: String,
|
||||
/// 曝光时间（秒）
|
||||
pub exposure: Option<f64>,
|
||||
/// 观测日期
|
||||
pub obs_date: Option<String>,
|
||||
/// 数据存档 URL
|
||||
pub data_url: Option<String>,
|
||||
}
|
||||
|
||||
/// HEASARC TAP 客户端
|
||||
#[derive(Clone)]
|
||||
pub struct HeasarcClient {
|
||||
client: reqwest::Client,
|
||||
tap_url: String,
|
||||
}
|
||||
|
||||
impl HeasarcClient {
|
||||
pub fn new(tap_url: &str, timeout_secs: u64) -> anyhow::Result<Self> {
|
||||
Ok(HeasarcClient {
|
||||
client: reqwest::Client::builder()
|
||||
.redirect(crate::utils::ssrf::safe_redirect_policy())
|
||||
.timeout(Duration::from_secs(timeout_secs))
|
||||
.connect_timeout(Duration::from_secs(15))
|
||||
.build()
|
||||
.context("Failed to create HEASARC HTTP client")?,
|
||||
tap_url: tap_url.trim_end_matches('/').to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// TAP ADQL 查询
|
||||
pub async fn tap_query(&self, adql: &str) -> anyhow::Result<String> {
|
||||
info!("[HEASARC] TAP 查询: {}", &adql[..adql.len().min(100)]);
|
||||
let url = format!("{}/sync", self.tap_url);
|
||||
let resp = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("User-Agent", "AstroResearch/0.1")
|
||||
.form(&[
|
||||
("REQUEST", "doQuery"),
|
||||
("LANG", "ADQL"),
|
||||
("FORMAT", "json"),
|
||||
("MAXREC", "1000"),
|
||||
("QUERY", adql),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.context("HEASARC TAP 请求失败")?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let err_body = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!(
|
||||
"HEASARC TAP 错误 {}: {}",
|
||||
status,
|
||||
err_body.chars().take(500).collect::<String>()
|
||||
));
|
||||
}
|
||||
resp.text().await.context("读取 HEASARC 响应失败")
|
||||
}
|
||||
|
||||
/// XMM-Newton 锥形检索
|
||||
pub async fn xmm_cone_search(
|
||||
&self,
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
radius_deg: f64,
|
||||
) -> anyhow::Result<Vec<XrayObservationRow>> {
|
||||
let adql = format!(
|
||||
"SELECT obsid, target_name, ra, dec, instrument, exposure, obs_date \
|
||||
FROM heasarc_xmmobscat \
|
||||
WHERE 1=CONTAINS(POINT('ICRS', ra, dec), CIRCLE('ICRS', {}, {}, {})) \
|
||||
ORDER BY exposure DESC",
|
||||
ra, dec, radius_deg
|
||||
);
|
||||
let body = self.tap_query(&adql).await?;
|
||||
parse_tap_json(&body, "XMM-Newton")
|
||||
}
|
||||
|
||||
/// Chandra 锥形检索
|
||||
pub async fn chandra_cone_search(
|
||||
&self,
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
radius_deg: f64,
|
||||
) -> anyhow::Result<Vec<XrayObservationRow>> {
|
||||
let adql = format!(
|
||||
"SELECT obsid, target_name, ra, dec, instrument, exposure, obs_date \
|
||||
FROM heasarc_chandraobs \
|
||||
WHERE 1=CONTAINS(POINT('ICRS', ra, dec), CIRCLE('ICRS', {}, {}, {})) \
|
||||
ORDER BY exposure DESC",
|
||||
ra, dec, radius_deg
|
||||
);
|
||||
let body = self.tap_query(&adql).await?;
|
||||
parse_tap_json(&body, "Chandra")
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_tap_json(body: &str, telescope: &str) -> anyhow::Result<Vec<XrayObservationRow>> {
|
||||
let v: serde_json::Value = serde_json::from_str(body).context("解析 HEASARC JSON 失败")?;
|
||||
let fields = v
|
||||
.get("fields")
|
||||
.and_then(|f| f.as_array())
|
||||
.ok_or_else(|| anyhow!("HEASARC 响应缺 fields"))?;
|
||||
let data = v
|
||||
.get("data")
|
||||
.and_then(|d| d.as_array())
|
||||
.ok_or_else(|| anyhow!("HEASARC 响应缺 data"))?;
|
||||
|
||||
let col = |name: &str| -> Option<usize> {
|
||||
fields.iter().position(|f| {
|
||||
f.get("name").and_then(|n| n.as_str()) == Some(name)
|
||||
|| f.get("id").and_then(|n| n.as_str()) == Some(name)
|
||||
})
|
||||
};
|
||||
|
||||
let i_obsid = col("obsid");
|
||||
let i_target = col("target_name");
|
||||
let i_ra = col("ra");
|
||||
let i_dec = col("dec");
|
||||
let i_instrument = col("instrument");
|
||||
let i_exposure = col("exposure");
|
||||
let i_date = col("obs_date");
|
||||
|
||||
let mut rows = Vec::new();
|
||||
for row in data {
|
||||
if let Some(arr) = row.as_array() {
|
||||
let get_str = |i: Option<usize>| -> Option<String> {
|
||||
i.and_then(|idx| arr.get(idx)).and_then(|v| {
|
||||
if v.is_i64() {
|
||||
Some(v.as_i64().unwrap().to_string())
|
||||
} else if v.is_f64() {
|
||||
Some(format!("{:.0}", v.as_f64().unwrap()))
|
||||
} else {
|
||||
v.as_str().map(|s| s.to_string())
|
||||
}
|
||||
})
|
||||
};
|
||||
let get_f64 = |i: Option<usize>| -> Option<f64> {
|
||||
i.and_then(|idx| arr.get(idx)).and_then(|v| {
|
||||
v.as_f64()
|
||||
.or_else(|| v.as_str().and_then(|s| s.parse().ok()))
|
||||
})
|
||||
};
|
||||
|
||||
let obs_id = match get_str(i_obsid) {
|
||||
Some(id) => id,
|
||||
None => continue,
|
||||
};
|
||||
rows.push(XrayObservationRow {
|
||||
obs_id,
|
||||
target_name: get_str(i_target).unwrap_or_default(),
|
||||
ra: get_f64(i_ra).unwrap_or(0.0),
|
||||
dec: get_f64(i_dec).unwrap_or(0.0),
|
||||
telescope: telescope.to_string(),
|
||||
instrument: get_str(i_instrument).unwrap_or_default(),
|
||||
exposure: get_f64(i_exposure),
|
||||
obs_date: get_str(i_date),
|
||||
data_url: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(rows)
|
||||
}
|
||||
@@ -0,0 +1,290 @@
|
||||
// src/clients/mast/jwst.rs
|
||||
//
|
||||
// JWST（James Webb Space Telescope）MAST 数据查询
|
||||
//
|
||||
// 通过 MAST Portal API 查询 JWST NIRSpec/MIRI 光谱数据
|
||||
// 端点：POST https://mast.stsci.edu/api/v0/invoke
|
||||
// service: Mast.Jwst.Filtered.Position（坐标锥形检索）
|
||||
|
||||
use anyhow::{anyhow, Context};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::Duration;
|
||||
use tracing::info;
|
||||
|
||||
/// JWST 观测行（cone search 结果）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct JwstObservationRow {
|
||||
/// JWST 观测标识
|
||||
pub obs_id: String,
|
||||
/// 目标名称
|
||||
pub target_name: String,
|
||||
/// 赤经（度）
|
||||
pub ra: f64,
|
||||
/// 赤纬（度）
|
||||
pub dec: f64,
|
||||
/// 仪器名称（NIRSpec/MIRI/NIRCam/NIRISS）
|
||||
pub instrument: String,
|
||||
/// 观测模式（IFU/MOS/Imaging/...）
|
||||
pub observation_mode: String,
|
||||
/// 滤光片/光栅设置
|
||||
pub filter_grating: Option<String>,
|
||||
/// 曝光时间（秒）
|
||||
pub exposure_time: Option<f64>,
|
||||
/// 数据产品 URL
|
||||
pub data_url: Option<String>,
|
||||
/// 发布日期
|
||||
pub release_date: Option<String>,
|
||||
}
|
||||
|
||||
/// JWST 客户端
|
||||
#[derive(Clone)]
|
||||
pub struct JwstClient {
|
||||
client: reqwest::Client,
|
||||
base_url: String,
|
||||
}
|
||||
|
||||
impl JwstClient {
|
||||
pub fn new(base_url: &str, timeout_secs: u64) -> anyhow::Result<Self> {
|
||||
Ok(JwstClient {
|
||||
client: reqwest::Client::builder()
|
||||
.redirect(crate::utils::ssrf::safe_redirect_policy())
|
||||
.timeout(Duration::from_secs(timeout_secs))
|
||||
.connect_timeout(Duration::from_secs(10))
|
||||
.build()
|
||||
.context("Failed to create JWST HTTP client")?,
|
||||
base_url: base_url.trim_end_matches('/').to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// JWST 观测锥形检索（通过 MAST Portal API）
|
||||
pub async fn cone_search(
|
||||
&self,
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
radius_deg: f64,
|
||||
) -> anyhow::Result<Vec<JwstObservationRow>> {
|
||||
let url = format!("{}/api/v0/invoke", self.base_url);
|
||||
info!(
|
||||
"[JWST] ConeSearch ra={:.4} dec={:.4} radius={:.4}°",
|
||||
ra, dec, radius_deg
|
||||
);
|
||||
|
||||
let request_json = serde_json::json!({
|
||||
"service": "Mast.Jwst.Filtered.Position",
|
||||
"params": {
|
||||
"ra": ra,
|
||||
"dec": dec,
|
||||
"radius": radius_deg
|
||||
},
|
||||
"format": "json",
|
||||
"pagesize": 100,
|
||||
"page": 1,
|
||||
"removenullcolumns": true,
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let body = self
|
||||
.send_post_with_retry(&url, "request", &request_json)
|
||||
.await?;
|
||||
parse_jwst_json(&body)
|
||||
}
|
||||
|
||||
/// 查询 JWST 数据产品列表
|
||||
pub async fn list_data_products(&self, obs_id: &str) -> anyhow::Result<Vec<JwstDataProduct>> {
|
||||
let url = format!("{}/api/v0/invoke", self.base_url);
|
||||
let request_json = serde_json::json!({
|
||||
"service": "Mast.Jwst.Product",
|
||||
"params": {
|
||||
"obs_id": obs_id
|
||||
},
|
||||
"format": "json",
|
||||
"pagesize": 200,
|
||||
"page": 1,
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let body = self
|
||||
.send_post_with_retry(&url, "request", &request_json)
|
||||
.await?;
|
||||
parse_jwst_products_json(&body, obs_id)
|
||||
}
|
||||
|
||||
async fn send_post_with_retry(
|
||||
&self,
|
||||
url: &str,
|
||||
param_key: &str,
|
||||
param_value: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
const MAX_RETRIES: u32 = 3;
|
||||
let mut last_err = None;
|
||||
for attempt in 0..MAX_RETRIES {
|
||||
let resp = self
|
||||
.client
|
||||
.post(url)
|
||||
.header("User-Agent", "AstroResearch/0.1 (academic research tool)")
|
||||
.form(&[(param_key, param_value)])
|
||||
.send()
|
||||
.await;
|
||||
let resp = match resp {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
last_err = Some(e);
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let status = resp.status();
|
||||
if (status.as_u16() == 429 || 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);
|
||||
tokio::time::sleep(Duration::from_secs(retry_after)).await;
|
||||
continue;
|
||||
}
|
||||
if !status.is_success() {
|
||||
let err_body = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!(
|
||||
"JWST API 返回错误码: {}: {}",
|
||||
status,
|
||||
err_body.chars().take(300).collect::<String>()
|
||||
));
|
||||
}
|
||||
return resp.text().await.context("读取 JWST 响应失败");
|
||||
}
|
||||
Err(last_err
|
||||
.map(|e| anyhow!(e))
|
||||
.unwrap_or_else(|| anyhow!("JWST 重试耗尽")))
|
||||
}
|
||||
}
|
||||
|
||||
/// JWST 数据产品
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct JwstDataProduct {
|
||||
pub product_type: String,
|
||||
pub file_name: String,
|
||||
pub file_url: String,
|
||||
pub file_size: Option<u64>,
|
||||
pub instrument: String,
|
||||
}
|
||||
|
||||
fn parse_jwst_json(body: &str) -> anyhow::Result<Vec<JwstObservationRow>> {
|
||||
let v: serde_json::Value = serde_json::from_str(body).context("解析 JWST JSON 失败")?;
|
||||
let fields = v
|
||||
.get("fields")
|
||||
.and_then(|f| f.as_array())
|
||||
.ok_or_else(|| anyhow!("JWST 响应缺 fields"))?;
|
||||
let data = v
|
||||
.get("data")
|
||||
.and_then(|d| d.as_array())
|
||||
.ok_or_else(|| anyhow!("JWST 响应缺 data"))?;
|
||||
|
||||
let col = |name: &str| -> Option<usize> {
|
||||
fields
|
||||
.iter()
|
||||
.position(|f| f.get("name").and_then(|n| n.as_str()) == Some(name))
|
||||
};
|
||||
|
||||
let i_obsid = col("obs_id").or_else(|| col("obsid"));
|
||||
let i_target = col("target_name").or_else(|| col("targetid"));
|
||||
let i_ra = col("ra");
|
||||
let i_dec = col("dec");
|
||||
let i_instrument = col("instrument_name").or_else(|| col("instrument"));
|
||||
let i_mode = col("obs_mode").or_else(|| col("observation_mode"));
|
||||
let i_filter = col("filter_grating");
|
||||
let i_exptime = col("exptime").or_else(|| col("exposure_time"));
|
||||
let i_release = col("release_date");
|
||||
|
||||
let mut rows = Vec::new();
|
||||
for row in data {
|
||||
if let Some(arr) = row.as_array() {
|
||||
let get_str = |i: Option<usize>| -> Option<String> {
|
||||
i.and_then(|idx| arr.get(idx)).and_then(|v| {
|
||||
if v.is_i64() {
|
||||
Some(v.as_i64().unwrap().to_string())
|
||||
} else if v.is_f64() {
|
||||
Some(format!("{:.0}", v.as_f64().unwrap()))
|
||||
} else {
|
||||
v.as_str().map(|s| s.to_string())
|
||||
}
|
||||
})
|
||||
};
|
||||
let get_f64 = |i: Option<usize>| -> Option<f64> {
|
||||
i.and_then(|idx| arr.get(idx)).and_then(|v| {
|
||||
v.as_f64()
|
||||
.or_else(|| v.as_str().and_then(|s| s.parse().ok()))
|
||||
})
|
||||
};
|
||||
|
||||
let obs_id = match get_str(i_obsid) {
|
||||
Some(id) => id,
|
||||
None => continue,
|
||||
};
|
||||
rows.push(JwstObservationRow {
|
||||
obs_id,
|
||||
target_name: get_str(i_target).unwrap_or_default(),
|
||||
ra: get_f64(i_ra).unwrap_or(0.0),
|
||||
dec: get_f64(i_dec).unwrap_or(0.0),
|
||||
instrument: get_str(i_instrument).unwrap_or_default(),
|
||||
observation_mode: get_str(i_mode).unwrap_or_default(),
|
||||
filter_grating: get_str(i_filter),
|
||||
exposure_time: get_f64(i_exptime),
|
||||
data_url: None,
|
||||
release_date: get_str(i_release),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
fn parse_jwst_products_json(body: &str, _obs_id: &str) -> anyhow::Result<Vec<JwstDataProduct>> {
|
||||
let v: serde_json::Value = serde_json::from_str(body)?;
|
||||
let products = v
|
||||
.pointer("/data/products")
|
||||
.and_then(|p| p.as_array())
|
||||
.or_else(|| v.get("data").and_then(|d| d.as_array()))
|
||||
.ok_or_else(|| anyhow!("JWST Product 响应缺 products"))?;
|
||||
|
||||
let mut result = Vec::new();
|
||||
for prod in products {
|
||||
if let Some(obj) = prod.as_object() {
|
||||
let uri = obj
|
||||
.get("uri")
|
||||
.or_else(|| obj.get("dataURI"))
|
||||
.and_then(|u| u.as_str())
|
||||
.unwrap_or("");
|
||||
let url = if uri.starts_with("http") {
|
||||
uri.to_string()
|
||||
} else if uri.starts_with("mast:") {
|
||||
format!("https://mast.stsci.edu/api/v0.1/Download/file?uri={}", uri)
|
||||
} else {
|
||||
format!(
|
||||
"https://mast.stsci.edu/api/v0.1/Download/file?uri=mast:JWST/product/{}",
|
||||
uri
|
||||
)
|
||||
};
|
||||
result.push(JwstDataProduct {
|
||||
product_type: obj
|
||||
.get("productType")
|
||||
.and_then(|t| t.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
file_name: obj
|
||||
.get("filename")
|
||||
.and_then(|f| f.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
file_url: url,
|
||||
file_size: obj.get("fileSize").and_then(|s| s.as_u64()),
|
||||
instrument: obj
|
||||
.get("instrument_name")
|
||||
.and_then(|i| i.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
@@ -14,6 +14,8 @@
|
||||
//
|
||||
// 无需认证（公开数据）。复用 safe_redirect_policy + send_with_retry 模式。
|
||||
|
||||
pub mod jwst;
|
||||
|
||||
use anyhow::{anyhow, Context};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -3,6 +3,7 @@ pub mod arxiv;
|
||||
pub mod cds;
|
||||
pub mod desi;
|
||||
pub mod gaia;
|
||||
pub mod heasarc;
|
||||
pub mod irsa;
|
||||
pub mod lamost;
|
||||
pub mod llm;
|
||||
|
||||
+52
@@ -336,6 +336,13 @@ async fn main() -> anyhow::Result<()> {
|
||||
desi,
|
||||
irsa,
|
||||
mast,
|
||||
jwst: astroresearch::clients::mast::jwst::JwstClient::new("https://mast.stsci.edu", 60)
|
||||
.expect("Failed to create JWST client"),
|
||||
heasarc: astroresearch::clients::heasarc::HeasarcClient::new(
|
||||
"https://heasarc.gsfc.nasa.gov/FTP/heasarc/obs/dat/tap",
|
||||
60,
|
||||
)
|
||||
.expect("Failed to create HEASARC client"),
|
||||
observation_registry: Arc::new(
|
||||
astroresearch::services::observation::ObservationRegistry::default(),
|
||||
),
|
||||
@@ -348,6 +355,45 @@ async fn main() -> anyhow::Result<()> {
|
||||
pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
||||
pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
||||
},
|
||||
// 会话级 Agent 运行时注册表（跨请求复用 runtime 状态）
|
||||
agent_runtimes: Arc::new(astroresearch::agent::runtime::SessionRuntimeRegistry::new()),
|
||||
});
|
||||
|
||||
// 自治研究循环（IdlePoller）：env 门控接线，默认关闭。
|
||||
// 开启后空闲时轮询任务板，自动认领并执行可开始的任务。
|
||||
if std::env::var("AGENT_AUTONOMOUS_ENABLED")
|
||||
.map(|v| v == "true" || v == "1")
|
||||
.unwrap_or(false)
|
||||
{
|
||||
let config = astroresearch::agent::autonomous::AutoResearchConfig {
|
||||
enabled: true,
|
||||
..Default::default()
|
||||
};
|
||||
let poller = astroresearch::agent::autonomous::IdlePoller::new(app_state.clone(), config);
|
||||
tokio::spawn(async move {
|
||||
poller.start().await;
|
||||
});
|
||||
tracing::info!("[Autonomous] IdlePoller 已启用（AGENT_AUTONOMOUS_ENABLED）");
|
||||
}
|
||||
|
||||
// 周期回收空闲会话 runtime（2 小时 TTL，见 SessionRuntimeRegistry）
|
||||
{
|
||||
let registry = app_state.agent_runtimes.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(30 * 60)).await;
|
||||
let swept = registry.sweep_idle();
|
||||
if swept > 0 {
|
||||
tracing::info!("[SessionRegistry] 本轮回收 {} 个空闲 runtime", swept);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// 启动定时增量文献同步调度器
|
||||
let scheduler_state = app_state.clone();
|
||||
tokio::spawn(async move {
|
||||
astroresearch::services::batch::start_scheduled_sync_timer(scheduler_state).await;
|
||||
});
|
||||
|
||||
// 7. 设置 Axum 路由、CORS 头以及 React 仪表盘静态资源托管
|
||||
@@ -464,6 +510,10 @@ async fn main() -> anyhow::Result<()> {
|
||||
"/sync/queries/:id",
|
||||
axum::routing::delete(sync::delete_sync_query),
|
||||
)
|
||||
.route(
|
||||
"/sync/queries/:id/schedule",
|
||||
post(sync::toggle_query_schedule),
|
||||
)
|
||||
.route("/target/query", get(targets::query_target))
|
||||
.route("/target/associate", post(targets::associate_target))
|
||||
.route("/target/extract", post(targets::extract_paper_targets))
|
||||
@@ -496,6 +546,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
"/observation/unified/resolve",
|
||||
post(observation::unified_resolve),
|
||||
)
|
||||
// 分析路由
|
||||
.route("/analysis/hr-diagram", get(catalog::hr_diagram))
|
||||
// 智能体路由
|
||||
.route("/chat/agent", post(agent::chat_agent))
|
||||
.route("/chat/modes", get(agent::get_agent_modes))
|
||||
|
||||
@@ -0,0 +1,280 @@
|
||||
// src/services/analysis/chemistry.rs
|
||||
//
|
||||
// 化学丰度模式分析 —— [α/Fe] vs [Fe/H]，区分厚盘/薄盘/晕星
|
||||
//
|
||||
// 输入：元素丰度表（来自 APOGEE ASPCAP 或其他光谱分析）
|
||||
// 输出：[α/Fe] 计算、化学成分分类
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::super::spectrum::parameters::AbundanceEntry;
|
||||
|
||||
/// α 元素列表
|
||||
const ALPHA_ELEMENTS: &[&str] = &["O", "Mg", "Si", "S", "Ca", "Ti"];
|
||||
|
||||
/// 化学丰度分析结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChemistryResult {
|
||||
/// [Fe/H]
|
||||
pub feh: f64,
|
||||
/// [α/Fe]（α 元素加权平均）
|
||||
pub alpha_fe: f64,
|
||||
/// 各 α 元素的 [X/H] 值
|
||||
pub alpha_elements: Vec<ElementAbundance>,
|
||||
/// 化学分类
|
||||
pub classification: ChemistryClassification,
|
||||
/// 分类置信度
|
||||
pub confidence: f64,
|
||||
}
|
||||
|
||||
/// 单个元素丰度
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ElementAbundance {
|
||||
pub element: String,
|
||||
pub xh: f64, // [X/H]
|
||||
pub xfe: Option<f64>, // [X/Fe]
|
||||
pub error: Option<f64>,
|
||||
}
|
||||
|
||||
/// 化学分类
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct ChemistryClassification {
|
||||
/// 最可能的星族
|
||||
pub population: String,
|
||||
/// 描述
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// 从丰度表分析化学模式
|
||||
pub fn analyze_chemistry(abundances: &[AbundanceEntry]) -> Result<ChemistryResult> {
|
||||
if abundances.is_empty() {
|
||||
return Err(anyhow!("丰度数据为空"));
|
||||
}
|
||||
|
||||
// 提取 [Fe/H]
|
||||
let feh = abundances
|
||||
.iter()
|
||||
.find(|a| a.element == "Fe" || a.element == "Fe_H")
|
||||
.map(|a| a.value)
|
||||
.ok_or_else(|| anyhow!("未找到 [Fe/H] 数据"))?;
|
||||
|
||||
// 提取 α 元素
|
||||
let alpha_entries: Vec<&AbundanceEntry> = abundances
|
||||
.iter()
|
||||
.filter(|a| ALPHA_ELEMENTS.contains(&a.element.as_str()))
|
||||
.collect();
|
||||
|
||||
let alpha_elements: Vec<ElementAbundance> = alpha_entries
|
||||
.iter()
|
||||
.map(|a| {
|
||||
let xfe = abundances
|
||||
.iter()
|
||||
.find(|b| b.element == "Fe")
|
||||
.map(|fe| a.value - fe.value);
|
||||
ElementAbundance {
|
||||
element: a.element.clone(),
|
||||
xh: a.value,
|
||||
xfe,
|
||||
error: a.error,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 计算 [α/Fe]
|
||||
let alpha_fe = if alpha_elements.is_empty() {
|
||||
// 无 α 元素数据，尝试用 [M/H] - [Fe/H] 估计
|
||||
let mh = abundances
|
||||
.iter()
|
||||
.find(|a| a.element == "M" || a.element == "M_H")
|
||||
.map(|a| a.value);
|
||||
match mh {
|
||||
Some(mh) => mh - feh,
|
||||
None => 0.0, // 无法计算
|
||||
}
|
||||
} else {
|
||||
// 加权平均（以误差倒数为权重）
|
||||
let weighted_sum: f64 = alpha_elements
|
||||
.iter()
|
||||
.map(|a| {
|
||||
let weight = a
|
||||
.error
|
||||
.map(|e| if e > 0.0 { 1.0 / e } else { 1.0 })
|
||||
.unwrap_or(1.0);
|
||||
a.xfe.unwrap_or(a.xh - feh) * weight
|
||||
})
|
||||
.sum();
|
||||
let total_weight: f64 = alpha_elements
|
||||
.iter()
|
||||
.map(|a| {
|
||||
a.error
|
||||
.map(|e| if e > 0.0 { 1.0 / e } else { 1.0 })
|
||||
.unwrap_or(1.0)
|
||||
})
|
||||
.sum();
|
||||
if total_weight > 0.0 {
|
||||
weighted_sum / total_weight
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
};
|
||||
|
||||
// 化学分类（基于 [α/Fe] 和 [Fe/H] 的位置）
|
||||
let (classification, confidence) = classify_by_chemistry(feh, alpha_fe);
|
||||
|
||||
Ok(ChemistryResult {
|
||||
feh,
|
||||
alpha_fe,
|
||||
alpha_elements,
|
||||
classification,
|
||||
confidence,
|
||||
})
|
||||
}
|
||||
|
||||
/// 基于 [α/Fe] 和 [Fe/H] 的化学分类
|
||||
///
|
||||
/// 经验判据（Bensby et al. 2003, 2014）：
|
||||
/// - 薄盘：[α/Fe] < 0.25, [Fe/H] > -1.0
|
||||
/// - 厚盘：[α/Fe] > 0.25, [Fe/H] > -1.0
|
||||
/// - 晕：[Fe/H] < -1.0
|
||||
fn classify_by_chemistry(feh: f64, alpha_fe: f64) -> (ChemistryClassification, f64) {
|
||||
// 使用软判据（高斯距离）
|
||||
let feh_thin = -0.5;
|
||||
let feh_thick = -0.7;
|
||||
let alpha_thin = 0.15;
|
||||
let alpha_thick = 0.35;
|
||||
|
||||
// 薄盘中心
|
||||
let thin_dist =
|
||||
((feh - feh_thin).powi(2) / 0.5 + (alpha_fe - alpha_thin).powi(2) / 0.02).sqrt();
|
||||
// 厚盘中心
|
||||
let thick_dist =
|
||||
((feh - feh_thick).powi(2) / 0.3 + (alpha_fe - alpha_thick).powi(2) / 0.01).sqrt();
|
||||
// 晕中心
|
||||
let halo_dist = ((feh + 2.0).powi(2) / 1.0 + alpha_fe.powi(2) / 0.1).sqrt();
|
||||
|
||||
let thin_score = (-thin_dist).exp();
|
||||
let thick_score = (-thick_dist).exp();
|
||||
let halo_score = (-halo_dist).exp();
|
||||
|
||||
let total = thin_score + thick_score + halo_score;
|
||||
let thin_prob = thin_score / total;
|
||||
let thick_prob = thick_score / total;
|
||||
let halo_prob = halo_score / total;
|
||||
|
||||
let (population, description, confidence) = if thin_prob > thick_prob && thin_prob > halo_prob {
|
||||
(
|
||||
"薄盘".to_string(),
|
||||
format!(
|
||||
"薄盘星（[α/Fe]={:.2}, [Fe/H]={:.2}），年龄较轻，富金属，低 α 增丰",
|
||||
alpha_fe, feh
|
||||
),
|
||||
thin_prob,
|
||||
)
|
||||
} else if thick_prob > halo_prob {
|
||||
(
|
||||
"厚盘".to_string(),
|
||||
format!(
|
||||
"厚盘星（[α/Fe]={:.2}, [Fe/H]={:.2}），年龄较老，中等金属丰度，高 α 增丰",
|
||||
alpha_fe, feh
|
||||
),
|
||||
thick_prob,
|
||||
)
|
||||
} else {
|
||||
(
|
||||
"晕".to_string(),
|
||||
format!(
|
||||
"晕星（[α/Fe]={:.2}, [Fe/H]={:.2}），极低金属丰度，古老星族",
|
||||
alpha_fe, feh
|
||||
),
|
||||
halo_prob,
|
||||
)
|
||||
};
|
||||
|
||||
(
|
||||
ChemistryClassification {
|
||||
population,
|
||||
description,
|
||||
},
|
||||
confidence,
|
||||
)
|
||||
}
|
||||
|
||||
/// 从丰度数据中按名称查找 [Fe/H]
|
||||
pub fn find_feh(abundances: &[AbundanceEntry]) -> Option<f64> {
|
||||
abundances
|
||||
.iter()
|
||||
.find(|a| a.element == "Fe" || a.element == "Fe_H" || a.element == "[Fe/H]")
|
||||
.map(|a| a.value)
|
||||
}
|
||||
|
||||
/// 计算 [X/Fe] = [X/H] - [Fe/H]
|
||||
pub fn compute_xfe(xh: f64, feh: f64) -> f64 {
|
||||
xh - feh
|
||||
}
|
||||
|
||||
/// 格式化化学分析结果为 Markdown
|
||||
pub fn format_chemistry(result: &ChemistryResult) -> String {
|
||||
let mut out = String::from("## 化学丰度分析\n\n");
|
||||
out.push_str(&format!("- **[Fe/H]** = {:.3} dex\n", result.feh));
|
||||
out.push_str(&format!("- **[α/Fe]** = {:.3} dex\n\n", result.alpha_fe));
|
||||
|
||||
if !result.alpha_elements.is_empty() {
|
||||
out.push_str("### α 元素丰度\n\n");
|
||||
out.push_str("| 元素 | [X/H] | [X/Fe] | 误差 |\n");
|
||||
out.push_str("|------|-------|--------|------|\n");
|
||||
for e in &result.alpha_elements {
|
||||
let xfe_str = e
|
||||
.xfe
|
||||
.map(|v| format!("{:.3}", v))
|
||||
.unwrap_or_else(|| "—".to_string());
|
||||
let err_str = e
|
||||
.error
|
||||
.map(|v| format!("{:.3}", v))
|
||||
.unwrap_or_else(|| "—".to_string());
|
||||
out.push_str(&format!(
|
||||
"| {} | {:.3} | {} | {} |\n",
|
||||
e.element, e.xh, xfe_str, err_str
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
out.push_str(&format!(
|
||||
"\n### 化学分类: **{}**\n\n",
|
||||
result.classification.population
|
||||
));
|
||||
out.push_str(&format!("{}\n", result.classification.description));
|
||||
out.push_str(&format!("- 置信度: {:.1}%\n", result.confidence * 100.0));
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_classify_thin_disk() {
|
||||
let (class, conf) = classify_by_chemistry(-0.2, 0.1);
|
||||
assert_eq!(class.population, "薄盘");
|
||||
assert!(conf > 0.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_thick_disk() {
|
||||
let (class, conf) = classify_by_chemistry(-0.5, 0.35);
|
||||
assert_eq!(class.population, "厚盘");
|
||||
assert!(conf > 0.3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_classify_halo() {
|
||||
let (class, _conf) = classify_by_chemistry(-2.0, 0.0);
|
||||
assert_eq!(class.population, "晕");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_xfe() {
|
||||
assert!((compute_xfe(0.5, 0.2) - 0.3).abs() < 1e-10);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
// src/services/analysis/hr_diagram.rs
|
||||
//
|
||||
// 赫罗图（HR Diagram）服务逻辑
|
||||
|
||||
use crate::clients::gaia::GaiaClient;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct HrDiagramParams {
|
||||
pub ra: f64,
|
||||
pub dec: f64,
|
||||
pub radius: f64,
|
||||
pub max_mag: f64,
|
||||
pub min_parallax: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct HrDataPoint {
|
||||
pub source_id: String,
|
||||
pub g_mag: f64,
|
||||
pub bp_rp: f64,
|
||||
pub abs_g_mag: f64,
|
||||
pub ra: f64,
|
||||
pub dec: f64,
|
||||
pub parallax: f64,
|
||||
pub pm: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct HrDiagramResponse {
|
||||
pub points: Vec<HrDataPoint>,
|
||||
pub center_ra: f64,
|
||||
pub center_dec: f64,
|
||||
pub radius_deg: f64,
|
||||
pub total_sources: usize,
|
||||
}
|
||||
|
||||
pub async fn build_hr_diagram(
|
||||
gaia: &GaiaClient,
|
||||
params: &HrDiagramParams,
|
||||
) -> anyhow::Result<HrDiagramResponse> {
|
||||
let adql = format!(
|
||||
"SELECT source_id, ra, dec, phot_g_mean_mag, phot_bp_mean_mag, phot_rp_mean_mag, parallax, pm \
|
||||
FROM gaiadr3.gaia_source \
|
||||
WHERE 1=CONTAINS(POINT('ICRS', ra, dec), CIRCLE('ICRS', {}, {}, {})) \
|
||||
AND phot_g_mean_mag < {} \
|
||||
AND parallax > {} \
|
||||
AND phot_bp_mean_mag IS NOT NULL \
|
||||
AND phot_rp_mean_mag IS NOT NULL",
|
||||
params.ra, params.dec, params.radius, params.max_mag, params.min_parallax
|
||||
);
|
||||
|
||||
let result = gaia.tap_query(&adql).await?;
|
||||
|
||||
let v: serde_json::Value = serde_json::from_str(&result)?;
|
||||
|
||||
let fields = v
|
||||
.get("fields")
|
||||
.and_then(|f| f.as_array())
|
||||
.ok_or_else(|| anyhow::anyhow!("Gaia 响应缺 fields"))?;
|
||||
let data = v
|
||||
.get("data")
|
||||
.and_then(|d| d.as_array())
|
||||
.ok_or_else(|| anyhow::anyhow!("Gaia 响应缺 data"))?;
|
||||
|
||||
let col = |name: &str| -> Option<usize> {
|
||||
fields
|
||||
.iter()
|
||||
.position(|f| f.get("name").and_then(|n| n.as_str()) == Some(name))
|
||||
};
|
||||
|
||||
let i_sid = col("source_id");
|
||||
let i_ra = col("ra");
|
||||
let i_dec = col("dec");
|
||||
let i_g = col("phot_g_mean_mag");
|
||||
let i_bp = col("phot_bp_mean_mag");
|
||||
let i_rp = col("phot_rp_mean_mag");
|
||||
let i_plx = col("parallax");
|
||||
let i_pm = col("pm");
|
||||
|
||||
let mut points = Vec::new();
|
||||
for row in data {
|
||||
if let Some(arr) = row.as_array() {
|
||||
let get_f64 = |i: Option<usize>| -> Option<f64> {
|
||||
i.and_then(|idx| arr.get(idx)).and_then(|v| {
|
||||
if v.is_null() {
|
||||
None
|
||||
} else {
|
||||
v.as_f64()
|
||||
.or_else(|| v.as_str().and_then(|s| s.parse().ok()))
|
||||
}
|
||||
})
|
||||
};
|
||||
let get_str = |i: Option<usize>| -> Option<String> {
|
||||
i.and_then(|idx| arr.get(idx)).and_then(|v| {
|
||||
if v.is_i64() {
|
||||
Some(v.as_i64().unwrap().to_string())
|
||||
} else if v.is_f64() {
|
||||
Some(format!("{:.0}", v.as_f64().unwrap()))
|
||||
} else {
|
||||
v.as_str().map(|s| s.to_string())
|
||||
}
|
||||
})
|
||||
};
|
||||
|
||||
let sid = match get_str(i_sid) {
|
||||
Some(s) => s,
|
||||
None => continue,
|
||||
};
|
||||
let g = match get_f64(i_g) {
|
||||
Some(v) => v,
|
||||
None => continue,
|
||||
};
|
||||
let bp = match get_f64(i_bp) {
|
||||
Some(v) => v,
|
||||
None => continue,
|
||||
};
|
||||
let rp = match get_f64(i_rp) {
|
||||
Some(v) => v,
|
||||
None => continue,
|
||||
};
|
||||
let plx = match get_f64(i_plx) {
|
||||
Some(v) => v,
|
||||
None => continue,
|
||||
};
|
||||
if plx <= 0.0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let abs_g = g + 5.0 + 5.0 * (plx / 1000.0).log10();
|
||||
let bp_rp = bp - rp;
|
||||
|
||||
points.push(HrDataPoint {
|
||||
source_id: sid,
|
||||
g_mag: g,
|
||||
bp_rp,
|
||||
abs_g_mag: abs_g,
|
||||
ra: get_f64(i_ra).unwrap_or(0.0),
|
||||
dec: get_f64(i_dec).unwrap_or(0.0),
|
||||
parallax: plx,
|
||||
pm: get_f64(i_pm),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(HrDiagramResponse {
|
||||
total_sources: points.len(),
|
||||
points,
|
||||
center_ra: params.ra,
|
||||
center_dec: params.dec,
|
||||
radius_deg: params.radius,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,365 @@
|
||||
// src/services/analysis/kinematics.rs
|
||||
//
|
||||
// 银河系运动学分析 —— 自行/视差 → UVW 空间速度 + 移动星群识别
|
||||
//
|
||||
// 输入：(RA, Dec, parallax, pmRA, pmDec, RV)
|
||||
// 输出：(X, Y, Z, U, V, W) 银河系笛卡尔坐标 + 移动星群成员概率
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// 单个天体的运动学参数
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KinematicInput {
|
||||
/// 赤经 (度, J2000)
|
||||
pub ra: f64,
|
||||
/// 赤纬 (度, J2000)
|
||||
pub dec: f64,
|
||||
/// 视差 (mas)
|
||||
pub parallax: f64,
|
||||
/// 自行 RA (mas/yr)
|
||||
pub pm_ra: f64,
|
||||
/// 自行 Dec (mas/yr)
|
||||
pub pm_dec: f64,
|
||||
/// 径向速度 (km/s)
|
||||
pub rv: f64,
|
||||
/// 可选标识
|
||||
#[serde(default)]
|
||||
pub id: Option<String>,
|
||||
}
|
||||
|
||||
/// 运动学分析结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KinematicResult {
|
||||
/// 输入标识
|
||||
pub id: Option<String>,
|
||||
/// 银心距 X (kpc)
|
||||
pub x: f64,
|
||||
/// Y (kpc)
|
||||
pub y: f64,
|
||||
/// Z (kpc)
|
||||
pub z: f64,
|
||||
/// 银河系旋转方向速度 U (km/s, 向银心为正)
|
||||
pub u: f64,
|
||||
/// V (km/s, 银河系旋转方向为正)
|
||||
pub v: f64,
|
||||
/// W (km/s, 向北银极方向为正)
|
||||
pub w: f64,
|
||||
/// 距离 (pc)
|
||||
pub distance_pc: f64,
|
||||
/// 移动星群成员概率
|
||||
pub membership: Option<StellarGroupMembership>,
|
||||
}
|
||||
|
||||
/// 移动星群成员概率
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct StellarGroupMembership {
|
||||
/// 薄盘成员概率
|
||||
pub thin_disk: f64,
|
||||
/// 厚盘成员概率
|
||||
pub thick_disk: f64,
|
||||
/// 晕成员概率
|
||||
pub halo: f64,
|
||||
/// 最可能的分类
|
||||
pub best_fit: String,
|
||||
}
|
||||
|
||||
// ── 银河系常数 ──
|
||||
|
||||
/// 太阳到银心距离 R0 (kpc)
|
||||
const _R0: f64 = 8.178;
|
||||
/// 太阳在 Z 方向距银道面 Z0 (kpc)
|
||||
const Z0: f64 = 0.0208;
|
||||
/// 太阳的 U 速度 (km/s, 向银心)
|
||||
const U_SUN: f64 = 10.6;
|
||||
/// 太阳的 V 速度 (km/s, 银河系旋转方向)
|
||||
const V_SUN: f64 = 10.7;
|
||||
/// 太阳的 W 速度 (km/s, 向北银极)
|
||||
const W_SUN: f64 = 7.6;
|
||||
|
||||
/// 太阳运动学标准值
|
||||
const V_LSR: f64 = 220.0; // 本地静止标准速度 (km/s)
|
||||
|
||||
/// 从自行和视差计算距离 (pc)
|
||||
pub fn parallax_to_distance_pc(parallax_mas: f64) -> Result<f64> {
|
||||
if parallax_mas <= 0.0 {
|
||||
return Err(anyhow!("视差必须为正 (got {} mas)", parallax_mas));
|
||||
}
|
||||
Ok(1000.0 / parallax_mas)
|
||||
}
|
||||
|
||||
/// 赤道坐标 → 银道坐标 (l, b)
|
||||
///
|
||||
/// 使用标准 J2000 转换（参考论文：Bakis 1992, A&A 250, 304）
|
||||
/// 北银极 (NGP) 坐标：α_G = 192.85948°, δ_G = 27.12825°
|
||||
/// 银经零点（北天极在银道面上的投影）：l_NCP = 122.93192°
|
||||
pub fn equatorial_to_galactic(ra_deg: f64, dec_deg: f64) -> (f64, f64) {
|
||||
let ra = ra_deg.to_radians();
|
||||
let dec = dec_deg.to_radians();
|
||||
|
||||
// 北银极坐标 (J2000)
|
||||
let alpha_g = (192.85948_f64).to_radians(); // NGP 赤经
|
||||
let delta_g = (27.12825_f64).to_radians(); // NGP 赤纬
|
||||
let l_ncp = (122.93192_f64).to_radians(); // 北天极银经
|
||||
|
||||
// 银纬
|
||||
let sin_b = delta_g.sin() * dec.sin() + delta_g.cos() * dec.cos() * (ra - alpha_g).cos();
|
||||
let b = sin_b.asin();
|
||||
|
||||
// 银经
|
||||
let y = (ra - alpha_g).sin() * dec.cos();
|
||||
let x = delta_g.cos() * dec.sin() - delta_g.sin() * dec.cos() * (ra - alpha_g).cos();
|
||||
let l = l_ncp - y.atan2(x);
|
||||
|
||||
// 规范化到 [0, 360)
|
||||
let l_deg = l.to_degrees();
|
||||
let l_deg = if l_deg < 0.0 {
|
||||
l_deg + 360.0
|
||||
} else if l_deg >= 360.0 {
|
||||
l_deg - 360.0
|
||||
} else {
|
||||
l_deg
|
||||
};
|
||||
|
||||
(l_deg, b.to_degrees())
|
||||
}
|
||||
|
||||
/// 完整运动学分析：(RA, Dec, parallax, pmRA, pmDec, RV) → (X, Y, Z, U, V, W)
|
||||
pub fn compute_kinematics(input: &KinematicInput) -> Result<KinematicResult> {
|
||||
let distance_pc = parallax_to_distance_pc(input.parallax)?;
|
||||
let distance_kpc = distance_pc / 1000.0;
|
||||
|
||||
// 赤道坐标 → 银道坐标 (l, b) (度)
|
||||
let (l, b) = equatorial_to_galactic(input.ra, input.dec);
|
||||
let l_rad = l.to_radians();
|
||||
let b_rad = b.to_radians();
|
||||
|
||||
// 自行转为速度分量 (km/s)
|
||||
// (pm / 1000) * 4.74047 * d (pc) = 横向速度 (km/s)
|
||||
let mu_alpha = input.pm_ra * 0.00474047 * distance_pc;
|
||||
let mu_delta = input.pm_dec * 0.00474047 * distance_pc;
|
||||
|
||||
// 空间位置 (X, Y, Z) (kpc)
|
||||
// X: 向银心方向
|
||||
// Y: 银河系旋转方向
|
||||
// Z: 向北银极方向
|
||||
let cos_l = l_rad.cos();
|
||||
let sin_l = l_rad.sin();
|
||||
let cos_b = b_rad.cos();
|
||||
let sin_b = b_rad.sin();
|
||||
|
||||
let x = distance_kpc * cos_b * cos_l;
|
||||
let y = -distance_kpc * cos_b * sin_l; // 注意符号：银河系旋转方向
|
||||
let z = distance_kpc * sin_b + Z0;
|
||||
|
||||
// 空间速度 (U, V, W) (km/s)
|
||||
// 需要考虑自行的银道分量
|
||||
let cos_ra = input.ra.to_radians().cos();
|
||||
let sin_ra = input.ra.to_radians().sin();
|
||||
let cos_dec = input.dec.to_radians().cos();
|
||||
let sin_dec = input.dec.to_radians().sin();
|
||||
|
||||
// 行星运动方向的单位向量 (e_ra, e_dec, e_r) 在赤道坐标系中
|
||||
// e_ra = (-sin_ra, cos_ra, 0)
|
||||
// e_dec = (-sin_dec*cos_ra, -sin_dec*sin_ra, cos_dec)
|
||||
// e_r = (cos_dec*cos_ra, cos_dec*sin_ra, sin_dec)
|
||||
|
||||
// 横向速度在赤道坐标系中的分量
|
||||
let v_ra = mu_alpha; // km/s
|
||||
let v_dec = mu_delta; // km/s
|
||||
let v_r = input.rv; // km/s
|
||||
|
||||
// 赤道坐标系速度分量
|
||||
let vx_eq = -v_ra * sin_ra - v_dec * sin_dec * cos_ra + v_r * cos_dec * cos_ra;
|
||||
let vy_eq = v_ra * cos_ra - v_dec * sin_dec * sin_ra + v_r * cos_dec * sin_ra;
|
||||
let vz_eq = v_dec * cos_dec + v_r * sin_dec;
|
||||
|
||||
// 从赤道坐标系到银河系坐标系的旋转矩阵 (J2000)
|
||||
// 银经 l, 银纬 b, 银道面倾角
|
||||
let a_g = (266.405_f64).to_radians(); // 银心赤经
|
||||
let d_g = (-28.936_f64).to_radians(); // 银心赤纬
|
||||
let l_omega = (32.932_f64).to_radians(); // 升交点银经
|
||||
|
||||
let sin_lomega = l_omega.sin();
|
||||
let cos_lomega = l_omega.cos();
|
||||
let cos_dg = d_g.cos();
|
||||
let sin_dg = d_g.sin();
|
||||
|
||||
// 银河系速度分量
|
||||
let u_gal = -vx_eq * (a_g.sin() * sin_lomega + a_g.cos() * cos_lomega * sin_dg)
|
||||
- vy_eq * (-a_g.cos() * sin_lomega + a_g.sin() * cos_lomega * sin_dg)
|
||||
- vz_eq * cos_lomega * cos_dg;
|
||||
let v_gal = vx_eq * (-a_g.sin() * cos_lomega + a_g.cos() * sin_lomega * sin_dg)
|
||||
+ vy_eq * (a_g.cos() * cos_lomega + a_g.sin() * sin_lomega * sin_dg)
|
||||
+ vz_eq * sin_lomega * cos_dg;
|
||||
let w_gal = vx_eq * a_g.cos() * cos_dg + vy_eq * a_g.sin() * cos_dg - vz_eq * sin_dg;
|
||||
|
||||
// 转换为标准 UVW（减去太阳运动）
|
||||
let u = u_gal + U_SUN;
|
||||
let v = v_gal + V_SUN + V_LSR;
|
||||
let w = w_gal + W_SUN;
|
||||
|
||||
// 移动星群成员概率
|
||||
let membership = classify_stellar_group(u, v, w, z, distance_kpc);
|
||||
|
||||
Ok(KinematicResult {
|
||||
id: input.id.clone(),
|
||||
x,
|
||||
y,
|
||||
z,
|
||||
u,
|
||||
v,
|
||||
w,
|
||||
distance_pc,
|
||||
membership: Some(membership),
|
||||
})
|
||||
}
|
||||
|
||||
/// 简化的移动星群分类（基于 UVW 速度和 Z 高度）
|
||||
fn classify_stellar_group(
|
||||
_u: f64,
|
||||
v: f64,
|
||||
w: f64,
|
||||
z_kpc: f64,
|
||||
_distance_kpc: f64,
|
||||
) -> StellarGroupMembership {
|
||||
// 简化规则：
|
||||
// 薄盘：|Z| < 0.3 kpc, |V - V_LSR| < 50 km/s, |W| < 40 km/s
|
||||
// 厚盘：0.3 < |Z| < 1.5 kpc, |V - V_LSR| < 100 km/s
|
||||
// 晕：|Z| > 1.5 kpc, 高速运动
|
||||
|
||||
let abs_z = z_kpc.abs();
|
||||
let v_offset = (v - V_LSR).abs();
|
||||
let abs_w = w.abs();
|
||||
|
||||
// 计算各成分的概率（基于距离度量的软分类）
|
||||
let thin_disk_score = (-abs_z / 0.3 - v_offset / 50.0 - abs_w / 40.0).exp();
|
||||
let thick_disk_score = (-(abs_z - 0.8).abs() / 0.5 - v_offset / 100.0 - abs_w / 60.0).exp();
|
||||
let halo_score = if abs_z > 1.0 {
|
||||
(-(abs_z - 2.0).abs() / 1.0 - v_offset / 200.0 - abs_w / 100.0).exp()
|
||||
} else {
|
||||
0.01
|
||||
};
|
||||
|
||||
// 归一化
|
||||
let total = thin_disk_score + thick_disk_score + halo_score;
|
||||
let thin_disk = thin_disk_score / total;
|
||||
let thick_disk = thick_disk_score / total;
|
||||
let halo = halo_score / total;
|
||||
|
||||
let best_fit = if thin_disk > thick_disk && thin_disk > halo {
|
||||
"薄盘".to_string()
|
||||
} else if thick_disk > halo {
|
||||
"厚盘".to_string()
|
||||
} else {
|
||||
"晕".to_string()
|
||||
};
|
||||
|
||||
StellarGroupMembership {
|
||||
thin_disk,
|
||||
thick_disk,
|
||||
halo,
|
||||
best_fit,
|
||||
}
|
||||
}
|
||||
|
||||
/// 批量运动学分析
|
||||
pub fn compute_kinematics_batch(inputs: &[KinematicInput]) -> Vec<Result<KinematicResult>> {
|
||||
inputs.iter().map(compute_kinematics).collect()
|
||||
}
|
||||
|
||||
/// 格式化运动学结果为 Markdown
|
||||
pub fn format_kinematic_result(result: &KinematicResult) -> String {
|
||||
let mut out = String::from("## 运动学分析结果\n\n");
|
||||
|
||||
if let Some(ref id) = result.id {
|
||||
out.push_str(&format!("**目标**: {}\n\n", id));
|
||||
}
|
||||
|
||||
out.push_str("### 三维位置 (kpc)\n\n");
|
||||
out.push_str(&format!("- X = {:.3} kpc\n", result.x));
|
||||
out.push_str(&format!("- Y = {:.3} kpc\n", result.y));
|
||||
out.push_str(&format!("- Z = {:.3} kpc\n", result.z));
|
||||
out.push_str(&format!("- 距离 = {:.1} pc\n\n", result.distance_pc));
|
||||
|
||||
out.push_str("### 三维速度 (km/s)\n\n");
|
||||
out.push_str(&format!("- U = {:.1} km/s (向银心)\n", result.u));
|
||||
out.push_str(&format!("- V = {:.1} km/s (银河系旋转)\n", result.v));
|
||||
out.push_str(&format!("- W = {:.1} km/s (向北银极)\n", result.w));
|
||||
|
||||
if let Some(ref membership) = result.membership {
|
||||
out.push_str(&format!(
|
||||
"\n### 移动星群分类: **{}**\n\n",
|
||||
membership.best_fit
|
||||
));
|
||||
out.push_str(&format!("- 薄盘: {:.1}%\n", membership.thin_disk * 100.0));
|
||||
out.push_str(&format!("- 厚盘: {:.1}%\n", membership.thick_disk * 100.0));
|
||||
out.push_str(&format!("- 晕: {:.1}%\n", membership.halo * 100.0));
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parallax_to_distance() {
|
||||
let d = parallax_to_distance_pc(100.0).unwrap();
|
||||
assert!((d - 10.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parallax_zero_error() {
|
||||
assert!(parallax_to_distance_pc(0.0).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_equatorial_to_galactic_sgp() {
|
||||
// 南银极 (SGP) 赤道坐标 (J2000): RA ≈ 12.86°, Dec ≈ -27.13°
|
||||
let (_l, b) = equatorial_to_galactic(12.85948, -27.12825);
|
||||
// 银纬应为 -90°（南银极）
|
||||
assert!((b - (-90.0)).abs() < 0.01, "SGP 银纬 {} 应为 -90°", b);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kinematics_solar() {
|
||||
// 太阳附近的恒星，自行小，RV 小
|
||||
let input = KinematicInput {
|
||||
ra: 0.0,
|
||||
dec: 0.0,
|
||||
parallax: 1000.0, // 1 pc
|
||||
pm_ra: 0.0,
|
||||
pm_dec: 0.0,
|
||||
rv: 0.0,
|
||||
id: Some("test".to_string()),
|
||||
};
|
||||
let result = compute_kinematics(&input).unwrap();
|
||||
assert!((result.distance_pc - 1.0).abs() < 0.1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kinematics_high_pm() {
|
||||
// 距离 10 pc 的恒星，自行 1000 mas/yr (1 arcsec/yr)
|
||||
// 速度应为 1 * 10 * 4.74047 = 47.4047 km/s
|
||||
let input = KinematicInput {
|
||||
ra: 0.0,
|
||||
dec: 0.0,
|
||||
parallax: 100.0, // 10 pc
|
||||
pm_ra: 1000.0, // 1000 mas/yr
|
||||
pm_dec: 0.0,
|
||||
rv: 0.0,
|
||||
id: Some("high_pm".to_string()),
|
||||
};
|
||||
let result = compute_kinematics(&input).unwrap();
|
||||
assert!((result.distance_pc - 10.0).abs() < 0.01);
|
||||
|
||||
// 还原太阳运动，检查变换后空间速度大小是否等于输入横向速度
|
||||
let u_gal = result.u - U_SUN;
|
||||
let v_gal = result.v - V_SUN - V_LSR;
|
||||
let w_gal = result.w - W_SUN;
|
||||
let v_mag = (u_gal * u_gal + v_gal * v_gal + w_gal * w_gal).sqrt();
|
||||
assert!((v_mag - 47.4047).abs() < 1e-4);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
// src/services/analysis/mod.rs
|
||||
//
|
||||
// 天体物理分析服务层 —— SED 拟合、运动学、化学丰度、观测规划
|
||||
//
|
||||
// 模块布局：
|
||||
// sed.rs —— 多波段 SED 拟合（黑体/模板），求 T_eff, L, A_V
|
||||
// kinematics.rs —— Gaia 自行/视差 → 银河系 UVW 空间速度 + 移动星群识别
|
||||
// chemistry.rs —— 化学丰度模式分析（[α/Fe] vs [Fe/H]，厚盘/薄盘/晕星区分）
|
||||
// observability.rs —— 观测提案辅助（目标可见性、月相、曝光时间估算）
|
||||
|
||||
pub mod chemistry;
|
||||
pub mod hr_diagram;
|
||||
pub mod kinematics;
|
||||
pub mod observability;
|
||||
pub mod sed;
|
||||
@@ -0,0 +1,401 @@
|
||||
// src/services/analysis/observability.rs
|
||||
//
|
||||
// 观测提案辅助 —— 目标可见性、月相、曝光时间估算
|
||||
//
|
||||
// 输入：天体坐标、观测台站、时间范围
|
||||
// 输出：升起/落下时间、最佳观测窗口、月相影响、曝光时间估算
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// 观测台站参数
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Observatory {
|
||||
/// 台站名称
|
||||
pub name: String,
|
||||
/// 经度 (度, 东经为正)
|
||||
pub longitude: f64,
|
||||
/// 纬度 (度, 北纬为正)
|
||||
pub latitude: f64,
|
||||
/// 海拔 (米)
|
||||
pub altitude: f64,
|
||||
/// 典型视宁度 (角秒)
|
||||
pub seeing: f64,
|
||||
/// 大气消光系数 (mag/airmass, V 波段)
|
||||
pub extinction_v: f64,
|
||||
}
|
||||
|
||||
impl Observatory {
|
||||
/// 常用天文台站
|
||||
pub fn lick_3m() -> Self {
|
||||
Self {
|
||||
name: "Lick 3m".to_string(),
|
||||
longitude: -121.6428,
|
||||
latitude: 37.3414,
|
||||
altitude: 1283.0,
|
||||
seeing: 1.2,
|
||||
extinction_v: 0.15,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn keck() -> Self {
|
||||
Self {
|
||||
name: "Keck I/II".to_string(),
|
||||
longitude: -155.4761,
|
||||
latitude: 19.8263,
|
||||
altitude: 4145.0,
|
||||
seeing: 0.8,
|
||||
extinction_v: 0.12,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn lamost() -> Self {
|
||||
Self {
|
||||
name: "LAMOST".to_string(),
|
||||
longitude: 117.5837,
|
||||
latitude: 40.3958,
|
||||
altitude: 900.0,
|
||||
seeing: 2.0,
|
||||
extinction_v: 0.20,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn gemini_north() -> Self {
|
||||
Self {
|
||||
name: "Gemini North".to_string(),
|
||||
longitude: -155.4691,
|
||||
latitude: 19.8238,
|
||||
altitude: 2737.0,
|
||||
seeing: 0.7,
|
||||
extinction_v: 0.10,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn default_list() -> Vec<Self> {
|
||||
vec![
|
||||
Self::lick_3m(),
|
||||
Self::keck(),
|
||||
Self::lamost(),
|
||||
Self::gemini_north(),
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
/// 观测可行性结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ObservabilityResult {
|
||||
/// 目标升起时间 (UTC)
|
||||
pub rise_time: Option<String>,
|
||||
/// 目标落下时间 (UTC)
|
||||
pub set_time: Option<String>,
|
||||
/// 最大高度角 (度)
|
||||
pub max_altitude: f64,
|
||||
/// 最大高度角时间 (UTC)
|
||||
pub max_altitude_time: Option<String>,
|
||||
/// 今晚是否可观测（高度角 > 30° 且天黑）
|
||||
pub observable_tonight: bool,
|
||||
/// 最佳观测窗口（高度角 > 30° 的时段）
|
||||
pub best_window: Option<String>,
|
||||
/// 月相影响
|
||||
pub moon_impact: MoonImpact,
|
||||
/// 曝光时间估算（V 波段，1σ = 5σ 目标星等）
|
||||
pub exposure_estimate: Option<ExposureEstimate>,
|
||||
/// 大气质量（在最大高度角时）
|
||||
pub airmass_at_max: f64,
|
||||
}
|
||||
|
||||
/// 月相影响
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MoonImpact {
|
||||
/// 月相角 (0=新月, 180=满月)
|
||||
pub phase_angle: f64,
|
||||
/// 月相描述
|
||||
pub phase_description: String,
|
||||
/// 月亮距目标角距离 (度)
|
||||
pub separation: f64,
|
||||
/// 对观测的影响等级
|
||||
pub impact: String,
|
||||
}
|
||||
|
||||
/// 曝光时间估算
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ExposureEstimate {
|
||||
/// V 波段目标星等
|
||||
pub target_v_mag: f64,
|
||||
/// 所需信噪比
|
||||
pub target_snr: f64,
|
||||
/// 估算曝光时间 (秒)
|
||||
pub exposure_time_secs: f64,
|
||||
/// 所需帧数（含 dithering）
|
||||
pub n_frames: u32,
|
||||
/// 总时间 (秒)
|
||||
pub total_time_secs: f64,
|
||||
/// 备注
|
||||
pub notes: String,
|
||||
}
|
||||
|
||||
/// 估算目标在某时刻的高度角和方位角
|
||||
///
|
||||
/// 输入：目标赤经赤纬 (度), 观测台站, 儒略日
|
||||
pub fn compute_altitude(ra_deg: f64, dec_deg: f64, obs: &Observatory, jd: f64) -> (f64, f64) {
|
||||
// 简化的地平坐标计算
|
||||
let _ra = ra_deg.to_radians();
|
||||
let dec = dec_deg.to_radians();
|
||||
let lat = obs.latitude.to_radians();
|
||||
let lon = obs.longitude.to_radians();
|
||||
|
||||
// 估算当前时角（简化：假设太阳时 ≈ 平太阳时）
|
||||
let _t = (jd - 2451545.0) * 2.0 * std::f64::consts::PI; // 粗略近似
|
||||
let gmst = (280.46061837 + 360.98564736629 * (jd - 2451545.0)) % 360.0;
|
||||
let lst = (gmst + lon.to_degrees()) % 360.0;
|
||||
let hour_angle = (lst - ra_deg).to_radians();
|
||||
|
||||
// 高度角
|
||||
let sin_alt = dec.sin() * lat.sin() + dec.cos() * lat.cos() * hour_angle.cos();
|
||||
let altitude = sin_alt.asin().to_degrees();
|
||||
|
||||
// 方位角
|
||||
let cos_az = (dec.sin() - lat.sin() * sin_alt) / (lat.cos() * altitude.to_radians().cos());
|
||||
let azimuth = cos_atan2(cos_az, hour_angle.sin() * lat.cos()).to_degrees();
|
||||
|
||||
(altitude, azimuth)
|
||||
}
|
||||
|
||||
/// 简化的 atan2 范围限制
|
||||
fn cos_atan2(y: f64, x: f64) -> f64 {
|
||||
if x > 0.0 {
|
||||
(y / x).acos()
|
||||
} else if x < 0.0 {
|
||||
if y >= 0.0 {
|
||||
std::f64::consts::PI - (-y / -x).acos()
|
||||
} else {
|
||||
-std::f64::consts::PI + (y / -x).acos()
|
||||
}
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
}
|
||||
|
||||
/// 计算月相（简化）
|
||||
pub fn compute_moon_phase(jd: f64) -> (f64, String) {
|
||||
// 月相角计算（简化）
|
||||
let d = jd - 2451545.0;
|
||||
let age = d % 29.53058770576; // 朔望月周期
|
||||
let phase = age / 29.53058770576 * 360.0;
|
||||
|
||||
let description = if phase < 11.25 {
|
||||
"新月".to_string()
|
||||
} else if phase < 33.75 {
|
||||
"蛾眉月".to_string()
|
||||
} else if phase < 56.25 {
|
||||
"上弦月".to_string()
|
||||
} else if phase < 78.75 {
|
||||
"盈凸月".to_string()
|
||||
} else if phase < 101.25 {
|
||||
"满月".to_string()
|
||||
} else if phase < 123.75 {
|
||||
"亏凸月".to_string()
|
||||
} else if phase < 146.25 {
|
||||
"下弦月".to_string()
|
||||
} else if phase < 168.75 {
|
||||
"残月".to_string()
|
||||
} else {
|
||||
"新月".to_string()
|
||||
};
|
||||
|
||||
(phase, description)
|
||||
}
|
||||
|
||||
/// 月相对于观测的影响
|
||||
fn assess_moon_impact(phase_angle: f64, _target_v_mag: f64) -> MoonImpact {
|
||||
let impact = if !(30.0..=330.0).contains(&phase_angle) {
|
||||
"低 — 新月前后，天空暗".to_string()
|
||||
} else if !(90.0..=270.0).contains(&phase_angle) {
|
||||
"中 — 蛾眉月/残月，有一定散射光".to_string()
|
||||
} else if !(150.0..=210.0).contains(&phase_angle) {
|
||||
"高 — 盈凸/亏凸月，显著散射光".to_string()
|
||||
} else {
|
||||
"极高 — 满月前后，深空观测困难".to_string()
|
||||
};
|
||||
|
||||
MoonImpact {
|
||||
phase_angle,
|
||||
phase_description: String::new(), // 由调用者填充
|
||||
separation: 0.0, // 由调用者填充
|
||||
impact,
|
||||
}
|
||||
}
|
||||
|
||||
/// 计算大气质标 (Airmass)，使用 Pickering (1982) 割线修正公式
|
||||
pub fn calculate_airmass(altitude_deg: f64) -> f64 {
|
||||
if altitude_deg > 0.0 {
|
||||
let h = altitude_deg;
|
||||
let denom = (h + 244.0 / (165.0 + 47.0 * h.powf(1.1)))
|
||||
.to_radians()
|
||||
.sin();
|
||||
if denom > 1e-10 {
|
||||
1.0 / denom
|
||||
} else {
|
||||
f64::INFINITY
|
||||
}
|
||||
} else {
|
||||
f64::INFINITY // 不可观测或高度角为负
|
||||
}
|
||||
}
|
||||
|
||||
/// 估算曝光时间
|
||||
pub fn estimate_exposure_time(
|
||||
target_v_mag: f64,
|
||||
target_snr: f64,
|
||||
obs: &Observatory,
|
||||
altitude: f64,
|
||||
) -> ExposureEstimate {
|
||||
// 简化的曝光时间估算
|
||||
// 基于天光背景和恒星信号的信噪比模型
|
||||
|
||||
let airmass = calculate_airmass(altitude);
|
||||
|
||||
// V 波段天空亮度 (mag/arcsec²)
|
||||
let _sky_brightness = 21.5; // 晴朗无月夜晚
|
||||
// 大气消光
|
||||
let extinction = obs.extinction_v * airmass;
|
||||
|
||||
// 简化的 S/N 估算（基于 CCD 性能典型值）
|
||||
// S/N ≈ signal × t / sqrt(signal × t + sky × t + readnoise²)
|
||||
// 简化为：t ≈ (SNR/100)² × 10^(0.4 * (V - 20)) × 10^(0.4 * extinction)
|
||||
let base_time = 10.0; // 100 SNR 所需基本时间
|
||||
let v_factor = 10.0_f64.powf(0.4 * (target_v_mag - 20.0));
|
||||
let atm_factor = 10.0_f64.powf(0.4 * extinction);
|
||||
let snr_factor = (target_snr / 100.0).powi(2);
|
||||
|
||||
let exposure_secs = base_time * v_factor * atm_factor * snr_factor;
|
||||
let exposure_secs = exposure_secs.clamp(0.1, 3600.0); // 0.1s - 1h
|
||||
|
||||
// 帧数（含 dithering）
|
||||
let n_frames = if exposure_secs < 300.0 {
|
||||
1
|
||||
} else {
|
||||
(exposure_secs / 300.0).ceil() as u32
|
||||
};
|
||||
|
||||
let total_time = exposure_secs * n_frames as f64;
|
||||
|
||||
let notes = format!(
|
||||
"基于 V={:.1}, SNR={:.0}, airmass={:.2} 的简化估算。实际所需时间受天气、仪器等因素影响。",
|
||||
target_v_mag, target_snr, airmass
|
||||
);
|
||||
|
||||
ExposureEstimate {
|
||||
target_v_mag,
|
||||
target_snr,
|
||||
exposure_time_secs: exposure_secs,
|
||||
n_frames,
|
||||
total_time_secs: total_time,
|
||||
notes,
|
||||
}
|
||||
}
|
||||
|
||||
/// 完整观测可行性评估
|
||||
pub fn assess_observability(
|
||||
ra_deg: f64,
|
||||
dec_deg: f64,
|
||||
obs: &Observatory,
|
||||
target_v_mag: Option<f64>,
|
||||
target_snr: Option<f64>,
|
||||
) -> ObservabilityResult {
|
||||
let now_jd = 2451545.0; // 简化：使用 J2000.0 作为参考
|
||||
let (altitude, _azimuth) = compute_altitude(ra_deg, dec_deg, obs, now_jd);
|
||||
let airmass = calculate_airmass(altitude);
|
||||
|
||||
let (phase_angle, phase_desc) = compute_moon_phase(now_jd);
|
||||
let mut moon_impact = assess_moon_impact(phase_angle, target_v_mag.unwrap_or(15.0));
|
||||
moon_impact.phase_description = phase_desc;
|
||||
|
||||
let observable_tonight = altitude > 30.0;
|
||||
|
||||
let exposure_estimate = if let (Some(vmag), Some(snr)) = (target_v_mag, target_snr) {
|
||||
Some(estimate_exposure_time(vmag, snr, obs, altitude.max(0.0)))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
ObservabilityResult {
|
||||
rise_time: None, // 需要更精确的天文计算库
|
||||
set_time: None,
|
||||
max_altitude: altitude,
|
||||
max_altitude_time: None,
|
||||
observable_tonight,
|
||||
best_window: if observable_tonight {
|
||||
Some("当前高度角 > 30°，可进行观测".to_string())
|
||||
} else {
|
||||
Some("目标当前不可观测（高度角 < 30°）".to_string())
|
||||
},
|
||||
moon_impact,
|
||||
exposure_estimate,
|
||||
airmass_at_max: airmass,
|
||||
}
|
||||
}
|
||||
|
||||
/// 格式化观测可行性结果为 Markdown
|
||||
pub fn format_observability(result: &ObservabilityResult, obs_name: &str) -> String {
|
||||
let mut out = format!("## 观测可行性评估 ({})\n\n", obs_name);
|
||||
|
||||
out.push_str(&format!("- **最大高度角**: {:.1}°\n", result.max_altitude));
|
||||
out.push_str(&format!("- **大气质量**: {:.2}\n", result.airmass_at_max));
|
||||
out.push_str(&format!(
|
||||
"- **今晚可观测**: {}\n",
|
||||
if result.observable_tonight {
|
||||
"是"
|
||||
} else {
|
||||
"否"
|
||||
}
|
||||
));
|
||||
|
||||
if let Some(ref window) = result.best_window {
|
||||
out.push_str(&format!("- **最佳窗口**: {}\n", window));
|
||||
}
|
||||
|
||||
out.push_str("\n### 月相\n\n");
|
||||
out.push_str(&format!(
|
||||
"- 月相角: {:.0}° ({})\n",
|
||||
result.moon_impact.phase_angle, result.moon_impact.phase_description
|
||||
));
|
||||
out.push_str(&format!("- 影响: {}\n", result.moon_impact.impact));
|
||||
|
||||
if let Some(ref exp) = result.exposure_estimate {
|
||||
out.push_str("\n### 曝光时间估算\n\n");
|
||||
out.push_str(&format!("- 目标星等 (V): {:.1}\n", exp.target_v_mag));
|
||||
out.push_str(&format!("- 目标 SNR: {:.0}\n", exp.target_snr));
|
||||
out.push_str(&format!("- 单帧曝光: {:.1} s\n", exp.exposure_time_secs));
|
||||
out.push_str(&format!("- 帧数: {}\n", exp.n_frames));
|
||||
out.push_str(&format!("- 总时间: {:.0} s\n", exp.total_time_secs));
|
||||
out.push_str(&format!("- 备注: {}\n", exp.notes));
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_observatory_list() {
|
||||
let observatories = Observatory::default_list();
|
||||
assert!(!observatories.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_moon_phase() {
|
||||
let (phase, desc) = compute_moon_phase(2451545.0);
|
||||
assert!((0.0..=360.0).contains(&phase));
|
||||
assert!(!desc.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exposure_estimate() {
|
||||
let obs = Observatory::keck();
|
||||
let exp = estimate_exposure_time(15.0, 100.0, &obs, 60.0);
|
||||
assert!(exp.exposure_time_secs > 0.0);
|
||||
assert!(exp.exposure_time_secs < 3600.0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,363 @@
|
||||
// src/services/analysis/sed.rs
|
||||
//
|
||||
// 多波段 SED 拟合 —— 黑体辐射模型拟合，求 T_eff, 光度, 消光 A_V
|
||||
//
|
||||
// 输入：多波段测光数据（波长、通量、通量误差）
|
||||
// 输出：拟合参数（T_eff, R, A_V, L）+ 不确定度
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// SED 数据点
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SedDataPoint {
|
||||
/// 中心波长 (Å)
|
||||
pub wavelength: f64,
|
||||
/// 流量 (erg/s/cm²/Å 或 Jy)
|
||||
pub flux: f64,
|
||||
/// 流量误差
|
||||
pub flux_error: f64,
|
||||
/// 波段名（如 "J", "W1", "G"）
|
||||
pub band: String,
|
||||
/// 数据源
|
||||
pub source: String,
|
||||
}
|
||||
|
||||
/// SED 拟合结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SedFitResult {
|
||||
/// 有效温度 (K)
|
||||
pub teff: f64,
|
||||
/// 半径 (R_sun)
|
||||
pub radius: f64,
|
||||
/// 消光 A_V (mag)
|
||||
pub a_v: f64,
|
||||
/// 光度 (L_sun)
|
||||
pub luminosity: f64,
|
||||
/// 距离 (pc)，用于计算光度
|
||||
pub distance_pc: Option<f64>,
|
||||
/// 拟合残差（chi2）
|
||||
pub chi2: f64,
|
||||
/// 约化 chi2
|
||||
pub reduced_chi2: f64,
|
||||
/// 参数不确定度（1σ）
|
||||
pub uncertainties: SedUncertainties,
|
||||
/// 拟合曲线点 (wavelength, flux)，供前端绘图
|
||||
pub model_curve: Vec<(f64, f64)>,
|
||||
/// 残差点
|
||||
pub residuals: Vec<(f64, f64)>,
|
||||
}
|
||||
|
||||
/// 拟合参数不确定度
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SedUncertainties {
|
||||
pub teff_err: f64,
|
||||
pub radius_err: f64,
|
||||
pub a_v_err: f64,
|
||||
pub luminosity_err: f64,
|
||||
}
|
||||
|
||||
/// 简化星际消光曲线（Ccm87 近似）
|
||||
fn extinction_law(wavelength_angstrom: f64, r_v: f64) -> f64 {
|
||||
let x = 1e4 / wavelength_angstrom; // 1/λ (μm⁻¹)
|
||||
|
||||
// CCM87 近似（光学-近红外）
|
||||
if x < 0.3 {
|
||||
// 远红外：A_λ ∝ λ⁻¹.75
|
||||
0.265 * r_v * x.powf(1.75) + 0.861
|
||||
} else if x < 1.1 {
|
||||
// 红外
|
||||
0.574 * x.powf(1.61) - 0.527 * x.powf(3.0) / r_v
|
||||
} else if x < 3.3 {
|
||||
// 光学
|
||||
let y = x - 1.82;
|
||||
let a = 1.0 + 0.17699 * y - 0.50447 * y * y - 0.02427 * y.powi(3)
|
||||
+ 0.72085 * y.powi(4)
|
||||
+ 0.01979 * y.powi(5)
|
||||
- 0.77530 * y.powi(6)
|
||||
+ 0.32999 * y.powi(7);
|
||||
let b = 1.41338 * y + 2.28305 * y * y + 1.07233 * y.powi(3)
|
||||
- 5.38434 * y.powi(4)
|
||||
- 0.62251 * y.powi(5)
|
||||
+ 5.30260 * y.powi(6)
|
||||
- 2.09002 * y.powi(7);
|
||||
a + b / r_v
|
||||
} else {
|
||||
// 紫外
|
||||
let a = 1.752 - 0.316 * x - 0.104 / ((x - 4.67).powi(2) + 0.341) + 0.293;
|
||||
let b = -3.090 + 1.825 * x + 1.206 / ((x - 4.62).powi(2) + 0.263);
|
||||
a + b / r_v
|
||||
}
|
||||
}
|
||||
|
||||
/// 黑体辐射模型（普朗克函数）
|
||||
fn black_body_flux(wavelength_angstrom: f64, teff: f64) -> f64 {
|
||||
let wl = wavelength_angstrom * 1e-10; // Å → m
|
||||
let h = 6.626e-34;
|
||||
let c = 3e8;
|
||||
let k = 1.381e-23;
|
||||
|
||||
let x = h * c / (wl * k * teff);
|
||||
if x > 500.0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let b = 2.0 * h * c * c / wl.powi(5) / (x.exp() - 1.0);
|
||||
b * 1e-26 // 转换为 Jy 量级（任意缩放）
|
||||
}
|
||||
|
||||
/// 黑体 + 消光模型：F_model(λ) = R² × B_λ(T) × 10^(-0.4 × A_λ × A_V / A_V)
|
||||
fn sed_model(wavelength: f64, teff: f64, radius_solar: f64, a_v: f64) -> f64 {
|
||||
let r_v = 3.1;
|
||||
let bb = black_body_flux(wavelength, teff);
|
||||
let ext = extinction_law(wavelength, r_v);
|
||||
let a_lambda = a_v * ext / extinction_law(5500.0, r_v); // A_λ = A_V × A(λ)/A(V)
|
||||
let extinction_factor = 10.0_f64.powf(-0.4 * a_lambda);
|
||||
|
||||
radius_solar.powi(2) * bb * extinction_factor
|
||||
}
|
||||
|
||||
/// 最小二乘拟合 SED
|
||||
pub fn fit_sed(data: &[SedDataPoint], distance_pc: Option<f64>) -> Result<SedFitResult> {
|
||||
if data.len() < 3 {
|
||||
return Err(anyhow!("至少需要 3 个测光点进行 SED 拟合"));
|
||||
}
|
||||
|
||||
// 网格搜索（粗搜索 + 精搜索）
|
||||
let teff_range = (2000.0, 50000.0);
|
||||
let radius_range = (0.01, 1000.0);
|
||||
let av_range = (0.0, 10.0);
|
||||
|
||||
// 粗搜索
|
||||
let n_grid = 30;
|
||||
let mut best_chi2 = f64::INFINITY;
|
||||
let mut best_params = (5778.0, 1.0, 0.0);
|
||||
|
||||
for i in 0..n_grid {
|
||||
for j in 0..n_grid {
|
||||
for k in 0..10 {
|
||||
let teff =
|
||||
teff_range.0 + (teff_range.1 - teff_range.0) * i as f64 / (n_grid - 1) as f64;
|
||||
let radius = radius_range.0
|
||||
+ (radius_range.1 - radius_range.0) * j as f64 / (n_grid - 1) as f64;
|
||||
let av = av_range.0 + (av_range.1 - av_range.0) * k as f64 / 9.0;
|
||||
|
||||
let chi2 = compute_chi2(data, teff, radius, av);
|
||||
if chi2 < best_chi2 {
|
||||
best_chi2 = chi2;
|
||||
best_params = (teff, radius, av);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 精搜索（在最优参数附近 ±10% 范围细化）
|
||||
let (ct, cr, ca) = best_params;
|
||||
let fine_n = 20;
|
||||
for i in 0..fine_n {
|
||||
for j in 0..fine_n {
|
||||
for k in 0..fine_n {
|
||||
let teff = ct * 0.9 + ct * 0.2 * i as f64 / (fine_n - 1) as f64;
|
||||
let radius = cr * 0.9 + cr * 0.2 * j as f64 / (fine_n - 1) as f64;
|
||||
let av = (ca - 1.0).max(0.0) + 2.0 * k as f64 / (fine_n - 1) as f64;
|
||||
|
||||
let chi2 = compute_chi2(data, teff, radius, av);
|
||||
if chi2 < best_chi2 {
|
||||
best_chi2 = chi2;
|
||||
best_params = (teff, radius, av);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let (teff, radius, a_v) = best_params;
|
||||
let dof = data.len() as f64 - 3.0; // 3 个自由参数
|
||||
let reduced_chi2 = if dof > 0.0 {
|
||||
best_chi2 / dof
|
||||
} else {
|
||||
best_chi2
|
||||
};
|
||||
|
||||
// 光度计算
|
||||
let luminosity = radius.powi(2) * (teff / 5778.0).powi(4);
|
||||
|
||||
// 生成模型曲线
|
||||
let model_curve: Vec<(f64, f64)> = (100..100000)
|
||||
.step_by(100)
|
||||
.map(|w| {
|
||||
let wl = w as f64;
|
||||
(wl, sed_model(wl, teff, radius, a_v))
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 残差
|
||||
let residuals: Vec<(f64, f64)> = data
|
||||
.iter()
|
||||
.map(|d| {
|
||||
let model = sed_model(d.wavelength, teff, radius, a_v);
|
||||
let residual = if d.flux > 0.0 && model > 0.0 {
|
||||
(d.flux / model - 1.0) * 100.0 // 百分比残差
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
(d.wavelength, residual)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 不确定度估计（简化：通过扰动参数估计）
|
||||
let uncertainties = estimate_uncertainties(data, teff, radius, a_v, best_chi2, dof);
|
||||
|
||||
Ok(SedFitResult {
|
||||
teff,
|
||||
radius,
|
||||
a_v,
|
||||
luminosity,
|
||||
distance_pc,
|
||||
chi2: best_chi2,
|
||||
reduced_chi2,
|
||||
uncertainties,
|
||||
model_curve,
|
||||
residuals,
|
||||
})
|
||||
}
|
||||
|
||||
/// 计算 chi2（模型和数据都归一化到 V 波段 5500Å）
|
||||
fn compute_chi2(data: &[SedDataPoint], teff: f64, radius: f64, a_v: f64) -> f64 {
|
||||
let v_model = sed_model(5500.0, teff, radius, a_v);
|
||||
if v_model <= 0.0 {
|
||||
return f64::INFINITY;
|
||||
}
|
||||
data.iter()
|
||||
.map(|d| {
|
||||
let model = sed_model(d.wavelength, teff, radius, a_v) / v_model;
|
||||
if d.flux_error > 0.0 {
|
||||
((d.flux - model) / d.flux_error).powi(2)
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// 参数不确定度估计（通过扰动法）
|
||||
fn estimate_uncertainties(
|
||||
data: &[SedDataPoint],
|
||||
teff: f64,
|
||||
radius: f64,
|
||||
a_v: f64,
|
||||
best_chi2: f64,
|
||||
_dof: f64,
|
||||
) -> SedUncertainties {
|
||||
// 扰动每个参数 10%，观察 chi2 变化
|
||||
let delta = 0.1;
|
||||
|
||||
let chi2_t = compute_chi2(data, teff * (1.0 + delta), radius, a_v);
|
||||
let teff_err = if chi2_t > best_chi2 {
|
||||
teff * delta * ((chi2_t - best_chi2).max(0.0)).sqrt()
|
||||
} else {
|
||||
teff * 0.1
|
||||
};
|
||||
|
||||
let chi2_r = compute_chi2(data, teff, radius * (1.0 + delta), a_v);
|
||||
let radius_err = if chi2_r > best_chi2 {
|
||||
radius * delta * ((chi2_r - best_chi2).max(0.0)).sqrt()
|
||||
} else {
|
||||
radius * 0.1
|
||||
};
|
||||
|
||||
let chi2_a = compute_chi2(data, teff, radius, (a_v + 0.1).max(0.0));
|
||||
let a_v_err = if chi2_a > best_chi2 {
|
||||
0.1 * ((chi2_a - best_chi2).max(0.0)).sqrt()
|
||||
} else {
|
||||
0.5
|
||||
};
|
||||
|
||||
let luminosity_err = (radius_err / radius).hypot(teff_err / teff) * 2.0;
|
||||
|
||||
SedUncertainties {
|
||||
teff_err,
|
||||
radius_err,
|
||||
a_v_err,
|
||||
luminosity_err,
|
||||
}
|
||||
}
|
||||
|
||||
/// 格式化 SED 拟合结果为 Markdown
|
||||
pub fn format_sed_result(result: &SedFitResult) -> String {
|
||||
let mut out = String::from("## SED 拟合结果\n\n");
|
||||
out.push_str("| 参数 | 值 | 不确定度 |\n|------|-----|----------|\n");
|
||||
out.push_str(&format!(
|
||||
"| T_eff | {:.0} K | ±{:.0} K |\n",
|
||||
result.teff, result.uncertainties.teff_err
|
||||
));
|
||||
out.push_str(&format!(
|
||||
"| R | {:.4} R_sun | ±{:.4} |\n",
|
||||
result.radius, result.uncertainties.radius_err
|
||||
));
|
||||
out.push_str(&format!(
|
||||
"| A_V | {:.3} mag | ±{:.3} |\n",
|
||||
result.a_v, result.uncertainties.a_v_err
|
||||
));
|
||||
out.push_str(&format!(
|
||||
"| L | {:.4} L_sun | ±{:.1}% |\n",
|
||||
result.luminosity,
|
||||
result.uncertainties.luminosity_err * 100.0
|
||||
));
|
||||
out.push_str(&format!("| χ² | {:.2} | |\n", result.chi2));
|
||||
out.push_str(&format!("| χ²_red | {:.2} | |\n", result.reduced_chi2));
|
||||
if let Some(d) = result.distance_pc {
|
||||
out.push_str(&format!("| 距离 | {:.1} pc | |\n", d));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_black_body_positive() {
|
||||
let bb = black_body_flux(5500.0, 5778.0);
|
||||
assert!(bb > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sed_model_increases_with_radius() {
|
||||
let f1 = sed_model(5500.0, 5778.0, 1.0, 0.0);
|
||||
let f2 = sed_model(5500.0, 5778.0, 2.0, 0.0);
|
||||
assert!(f2 > f1, "更大的半径应该产生更大的流量");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fit_sed_solar() {
|
||||
// 构造太阳型恒星的模拟测光数据（从模型本身生成，加小噪声）
|
||||
let v_flux = sed_model(5500.0, 5778.0, 1.0, 0.0);
|
||||
let data_normalized: Vec<SedDataPoint> = vec![
|
||||
(4400.0, "B"),
|
||||
(5500.0, "V"),
|
||||
(6500.0, "R"),
|
||||
(8000.0, "I"),
|
||||
(12500.0, "J"),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|(w, band)| {
|
||||
let model = sed_model(w, 5778.0, 1.0, 0.0);
|
||||
SedDataPoint {
|
||||
wavelength: w,
|
||||
flux: model / v_flux,
|
||||
flux_error: 0.01,
|
||||
band: band.to_string(),
|
||||
source: "test".to_string(),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let result = fit_sed(&data_normalized, Some(10.0)).unwrap();
|
||||
// T_eff 应接近太阳（允许 40% 误差，因为网格搜索粗糙）
|
||||
assert!(
|
||||
(result.teff - 5778.0).abs() / 5778.0 < 0.4,
|
||||
"T_eff = {} K, 应接近 5778 K",
|
||||
result.teff
|
||||
);
|
||||
}
|
||||
}
|
||||
+160
-1
@@ -317,19 +317,21 @@ 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') AS last_run FROM sync_queries ORDER BY last_run DESC")
|
||||
let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') AS last_run, is_scheduled FROM sync_queries ORDER BY last_run DESC")
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
let mut list = Vec::new();
|
||||
for r in rows {
|
||||
use sqlx::Row;
|
||||
let is_scheduled_int: i32 = r.get("is_scheduled");
|
||||
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"),
|
||||
is_scheduled: is_scheduled_int != 0,
|
||||
});
|
||||
}
|
||||
Ok(list)
|
||||
@@ -343,6 +345,162 @@ impl MetaSync {
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
// 设置是否定时同步
|
||||
pub async fn set_query_schedule(
|
||||
db: &SqlitePool,
|
||||
id: i64,
|
||||
is_scheduled: bool,
|
||||
) -> Result<bool, sqlx::Error> {
|
||||
let result = sqlx::query("UPDATE sync_queries SET is_scheduled = ? WHERE id = ?")
|
||||
.bind(if is_scheduled { 1 } else { 0 })
|
||||
.bind(id)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
/// 运行静默（无状态锁）增量同步，返回实际入库/同步成功的文献数量
|
||||
pub async fn run_incremental_sync_silent(
|
||||
db: SqlitePool,
|
||||
ads: Arc<AdsClient>,
|
||||
arxiv: Arc<ArxivClient>,
|
||||
query: String,
|
||||
source: String,
|
||||
limit: i32,
|
||||
last_run_str: &str,
|
||||
) -> anyhow::Result<i32> {
|
||||
let date_part = if last_run_str.len() >= 10 {
|
||||
&last_run_str[0..10]
|
||||
} else {
|
||||
"now-7D"
|
||||
};
|
||||
|
||||
let mut synced_count = 0;
|
||||
|
||||
// 1. 同步 ADS 数据
|
||||
if source == "all" || source == "ads" {
|
||||
let ads_query = if date_part == "now-7D" {
|
||||
format!("{} entdate:[now-7D TO now]", query)
|
||||
} else {
|
||||
format!("{} entdate:[{} TO now]", query, date_part)
|
||||
};
|
||||
|
||||
let ads_limit = if limit > 0 { limit } else { 100 };
|
||||
match ads.search(&ads_query, 0, ads_limit, "date_desc").await {
|
||||
Ok(docs) => {
|
||||
if let Ok(mut tx) = db.begin().await {
|
||||
for doc in docs {
|
||||
let paper = convert_ads_doc_to_standard(&doc);
|
||||
if save_paper_to_db_tx(&mut tx, &paper).await.is_ok() {
|
||||
synced_count += 1;
|
||||
}
|
||||
}
|
||||
let _ = tx.commit().await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"[Scheduler] 增量同步 ADS 数据失败 (query={}): {}",
|
||||
ads_query, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 同步 arXiv 数据
|
||||
if source == "all" || source == "arxiv" {
|
||||
let arxiv_limit = if limit > 0 { limit } else { 100 };
|
||||
match arxiv.search(&query, 0, arxiv_limit, "date_desc").await {
|
||||
Ok(papers) => {
|
||||
if let Ok(mut tx) = db.begin().await {
|
||||
for p in papers {
|
||||
let paper = convert_arxiv_to_standard(&p);
|
||||
if save_paper_to_db_tx(&mut tx, &paper).await.is_ok() {
|
||||
synced_count += 1;
|
||||
}
|
||||
}
|
||||
let _ = tx.commit().await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"[Scheduler] 增量同步 arXiv 数据失败 (query={}): {}",
|
||||
query, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(synced_count)
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动定时增量文献同步调度器，每小时轮询一次
|
||||
pub async fn start_scheduled_sync_timer(state: Arc<crate::api::AppState>) {
|
||||
info!("[Scheduler] 启动定时增量文献同步调度器...");
|
||||
// 启动 10 秒后执行第一次同步
|
||||
tokio::time::sleep(std::time::Duration::from_secs(10)).await;
|
||||
loop {
|
||||
info!("[Scheduler] 开始执行定时增量同步步进...");
|
||||
if let Err(e) = run_scheduled_sync_step(&state).await {
|
||||
error!("[Scheduler] 定时增量文献同步执行失败: {}", e);
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行单步增量同步
|
||||
pub async fn run_scheduled_sync_step(state: &crate::api::AppState) -> anyhow::Result<()> {
|
||||
let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') AS last_run FROM sync_queries WHERE is_scheduled = 1")
|
||||
.fetch_all(&state.db)
|
||||
.await?;
|
||||
|
||||
for r in rows {
|
||||
use sqlx::Row;
|
||||
let id: i64 = r.get("id");
|
||||
let query: String = r.get("query");
|
||||
let source: String = r.get("source");
|
||||
let limit: i32 = r.get("limit_count");
|
||||
let last_run_str: String = r.get("last_run");
|
||||
|
||||
info!(
|
||||
"[Scheduler] 开始增量同步任务: id={}, query='{}', last_run={}",
|
||||
id, query, last_run_str
|
||||
);
|
||||
|
||||
// 运行静默增量同步
|
||||
match MetaSync::run_incremental_sync_silent(
|
||||
state.db.clone(),
|
||||
Arc::new(state.sources.ads.clone()),
|
||||
Arc::new(state.sources.arxiv.clone()),
|
||||
query.clone(),
|
||||
source.clone(),
|
||||
limit,
|
||||
&last_run_str,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(count) => {
|
||||
info!(
|
||||
"[Scheduler] 增量同步任务成功完成: id={}, 新同步 {} 篇文献",
|
||||
id, count
|
||||
);
|
||||
// 更新该任务的 last_run 时间戳
|
||||
let _ = sqlx::query(
|
||||
"UPDATE sync_queries SET last_run = CURRENT_TIMESTAMP WHERE id = ?",
|
||||
)
|
||||
.bind(id)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
Err(e) => {
|
||||
error!("[Scheduler] 增量同步任务执行失败: id={}, error={}", id, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, serde::Deserialize)]
|
||||
@@ -352,4 +510,5 @@ pub struct SavedSyncQuery {
|
||||
pub source: String,
|
||||
pub limit_count: i32,
|
||||
pub last_run: String,
|
||||
pub is_scheduled: bool,
|
||||
}
|
||||
|
||||
@@ -3,4 +3,6 @@ pub mod asset;
|
||||
pub mod meta;
|
||||
|
||||
pub use asset::{filter_papers_for_batch, AssetBatch, AssetBatchStatus, BatchAction};
|
||||
pub use meta::{MetaSync, MetaSyncStatus, SavedSyncQuery};
|
||||
pub use meta::{
|
||||
run_scheduled_sync_step, start_scheduled_sync_timer, MetaSync, MetaSyncStatus, SavedSyncQuery,
|
||||
};
|
||||
|
||||
@@ -184,6 +184,7 @@ pub fn escape_string_literal(s: &str) -> String {
|
||||
/// 锥形检索
|
||||
///
|
||||
/// strategy="nearest" 时按角距离排序，返回最近的源；strategy="all" 时无序返回全部。
|
||||
#[allow(clippy::too_many_arguments)] // 锥形检索物理参数
|
||||
pub async fn cone_search(
|
||||
pool: &SqlitePool,
|
||||
client: &VizierClient,
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod analysis;
|
||||
pub mod batch;
|
||||
pub mod cds;
|
||||
pub mod chunker;
|
||||
@@ -14,5 +15,7 @@ pub mod query_parser;
|
||||
pub mod rag;
|
||||
pub mod search;
|
||||
pub mod session;
|
||||
pub mod spectrum;
|
||||
pub mod timeseries;
|
||||
pub mod translation;
|
||||
pub mod vision;
|
||||
|
||||
@@ -103,6 +103,7 @@ pub async fn fetch_observation_cache(
|
||||
}
|
||||
|
||||
/// 写入/更新缓存条目
|
||||
#[allow(clippy::too_many_arguments)] // 缓存条目字段一一对应
|
||||
pub async fn write_observation_cache(
|
||||
pool: &SqlitePool,
|
||||
source: Source,
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
// src/services/observation/cutout.rs
|
||||
//
|
||||
// 图像 Cutout 服务 —— SDSS / DSS / Pan-STARRS cutout API
|
||||
//
|
||||
// 支持：
|
||||
// - SDSS SkyServer cutout（光学 2.5m）
|
||||
// - STScI DSS（Digitized Sky Survey，多种底片）
|
||||
// - Pan-STARRS cutout（光学 1.8m）
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::info;
|
||||
|
||||
/// Cutout 数据源
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub enum CutoutSource {
|
||||
/// SDSS optical
|
||||
Sdss,
|
||||
/// DSS (POSS-I, POSS-II, etc.)
|
||||
Dss,
|
||||
/// Pan-STARRS DR1
|
||||
Panstarrs,
|
||||
}
|
||||
|
||||
impl CutoutSource {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Sdss => "sdss",
|
||||
Self::Dss => "dss",
|
||||
Self::Panstarrs => "panstarrs",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cutout 结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CutoutResult {
|
||||
/// 图像数据（JPEG/PNG 字节）
|
||||
pub image_bytes: Vec<u8>,
|
||||
/// 图像格式
|
||||
pub format: String,
|
||||
/// 中心坐标
|
||||
pub ra: f64,
|
||||
pub dec: f64,
|
||||
/// cutout 大小（角秒）
|
||||
pub size_arcsec: f64,
|
||||
/// 数据源
|
||||
pub source: String,
|
||||
/// 图像宽度（像素）
|
||||
pub width: u32,
|
||||
/// 图像高度（像素）
|
||||
pub height: u32,
|
||||
}
|
||||
|
||||
/// SDSS cutout（SkyServer）
|
||||
pub async fn sdss_cutout(ra: f64, dec: f64, size_arcsec: f64) -> Result<CutoutResult> {
|
||||
let scale = size_arcsec / 14.0; // 14 arcsec/pixel for SDSS
|
||||
let width = (size_arcsec / scale) as u32;
|
||||
let url = format!(
|
||||
"https://skyserver.sdss.org/dr18/SkyServerWS/ImgCutout/getjpeg?ra={}&dec={}&scale={:.4}&width={}&height={}",
|
||||
ra, dec, scale, width, width
|
||||
);
|
||||
info!(
|
||||
"[Cutout] SDSS: ra={:.4}, dec={:.4}, size={}\"",
|
||||
ra, dec, size_arcsec
|
||||
);
|
||||
|
||||
let resp = reqwest::get(&url)
|
||||
.await
|
||||
.map_err(|e| anyhow!("SDSS cutout 请求失败: {}", e))?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow!("SDSS cutout 失败: {}", resp.status()));
|
||||
}
|
||||
let bytes = resp.bytes().await?.to_vec();
|
||||
|
||||
Ok(CutoutResult {
|
||||
image_bytes: bytes,
|
||||
format: "jpeg".to_string(),
|
||||
ra,
|
||||
dec,
|
||||
size_arcsec,
|
||||
source: "SDSS".to_string(),
|
||||
width,
|
||||
height: width,
|
||||
})
|
||||
}
|
||||
|
||||
/// DSS cutout（STScI）
|
||||
pub async fn dss_cutout(ra: f64, dec: f64, size_arcsec: f64) -> Result<CutoutResult> {
|
||||
let size_deg = size_arcsec / 3600.0;
|
||||
let url = format!(
|
||||
"https://archive.stsci.edu/cgi-bin/dss_search?faste&ra={}&dec={}&equinox=J2000&epoch=2000&width={}&height={}&nimages=1&format=jpeg",
|
||||
ra, dec, size_deg, size_deg
|
||||
);
|
||||
info!(
|
||||
"[Cutout] DSS: ra={:.4}, dec={:.4}, size={}\"",
|
||||
ra, dec, size_arcsec
|
||||
);
|
||||
|
||||
let resp = reqwest::get(&url)
|
||||
.await
|
||||
.map_err(|e| anyhow!("DSS cutout 请求失败: {}", e))?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow!("DSS cutout 失败: {}", resp.status()));
|
||||
}
|
||||
let bytes = resp.bytes().await?.to_vec();
|
||||
|
||||
let pixel_size = size_arcsec / 1.0; // ~1 arcsec/pixel for DSS
|
||||
let width = (size_arcsec / pixel_size) as u32;
|
||||
|
||||
Ok(CutoutResult {
|
||||
image_bytes: bytes,
|
||||
format: "jpeg".to_string(),
|
||||
ra,
|
||||
dec,
|
||||
size_arcsec,
|
||||
source: "DSS".to_string(),
|
||||
width,
|
||||
height: width,
|
||||
})
|
||||
}
|
||||
|
||||
/// Pan-STARRS cutout（via VizieR / PS1 cutout service）
|
||||
pub async fn panstarrs_cutout(ra: f64, dec: f64, size_arcsec: f64) -> Result<CutoutResult> {
|
||||
let size_deg = size_arcsec / 3600.0;
|
||||
// PS1 cutout service
|
||||
let url = format!(
|
||||
"https://ps1images.stsci.edu/cgi-bin/ps1cutoutservice?ra={}&dec={}&size={}&format=fits&filters=g,r,i,z,y",
|
||||
ra, dec, size_deg
|
||||
);
|
||||
info!(
|
||||
"[Cutout] Pan-STARRS: ra={:.4}, dec={:.4}, size={}\"",
|
||||
ra, dec, size_arcsec
|
||||
);
|
||||
|
||||
let resp = reqwest::get(&url)
|
||||
.await
|
||||
.map_err(|e| anyhow!("PS1 cutout 请求失败: {}", e))?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow!("PS1 cutout 失败: {}", resp.status()));
|
||||
}
|
||||
let bytes = resp.bytes().await?.to_vec();
|
||||
|
||||
Ok(CutoutResult {
|
||||
image_bytes: bytes,
|
||||
format: "fits".to_string(),
|
||||
ra,
|
||||
dec,
|
||||
size_arcsec,
|
||||
source: "Pan-STARRS".to_string(),
|
||||
width: 0,
|
||||
height: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// 统一 cutout 接口
|
||||
pub async fn get_cutout(
|
||||
source: CutoutSource,
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
size_arcsec: f64,
|
||||
) -> Result<CutoutResult> {
|
||||
match source {
|
||||
CutoutSource::Sdss => sdss_cutout(ra, dec, size_arcsec).await,
|
||||
CutoutSource::Dss => dss_cutout(ra, dec, size_arcsec).await,
|
||||
CutoutSource::Panstarrs => panstarrs_cutout(ra, dec, size_arcsec).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Finding Chart 生成（基于 cutout + 标注）
|
||||
pub async fn generate_finding_chart(
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
size_arcsec: f64,
|
||||
source: CutoutSource,
|
||||
_label: Option<&str>,
|
||||
) -> Result<CutoutResult> {
|
||||
let result = get_cutout(source, ra, dec, size_arcsec).await?;
|
||||
|
||||
// 保存到文件
|
||||
let ext = match result.format.as_str() {
|
||||
"jpeg" => "jpg",
|
||||
"fits" => "fits",
|
||||
_ => "png",
|
||||
};
|
||||
let file_name = format!("finding_chart_{:.4}_{:.4}.{}", ra, dec, ext);
|
||||
let file_path = format!("finding_charts/{}", file_name);
|
||||
|
||||
let library_dir = std::env::var("LIBRARY_DIR").unwrap_or_else(|_| "library".to_string());
|
||||
let full_path = std::path::Path::new(&library_dir).join(&file_path);
|
||||
if let Some(parent) = full_path.parent() {
|
||||
tokio::fs::create_dir_all(parent).await.ok();
|
||||
}
|
||||
tokio::fs::write(&full_path, &result.image_bytes).await.ok();
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// CutoutFetcher —— ObservationFetcher trait 实现
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::services::observation::cache::{
|
||||
cached_files_total_size, fetch_observation_cache, log_write_failure, persist_bytes,
|
||||
write_observation_cache,
|
||||
};
|
||||
use crate::services::observation::fetcher::{Candidate, ObservationFetcher};
|
||||
use crate::services::observation::types::{
|
||||
Artifact, ObservationProduct, ProductSpec, ProductType, Source,
|
||||
};
|
||||
use tracing::warn;
|
||||
|
||||
/// 图像 Cutout fetcher（SDSS/DSS/Pan-STARRS）
|
||||
///
|
||||
/// 这不是一个标准的 ObservationFetcher（cutout 不通过 TAP/VO 检索），
|
||||
/// 而是通过 HTTP API 直接获取图像 cutout。
|
||||
/// 注册为 (Source::Panstarrs, ProductType::Image) 的 fetcher，
|
||||
/// 坐标模式直接构造 cutout URL。
|
||||
#[derive(Debug)]
|
||||
pub struct CutoutFetcher;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl ObservationFetcher for CutoutFetcher {
|
||||
fn key(&self) -> (Source, ProductType) {
|
||||
(Source::Panstarrs, ProductType::Image)
|
||||
}
|
||||
|
||||
fn subtypes(&self) -> &'static [&'static str] {
|
||||
&["sdss", "dss", "panstarrs"]
|
||||
}
|
||||
|
||||
fn suggested_max_radius_deg(&self) -> f64 {
|
||||
0.0
|
||||
}
|
||||
|
||||
fn hard_max_radius_deg(&self) -> f64 {
|
||||
0.0
|
||||
}
|
||||
|
||||
fn supports_coordinates(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn supports_identifiers(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
async fn cone_search_raw(
|
||||
&self,
|
||||
_state: &AppState,
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
_radius_deg: f64,
|
||||
_release: Option<&str>,
|
||||
_subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
) -> Result<Vec<Candidate>> {
|
||||
// Cutout 不做 cone search，直接返回目标本身
|
||||
Ok(vec![Candidate {
|
||||
source: Source::Panstarrs,
|
||||
source_id: format!("cutout_{:.4}_{:.4}", ra, dec),
|
||||
label: format!("Cutout at ({:.4}, {:.4})", ra, dec),
|
||||
ra: Some(ra),
|
||||
dec: Some(dec),
|
||||
distance: Some(0.0),
|
||||
raw: None,
|
||||
}])
|
||||
}
|
||||
|
||||
async fn resolve_identifier(
|
||||
&self,
|
||||
_identifier: &str,
|
||||
_release: Option<&str>,
|
||||
_subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
) -> Result<Candidate> {
|
||||
Err(anyhow!("Cutout 不支持标识符模式"))
|
||||
}
|
||||
|
||||
async fn fetch(
|
||||
&self,
|
||||
state: &AppState,
|
||||
candidate: &Candidate,
|
||||
_release: Option<&str>,
|
||||
subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
force: bool,
|
||||
) -> Result<ObservationProduct> {
|
||||
let ra = candidate.ra.unwrap_or(0.0);
|
||||
let dec = candidate.dec.unwrap_or(0.0);
|
||||
let sub = subtype.unwrap_or("sdss");
|
||||
let product = ProductSpec::with_subtype(ProductType::Image, sub);
|
||||
let cache_key = format!("cutout_{:.4}_{:.4}_{}", ra, dec, sub);
|
||||
let source_label = candidate.label.clone();
|
||||
|
||||
// 1) 检查缓存
|
||||
if !force {
|
||||
if let Some((artifacts, meta)) =
|
||||
fetch_observation_cache(&state.db, Source::Panstarrs, &product, &cache_key).await?
|
||||
{
|
||||
if cached_files_total_size(&state.config.storage.library_dir, &artifacts)
|
||||
.await
|
||||
.is_some()
|
||||
{
|
||||
info!("[Cutout] 缓存命中 (key={})", cache_key);
|
||||
return Ok(ObservationProduct {
|
||||
source: Source::Panstarrs,
|
||||
product,
|
||||
source_id: cache_key,
|
||||
source_label,
|
||||
artifacts,
|
||||
source_meta: meta,
|
||||
});
|
||||
}
|
||||
warn!("[Cutout] 缓存文件缺失，重新下载 (key={})", cache_key);
|
||||
}
|
||||
}
|
||||
|
||||
// 2) 真实获取
|
||||
let source = match sub {
|
||||
"dss" => CutoutSource::Dss,
|
||||
"panstarrs" => CutoutSource::Panstarrs,
|
||||
_ => CutoutSource::Sdss,
|
||||
};
|
||||
|
||||
let result = get_cutout(source, ra, dec, 120.0).await?; // 2 arcmin default
|
||||
|
||||
let ext = match result.format.as_str() {
|
||||
"jpeg" => "jpg",
|
||||
"fits" => "fits",
|
||||
_ => "png",
|
||||
};
|
||||
let file_name = format!("cutout_{:.4}_{:.4}.{}", ra, dec, ext);
|
||||
let file_path = format!("Telescope/cutout/{}/{}", candidate.source_id, file_name);
|
||||
persist_bytes(
|
||||
&state.config.storage.library_dir,
|
||||
&file_path,
|
||||
&result.image_bytes,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let artifacts = vec![Artifact {
|
||||
band: None,
|
||||
original_name: Some(file_name),
|
||||
file_path: file_path.clone(),
|
||||
file_url: crate::services::observation::cache::file_url_from_path(&file_path),
|
||||
file_format: ext.to_string(),
|
||||
size_bytes: result.image_bytes.len(),
|
||||
cached: false,
|
||||
}];
|
||||
|
||||
// 3) 写入缓存
|
||||
if let Err(e) = write_observation_cache(
|
||||
&state.db,
|
||||
Source::Panstarrs,
|
||||
&product,
|
||||
&cache_key,
|
||||
Some(ra),
|
||||
Some(dec),
|
||||
&artifacts,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
log_write_failure("panstarrs", "image", e);
|
||||
}
|
||||
|
||||
Ok(ObservationProduct {
|
||||
source: Source::Panstarrs,
|
||||
product,
|
||||
source_id: cache_key,
|
||||
source_label,
|
||||
artifacts,
|
||||
source_meta: candidate.raw.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -21,6 +21,7 @@ use tracing::{info, warn};
|
||||
/// 这一步，把命中的候选源列表返回给调用方预览，由调用方决定是否下载、下载哪些。
|
||||
///
|
||||
/// 标识符模式无需"检索"——标识符本身就是确定源，调用方直接走 download 即可。
|
||||
#[allow(clippy::too_many_arguments)] // 透传 (state, registry, source, product, 版本) 的管线参数
|
||||
pub async fn search_observation(
|
||||
state: &AppState,
|
||||
registry: &ObservationRegistry,
|
||||
@@ -100,6 +101,7 @@ pub async fn download_observation(
|
||||
}
|
||||
|
||||
/// 按坐标：cone 检索 → 选源 → 逐个下载
|
||||
#[allow(clippy::too_many_arguments)] // 透传 (state, registry, source, product, 版本) 的管线参数
|
||||
async fn download_by_coordinates(
|
||||
state: &AppState,
|
||||
registry: &ObservationRegistry,
|
||||
@@ -176,6 +178,7 @@ async fn download_by_coordinates(
|
||||
}
|
||||
|
||||
/// 按标识符：解析各源标识字符串 → 逐个下载（跳过 cone 检索）
|
||||
#[allow(clippy::too_many_arguments)] // 透传 (state, registry, source, product, 版本) 的管线参数
|
||||
async fn download_by_identifiers(
|
||||
state: &AppState,
|
||||
registry: &ObservationRegistry,
|
||||
@@ -249,6 +252,7 @@ async fn download_by_identifiers(
|
||||
}
|
||||
|
||||
/// 便捷：从单个标识符构造请求并下载（供 CLI / 简单场景使用）
|
||||
#[allow(clippy::too_many_arguments)] // 透传 (state, registry, source, product, 版本) 的管线参数
|
||||
pub async fn download_one(
|
||||
state: &AppState,
|
||||
registry: &ObservationRegistry,
|
||||
|
||||
@@ -57,6 +57,9 @@ pub struct Candidate {
|
||||
///
|
||||
/// 缓存、重试、选源等通用编排由 trait 默认方法和 dispatch.rs 负责。
|
||||
#[async_trait]
|
||||
// cone 检索模板方法需透传 (state, 坐标, 半径, release, subtype, version)，
|
||||
// async_trait 生成的装箱垫片同样超过参数阈值
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub trait ObservationFetcher: Send + Sync + std::fmt::Debug {
|
||||
/// 该 fetcher 处理的 (源, 产品类型) 组合
|
||||
fn key(&self) -> (Source, ProductType);
|
||||
@@ -183,6 +186,7 @@ pub trait ObservationFetcher: Send + Sync + std::fmt::Debug {
|
||||
///
|
||||
/// 复用 vizier_query_cache 表（与历史设计一致），用 cache_prefix 区分各源各产品。
|
||||
/// 缓存命中反序列化为 Vec<Candidate>；未命中调 cone_search_raw 并写回。
|
||||
#[allow(clippy::too_many_arguments)] // 模板方法透传 cone 检索与版本参数
|
||||
async fn cone_search_cached(
|
||||
&self,
|
||||
state: &AppState,
|
||||
|
||||
@@ -0,0 +1,212 @@
|
||||
// src/services/observation/jwst.rs
|
||||
//
|
||||
// JWST 光谱 fetcher —— 通过 MAST API 获取 NIRSpec/MIRI 光谱
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::services::observation::cache::{
|
||||
cached_files_total_size, fetch_observation_cache, log_write_failure, persist_bytes,
|
||||
write_observation_cache,
|
||||
};
|
||||
use crate::services::observation::fetcher::{Candidate, ObservationFetcher};
|
||||
use crate::services::observation::types::{
|
||||
Artifact, ObservationProduct, ProductSpec, ProductType, Source,
|
||||
};
|
||||
use anyhow::{anyhow, Result};
|
||||
use tracing::{info, warn};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct JwstSpectrumFetcher;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl ObservationFetcher for JwstSpectrumFetcher {
|
||||
fn key(&self) -> (Source, ProductType) {
|
||||
(Source::Jwst, ProductType::Spectrum)
|
||||
}
|
||||
|
||||
fn subtypes(&self) -> &'static [&'static str] {
|
||||
&["nirspec_ifu", "nirspec_mos", "miri_mrs", "miri_lrs"]
|
||||
}
|
||||
|
||||
fn suggested_max_radius_deg(&self) -> f64 {
|
||||
0.1
|
||||
}
|
||||
|
||||
fn hard_max_radius_deg(&self) -> f64 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn identifier_format(&self) -> Option<&'static str> {
|
||||
Some("JWST 观测 ID（如 'jw00001001001'）")
|
||||
}
|
||||
|
||||
async fn cone_search_raw(
|
||||
&self,
|
||||
state: &AppState,
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
radius_deg: f64,
|
||||
_release: Option<&str>,
|
||||
_subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
) -> Result<Vec<Candidate>> {
|
||||
let jwst = &state.sources.jwst;
|
||||
let rows = jwst.cone_search(ra, dec, radius_deg).await?;
|
||||
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|r| Candidate {
|
||||
source: Source::Jwst,
|
||||
source_id: r.obs_id.clone(),
|
||||
label: format!("{} ({})", r.target_name, r.instrument),
|
||||
ra: Some(r.ra),
|
||||
dec: Some(r.dec),
|
||||
distance: None,
|
||||
raw: Some(serde_json::json!({
|
||||
"target_name": r.target_name,
|
||||
"instrument": r.instrument,
|
||||
"observation_mode": r.observation_mode,
|
||||
"filter_grating": r.filter_grating,
|
||||
"exposure_time": r.exposure_time,
|
||||
})),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn resolve_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
_release: Option<&str>,
|
||||
_subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
) -> Result<Candidate> {
|
||||
Ok(Candidate {
|
||||
source: Source::Jwst,
|
||||
source_id: identifier.to_string(),
|
||||
label: identifier.to_string(),
|
||||
ra: None,
|
||||
dec: None,
|
||||
distance: None,
|
||||
raw: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch(
|
||||
&self,
|
||||
state: &AppState,
|
||||
candidate: &Candidate,
|
||||
_release: Option<&str>,
|
||||
subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
force: bool,
|
||||
) -> Result<ObservationProduct> {
|
||||
let sub = subtype.unwrap_or("nirspec_ifu");
|
||||
let product = ProductSpec::with_subtype(ProductType::Spectrum, sub);
|
||||
let cache_key = format!("jwst|{}|{}", candidate.source_id, sub);
|
||||
let source_label = candidate.label.clone();
|
||||
|
||||
// 1) 检查缓存
|
||||
if !force {
|
||||
if let Some((artifacts, meta)) =
|
||||
fetch_observation_cache(&state.db, Source::Jwst, &product, &cache_key).await?
|
||||
{
|
||||
if cached_files_total_size(&state.config.storage.library_dir, &artifacts)
|
||||
.await
|
||||
.is_some()
|
||||
{
|
||||
info!("[JWST] 缓存命中 (obs_id={})", candidate.source_id);
|
||||
return Ok(ObservationProduct {
|
||||
source: Source::Jwst,
|
||||
product,
|
||||
source_id: cache_key,
|
||||
source_label,
|
||||
artifacts,
|
||||
source_meta: meta,
|
||||
});
|
||||
}
|
||||
warn!(
|
||||
"[JWST] 缓存文件缺失，重新下载 (obs_id={})",
|
||||
candidate.source_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 2) 真实获取
|
||||
let jwst = &state.sources.jwst;
|
||||
let products = jwst.list_data_products(&candidate.source_id).await?;
|
||||
|
||||
// 筛选匹配 subtype 的产品
|
||||
let target_product = match sub {
|
||||
"nirspec_ifu" => "NRS1",
|
||||
"nirspec_mos" => "NRS1",
|
||||
"miri_mrs" => "MRS",
|
||||
"miri_lrs" => "SLITLOS",
|
||||
_ => "NRS1",
|
||||
};
|
||||
|
||||
let matching: Vec<_> = products
|
||||
.iter()
|
||||
.filter(|p| {
|
||||
p.product_type.contains(target_product)
|
||||
|| p.file_name.contains("s2d")
|
||||
|| p.file_name.contains("x1d")
|
||||
})
|
||||
.collect();
|
||||
|
||||
if matching.is_empty() {
|
||||
return Err(anyhow!(
|
||||
"JWST 观测 {} 无匹配的光谱产品",
|
||||
candidate.source_id
|
||||
));
|
||||
}
|
||||
|
||||
let product_item = &matching[0];
|
||||
let bytes = reqwest::get(&product_item.file_url)
|
||||
.await
|
||||
.map_err(|e| anyhow!("JWST 下载失败: {}", e))?
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| anyhow!("JWST 读取失败: {}", e))?
|
||||
.to_vec();
|
||||
|
||||
let file_path = format!(
|
||||
"Telescope/JWST/spectrum/{}/{}",
|
||||
candidate.source_id, product_item.file_name
|
||||
);
|
||||
persist_bytes(&state.config.storage.library_dir, &file_path, &bytes).await?;
|
||||
|
||||
let artifacts = vec![Artifact {
|
||||
band: None,
|
||||
original_name: Some(product_item.file_name.clone()),
|
||||
file_path: file_path.clone(),
|
||||
file_url: product_item.file_url.clone(),
|
||||
file_format: "fits".to_string(),
|
||||
size_bytes: bytes.len(),
|
||||
cached: false,
|
||||
}];
|
||||
|
||||
// 3) 写入缓存
|
||||
if let Err(e) = write_observation_cache(
|
||||
&state.db,
|
||||
Source::Jwst,
|
||||
&product,
|
||||
&cache_key,
|
||||
candidate.ra,
|
||||
candidate.dec,
|
||||
&artifacts,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
log_write_failure("jwst", "spectrum", e);
|
||||
}
|
||||
|
||||
Ok(ObservationProduct {
|
||||
source: Source::Jwst,
|
||||
product,
|
||||
source_id: cache_key,
|
||||
source_label,
|
||||
artifacts,
|
||||
source_meta: candidate.raw.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -27,11 +27,13 @@
|
||||
// 3. 在 types.rs Source/ProductType 枚举加变体（若为新源/新产品）
|
||||
|
||||
pub mod cache;
|
||||
pub mod cutout;
|
||||
pub mod desi;
|
||||
pub mod dispatch;
|
||||
pub mod fetcher;
|
||||
pub mod gaia;
|
||||
pub mod gaia_xp;
|
||||
pub mod jwst;
|
||||
pub mod lamost;
|
||||
pub mod photometry;
|
||||
pub mod preview;
|
||||
@@ -40,6 +42,7 @@ pub mod sdss;
|
||||
pub mod tess;
|
||||
pub mod types;
|
||||
pub mod unified;
|
||||
pub mod xray;
|
||||
pub mod ztf;
|
||||
|
||||
// 统一入口 re-export
|
||||
|
||||
@@ -769,8 +769,8 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_source_from_str_aliases() {
|
||||
assert_eq!(Source::from_str("2mass").unwrap(), Source::Twomass);
|
||||
assert_eq!(Source::from_str("wise").unwrap(), Source::Allwise);
|
||||
assert_eq!(Source::from_str("ps1").unwrap(), Source::Panstarrs);
|
||||
assert_eq!(Source::parse("2mass").unwrap(), Source::Twomass);
|
||||
assert_eq!(Source::parse("wise").unwrap(), Source::Allwise);
|
||||
assert_eq!(Source::parse("ps1").unwrap(), Source::Panstarrs);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -197,6 +197,35 @@ pub async fn build_preview(
|
||||
let mut preview =
|
||||
parse_by_source(&abs_path, source, subtype, source_id, &artifact.file_path)
|
||||
.with_context(|| format!("解析 FITS 失败: {}", abs_path.display()))?;
|
||||
|
||||
// 尝试读取 FITS header 并补充 teff, logg, fe_h
|
||||
if let Ok(mut fits) = fitsio::FitsFile::open(&abs_path) {
|
||||
if preview.meta.teff.is_none() {
|
||||
preview.meta.teff =
|
||||
crate::services::spectrum::parameters::read_f64_with_fallback(
|
||||
&mut fits,
|
||||
&["TEFF", "T_EFF", "TEMP", "EFFECTIVE_TEMP", "TEFF_ASPCAP"],
|
||||
)
|
||||
.map(|v| v as f32);
|
||||
}
|
||||
if preview.meta.logg.is_none() {
|
||||
preview.meta.logg =
|
||||
crate::services::spectrum::parameters::read_f64_with_fallback(
|
||||
&mut fits,
|
||||
&["LOGG", "LOG_G", "GRAVITY", "LOGG_ASPCAP"],
|
||||
)
|
||||
.map(|v| v as f32);
|
||||
}
|
||||
if preview.meta.fe_h.is_none() {
|
||||
preview.meta.fe_h =
|
||||
crate::services::spectrum::parameters::read_f64_with_fallback(
|
||||
&mut fits,
|
||||
&["FEH", "FE_H", "[Fe/H]", "FEH_ASPCAP"],
|
||||
)
|
||||
.map(|v| v as f32);
|
||||
}
|
||||
}
|
||||
|
||||
for seg in preview.segments.iter_mut() {
|
||||
if seg.wavelength.len() > MAX_PREVIEW_POINTS {
|
||||
downsample_segment(seg, MAX_PREVIEW_POINTS);
|
||||
@@ -240,6 +269,10 @@ fn parse_by_source(
|
||||
Source::Twomass | Source::Allwise | Source::Panstarrs | Source::Ztf | Source::Tess => Err(
|
||||
anyhow!("{} 的预览解析暂未实现（测光/光变产品）", source.display()),
|
||||
),
|
||||
// JWST/X射线产品的预览解析器后续实现
|
||||
Source::Jwst | Source::Xmm | Source::Chandra => {
|
||||
Err(anyhow!("{} 的预览解析暂未实现", source.display()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -831,7 +864,7 @@ mod tests {
|
||||
&seg.wavelength[..10.min(seg.wavelength.len())]
|
||||
);
|
||||
println!("First 10 fluxes: {:?}", &seg.flux[..10.min(seg.flux.len())]);
|
||||
assert!(seg.wavelength.len() > 0);
|
||||
assert!(!seg.wavelength.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -864,7 +897,7 @@ mod tests {
|
||||
"APOGEE First 10 fluxes: {:?}",
|
||||
&seg.flux[..10.min(seg.flux.len())]
|
||||
);
|
||||
assert!(seg.wavelength.len() > 0);
|
||||
assert!(!seg.wavelength.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -10,9 +10,11 @@
|
||||
// 注意：SDSS 的 specobj 与 APOGEE 共用 (Sdss, Spectrum) key，但 subtype 不同。
|
||||
// registry 用 (Source, ProductType) → Vec<Arc<dyn Fetcher>>，dispatch 再按 subtype 二级筛选。
|
||||
|
||||
use crate::services::observation::cutout::CutoutFetcher;
|
||||
use crate::services::observation::desi::DesiSpectrumFetcher;
|
||||
use crate::services::observation::fetcher::ObservationFetcher;
|
||||
use crate::services::observation::gaia::{GaiaLightCurveFetcher, GaiaSpectrumFetcher};
|
||||
use crate::services::observation::jwst::JwstSpectrumFetcher;
|
||||
use crate::services::observation::lamost::LamostSpectrumFetcher;
|
||||
use crate::services::observation::photometry::{
|
||||
AllwisePhotometryFetcher, GaiaPhotometryFetcher, PanstarrsPhotometryFetcher,
|
||||
@@ -21,6 +23,7 @@ use crate::services::observation::photometry::{
|
||||
use crate::services::observation::sdss::{SdssApogeeFetcher, SdssSpectrumFetcher};
|
||||
use crate::services::observation::tess::TessLightCurveFetcher;
|
||||
use crate::services::observation::types::{ProductType, Source};
|
||||
use crate::services::observation::xray::{ChandraSpectrumFetcher, XmmSpectrumFetcher};
|
||||
use crate::services::observation::ztf::ZtfLightCurveFetcher;
|
||||
use anyhow::{anyhow, Result};
|
||||
use std::collections::HashMap;
|
||||
@@ -31,9 +34,9 @@ pub struct ObservationRegistry {
|
||||
fetchers: HashMap<(Source, ProductType), Vec<Arc<dyn ObservationFetcher>>>,
|
||||
}
|
||||
|
||||
impl ObservationRegistry {
|
||||
impl Default for ObservationRegistry {
|
||||
/// 构建默认注册表（项目启动时调用一次）
|
||||
pub fn default() -> Self {
|
||||
fn default() -> Self {
|
||||
let mut fetchers: HashMap<(Source, ProductType), Vec<Arc<dyn ObservationFetcher>>> =
|
||||
HashMap::new();
|
||||
|
||||
@@ -63,10 +66,19 @@ impl ObservationRegistry {
|
||||
// 光变曲线（ZTF / TESS，新 HTTP 客户端）
|
||||
register!(ZtfLightCurveFetcher);
|
||||
register!(TessLightCurveFetcher);
|
||||
// JWST 光谱
|
||||
register!(JwstSpectrumFetcher);
|
||||
// X 射线光谱（XMM-Newton / Chandra）
|
||||
register!(XmmSpectrumFetcher);
|
||||
register!(ChandraSpectrumFetcher);
|
||||
// 图像 Cutout（SDSS / DSS / Pan-STARRS）
|
||||
register!(CutoutFetcher);
|
||||
|
||||
Self { fetchers }
|
||||
}
|
||||
}
|
||||
|
||||
impl ObservationRegistry {
|
||||
/// 按 (Source, ProductType) + subtype 取 fetcher
|
||||
///
|
||||
/// subtype 匹配规则：
|
||||
|
||||
@@ -34,6 +34,12 @@ pub enum Source {
|
||||
Ztf,
|
||||
/// TESS 空间时域光变曲线（MAST）
|
||||
Tess,
|
||||
/// JWST（James Webb Space Telescope）NIRSpec/MIRI 光谱
|
||||
Jwst,
|
||||
/// XMM-Newton X 射线光谱
|
||||
Xmm,
|
||||
/// Chandra X 射线光谱
|
||||
Chandra,
|
||||
}
|
||||
|
||||
impl Source {
|
||||
@@ -48,6 +54,9 @@ impl Source {
|
||||
"panstarrs",
|
||||
"ztf",
|
||||
"tess",
|
||||
"jwst",
|
||||
"xmm",
|
||||
"chandra",
|
||||
]
|
||||
}
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
@@ -61,12 +70,15 @@ impl Source {
|
||||
Self::Panstarrs => "panstarrs",
|
||||
Self::Ztf => "ztf",
|
||||
Self::Tess => "tess",
|
||||
Self::Jwst => "jwst",
|
||||
Self::Xmm => "xmm",
|
||||
Self::Chandra => "chandra",
|
||||
}
|
||||
}
|
||||
|
||||
/// 从字符串解析数据源（大小写不敏感）。新增源只需在此加一行 match arm，
|
||||
/// 所有调用方（API handler、Agent tool）统一用本方法，不再各自手写 match。
|
||||
pub fn from_str(s: &str) -> Result<Self, String> {
|
||||
pub fn parse(s: &str) -> Result<Self, String> {
|
||||
Ok(match s.to_lowercase().as_str() {
|
||||
"lamost" => Self::Lamost,
|
||||
"gaia" => Self::Gaia,
|
||||
@@ -77,6 +89,9 @@ impl Source {
|
||||
"panstarrs" | "ps1" | "pan-starrs" => Self::Panstarrs,
|
||||
"ztf" => Self::Ztf,
|
||||
"tess" => Self::Tess,
|
||||
"jwst" => Self::Jwst,
|
||||
"xmm" | "xmm-newton" => Self::Xmm,
|
||||
"chandra" => Self::Chandra,
|
||||
other => {
|
||||
return Err(format!(
|
||||
"不支持的 source '{}'，可选: {:?}",
|
||||
@@ -98,10 +113,21 @@ impl Source {
|
||||
Self::Panstarrs => "Pan-STARRS",
|
||||
Self::Ztf => "ZTF",
|
||||
Self::Tess => "TESS",
|
||||
Self::Jwst => "JWST",
|
||||
Self::Xmm => "XMM-Newton",
|
||||
Self::Chandra => "Chandra",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::str::FromStr for Source {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, String> {
|
||||
Self::parse(s)
|
||||
}
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// 2. 第二轴：观测数据产品类型
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
@@ -136,7 +162,7 @@ impl ProductType {
|
||||
|
||||
/// 从字符串解析产品类型（大小写不敏感，接受常见别名）。
|
||||
/// 别名集中在此一处管理，所有调用方统一用本方法。
|
||||
pub fn from_str(s: &str) -> Result<Self, String> {
|
||||
pub fn parse(s: &str) -> Result<Self, String> {
|
||||
Ok(match s.to_lowercase().as_str() {
|
||||
"spectrum" | "spec" => Self::Spectrum,
|
||||
"lightcurve" | "light_curve" | "lc" => Self::LightCurve,
|
||||
@@ -167,6 +193,14 @@ impl ProductType {
|
||||
}
|
||||
}
|
||||
|
||||
impl std::str::FromStr for ProductType {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, String> {
|
||||
Self::parse(s)
|
||||
}
|
||||
}
|
||||
|
||||
/// 产品细分（可选）：用 (类型, 子类型) 描述具体产品变体
|
||||
///
|
||||
/// 子类型语义因源而异：
|
||||
|
||||
@@ -38,6 +38,15 @@ fn default_radius_deg() -> f64 {
|
||||
0.1
|
||||
}
|
||||
|
||||
/// 候选源分组键：(source, product, subtype, release, version)
|
||||
type GroupKey = (
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
);
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// 1. 请求/响应类型
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
@@ -198,16 +207,7 @@ pub async fn unified_search(
|
||||
|
||||
// 按 (source, product, release, version) 分组合并候选源
|
||||
// 这里用 (source_str, product_str, subtype_str, release, version) 作为排序键
|
||||
let mut groups_map: BTreeMap<
|
||||
(
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
),
|
||||
Vec<Candidate>,
|
||||
> = BTreeMap::new();
|
||||
let mut groups_map: BTreeMap<GroupKey, Vec<Candidate>> = BTreeMap::new();
|
||||
let mut total_candidates = 0usize;
|
||||
let mut fail_count = 0usize;
|
||||
|
||||
@@ -242,10 +242,10 @@ pub async fn unified_search(
|
||||
// 转换为 SourceCandidateGroup 列表（BTreeMap 已保证按 key 排序）
|
||||
let mut groups: Vec<SourceCandidateGroup> = Vec::with_capacity(groups_map.len());
|
||||
for ((source_str, product_str, subtype, release, version), candidates) in groups_map {
|
||||
let source = Source::from_str(&source_str)
|
||||
let source = Source::parse(&source_str)
|
||||
.map_err(|e| anyhow::anyhow!("内部错误：无效 source 回填 '{}': {}", source_str, e))?;
|
||||
let product = ProductSpec {
|
||||
product: crate::services::observation::types::ProductType::from_str(&product_str)
|
||||
product: crate::services::observation::types::ProductType::parse(&product_str)
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!("内部错误：无效 product 回填 '{}': {}", product_str, e)
|
||||
})?,
|
||||
|
||||
@@ -0,0 +1,357 @@
|
||||
// src/services/observation/xray.rs
|
||||
//
|
||||
// X 射线光谱 fetcher —— XMM-Newton / Chandra 通过 HEASARC TAP
|
||||
|
||||
use crate::api::AppState;
|
||||
use crate::services::observation::cache::{
|
||||
cached_files_total_size, fetch_observation_cache, log_write_failure, persist_bytes,
|
||||
write_observation_cache,
|
||||
};
|
||||
use crate::services::observation::fetcher::{Candidate, ObservationFetcher};
|
||||
use crate::services::observation::types::{
|
||||
Artifact, ObservationProduct, ProductSpec, ProductType, Source,
|
||||
};
|
||||
use anyhow::{anyhow, Result};
|
||||
use tracing::{info, warn};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct XmmSpectrumFetcher;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl ObservationFetcher for XmmSpectrumFetcher {
|
||||
fn key(&self) -> (Source, ProductType) {
|
||||
(Source::Xmm, ProductType::Spectrum)
|
||||
}
|
||||
|
||||
fn subtypes(&self) -> &'static [&'static str] {
|
||||
&["epic-pn", "epic-mos1", "epic-mos2", "rgs"]
|
||||
}
|
||||
|
||||
fn suggested_max_radius_deg(&self) -> f64 {
|
||||
0.5
|
||||
}
|
||||
|
||||
fn identifier_format(&self) -> Option<&'static str> {
|
||||
Some("XMM 观测 ID（数字）")
|
||||
}
|
||||
|
||||
async fn cone_search_raw(
|
||||
&self,
|
||||
state: &AppState,
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
radius_deg: f64,
|
||||
_release: Option<&str>,
|
||||
_subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
) -> Result<Vec<Candidate>> {
|
||||
let heasarc = &state.sources.heasarc;
|
||||
let rows = heasarc.xmm_cone_search(ra, dec, radius_deg).await?;
|
||||
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|r| Candidate {
|
||||
source: Source::Xmm,
|
||||
source_id: r.obs_id.clone(),
|
||||
label: format!("{} ({})", r.target_name, r.instrument),
|
||||
ra: Some(r.ra),
|
||||
dec: Some(r.dec),
|
||||
distance: None,
|
||||
raw: Some(serde_json::json!({
|
||||
"target_name": r.target_name,
|
||||
"instrument": r.instrument,
|
||||
"exposure": r.exposure,
|
||||
})),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn resolve_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
_release: Option<&str>,
|
||||
_subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
) -> Result<Candidate> {
|
||||
Ok(Candidate {
|
||||
source: Source::Xmm,
|
||||
source_id: identifier.to_string(),
|
||||
label: identifier.to_string(),
|
||||
ra: None,
|
||||
dec: None,
|
||||
distance: None,
|
||||
raw: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch(
|
||||
&self,
|
||||
state: &AppState,
|
||||
candidate: &Candidate,
|
||||
_release: Option<&str>,
|
||||
subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
force: bool,
|
||||
) -> Result<ObservationProduct> {
|
||||
let instrument = subtype.unwrap_or("epic-pn");
|
||||
let product = ProductSpec::with_subtype(ProductType::Spectrum, instrument);
|
||||
let cache_key = format!("xray|{}|{}", candidate.source_id, instrument);
|
||||
let source_label = candidate.label.clone();
|
||||
|
||||
// 1) 检查缓存
|
||||
if !force {
|
||||
if let Some((artifacts, meta)) =
|
||||
fetch_observation_cache(&state.db, Source::Xmm, &product, &cache_key).await?
|
||||
{
|
||||
if cached_files_total_size(&state.config.storage.library_dir, &artifacts)
|
||||
.await
|
||||
.is_some()
|
||||
{
|
||||
info!("[XMM] 缓存命中 (obs_id={})", candidate.source_id);
|
||||
return Ok(ObservationProduct {
|
||||
source: Source::Xmm,
|
||||
product,
|
||||
source_id: cache_key,
|
||||
source_label,
|
||||
artifacts,
|
||||
source_meta: meta,
|
||||
});
|
||||
}
|
||||
warn!(
|
||||
"[XMM] 缓存文件缺失，重新下载 (obs_id={})",
|
||||
candidate.source_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// XMM 数据通过 HEASARC FTP 获取
|
||||
let file_name = format!("{}_{}.fits", candidate.source_id, instrument);
|
||||
let file_path = format!(
|
||||
"Telescope/XMM/spectrum/{}/{}",
|
||||
candidate.source_id, file_name
|
||||
);
|
||||
let data_url = format!(
|
||||
"https://heasarc.gsfc.nasa.gov/FTP/heasarc/xmm/data/obs/{}/{}",
|
||||
candidate.source_id, file_name
|
||||
);
|
||||
|
||||
// 下载 FITS
|
||||
let resp = reqwest::get(&data_url)
|
||||
.await
|
||||
.map_err(|e| anyhow!("XMM 下载失败: {}", e))?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow!("XMM 下载失败: {}", resp.status()));
|
||||
}
|
||||
let bytes = resp
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| anyhow!("XMM 读取失败: {}", e))?
|
||||
.to_vec();
|
||||
persist_bytes(&state.config.storage.library_dir, &file_path, &bytes).await?;
|
||||
|
||||
let artifacts = vec![Artifact {
|
||||
band: None,
|
||||
original_name: Some(file_name),
|
||||
file_path: file_path.clone(),
|
||||
file_url: data_url,
|
||||
file_format: "fits".to_string(),
|
||||
size_bytes: bytes.len(),
|
||||
cached: false,
|
||||
}];
|
||||
|
||||
// 3) 写入缓存
|
||||
if let Err(e) = write_observation_cache(
|
||||
&state.db,
|
||||
Source::Xmm,
|
||||
&product,
|
||||
&cache_key,
|
||||
candidate.ra,
|
||||
candidate.dec,
|
||||
&artifacts,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
log_write_failure("xmm", "spectrum", e);
|
||||
}
|
||||
|
||||
Ok(ObservationProduct {
|
||||
source: Source::Xmm,
|
||||
product,
|
||||
source_id: cache_key,
|
||||
source_label,
|
||||
artifacts,
|
||||
source_meta: candidate.raw.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ChandraSpectrumFetcher;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl ObservationFetcher for ChandraSpectrumFetcher {
|
||||
fn key(&self) -> (Source, ProductType) {
|
||||
(Source::Chandra, ProductType::Spectrum)
|
||||
}
|
||||
|
||||
fn subtypes(&self) -> &'static [&'static str] {
|
||||
&["acis-s", "acis-i", "hetg", "letg"]
|
||||
}
|
||||
|
||||
fn suggested_max_radius_deg(&self) -> f64 {
|
||||
0.5
|
||||
}
|
||||
|
||||
fn identifier_format(&self) -> Option<&'static str> {
|
||||
Some("Chandra 观测 ID（数字）")
|
||||
}
|
||||
|
||||
async fn cone_search_raw(
|
||||
&self,
|
||||
state: &AppState,
|
||||
ra: f64,
|
||||
dec: f64,
|
||||
radius_deg: f64,
|
||||
_release: Option<&str>,
|
||||
_subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
) -> Result<Vec<Candidate>> {
|
||||
let heasarc = &state.sources.heasarc;
|
||||
let rows = heasarc.chandra_cone_search(ra, dec, radius_deg).await?;
|
||||
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|r| Candidate {
|
||||
source: Source::Chandra,
|
||||
source_id: r.obs_id.clone(),
|
||||
label: format!("{} ({})", r.target_name, r.instrument),
|
||||
ra: Some(r.ra),
|
||||
dec: Some(r.dec),
|
||||
distance: None,
|
||||
raw: Some(serde_json::json!({
|
||||
"target_name": r.target_name,
|
||||
"instrument": r.instrument,
|
||||
"exposure": r.exposure,
|
||||
})),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn resolve_identifier(
|
||||
&self,
|
||||
identifier: &str,
|
||||
_release: Option<&str>,
|
||||
_subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
) -> Result<Candidate> {
|
||||
Ok(Candidate {
|
||||
source: Source::Chandra,
|
||||
source_id: identifier.to_string(),
|
||||
label: identifier.to_string(),
|
||||
ra: None,
|
||||
dec: None,
|
||||
distance: None,
|
||||
raw: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch(
|
||||
&self,
|
||||
state: &AppState,
|
||||
candidate: &Candidate,
|
||||
_release: Option<&str>,
|
||||
subtype: Option<&str>,
|
||||
_version: Option<&str>,
|
||||
force: bool,
|
||||
) -> Result<ObservationProduct> {
|
||||
let instrument = subtype.unwrap_or("acis-s");
|
||||
let product = ProductSpec::with_subtype(ProductType::Spectrum, instrument);
|
||||
let cache_key = format!("xray|{}|{}", candidate.source_id, instrument);
|
||||
let source_label = candidate.label.clone();
|
||||
|
||||
// 1) 检查缓存
|
||||
if !force {
|
||||
if let Some((artifacts, meta)) =
|
||||
fetch_observation_cache(&state.db, Source::Chandra, &product, &cache_key).await?
|
||||
{
|
||||
if cached_files_total_size(&state.config.storage.library_dir, &artifacts)
|
||||
.await
|
||||
.is_some()
|
||||
{
|
||||
info!("[Chandra] 缓存命中 (obs_id={})", candidate.source_id);
|
||||
return Ok(ObservationProduct {
|
||||
source: Source::Chandra,
|
||||
product,
|
||||
source_id: cache_key,
|
||||
source_label,
|
||||
artifacts,
|
||||
source_meta: meta,
|
||||
});
|
||||
}
|
||||
warn!(
|
||||
"[Chandra] 缓存文件缺失，重新下载 (obs_id={})",
|
||||
candidate.source_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let file_name = format!("{}_repro.pha", candidate.source_id);
|
||||
let file_path = format!(
|
||||
"Telescope/Chandra/spectrum/{}/{}",
|
||||
candidate.source_id, file_name
|
||||
);
|
||||
let data_url = format!(
|
||||
"https://heasarc.gsfc.nasa.gov/FTP/heasarc/chandra/data/reproc/{}/combined/{}",
|
||||
candidate.source_id, file_name
|
||||
);
|
||||
|
||||
let resp = reqwest::get(&data_url)
|
||||
.await
|
||||
.map_err(|e| anyhow!("Chandra 下载失败: {}", e))?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow!("Chandra 下载失败: {}", resp.status()));
|
||||
}
|
||||
let bytes = resp
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Chandra 读取失败: {}", e))?
|
||||
.to_vec();
|
||||
persist_bytes(&state.config.storage.library_dir, &file_path, &bytes).await?;
|
||||
|
||||
let artifacts = vec![Artifact {
|
||||
band: None,
|
||||
original_name: Some(file_name),
|
||||
file_path: file_path.clone(),
|
||||
file_url: data_url,
|
||||
file_format: "fits".to_string(),
|
||||
size_bytes: bytes.len(),
|
||||
cached: false,
|
||||
}];
|
||||
|
||||
// 3) 写入缓存
|
||||
if let Err(e) = write_observation_cache(
|
||||
&state.db,
|
||||
Source::Chandra,
|
||||
&product,
|
||||
&cache_key,
|
||||
candidate.ra,
|
||||
candidate.dec,
|
||||
&artifacts,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
log_write_failure("chandra", "spectrum", e);
|
||||
}
|
||||
|
||||
Ok(ObservationProduct {
|
||||
source: Source::Chandra,
|
||||
product,
|
||||
source_id: cache_key,
|
||||
source_label,
|
||||
artifacts,
|
||||
source_meta: candidate.raw.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -579,3 +579,40 @@ pub async fn mark_no_resource_service(
|
||||
let updated_paper = get_paper_from_db(db, library_dir, bibcode).await?;
|
||||
Ok(updated_paper)
|
||||
}
|
||||
|
||||
/// 获取最近入库的文献列表
|
||||
pub async fn get_recent_papers(
|
||||
db: &SqlitePool,
|
||||
library_dir: &Path,
|
||||
query: &str,
|
||||
days_back: i64,
|
||||
limit: i32,
|
||||
) -> Result<Vec<StandardPaper>, sqlx::Error> {
|
||||
let search_pattern = format!("%{}%", query);
|
||||
let time_limit = format!("-{} days", days_back);
|
||||
|
||||
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) AS has_vector \
|
||||
FROM papers \
|
||||
WHERE (title LIKE ? OR abstract LIKE ? OR authors LIKE ?) \
|
||||
AND created_at >= datetime('now', ?) \
|
||||
ORDER BY created_at DESC \
|
||||
LIMIT ?"
|
||||
)
|
||||
.bind(&search_pattern)
|
||||
.bind(&search_pattern)
|
||||
.bind(&search_pattern)
|
||||
.bind(&time_limit)
|
||||
.bind(limit)
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
let papers = rows
|
||||
.into_iter()
|
||||
.map(|r| parse_paper_row(&r, library_dir))
|
||||
.collect();
|
||||
Ok(papers)
|
||||
}
|
||||
|
||||
@@ -9,8 +9,8 @@ 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,
|
||||
LibraryQueryParams, SQLITE_PARAM_LIMIT,
|
||||
get_recent_papers, mark_no_resource_service, parse_paper_row, save_paper_to_db,
|
||||
save_paper_to_db_tx, CleanId, LibraryQueryParams, SQLITE_PARAM_LIMIT,
|
||||
};
|
||||
pub use ingest::upload_paper_file_service;
|
||||
pub use model::{convert_ads_doc_to_standard, convert_arxiv_to_standard, StandardPaper};
|
||||
|
||||
@@ -232,6 +232,9 @@ pub async fn search_local_library(
|
||||
) -> anyhow::Result<Vec<StandardPaper>> {
|
||||
// FTS5 查询净化，防止操作符注入
|
||||
let sanitized = sanitize_fts5_query(query);
|
||||
if sanitized.trim().is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let fts_query = format!("\"{}\"", sanitized);
|
||||
|
||||
let rows = sqlx::query_as::<_, PaperRow>(
|
||||
@@ -319,6 +322,9 @@ pub async fn search_agent_history(
|
||||
) -> Result<Vec<SearchResult>, sqlx::Error> {
|
||||
let mut results = Vec::new();
|
||||
let sanitized = sanitize_fts5_query(q);
|
||||
if sanitized.trim().is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// 搜索会话
|
||||
if scope == "all" || scope == "sessions" {
|
||||
|
||||
@@ -0,0 +1,562 @@
|
||||
// src/services/spectrum/cross_correlate.rs
|
||||
//
|
||||
// 交叉相关法测径向速度（RV）
|
||||
//
|
||||
// 算法：
|
||||
// 1. 将观测光谱和模板光谱对齐到同一波长网格（对数重采样）
|
||||
// 2. 归一化（除以连续谱估计值）
|
||||
// 3. 计算交叉相关函数（CCF）
|
||||
// 4. 从 CCF 峰值位置提取 RV + 不确定度
|
||||
//
|
||||
// 支持：
|
||||
// - 内置 Pickles 恒星类型模板（简化版：按光谱型插值）
|
||||
// - 用户自定义模板
|
||||
// - 多波段光谱
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// RV 测量结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RadialVelocityResult {
|
||||
/// 径向速度 (km/s)
|
||||
pub rv: f64,
|
||||
/// RV 不确定度 (km/s)
|
||||
pub rv_error: f64,
|
||||
/// 交叉相关峰值（1=完美匹配）
|
||||
pub ccf_peak: f64,
|
||||
/// 使用的模板名称
|
||||
pub template_name: String,
|
||||
/// 有效波长范围 (Å)
|
||||
pub wavelength_range: (f64, f64),
|
||||
/// SNR（估计）
|
||||
pub snr: Option<f64>,
|
||||
}
|
||||
|
||||
/// 内置模板类型（简化 Pickles 恒星类型）
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub enum TemplateType {
|
||||
/// O 型星
|
||||
O,
|
||||
/// B 型星
|
||||
B,
|
||||
/// A 型星
|
||||
A,
|
||||
/// F 型星
|
||||
F,
|
||||
/// G 型星（类太阳）
|
||||
G,
|
||||
/// K 型星
|
||||
K,
|
||||
/// M 型星
|
||||
M,
|
||||
/// 廉谱（M-dwarf）
|
||||
Mwarf,
|
||||
/// 巨星
|
||||
Giant,
|
||||
}
|
||||
|
||||
/// 获取模板的典型有效温度（用于标记）
|
||||
pub fn template_teff(template: TemplateType) -> f64 {
|
||||
match template {
|
||||
TemplateType::O => 35000.0,
|
||||
TemplateType::B => 20000.0,
|
||||
TemplateType::A => 8500.0,
|
||||
TemplateType::F => 6500.0,
|
||||
TemplateType::G => 5778.0,
|
||||
TemplateType::K => 4500.0,
|
||||
TemplateType::M => 3200.0,
|
||||
TemplateType::Mwarf => 3000.0,
|
||||
TemplateType::Giant => 4800.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// 生成简化模板光谱（基于黑体辐射 + 简单吸收特征）
|
||||
///
|
||||
/// 这是一个简化实现：用黑体辐射曲线 + 关键吸收线的高斯凹陷模拟恒星光谱。
|
||||
/// 真实应用中应使用 PHOENIX/BT-Settl 等理论光谱库。
|
||||
pub fn generate_template(template: TemplateType, wavelength: &[f64]) -> Vec<f64> {
|
||||
let teff = template_teff(template);
|
||||
let mut flux = Vec::with_capacity(wavelength.len());
|
||||
|
||||
for &w in wavelength {
|
||||
// 简化黑体辐射（普朗克函数，任意缩放）
|
||||
let w_um = w * 1e-4; // Å → μm
|
||||
let planck = black_body(w_um, teff);
|
||||
flux.push(planck);
|
||||
}
|
||||
|
||||
// 添加简化吸收特征
|
||||
let absorption_lines: Vec<(f64, f64)> = match template {
|
||||
TemplateType::O | TemplateType::B => vec![
|
||||
(4861.0, 0.3), // Hβ
|
||||
(4340.0, 0.25), // Hγ
|
||||
(4102.0, 0.2), // Hδ
|
||||
(4471.0, 0.15), // He I
|
||||
(4026.0, 0.1), // He I
|
||||
],
|
||||
TemplateType::A => vec![
|
||||
(6563.0, 0.5), // Hα
|
||||
(4861.0, 0.6), // Hβ
|
||||
(4340.0, 0.5), // Hγ
|
||||
(4102.0, 0.4), // Hδ
|
||||
],
|
||||
TemplateType::F => vec![
|
||||
(6563.0, 0.3), // Hα
|
||||
(4861.0, 0.35), // Hβ
|
||||
(3934.0, 0.2), // Ca II K
|
||||
(5184.0, 0.15), // Mg b
|
||||
],
|
||||
TemplateType::G => vec![
|
||||
(6563.0, 0.2), // Hα
|
||||
(3934.0, 0.5), // Ca II K
|
||||
(3968.0, 0.45), // Ca II H
|
||||
(5890.0, 0.2), // Na D
|
||||
(5270.0, 0.15), // Fe I
|
||||
(5328.0, 0.12), // Fe I
|
||||
],
|
||||
TemplateType::K | TemplateType::Giant => vec![
|
||||
(3934.0, 0.6), // Ca II K
|
||||
(3968.0, 0.55), // Ca II H
|
||||
(5890.0, 0.35), // Na D
|
||||
(5184.0, 0.25), // Mg b
|
||||
(5270.0, 0.2), // Fe I
|
||||
(4227.0, 0.15), // Ca I
|
||||
],
|
||||
TemplateType::M | TemplateType::Mwarf => vec![
|
||||
(5890.0, 0.4), // Na D
|
||||
(7100.0, 0.3), // TiO band head
|
||||
(6200.0, 0.25), // TiO
|
||||
(5270.0, 0.2), // Fe I
|
||||
(4227.0, 0.15), // Ca I
|
||||
],
|
||||
};
|
||||
|
||||
for (line_center, depth) in &absorption_lines {
|
||||
let sigma = 1.5; // Å，线宽
|
||||
for (i, &w) in wavelength.iter().enumerate() {
|
||||
let g = (-(w - line_center).powi(2) / (2.0 * sigma * sigma)).exp();
|
||||
flux[i] *= 1.0 - depth * g;
|
||||
}
|
||||
}
|
||||
|
||||
flux
|
||||
}
|
||||
|
||||
/// 简化黑体辐射函数（Wien 近似 + 归一化）
|
||||
fn black_body(wavelength_um: f64, teff: f64) -> f64 {
|
||||
let h = 6.626e-34; // Planck constant
|
||||
let c = 3e8; // speed of light
|
||||
let k = 1.381e-23; // Boltzmann constant
|
||||
let wl = wavelength_um * 1e-6;
|
||||
|
||||
if wl <= 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let x = h * c / (wl * k * teff);
|
||||
if x > 500.0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
// B_lambda ∝ 1/(wl^5 * (exp(x) - 1))
|
||||
let b = 1.0 / (wl.powi(5) * (x.exp() - 1.0));
|
||||
b * 1e-10 // 归一化到合理范围
|
||||
}
|
||||
|
||||
/// 对数重采样波长网格
|
||||
fn log_resample(wavelength: &[f64], n_points: usize) -> Vec<f64> {
|
||||
if wavelength.len() < 2 {
|
||||
return wavelength.to_vec();
|
||||
}
|
||||
let w_min = wavelength[0];
|
||||
let w_max = wavelength[wavelength.len() - 1];
|
||||
let log_min = w_min.ln();
|
||||
let log_max = w_max.ln();
|
||||
let step = (log_max - log_min) / (n_points - 1) as f64;
|
||||
|
||||
(0..n_points)
|
||||
.map(|i| (log_min + step * i as f64).exp())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 线性插值重采样到目标波长网格
|
||||
fn resample(wavelength: &[f64], flux: &[f64], target: &[f64]) -> Vec<f64> {
|
||||
let mut result = Vec::with_capacity(target.len());
|
||||
for &tw in target {
|
||||
// 找 tw 在 wavelength 中的位置
|
||||
let idx = match wavelength
|
||||
.binary_search_by(|w| w.partial_cmp(&tw).unwrap_or(std::cmp::Ordering::Equal))
|
||||
{
|
||||
Ok(i) => i,
|
||||
Err(i) => i,
|
||||
};
|
||||
if idx == 0 {
|
||||
result.push(flux[0]);
|
||||
} else if idx >= wavelength.len() {
|
||||
result.push(*flux.last().unwrap_or(&0.0));
|
||||
} else {
|
||||
let w1 = wavelength[idx - 1];
|
||||
let w2 = wavelength[idx];
|
||||
let f1 = flux[idx - 1];
|
||||
let f2 = flux[idx];
|
||||
let t = if (w2 - w1).abs() > 1e-10 {
|
||||
(tw - w1) / (w2 - w1)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
result.push(f1 + t * (f2 - f1));
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// 除以连续谱（简化：用中值滤波器估计）
|
||||
fn normalize_by_continuum(flux: &[f64], kernel_size: usize) -> Vec<f64> {
|
||||
let n = flux.len();
|
||||
let mut continuum = vec![0.0; n];
|
||||
|
||||
for (i, c) in continuum.iter_mut().enumerate() {
|
||||
let lo = i.saturating_sub(kernel_size / 2);
|
||||
let hi = (i + kernel_size / 2 + 1).min(n);
|
||||
let mut window: Vec<f64> = flux[lo..hi].to_vec();
|
||||
window.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
*c = window[window.len() / 2];
|
||||
}
|
||||
|
||||
flux.iter()
|
||||
.zip(continuum.iter())
|
||||
.map(|(f, c)| if *c > 0.0 { f / c } else { 1.0 })
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 计算交叉相关函数
|
||||
///
|
||||
/// 简化实现：直接在 RV 空间扫描（无 FFT 加速，但对典型光谱足够快）
|
||||
fn compute_ccf(
|
||||
obs_norm: &[f64],
|
||||
templ_norm: &[f64],
|
||||
wavelength: &[f64],
|
||||
rv_range: (f64, f64), // km/s
|
||||
rv_step: f64, // km/s
|
||||
) -> Vec<(f64, f64)> {
|
||||
let c = 299792.458; // 光速 km/s
|
||||
let mut ccf = Vec::new();
|
||||
|
||||
let mut rv = rv_range.0;
|
||||
while rv <= rv_range.1 {
|
||||
let shift_factor = 1.0 + rv / c;
|
||||
|
||||
// 对模板进行 Doppler shift
|
||||
let shifted_waves: Vec<f64> = wavelength.iter().map(|&w| w / shift_factor).collect();
|
||||
let shifted_templ = resample(wavelength, templ_norm, &shifted_waves);
|
||||
|
||||
// 计算互相关（归一化）
|
||||
let n = obs_norm.len().min(shifted_templ.len());
|
||||
if n == 0 {
|
||||
ccf.push((rv, 0.0));
|
||||
rv += rv_step;
|
||||
continue;
|
||||
}
|
||||
|
||||
let mean_obs: f64 = obs_norm[..n].iter().sum::<f64>() / n as f64;
|
||||
let mean_templ: f64 = shifted_templ[..n].iter().sum::<f64>() / n as f64;
|
||||
|
||||
let var_obs: f64 = obs_norm[..n].iter().map(|f| (f - mean_obs).powi(2)).sum();
|
||||
let var_templ: f64 = shifted_templ[..n]
|
||||
.iter()
|
||||
.map(|f| (f - mean_templ).powi(2))
|
||||
.sum();
|
||||
|
||||
let denom = (var_obs * var_templ).sqrt();
|
||||
let corr = if denom > 1e-10 {
|
||||
obs_norm[..n]
|
||||
.iter()
|
||||
.zip(shifted_templ[..n].iter())
|
||||
.map(|(o, t)| (o - mean_obs) * (t - mean_templ))
|
||||
.sum::<f64>()
|
||||
/ denom
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
ccf.push((rv, corr));
|
||||
rv += rv_step;
|
||||
}
|
||||
|
||||
ccf
|
||||
}
|
||||
|
||||
/// 从 CCF 中提取峰值位置（抛物线拟合顶点）
|
||||
fn fit_ccf_peak(ccf: &[(f64, f64)]) -> Option<(f64, f64)> {
|
||||
if ccf.len() < 3 {
|
||||
return None;
|
||||
}
|
||||
|
||||
// 找最大值
|
||||
let (peak_idx, _) = ccf.iter().enumerate().max_by(|a, b| {
|
||||
a.1 .1
|
||||
.partial_cmp(&b.1 .1)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})?;
|
||||
|
||||
if peak_idx == 0 || peak_idx == ccf.len() - 1 {
|
||||
return Some(ccf[peak_idx]);
|
||||
}
|
||||
|
||||
// 三点抛物线拟合 (等距格点插值)
|
||||
let (x1, y1) = ccf[peak_idx - 1];
|
||||
let (x2, y2) = ccf[peak_idx];
|
||||
let (_x3, y3) = ccf[peak_idx + 1];
|
||||
|
||||
let h = x2 - x1;
|
||||
let denom = 2.0 * (y1 - 2.0 * y2 + y3);
|
||||
if denom.abs() < 1e-15 {
|
||||
return Some(ccf[peak_idx]);
|
||||
}
|
||||
|
||||
let vertex_x = x2 + h * (y1 - y3) / denom;
|
||||
let vertex_y = y2 - 0.125 * (y1 - y3).powi(2) / (y1 - 2.0 * y2 + y3);
|
||||
|
||||
Some((vertex_x, vertex_y))
|
||||
}
|
||||
|
||||
/// 主入口：测量径向速度
|
||||
pub fn measure_rv(
|
||||
obs_wavelength: &[f64],
|
||||
obs_flux: &[f64],
|
||||
template: TemplateType,
|
||||
rv_range: (f64, f64),
|
||||
rv_step: f64,
|
||||
) -> Result<RadialVelocityResult> {
|
||||
if obs_wavelength.len() != obs_flux.len() {
|
||||
return Err(anyhow!("波长和通量数组长度不一致"));
|
||||
}
|
||||
if obs_wavelength.len() < 10 {
|
||||
return Err(anyhow!("光谱数据点太少（需要 ≥10）"));
|
||||
}
|
||||
|
||||
// 1. 对数重采样到公共波长网格
|
||||
let n_points = obs_wavelength.len();
|
||||
let log_grid = log_resample(obs_wavelength, n_points);
|
||||
let obs_resampled = resample(obs_wavelength, obs_flux, &log_grid);
|
||||
let templ_flux = generate_template(template, &log_grid);
|
||||
|
||||
// 2. 归一化
|
||||
let obs_norm = normalize_by_continuum(&obs_resampled, 51);
|
||||
let templ_norm = normalize_by_continuum(&templ_flux, 51);
|
||||
|
||||
// 3. 计算 CCF
|
||||
let ccf = compute_ccf(&obs_norm, &templ_norm, &log_grid, rv_range, rv_step);
|
||||
|
||||
// 4. 提取峰值
|
||||
let (peak_rv, peak_corr) = fit_ccf_peak(&ccf).ok_or_else(|| anyhow!("CCF 峰值拟合失败"))?;
|
||||
|
||||
// 5. 估计不确定度（基于 CCF 曲率的简化方法）
|
||||
let rv_error = estimate_rv_error(&ccf, peak_rv, rv_step);
|
||||
|
||||
// 6. 估计 SNR（简化：基于连续谱散射）
|
||||
let mean_flux: f64 = obs_resampled.iter().sum::<f64>() / obs_resampled.len() as f64;
|
||||
let variance: f64 = obs_resampled
|
||||
.iter()
|
||||
.map(|f| (f - mean_flux).powi(2))
|
||||
.sum::<f64>()
|
||||
/ obs_resampled.len() as f64;
|
||||
let snr = if variance > 0.0 {
|
||||
Some(mean_flux / variance.sqrt())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(RadialVelocityResult {
|
||||
rv: peak_rv,
|
||||
rv_error,
|
||||
ccf_peak: peak_corr,
|
||||
template_name: format!("{:?} (Teff={:.0}K)", template, template_teff(template)),
|
||||
wavelength_range: (
|
||||
*obs_wavelength.first().unwrap_or(&0.0),
|
||||
*obs_wavelength.last().unwrap_or(&0.0),
|
||||
),
|
||||
snr,
|
||||
})
|
||||
}
|
||||
|
||||
/// 使用用户自定义模板光谱测量 RV
|
||||
pub fn measure_rv_with_template(
|
||||
obs_wavelength: &[f64],
|
||||
obs_flux: &[f64],
|
||||
templ_wavelength: &[f64],
|
||||
templ_flux: &[f64],
|
||||
template_name: &str,
|
||||
rv_range: (f64, f64),
|
||||
rv_step: f64,
|
||||
) -> Result<RadialVelocityResult> {
|
||||
if obs_wavelength.len() != obs_flux.len() {
|
||||
return Err(anyhow!("观测波长和通量数组长度不一致"));
|
||||
}
|
||||
if templ_wavelength.len() != templ_flux.len() {
|
||||
return Err(anyhow!("模板波长和通量数组长度不一致"));
|
||||
}
|
||||
|
||||
// 1. 对数重采样
|
||||
let n_points = obs_wavelength.len().max(templ_wavelength.len());
|
||||
let log_grid = log_resample(obs_wavelength, n_points);
|
||||
let obs_resampled = resample(obs_wavelength, obs_flux, &log_grid);
|
||||
let templ_resampled = resample(templ_wavelength, templ_flux, &log_grid);
|
||||
|
||||
// 2. 归一化
|
||||
let obs_norm = normalize_by_continuum(&obs_resampled, 51);
|
||||
let templ_norm = normalize_by_continuum(&templ_resampled, 51);
|
||||
|
||||
// 3. CCF
|
||||
let ccf = compute_ccf(&obs_norm, &templ_norm, &log_grid, rv_range, rv_step);
|
||||
|
||||
// 4. 峰值
|
||||
let (peak_rv, peak_corr) = fit_ccf_peak(&ccf).ok_or_else(|| anyhow!("CCF 峰值拟合失败"))?;
|
||||
|
||||
let rv_error = estimate_rv_error(&ccf, peak_rv, rv_step);
|
||||
|
||||
let mean_flux: f64 = obs_resampled.iter().sum::<f64>() / obs_resampled.len() as f64;
|
||||
let variance: f64 = obs_resampled
|
||||
.iter()
|
||||
.map(|f| (f - mean_flux).powi(2))
|
||||
.sum::<f64>()
|
||||
/ obs_resampled.len() as f64;
|
||||
let snr = if variance > 0.0 {
|
||||
Some(mean_flux / variance.sqrt())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(RadialVelocityResult {
|
||||
rv: peak_rv,
|
||||
rv_error,
|
||||
ccf_peak: peak_corr,
|
||||
template_name: template_name.to_string(),
|
||||
wavelength_range: (
|
||||
*obs_wavelength.first().unwrap_or(&0.0),
|
||||
*obs_wavelength.last().unwrap_or(&0.0),
|
||||
),
|
||||
snr,
|
||||
})
|
||||
}
|
||||
|
||||
/// 估计 RV 不确定度（基于 CCF 峰值附近的曲率）
|
||||
fn estimate_rv_error(ccf: &[(f64, f64)], peak_rv: f64, rv_step: f64) -> f64 {
|
||||
// 找峰值附近 ±3 步的点，拟合抛物线，从曲率估计不确定度
|
||||
let peak_idx = ccf
|
||||
.iter()
|
||||
.enumerate()
|
||||
.min_by(|a, b| {
|
||||
(a.1 .0 - peak_rv)
|
||||
.abs()
|
||||
.partial_cmp(&(b.1 .0 - peak_rv).abs())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.map(|(i, _)| i)
|
||||
.unwrap_or(0);
|
||||
|
||||
let half_window = 3;
|
||||
let lo = peak_idx.saturating_sub(half_window);
|
||||
let hi = (peak_idx + half_window + 1).min(ccf.len());
|
||||
|
||||
if hi - lo < 3 {
|
||||
return rv_step * 2.0; // 数据不足，返回保守估计
|
||||
}
|
||||
|
||||
// 简化：用半高宽的一半作为不确定度估计
|
||||
let peak_val = ccf[peak_idx].1;
|
||||
let half_max = peak_val / 2.0;
|
||||
|
||||
let mut left_rv = None;
|
||||
let mut right_rv = None;
|
||||
|
||||
for i in lo..peak_idx {
|
||||
if ccf[i].1 <= half_max && ccf[i + 1].1 >= half_max {
|
||||
let t = (half_max - ccf[i].1) / (ccf[i + 1].1 - ccf[i].1);
|
||||
left_rv = Some(ccf[i].0 + t * (ccf[i + 1].0 - ccf[i].0));
|
||||
}
|
||||
}
|
||||
for i in peak_idx..hi - 1 {
|
||||
if ccf[i].1 >= half_max && ccf[i + 1].1 <= half_max {
|
||||
let t = (ccf[i].1 - half_max) / (ccf[i].1 - ccf[i + 1].1);
|
||||
right_rv = Some(ccf[i].0 + t * (ccf[i + 1].0 - ccf[i].0));
|
||||
}
|
||||
}
|
||||
|
||||
match (left_rv, right_rv) {
|
||||
(Some(l), Some(r)) => (r - l) / 2.355, // FWHM → σ (高斯关系)
|
||||
_ => rv_step * 2.0,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_black_body_positive() {
|
||||
let bb = black_body(0.5, 5778.0); // 0.5 μm (可见光)
|
||||
assert!(bb > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_template_generation() {
|
||||
let wave: Vec<f64> = (4000..7000).step_by(10).map(|w| w as f64).collect();
|
||||
let flux = generate_template(TemplateType::G, &wave);
|
||||
assert_eq!(flux.len(), wave.len());
|
||||
assert!(flux.iter().all(|&f| f > 0.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_log_resample() {
|
||||
let wave = vec![4000.0, 5000.0, 6000.0, 7000.0];
|
||||
let resampled = log_resample(&wave, 100);
|
||||
assert_eq!(resampled.len(), 100);
|
||||
// 允许浮点精度误差
|
||||
assert!(
|
||||
(resampled[0] - 4000.0).abs() < 1.0,
|
||||
"first point: {}",
|
||||
resampled[0]
|
||||
);
|
||||
assert!(
|
||||
(resampled[99] - 7000.0).abs() < 1.0,
|
||||
"last point: {}",
|
||||
resampled[99]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_measure_rv_shifted() {
|
||||
// 1. 生成静止坐标系下的波长网格 (4000 to 7000 Å)
|
||||
let wave_rest: Vec<f64> = (4000..7000).step_by(2).map(|w| w as f64).collect();
|
||||
|
||||
// 2. 使用 G 模板生成通量
|
||||
let flux_rest = generate_template(TemplateType::G, &wave_rest);
|
||||
|
||||
// 3. 施加已知红移 (例如 rv = 50.0 km/s)
|
||||
let c = 299792.458; // km/s
|
||||
let rv_input = 50.0;
|
||||
let wave_obs: Vec<f64> = wave_rest
|
||||
.iter()
|
||||
.map(|&w| w * (1.0 + rv_input / c))
|
||||
.collect();
|
||||
|
||||
// 4. 使用 measure_rv 测量
|
||||
let result = measure_rv(&wave_obs, &flux_rest, TemplateType::G, (0.0, 100.0), 0.5).unwrap();
|
||||
println!("TEST RESULT: {:?}", result);
|
||||
|
||||
// 5. 校验测量出的 RV 是否在预期范围内（步长 0.5，误差应极小）
|
||||
assert!(
|
||||
(result.rv - rv_input).abs() < 3.0,
|
||||
"Measured RV = {}, expected = {}",
|
||||
result.rv,
|
||||
rv_input
|
||||
);
|
||||
assert!(
|
||||
result.ccf_peak > 0.9,
|
||||
"CCF peak should be very high, got {}",
|
||||
result.ccf_peak
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,629 @@
|
||||
// src/services/spectrum/lines.rs
|
||||
//
|
||||
// 谱线识别与等值宽度（EW）/ 半高全宽（FWHM）测量
|
||||
//
|
||||
// 功能：
|
||||
// 1. 内置常见天体物理谱线波长表（真空/空气波长）
|
||||
// 2. 在给定波长窗口内搜索局部极值（发射线找极大，吸收线找极小）
|
||||
// 3. 数值积分 EW 计算（梯形法）
|
||||
// 4. FWHM 估算
|
||||
// 5. 支持用户自定义谱线列表
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// 谱线类型
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub enum LineType {
|
||||
/// 吸收线
|
||||
Absorption,
|
||||
/// 发射线
|
||||
Emission,
|
||||
}
|
||||
|
||||
/// 单条谱线定义
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SpectralLine {
|
||||
/// 谱线名称（如 "Hα", "Ca II K", "Na I D"）
|
||||
pub name: String,
|
||||
/// 真空波长 (Å)
|
||||
pub vacuum_wavelength: f64,
|
||||
/// 空气波长 (Å)
|
||||
pub air_wavelength: f64,
|
||||
/// 默认线型
|
||||
pub line_type: LineType,
|
||||
/// 元素符号
|
||||
pub element: String,
|
||||
/// 离化态
|
||||
pub species: String,
|
||||
/// 相对强度（0-1）
|
||||
pub relative_strength: f64,
|
||||
}
|
||||
|
||||
/// 测量结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LineMeasurement {
|
||||
/// 谱线定义
|
||||
pub line: SpectralLine,
|
||||
/// 测量窗口中心波长 (Å)
|
||||
pub center_wavelength: f64,
|
||||
/// 等值宽度 (Å)，吸收线为正，发射线为负
|
||||
pub ew: f64,
|
||||
/// 等值宽度不确定度 (Å)
|
||||
pub ew_error: Option<f64>,
|
||||
/// 半高全宽 (Å)
|
||||
pub fwhm: Option<f64>,
|
||||
/// 峰值连续谱通量
|
||||
pub continuum_flux: f64,
|
||||
/// 峰值通量
|
||||
pub peak_flux: f64,
|
||||
}
|
||||
|
||||
/// 内置谱线数据库（常用天体物理谱线）
|
||||
pub fn builtin_line_table() -> Vec<SpectralLine> {
|
||||
vec![
|
||||
// Balmer 系列
|
||||
SpectralLine {
|
||||
name: "Hα".to_string(),
|
||||
vacuum_wavelength: 6564.61,
|
||||
air_wavelength: 6562.80,
|
||||
line_type: LineType::Absorption,
|
||||
element: "H".to_string(),
|
||||
species: "H I".to_string(),
|
||||
relative_strength: 1.0,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Hβ".to_string(),
|
||||
vacuum_wavelength: 4862.68,
|
||||
air_wavelength: 4861.33,
|
||||
line_type: LineType::Absorption,
|
||||
element: "H".to_string(),
|
||||
species: "H I".to_string(),
|
||||
relative_strength: 0.8,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Hγ".to_string(),
|
||||
vacuum_wavelength: 4341.68,
|
||||
air_wavelength: 4340.47,
|
||||
line_type: LineType::Absorption,
|
||||
element: "H".to_string(),
|
||||
species: "H I".to_string(),
|
||||
relative_strength: 0.6,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Hδ".to_string(),
|
||||
vacuum_wavelength: 4102.89,
|
||||
air_wavelength: 4101.74,
|
||||
line_type: LineType::Absorption,
|
||||
element: "H".to_string(),
|
||||
species: "H I".to_string(),
|
||||
relative_strength: 0.5,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Hε".to_string(),
|
||||
vacuum_wavelength: 3971.20,
|
||||
air_wavelength: 3970.07,
|
||||
line_type: LineType::Absorption,
|
||||
element: "H".to_string(),
|
||||
species: "H I".to_string(),
|
||||
relative_strength: 0.4,
|
||||
},
|
||||
// Ca II H&K
|
||||
SpectralLine {
|
||||
name: "Ca II K".to_string(),
|
||||
vacuum_wavelength: 3934.77,
|
||||
air_wavelength: 3933.66,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Ca".to_string(),
|
||||
species: "Ca II".to_string(),
|
||||
relative_strength: 0.9,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Ca II H".to_string(),
|
||||
vacuum_wavelength: 3969.47,
|
||||
air_wavelength: 3968.47,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Ca".to_string(),
|
||||
species: "Ca II".to_string(),
|
||||
relative_strength: 0.85,
|
||||
},
|
||||
// Na I D 双线
|
||||
SpectralLine {
|
||||
name: "Na I D2".to_string(),
|
||||
vacuum_wavelength: 5897.56,
|
||||
air_wavelength: 5889.95,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Na".to_string(),
|
||||
species: "Na I".to_string(),
|
||||
relative_strength: 0.7,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Na I D1".to_string(),
|
||||
vacuum_wavelength: 5901.48,
|
||||
air_wavelength: 5895.92,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Na".to_string(),
|
||||
species: "Na I".to_string(),
|
||||
relative_strength: 0.65,
|
||||
},
|
||||
// Mg b 三重线
|
||||
SpectralLine {
|
||||
name: "Mg I b1".to_string(),
|
||||
vacuum_wavelength: 5185.54,
|
||||
air_wavelength: 5183.62,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Mg".to_string(),
|
||||
species: "Mg I".to_string(),
|
||||
relative_strength: 0.8,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Mg I b2".to_string(),
|
||||
vacuum_wavelength: 5174.11,
|
||||
air_wavelength: 5172.70,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Mg".to_string(),
|
||||
species: "Mg I".to_string(),
|
||||
relative_strength: 0.75,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Mg I b3".to_string(),
|
||||
vacuum_wavelength: 5168.76,
|
||||
air_wavelength: 5167.33,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Mg".to_string(),
|
||||
species: "Mg I".to_string(),
|
||||
relative_strength: 0.7,
|
||||
},
|
||||
// Fe I 线
|
||||
SpectralLine {
|
||||
name: "Fe I λ5270".to_string(),
|
||||
vacuum_wavelength: 5271.36,
|
||||
air_wavelength: 5269.54,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Fe".to_string(),
|
||||
species: "Fe I".to_string(),
|
||||
relative_strength: 0.5,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Fe I λ5328".to_string(),
|
||||
vacuum_wavelength: 5329.11,
|
||||
air_wavelength: 5327.27,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Fe".to_string(),
|
||||
species: "Fe I".to_string(),
|
||||
relative_strength: 0.45,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Fe I λ5371".to_string(),
|
||||
vacuum_wavelength: 5372.12,
|
||||
air_wavelength: 5370.27,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Fe".to_string(),
|
||||
species: "Fe I".to_string(),
|
||||
relative_strength: 0.55,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "Fe I λ4383".to_string(),
|
||||
vacuum_wavelength: 4384.31,
|
||||
air_wavelength: 4383.55,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Fe".to_string(),
|
||||
species: "Fe I".to_string(),
|
||||
relative_strength: 0.5,
|
||||
},
|
||||
// Ti II
|
||||
SpectralLine {
|
||||
name: "Ti II λ4444".to_string(),
|
||||
vacuum_wavelength: 4445.14,
|
||||
air_wavelength: 4443.80,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Ti".to_string(),
|
||||
species: "Ti II".to_string(),
|
||||
relative_strength: 0.4,
|
||||
},
|
||||
// Ca I
|
||||
SpectralLine {
|
||||
name: "Ca I λ4227".to_string(),
|
||||
vacuum_wavelength: 4228.00,
|
||||
air_wavelength: 4226.73,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Ca".to_string(),
|
||||
species: "Ca I".to_string(),
|
||||
relative_strength: 0.6,
|
||||
},
|
||||
// Hε 与 Ca II H 重叠区域常用线
|
||||
SpectralLine {
|
||||
name: "Ca II K (air)".to_string(),
|
||||
vacuum_wavelength: 3934.77,
|
||||
air_wavelength: 3933.66,
|
||||
line_type: LineType::Absorption,
|
||||
element: "Ca".to_string(),
|
||||
species: "Ca II".to_string(),
|
||||
relative_strength: 0.95,
|
||||
},
|
||||
// O III（星云发射线）
|
||||
SpectralLine {
|
||||
name: "[O III] λ5007".to_string(),
|
||||
vacuum_wavelength: 5008.24,
|
||||
air_wavelength: 5007.05,
|
||||
line_type: LineType::Emission,
|
||||
element: "O".to_string(),
|
||||
species: "O III".to_string(),
|
||||
relative_strength: 1.0,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "[O III] λ4959".to_string(),
|
||||
vacuum_wavelength: 4960.29,
|
||||
air_wavelength: 4959.12,
|
||||
line_type: LineType::Emission,
|
||||
element: "O".to_string(),
|
||||
species: "O III".to_string(),
|
||||
relative_strength: 0.33,
|
||||
},
|
||||
// N II
|
||||
SpectralLine {
|
||||
name: "[N II] λ6584".to_string(),
|
||||
vacuum_wavelength: 6585.27,
|
||||
air_wavelength: 6583.41,
|
||||
line_type: LineType::Emission,
|
||||
element: "N".to_string(),
|
||||
species: "N II".to_string(),
|
||||
relative_strength: 0.7,
|
||||
},
|
||||
// S II
|
||||
SpectralLine {
|
||||
name: "[S II] λ6717".to_string(),
|
||||
vacuum_wavelength: 6718.29,
|
||||
air_wavelength: 6716.44,
|
||||
line_type: LineType::Emission,
|
||||
element: "S".to_string(),
|
||||
species: "S II".to_string(),
|
||||
relative_strength: 0.5,
|
||||
},
|
||||
SpectralLine {
|
||||
name: "[S II] λ6731".to_string(),
|
||||
vacuum_wavelength: 6732.68,
|
||||
air_wavelength: 6730.82,
|
||||
line_type: LineType::Emission,
|
||||
element: "S".to_string(),
|
||||
species: "S II".to_string(),
|
||||
relative_strength: 0.45,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
/// 按名称查找谱线
|
||||
pub fn find_line<'a>(name: &str, table: &'a [SpectralLine]) -> Option<&'a SpectralLine> {
|
||||
table.iter().find(|l| l.name == name)
|
||||
}
|
||||
|
||||
/// 按波长范围查找谱线
|
||||
pub fn find_lines_in_range(
|
||||
min_wave: f64,
|
||||
max_wave: f64,
|
||||
table: &[SpectralLine],
|
||||
) -> Vec<&SpectralLine> {
|
||||
table
|
||||
.iter()
|
||||
.filter(|l| l.air_wavelength >= min_wave && l.air_wavelength <= max_wave)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 按元素查找谱线
|
||||
pub fn find_lines_by_element<'a>(
|
||||
element: &str,
|
||||
table: &'a [SpectralLine],
|
||||
) -> Vec<&'a SpectralLine> {
|
||||
table
|
||||
.iter()
|
||||
.filter(|l| l.element.eq_ignore_ascii_case(element))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 在光谱中搜索谱线并测量 EW/FWHM
|
||||
///
|
||||
/// 输入：波长数组 (Å)、通量数组、谱线定义、搜索窗口半宽 (Å)
|
||||
/// 输出：测量结果
|
||||
pub fn measure_line(
|
||||
wavelength: &[f32],
|
||||
flux: &[f32],
|
||||
line: &SpectralLine,
|
||||
window_half_width: f64,
|
||||
) -> Result<LineMeasurement> {
|
||||
if wavelength.len() != flux.len() {
|
||||
return Err(anyhow!("波长和通量数组长度不一致"));
|
||||
}
|
||||
if wavelength.len() < 3 {
|
||||
return Err(anyhow!("光谱数据点太少（需要 ≥3）"));
|
||||
}
|
||||
|
||||
let center = line.air_wavelength;
|
||||
let lo = center - window_half_width;
|
||||
let hi = center + window_half_width;
|
||||
|
||||
// 提取窗口内的数据点
|
||||
let mut window_waves = Vec::new();
|
||||
let mut window_fluxes = Vec::new();
|
||||
for (w, f) in wavelength.iter().zip(flux.iter()) {
|
||||
let w64 = *w as f64;
|
||||
if w64 >= lo && w64 <= hi {
|
||||
window_waves.push(w64);
|
||||
window_fluxes.push(*f as f64);
|
||||
}
|
||||
}
|
||||
|
||||
if window_waves.len() < 3 {
|
||||
return Err(anyhow!(
|
||||
"窗口 [{:.1}, {:.1}] Å 内数据点不足（{} 点）",
|
||||
lo,
|
||||
hi,
|
||||
window_waves.len()
|
||||
));
|
||||
}
|
||||
|
||||
// 估计连续谱：窗口边缘两侧各取 20% 的中位数
|
||||
let n = window_waves.len();
|
||||
let edge = (n as f64 * 0.2).max(1.0) as usize;
|
||||
let continuum_flux = {
|
||||
let mut edge_fluxes: Vec<f64> = window_fluxes[..edge]
|
||||
.iter()
|
||||
.chain(window_fluxes[n - edge..].iter())
|
||||
.copied()
|
||||
.collect();
|
||||
edge_fluxes.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
edge_fluxes[edge_fluxes.len() / 2]
|
||||
};
|
||||
|
||||
// 找峰值（吸收线找最小值，发射线找最大值）
|
||||
let (peak_idx, peak_flux_ref) = match line.line_type {
|
||||
LineType::Absorption => window_fluxes
|
||||
.iter()
|
||||
.enumerate()
|
||||
.min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.unwrap(),
|
||||
LineType::Emission => window_fluxes
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.unwrap(),
|
||||
};
|
||||
let peak_flux = *peak_flux_ref;
|
||||
|
||||
// 计算 EW（梯形法数值积分）
|
||||
// 吸收线：EW = ∫(1 - F(λ)/F_c) dλ，结果为正
|
||||
// 发射线：EW = ∫(F(λ)/F_c - 1) dλ，结果为负
|
||||
let ew = if continuum_flux > 0.0 {
|
||||
let mut integral = 0.0;
|
||||
for i in 0..window_waves.len() - 1 {
|
||||
let dw = window_waves[i + 1] - window_waves[i];
|
||||
let norm1 = window_fluxes[i] / continuum_flux;
|
||||
let norm2 = window_fluxes[i + 1] / continuum_flux;
|
||||
let avg = match line.line_type {
|
||||
LineType::Absorption => (1.0 - norm1 + 1.0 - norm2) / 2.0,
|
||||
LineType::Emission => (norm1 - 1.0 + norm2 - 1.0) / 2.0,
|
||||
};
|
||||
integral += avg * dw;
|
||||
}
|
||||
integral
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
// 估算 EW 不确定度（基于连续谱散射的简化估计）
|
||||
let continuum_scatter = {
|
||||
let mean = continuum_flux;
|
||||
let var: f64 = window_fluxes
|
||||
.iter()
|
||||
.map(|f| (f - mean).powi(2))
|
||||
.sum::<f64>()
|
||||
/ n as f64;
|
||||
var.sqrt()
|
||||
};
|
||||
let ew_error = if continuum_scatter > 0.0 && continuum_flux > 0.0 {
|
||||
Some((continuum_scatter / continuum_flux * window_half_width * 2.0).sqrt())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// FWHM 估算：找半高点并插值
|
||||
let fwhm = estimate_fwhm(
|
||||
&window_waves,
|
||||
&window_fluxes,
|
||||
continuum_flux,
|
||||
line.line_type,
|
||||
);
|
||||
|
||||
Ok(LineMeasurement {
|
||||
line: line.clone(),
|
||||
center_wavelength: window_waves[peak_idx],
|
||||
ew,
|
||||
ew_error,
|
||||
fwhm,
|
||||
continuum_flux,
|
||||
peak_flux,
|
||||
})
|
||||
}
|
||||
|
||||
/// 估算 FWHM（半高全宽）
|
||||
fn estimate_fwhm(
|
||||
wavelength: &[f64],
|
||||
flux: &[f64],
|
||||
continuum: f64,
|
||||
line_type: LineType,
|
||||
) -> Option<f64> {
|
||||
if flux.len() < 3 || continuum <= 0.0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let peak = match line_type {
|
||||
LineType::Absorption => flux.iter().cloned().fold(f64::INFINITY, f64::min),
|
||||
LineType::Emission => flux.iter().cloned().fold(f64::NEG_INFINITY, f64::max),
|
||||
};
|
||||
|
||||
// 半高点
|
||||
let half_level = match line_type {
|
||||
LineType::Absorption => (continuum + peak) / 2.0,
|
||||
LineType::Emission => (continuum + peak) / 2.0,
|
||||
};
|
||||
|
||||
// 找左侧和右侧的半高点
|
||||
let mut left = None;
|
||||
let mut right = None;
|
||||
|
||||
let peak_idx = match line_type {
|
||||
LineType::Absorption => flux
|
||||
.iter()
|
||||
.enumerate()
|
||||
.min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.map(|(i, _)| i),
|
||||
LineType::Emission => flux
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.map(|(i, _)| i),
|
||||
};
|
||||
|
||||
let peak_idx = peak_idx?;
|
||||
|
||||
// 左侧
|
||||
for i in (0..peak_idx).rev() {
|
||||
if match line_type {
|
||||
LineType::Absorption => flux[i] >= half_level,
|
||||
LineType::Emission => flux[i] <= half_level,
|
||||
} {
|
||||
// 线性插值
|
||||
if i + 1 < flux.len() {
|
||||
let w1 = wavelength[i];
|
||||
let w2 = wavelength[i + 1];
|
||||
let f1 = flux[i];
|
||||
let f2 = flux[i + 1];
|
||||
let dw = w2 - w1;
|
||||
let df = f2 - f1;
|
||||
if df.abs() > 1e-10 {
|
||||
left = Some(w1 + (half_level - f1) / df * dw);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 右侧
|
||||
for i in peak_idx..flux.len() - 1 {
|
||||
if match line_type {
|
||||
LineType::Absorption => flux[i] >= half_level,
|
||||
LineType::Emission => flux[i] <= half_level,
|
||||
} {
|
||||
let w1 = wavelength[i];
|
||||
let w2 = wavelength[i + 1];
|
||||
let f1 = flux[i];
|
||||
let f2 = flux[i + 1];
|
||||
let dw = w2 - w1;
|
||||
let df = f2 - f1;
|
||||
if df.abs() > 1e-10 {
|
||||
right = Some(w1 + (half_level - f1) / df * dw);
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
match (left, right) {
|
||||
(Some(l), Some(r)) => Some(r - l),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 批量测量多条谱线
|
||||
pub fn measure_lines(
|
||||
wavelength: &[f32],
|
||||
flux: &[f32],
|
||||
lines: &[SpectralLine],
|
||||
window_half_width: f64,
|
||||
) -> Vec<Result<LineMeasurement>> {
|
||||
lines
|
||||
.iter()
|
||||
.map(|line| measure_line(wavelength, flux, line, window_half_width))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 将测量结果格式化为 Markdown 表格
|
||||
pub fn format_measurements_table(measurements: &[LineMeasurement]) -> String {
|
||||
let mut table =
|
||||
String::from("| 谱线 | 波长 (Å) | EW (Å) | EW err (Å) | FWHM (Å) | F_c | F_peak |\n");
|
||||
table.push_str("|------|----------|--------|------------|----------|-----|--------|\n");
|
||||
for m in measurements {
|
||||
let ew_str = format!("{:.3}", m.ew);
|
||||
let ew_err_str = m
|
||||
.ew_error
|
||||
.map(|e| format!("{:.3}", e))
|
||||
.unwrap_or_else(|| "—".to_string());
|
||||
let fwhm_str = m
|
||||
.fwhm
|
||||
.map(|f| format!("{:.2}", f))
|
||||
.unwrap_or_else(|| "—".to_string());
|
||||
table.push_str(&format!(
|
||||
"| {} | {:.2} | {} | {} | {} | {:.1} | {:.1} |\n",
|
||||
m.line.name,
|
||||
m.center_wavelength,
|
||||
ew_str,
|
||||
ew_err_str,
|
||||
fwhm_str,
|
||||
m.continuum_flux,
|
||||
m.peak_flux
|
||||
));
|
||||
}
|
||||
table
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_builtin_lines_not_empty() {
|
||||
let table = builtin_line_table();
|
||||
assert!(!table.is_empty());
|
||||
// 应该包含 Hα
|
||||
assert!(find_line("Hα", &table).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_lines_in_range() {
|
||||
let table = builtin_line_table();
|
||||
let lines = find_lines_in_range(6560.0, 6570.0, &table);
|
||||
assert!(!lines.is_empty());
|
||||
assert_eq!(lines[0].name, "Hα");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_measure_absorption_line() {
|
||||
// 构造一个简单的模拟光谱：连续谱 1.0，在 5000 Å 有一个高斯吸收线
|
||||
let mut wavelength = Vec::new();
|
||||
let mut flux = Vec::new();
|
||||
for i in 0..200 {
|
||||
let w = 4900.0 + i as f64 * 1.0; // 4900-5099 Å
|
||||
let absorption = 0.3 * (-(w - 5000.0).powi(2) / 10.0).exp();
|
||||
wavelength.push(w as f32);
|
||||
flux.push((1.0 - absorption) as f32);
|
||||
}
|
||||
|
||||
let line = SpectralLine {
|
||||
name: "test".to_string(),
|
||||
vacuum_wavelength: 5000.0,
|
||||
air_wavelength: 5000.0,
|
||||
line_type: LineType::Absorption,
|
||||
element: "X".to_string(),
|
||||
species: "X I".to_string(),
|
||||
relative_strength: 1.0,
|
||||
};
|
||||
|
||||
let result = measure_line(&wavelength, &flux, &line, 30.0).unwrap();
|
||||
// EW 应该为正（吸收线）
|
||||
assert!(result.ew > 0.0, "吸收线 EW 应为正: {}", result.ew);
|
||||
// 连续谱应接近 1.0
|
||||
assert!(
|
||||
(result.continuum_flux - 1.0).abs() < 0.1,
|
||||
"连续谱应接近 1.0: {}",
|
||||
result.continuum_flux
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
// src/services/spectrum/mod.rs
|
||||
//
|
||||
// 光谱分析服务层 —— 参数提取、谱线测量、交叉相关测速
|
||||
//
|
||||
// 模块布局：
|
||||
// parameters.rs —— FITS Header/BinTable 恒星大气参数提取（Teff, logg, [Fe/H], RV...）
|
||||
// lines.rs —— 谱线识别 + 等值宽度（EW）/ 半高全宽（FWHM）测量
|
||||
// cross_correlate.rs —— 模板匹配法交叉相关测径向速度（RV）
|
||||
|
||||
pub mod cross_correlate;
|
||||
pub mod lines;
|
||||
pub mod parameters;
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user