lc-agents 0.17.0

// lc-agents/src/executor/hooks.rs
//! Completion hooks run around each LLM call (rate limiting / quota).

use super::AgentError;
use crate::hooks::{AgentHook, CompletionAction, CompletionContext, CompletionResult};
use crate::types::AgentOutput;
use lc_core::language_models::TokenUsage;
use lc_schema::Message;
use std::collections::HashMap;
use std::sync::Arc;

/// P2-9: LLM 调用前跑一遍 completion hooks(限流/配额)。
///
/// 构造 [`CompletionContext`] 后逐个调用 `on_before_completion`:
/// - `Continue` → 放行;
/// - `Modify` → 执行器无法改写 Agent 自建 prompt,记 warn 后继续;
/// - `Reject { reason }` → 转为 `AgentError` 中止本轮。
pub(crate) fn run_before_completion_hooks(
    hooks: &[Arc<dyn AgentHook>],
    inputs: &HashMap<String, String>,
) -> Result<(), AgentError> {
    let messages = inputs
        .values()
        .map(|v| Message::human(v.clone()))
        .collect::<Vec<_>>();
    let mut ctx = CompletionContext {
        messages,
        model: "agent".to_string(),
        metadata: HashMap::new(),
    };
    for hook in hooks {
        match hook.on_before_completion(&mut ctx) {
            CompletionAction::Continue => {}
            CompletionAction::Modify { .. } => {
                log::warn!(
                    target: "lc_agents::security",
                    "CompletionAction::Modify ignored at AgentExecutor level (agent builds its own prompt)"
                );
            }
            CompletionAction::Reject { reason } => {
                return Err(AgentError::Other(format!(
                    "LLM call rejected by hook: {reason}"
                )));
            }
        }
    }
    Ok(())
}

/// P2-9: LLM 调用后跑一遍 completion hooks(累计 token 用量)。
///
/// 构造 [`CompletionResult`] 后逐个调用 `on_after_completion`;hook 报错只记
/// warn 不中止执行(与 `on_after_tool_call` 的容错策略一致)。
pub(crate) fn run_after_completion_hooks(
    hooks: &[Arc<dyn AgentHook>],
    output: &AgentOutput,
    token_usage: Option<&TokenUsage>,
) {
    let message = Message::ai(match output {
        AgentOutput::Finish(finish) => finish.output().unwrap_or("").to_string(),
        _ => String::new(),
    });
    let mut ctx = CompletionResult {
        message,
        tokens_used: token_usage.cloned(),
    };
    for hook in hooks {
        if let Err(e) = hook.on_after_completion(&mut ctx) {
            log::warn!("Hook on_after_completion error: {}", e);
        }
    }
}