lc_testkit/
recording.rs

1//! `RecordingProvider`:真实调用一次,把请求/响应对追加到 JSONL 录制文件。
2//!
3//! 录制是**旁路**:真实调用失败就返回失败、不写录播;真实调用成功但写盘失败
4//! 仅 `log::warn!`,不阻断真实结果。
5
6use std::io::Write;
7use std::path::Path;
8use std::pin::Pin;
9use std::sync::{Arc, Mutex};
10
11use async_trait::async_trait;
12use futures_util::{Stream, StreamExt};
13use lc_core::language_models::{
14    BaseChatModel, BaseLanguageModel, LLMResult, StreamChunk, TokenUsage,
15};
16use lc_core::runnables::{Runnable, RunnableConfig};
17use lc_core::tools::ToolDefinition;
18use lc_providers::ProviderError;
19use lc_schema::Message;
20use serde::{Deserialize, Serialize};
21
22use crate::error::TestkitError;
23
24/// 一次录制的请求/响应对,序列化为 JSONL 一行。
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct RecordedExchange {
27    /// 请求(含 system/user/assistant/tool 历史)。
28    pub messages: Vec<Message>,
29    /// 完整响应。
30    pub response: LLMResult,
31    /// 请求侧绑定的工具定义(`bind_tools` 绑定后非空)。
32    ///
33    /// `#[serde(default)]` 让旧 fixture(无 `tools` 字段)读成 `None`,零改动兼容;
34    /// `skip_serializing_if` 让未绑定工具的录播文件保持旧格式。
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub tools: Option<Vec<ToolDefinition>>,
37}
38
39/// 追加写录制文件的共享句柄(append 模式,std Mutex 保护)。
40pub struct Recorder {
41    file: Mutex<std::fs::File>,
42}
43
44impl Recorder {
45    /// 打开/创建录制文件。打不开 → 构造期直接 `Err`(fail fast)。
46    pub fn new(path: impl AsRef<Path>) -> std::io::Result<Self> {
47        let file = std::fs::OpenOptions::new()
48            .create(true)
49            .append(true)
50            .open(path)?;
51        Ok(Self {
52            file: Mutex::new(file),
53        })
54    }
55
56    /// Best-effort 追加一条录播:失败只 `log::warn!`,绝不向上传播。
57    pub fn record(&self, exchange: &RecordedExchange) {
58        let line = match serde_json::to_string(exchange) {
59            Ok(line) => line,
60            Err(e) => {
61                log::warn!("lc-testkit: failed to serialize recording: {e}");
62                return;
63            }
64        };
65        let Ok(mut file) = self.file.lock() else {
66            log::warn!("lc-testkit: recording lock poisoned");
67            return;
68        };
69        if let Err(e) = writeln!(file, "{line}") {
70            log::warn!("lc-testkit: failed to append recording: {e}");
71        }
72    }
73}
74
75/// 把内层模型错误映射为 `TestkitError`(经 `ProviderError` 无损透传)。
76fn to_testkit<E: Into<ProviderError>>(e: E) -> TestkitError {
77    TestkitError::Inner(e.into())
78}
79
80/// 包裹任意 `BaseChatModel`:成功响应后把请求/响应对追加到 JSONL。
81pub struct RecordingProvider<M> {
82    inner: M,
83    recorder: Arc<Recorder>,
84    model_name: String,
85    /// 当前绑定的工具定义(`bind_tools` 设置,chat 时录进 exchange)。
86    tools: Option<Vec<ToolDefinition>>,
87}
88
89impl<M> RecordingProvider<M>
90where
91    M: BaseChatModel + Send + Sync + 'static,
92    M::Error: Into<ProviderError>,
93{
94    /// 用内层模型 + 录制文件构造。文件打不开 → `Err`。
95    pub fn new(inner: M, path: impl AsRef<Path>) -> std::io::Result<Self> {
96        let model_name = format!("{}-recorded", inner.model_name());
97        let recorder = Arc::new(Recorder::new(path)?);
98        Ok(Self {
99            inner,
100            recorder,
101            model_name,
102            tools: None,
103        })
104    }
105
106    /// 访问内层模型。
107    pub fn inner(&self) -> &M {
108        &self.inner
109    }
110
111    /// 绑定工具:返回一个记录该工具集的新实例。
112    ///
113    /// 后续 `chat` 会把工具定义录进 `RecordedExchange.tools`,让回放侧
114    /// 能按工具名路由(见 [`crate::ReplayStrategy::ByToolName`])。
115    pub fn bind_tools(&self, tools: Vec<ToolDefinition>) -> Self
116    where
117        M: Clone,
118    {
119        Self {
120            inner: self.inner.clone(),
121            recorder: self.recorder.clone(),
122            model_name: self.model_name.clone(),
123            tools: Some(tools),
124        }
125    }
126}
127
128#[async_trait]
129impl<M> Runnable<Vec<Message>, LLMResult> for RecordingProvider<M>
130where
131    M: BaseChatModel + Clone + Send + Sync + 'static,
132    M::Error: Into<ProviderError>,
133{
134    type Error = TestkitError;
135
136    async fn invoke(
137        &self,
138        input: Vec<Message>,
139        config: Option<RunnableConfig>,
140    ) -> Result<LLMResult, Self::Error> {
141        self.chat(input, config).await
142    }
143}
144
145impl<M> BaseLanguageModel<Vec<Message>, LLMResult> for RecordingProvider<M>
146where
147    M: BaseChatModel + Clone + Send + Sync + 'static,
148    M::Error: Into<ProviderError>,
149{
150    fn model_name(&self) -> &str {
151        &self.model_name
152    }
153
154    fn get_num_tokens(&self, text: &str) -> usize {
155        self.inner.get_num_tokens(text)
156    }
157
158    fn temperature(&self) -> Option<f32> {
159        self.inner.temperature()
160    }
161
162    fn max_tokens(&self) -> Option<usize> {
163        self.inner.max_tokens()
164    }
165
166    fn with_temperature(mut self, temp: f32) -> Self {
167        self.inner = self.inner.with_temperature(temp);
168        self
169    }
170
171    fn with_max_tokens(mut self, max: usize) -> Self {
172        self.inner = self.inner.with_max_tokens(max);
173        self
174    }
175}
176
177#[async_trait]
178impl<M> BaseChatModel for RecordingProvider<M>
179where
180    M: BaseChatModel + Clone + Send + Sync + 'static,
181    M::Error: Into<ProviderError>,
182{
183    async fn chat(
184        &self,
185        messages: Vec<Message>,
186        config: Option<RunnableConfig>,
187    ) -> Result<LLMResult, Self::Error> {
188        let response = self
189            .inner
190            .chat(messages.clone(), config)
191            .await
192            .map_err(to_testkit)?;
193        self.recorder.record(&RecordedExchange {
194            messages,
195            response: response.clone(),
196            tools: self.tools.clone(),
197        });
198        Ok(response)
199    }
200
201    async fn stream_chat(
202        &self,
203        messages: Vec<Message>,
204        config: Option<RunnableConfig>,
205    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
206    {
207        let mut stream = self
208            .inner
209            .stream_chat(messages.clone(), config)
210            .await
211            .map_err(to_testkit)?;
212        let mut full = String::new();
213        let mut usage: Option<TokenUsage> = None;
214        while let Some(chunk) = stream.next().await {
215            let chunk = chunk.map_err(to_testkit)?;
216            full.push_str(&chunk.text);
217            if chunk.token_usage.is_some() {
218                usage = chunk.token_usage;
219            }
220        }
221        let response = LLMResult {
222            content: full.clone(),
223            model: self.model_name.clone(),
224            token_usage: usage.clone(),
225            ..Default::default()
226        };
227        self.recorder.record(&RecordedExchange {
228            messages,
229            response,
230            tools: self.tools.clone(),
231        });
232        let stream = futures_util::stream::iter(vec![Ok(StreamChunk {
233            text: full,
234            token_usage: usage,
235        })]);
236        Ok(Box::pin(stream))
237    }
238
239    fn bind_tools(
240        &self,
241        tools: Vec<ToolDefinition>,
242    ) -> Option<Box<dyn BaseChatModel<Error = Self::Error> + Send + Sync>> {
243        // 委托给 inherent `bind_tools`:克隆内层模型、记录工具集、返回新实例。
244        Some(Box::new(self.bind_tools(tools)))
245    }
246}
lc_testkit/recording.rs

lc_testkit/
recording.rs
