103 lines
2.1 KiB
Rust
103 lines
2.1 KiB
Rust
use std::{collections::HashMap, str::FromStr, vec};
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
#[derive(Clone, Default, Debug, Serialize, Deserialize)]
|
|
pub struct Content {
|
|
pub role: Option<Role>,
|
|
pub parts: Option<Vec<Part>>,
|
|
}
|
|
|
|
impl Content {
|
|
pub fn get_text(&self) -> Option<String> {
|
|
self.parts.as_ref().map(|parts| {
|
|
parts
|
|
.iter()
|
|
.filter_map(|part| match part {
|
|
Part::Text(text) => Some(text.clone()),
|
|
_ => None,
|
|
})
|
|
.collect::<String>()
|
|
})
|
|
}
|
|
|
|
pub fn builder() -> ContentBuilder {
|
|
ContentBuilder::new()
|
|
}
|
|
}
|
|
|
|
pub struct ContentBuilder {
|
|
content: Content,
|
|
}
|
|
|
|
impl ContentBuilder {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
content: Default::default(),
|
|
}
|
|
}
|
|
|
|
pub fn add_part(mut self, part: Part) -> Self {
|
|
match &mut self.content.parts {
|
|
Some(parts) => parts.push(part),
|
|
None => self.content.parts = Some(vec![part]),
|
|
}
|
|
self
|
|
}
|
|
|
|
pub fn role(mut self, role: Role) -> Self {
|
|
self.content.role = Some(role);
|
|
self
|
|
}
|
|
|
|
pub fn build(self) -> Content {
|
|
self.content
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
|
|
#[serde(rename_all = "lowercase")]
|
|
pub enum Role {
|
|
User,
|
|
Model,
|
|
}
|
|
|
|
impl ToString for Role {
|
|
fn to_string(&self) -> String {
|
|
match self {
|
|
Role::User => "user".to_string(),
|
|
Role::Model => "model".to_string(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl FromStr for Role {
|
|
type Err = ();
|
|
|
|
fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
|
|
match s {
|
|
"user" => Ok(Role::User),
|
|
"model" => Ok(Role::Model),
|
|
_ => Err(()),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, Serialize, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub enum Part {
|
|
Text(String),
|
|
InlineData {
|
|
mime_type: String,
|
|
data: String,
|
|
},
|
|
FileData {
|
|
mime_type: String,
|
|
file_uri: String,
|
|
},
|
|
FunctionCall {
|
|
name: String,
|
|
args: HashMap<String, String>,
|
|
},
|
|
}
|