openai_ergonomic/builders/
threads.rs

1//! Threads API builders.
2//!
3//! Provides ergonomic builders for creating assistant threads and messages with
4//! attachments and metadata support.
5
6use std::collections::HashMap;
7
8use openai_client_base::models::create_message_request::Role as MessageRole;
9use openai_client_base::models::{
10    assistant_tools_code, assistant_tools_file_search_type_only, AssistantToolsCode,
11    AssistantToolsFileSearchTypeOnly, CreateMessageRequest, CreateMessageRequestAttachmentsInner,
12    CreateMessageRequestAttachmentsInnerToolsInner, CreateThreadRequest,
13};
14use serde_json::Value;
15
16use crate::{Builder, Result};
17
18/// Attachment tools that can be associated with a message file.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum AttachmentTool {
21    /// Make the attachment available to the code interpreter tool.
22    CodeInterpreter,
23    /// Make the attachment available to the file search tool.
24    FileSearch,
25}
26
27impl AttachmentTool {
28    fn to_api(self) -> CreateMessageRequestAttachmentsInnerToolsInner {
29        match self {
30            Self::CodeInterpreter => {
31                CreateMessageRequestAttachmentsInnerToolsInner::AssistantToolsCode(Box::new(
32                    AssistantToolsCode::new(assistant_tools_code::Type::CodeInterpreter),
33                ))
34            }
35            Self::FileSearch => {
36                CreateMessageRequestAttachmentsInnerToolsInner::AssistantToolsFileSearchTypeOnly(
37                    Box::new(AssistantToolsFileSearchTypeOnly::new(
38                        assistant_tools_file_search_type_only::Type::FileSearch,
39                    )),
40                )
41            }
42        }
43    }
44}
45
46/// Attachment to include with a thread message.
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct MessageAttachment {
49    file_id: String,
50    tools: Vec<AttachmentTool>,
51}
52
53impl MessageAttachment {
54    /// Attach a file for the code interpreter tool.
55    #[must_use]
56    pub fn for_code_interpreter(file_id: impl Into<String>) -> Self {
57        Self {
58            file_id: file_id.into(),
59            tools: vec![AttachmentTool::CodeInterpreter],
60        }
61    }
62
63    /// Attach a file for the file search tool.
64    #[must_use]
65    pub fn for_file_search(file_id: impl Into<String>) -> Self {
66        Self {
67            file_id: file_id.into(),
68            tools: vec![AttachmentTool::FileSearch],
69        }
70    }
71
72    /// Add an additional tool that should receive this attachment.
73    #[must_use]
74    pub fn with_tool(mut self, tool: AttachmentTool) -> Self {
75        if !self.tools.contains(&tool) {
76            self.tools.push(tool);
77        }
78        self
79    }
80
81    fn into_api(self) -> CreateMessageRequestAttachmentsInner {
82        let mut inner = CreateMessageRequestAttachmentsInner::new();
83        inner.file_id = Some(self.file_id);
84        if !self.tools.is_empty() {
85            let tools = self.tools.into_iter().map(AttachmentTool::to_api).collect();
86            inner.tools = Some(tools);
87        }
88        inner
89    }
90}
91
92#[derive(Debug, Clone, Default, PartialEq, Eq)]
93enum MetadataState {
94    #[default]
95    Unset,
96    Present(HashMap<String, String>),
97    ExplicitNull,
98}
99
100impl MetadataState {
101    fn upsert(&mut self, key: String, value: String) {
102        match self {
103            MetadataState::Unset | MetadataState::ExplicitNull => {
104                let mut map = HashMap::new();
105                map.insert(key, value);
106                *self = MetadataState::Present(map);
107            }
108            MetadataState::Present(map) => {
109                map.insert(key, value);
110            }
111        }
112    }
113
114    fn replace(&mut self, metadata: HashMap<String, String>) {
115        *self = MetadataState::Present(metadata);
116    }
117
118    fn clear(&mut self) {
119        *self = MetadataState::ExplicitNull;
120    }
121
122    #[allow(clippy::option_option)]
123    fn into_option(self) -> Option<Option<HashMap<String, String>>> {
124        match self {
125            MetadataState::Unset => None,
126            MetadataState::Present(map) if map.is_empty() => None,
127            MetadataState::Present(map) => Some(Some(map)),
128            MetadataState::ExplicitNull => Some(None),
129        }
130    }
131}
132
133/// Builder for messages that seed a thread.
134#[derive(Debug, Clone, Default)]
135pub struct ThreadMessageBuilder {
136    role: MessageRole,
137    content: String,
138    attachments: Vec<MessageAttachment>,
139    metadata: MetadataState,
140}
141
142impl ThreadMessageBuilder {
143    /// Create a user message with the provided text content.
144    #[must_use]
145    pub fn user(content: impl Into<String>) -> Self {
146        Self {
147            role: MessageRole::User,
148            content: content.into(),
149            attachments: Vec::new(),
150            metadata: MetadataState::Unset,
151        }
152    }
153
154    /// Create an assistant-authored message.
155    #[must_use]
156    pub fn assistant(content: impl Into<String>) -> Self {
157        Self {
158            role: MessageRole::Assistant,
159            content: content.into(),
160            attachments: Vec::new(),
161            metadata: MetadataState::Unset,
162        }
163    }
164
165    /// Set the message content explicitly.
166    #[must_use]
167    pub fn content(mut self, content: impl Into<String>) -> Self {
168        self.content = content.into();
169        self
170    }
171
172    /// Attach a file to the message.
173    #[must_use]
174    pub fn attachment(mut self, attachment: MessageAttachment) -> Self {
175        self.attachments.push(attachment);
176        self
177    }
178
179    /// Attach multiple files to the message.
180    #[must_use]
181    pub fn attachments<I>(mut self, attachments: I) -> Self
182    where
183        I: IntoIterator<Item = MessageAttachment>,
184    {
185        self.attachments.extend(attachments);
186        self
187    }
188
189    /// Set metadata for the message.
190    #[must_use]
191    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
192        self.metadata.upsert(key.into(), value.into());
193        self
194    }
195
196    /// Replace metadata with a full map.
197    #[must_use]
198    pub fn metadata_map(mut self, metadata: HashMap<String, String>) -> Self {
199        self.metadata.replace(metadata);
200        self
201    }
202
203    /// Remove metadata by sending an explicit null.
204    #[must_use]
205    pub fn clear_metadata(mut self) -> Self {
206        self.metadata.clear();
207        self
208    }
209}
210
211impl Builder<CreateMessageRequest> for ThreadMessageBuilder {
212    fn build(self) -> Result<CreateMessageRequest> {
213        let mut request = CreateMessageRequest::new(self.role, Value::String(self.content));
214        if !self.attachments.is_empty() {
215            let attachments = self
216                .attachments
217                .into_iter()
218                .map(MessageAttachment::into_api)
219                .collect();
220            request.attachments = Some(Some(attachments));
221        }
222        request.metadata = self.metadata.into_option();
223        Ok(request)
224    }
225}
226
227impl ThreadMessageBuilder {
228    /// Build the message, panicking only if serialization fails (not expected).
229    #[must_use]
230    pub fn finish(self) -> CreateMessageRequest {
231        self.build()
232            .expect("thread message builder should be infallible")
233    }
234}
235
236/// Builder for creating a thread with initial messages and metadata.
237///
238/// # Examples
239///
240/// ```rust
241/// use openai_ergonomic::{Builder};
242/// use openai_ergonomic::builders::threads::{
243///     MessageAttachment, ThreadMessageBuilder, ThreadRequestBuilder,
244/// };
245///
246/// let thread_request = ThreadRequestBuilder::new()
247///     .user_message("Summarise the attached doc")
248///     .message_builder(
249///         ThreadMessageBuilder::assistant("Sure, I'll reference it.")
250///             .attachment(MessageAttachment::for_file_search("file-xyz")),
251///     )
252///     .unwrap()
253///     .build()
254///     .unwrap();
255///
256/// assert!(thread_request.messages.is_some());
257/// ```
258#[derive(Debug, Clone, Default)]
259pub struct ThreadRequestBuilder {
260    messages: Vec<CreateMessageRequest>,
261    metadata: MetadataState,
262}
263
264impl ThreadRequestBuilder {
265    /// Create a new empty thread builder.
266    #[must_use]
267    pub fn new() -> Self {
268        Self::default()
269    }
270
271    /// Seed the thread with an initial user message.
272    #[must_use]
273    pub fn user_message(mut self, content: impl Into<String>) -> Self {
274        self.messages
275            .push(ThreadMessageBuilder::user(content).finish());
276        self
277    }
278
279    /// Seed the thread with an assistant message.
280    #[must_use]
281    pub fn assistant_message(mut self, content: impl Into<String>) -> Self {
282        self.messages
283            .push(ThreadMessageBuilder::assistant(content).finish());
284        self
285    }
286
287    /// Add a fully configured message request.
288    #[must_use]
289    pub fn message_request(mut self, message: CreateMessageRequest) -> Self {
290        self.messages.push(message);
291        self
292    }
293
294    /// Add a thread message builder.
295    pub fn message_builder(mut self, builder: ThreadMessageBuilder) -> Result<Self> {
296        self.messages.push(builder.build()?);
297        Ok(self)
298    }
299
300    /// Add metadata to the thread.
301    #[must_use]
302    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
303        self.metadata.upsert(key.into(), value.into());
304        self
305    }
306
307    /// Replace thread metadata with a full map.
308    #[must_use]
309    pub fn metadata_map(mut self, metadata: HashMap<String, String>) -> Self {
310        self.metadata.replace(metadata);
311        self
312    }
313
314    /// Remove metadata by sending an explicit null.
315    #[must_use]
316    pub fn clear_metadata(mut self) -> Self {
317        self.metadata.clear();
318        self
319    }
320
321    /// Access the configured messages.
322    #[must_use]
323    pub fn messages(&self) -> &[CreateMessageRequest] {
324        &self.messages
325    }
326}
327
328impl Builder<CreateThreadRequest> for ThreadRequestBuilder {
329    fn build(self) -> Result<CreateThreadRequest> {
330        let mut request = CreateThreadRequest::new();
331        if !self.messages.is_empty() {
332            request.messages = Some(self.messages);
333        }
334        request.metadata = self.metadata.into_option();
335        Ok(request)
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn builds_basic_user_message() {
345        let builder = ThreadMessageBuilder::user("Hello");
346        let message = builder.build().expect("builder should succeed");
347
348        assert_eq!(message.role, MessageRole::User);
349        assert_eq!(message.content, Value::String("Hello".to_string()));
350        assert!(message.attachments.is_none());
351        assert!(message.metadata.is_none());
352    }
353
354    #[test]
355    fn builds_message_with_attachment() {
356        let attachment = MessageAttachment::for_code_interpreter("file-123");
357        let message = ThreadMessageBuilder::user("process this")
358            .attachment(attachment)
359            .build()
360            .expect("builder should succeed");
361
362        let attachments = message.attachments.unwrap().unwrap();
363        assert_eq!(attachments.len(), 1);
364        assert_eq!(attachments[0].file_id.as_deref(), Some("file-123"));
365        assert!(attachments[0].tools.as_ref().is_some());
366    }
367
368    #[test]
369    fn builds_thread_with_metadata() {
370        let thread = ThreadRequestBuilder::new()
371            .user_message("Hi there")
372            .metadata("topic", "support")
373            .build()
374            .expect("builder should succeed");
375
376        assert!(thread.messages.is_some());
377        let metadata = thread.metadata.unwrap().unwrap();
378        assert_eq!(metadata.get("topic"), Some(&"support".to_string()));
379    }
380
381    #[test]
382    fn can_explicitly_clear_metadata() {
383        let thread = ThreadRequestBuilder::new()
384            .metadata("foo", "bar")
385            .clear_metadata()
386            .build()
387            .expect("builder should succeed");
388
389        assert!(thread.metadata.is_some());
390        assert!(thread.metadata.unwrap().is_none());
391    }
392
393    #[test]
394    fn accepts_custom_message_builder() {
395        let message_builder = ThreadMessageBuilder::assistant("Hello").metadata("tone", "friendly");
396        let thread = ThreadRequestBuilder::new()
397            .message_builder(message_builder)
398            .expect("builder should succeed")
399            .build()
400            .expect("thread build should succeed");
401
402        let message = thread.messages.unwrap();
403        assert_eq!(message.len(), 1);
404        assert_eq!(message[0].role, MessageRole::Assistant);
405        let metadata = message[0].metadata.clone().unwrap().unwrap();
406        assert_eq!(metadata.get("tone"), Some(&"friendly".to_string()));
407    }
408}
openai_ergonomic/builders/threads.rs

openai_ergonomic/builders/
threads.rs
