1use lc_schema::Message;
10use serde::de::DeserializeOwned;
11
12use crate::language_models::BaseChatModel;
13use crate::tools::ToolDefinition;
14
15#[derive(Debug, thiserror::Error)]
19#[non_exhaustive]
20pub enum StructuredJudgeError {
21 #[error("LLM call failed: {0}")]
23 Call(String),
24 #[error("structured parse failed: {0}")]
26 Parse(String),
27}
28
29pub async fn structured_call<M, T, F>(
40 judge: &M,
41 tool: ToolDefinition,
42 messages: Vec<Message>,
43 text_fallback: F,
44) -> Result<T, StructuredJudgeError>
45where
46 M: BaseChatModel,
47 T: DeserializeOwned,
48 F: FnOnce(&str) -> Result<T, StructuredJudgeError>,
49{
50 if let Some(bound) = judge.bind_tools(vec![tool]) {
51 let result = bound
52 .chat(messages, None)
53 .await
54 .map_err(|e| StructuredJudgeError::Call(e.to_string()))?;
55 match result.tool_calls {
56 Some(calls) => {
57 let call = calls.first().ok_or_else(|| {
58 StructuredJudgeError::Parse("judge returned empty tool_calls".to_string())
59 })?;
60 let parsed = call.parse_arguments::<T>().map_err(|e| {
61 StructuredJudgeError::Parse(format!(
62 "failed to parse judge structured arguments: {}",
63 e
64 ))
65 })?;
66 Ok(parsed)
67 }
68 None => {
69 log::warn!(
70 "judge model bound tools but returned plain text; falling back to text parsing"
71 );
72 text_fallback(&result.content)
73 }
74 }
75 } else {
76 log::warn!("judge model does not support bind_tools; falling back to text parsing");
77 let result = judge
78 .chat(messages, None)
79 .await
80 .map_err(|e| StructuredJudgeError::Call(e.to_string()))?;
81 text_fallback(&result.content)
82 }
83}
84
85pub fn truncate(s: &str, max: usize) -> String {
87 if s.chars().count() <= max {
88 s.to_string()
89 } else {
90 let truncated: String = s.chars().take(max).collect();
91 format!("{}...", truncated)
92 }
93}
94
95#[cfg(test)]
96mod tests {
97 use super::*;
98 use async_trait::async_trait;
99 use futures_util::Stream;
100 use lc_schema::Message;
101 use std::pin::Pin;
102 use std::sync::atomic::{AtomicUsize, Ordering};
103 use std::sync::Arc;
104
105 use crate::language_models::{LLMResult, StreamChunk};
106 use crate::{BaseLanguageModel, Runnable, RunnableConfig};
107
108 #[derive(Debug)]
109 struct JudgeError(String);
110 impl std::fmt::Display for JudgeError {
111 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112 write!(f, "{}", self.0)
113 }
114 }
115 impl std::error::Error for JudgeError {}
116
117 struct SeqMockJudge {
119 replies: Vec<String>,
120 call: Arc<AtomicUsize>,
121 }
122 impl SeqMockJudge {
123 fn new(replies: Vec<String>) -> Self {
124 Self {
125 replies,
126 call: Arc::new(AtomicUsize::new(0)),
127 }
128 }
129 }
130
131 #[async_trait]
132 impl Runnable<Vec<Message>, LLMResult> for SeqMockJudge {
133 type Error = JudgeError;
134 async fn invoke(
135 &self,
136 _input: Vec<Message>,
137 _config: Option<RunnableConfig>,
138 ) -> Result<LLMResult, Self::Error> {
139 Err(JudgeError("use chat".into()))
140 }
141 }
142
143 #[async_trait]
144 impl BaseLanguageModel<Vec<Message>, LLMResult> for SeqMockJudge {
145 fn model_name(&self) -> &str {
146 "seq-mock"
147 }
148 fn get_num_tokens(&self, t: &str) -> usize {
149 t.len()
150 }
151 fn with_temperature(self, _: f32) -> Self {
152 self
153 }
154 fn with_max_tokens(self, _: usize) -> Self {
155 self
156 }
157 }
158
159 #[async_trait]
160 impl BaseChatModel for SeqMockJudge {
161 async fn chat(
162 &self,
163 _messages: Vec<Message>,
164 _config: Option<RunnableConfig>,
165 ) -> Result<LLMResult, Self::Error> {
166 let idx = self.call.fetch_add(1, Ordering::SeqCst);
167 let reply = self.replies.get(idx).cloned().unwrap_or_default();
168 Ok(LLMResult {
169 content: reply,
170 model: "seq-mock".to_string(),
171 token_usage: None,
172 tool_calls: None,
173 thinking_content: None,
174 })
175 }
176 async fn stream_chat(
177 &self,
178 _messages: Vec<Message>,
179 _config: Option<RunnableConfig>,
180 ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
181 {
182 Err(JudgeError("not supported".into()))
183 }
184 }
185
186 #[derive(serde::Deserialize, Debug)]
187 struct MockArgs {
188 verdict: String,
189 }
190
191 fn mock_tool() -> ToolDefinition {
192 ToolDefinition::new("mock_judge", "返回判定。")
193 }
194
195 #[tokio::test]
196 async fn test_fallback_on_text_only_model() {
197 let judge = SeqMockJudge::new(vec!["yes".into()]);
199 let messages = vec![Message::human("判断")];
200 let out = structured_call(&judge, mock_tool(), messages, |raw| {
201 Ok(MockArgs {
202 verdict: raw.trim().to_string(),
203 })
204 })
205 .await
206 .unwrap();
207 assert_eq!(out.verdict, "yes");
208 }
209
210 #[tokio::test]
211 async fn test_parse_error_raised_not_silently_defaulted() {
212 let judge = SeqMockJudge::new(vec!["没法判断".into()]);
214 let messages = vec![Message::human("判断")];
215 let err = structured_call(
216 &judge,
217 mock_tool(),
218 messages,
219 |_raw: &str| -> Result<MockArgs, StructuredJudgeError> {
220 Err(StructuredJudgeError::Parse("parse failed".into()))
221 },
222 )
223 .await
224 .unwrap_err();
225 assert!(matches!(err, StructuredJudgeError::Parse(_)));
226 }
227}