feat: 拆分 worker 架构并增强 assistant 安全保护
- 新增 daemon/channel/assistant/function/scheduler/logging worker 架构和 JetStream 消息协议 - 修复邮件通知误入 assistant、迁移 checksum、日志 subject 和工具调用无限循环风险 - 新增 assistant 工具轮次熔断、工具结果截断和提示词软收敛规则 - 补充邮件推送关键字、可信发件人配置、升级日志和验证测试
This commit is contained in:
@@ -0,0 +1,19 @@
|
||||
[package]
|
||||
name = "ias-assistant"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
ias-ai = { path = "../ias-ai" }
|
||||
ias-common = { path = "../ias-common" }
|
||||
ias-context = { path = "../ias-context" }
|
||||
anyhow = "1"
|
||||
serde_json = "1"
|
||||
sqlx = { version = "0.9.0", features = [
|
||||
"runtime-tokio",
|
||||
"tls-rustls",
|
||||
"postgres",
|
||||
"chrono",
|
||||
] }
|
||||
tokio = { version = "1.0", features = ["rt-multi-thread", "macros", "sync", "time"] }
|
||||
uuid = { version = "1", features = ["v4", "serde"] }
|
||||
@@ -0,0 +1,877 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use ias_ai::core::request_llm;
|
||||
use ias_common::model::ai::ChatCompletionRequestMessage;
|
||||
use ias_common::mq::{
|
||||
ConsumerSpec, FunctionCallPayload, FunctionResultPayload, InboundPayload, JetStreamBus,
|
||||
JetStreamConfig, LogEventPayload, MessageEnvelope, OutboundCommandKind, OutboundCommandPayload,
|
||||
StreamKind, StreamNames, function_call_subject, inbound_subject, log_subject, outbound_subject,
|
||||
};
|
||||
use ias_common::queue::ChannelMessage;
|
||||
use ias_context::{ContextManager, NewSessionReason};
|
||||
use sqlx::PgPool;
|
||||
use tokio::sync::Mutex;
|
||||
use uuid::Uuid;
|
||||
|
||||
type TurnMemory = Arc<Mutex<HashMap<String, String>>>;
|
||||
|
||||
const DEFAULT_MAX_TOOL_ROUNDS: i32 = 8;
|
||||
const DEFAULT_MAX_TOOL_CALLS_PER_ROUND: usize = 5;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TurnRegistry {
|
||||
pool: PgPool,
|
||||
memory: TurnMemory,
|
||||
}
|
||||
|
||||
#[derive(Debug, sqlx::FromRow)]
|
||||
struct TurnRow {
|
||||
turn_id: String,
|
||||
status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, Copy, sqlx::FromRow)]
|
||||
struct ToolBudget {
|
||||
tool_rounds: i32,
|
||||
tool_calls: i32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum TurnState {
|
||||
Active,
|
||||
Missing,
|
||||
Stale,
|
||||
}
|
||||
|
||||
impl TurnRegistry {
|
||||
fn new(pool: PgPool) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
memory: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_turn(&self, channel: &mut ChannelMessage, trace_id: &str) -> anyhow::Result<()> {
|
||||
let turn_id = Uuid::new_v4().to_string();
|
||||
let key = operation_key(channel);
|
||||
channel.turn_id = Some(turn_id.clone());
|
||||
self.memory
|
||||
.lock()
|
||||
.await
|
||||
.insert(key.clone(), turn_id.clone());
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO ias_assistant_turns (
|
||||
operation_key,
|
||||
channel,
|
||||
account_id,
|
||||
from_user_id,
|
||||
turn_id,
|
||||
trace_id,
|
||||
status,
|
||||
tool_rounds,
|
||||
tool_calls,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, 'active', 0, 0, NOW())
|
||||
ON CONFLICT (operation_key) DO UPDATE SET
|
||||
channel = EXCLUDED.channel,
|
||||
account_id = EXCLUDED.account_id,
|
||||
from_user_id = EXCLUDED.from_user_id,
|
||||
turn_id = EXCLUDED.turn_id,
|
||||
trace_id = EXCLUDED.trace_id,
|
||||
status = 'active',
|
||||
tool_rounds = 0,
|
||||
tool_calls = 0,
|
||||
last_tool_at = NULL,
|
||||
updated_at = NOW()
|
||||
"#,
|
||||
)
|
||||
.bind(&key)
|
||||
.bind(&channel.channel)
|
||||
.bind(&channel.account_id)
|
||||
.bind(&channel.from_user_id)
|
||||
.bind(&turn_id)
|
||||
.bind(trace_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cancel_turn(&self, channel: &ChannelMessage, trace_id: &str) -> anyhow::Result<()> {
|
||||
let key = operation_key(channel);
|
||||
self.memory.lock().await.remove(&key);
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO ias_assistant_turns (
|
||||
operation_key,
|
||||
channel,
|
||||
account_id,
|
||||
from_user_id,
|
||||
turn_id,
|
||||
trace_id,
|
||||
status,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, '', $5, 'cancelled', NOW())
|
||||
ON CONFLICT (operation_key) DO UPDATE SET
|
||||
trace_id = EXCLUDED.trace_id,
|
||||
status = 'cancelled',
|
||||
updated_at = NOW()
|
||||
"#,
|
||||
)
|
||||
.bind(&key)
|
||||
.bind(&channel.channel)
|
||||
.bind(&channel.account_id)
|
||||
.bind(&channel.from_user_id)
|
||||
.bind(trace_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn recover_missing_turn(
|
||||
&self,
|
||||
channel: &ChannelMessage,
|
||||
trace_id: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let Some(turn_id) = channel.turn_id.as_deref() else {
|
||||
return Ok(());
|
||||
};
|
||||
let key = operation_key(channel);
|
||||
self.memory
|
||||
.lock()
|
||||
.await
|
||||
.insert(key.clone(), turn_id.to_string());
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO ias_assistant_turns (
|
||||
operation_key,
|
||||
channel,
|
||||
account_id,
|
||||
from_user_id,
|
||||
turn_id,
|
||||
trace_id,
|
||||
status,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, 'active', NOW())
|
||||
ON CONFLICT (operation_key) DO NOTHING
|
||||
"#,
|
||||
)
|
||||
.bind(&key)
|
||||
.bind(&channel.channel)
|
||||
.bind(&channel.account_id)
|
||||
.bind(&channel.from_user_id)
|
||||
.bind(turn_id)
|
||||
.bind(trace_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finish_turn(&self, channel: &ChannelMessage) -> anyhow::Result<()> {
|
||||
let key = operation_key(channel);
|
||||
self.memory.lock().await.remove(&key);
|
||||
let Some(turn_id) = channel.turn_id.as_deref() else {
|
||||
return Ok(());
|
||||
};
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE ias_assistant_turns
|
||||
SET status = 'completed', updated_at = NOW()
|
||||
WHERE operation_key = $1
|
||||
AND turn_id = $2
|
||||
AND status = 'active'
|
||||
"#,
|
||||
)
|
||||
.bind(&key)
|
||||
.bind(turn_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn state(&self, channel: &ChannelMessage) -> anyhow::Result<TurnState> {
|
||||
let Some(turn_id) = channel.turn_id.as_deref() else {
|
||||
return Ok(TurnState::Stale);
|
||||
};
|
||||
let key = operation_key(channel);
|
||||
if let Some(active) = self.memory.lock().await.get(&key).cloned() {
|
||||
return Ok(if active == turn_id {
|
||||
TurnState::Active
|
||||
} else {
|
||||
TurnState::Stale
|
||||
});
|
||||
}
|
||||
|
||||
let row = sqlx::query_as::<_, TurnRow>(
|
||||
r#"
|
||||
SELECT turn_id, status
|
||||
FROM ias_assistant_turns
|
||||
WHERE operation_key = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&key)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
match row {
|
||||
None => Ok(TurnState::Missing),
|
||||
Some(row) if row.turn_id == turn_id && row.status == "active" => {
|
||||
self.memory.lock().await.insert(key, turn_id.to_string());
|
||||
Ok(TurnState::Active)
|
||||
}
|
||||
Some(_) => Ok(TurnState::Stale),
|
||||
}
|
||||
}
|
||||
|
||||
async fn tool_budget(&self, channel: &ChannelMessage) -> anyhow::Result<ToolBudget> {
|
||||
let Some(turn_id) = channel.turn_id.as_deref() else {
|
||||
return Ok(ToolBudget::default());
|
||||
};
|
||||
let key = operation_key(channel);
|
||||
let budget = sqlx::query_as::<_, ToolBudget>(
|
||||
r#"
|
||||
SELECT tool_rounds, tool_calls
|
||||
FROM ias_assistant_turns
|
||||
WHERE operation_key = $1
|
||||
AND turn_id = $2
|
||||
AND status = 'active'
|
||||
"#,
|
||||
)
|
||||
.bind(&key)
|
||||
.bind(turn_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(budget)
|
||||
}
|
||||
|
||||
async fn next_tool_round(
|
||||
&self,
|
||||
channel: &ChannelMessage,
|
||||
tool_call_count: usize,
|
||||
) -> anyhow::Result<i32> {
|
||||
let Some(turn_id) = channel.turn_id.as_deref() else {
|
||||
return Ok(i32::MAX);
|
||||
};
|
||||
let key = operation_key(channel);
|
||||
let count = tool_call_count.min(i32::MAX as usize) as i32;
|
||||
let rounds = sqlx::query_scalar::<_, i32>(
|
||||
r#"
|
||||
UPDATE ias_assistant_turns
|
||||
SET
|
||||
tool_rounds = tool_rounds + 1,
|
||||
tool_calls = tool_calls + $3,
|
||||
last_tool_at = NOW(),
|
||||
updated_at = NOW()
|
||||
WHERE operation_key = $1
|
||||
AND turn_id = $2
|
||||
AND status = 'active'
|
||||
RETURNING tool_rounds
|
||||
"#,
|
||||
)
|
||||
.bind(&key)
|
||||
.bind(turn_id)
|
||||
.bind(count)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
.unwrap_or(i32::MAX);
|
||||
|
||||
Ok(rounds)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_worker(pool: PgPool) -> anyhow::Result<()> {
|
||||
let config = JetStreamConfig::from_env();
|
||||
loop {
|
||||
match run_once(pool.clone(), config.clone()).await {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(error) => {
|
||||
eprintln!("Assistant Worker 连接或消费失败: {error:#}");
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_once(pool: PgPool, config: JetStreamConfig) -> anyhow::Result<()> {
|
||||
let bus = JetStreamBus::connect(config.clone()).await?;
|
||||
bus.ensure_streams().await?;
|
||||
let names = StreamNames::new(config.stream_prefix.clone());
|
||||
let context = ContextManager::new(pool.clone());
|
||||
let turns = TurnRegistry::new(pool);
|
||||
|
||||
let inbound_bus = bus.clone();
|
||||
let inbound_context = context.clone();
|
||||
let inbound_turns = turns.clone();
|
||||
let inbound_spec = ConsumerSpec {
|
||||
stream: names.name(StreamKind::Inbound),
|
||||
durable_name: "assistant-inbound".to_string(),
|
||||
filter_subject: "inbound.>".to_string(),
|
||||
ack_wait: Duration::from_secs(config.ack_wait_secs.max(180)),
|
||||
max_deliver: config.max_deliver,
|
||||
};
|
||||
|
||||
let result_bus = bus.clone();
|
||||
let result_context = context.clone();
|
||||
let result_turns = turns.clone();
|
||||
let result_spec = ConsumerSpec {
|
||||
stream: names.name(StreamKind::Function),
|
||||
durable_name: "assistant-function-results".to_string(),
|
||||
filter_subject: "function.result.>".to_string(),
|
||||
ack_wait: Duration::from_secs(config.ack_wait_secs.max(180)),
|
||||
max_deliver: config.max_deliver,
|
||||
};
|
||||
|
||||
let inbound_consumer_bus = bus.clone();
|
||||
let result_consumer_bus = bus.clone();
|
||||
let inbound_future = inbound_consumer_bus.consume_envelopes::<InboundPayload, _, _>(
|
||||
inbound_spec,
|
||||
move |envelope| {
|
||||
let bus = inbound_bus.clone();
|
||||
let context = inbound_context.clone();
|
||||
let active_turns = inbound_turns.clone();
|
||||
async move { handle_inbound(bus, context, active_turns, envelope).await }
|
||||
},
|
||||
);
|
||||
let result_publish_bus = result_bus.clone();
|
||||
let result_future = result_consumer_bus.consume_envelopes::<FunctionResultPayload, _, _>(
|
||||
result_spec,
|
||||
move |envelope| {
|
||||
let bus = result_publish_bus.clone();
|
||||
let context = result_context.clone();
|
||||
let active_turns = result_turns.clone();
|
||||
async move { handle_function_result(bus, context, active_turns, envelope).await }
|
||||
},
|
||||
);
|
||||
tokio::try_join!(inbound_future, result_future)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_inbound(
|
||||
bus: JetStreamBus,
|
||||
context: ContextManager,
|
||||
turns: TurnRegistry,
|
||||
envelope: MessageEnvelope<InboundPayload>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut channel = envelope.payload.into_channel_message();
|
||||
let trace_id = envelope.trace_id;
|
||||
|
||||
if channel.content.trim().eq_ignore_ascii_case("/new") {
|
||||
turns.cancel_turn(&channel, &trace_id).await?;
|
||||
context
|
||||
.start_new_session(&channel, NewSessionReason::UserCommand)
|
||||
.await?;
|
||||
publish_typing_stop(&bus, &channel, &trace_id).await?;
|
||||
publish_text(&bus, &channel, "已开启新的上下文。", &trace_id).await?;
|
||||
publish_log(
|
||||
&bus,
|
||||
"info",
|
||||
"worker.assistant",
|
||||
"用户通过 /new 开启新上下文",
|
||||
&trace_id,
|
||||
serde_json::json!({
|
||||
"channel": channel.channel,
|
||||
"account_id": channel.account_id,
|
||||
"from_user_id": channel.from_user_id,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
turns.start_turn(&mut channel, &trace_id).await?;
|
||||
publish_typing_start(&bus, &channel, &trace_id).await?;
|
||||
|
||||
let messages = match context.record_user_message(&channel).await {
|
||||
Ok(snapshot) => snapshot.messages,
|
||||
Err(error) => {
|
||||
publish_log(
|
||||
&bus,
|
||||
"error",
|
||||
"worker.assistant",
|
||||
"记录用户上下文失败，退回单轮请求",
|
||||
&trace_id,
|
||||
serde_json::json!({"error": error.to_string()}),
|
||||
)
|
||||
.await;
|
||||
vec![ChatCompletionRequestMessage::new(
|
||||
"user".to_string(),
|
||||
channel.content.clone(),
|
||||
)]
|
||||
}
|
||||
};
|
||||
|
||||
request_and_route_llm(bus, context, turns, channel, messages, trace_id).await
|
||||
}
|
||||
|
||||
async fn handle_function_result(
|
||||
bus: JetStreamBus,
|
||||
context: ContextManager,
|
||||
turns: TurnRegistry,
|
||||
envelope: MessageEnvelope<FunctionResultPayload>,
|
||||
) -> anyhow::Result<()> {
|
||||
let result = envelope.payload;
|
||||
match turns.state(&result.channel).await? {
|
||||
TurnState::Active => {}
|
||||
TurnState::Missing => {
|
||||
turns
|
||||
.recover_missing_turn(&result.channel, &envelope.trace_id)
|
||||
.await?;
|
||||
publish_log(
|
||||
&bus,
|
||||
"warn",
|
||||
"worker.assistant",
|
||||
"恢复缺失的 active turn 后继续处理工具结果",
|
||||
&envelope.trace_id,
|
||||
serde_json::json!({"call_id": result.call_id, "tool": result.name}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
TurnState::Stale => {
|
||||
publish_log(
|
||||
&bus,
|
||||
"warn",
|
||||
"worker.assistant",
|
||||
"跳过已被打断的工具结果",
|
||||
&envelope.trace_id,
|
||||
serde_json::json!({"call_id": result.call_id, "tool": result.name}),
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
publish_typing_refresh(&bus, &result.channel, &envelope.trace_id).await?;
|
||||
let snapshot = context
|
||||
.record_tool_message(
|
||||
&result.channel,
|
||||
result.call_id.clone(),
|
||||
result.name.clone(),
|
||||
result.result.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
request_and_route_llm(
|
||||
bus,
|
||||
context,
|
||||
turns,
|
||||
result.channel,
|
||||
snapshot.messages,
|
||||
envelope.trace_id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn request_and_route_llm(
|
||||
bus: JetStreamBus,
|
||||
context: ContextManager,
|
||||
turns: TurnRegistry,
|
||||
channel: ChannelMessage,
|
||||
messages: Vec<ChatCompletionRequestMessage>,
|
||||
trace_id: String,
|
||||
) -> anyhow::Result<()> {
|
||||
match turns.state(&channel).await? {
|
||||
TurnState::Active => {}
|
||||
TurnState::Missing => {
|
||||
turns.recover_missing_turn(&channel, &trace_id).await?;
|
||||
}
|
||||
TurnState::Stale => return Ok(()),
|
||||
}
|
||||
|
||||
let budget = turns.tool_budget(&channel).await?;
|
||||
let messages = with_tool_budget_message(messages, budget);
|
||||
let Some(response) = request_llm(messages).await? else {
|
||||
publish_typing_stop(&bus, &channel, &trace_id).await?;
|
||||
turns.finish_turn(&channel).await?;
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
for choice in response.choices {
|
||||
if turns.state(&channel).await? == TurnState::Stale {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let raw_message = serde_json::to_value(&choice.message).ok();
|
||||
let content = choice.message.content.unwrap_or_default();
|
||||
let reasoning = choice.message.reasoning_content.unwrap_or_default();
|
||||
let tool_calls = choice.message.tool_calls.clone();
|
||||
let has_tool_calls = tool_calls.as_ref().is_some_and(|calls| !calls.is_empty());
|
||||
|
||||
context
|
||||
.record_assistant_message(
|
||||
&channel,
|
||||
content.clone(),
|
||||
tool_calls.clone(),
|
||||
Some(reasoning.clone()),
|
||||
raw_message,
|
||||
)
|
||||
.await?;
|
||||
|
||||
if has_tool_calls {
|
||||
let tool_calls = tool_calls.unwrap_or_default();
|
||||
let max_rounds = max_tool_rounds();
|
||||
let max_calls = max_tool_calls_per_round();
|
||||
let tool_round = turns.next_tool_round(&channel, tool_calls.len()).await?;
|
||||
if tool_round > max_rounds {
|
||||
publish_log(
|
||||
&bus,
|
||||
"error",
|
||||
"worker.assistant",
|
||||
"工具调用轮次超过上限，已熔断当前 turn",
|
||||
&trace_id,
|
||||
serde_json::json!({
|
||||
"channel": channel.channel,
|
||||
"account_id": channel.account_id,
|
||||
"from_user_id": channel.from_user_id,
|
||||
"turn_id": channel.turn_id,
|
||||
"tool_round": tool_round,
|
||||
"max_tool_rounds": max_rounds,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
publish_typing_stop(&bus, &channel, &trace_id).await?;
|
||||
publish_text(
|
||||
&bus,
|
||||
&channel,
|
||||
"工具调用次数异常偏高，已自动停止本轮任务以避免继续消耗额度。请缩小查询范围后重试。",
|
||||
&trace_id,
|
||||
)
|
||||
.await?;
|
||||
turns.finish_turn(&channel).await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if tool_calls.len() > max_calls {
|
||||
publish_log(
|
||||
&bus,
|
||||
"error",
|
||||
"worker.assistant",
|
||||
"单轮工具调用数量超过上限，已熔断当前 turn",
|
||||
&trace_id,
|
||||
serde_json::json!({
|
||||
"channel": channel.channel,
|
||||
"account_id": channel.account_id,
|
||||
"from_user_id": channel.from_user_id,
|
||||
"turn_id": channel.turn_id,
|
||||
"tool_call_count": tool_calls.len(),
|
||||
"max_tool_calls_per_round": max_calls,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
publish_typing_stop(&bus, &channel, &trace_id).await?;
|
||||
publish_text(
|
||||
&bus,
|
||||
&channel,
|
||||
"本轮请求生成了过多工具调用，已自动停止以避免继续消耗额度。请缩小查询范围后重试。",
|
||||
&trace_id,
|
||||
)
|
||||
.await?;
|
||||
turns.finish_turn(&channel).await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
publish_typing_refresh(&bus, &channel, &trace_id).await?;
|
||||
for tool_call in tool_calls {
|
||||
let Some(function) = tool_call.function else {
|
||||
continue;
|
||||
};
|
||||
let payload = FunctionCallPayload {
|
||||
call_id: tool_call.id,
|
||||
name: function.name,
|
||||
arguments: function.arguments,
|
||||
channel: channel.clone(),
|
||||
};
|
||||
let subject = function_call_subject(&payload.name);
|
||||
let envelope = MessageEnvelope::new(
|
||||
subject.as_str(),
|
||||
format!("function-call/{}", payload.call_id),
|
||||
trace_id.clone(),
|
||||
payload,
|
||||
);
|
||||
bus.publish_envelope(subject.as_str(), &envelope).await?;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let text = if !content.trim().is_empty() {
|
||||
content.trim()
|
||||
} else {
|
||||
reasoning.trim()
|
||||
};
|
||||
if !text.is_empty() {
|
||||
publish_text(&bus, &channel, text, &trace_id).await?;
|
||||
}
|
||||
publish_typing_stop(&bus, &channel, &trace_id).await?;
|
||||
turns.finish_turn(&channel).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn publish_text(
|
||||
bus: &JetStreamBus,
|
||||
channel: &ChannelMessage,
|
||||
text: &str,
|
||||
trace_id: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let payload = OutboundCommandPayload::text(
|
||||
channel.channel.clone(),
|
||||
channel.account_id.clone(),
|
||||
channel.from_user_id.clone(),
|
||||
text.to_string(),
|
||||
channel.context.clone(),
|
||||
channel.turn_id.clone(),
|
||||
);
|
||||
publish_outbound(bus, payload, trace_id).await
|
||||
}
|
||||
|
||||
async fn publish_typing_start(
|
||||
bus: &JetStreamBus,
|
||||
channel: &ChannelMessage,
|
||||
trace_id: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let payload = OutboundCommandPayload::typing(
|
||||
OutboundCommandKind::TypingStart,
|
||||
channel.channel.clone(),
|
||||
channel.account_id.clone(),
|
||||
channel.from_user_id.clone(),
|
||||
channel.context.clone(),
|
||||
None,
|
||||
channel.turn_id.clone(),
|
||||
);
|
||||
publish_outbound(bus, payload, trace_id).await
|
||||
}
|
||||
|
||||
async fn publish_typing_refresh(
|
||||
bus: &JetStreamBus,
|
||||
channel: &ChannelMessage,
|
||||
trace_id: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let payload = OutboundCommandPayload::typing(
|
||||
OutboundCommandKind::TypingRefresh,
|
||||
channel.channel.clone(),
|
||||
channel.account_id.clone(),
|
||||
channel.from_user_id.clone(),
|
||||
channel.context.clone(),
|
||||
None,
|
||||
channel.turn_id.clone(),
|
||||
);
|
||||
publish_outbound(bus, payload, trace_id).await
|
||||
}
|
||||
|
||||
async fn publish_typing_stop(
|
||||
bus: &JetStreamBus,
|
||||
channel: &ChannelMessage,
|
||||
trace_id: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let payload = OutboundCommandPayload::typing(
|
||||
OutboundCommandKind::TypingStop,
|
||||
channel.channel.clone(),
|
||||
channel.account_id.clone(),
|
||||
channel.from_user_id.clone(),
|
||||
channel.context.clone(),
|
||||
None,
|
||||
channel.turn_id.clone(),
|
||||
);
|
||||
publish_outbound(bus, payload, trace_id).await
|
||||
}
|
||||
|
||||
async fn publish_outbound(
|
||||
bus: &JetStreamBus,
|
||||
payload: OutboundCommandPayload,
|
||||
trace_id: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let subject = outbound_subject(&payload.channel, &payload.account_id);
|
||||
let envelope = MessageEnvelope::new(
|
||||
subject.as_str(),
|
||||
payload.command_id.clone(),
|
||||
trace_id.to_string(),
|
||||
payload,
|
||||
);
|
||||
bus.publish_envelope(subject.as_str(), &envelope).await
|
||||
}
|
||||
|
||||
async fn publish_log(
|
||||
bus: &JetStreamBus,
|
||||
level: &str,
|
||||
source: &str,
|
||||
message: &str,
|
||||
trace_id: &str,
|
||||
fields: serde_json::Value,
|
||||
) {
|
||||
let payload = LogEventPayload::new(level, source, message, trace_id.to_string(), fields);
|
||||
let subject = log_subject(source);
|
||||
let envelope = MessageEnvelope::new(
|
||||
subject.as_str(),
|
||||
payload.event_id.clone(),
|
||||
trace_id.to_string(),
|
||||
payload,
|
||||
);
|
||||
if let Err(error) = bus.publish_envelope(subject.as_str(), &envelope).await {
|
||||
eprintln!("发布 assistant 日志事件失败: {error:#}");
|
||||
}
|
||||
}
|
||||
|
||||
fn with_tool_budget_message(
|
||||
mut messages: Vec<ChatCompletionRequestMessage>,
|
||||
budget: ToolBudget,
|
||||
) -> Vec<ChatCompletionRequestMessage> {
|
||||
let insert_at = messages
|
||||
.iter()
|
||||
.position(|message| message.role != "system")
|
||||
.unwrap_or(messages.len());
|
||||
messages.insert(
|
||||
insert_at,
|
||||
ChatCompletionRequestMessage::new("system".to_string(), tool_budget_prompt(budget)),
|
||||
);
|
||||
messages
|
||||
}
|
||||
|
||||
fn tool_budget_prompt(budget: ToolBudget) -> String {
|
||||
tool_budget_prompt_with_limits(budget, max_tool_rounds(), max_tool_calls_per_round())
|
||||
}
|
||||
|
||||
fn tool_budget_prompt_with_limits(
|
||||
budget: ToolBudget,
|
||||
max_rounds: i32,
|
||||
max_calls_per_round: usize,
|
||||
) -> String {
|
||||
let remaining_rounds = max_rounds.saturating_sub(budget.tool_rounds);
|
||||
let round_guidance = if remaining_rounds <= 2 {
|
||||
"已接近工具调用上限；除非只需一次必要的新查询，否则停止调用工具，直接总结当前结果或请用户缩小范围。"
|
||||
} else {
|
||||
"优先在 1 到 2 轮工具往返内完成任务；已有足够信息时必须直接回答。"
|
||||
};
|
||||
|
||||
format!(
|
||||
"本轮工具调用预算（软约束，硬熔断仍由系统执行）：\n\
|
||||
- 已使用 {}/{} 轮工具往返，累计 {} 次工具调用；每轮最多 {} 次工具调用。\n\
|
||||
- 每次调用工具前确认它会带来新的、必要的信息；不要重复调用相同能力和相同参数。\n\
|
||||
- 邮件类宽泛请求默认只查列表并汇总发件人、主题、时间和预览；不要逐封读取详情。同一邮件 ID 本轮最多读取一次。\n\
|
||||
- 工具结果截断、返回过大或仍不完整时，先总结当前结果并请用户缩小范围，不要循环查看剩余内容。\n\
|
||||
- {round_guidance}",
|
||||
budget.tool_rounds, max_rounds, budget.tool_calls, max_calls_per_round
|
||||
)
|
||||
}
|
||||
|
||||
fn operation_key(channel: &ChannelMessage) -> String {
|
||||
format!(
|
||||
"{}:{}:{}",
|
||||
channel.channel, channel.account_id, channel.from_user_id
|
||||
)
|
||||
}
|
||||
|
||||
fn max_tool_rounds() -> i32 {
|
||||
read_i32_env("IAS_ASSISTANT_MAX_TOOL_ROUNDS", DEFAULT_MAX_TOOL_ROUNDS)
|
||||
}
|
||||
|
||||
fn max_tool_calls_per_round() -> usize {
|
||||
read_usize_env(
|
||||
"IAS_ASSISTANT_MAX_TOOL_CALLS_PER_ROUND",
|
||||
DEFAULT_MAX_TOOL_CALLS_PER_ROUND,
|
||||
)
|
||||
}
|
||||
|
||||
fn read_i32_env(name: &str, default: i32) -> i32 {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<i32>().ok())
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn read_usize_env(name: &str, default: usize) -> usize {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn inbound_message_id(channel: &str, account_id: &str, external_message_id: &str) -> String {
|
||||
let subject = inbound_subject(channel, account_id);
|
||||
format!("{}/{}", subject.as_str(), external_message_id)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
ToolBudget, inbound_message_id, operation_key, tool_budget_prompt_with_limits,
|
||||
with_tool_budget_message,
|
||||
};
|
||||
use ias_common::model::ai::ChatCompletionRequestMessage;
|
||||
use ias_common::queue::ChannelMessage;
|
||||
|
||||
#[test]
|
||||
fn operation_key_is_stable_for_active_turn_idempotency() {
|
||||
let channel = ChannelMessage {
|
||||
channel: "wechat".to_string(),
|
||||
account_id: "acct".to_string(),
|
||||
from_user_id: "user".to_string(),
|
||||
content: "hello".to_string(),
|
||||
context: Some("ctx".to_string()),
|
||||
turn_id: None,
|
||||
external_message_id: None,
|
||||
session_id: None,
|
||||
group_id: None,
|
||||
create_time_ms: None,
|
||||
};
|
||||
|
||||
assert_eq!(operation_key(&channel), "wechat:acct:user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inbound_message_id_uses_subject_and_external_id() {
|
||||
assert_eq!(
|
||||
inbound_message_id("wechat", "acct", "msg-1"),
|
||||
"inbound.wechat.acct/msg-1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_budget_prompt_pushes_model_to_stop_near_limit() {
|
||||
let prompt = tool_budget_prompt_with_limits(
|
||||
ToolBudget {
|
||||
tool_rounds: 6,
|
||||
tool_calls: 14,
|
||||
},
|
||||
8,
|
||||
5,
|
||||
);
|
||||
|
||||
assert!(prompt.contains("已使用 6/8 轮工具往返"));
|
||||
assert!(prompt.contains("同一邮件 ID 本轮最多读取一次"));
|
||||
assert!(prompt.contains("不要重复调用相同能力和相同参数"));
|
||||
assert!(prompt.contains("已接近工具调用上限"));
|
||||
assert!(prompt.contains("直接总结当前结果或请用户缩小范围"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_budget_message_is_inserted_after_system_messages() {
|
||||
let messages = vec![
|
||||
ChatCompletionRequestMessage::new("system".to_string(), "固定规则".to_string()),
|
||||
ChatCompletionRequestMessage::new("user".to_string(), "查邮件".to_string()),
|
||||
];
|
||||
|
||||
let messages = with_tool_budget_message(
|
||||
messages,
|
||||
ToolBudget {
|
||||
tool_rounds: 0,
|
||||
tool_calls: 0,
|
||||
},
|
||||
);
|
||||
|
||||
assert_eq!(messages.len(), 3);
|
||||
assert_eq!(messages[0].content, "固定规则");
|
||||
assert_eq!(messages[1].role, "system");
|
||||
assert!(messages[1].content.contains("本轮工具调用预算"));
|
||||
assert_eq!(messages[2].content, "查邮件");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user