use std::mem; use anyhow::{Result, anyhow, bail};use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest};use serde::{Deserialize, Deserializer, Serialize, Serializer}; pub const API_URL: &str = "https://generativelanguage.googleapis.com"; pub async fn stream_generate_content(    client: &dyn HttpClient,    api_url: &str,    api_key: &str,    mut request: GenerateContentRequest,) -> Result<BoxStream<'static, Result<GenerateContentResponse>>> {    validate_generate_content_request(&request)?;     // The `model` field is emptied as it is provided as a path parameter.    let model_id = mem::take(&mut request.model.model_id);     let uri =        format!("{api_url}/v1beta/models/{model_id}:streamGenerateContent?alt=sse&key={api_key}",);     let request_builder = HttpRequest::builder()        .method(Method::POST)        .uri(uri)        .header("Content-Type", "application/json");     let request = request_builder.body(AsyncBody::from(serde_json::to_string(&request)?))?;    let mut response = client.send(request).await?;    if response.status().is_success() {        let reader = BufReader::new(response.into_body());        Ok(reader            .lines()            .filter_map(|line| async move {                match line {                    Ok(line) => {                        if let Some(line) = line.strip_prefix("data: ") {                            match serde_json::from_str(line) {                                Ok(response) => Some(Ok(response)),                                Err(error) => Some(Err(anyhow!(format!(                                    "Error parsing JSON: {error:?}\n{line:?}"                                )))),                            }                        } else {                            None                        }                    }                    Err(error) => Some(Err(anyhow!(error))),                }            })            .boxed())    } else {        let mut text = String::new();        response.body_mut().read_to_string(&mut text).await?;        Err(anyhow!(            "error during streamGenerateContent, status code: {:?}, body: {}",            response.status(),            text        ))    }} pub async fn count_tokens(    client: &dyn HttpClient,    api_url: &str,    api_key: &str,    request: CountTokensRequest,) -> Result<CountTokensResponse> {    validate_generate_content_request(&request.generate_content_request)?;     let uri = format!(        "{api_url}/v1beta/models/{model_id}:countTokens?key={api_key}",        model_id = &request.generate_content_request.model.model_id,    );     let request = serde_json::to_string(&request)?;    let request_builder = HttpRequest::builder()        .method(Method::POST)        .uri(&uri)        .header("Content-Type", "application/json");    let http_request = request_builder.body(AsyncBody::from(request))?;     let mut response = client.send(http_request).await?;    let mut text = String::new();    response.body_mut().read_to_string(&mut text).await?;    anyhow::ensure!(        response.status().is_success(),        "error during countTokens, status code: {:?}, body: {}",        response.status(),        text    );    Ok(serde_json::from_str::<CountTokensResponse>(&text)?)} pub fn validate_generate_content_request(request: &GenerateContentRequest) -> Result<()> {    if request.model.is_empty() {        bail!("Model must be specified");    }     if request.contents.is_empty() {        bail!("Request must contain at least one content item");    }     if let Some(user_content) = request        .contents        .iter()        .find(|content| content.role == Role::User)    {        if user_content.parts.is_empty() {            bail!("User content must contain at least one part");        }    }     Ok(())} #[derive(Debug, Serialize, Deserialize)]pub enum Task {    #[serde(rename = "generateContent")]    GenerateContent,    #[serde(rename = "streamGenerateContent")]    StreamGenerateContent,    #[serde(rename = "countTokens")]    CountTokens,    #[serde(rename = "embedContent")]    EmbedContent,    #[serde(rename = "batchEmbedContents")]    BatchEmbedContents,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct GenerateContentRequest {    #[serde(default, skip_serializing_if = "ModelName::is_empty")]    pub model: ModelName,    pub contents: Vec<Content>,    #[serde(skip_serializing_if = "Option::is_none")]    pub system_instruction: Option<SystemInstruction>,    #[serde(skip_serializing_if = "Option::is_none")]    pub generation_config: Option<GenerationConfig>,    #[serde(skip_serializing_if = "Option::is_none")]    pub safety_settings: Option<Vec<SafetySetting>>,    #[serde(skip_serializing_if = "Option::is_none")]    pub tools: Option<Vec<Tool>>,    #[serde(skip_serializing_if = "Option::is_none")]    pub tool_config: Option<ToolConfig>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct GenerateContentResponse {    #[serde(skip_serializing_if = "Option::is_none")]    pub candidates: Option<Vec<GenerateContentCandidate>>,    #[serde(skip_serializing_if = "Option::is_none")]    pub prompt_feedback: Option<PromptFeedback>,    #[serde(skip_serializing_if = "Option::is_none")]    pub usage_metadata: Option<UsageMetadata>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct GenerateContentCandidate {    #[serde(skip_serializing_if = "Option::is_none")]    pub index: Option<usize>,    pub content: Content,    #[serde(skip_serializing_if = "Option::is_none")]    pub finish_reason: Option<String>,    #[serde(skip_serializing_if = "Option::is_none")]    pub finish_message: Option<String>,    #[serde(skip_serializing_if = "Option::is_none")]    pub safety_ratings: Option<Vec<SafetyRating>>,    #[serde(skip_serializing_if = "Option::is_none")]    pub citation_metadata: Option<CitationMetadata>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct Content {    #[serde(default)]    pub parts: Vec<Part>,    pub role: Role,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct SystemInstruction {    pub parts: Vec<Part>,} #[derive(Debug, PartialEq, Deserialize, Serialize)]#[serde(rename_all = "camelCase")]pub enum Role {    User,    Model,} #[derive(Debug, Serialize, Deserialize)]#[serde(untagged)]pub enum Part {    TextPart(TextPart),    InlineDataPart(InlineDataPart),    FunctionCallPart(FunctionCallPart),    FunctionResponsePart(FunctionResponsePart),    ThoughtPart(ThoughtPart),} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct TextPart {    pub text: String,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct InlineDataPart {    pub inline_data: GenerativeContentBlob,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct GenerativeContentBlob {    pub mime_type: String,    pub data: String,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct FunctionCallPart {    pub function_call: FunctionCall,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct FunctionResponsePart {    pub function_response: FunctionResponse,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct ThoughtPart {    pub thought: bool,    pub thought_signature: String,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct CitationSource {    #[serde(skip_serializing_if = "Option::is_none")]    pub start_index: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub end_index: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub uri: Option<String>,    #[serde(skip_serializing_if = "Option::is_none")]    pub license: Option<String>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct CitationMetadata {    pub citation_sources: Vec<CitationSource>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct PromptFeedback {    #[serde(skip_serializing_if = "Option::is_none")]    pub block_reason: Option<String>,    pub safety_ratings: Vec<SafetyRating>,    #[serde(skip_serializing_if = "Option::is_none")]    pub block_reason_message: Option<String>,} #[derive(Debug, Serialize, Deserialize, Default)]#[serde(rename_all = "camelCase")]pub struct UsageMetadata {    #[serde(skip_serializing_if = "Option::is_none")]    pub prompt_token_count: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub cached_content_token_count: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub candidates_token_count: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub tool_use_prompt_token_count: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub thoughts_token_count: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub total_token_count: Option<usize>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct ThinkingConfig {    pub thinking_budget: u32,} #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]pub enum GoogleModelMode {    #[default]    Default,    Thinking {        budget_tokens: Option<u32>,    },} #[derive(Debug, Deserialize, Serialize)]#[serde(rename_all = "camelCase")]pub struct GenerationConfig {    #[serde(skip_serializing_if = "Option::is_none")]    pub candidate_count: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub stop_sequences: Option<Vec<String>>,    #[serde(skip_serializing_if = "Option::is_none")]    pub max_output_tokens: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub temperature: Option<f64>,    #[serde(skip_serializing_if = "Option::is_none")]    pub top_p: Option<f64>,    #[serde(skip_serializing_if = "Option::is_none")]    pub top_k: Option<usize>,    #[serde(skip_serializing_if = "Option::is_none")]    pub thinking_config: Option<ThinkingConfig>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct SafetySetting {    pub category: HarmCategory,    pub threshold: HarmBlockThreshold,} #[derive(Debug, Serialize, Deserialize)]pub enum HarmCategory {    #[serde(rename = "HARM_CATEGORY_UNSPECIFIED")]    Unspecified,    #[serde(rename = "HARM_CATEGORY_DEROGATORY")]    Derogatory,    #[serde(rename = "HARM_CATEGORY_TOXICITY")]    Toxicity,    #[serde(rename = "HARM_CATEGORY_VIOLENCE")]    Violence,    #[serde(rename = "HARM_CATEGORY_SEXUAL")]    Sexual,    #[serde(rename = "HARM_CATEGORY_MEDICAL")]    Medical,    #[serde(rename = "HARM_CATEGORY_DANGEROUS")]    Dangerous,    #[serde(rename = "HARM_CATEGORY_HARASSMENT")]    Harassment,    #[serde(rename = "HARM_CATEGORY_HATE_SPEECH")]    HateSpeech,    #[serde(rename = "HARM_CATEGORY_SEXUALLY_EXPLICIT")]    SexuallyExplicit,    #[serde(rename = "HARM_CATEGORY_DANGEROUS_CONTENT")]    DangerousContent,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "SCREAMING_SNAKE_CASE")]pub enum HarmBlockThreshold {    #[serde(rename = "HARM_BLOCK_THRESHOLD_UNSPECIFIED")]    Unspecified,    BlockLowAndAbove,    BlockMediumAndAbove,    BlockOnlyHigh,    BlockNone,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "SCREAMING_SNAKE_CASE")]pub enum HarmProbability {    #[serde(rename = "HARM_PROBABILITY_UNSPECIFIED")]    Unspecified,    Negligible,    Low,    Medium,    High,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct SafetyRating {    pub category: HarmCategory,    pub probability: HarmProbability,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct CountTokensRequest {    pub generate_content_request: GenerateContentRequest,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct CountTokensResponse {    pub total_tokens: usize,} #[derive(Debug, Serialize, Deserialize)]pub struct FunctionCall {    pub name: String,    pub args: serde_json::Value,} #[derive(Debug, Serialize, Deserialize)]pub struct FunctionResponse {    pub name: String,    pub response: serde_json::Value,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct Tool {    pub function_declarations: Vec<FunctionDeclaration>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct ToolConfig {    pub function_calling_config: FunctionCallingConfig,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "camelCase")]pub struct FunctionCallingConfig {    pub mode: FunctionCallingMode,    #[serde(skip_serializing_if = "Option::is_none")]    pub allowed_function_names: Option<Vec<String>>,} #[derive(Debug, Serialize, Deserialize)]#[serde(rename_all = "lowercase")]pub enum FunctionCallingMode {    Auto,    Any,    None,} #[derive(Debug, Serialize, Deserialize)]pub struct FunctionDeclaration {    pub name: String,    pub description: String,    pub parameters: serde_json::Value,} #[derive(Debug, Default)]pub struct ModelName {    pub model_id: String,} impl ModelName {    pub fn is_empty(&self) -> bool {        self.model_id.is_empty()    }} const MODEL_NAME_PREFIX: &str = "models/"; impl Serialize for ModelName {    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>    where        S: Serializer,    {        serializer.serialize_str(&format!("{MODEL_NAME_PREFIX}{}", &self.model_id))    }} impl<'de> Deserialize<'de> for ModelName {    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>    where        D: Deserializer<'de>,    {        let string = String::deserialize(deserializer)?;        if let Some(id) = string.strip_prefix(MODEL_NAME_PREFIX) {            Ok(Self {                model_id: id.to_string(),            })        } else {            return Err(serde::de::Error::custom(format!(                "Expected model name to begin with {}, got: {}",                MODEL_NAME_PREFIX, string            )));        }    }} #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]#[derive(Clone, Default, Debug, Deserialize, Serialize, PartialEq, Eq, strum::EnumIter)]pub enum Model {    #[serde(rename = "gemini-1.5-pro")]    Gemini15Pro,    #[serde(rename = "gemini-1.5-flash")]    Gemini15Flash,    #[serde(rename = "gemini-2.0-pro-exp")]    Gemini20Pro,    #[serde(rename = "gemini-2.0-flash")]    #[default]    Gemini20Flash,

Lines 1–500 of 596. Load more