Skip to main content

llm_trait/
request.rs

1//! Chat request types.
2
3use serde_json::Value;
4
5use super::message::ChatMessage;
6use super::reasoning::ReasoningConfig;
7
8/// Response format configuration.
9#[derive(Clone, Debug)]
10pub enum ResponseFormat {
11    JsonObject,
12    JsonSchema { name: String, schema: Value },
13}
14
15impl ResponseFormat {
16    pub fn to_api_value(&self) -> Value {
17        match self {
18            ResponseFormat::JsonObject => {
19                serde_json::json!({ "type": "json_object" })
20            }
21            ResponseFormat::JsonSchema { name, schema } => {
22                serde_json::json!({
23                    "type": "json_schema",
24                    "json_schema": {
25                        "name": name,
26                        "schema": schema,
27                    }
28                })
29            }
30        }
31    }
32}
33
34/// Unified chat request format.
35#[derive(Debug, Clone)]
36pub struct ChatRequest {
37    /// Message list (supports multimodal content)
38    pub messages: Vec<ChatMessage>,
39
40    /// Available tools
41    pub tools: Vec<Value>,
42
43    /// Reasoning/thinking configuration
44    pub reasoning: Option<ReasoningConfig>,
45
46    /// Response format configuration
47    pub response_format: Option<ResponseFormat>,
48}
49
50impl ChatRequest {
51    /// Create a simple text request.
52    pub fn new(messages: Vec<ChatMessage>) -> Self {
53        Self {
54            messages,
55            tools: Vec::new(),
56            reasoning: None,
57            response_format: None,
58        }
59    }
60
61    /// Add tools.
62    pub fn with_tools(mut self, tools: Vec<Value>) -> Self {
63        self.tools = tools;
64        self
65    }
66
67    /// Add reasoning configuration.
68    pub fn with_reasoning(mut self, reasoning: ReasoningConfig) -> Self {
69        self.reasoning = Some(reasoning);
70        self
71    }
72
73    /// Set response format.
74    pub fn with_response_format(mut self, format: ResponseFormat) -> Self {
75        self.response_format = Some(format);
76        self
77    }
78}
79
80#[cfg(test)]
81mod response_format_tests {
82    use super::*;
83
84    #[test]
85    fn json_object_to_api_value() {
86        assert_eq!(
87            ResponseFormat::JsonObject.to_api_value(),
88            serde_json::json!({ "type": "json_object" })
89        );
90    }
91
92    #[test]
93    fn json_schema_to_api_value() {
94        let schema = serde_json::json!({"type": "object"});
95        let value = ResponseFormat::JsonSchema {
96            name: "answer".to_string(),
97            schema: schema.clone(),
98        }
99        .to_api_value();
100
101        assert_eq!(value["type"], "json_schema");
102        assert_eq!(value["json_schema"]["name"], "answer");
103        assert_eq!(value["json_schema"]["schema"], schema);
104    }
105
106    #[test]
107    fn with_response_format_stores_the_format() {
108        let req = ChatRequest::new(vec![ChatMessage::user("hi")])
109            .with_response_format(ResponseFormat::JsonObject);
110        assert!(matches!(
111            req.response_format,
112            Some(ResponseFormat::JsonObject)
113        ));
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn chat_request_new() {
123        let req = ChatRequest::new(vec![ChatMessage::user("hello")]);
124        assert_eq!(req.messages.len(), 1);
125        assert!(req.tools.is_empty());
126        assert!(req.reasoning.is_none());
127        assert!(req.response_format.is_none());
128    }
129
130    #[test]
131    fn chat_request_with_tools() {
132        let tools = vec![serde_json::json!({"type": "function", "function": {"name": "test"}})];
133        let req = ChatRequest::new(vec![ChatMessage::user("hello")]).with_tools(tools.clone());
134        assert_eq!(req.tools.len(), 1);
135        assert_eq!(req.tools[0], tools[0]);
136    }
137
138    #[test]
139    fn chat_request_with_reasoning() {
140        let reasoning = ReasoningConfig {
141            enabled: Some(true),
142            budget_tokens: Some(4096),
143            effort: None,
144        };
145        let req = ChatRequest::new(vec![ChatMessage::user("hello")]).with_reasoning(reasoning);
146        assert!(req.reasoning.is_some());
147        assert_eq!(req.reasoning.as_ref().unwrap().budget_tokens, Some(4096));
148    }
149
150    #[test]
151    fn chat_request_builder_chain() {
152        let tools = vec![serde_json::json!({"name": "tool1"})];
153        let reasoning = ReasoningConfig {
154            enabled: Some(true),
155            ..Default::default()
156        };
157        let req = ChatRequest::new(vec![ChatMessage::user("hello")])
158            .with_tools(tools)
159            .with_reasoning(reasoning);
160        assert_eq!(req.tools.len(), 1);
161        assert!(req.reasoning.is_some());
162    }
163}