use lc_schema::Message;
use serde::de::DeserializeOwned;
use crate::language_models::BaseChatModel;
use crate::tools::ToolDefinition;
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum StructuredJudgeError {
#[error("LLM call failed: {0}")]
Call(String),
#[error("structured parse failed: {0}")]
Parse(String),
}
pub async fn structured_call<M, T, F>(
judge: &M,
tool: ToolDefinition,
messages: Vec<Message>,
text_fallback: F,
) -> Result<T, StructuredJudgeError>
where
M: BaseChatModel,
T: DeserializeOwned,
F: FnOnce(&str) -> Result<T, StructuredJudgeError>,
{
if let Some(bound) = judge.bind_tools(vec![tool]) {
let result = bound
.chat(messages, None)
.await
.map_err(|e| StructuredJudgeError::Call(e.to_string()))?;
match result.tool_calls {
Some(calls) => {
let call = calls.first().ok_or_else(|| {
StructuredJudgeError::Parse("judge returned empty tool_calls".to_string())
})?;
let parsed = call.parse_arguments::<T>().map_err(|e| {
StructuredJudgeError::Parse(format!(
"failed to parse judge structured arguments: {}",
e
))
})?;
Ok(parsed)
}
None => {
log::warn!(
"judge model bound tools but returned plain text; falling back to text parsing"
);
text_fallback(&result.content)
}
}
} else {
log::warn!("judge model does not support bind_tools; falling back to text parsing");
let result = judge
.chat(messages, None)
.await
.map_err(|e| StructuredJudgeError::Call(e.to_string()))?;
text_fallback(&result.content)
}
}
pub fn truncate(s: &str, max: usize) -> String {
if s.chars().count() <= max {
s.to_string()
} else {
let truncated: String = s.chars().take(max).collect();
format!("{}...", truncated)
}
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
use futures_util::Stream;
use lc_schema::Message;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use crate::language_models::{LLMResult, StreamChunk};
use crate::{BaseLanguageModel, Runnable, RunnableConfig};
#[derive(Debug)]
struct JudgeError(String);
impl std::fmt::Display for JudgeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl std::error::Error for JudgeError {}
struct SeqMockJudge {
replies: Vec<String>,
call: Arc<AtomicUsize>,
}
impl SeqMockJudge {
fn new(replies: Vec<String>) -> Self {
Self {
replies,
call: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl Runnable<Vec<Message>, LLMResult> for SeqMockJudge {
type Error = JudgeError;
async fn invoke(
&self,
_input: Vec<Message>,
_config: Option<RunnableConfig>,
) -> Result<LLMResult, Self::Error> {
Err(JudgeError("use chat".into()))
}
}
#[async_trait]
impl BaseLanguageModel<Vec<Message>, LLMResult> for SeqMockJudge {
fn model_name(&self) -> &str {
"seq-mock"
}
fn get_num_tokens(&self, t: &str) -> usize {
t.len()
}
fn with_temperature(self, _: f32) -> Self {
self
}
fn with_max_tokens(self, _: usize) -> Self {
self
}
}
#[async_trait]
impl BaseChatModel for SeqMockJudge {
async fn chat(
&self,
_messages: Vec<Message>,
_config: Option<RunnableConfig>,
) -> Result<LLMResult, Self::Error> {
let idx = self.call.fetch_add(1, Ordering::SeqCst);
let reply = self.replies.get(idx).cloned().unwrap_or_default();
Ok(LLMResult {
content: reply,
model: "seq-mock".to_string(),
token_usage: None,
tool_calls: None,
thinking_content: None,
})
}
async fn stream_chat(
&self,
_messages: Vec<Message>,
_config: Option<RunnableConfig>,
) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
{
Err(JudgeError("not supported".into()))
}
}
#[derive(serde::Deserialize, Debug)]
struct MockArgs {
verdict: String,
}
fn mock_tool() -> ToolDefinition {
ToolDefinition::new("mock_judge", "返回判定。")
}
#[tokio::test]
async fn test_fallback_on_text_only_model() {
let judge = SeqMockJudge::new(vec!["yes".into()]);
let messages = vec![Message::human("判断")];
let out = structured_call(&judge, mock_tool(), messages, |raw| {
Ok(MockArgs {
verdict: raw.trim().to_string(),
})
})
.await
.unwrap();
assert_eq!(out.verdict, "yes");
}
#[tokio::test]
async fn test_parse_error_raised_not_silently_defaulted() {
let judge = SeqMockJudge::new(vec!["没法判断".into()]);
let messages = vec![Message::human("判断")];
let err = structured_call(
&judge,
mock_tool(),
messages,
|_raw: &str| -> Result<MockArgs, StructuredJudgeError> {
Err(StructuredJudgeError::Parse("parse failed".into()))
},
)
.await
.unwrap_err();
assert!(matches!(err, StructuredJudgeError::Parse(_)));
}
}