artificial_core/provider/
chat_complete.rs1use std::{future::Future, pin::Pin};
2
3use crate::{
4 error::Result,
5 generic::{GenericChatCompletionResponse, GenericFunctionSpec, GenericMessage},
6 model::Model,
7};
8use futures_core::stream::Stream;
9
10pub trait ChatCompletionProvider: Send + Sync {
21 type Message: Send + Sync + 'static;
23
24 fn chat_complete<'p, M>(
27 &self,
28 params: ChatCompleteParameters<M>,
29 ) -> Pin<
30 Box<dyn Future<Output = Result<GenericChatCompletionResponse<GenericMessage>>> + Send + 'p>,
31 >
32 where
33 M: Into<Self::Message> + Send + Sync + 'p;
34}
35
36pub trait StreamingChatProvider: ChatCompletionProvider {
43 type Delta<'s>: Stream<Item = Result<String>> + Send + 's
46 where
47 Self: 's;
48
49 fn chat_complete_stream<'p, M>(&self, params: ChatCompleteParameters<M>) -> Self::Delta<'p>
51 where
52 M: Into<Self::Message> + Send + Sync + 'p;
53}
54
55pub struct ChatCompleteParameters<M> {
56 pub messages: Vec<M>,
57 pub model: Model,
58 pub tools: Option<Vec<GenericFunctionSpec>>,
59 pub temperature: Option<f64>,
60 pub response_format: Option<serde_json::Value>,
61}
62
63impl<M> ChatCompleteParameters<M> {
64 pub fn new(messages: Vec<M>, model: Model) -> Self {
65 Self {
66 messages,
67 model,
68 tools: None,
69 temperature: None,
70 response_format: None,
71 }
72 }
73
74 pub fn messages(&self) -> &Vec<M> {
75 &self.messages
76 }
77
78 pub fn model(&self) -> Model {
79 self.model.clone()
80 }
81
82 pub fn tools(&self) -> Option<&Vec<GenericFunctionSpec>> {
83 self.tools.as_ref()
84 }
85
86 pub fn with_temperature(mut self, temperature: f64) -> Self {
87 self.temperature = Some(temperature);
88 self
89 }
90
91 pub fn with_response_format(mut self, response_format: serde_json::Value) -> Self {
92 self.response_format = Some(response_format);
93 self
94 }
95
96 pub fn with_tools(mut self, tools: Vec<GenericFunctionSpec>) -> Self {
97 self.tools = Some(tools);
98 self
99 }
100}