swiftide_core/chat_completion/
chat_completion_request.rs1use std::collections::HashSet;
2
3use derive_builder::Builder;
4
5use super::{chat_message::ChatMessage, tools::ToolSpec, traits::Tool};
6
7#[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 pub fn messages(&self) -> &[ChatMessage] {
24 self.messages.as_slice()
25 }
26
27 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 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 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 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 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 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}