#![allow(clippy::enum_variant_names)]
use serde::{Deserialize, Serialize};
use crate::{File, FileHandle, FilesError};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
User,
Model,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum Part {
Text {
text: String,
#[serde(skip_serializing_if = "Option::is_none")]
thought: Option<bool>,
#[serde(rename = "thoughtSignature", skip_serializing_if = "Option::is_none")]
thought_signature: Option<String>,
},
InlineData {
#[serde(rename = "inlineData")]
inline_data: Blob,
#[serde(skip_serializing_if = "Option::is_none")]
media_resolution: Option<super::generation::model::MediaResolution>,
},
FunctionCall {
#[serde(rename = "functionCall")]
function_call: super::tools::FunctionCall,
#[serde(rename = "thoughtSignature", skip_serializing_if = "Option::is_none")]
thought_signature: Option<String>,
},
FunctionResponse {
#[serde(rename = "functionResponse")]
function_response: super::tools::FunctionResponse,
},
FileData {
#[serde(rename = "fileData")]
file_data: FileData,
},
ExecutableCode {
#[serde(rename = "executableCode")]
executable_code: super::tools::ExecutableCode,
},
CodeExecutionResult {
#[serde(rename = "codeExecutionResult")]
code_execution_result: super::tools::CodeExecutionResult,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct FileData {
pub mime_type: String,
pub file_uri: String,
}
impl TryFrom<&FileHandle> for FileData {
type Error = FilesError;
fn try_from(file_handle: &FileHandle) -> Result<Self, Self::Error> {
let File { mime_type, uri, .. } = file_handle.get_file_meta();
let none_fields: Vec<_> = [
mime_type.is_none().then_some("mime_type"),
uri.is_none().then_some("uri"),
]
.into_iter()
.flatten()
.map(String::from)
.collect();
if !none_fields.is_empty() {
return Err(FilesError::Incomplete {
fields: none_fields,
});
}
Ok(Self {
mime_type: mime_type
.as_ref()
.expect("Some-ness checked above")
.to_string(),
file_uri: uri.as_ref().expect("Some-ness checked above").to_string(),
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct Blob {
pub mime_type: String,
pub data: String,
}
impl Blob {
pub fn new(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
Self {
mime_type: mime_type.into(),
data: data.into(),
}
}
}
#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct Content {
#[serde(skip_serializing_if = "Option::is_none")]
pub parts: Option<Vec<Part>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub role: Option<Role>,
}
impl Content {
pub fn text(text: impl Into<String>) -> Self {
Self {
parts: Some(vec![Part::Text {
text: text.into(),
thought: None,
thought_signature: None,
}]),
role: None,
}
}
pub fn function_call(function_call: super::tools::FunctionCall) -> Self {
Self {
parts: Some(vec![Part::FunctionCall {
function_call,
thought_signature: None,
}]),
role: None,
}
}
pub fn function_call_with_thought(
function_call: super::tools::FunctionCall,
thought_signature: impl Into<String>,
) -> Self {
Self {
parts: Some(vec![Part::FunctionCall {
function_call,
thought_signature: Some(thought_signature.into()),
}]),
role: None,
}
}
pub fn text_with_thought_signature(
text: impl Into<String>,
thought_signature: impl Into<String>,
) -> Self {
Self {
parts: Some(vec![Part::Text {
text: text.into(),
thought: None,
thought_signature: Some(thought_signature.into()),
}]),
role: None,
}
}
pub fn thought_with_signature(
text: impl Into<String>,
thought_signature: impl Into<String>,
) -> Self {
Self {
parts: Some(vec![Part::Text {
text: text.into(),
thought: Some(true),
thought_signature: Some(thought_signature.into()),
}]),
role: None,
}
}
pub fn function_response(function_response: super::tools::FunctionResponse) -> Self {
Self {
parts: Some(vec![Part::FunctionResponse { function_response }]),
role: None,
}
}
pub fn function_response_json(name: impl Into<String>, response: serde_json::Value) -> Self {
Self {
parts: Some(vec![Part::FunctionResponse {
function_response: super::tools::FunctionResponse::new(name, response),
}]),
role: None,
}
}
pub fn inline_data(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
Self {
parts: Some(vec![Part::InlineData {
inline_data: Blob::new(mime_type, data),
media_resolution: None,
}]),
role: None,
}
}
pub fn inline_data_with_resolution(
mime_type: impl Into<String>,
data: impl Into<String>,
resolution: super::generation::model::MediaResolutionLevel,
) -> Self {
Self {
parts: Some(vec![Part::InlineData {
inline_data: Blob::new(mime_type, data),
media_resolution: Some(super::generation::model::MediaResolution {
level: resolution,
}),
}]),
role: None,
}
}
pub fn text_with_file(
text: impl Into<String>,
file_handle: &FileHandle,
) -> Result<Self, FilesError> {
Ok(Self {
parts: Some(vec![
Part::Text {
text: text.into(),
thought: None,
thought_signature: None,
},
Part::FileData {
file_data: FileData::try_from(file_handle)?,
},
]),
role: None,
})
}
pub fn with_role(mut self, role: Role) -> Self {
self.role = Some(role);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub content: Content,
pub role: Role,
}
impl Message {
pub fn user(text: impl Into<String>) -> Self {
Self {
content: Content::text(text).with_role(Role::User),
role: Role::User,
}
}
pub fn model(text: impl Into<String>) -> Self {
Self {
content: Content::text(text).with_role(Role::Model),
role: Role::Model,
}
}
pub fn embed(text: impl Into<String>) -> Self {
Self {
content: Content::text(text),
role: Role::Model,
}
}
pub fn function(name: impl Into<String>, response: serde_json::Value) -> Self {
Self {
content: Content::function_response_json(name, response).with_role(Role::Model),
role: Role::Model,
}
}
pub fn function_str(
name: impl Into<String>,
response: impl Into<String>,
) -> Result<Self, serde_json::Error> {
let response_str = response.into();
let json = serde_json::from_str(&response_str)?;
Ok(Self {
content: Content::function_response_json(name, json).with_role(Role::Model),
role: Role::Model,
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum Modality {
ModalityUnspecified,
Document,
Text,
Image,
Audio,
Video,
#[serde(untagged)]
Other(String),
}