1use std::marker::PhantomData;
2
3use crate::{
4 Client, Result,
5 types::{self, Response},
6};
7
8pub struct Chat<T> {
10 model: Box<str>,
11 client: Client,
12 system_instruction: Option<Box<str>>,
13 safety_settings: Vec<types::SafetySettings>,
14 history: Vec<types::Content>,
15 config: Option<types::GenerationConfig>,
16 phantom: PhantomData<T>,
17}
18
19impl<T> Chat<T> {
20 pub fn new(client: &Client, model: &str) -> Self {
21 Self {
22 model: model.into(),
23 client: client.clone(),
24 system_instruction: None,
25 safety_settings: Vec::new(),
26 history: Vec::new(),
27 config: None,
28 phantom: PhantomData,
29 }
30 }
31
32 pub fn config(&mut self) -> &types::GenerationConfig {
33 self.config.get_or_insert_default()
34 }
35
36 pub fn to_json(mut self) -> Chat<Json> {
37 self.config_mut().response_mime_type = Some("application/json".into());
38 Chat {
39 model: self.model,
40 client: self.client,
41 system_instruction: self.system_instruction,
42 safety_settings: self.safety_settings,
43 history: self.history,
44 config: self.config,
45 phantom: PhantomData,
46 }
47 }
48
49 pub fn config_mut(&mut self) -> &mut types::GenerationConfig {
50 self.config.get_or_insert_default()
51 }
52
53 pub fn history(&self) -> &[types::Content] {
54 &self.history
55 }
56
57 pub fn history_mut(&mut self) -> &mut Vec<types::Content> {
58 &mut self.history
59 }
60
61 pub fn safety_settings(&mut self, safety_settings: Vec<types::SafetySettings>) {
62 self.safety_settings = safety_settings;
63 }
64
65 pub fn system_instruction(mut self, instruction: &str) -> Self {
66 self.system_instruction = Some(Box::from(instruction));
67 self
68 }
69
70 pub async fn generate_content(&mut self) -> Result<Response> {
71 let mut generate_content = self.client.generate_content(&self.model);
72
73 if let Some(system_instruction) = &self.system_instruction {
74 generate_content.system_instruction(system_instruction);
75 }
76
77 if let Some(config) = &self.config {
78 generate_content.config(config.clone());
79 }
80
81 generate_content.contents(self.history.clone());
82 generate_content.safety_settings(self.safety_settings.clone());
83 generate_content.await
84 }
85
86 pub async fn send_message(&mut self, message: &str) -> Result<Response> {
87 self.history.push(types::Content {
88 role: types::Role::User,
89 parts: vec![types::Part::text(message)],
90 });
91
92 self.generate_content().await
93 }
94 pub async fn send_parted_messages(&mut self, parts: Vec<types::Part>) -> Result<Response> {
95 self.history.push(types::Content {
96 role: types::Role::User,
97 parts,
98 });
99
100 self.generate_content().await
101 }
102}
103
104impl Chat<Json> {
105 pub fn response_schema(mut self, schema: types::Schema) -> Self {
106 self.config_mut().response_schema = Some(schema);
107 self
108 }
109
110 pub async fn json<T: serde::de::DeserializeOwned>(&mut self, message: &str) -> Result<T> {
111 let response = self.send_message(message).await?;
112 let json = format!("{response}");
113 serde_json::from_str(&json).map_err(Into::into)
114 }
115}
116
117pub struct Text {}
118
119pub struct Json {}