rskit_llm_common/
openai.rs

1//! `OpenAI` chat-completions dialect.
2
3use rskit_errors::{AppError, AppResult, ErrorCode};
4use rskit_llm::types::{
5    AssistantMessage, CompletionRequest, CompletionResponse, ContentPart, FinishReason, Message,
6    ToolUseBlock, Usage,
7};
8use serde::{Deserialize, Serialize};
9
10use crate as common;
11use crate::{StreamChunk, StreamToolCall};
12
13/// Converts an rskit [`CompletionRequest`] into the `OpenAI` wire format
14/// and parses the response back into an rskit [`CompletionResponse`].
15pub struct OpenAiDialect;
16
17// --- Wire types ---
18
19#[derive(Serialize)]
20struct ChatRequest {
21    model: String,
22    messages: Vec<WireMessage>,
23    #[serde(skip_serializing_if = "Option::is_none")]
24    max_tokens: Option<u32>,
25    #[serde(skip_serializing_if = "Option::is_none")]
26    temperature: Option<f32>,
27    stream: bool,
28}
29
30#[derive(Serialize)]
31struct WireMessage {
32    role: String,
33    content: String,
34}
35
36#[derive(Deserialize)]
37struct ChatResponse {
38    model: String,
39    choices: Vec<Choice>,
40    usage: WireUsage,
41}
42
43#[derive(Deserialize)]
44struct Choice {
45    message: ResponseMessage,
46    #[serde(default)]
47    finish_reason: Option<String>,
48}
49
50#[derive(Deserialize)]
51struct ResponseMessage {
52    content: Option<String>,
53    #[serde(default)]
54    tool_calls: Vec<ResponseToolCall>,
55}
56
57#[derive(Deserialize)]
58struct ResponseToolCall {
59    #[serde(default)]
60    id: String,
61    function: ResponseFunction,
62}
63
64#[derive(Deserialize)]
65struct ResponseFunction {
66    name: String,
67    arguments: String,
68}
69
70#[derive(Deserialize)]
71struct WireUsage {
72    prompt_tokens: u32,
73    completion_tokens: u32,
74}
75
76// --- Conversion helpers ---
77
78fn to_wire(msg: &Message) -> WireMessage {
79    match msg {
80        Message::System(s) => WireMessage {
81            role: "system".to_string(),
82            content: s.content.clone(),
83        },
84        Message::User(u) => WireMessage {
85            role: "user".to_string(),
86            content: rskit_llm::types::text_of(&u.content),
87        },
88        Message::Assistant(a) => WireMessage {
89            role: "assistant".to_string(),
90            content: rskit_llm::types::text_of(&a.content),
91        },
92        Message::Tool(tr) => WireMessage {
93            role: "tool".to_string(),
94            content: tr.content.clone(),
95        },
96    }
97}
98
99fn parse_finish_reason(reason: Option<&str>) -> FinishReason {
100    match reason {
101        Some("length") => FinishReason::Length,
102        Some("tool_calls") => FinishReason::ToolUse,
103        Some("content_filter") => FinishReason::ContentFilter,
104        Some("error") => FinishReason::Error,
105        Some("cancelled") => FinishReason::Cancelled,
106        _ => FinishReason::Stop,
107    }
108}
109
110fn parse_tool_call(tool_call: ResponseToolCall, index: usize) -> AppResult<ToolUseBlock> {
111    Ok(ToolUseBlock {
112        id: if tool_call.id.is_empty() {
113            format!("tool_call_{index}")
114        } else {
115            tool_call.id
116        },
117        name: tool_call.function.name,
118        input: common::parse_input_json(&tool_call.function.arguments)?,
119    })
120}
121
122impl OpenAiDialect {
123    /// Build the JSON body for a chat completion request.
124    pub fn build_body(req: &CompletionRequest) -> AppResult<serde_json::Value> {
125        let wire = ChatRequest {
126            model: req.model.clone(),
127            messages: req.messages.iter().map(to_wire).collect(),
128            max_tokens: req.max_tokens,
129            temperature: req.temperature,
130            stream: false,
131        };
132        serde_json::to_value(&wire).map_err(|e| {
133            AppError::new(
134                ErrorCode::Internal,
135                format!("failed to serialize OpenAI request: {e}"),
136            )
137        })
138    }
139
140    /// Parse a successful chat-completion response body.
141    pub fn parse_response(body: &str) -> AppResult<CompletionResponse> {
142        let resp: ChatResponse = serde_json::from_str(body).map_err(|e| {
143            AppError::new(
144                ErrorCode::ExternalService,
145                format!("failed to parse OpenAI response: {e}"),
146            )
147        })?;
148
149        let choice = resp.choices.into_iter().next();
150        let (content, tool_calls, stop_reason) = match choice {
151            Some(choice) => {
152                let tool_calls = choice
153                    .message
154                    .tool_calls
155                    .into_iter()
156                    .enumerate()
157                    .map(|(index, tool_call)| parse_tool_call(tool_call, index))
158                    .collect::<AppResult<Vec<_>>>()?;
159                (
160                    choice.message.content.unwrap_or_default(),
161                    tool_calls,
162                    parse_finish_reason(choice.finish_reason.as_deref()),
163                )
164            }
165            None => (String::new(), Vec::new(), FinishReason::Stop),
166        };
167
168        Ok(CompletionResponse {
169            message: AssistantMessage {
170                content: vec![ContentPart::Text { text: content }],
171                tool_calls,
172                usage: None,
173            },
174            model: resp.model,
175            usage: Usage {
176                input_tokens: u64::from(resp.usage.prompt_tokens),
177                output_tokens: u64::from(resp.usage.completion_tokens),
178                cached_tokens: 0,
179                reasoning_tokens: 0,
180            },
181            stop_reason: Some(stop_reason),
182        })
183    }
184
185    /// Parse an error response body into an [`AppError`].
186    pub fn parse_error(status: u16, body: &str) -> AppError {
187        common::parse_openai_error(status, body).into()
188    }
189
190    /// Parse a single SSE data payload from the streaming chat-completions API.
191    pub fn parse_stream_chunk(data: &[u8]) -> AppResult<StreamChunk> {
192        let s = std::str::from_utf8(data).map_err(|e| {
193            AppError::new(
194                ErrorCode::InvalidFormat,
195                format!("invalid UTF-8 in OpenAI stream chunk: {e}"),
196            )
197        })?;
198
199        if s.trim() == "[DONE]" {
200            return Ok(StreamChunk {
201                done: true,
202                ..Default::default()
203            });
204        }
205
206        let v: serde_json::Value = serde_json::from_str(s).map_err(|e| {
207            AppError::new(
208                ErrorCode::ExternalService,
209                format!("failed to parse OpenAI stream chunk: {e}"),
210            )
211        })?;
212
213        let Some(choice) = v.get("choices").and_then(|c| c.get(0)) else {
214            return Ok(StreamChunk::default());
215        };
216
217        let delta = choice.get("delta").unwrap_or(&serde_json::Value::Null);
218
219        let content = delta
220            .get("content")
221            .and_then(|c| c.as_str())
222            .unwrap_or("")
223            .to_string();
224
225        let mut tool_calls = Vec::new();
226        if let Some(tcs) = delta.get("tool_calls").and_then(|t| t.as_array()) {
227            for tc in tcs {
228                let index = tc
229                    .get("index")
230                    .and_then(serde_json::Value::as_u64)
231                    .map_or(0, |value| usize::try_from(value).unwrap_or(usize::MAX));
232                let id = tc
233                    .get("id")
234                    .and_then(|v| v.as_str())
235                    .unwrap_or("")
236                    .to_string();
237                let func = tc.get("function").unwrap_or(&serde_json::Value::Null);
238                let name = func
239                    .get("name")
240                    .and_then(|v| v.as_str())
241                    .unwrap_or("")
242                    .to_string();
243                let input_delta = func
244                    .get("arguments")
245                    .and_then(|v| v.as_str())
246                    .unwrap_or("")
247                    .to_string();
248                tool_calls.push(StreamToolCall {
249                    index,
250                    id,
251                    name,
252                    input_delta,
253                });
254            }
255        }
256
257        let done = choice
258            .get("finish_reason")
259            .and_then(|v| v.as_str())
260            .is_some_and(|reason| !reason.is_empty());
261
262        Ok(StreamChunk {
263            content,
264            tool_calls,
265            done,
266        })
267    }
268
269    /// Returns the chat completions endpoint path.
270    pub const fn endpoint() -> &'static str {
271        "/chat/completions"
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278    use rskit_llm::types;
279
280    #[test]
281    fn build_body_produces_valid_json() {
282        let req = CompletionRequest {
283            model: "gpt-4o".to_string(),
284            messages: vec![types::user("hello")],
285            max_tokens: Some(100),
286            temperature: Some(0.7),
287            stream: false,
288            tools: None,
289            tool_choice: None,
290        };
291        let body = OpenAiDialect::build_body(&req).unwrap();
292        assert_eq!(body["model"], "gpt-4o");
293        assert_eq!(body["messages"][0]["role"], "user");
294        assert_eq!(body["messages"][0]["content"], "hello");
295        assert_eq!(body["max_tokens"], 100);
296    }
297
298    #[test]
299    fn build_body_maps_all_message_roles() {
300        let req = CompletionRequest {
301            model: "gpt-4o".to_string(),
302            messages: vec![
303                types::system("be brief"),
304                types::assistant("previous"),
305                types::tool_result_msg("call_1", "result", false),
306            ],
307            max_tokens: None,
308            temperature: None,
309            stream: false,
310            tools: None,
311            tool_choice: None,
312        };
313
314        let body = OpenAiDialect::build_body(&req).unwrap();
315
316        assert_eq!(body["messages"][0]["role"], "system");
317        assert_eq!(body["messages"][1]["role"], "assistant");
318        assert_eq!(body["messages"][2]["role"], "tool");
319    }
320
321    #[test]
322    fn parse_response_extracts_content() {
323        let body = r#"{
324            "id": "chatcmpl-1",
325            "model": "gpt-4o",
326            "choices": [{"message": {"content": "Hi there!"}, "finish_reason": "stop"}],
327            "usage": {"prompt_tokens": 10, "completion_tokens": 5}
328        }"#;
329        let resp = OpenAiDialect::parse_response(body).unwrap();
330        assert_eq!(resp.text(), "Hi there!");
331        assert_eq!(resp.model, "gpt-4o");
332        assert_eq!(resp.usage.input_tokens, 10);
333        assert_eq!(resp.usage.output_tokens, 5);
334        assert_eq!(resp.stop_reason, Some(FinishReason::Stop));
335    }
336
337    #[test]
338    fn parse_response_handles_empty_choices_tool_calls_and_finish_reasons() {
339        let empty = OpenAiDialect::parse_response(
340            r#"{"model":"gpt-4o","choices":[],"usage":{"prompt_tokens":0,"completion_tokens":0}}"#,
341        )
342        .unwrap();
343        assert!(empty.text().is_empty());
344        assert_eq!(empty.stop_reason, Some(FinishReason::Stop));
345
346        let body = r#"{
347            "model": "gpt-4o",
348            "choices": [{
349                "message": {
350                    "content": null,
351                    "tool_calls": [{"id":"","function":{"name":"lookup","arguments":"{\"q\":\"rust\"}"}}]
352                },
353                "finish_reason": "tool_calls"
354            }],
355            "usage": {"prompt_tokens": 1, "completion_tokens": 2}
356        }"#;
357        let resp = OpenAiDialect::parse_response(body).unwrap();
358        assert_eq!(resp.stop_reason, Some(FinishReason::ToolUse));
359        assert_eq!(resp.message.tool_calls[0].id, "tool_call_0");
360        assert_eq!(resp.message.tool_calls[0].input["q"], "rust");
361
362        for (wire, expected) in [
363            ("length", FinishReason::Length),
364            ("content_filter", FinishReason::ContentFilter),
365            ("error", FinishReason::Error),
366            ("cancelled", FinishReason::Cancelled),
367        ] {
368            let body = format!(
369                r#"{{"model":"gpt-4o","choices":[{{"message":{{"content":"x"}},"finish_reason":"{wire}"}}],"usage":{{"prompt_tokens":1,"completion_tokens":1}}}}"#
370            );
371            let resp = OpenAiDialect::parse_response(&body).unwrap();
372            assert_eq!(resp.stop_reason, Some(expected));
373        }
374    }
375
376    #[test]
377    fn parse_response_rejects_invalid_json_and_tool_arguments() {
378        assert_eq!(
379            OpenAiDialect::parse_response("{").unwrap_err().code(),
380            ErrorCode::ExternalService
381        );
382        let body = r#"{
383            "model": "gpt-4o",
384            "choices": [{"message": {"tool_calls": [{"function":{"name":"lookup","arguments":"not-json"}}]}}],
385            "usage": {"prompt_tokens": 1, "completion_tokens": 2}
386        }"#;
387        assert_eq!(
388            OpenAiDialect::parse_response(body).unwrap_err().code(),
389            ErrorCode::InvalidFormat
390        );
391    }
392
393    #[test]
394    fn endpoint_is_correct() {
395        assert_eq!(OpenAiDialect::endpoint(), "/chat/completions");
396    }
397
398    #[test]
399    fn stream_chunk_content_delta() {
400        let data = br#"{"choices":[{"delta":{"content":"Hello"},"finish_reason":null}]}"#;
401        let chunk = OpenAiDialect::parse_stream_chunk(data).unwrap();
402        assert_eq!(chunk.content, "Hello");
403        assert!(chunk.tool_calls.is_empty());
404        assert!(!chunk.done);
405    }
406
407    #[test]
408    fn stream_chunk_tool_call() {
409        let data = br#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"loc\":"}}]},"finish_reason":null}]}"#;
410        let chunk = OpenAiDialect::parse_stream_chunk(data).unwrap();
411        assert!(chunk.content.is_empty());
412        assert_eq!(chunk.tool_calls.len(), 1);
413        assert_eq!(chunk.tool_calls[0].id, "call_1");
414        assert_eq!(chunk.tool_calls[0].name, "get_weather");
415        assert_eq!(chunk.tool_calls[0].input_delta, r#"{"loc":"#);
416        assert!(!chunk.done);
417    }
418
419    #[test]
420    fn stream_chunk_done_signal() {
421        let data = b"[DONE]";
422        let chunk = OpenAiDialect::parse_stream_chunk(data).unwrap();
423        assert!(chunk.done);
424        assert!(chunk.content.is_empty());
425    }
426
427    #[test]
428    fn stream_chunk_finish_reason_stop() {
429        let data = br#"{"choices":[{"delta":{},"finish_reason":"stop"}]}"#;
430        let chunk = OpenAiDialect::parse_stream_chunk(data).unwrap();
431        assert!(chunk.done);
432    }
433
434    #[test]
435    fn stream_chunk_empty_choices() {
436        let data = br#"{"choices":[]}"#;
437        let chunk = OpenAiDialect::parse_stream_chunk(data).unwrap();
438        assert!(chunk.content.is_empty());
439        assert!(!chunk.done);
440    }
441
442    #[test]
443    fn stream_chunk_rejects_invalid_utf8_or_json_and_defaults_tool_fields() {
444        assert_eq!(
445            OpenAiDialect::parse_stream_chunk(&[0xff])
446                .unwrap_err()
447                .code(),
448            ErrorCode::InvalidFormat
449        );
450        assert_eq!(
451            OpenAiDialect::parse_stream_chunk(b"{").unwrap_err().code(),
452            ErrorCode::ExternalService
453        );
454
455        let data =
456            br#"{"choices":[{"delta":{"tool_calls":[{"index":18446744073709551615,"function":{}}]},"finish_reason":"tool_calls"}]}"#;
457        let chunk = OpenAiDialect::parse_stream_chunk(data).unwrap();
458        assert!(chunk.done);
459        assert_eq!(chunk.tool_calls[0].index, usize::MAX);
460        assert!(chunk.tool_calls[0].id.is_empty());
461        assert!(chunk.tool_calls[0].name.is_empty());
462        assert!(chunk.tool_calls[0].input_delta.is_empty());
463    }
464}
rskit_llm_common/openai.rs

rskit_llm_common/
openai.rs
