gemini-rs/src/types/common.rs

use std::{fmt::Display, str::FromStr, vec};

use serde::{Deserialize, Serialize};
use serde_json::Value;

#[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| part.text.clone())
                .collect::<String>()
        })
    }

    pub fn builder() -> ContentBuilder {
        ContentBuilder::default()
    }
}

pub struct SystemPrompt {
    content: Content,
}

impl<T: AsRef<str>> From<T> for SystemPrompt {
    fn from(value: T) -> Self {
        SystemPrompt {
            content: Content {
                parts: Some(vec![Part::text(value.as_ref())]),
                role: None,
            },
        }
    }
}

impl From<SystemPrompt> for Content {
    fn from(value: SystemPrompt) -> Self {
        value.content
    }
}

#[derive(Default)]
pub struct ContentBuilder {
    content: Content,
}

impl ContentBuilder {
    pub fn add_text_part<T: Into<String>>(self, text: T) -> Self {
        self.add_part(Part::text(text))
    }

    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, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    User,
    Model,
}

impl Display for Role {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let role_str = match self {
            Role::User => "user",
            Role::Model => "model",
        };
        f.write_str(role_str)
    }
}

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, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Part {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub text: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub inline_data: Option<InlineData>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub file_data: Option<FileData>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub function_call: Option<FunctionCallData>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub function_response: Option<FunctionResponseData>,
    /// Opaque signature emitted by Gemini 2.5+ thinking models on function-call
    /// parts. Must be echoed back on the same part — and on the corresponding
    /// `functionResponse` part — in subsequent turns, otherwise Vertex rejects
    /// the request with "function call X is missing a thought_signature".
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thought_signature: Option<String>,
    /// `Some(true)` on a response part marks it as a reasoning/thinking token
    /// (only emitted when `thinkingConfig.includeThoughts` is set on the
    /// request). Absent on regular text and on outgoing requests.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thought: Option<bool>,
}

impl Part {
    pub fn text(text: impl Into<String>) -> Self {
        Part {
            text: Some(text.into()),
            ..Default::default()
        }
    }

    pub fn function_call(name: impl Into<String>, args: Value) -> Self {
        Part {
            function_call: Some(FunctionCallData {
                name: name.into(),
                args,
            }),
            ..Default::default()
        }
    }

    pub fn function_response(name: impl Into<String>, response: Value) -> Self {
        Part {
            function_response: Some(FunctionResponseData {
                name: name.into(),
                response,
            }),
            ..Default::default()
        }
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct FunctionCallData {
    pub name: String,
    pub args: Value,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct FunctionResponseData {
    pub name: String,
    pub response: Value,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InlineData {
    pub mime_type: String,
    pub data: String,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FileData {
    pub mime_type: String,
    pub file_uri: String,
}