lc_testkit/
replay.rs

1//! `ReplayProvider`:从录制文件按 FIFO 顺序回放,零网络。
2//!
3//! 回放**不做消息匹配**(LLMChain 渲染出的 prompt 逐次可变),只按顺序弹出录播。
4//! 队列耗尽返回 [`TestkitError::ReplayExhausted`]。
5
6use std::collections::VecDeque;
7use std::io::BufRead;
8use std::path::Path;
9use std::pin::Pin;
10use std::sync::{Arc, Mutex};
11
12use async_trait::async_trait;
13use futures_util::Stream;
14use lc_core::language_models::{BaseChatModel, BaseLanguageModel, LLMResult};
15use lc_core::runnables::{Runnable, RunnableConfig};
16use lc_core::tools::ToolDefinition;
17use lc_schema::Message;
18
19use crate::error::TestkitError;
20use crate::recording::RecordedExchange;
21
22/// 回放策略:决定并发/乱序场景下一条请求取哪条录播。
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
24pub enum ReplayStrategy {
25    /// 并发 FIFO:按录制顺序弹出,原子 pop 无 UB,但"哪条请求拿到哪条响应"不确定。
26    ///
27    /// 适合串行/顺序稳定的链;并行场景配**顺序无关断言**(只断言回放全部耗尽、
28    /// 各响应结构非空),不断言"第 N 条是工具 A 的响应"。
29    #[default]
30    Fifo,
31    /// 按工具名路由:绑定工具时携带名字(`bind_tools`),回放时从队列里挑
32    /// 工具名匹配的第一条录播(请求侧 `exchange.tools` 或响应侧 `tool_calls`)。
33    ///
34    /// 覆盖"不同工具响应不同、需要精确对应"的并行场景——每个并行分支各绑
35    /// 各的工具,拿到各自正确的响应。比"整段消息签名匹配"弱但可复现。
36    ByToolName,
37}
38
39/// 从录制文件回放的零网络 `BaseChatModel`。
40///
41/// 队列用 `Arc<Mutex<VecDeque>>` 共享,`bind_tools` 返回携带工具集的新实例
42/// 时与原件共享同一队列(FIFO 顺序在分支间一致)。
43#[derive(Clone)]
44pub struct ReplayProvider {
45    queue: Arc<Mutex<VecDeque<RecordedExchange>>>,
46    model_name: String,
47    strategy: ReplayStrategy,
48    /// 本实例绑定的工具定义(`ByToolName` 路由的匹配键)。
49    tools: Option<Vec<ToolDefinition>>,
50}
51
52impl ReplayProvider {
53    /// 读 JSONL 录制文件(缺文件/坏行 → `Err`)。
54    pub fn from_file(path: impl AsRef<Path>) -> Result<Self, TestkitError> {
55        let file = std::fs::File::open(path)?;
56        let reader = std::io::BufReader::new(file);
57        let mut queue = VecDeque::new();
58        for line in reader.lines() {
59            let line = line?.trim().to_string();
60            if line.is_empty() {
61                continue;
62            }
63            let exchange: RecordedExchange = serde_json::from_str(&line).map_err(|e| {
64                TestkitError::Io(std::io::Error::new(
65                    std::io::ErrorKind::InvalidData,
66                    format!("invalid recording line: {e}"),
67                ))
68            })?;
69            queue.push_back(exchange);
70        }
71        Ok(Self {
72            queue: Arc::new(Mutex::new(queue)),
73            model_name: "replay".to_string(),
74            strategy: ReplayStrategy::Fifo,
75            tools: None,
76        })
77    }
78
79    /// 内存构造(手写录播 = MockProvider 的等价物)。
80    pub fn from_exchanges(exchanges: Vec<RecordedExchange>) -> Self {
81        Self {
82            queue: Arc::new(Mutex::new(exchanges.into())),
83            model_name: "replay".to_string(),
84            strategy: ReplayStrategy::Fifo,
85            tools: None,
86        }
87    }
88
89    /// 单一固定响应:任意请求都返回同一 `response`(最简 mock)。
90    pub fn single(response: LLMResult) -> Self {
91        Self::from_exchanges(vec![RecordedExchange {
92            messages: Vec::new(),
93            response,
94            tools: None,
95        }])
96    }
97
98    /// 设置回放策略(默认 FIFO;`ByToolName` 按工具名路由)。
99    pub fn with_strategy(mut self, strategy: ReplayStrategy) -> Self {
100        self.strategy = strategy;
101        self
102    }
103
104    /// 绑定工具:返回携带该工具集的新实例(共享同一回放队列)。
105    ///
106    /// 回放侧绑工具只是让 agent"绑了工具的循环"前提成立——录播里已含
107    /// 请求侧 `tools` 与响应侧 `tool_calls`,回放逻辑无需改。`ByToolName`
108    /// 策略下,`tools` 同时是路由的匹配键。
109    pub fn bind_tools(&self, tools: Vec<ToolDefinition>) -> Self {
110        Self {
111            queue: self.queue.clone(),
112            model_name: self.model_name.clone(),
113            strategy: self.strategy,
114            tools: Some(tools),
115        }
116    }
117
118    /// 剩余录播条数。
119    pub fn len(&self) -> usize {
120        self.queue.lock().unwrap_or_else(|e| e.into_inner()).len()
121    }
122
123    /// 是否已无剩余录播。
124    pub fn is_empty(&self) -> bool {
125        self.len() == 0
126    }
127}
128
129/// `ByToolName` 的匹配:请求侧绑定的工具名或响应侧请求的工具调用名命中即真。
130fn exchange_matches(exchange: &RecordedExchange, tool_name: &str) -> bool {
131    if let Some(tools) = &exchange.tools {
132        if tools.iter().any(|t| t.function.name == tool_name) {
133            return true;
134        }
135    }
136    if let Some(calls) = &exchange.response.tool_calls {
137        if calls.iter().any(|c| c.name() == tool_name) {
138            return true;
139        }
140    }
141    false
142}
143
144#[async_trait]
145impl Runnable<Vec<Message>, LLMResult> for ReplayProvider {
146    type Error = TestkitError;
147
148    async fn invoke(
149        &self,
150        input: Vec<Message>,
151        config: Option<RunnableConfig>,
152    ) -> Result<LLMResult, Self::Error> {
153        self.chat(input, config).await
154    }
155}
156
157impl BaseLanguageModel<Vec<Message>, LLMResult> for ReplayProvider {
158    fn model_name(&self) -> &str {
159        &self.model_name
160    }
161
162    fn get_num_tokens(&self, text: &str) -> usize {
163        // 估算:约 4 字符/ token。
164        text.chars().count() / 4 + 1
165    }
166
167    fn temperature(&self) -> Option<f32> {
168        None
169    }
170
171    fn max_tokens(&self) -> Option<usize> {
172        None
173    }
174
175    fn with_temperature(self, _temp: f32) -> Self {
176        self
177    }
178
179    fn with_max_tokens(self, _max: usize) -> Self {
180        self
181    }
182}
183
184#[async_trait]
185impl BaseChatModel for ReplayProvider {
186    async fn chat(
187        &self,
188        messages: Vec<Message>,
189        _config: Option<RunnableConfig>,
190    ) -> Result<LLMResult, Self::Error> {
191        let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
192        let exchange = match self.strategy {
193            ReplayStrategy::Fifo => queue.pop_front(),
194            ReplayStrategy::ByToolName => {
195                let want = self
196                    .tools
197                    .as_ref()
198                    .and_then(|tools| tools.first().map(|t| t.function.name.clone()));
199                match want {
200                    // 按工具名从队列里挑匹配录播;未绑定工具 → 退化为 FIFO。
201                    Some(name) => queue
202                        .iter()
203                        .position(|ex| exchange_matches(ex, &name))
204                        .map(|i| queue.remove(i).expect("position 必有元素")),
205                    None => queue.pop_front(),
206                }
207            }
208        };
209        let Some(exchange) = exchange else {
210            return Err(TestkitError::ReplayExhausted {
211                requested: messages.len(),
212            });
213        };
214        Ok(exchange.response)
215    }
216
217    async fn stream_chat(
218        &self,
219        messages: Vec<Message>,
220        config: Option<RunnableConfig>,
221    ) -> Result<Pin<Box<dyn Stream<Item = Result<String, Self::Error>> + Send>>, Self::Error> {
222        let response = self.chat(messages, config).await?;
223        let stream = futures_util::stream::iter(vec![Ok(response.content)]);
224        Ok(Box::pin(stream))
225    }
226
227    fn bind_tools(
228        &self,
229        tools: Vec<ToolDefinition>,
230    ) -> Option<Box<dyn BaseChatModel<Error = Self::Error> + Send + Sync>> {
231        // 委托给 inherent `bind_tools`:共享队列、记录工具集、返回新实例。
232        Some(Box::new(self.bind_tools(tools)))
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239    use lc_core::language_models::TokenUsage;
240    use lc_core::tools::{ToolCall, ToolDefinition};
241
242    fn exchange(content: &str) -> RecordedExchange {
243        RecordedExchange {
244            messages: vec![Message::system("ping")],
245            response: LLMResult {
246                content: content.to_string(),
247                model: "replay".to_string(),
248                token_usage: Some(TokenUsage {
249                    prompt_tokens: 1,
250                    completion_tokens: 2,
251                    total_tokens: 3,
252                }),
253                ..Default::default()
254            },
255            tools: None,
256        }
257    }
258
259    /// 带响应侧工具调用名的录播(模拟模型请求调用某工具)。
260    fn exchange_with_tool_call(tool_name: &str, content: &str) -> RecordedExchange {
261        let mut response = exchange(content).response;
262        response.tool_calls = Some(vec![ToolCall::builder("call_1")
263            .name(tool_name)
264            .arguments("{}".to_string())
265            .build()]);
266        RecordedExchange {
267            response,
268            ..exchange(content)
269        }
270    }
271
272    /// 带请求侧工具定义名的录播(模拟绑定了某工具后发起请求)。
273    fn exchange_with_bound_tool(tool_name: &str, content: &str) -> RecordedExchange {
274        RecordedExchange {
275            tools: Some(vec![ToolDefinition::new(tool_name, "a tool")]),
276            ..exchange(content)
277        }
278    }
279
280    #[tokio::test]
281    async fn single_returns_fixed_response_for_any_request() {
282        let provider = ReplayProvider::single(exchange("hello").response);
283        let result = provider
284            .chat(vec![Message::system("any")], None)
285            .await
286            .unwrap();
287        assert_eq!(result.content, "hello");
288    }
289
290    #[tokio::test]
291    async fn replay_is_fifo_ordered() {
292        let provider = ReplayProvider::from_exchanges(vec![exchange("first"), exchange("second")]);
293        let first = provider
294            .chat(vec![Message::system("a")], None)
295            .await
296            .unwrap();
297        let second = provider
298            .chat(vec![Message::system("b")], None)
299            .await
300            .unwrap();
301        assert_eq!(first.content, "first");
302        assert_eq!(second.content, "second");
303    }
304
305    #[tokio::test]
306    async fn replay_exhausted_returns_error() {
307        let provider = ReplayProvider::from_exchanges(vec![exchange("only")]);
308        provider
309            .chat(vec![Message::system("a")], None)
310            .await
311            .unwrap();
312        let err = provider
313            .chat(vec![Message::system("b")], None)
314            .await
315            .unwrap_err();
316        assert!(matches!(
317            err,
318            TestkitError::ReplayExhausted { requested: 1 }
319        ));
320    }
321
322    #[test]
323    fn bind_tools_returns_some_and_carries_tools() {
324        let provider = ReplayProvider::from_exchanges(vec![exchange("x")]);
325        // inherent bind_tools 返回携带工具集的新实例(共享队列)。
326        let bound = provider.bind_tools(vec![ToolDefinition::new("calculator", "calc")]);
327        assert!(bound.tools.is_some());
328        assert_eq!(bound.tools.as_ref().unwrap()[0].function.name, "calculator");
329        assert_eq!(provider.len(), 1);
330        assert_eq!(bound.len(), 1);
331        // trait bind_tools(agent 通过 `Box<dyn BaseChatModel>` 调用)恒为 Some。
332        let trait_bound = BaseChatModel::bind_tools(&provider, vec![ToolDefinition::new("x", "y")]);
333        assert!(trait_bound.is_some());
334    }
335
336    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
337    async fn parallel_replay_fifo_is_order_independent() {
338        // 并发 FIFO:两条请求同时到达,哪条拿哪条响应不确定,但都能拿到、
339        // 且队列最终耗尽。断言用"集合相等"而非"第 N 条是 xx"。
340        let provider = ReplayProvider::from_exchanges(vec![
341            exchange("first"),
342            exchange("second"),
343            exchange("third"),
344        ]);
345        let provider = std::sync::Arc::new(provider);
346
347        let mut handles = Vec::new();
348        for _ in 0..3 {
349            let p = provider.clone();
350            handles.push(tokio::spawn(async move {
351                p.chat(vec![Message::system("parallel")], None)
352                    .await
353                    .expect("并发回放不应失败")
354            }));
355        }
356        let mut contents: Vec<String> = Vec::new();
357        for handle in handles {
358            contents.push(handle.await.unwrap().content);
359        }
360        contents.sort();
361        assert_eq!(
362            contents,
363            vec![
364                "first".to_string(),
365                "second".to_string(),
366                "third".to_string()
367            ]
368        );
369        assert!(provider.is_empty(), "并发回放应恰好耗尽全部录播");
370    }
371
372    #[tokio::test]
373    async fn by_tool_name_routes_to_matching_exchange() {
374        // 每个并行分支各绑各的工具,应各拿到各自正确的响应——顺序无关。
375        let provider = ReplayProvider::from_exchanges(vec![
376            exchange_with_tool_call("search", "search result"),
377            exchange_with_tool_call("calc", "calc result"),
378        ])
379        .with_strategy(ReplayStrategy::ByToolName);
380
381        let search =
382            BaseChatModel::bind_tools(&provider, vec![ToolDefinition::new("search", "s")]).unwrap();
383        let calc =
384            BaseChatModel::bind_tools(&provider, vec![ToolDefinition::new("calc", "c")]).unwrap();
385
386        let calc_res = calc.chat(vec![Message::system("q")], None).await.unwrap();
387        let search_res = search.chat(vec![Message::system("q")], None).await.unwrap();
388
389        assert_eq!(search_res.content, "search result");
390        assert_eq!(calc_res.content, "calc result");
391        assert!(provider.is_empty());
392    }
393
394    #[tokio::test]
395    async fn by_tool_name_matches_request_side_tools() {
396        // 请求侧 `exchange.tools` 也参与匹配(录制自绑定工具后发起的请求)。
397        let provider = ReplayProvider::from_exchanges(vec![
398            exchange_with_bound_tool("weather", "sunny"),
399            exchange_with_bound_tool("news", "headlines"),
400        ])
401        .with_strategy(ReplayStrategy::ByToolName);
402
403        let weather =
404            BaseChatModel::bind_tools(&provider, vec![ToolDefinition::new("weather", "w")])
405                .unwrap();
406        let res = weather
407            .chat(vec![Message::system("q")], None)
408            .await
409            .unwrap();
410        assert_eq!(res.content, "sunny");
411    }
412}
lc_testkit/replay.rs

lc_testkit/
replay.rs
