llm_kernel/llm/
types.rs

1//! Core types for the LLM client module.
2#![deny(missing_docs)]
3
4use std::fmt;
5use std::pin::Pin;
6
7use serde::{Deserialize, Serialize};
8
9/// Role of a message sender in a chat conversation.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum MessageRole {
13    /// System-level instruction message.
14    System,
15    /// User input message.
16    User,
17    /// Assistant response message.
18    Assistant,
19    /// Tool/function result message.
20    Tool,
21}
22
23impl fmt::Display for MessageRole {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        match self {
26            Self::System => write!(f, "system"),
27            Self::User => write!(f, "user"),
28            Self::Assistant => write!(f, "assistant"),
29            Self::Tool => write!(f, "tool"),
30        }
31    }
32}
33
34/// A single content part in a multimodal chat message.
35///
36/// Supports text, image URLs, and base64-encoded images.
37/// Single-text messages serialize as a plain string for backward compatibility
38/// with OpenAI and Anthropic APIs.
39#[derive(Debug, Clone, Serialize, Deserialize)]
40#[serde(tag = "type", rename_all = "snake_case")]
41pub enum ContentPart {
42    /// Plain text content.
43    Text {
44        /// The text string.
45        text: String,
46    },
47    /// Image specified by URL.
48    ImageUrl {
49        /// URL pointing to the image.
50        url: String,
51    },
52    /// Image specified as base64-encoded data.
53    ImageBase64 {
54        /// MIME type (e.g. `"image/png"`).
55        media_type: String,
56        /// Base64-encoded image data.
57        data: String,
58    },
59}
60
61impl ContentPart {
62    /// Create a text content part.
63    pub fn text(s: impl Into<String>) -> Self {
64        Self::Text { text: s.into() }
65    }
66
67    /// Create an image URL content part.
68    pub fn image_url(url: impl Into<String>) -> Self {
69        Self::ImageUrl { url: url.into() }
70    }
71
72    /// Extract text content, if this is a text part.
73    pub fn as_text(&self) -> Option<&str> {
74        match self {
75            Self::Text { text } => Some(text),
76            _ => None,
77        }
78    }
79}
80
81/// Serde helper: serialize `Vec<ContentPart>` as a plain string when there's
82/// a single text entry, or as an array otherwise.
83mod content_vec_serde {
84    use super::ContentPart;
85    use serde::{Deserialize, Deserializer, Serialize, Serializer};
86
87    pub fn serialize<S: Serializer>(parts: &[ContentPart], s: S) -> Result<S::Ok, S::Error> {
88        if parts.len() == 1
89            && let ContentPart::Text { text } = &parts[0]
90        {
91            return s.serialize_str(text);
92        }
93        parts.serialize(s)
94    }
95
96    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<ContentPart>, D::Error> {
97        #[derive(Deserialize)]
98        #[serde(untagged)]
99        enum StringOrParts {
100            S(String),
101            P(Vec<ContentPart>),
102        }
103        match StringOrParts::deserialize(d)? {
104            StringOrParts::S(s) => Ok(vec![ContentPart::text(s)]),
105            StringOrParts::P(v) => Ok(v),
106        }
107    }
108}
109
110/// A single message in a chat conversation.
111///
112/// Implements [`Default`] for forward-compatible struct-update syntax.
113/// Prefer the `ChatMessage::system` / `::user` / `::assistant` / `::tool`
114/// constructors for clarity.
115#[derive(Debug, Clone, Serialize, Deserialize)]
116pub struct ChatMessage {
117    /// Role of the message sender.
118    pub role: MessageRole,
119    /// Content parts (text, images). Serializes as a plain string when
120    /// containing a single text part for backward compatibility.
121    #[serde(with = "content_vec_serde")]
122    pub content: Vec<ContentPart>,
123}
124
125impl Default for ChatMessage {
126    fn default() -> Self {
127        Self {
128            role: MessageRole::User,
129            content: Vec::new(),
130        }
131    }
132}
133
134impl ChatMessage {
135    /// Create a system message with text content.
136    pub fn system(content: impl Into<String>) -> Self {
137        Self {
138            role: MessageRole::System,
139            content: vec![ContentPart::text(content)],
140        }
141    }
142
143    /// Create a user message with text content.
144    pub fn user(content: impl Into<String>) -> Self {
145        Self {
146            role: MessageRole::User,
147            content: vec![ContentPart::text(content)],
148        }
149    }
150
151    /// Create an assistant message with text content.
152    pub fn assistant(content: impl Into<String>) -> Self {
153        Self {
154            role: MessageRole::Assistant,
155            content: vec![ContentPart::text(content)],
156        }
157    }
158
159    /// Create a tool result message.
160    pub fn tool(content: impl Into<String>) -> Self {
161        Self {
162            role: MessageRole::Tool,
163            content: vec![ContentPart::text(content)],
164        }
165    }
166
167    /// Create a user message with multimodal content parts.
168    pub fn user_multimodal(parts: Vec<ContentPart>) -> Self {
169        Self {
170            role: MessageRole::User,
171            content: parts,
172        }
173    }
174
175    /// Extract all text from this message's content parts.
176    pub fn text_content(&self) -> String {
177        self.content
178            .iter()
179            .filter_map(|p| p.as_text())
180            .collect::<Vec<_>>()
181            .join("")
182    }
183}
184
185/// Configuration for a specific LLM model and provider.
186#[derive(Debug, Clone, Serialize, Deserialize)]
187pub struct ModelConfig {
188    /// Provider name (e.g. `"openai"`, `"anthropic"`).
189    pub provider: String,
190    /// Model identifier (e.g. `"gpt-4o"`, `"claude-sonnet-4-6"`).
191    pub model: String,
192    /// Environment variable name holding the API key.
193    pub api_key_env: String,
194    /// Optional base URL override for the provider API.
195    pub base_url: Option<String>,
196    /// Sampling temperature (0.0–2.0).
197    pub temperature: f32,
198    /// Maximum tokens to generate in the response.
199    pub max_tokens: Option<u32>,
200}
201
202impl Default for ModelConfig {
203    fn default() -> Self {
204        Self {
205            provider: "openai".into(),
206            model: "gpt-4o".into(),
207            api_key_env: "OPENAI_API_KEY".into(),
208            base_url: None,
209            temperature: 0.7,
210            max_tokens: Some(4096),
211        }
212    }
213}
214
215/// Desired output format for the LLM response.
216#[derive(Debug, Clone, Serialize, Deserialize)]
217#[serde(tag = "type", rename_all = "snake_case")]
218pub enum ResponseFormat {
219    /// Plain text response (default).
220    Text,
221    /// JSON object response.
222    Json,
223    /// JSON response conforming to the given schema.
224    JsonSchema {
225        /// JSON Schema the response must satisfy.
226        schema: serde_json::Value,
227    },
228}
229
230/// A chat completion request to an LLM provider.
231///
232/// This struct implements [`Default`] so callers can use struct-update syntax
233/// to stay forward-compatible with future field additions:
234///
235/// ```rust,ignore
236/// let req = LLMRequest {
237///     system: Some("...".into()),
238///     messages: vec![ChatMessage::user("hi")],
239///     ..LLMRequest::default()
240/// };
241/// ```
242///
243/// New fields added to `LLMRequest` in future non-breaking releases are
244/// absorbed by `..LLMRequest::default()` and will not break such call sites
245/// (unlike full struct literals, which must enumerate every field). For the
246/// fluent equivalent, see [`LLMRequest::builder`].
247#[derive(Debug, Clone, Serialize, Deserialize)]
248pub struct LLMRequest {
249    /// Optional system prompt prepended to the conversation.
250    pub system: Option<String>,
251    /// Ordered list of chat messages forming the conversation.
252    pub messages: Vec<ChatMessage>,
253    /// Sampling temperature (0.0–2.0).
254    pub temperature: f32,
255    /// Maximum tokens to generate. `None` uses the provider default.
256    pub max_tokens: Option<u32>,
257    /// Model override for this request. `None` uses the client default.
258    pub model: Option<String>,
259    /// Desired response format. `None` uses the provider default.
260    ///
261    /// Forwarded to the provider by [`OpenAIClient`](crate::llm::OpenAIClient)
262    /// (OpenAI `response_format`) and, for [`ResponseFormat::JsonSchema`], by
263    /// [`AnthropicClient`](crate::llm::AnthropicClient) (Anthropic
264    /// `output_config.format`). [`ResponseFormat::Json`] without a schema has no
265    /// native Anthropic equivalent and is a no-op there.
266    #[serde(skip_serializing_if = "Option::is_none")]
267    pub response_format: Option<ResponseFormat>,
268    /// Tool definitions available to the model for this request.
269    ///
270    /// Forwarded to both OpenAI (`tools` with `type: "function"`) and Anthropic
271    /// (`tools` with `input_schema`). Any tool calls the model makes are returned
272    /// in [`LLMResponse::tool_calls`].
273    #[serde(skip_serializing_if = "Option::is_none")]
274    pub tools: Option<Vec<crate::llm::ToolDefinition>>,
275}
276
277impl Default for LLMRequest {
278    fn default() -> Self {
279        Self {
280            system: None,
281            messages: Vec::new(),
282            // Matches `LLMRequestBuilder::build()`'s `unwrap_or(0.7)` so the two
283            // construction paths agree. Keep these coupled.
284            temperature: 0.7,
285            max_tokens: None,
286            model: None,
287            response_format: None,
288            tools: None,
289        }
290    }
291}
292
293impl LLMRequest {
294    /// Create a new builder for constructing an `LLMRequest`.
295    pub fn builder() -> LLMRequestBuilder {
296        LLMRequestBuilder::default()
297    }
298
299    /// Convert into OpenAI-format messages, consuming the request.
300    ///
301    /// Prepends a system message if `self.system` is set.
302    pub fn into_openai_messages(self) -> Vec<(String, String)> {
303        let mut out = Vec::with_capacity(self.messages.len() + 1);
304        if let Some(system) = self.system {
305            out.push(("system".into(), system));
306        }
307        for msg in self.messages {
308            out.push((msg.role.to_string(), msg.text_content()));
309        }
310        out
311    }
312
313    /// Convert into Anthropic-format messages, consuming the request.
314    ///
315    /// Returns only user/assistant messages (system is handled separately by Anthropic API).
316    pub fn into_anthropic_messages(self) -> Vec<(String, String)> {
317        self.messages
318            .into_iter()
319            .map(|m| (m.role.to_string(), m.text_content()))
320            .collect()
321    }
322}
323
324/// Builder for constructing `LLMRequest` instances with a fluent API.
325#[derive(Debug, Clone, Default)]
326pub struct LLMRequestBuilder {
327    system: Option<String>,
328    messages: Vec<ChatMessage>,
329    temperature: Option<f32>,
330    max_tokens: Option<u32>,
331    model: Option<String>,
332    response_format: Option<ResponseFormat>,
333    tools: Option<Vec<crate::llm::ToolDefinition>>,
334}
335
336impl LLMRequestBuilder {
337    /// Set the system prompt.
338    pub fn system(mut self, prompt: impl Into<String>) -> Self {
339        self.system = Some(prompt.into());
340        self
341    }
342
343    /// Append a user message.
344    pub fn user_message(mut self, content: impl Into<String>) -> Self {
345        self.messages.push(ChatMessage::user(content));
346        self
347    }
348
349    /// Append an assistant message.
350    pub fn assistant_message(mut self, content: impl Into<String>) -> Self {
351        self.messages.push(ChatMessage::assistant(content));
352        self
353    }
354
355    /// Append a raw `ChatMessage`.
356    pub fn message(mut self, msg: ChatMessage) -> Self {
357        self.messages.push(msg);
358        self
359    }
360
361    /// Replace the message list with the provided messages.
362    ///
363    /// Convenience for callers that already hold a `Vec<ChatMessage>` (e.g. a
364    /// pre-built conversation), avoiding repeated `.message()` calls.
365    pub fn messages(mut self, messages: Vec<ChatMessage>) -> Self {
366        self.messages = messages;
367        self
368    }
369
370    /// Set the sampling temperature.
371    pub fn temperature(mut self, temp: f32) -> Self {
372        self.temperature = Some(temp);
373        self
374    }
375
376    /// Set the maximum tokens to generate.
377    pub fn max_tokens(mut self, tokens: u32) -> Self {
378        self.max_tokens = Some(tokens);
379        self
380    }
381
382    /// Set the maximum tokens to generate, or `None` to use the provider default.
383    ///
384    /// Convenience for callers that already hold an `Option<u32>` (e.g. a
385    /// config field), avoiding a conditional chain.
386    pub fn maybe_max_tokens(mut self, tokens: Option<u32>) -> Self {
387        self.max_tokens = tokens;
388        self
389    }
390
391    /// Override the model for this request.
392    pub fn model(mut self, model: impl Into<String>) -> Self {
393        self.model = Some(model.into());
394        self
395    }
396
397    /// Set the desired response format.
398    pub fn response_format(mut self, format: ResponseFormat) -> Self {
399        self.response_format = Some(format);
400        self
401    }
402
403    /// Set the tool definitions available to the model.
404    pub fn tools(mut self, tools: Vec<crate::llm::ToolDefinition>) -> Self {
405        self.tools = Some(tools);
406        self
407    }
408
409    /// Build the `LLMRequest`.
410    pub fn build(self) -> LLMRequest {
411        LLMRequest {
412            system: self.system,
413            messages: self.messages,
414            temperature: self.temperature.unwrap_or(0.7),
415            max_tokens: self.max_tokens,
416            model: self.model,
417            response_format: self.response_format,
418            tools: self.tools,
419        }
420    }
421}
422
423/// A chat completion response from an LLM provider.
424///
425/// Implements [`Default`] for forward-compatible struct-update syntax
426/// (`LLMResponse { ..LLMResponse::default() }`).
427#[derive(Debug, Clone, Default, Serialize, Deserialize)]
428pub struct LLMResponse {
429    /// Generated text content.
430    pub content: String,
431    /// Model that produced this response.
432    pub model: String,
433    /// Token usage statistics.
434    pub usage: TokenUsage,
435    /// Tool calls the model requested this turn.
436    ///
437    /// Empty unless the request supplied [`LLMRequest::tools`] and the model
438    /// chose to call one. Each entry carries the provider-assigned call `id`,
439    /// tool `name`, and JSON-encoded `arguments`.
440    #[serde(default, skip_serializing_if = "Vec::is_empty")]
441    pub tool_calls: Vec<crate::llm::ToolCall>,
442    /// Reason the generation stopped (e.g. `"stop"`, `"length"`, `"tool_calls"`).
443    #[serde(default, skip_serializing_if = "Option::is_none")]
444    pub finish_reason: Option<String>,
445    /// Provider-assigned response ID (useful for logging and deduplication).
446    #[serde(default, skip_serializing_if = "Option::is_none")]
447    pub id: Option<String>,
448    /// Unix timestamp (seconds) when the response was created.
449    #[serde(default, skip_serializing_if = "Option::is_none")]
450    pub created: Option<u64>,
451}
452
453/// Token usage statistics from an LLM response.
454#[derive(Debug, Clone, Default, Serialize, Deserialize)]
455pub struct TokenUsage {
456    /// Number of tokens in the prompt.
457    pub prompt_tokens: u32,
458    /// Number of tokens in the completion.
459    pub completion_tokens: u32,
460    /// Total tokens (prompt + completion).
461    pub total_tokens: u32,
462}
463
464/// A single event in an LLM streaming response.
465#[derive(Debug, Clone)]
466pub enum StreamEvent {
467    /// Partial text content arrived.
468    Delta {
469        /// The partial text chunk.
470        content: String,
471    },
472    /// Final token usage statistics.
473    Usage(TokenUsage),
474    /// Stream has ended.
475    Done,
476}
477
478/// Type alias for a boxed streaming response.
479#[cfg(feature = "client-async")]
480pub type LLMStream =
481    Pin<Box<dyn futures_core::Stream<Item = crate::error::Result<StreamEvent>> + Send>>;
482
483#[cfg(test)]
484mod tests {
485    use super::*;
486
487    #[test]
488    fn message_role_display() {
489        assert_eq!(MessageRole::System.to_string(), "system");
490        assert_eq!(MessageRole::User.to_string(), "user");
491        assert_eq!(MessageRole::Assistant.to_string(), "assistant");
492        assert_eq!(MessageRole::Tool.to_string(), "tool");
493    }
494
495    #[test]
496    fn message_role_serde_roundtrip() {
497        let json = serde_json::to_string(&MessageRole::User).unwrap();
498        assert_eq!(json, "\"user\"");
499        let back: MessageRole = serde_json::from_str(&json).unwrap();
500        assert_eq!(back, MessageRole::User);
501    }
502
503    #[test]
504    fn chat_message_constructors() {
505        let sys = ChatMessage::system("instructions");
506        assert_eq!(sys.role, MessageRole::System);
507
508        let user = ChatMessage::user("hello");
509        assert_eq!(user.role, MessageRole::User);
510
511        let asst = ChatMessage::assistant("hi there");
512        assert_eq!(asst.role, MessageRole::Assistant);
513
514        let tool = ChatMessage::tool("result");
515        assert_eq!(tool.role, MessageRole::Tool);
516    }
517
518    #[test]
519    fn single_text_serializes_as_string() {
520        let msg = ChatMessage::user("hello");
521        let json = serde_json::to_string(&msg).unwrap();
522        assert!(json.contains("\"content\":\"hello\""), "got: {json}");
523    }
524
525    #[test]
526    fn multipart_serializes_as_array() {
527        let msg = ChatMessage::user_multimodal(vec![
528            ContentPart::text("describe this"),
529            ContentPart::image_url("https://example.com/img.png"),
530        ]);
531        let json = serde_json::to_string(&msg).unwrap();
532        assert!(
533            json.contains("\"content\":["),
534            "expected array serialization, got: {json}"
535        );
536    }
537
538    #[test]
539    fn single_text_deserialize_from_string() {
540        let json = r#"{"role":"user","content":"hello"}"#;
541        let msg: ChatMessage = serde_json::from_str(json).unwrap();
542        assert_eq!(msg.role, MessageRole::User);
543        assert_eq!(msg.content.len(), 1);
544        assert_eq!(msg.text_content(), "hello");
545    }
546
547    #[test]
548    fn multipart_deserialize_from_array() {
549        let json = r#"{"role":"user","content":[{"type":"text","text":"hi"},{"type":"image_url","url":"https://x.com/img.png"}]}"#;
550        let msg: ChatMessage = serde_json::from_str(json).unwrap();
551        assert_eq!(msg.content.len(), 2);
552    }
553
554    #[test]
555    fn content_part_text_helper() {
556        let p = ContentPart::text("hello");
557        assert_eq!(p.as_text(), Some("hello"));
558    }
559
560    #[test]
561    fn response_format_json_serialization() {
562        let fmt = ResponseFormat::Json;
563        let json = serde_json::to_string(&fmt).unwrap();
564        assert!(json.contains("\"type\":\"json\""), "got: {json}");
565    }
566
567    #[test]
568    fn response_format_text_serialization() {
569        let fmt = ResponseFormat::Text;
570        let json = serde_json::to_string(&fmt).unwrap();
571        assert!(json.contains("\"type\":\"text\""), "got: {json}");
572    }
573
574    #[test]
575    fn response_format_json_schema() {
576        let fmt = ResponseFormat::JsonSchema {
577            schema: serde_json::json!({"type": "object"}),
578        };
579        let json = serde_json::to_string(&fmt).unwrap();
580        assert!(json.contains("json_schema"), "got: {json}");
581    }
582
583    #[test]
584    fn builder_basic() {
585        let req = LLMRequest::builder()
586            .system("you are helpful")
587            .user_message("hello")
588            .temperature(0.5)
589            .build();
590        assert_eq!(req.system.as_deref(), Some("you are helpful"));
591        assert_eq!(req.messages.len(), 1);
592        assert_eq!(req.temperature, 0.5);
593    }
594
595    #[test]
596    fn builder_with_model_and_format() {
597        let req = LLMRequest::builder()
598            .user_message("test")
599            .model("gpt-4o-mini")
600            .response_format(ResponseFormat::Json)
601            .max_tokens(100)
602            .build();
603        assert_eq!(req.model.as_deref(), Some("gpt-4o-mini"));
604        assert!(matches!(req.response_format, Some(ResponseFormat::Json)));
605        assert_eq!(req.max_tokens, Some(100));
606    }
607
608    #[test]
609    fn builder_with_tools() {
610        use crate::llm::ToolDefinition;
611        let req = LLMRequest::builder()
612            .user_message("what's the weather?")
613            .tools(vec![ToolDefinition {
614                name: "get_weather".into(),
615                description: "Get weather".into(),
616                input_schema: serde_json::json!({"type": "object"}),
617            }])
618            .build();
619        assert!(req.tools.is_some());
620        assert_eq!(req.tools.unwrap().len(), 1);
621    }
622
623    /// `LLMRequest::default()` and `LLMRequest::builder().build()` must agree on
624    /// every field. The temperature default (0.7) in particular is duplicated
625    /// between the manual `Default` impl and the builder's `unwrap_or(0.7)`;
626    /// this test couples them so a future edit to one without the other is caught.
627    #[test]
628    fn default_matches_builder_default() {
629        let from_default = LLMRequest::default();
630        let from_builder = LLMRequest::builder().build();
631        assert_eq!(from_default.temperature, from_builder.temperature);
632        assert_eq!(from_default.temperature, 0.7);
633        assert!(from_default.system.is_none());
634        assert!(from_default.messages.is_empty());
635        assert!(from_default.max_tokens.is_none());
636        assert!(from_default.model.is_none());
637        assert!(from_default.response_format.is_none());
638        assert!(from_default.tools.is_none());
639    }
640
641    #[test]
642    fn builder_messages_setter_replaces_list() {
643        let conv = vec![ChatMessage::user("first"), ChatMessage::assistant("second")];
644        let req = LLMRequest::builder().messages(conv).build();
645        assert_eq!(req.messages.len(), 2);
646        assert_eq!(req.messages[0].role, MessageRole::User);
647        assert_eq!(req.messages[1].role, MessageRole::Assistant);
648    }
649
650    #[test]
651    fn builder_maybe_max_tokens_accepts_option() {
652        // Some — sets the value
653        let req = LLMRequest::builder().maybe_max_tokens(Some(512)).build();
654        assert_eq!(req.max_tokens, Some(512));
655        // None — explicitly defers to provider default
656        let req = LLMRequest::builder().maybe_max_tokens(None).build();
657        assert_eq!(req.max_tokens, None);
658    }
659
660    #[test]
661    fn into_openai_messages_with_system() {
662        let req = LLMRequest::builder()
663            .system("be helpful")
664            .user_message("hi")
665            .assistant_message("hello")
666            .build();
667        let msgs = req.into_openai_messages();
668        assert_eq!(msgs.len(), 3);
669        assert_eq!(msgs[0].0, "system");
670        assert_eq!(msgs[1].0, "user");
671        assert_eq!(msgs[2].0, "assistant");
672    }
673
674    #[test]
675    fn into_anthropic_messages_excludes_system() {
676        let req = LLMRequest::builder()
677            .system("be helpful")
678            .user_message("hi")
679            .build();
680        let msgs = req.into_anthropic_messages();
681        assert_eq!(msgs.len(), 1);
682        assert_eq!(msgs[0].0, "user");
683    }
684
685    #[test]
686    fn text_content_extracts_text() {
687        let msg = ChatMessage::user_multimodal(vec![
688            ContentPart::text("hello "),
689            ContentPart::image_url("http://x.com/i.png"),
690            ContentPart::text("world"),
691        ]);
692        assert_eq!(msg.text_content(), "hello world");
693    }
694}
llm_kernel/llm/types.rs

llm_kernel/llm/
types.rs
