swiftide_core/chat_completion/
chat_completion_request.rs

1use std::{borrow::Cow, collections::BTreeSet};
2
3use derive_builder::Builder;
4
5use super::{chat_message::ChatMessage, tools::ToolSpec, traits::Tool};
6
7/// A chat completion request represents a series of chat messages and tool interactions that can
8/// be send to any LLM.
9#[derive(Builder, Clone, PartialEq, Debug)]
10#[builder(setter(into, strip_option))]
11pub struct ChatCompletionRequest<'a> {
12    pub messages: Cow<'a, [ChatMessage]>,
13    #[builder(default, setter(custom))]
14    pub tools_spec: BTreeSet<ToolSpec>,
15}
16
17impl<'a> ChatCompletionRequest<'a> {
18    pub fn builder() -> ChatCompletionRequestBuilder<'a> {
19        ChatCompletionRequestBuilder::default()
20    }
21
22    /// Returns the chat messages included in the request.
23    pub fn messages(&self) -> &[ChatMessage] {
24        self.messages.as_ref()
25    }
26
27    /// Returns the tool specifications currently attached to the request.
28    pub fn tools_spec(&self) -> &BTreeSet<ToolSpec> {
29        &self.tools_spec
30    }
31
32    /// Returns an owned request with `'static` data.
33    pub fn to_owned(&self) -> ChatCompletionRequest<'static> {
34        ChatCompletionRequest {
35            messages: Cow::Owned(self.messages.iter().map(ChatMessage::to_owned).collect()),
36            tools_spec: self.tools_spec.clone(),
37        }
38    }
39}
40
41impl From<Vec<ChatMessage>> for ChatCompletionRequest<'_> {
42    fn from(messages: Vec<ChatMessage>) -> Self {
43        ChatCompletionRequest {
44            messages: Cow::Owned(messages),
45            tools_spec: BTreeSet::new(),
46        }
47    }
48}
49
50impl<'a> From<&'a [ChatMessage]> for ChatCompletionRequest<'a> {
51    fn from(messages: &'a [ChatMessage]) -> Self {
52        ChatCompletionRequest {
53            messages: Cow::Borrowed(messages),
54            tools_spec: BTreeSet::new(),
55        }
56    }
57}
58
59impl ChatCompletionRequestBuilder<'_> {
60    #[deprecated(note = "Use `tools` with real Tool instances instead")]
61    pub fn tools_spec<I>(&mut self, tools_spec: I) -> &mut Self
62    where
63        I: IntoIterator<Item = ToolSpec>,
64    {
65        self.tools_spec = Some(tools_spec.into_iter().collect());
66        self
67    }
68
69    /// Adds multiple tools by deriving their specs from the provided instances.
70    pub fn tools<I, T>(&mut self, tools: I) -> &mut Self
71    where
72        I: IntoIterator<Item = T>,
73        T: Into<Box<dyn Tool>>,
74    {
75        let specs = tools.into_iter().map(|tool| {
76            let boxed: Box<dyn Tool> = tool.into();
77            boxed.tool_spec()
78        });
79        self.tool_specs(specs)
80    }
81
82    /// Adds a single tool instance to the request by deriving its spec.
83    pub fn tool<T>(&mut self, tool: T) -> &mut Self
84    where
85        T: Into<Box<dyn Tool>>,
86    {
87        let boxed: Box<dyn Tool> = tool.into();
88        self.tool_specs(std::iter::once(boxed.tool_spec()))
89    }
90
91    /// Extends the request with additional tool specifications.
92    pub fn tool_specs<I>(&mut self, specs: I) -> &mut Self
93    where
94        I: IntoIterator<Item = ToolSpec>,
95    {
96        let entry = self.tools_spec.get_or_insert_with(BTreeSet::new);
97        entry.extend(specs);
98        self
99    }
100
101    /// Adds a single chat message to the request
102    pub fn message(&mut self, message: impl Into<ChatMessage>) -> &mut Self {
103        let mut messages = self
104            .messages
105            .take()
106            .map(Cow::into_owned)
107            .unwrap_or_default();
108        messages.push(message.into());
109
110        self.messages = Some(Cow::Owned(messages));
111        self
112    }
113
114    /// Extends the request with multiple chat messages.
115    pub fn messages_iter<I>(&mut self, messages: I) -> &mut Self
116    where
117        I: IntoIterator<Item = ChatMessage>,
118    {
119        let mut new_messages = self
120            .messages
121            .take()
122            .map(Cow::into_owned)
123            .unwrap_or_default();
124        new_messages.extend(messages);
125        self.messages = Some(Cow::Owned(new_messages));
126        self
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::ChatCompletionRequest;
133    use crate::chat_completion::{ChatMessage, ToolSpec};
134    use schemars::Schema;
135    use serde_json::json;
136
137    #[test]
138    fn tool_specs_are_stored_in_deterministic_order() {
139        let zebra = ToolSpec::builder()
140            .name("zebra")
141            .description("later alphabetically")
142            .parameters_schema(schema_from_json(json!({
143                "type": "object",
144                "properties": {
145                    "b": { "type": "string" },
146                    "a": { "type": "string" }
147                }
148            })))
149            .build()
150            .unwrap();
151
152        let alpha = ToolSpec::builder()
153            .name("alpha")
154            .description("earlier alphabetically")
155            .parameters_schema(schema_from_json(json!({
156                "properties": {
157                    "z": { "type": "string" },
158                    "m": { "type": "string" }
159                },
160                "type": "object"
161            })))
162            .build()
163            .unwrap();
164
165        let request = ChatCompletionRequest::builder()
166            .messages(vec![ChatMessage::User("hi".into())])
167            .tool_specs([zebra, alpha])
168            .build()
169            .unwrap();
170
171        let names = request
172            .tools_spec()
173            .iter()
174            .map(|spec| spec.name.as_str())
175            .collect::<Vec<_>>();
176
177        assert_eq!(names, vec!["alpha", "zebra"]);
178    }
179
180    fn schema_from_json(value: serde_json::Value) -> Schema {
181        serde_json::from_value(value).expect("valid schema")
182    }
183}
swiftide_core/chat_completion/chat_completion_request.rs

swiftide_core/chat_completion/
chat_completion_request.rs
