swiftide_core/chat_completion/
chat_completion_request.rs

1use std::collections::HashSet;
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 {
12    pub messages: Vec<ChatMessage>,
13    #[builder(default, setter(custom))]
14    pub tools_spec: HashSet<ToolSpec>,
15}
16
17impl ChatCompletionRequest {
18    pub fn builder() -> ChatCompletionRequestBuilder {
19        ChatCompletionRequestBuilder::default()
20    }
21
22    /// Returns the chat messages included in the request.
23    pub fn messages(&self) -> &[ChatMessage] {
24        self.messages.as_slice()
25    }
26
27    /// Returns the tool specifications currently attached to the request.
28    pub fn tools_spec(&self) -> &HashSet<ToolSpec> {
29        &self.tools_spec
30    }
31}
32
33impl From<Vec<ChatMessage>> for ChatCompletionRequest {
34    fn from(messages: Vec<ChatMessage>) -> Self {
35        ChatCompletionRequest {
36            messages,
37            tools_spec: HashSet::new(),
38        }
39    }
40}
41
42impl ChatCompletionRequestBuilder {
43    #[deprecated(note = "Use `tools` with real Tool instances instead")]
44    pub fn tools_spec(&mut self, tools_spec: HashSet<ToolSpec>) -> &mut Self {
45        self.tools_spec = Some(tools_spec);
46        self
47    }
48
49    /// Adds multiple tools by deriving their specs from the provided instances.
50    pub fn tools<I, T>(&mut self, tools: I) -> &mut Self
51    where
52        I: IntoIterator<Item = T>,
53        T: Into<Box<dyn Tool>>,
54    {
55        let specs = tools.into_iter().map(|tool| {
56            let boxed: Box<dyn Tool> = tool.into();
57            boxed.tool_spec()
58        });
59        self.tool_specs(specs)
60    }
61
62    /// Adds a single tool instance to the request by deriving its spec.
63    pub fn tool<T>(&mut self, tool: T) -> &mut Self
64    where
65        T: Into<Box<dyn Tool>>,
66    {
67        let boxed: Box<dyn Tool> = tool.into();
68        self.tool_specs(std::iter::once(boxed.tool_spec()))
69    }
70
71    /// Extends the request with additional tool specifications.
72    pub fn tool_specs<I>(&mut self, specs: I) -> &mut Self
73    where
74        I: IntoIterator<Item = ToolSpec>,
75    {
76        let entry = self.tools_spec.get_or_insert_with(HashSet::new);
77        entry.extend(specs);
78        self
79    }
80
81    /// Appends a single chat message to the request.
82    pub fn message(&mut self, message: impl Into<ChatMessage>) -> &mut Self {
83        self.messages
84            .get_or_insert_with(Vec::new)
85            .push(message.into());
86        self
87    }
88
89    /// Extends the request with multiple chat messages.
90    pub fn messages_iter<I>(&mut self, messages: I) -> &mut Self
91    where
92        I: IntoIterator<Item = ChatMessage>,
93    {
94        let entry = self.messages.get_or_insert_with(Vec::new);
95        entry.extend(messages);
96        self
97    }
98}