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(crate) 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(crate) 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///
326/// # Example
327///
328/// ```no_run
329/// use llm_kernel::llm::LLMRequest;
330///
331/// let request = LLMRequest::builder()
332///     .system("You are concise.")
333///     .user_message("Summarise Rust ownership in one line.")
334///     .temperature(0.0)
335///     .build();
336/// ```
337#[derive(Debug, Clone, Default)]
338pub struct LLMRequestBuilder {
339    system: Option<String>,
340    messages: Vec<ChatMessage>,
341    temperature: Option<f32>,
342    max_tokens: Option<u32>,
343    model: Option<String>,
344    response_format: Option<ResponseFormat>,
345    tools: Option<Vec<crate::llm::ToolDefinition>>,
346}
347
348impl LLMRequestBuilder {
349    /// Set the system prompt.
350    pub fn system(mut self, prompt: impl Into<String>) -> Self {
351        self.system = Some(prompt.into());
352        self
353    }
354
355    /// Append a user message.
356    pub fn user_message(mut self, content: impl Into<String>) -> Self {
357        self.messages.push(ChatMessage::user(content));
358        self
359    }
360
361    /// Append an assistant message.
362    pub fn assistant_message(mut self, content: impl Into<String>) -> Self {
363        self.messages.push(ChatMessage::assistant(content));
364        self
365    }
366
367    /// Append a raw `ChatMessage`.
368    pub fn message(mut self, msg: ChatMessage) -> Self {
369        self.messages.push(msg);
370        self
371    }
372
373    /// Replace the message list with the provided messages.
374    ///
375    /// Convenience for callers that already hold a `Vec<ChatMessage>` (e.g. a
376    /// pre-built conversation), avoiding repeated `.message()` calls.
377    pub fn messages(mut self, messages: Vec<ChatMessage>) -> Self {
378        self.messages = messages;
379        self
380    }
381
382    /// Set the sampling temperature.
383    pub fn temperature(mut self, temp: f32) -> Self {
384        self.temperature = Some(temp);
385        self
386    }
387
388    /// Set the maximum tokens to generate.
389    pub fn max_tokens(mut self, tokens: u32) -> Self {
390        self.max_tokens = Some(tokens);
391        self
392    }
393
394    /// Set the maximum tokens to generate, or `None` to use the provider default.
395    ///
396    /// Convenience for callers that already hold an `Option<u32>` (e.g. a
397    /// config field), avoiding a conditional chain.
398    pub fn maybe_max_tokens(mut self, tokens: Option<u32>) -> Self {
399        self.max_tokens = tokens;
400        self
401    }
402
403    /// Override the model for this request.
404    pub fn model(mut self, model: impl Into<String>) -> Self {
405        self.model = Some(model.into());
406        self
407    }
408
409    /// Set the desired response format.
410    pub fn response_format(mut self, format: ResponseFormat) -> Self {
411        self.response_format = Some(format);
412        self
413    }
414
415    /// Set the tool definitions available to the model.
416    pub fn tools(mut self, tools: Vec<crate::llm::ToolDefinition>) -> Self {
417        self.tools = Some(tools);
418        self
419    }
420
421    /// Build the `LLMRequest`.
422    pub fn build(self) -> LLMRequest {
423        LLMRequest {
424            system: self.system,
425            messages: self.messages,
426            temperature: self.temperature.unwrap_or(0.7),
427            max_tokens: self.max_tokens,
428            model: self.model,
429            response_format: self.response_format,
430            tools: self.tools,
431        }
432    }
433}
434
435/// A chat completion response from an LLM provider.
436///
437/// Implements [`Default`] for forward-compatible struct-update syntax
438/// (`LLMResponse { ..LLMResponse::default() }`).
439#[derive(Debug, Clone, Default, Serialize, Deserialize)]
440pub struct LLMResponse {
441    /// Generated text content.
442    pub content: String,
443    /// Model that produced this response.
444    pub model: String,
445    /// Token usage statistics.
446    pub usage: TokenUsage,
447    /// Tool calls the model requested this turn.
448    ///
449    /// Empty unless the request supplied [`LLMRequest::tools`] and the model
450    /// chose to call one. Each entry carries the provider-assigned call `id`,
451    /// tool `name`, and JSON-encoded `arguments`.
452    #[serde(default, skip_serializing_if = "Vec::is_empty")]
453    pub tool_calls: Vec<crate::llm::ToolCall>,
454    /// Reason the generation stopped (e.g. `"stop"`, `"length"`, `"tool_calls"`).
455    #[serde(default, skip_serializing_if = "Option::is_none")]
456    pub finish_reason: Option<String>,
457    /// Provider-assigned response ID (useful for logging and deduplication).
458    #[serde(default, skip_serializing_if = "Option::is_none")]
459    pub id: Option<String>,
460    /// Unix timestamp (seconds) when the response was created.
461    #[serde(default, skip_serializing_if = "Option::is_none")]
462    pub created: Option<u64>,
463}
464
465/// Token usage statistics from an LLM response.
466#[derive(Debug, Clone, Default, Serialize, Deserialize)]
467pub struct TokenUsage {
468    /// Number of tokens in the prompt.
469    pub prompt_tokens: u32,
470    /// Number of tokens in the completion.
471    pub completion_tokens: u32,
472    /// Total tokens (prompt + completion).
473    pub total_tokens: u32,
474}
475
476/// A single event in an LLM streaming response.
477#[derive(Debug, Clone)]
478pub enum StreamEvent {
479    /// Partial text content arrived.
480    Delta {
481        /// The partial text chunk.
482        content: String,
483    },
484    /// Final token usage statistics.
485    Usage(TokenUsage),
486    /// Stream has ended.
487    Done,
488}
489
490/// Type alias for a boxed streaming response.
491#[cfg(feature = "client-async")]
492pub type LLMStream =
493    Pin<Box<dyn futures_core::Stream<Item = crate::error::Result<StreamEvent>> + Send>>;
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498
499    #[test]
500    fn message_role_display() {
501        assert_eq!(MessageRole::System.to_string(), "system");
502        assert_eq!(MessageRole::User.to_string(), "user");
503        assert_eq!(MessageRole::Assistant.to_string(), "assistant");
504        assert_eq!(MessageRole::Tool.to_string(), "tool");
505    }
506
507    #[test]
508    fn message_role_serde_roundtrip() {
509        let json = serde_json::to_string(&MessageRole::User).unwrap();
510        assert_eq!(json, "\"user\"");
511        let back: MessageRole = serde_json::from_str(&json).unwrap();
512        assert_eq!(back, MessageRole::User);
513    }
514
515    #[test]
516    fn chat_message_constructors() {
517        let sys = ChatMessage::system("instructions");
518        assert_eq!(sys.role, MessageRole::System);
519
520        let user = ChatMessage::user("hello");
521        assert_eq!(user.role, MessageRole::User);
522
523        let asst = ChatMessage::assistant("hi there");
524        assert_eq!(asst.role, MessageRole::Assistant);
525
526        let tool = ChatMessage::tool("result");
527        assert_eq!(tool.role, MessageRole::Tool);
528    }
529
530    #[test]
531    fn single_text_serializes_as_string() {
532        let msg = ChatMessage::user("hello");
533        let json = serde_json::to_string(&msg).unwrap();
534        assert!(json.contains("\"content\":\"hello\""), "got: {json}");
535    }
536
537    #[test]
538    fn multipart_serializes_as_array() {
539        let msg = ChatMessage::user_multimodal(vec![
540            ContentPart::text("describe this"),
541            ContentPart::image_url("https://example.com/img.png"),
542        ]);
543        let json = serde_json::to_string(&msg).unwrap();
544        assert!(
545            json.contains("\"content\":["),
546            "expected array serialization, got: {json}"
547        );
548    }
549
550    #[test]
551    fn single_text_deserialize_from_string() {
552        let json = r#"{"role":"user","content":"hello"}"#;
553        let msg: ChatMessage = serde_json::from_str(json).unwrap();
554        assert_eq!(msg.role, MessageRole::User);
555        assert_eq!(msg.content.len(), 1);
556        assert_eq!(msg.text_content(), "hello");
557    }
558
559    #[test]
560    fn multipart_deserialize_from_array() {
561        let json = r#"{"role":"user","content":[{"type":"text","text":"hi"},{"type":"image_url","url":"https://x.com/img.png"}]}"#;
562        let msg: ChatMessage = serde_json::from_str(json).unwrap();
563        assert_eq!(msg.content.len(), 2);
564    }
565
566    #[test]
567    fn content_part_text_helper() {
568        let p = ContentPart::text("hello");
569        assert_eq!(p.as_text(), Some("hello"));
570    }
571
572    #[test]
573    fn response_format_json_serialization() {
574        let fmt = ResponseFormat::Json;
575        let json = serde_json::to_string(&fmt).unwrap();
576        assert!(json.contains("\"type\":\"json\""), "got: {json}");
577    }
578
579    #[test]
580    fn response_format_text_serialization() {
581        let fmt = ResponseFormat::Text;
582        let json = serde_json::to_string(&fmt).unwrap();
583        assert!(json.contains("\"type\":\"text\""), "got: {json}");
584    }
585
586    #[test]
587    fn response_format_json_schema() {
588        let fmt = ResponseFormat::JsonSchema {
589            schema: serde_json::json!({"type": "object"}),
590        };
591        let json = serde_json::to_string(&fmt).unwrap();
592        assert!(json.contains("json_schema"), "got: {json}");
593    }
594
595    #[test]
596    fn builder_basic() {
597        let req = LLMRequest::builder()
598            .system("you are helpful")
599            .user_message("hello")
600            .temperature(0.5)
601            .build();
602        assert_eq!(req.system.as_deref(), Some("you are helpful"));
603        assert_eq!(req.messages.len(), 1);
604        assert_eq!(req.temperature, 0.5);
605    }
606
607    #[test]
608    fn builder_with_model_and_format() {
609        let req = LLMRequest::builder()
610            .user_message("test")
611            .model("gpt-4o-mini")
612            .response_format(ResponseFormat::Json)
613            .max_tokens(100)
614            .build();
615        assert_eq!(req.model.as_deref(), Some("gpt-4o-mini"));
616        assert!(matches!(req.response_format, Some(ResponseFormat::Json)));
617        assert_eq!(req.max_tokens, Some(100));
618    }
619
620    #[test]
621    fn builder_with_tools() {
622        use crate::llm::ToolDefinition;
623        let req = LLMRequest::builder()
624            .user_message("what's the weather?")
625            .tools(vec![ToolDefinition {
626                name: "get_weather".into(),
627                description: "Get weather".into(),
628                input_schema: serde_json::json!({"type": "object"}),
629            }])
630            .build();
631        assert!(req.tools.is_some());
632        assert_eq!(req.tools.unwrap().len(), 1);
633    }
634
635    /// `LLMRequest::default()` and `LLMRequest::builder().build()` must agree on
636    /// every field. The temperature default (0.7) in particular is duplicated
637    /// between the manual `Default` impl and the builder's `unwrap_or(0.7)`;
638    /// this test couples them so a future edit to one without the other is caught.
639    #[test]
640    fn default_matches_builder_default() {
641        let from_default = LLMRequest::default();
642        let from_builder = LLMRequest::builder().build();
643        assert_eq!(from_default.temperature, from_builder.temperature);
644        assert_eq!(from_default.temperature, 0.7);
645        assert!(from_default.system.is_none());
646        assert!(from_default.messages.is_empty());
647        assert!(from_default.max_tokens.is_none());
648        assert!(from_default.model.is_none());
649        assert!(from_default.response_format.is_none());
650        assert!(from_default.tools.is_none());
651    }
652
653    #[test]
654    fn builder_messages_setter_replaces_list() {
655        let conv = vec![ChatMessage::user("first"), ChatMessage::assistant("second")];
656        let req = LLMRequest::builder().messages(conv).build();
657        assert_eq!(req.messages.len(), 2);
658        assert_eq!(req.messages[0].role, MessageRole::User);
659        assert_eq!(req.messages[1].role, MessageRole::Assistant);
660    }
661
662    #[test]
663    fn builder_maybe_max_tokens_accepts_option() {
664        // Some — sets the value
665        let req = LLMRequest::builder().maybe_max_tokens(Some(512)).build();
666        assert_eq!(req.max_tokens, Some(512));
667        // None — explicitly defers to provider default
668        let req = LLMRequest::builder().maybe_max_tokens(None).build();
669        assert_eq!(req.max_tokens, None);
670    }
671
672    #[test]
673    fn into_openai_messages_with_system() {
674        let req = LLMRequest::builder()
675            .system("be helpful")
676            .user_message("hi")
677            .assistant_message("hello")
678            .build();
679        let msgs = req.into_openai_messages();
680        assert_eq!(msgs.len(), 3);
681        assert_eq!(msgs[0].0, "system");
682        assert_eq!(msgs[1].0, "user");
683        assert_eq!(msgs[2].0, "assistant");
684    }
685
686    #[test]
687    fn into_anthropic_messages_excludes_system() {
688        let req = LLMRequest::builder()
689            .system("be helpful")
690            .user_message("hi")
691            .build();
692        let msgs = req.into_anthropic_messages();
693        assert_eq!(msgs.len(), 1);
694        assert_eq!(msgs[0].0, "user");
695    }
696
697    #[test]
698    fn text_content_extracts_text() {
699        let msg = ChatMessage::user_multimodal(vec![
700            ContentPart::text("hello "),
701            ContentPart::image_url("http://x.com/i.png"),
702            ContentPart::text("world"),
703        ]);
704        assert_eq!(msg.text_content(), "hello world");
705    }
706}
llm_kernel/llm/types.rs

llm_kernel/llm/
types.rs
