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