logiguard fork v3: full patch set on verified 8c74db0 tree

Includes prior-session patches (carry forward so the app compiles):
  - crates/gpui/build.rs: cross-compile manifest fix
  - crates/gpui/src/platform.rs: PlatformWindow::activate_with_token trait method
  - crates/gpui/src/window.rs: Window::activate_with_token public API
  - crates/gpui_linux/src/linux/wayland/window.rs: WaylandWindow::activate_with_token + activate() keyboard-serial fix

Plus the focus-serial fix:
  - serial.rs: SerialKind::KeyboardEnter
  - client.rs: store wl_keyboard.enter serial; latest_serial_of()

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Mohamad Khani
2026-07-14 01:52:12 +03:30
commit b9819977a5
3984 changed files with 1487015 additions and 0 deletions

View File

@@ -0,0 +1,26 @@
[package]
name = "google_ai"
version = "0.1.0"
edition.workspace = true
publish.workspace = true
license = "GPL-3.0-or-later"
[lints]
workspace = true
[lib]
path = "src/google_ai.rs"
[features]
schemars = ["dep:schemars"]
[dependencies]
anyhow.workspace = true
futures.workspace = true
http_client.workspace = true
language_model_core.workspace = true
log.workspace = true
schemars = { workspace = true, optional = true }
serde.workspace = true
serde_json.workspace = true
strum.workspace = true

View File

@@ -0,0 +1 @@
../../LICENSE-GPL

View File

@@ -0,0 +1,470 @@
use anyhow::Result;
use futures::{Stream, StreamExt};
use language_model_core::{
LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelRequest,
LanguageModelToolChoice, LanguageModelToolUse, LanguageModelToolUseId, MessageContent, Role,
StopReason, TokenUsage,
};
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{self, AtomicU64};
use crate::{
Content, FunctionCallingConfig, FunctionCallingMode, FunctionDeclaration,
GenerateContentResponse, GenerationConfig, GenerativeContentBlob, GoogleModelMode,
InlineDataPart, ModelName, Part, SystemInstruction, TextPart, ThinkingConfig, ToolConfig,
UsageMetadata,
};
pub fn into_google(
mut request: LanguageModelRequest,
model_id: String,
mode: GoogleModelMode,
) -> crate::GenerateContentRequest {
fn map_content(content: Vec<MessageContent>) -> Vec<Part> {
content
.into_iter()
.flat_map(|content| match content {
MessageContent::Text(text) => {
if !text.is_empty() {
vec![Part::TextPart(TextPart { text })]
} else {
vec![]
}
}
MessageContent::Thinking {
text: _,
signature: Some(signature),
} => {
if !signature.is_empty() {
vec![Part::ThoughtPart(crate::ThoughtPart {
thought: true,
thought_signature: signature,
})]
} else {
vec![]
}
}
MessageContent::Thinking { .. } => {
vec![]
}
MessageContent::RedactedThinking(_) => vec![],
MessageContent::Image(image) => {
vec![Part::InlineDataPart(InlineDataPart {
inline_data: GenerativeContentBlob {
mime_type: "image/png".to_string(),
data: image.source.to_string(),
},
})]
}
MessageContent::ToolUse(tool_use) => {
// Normalize empty string signatures to None
let thought_signature = tool_use.thought_signature.filter(|s| !s.is_empty());
vec![Part::FunctionCallPart(crate::FunctionCallPart {
function_call: crate::FunctionCall {
name: tool_use.name.to_string(),
args: tool_use.input,
},
thought_signature,
})]
}
MessageContent::ToolResult(tool_result) => {
let mut text_output = String::new();
let mut images: Vec<InlineDataPart> = Vec::new();
for part in tool_result.content {
match part {
language_model_core::LanguageModelToolResultContent::Text(text) => {
text_output.push_str(&text);
}
language_model_core::LanguageModelToolResultContent::Image(image) => {
images.push(InlineDataPart {
inline_data: GenerativeContentBlob {
mime_type: "image/png".to_string(),
data: image.source.to_string(),
},
});
}
}
}
let output = if text_output.is_empty() && !images.is_empty() {
"Tool responded with an image".to_string()
} else {
text_output
};
let mut parts = vec![Part::FunctionResponsePart(crate::FunctionResponsePart {
function_response: crate::FunctionResponse {
name: tool_result.tool_name.to_string(),
// The API expects a valid JSON object
response: serde_json::json!({
"output": output
}),
},
})];
parts.extend(images.into_iter().map(Part::InlineDataPart));
parts
}
})
.collect()
}
let system_instructions = if request
.messages
.first()
.is_some_and(|msg| matches!(msg.role, Role::System))
{
let message = request.messages.remove(0);
Some(SystemInstruction {
parts: map_content(message.content),
})
} else {
None
};
crate::GenerateContentRequest {
model: ModelName { model_id },
system_instruction: system_instructions,
contents: request
.messages
.into_iter()
.filter_map(|message| {
let parts = map_content(message.content);
if parts.is_empty() {
None
} else {
Some(Content {
parts,
role: match message.role {
Role::User => crate::Role::User,
Role::Assistant => crate::Role::Model,
Role::System => crate::Role::User, // Google AI doesn't have a system role
},
})
}
})
.collect(),
generation_config: Some(GenerationConfig {
candidate_count: Some(1),
stop_sequences: Some(request.stop),
max_output_tokens: None,
temperature: request.temperature.map(|t| t as f64).or(Some(1.0)),
thinking_config: match (request.thinking_allowed, mode) {
(true, GoogleModelMode::Thinking { budget_tokens }) => {
budget_tokens.map(|thinking_budget| ThinkingConfig { thinking_budget })
}
_ => None,
},
top_p: None,
top_k: None,
}),
safety_settings: None,
tools: (!request.tools.is_empty()).then(|| {
vec![crate::Tool {
function_declarations: request
.tools
.into_iter()
.map(|tool| FunctionDeclaration {
name: tool.name,
description: tool.description,
parameters: tool.input_schema,
})
.collect(),
}]
}),
tool_config: request.tool_choice.map(|choice| ToolConfig {
function_calling_config: FunctionCallingConfig {
mode: match choice {
LanguageModelToolChoice::Auto => FunctionCallingMode::Auto,
LanguageModelToolChoice::Any => FunctionCallingMode::Any,
LanguageModelToolChoice::None => FunctionCallingMode::None,
},
allowed_function_names: None,
},
}),
}
}
pub struct GoogleEventMapper {
usage: UsageMetadata,
stop_reason: StopReason,
}
impl GoogleEventMapper {
pub fn new() -> Self {
Self {
usage: UsageMetadata::default(),
stop_reason: StopReason::EndTurn,
}
}
pub fn map_stream(
mut self,
events: Pin<Box<dyn Send + Stream<Item = Result<GenerateContentResponse>>>>,
) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
{
events
.map(Some)
.chain(futures::stream::once(async { None }))
.flat_map(move |event| {
futures::stream::iter(match event {
Some(Ok(event)) => self.map_event(event),
Some(Err(error)) => {
vec![Err(LanguageModelCompletionError::from(error))]
}
None => vec![Ok(LanguageModelCompletionEvent::Stop(self.stop_reason))],
})
})
}
pub fn map_event(
&mut self,
event: GenerateContentResponse,
) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
static TOOL_CALL_COUNTER: AtomicU64 = AtomicU64::new(0);
let mut events: Vec<_> = Vec::new();
let mut wants_to_use_tool = false;
if let Some(usage_metadata) = event.usage_metadata {
update_usage(&mut self.usage, &usage_metadata);
events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(
convert_usage(&self.usage),
)))
}
if let Some(prompt_feedback) = event.prompt_feedback
&& let Some(block_reason) = prompt_feedback.block_reason.as_deref()
{
self.stop_reason = match block_reason {
"SAFETY" | "OTHER" | "BLOCKLIST" | "PROHIBITED_CONTENT" | "IMAGE_SAFETY" => {
StopReason::Refusal
}
_ => {
log::error!("Unexpected Google block_reason: {block_reason}");
StopReason::Refusal
}
};
events.push(Ok(LanguageModelCompletionEvent::Stop(self.stop_reason)));
return events;
}
if let Some(candidates) = event.candidates {
for candidate in candidates {
if let Some(finish_reason) = candidate.finish_reason.as_deref() {
self.stop_reason = match finish_reason {
"STOP" => StopReason::EndTurn,
"MAX_TOKENS" => StopReason::MaxTokens,
_ => {
log::error!("Unexpected google finish_reason: {finish_reason}");
StopReason::EndTurn
}
};
}
candidate
.content
.parts
.into_iter()
.for_each(|part| match part {
Part::TextPart(text_part) => {
events.push(Ok(LanguageModelCompletionEvent::Text(text_part.text)))
}
Part::InlineDataPart(_) => {}
Part::FunctionCallPart(function_call_part) => {
wants_to_use_tool = true;
let name: Arc<str> = function_call_part.function_call.name.into();
let next_tool_id =
TOOL_CALL_COUNTER.fetch_add(1, atomic::Ordering::SeqCst);
let id: LanguageModelToolUseId =
format!("{}-{}", name, next_tool_id).into();
// Normalize empty string signatures to None
let thought_signature = function_call_part
.thought_signature
.filter(|s| !s.is_empty());
events.push(Ok(LanguageModelCompletionEvent::ToolUse(
LanguageModelToolUse {
id,
name,
is_input_complete: true,
raw_input: function_call_part.function_call.args.to_string(),
input: function_call_part.function_call.args,
thought_signature,
},
)));
}
Part::FunctionResponsePart(_) => {}
Part::ThoughtPart(part) => {
events.push(Ok(LanguageModelCompletionEvent::Thinking {
text: "(Encrypted thought)".to_string(), // TODO: Can we populate this from thought summaries?
signature: Some(part.thought_signature),
}));
}
});
}
}
// Even when Gemini wants to use a Tool, the API
// responds with `finish_reason: STOP`
if wants_to_use_tool {
self.stop_reason = StopReason::ToolUse;
events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)));
}
events
}
}
fn update_usage(usage: &mut UsageMetadata, new: &UsageMetadata) {
if let Some(prompt_token_count) = new.prompt_token_count {
usage.prompt_token_count = Some(prompt_token_count);
}
if let Some(cached_content_token_count) = new.cached_content_token_count {
usage.cached_content_token_count = Some(cached_content_token_count);
}
if let Some(candidates_token_count) = new.candidates_token_count {
usage.candidates_token_count = Some(candidates_token_count);
}
if let Some(tool_use_prompt_token_count) = new.tool_use_prompt_token_count {
usage.tool_use_prompt_token_count = Some(tool_use_prompt_token_count);
}
if let Some(thoughts_token_count) = new.thoughts_token_count {
usage.thoughts_token_count = Some(thoughts_token_count);
}
if let Some(total_token_count) = new.total_token_count {
usage.total_token_count = Some(total_token_count);
}
}
fn convert_usage(usage: &UsageMetadata) -> TokenUsage {
let prompt_tokens = usage.prompt_token_count.unwrap_or(0);
let cached_tokens = usage.cached_content_token_count.unwrap_or(0);
let input_tokens = prompt_tokens - cached_tokens;
let output_tokens = usage.candidates_token_count.unwrap_or(0);
TokenUsage {
input_tokens,
output_tokens,
cache_read_input_tokens: cached_tokens,
cache_creation_input_tokens: 0,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
Content, FunctionCall, FunctionCallPart, GenerateContentCandidate, GenerateContentResponse,
Part, Role as GoogleRole,
};
use serde_json::json;
#[test]
fn test_function_call_with_signature_creates_tool_use_with_signature() {
let mut mapper = GoogleEventMapper::new();
let response = GenerateContentResponse {
candidates: Some(vec![GenerateContentCandidate {
index: Some(0),
content: Content {
parts: vec![Part::FunctionCallPart(FunctionCallPart {
function_call: FunctionCall {
name: "test_function".to_string(),
args: json!({"arg": "value"}),
},
thought_signature: Some("test_signature_123".to_string()),
})],
role: GoogleRole::Model,
},
finish_reason: None,
finish_message: None,
safety_ratings: None,
citation_metadata: None,
}]),
prompt_feedback: None,
usage_metadata: None,
};
let events = mapper.map_event(response);
assert_eq!(events.len(), 2);
if let Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) = &events[0] {
assert_eq!(tool_use.name.as_ref(), "test_function");
assert_eq!(
tool_use.thought_signature.as_deref(),
Some("test_signature_123")
);
} else {
panic!("Expected ToolUse event");
}
}
#[test]
fn test_function_call_without_signature_has_none() {
let mut mapper = GoogleEventMapper::new();
let response = GenerateContentResponse {
candidates: Some(vec![GenerateContentCandidate {
index: Some(0),
content: Content {
parts: vec![Part::FunctionCallPart(FunctionCallPart {
function_call: FunctionCall {
name: "test_function".to_string(),
args: json!({"arg": "value"}),
},
thought_signature: None,
})],
role: GoogleRole::Model,
},
finish_reason: None,
finish_message: None,
safety_ratings: None,
citation_metadata: None,
}]),
prompt_feedback: None,
usage_metadata: None,
};
let events = mapper.map_event(response);
assert_eq!(events.len(), 2);
if let Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) = &events[0] {
assert!(tool_use.thought_signature.is_none());
} else {
panic!("Expected ToolUse event");
}
}
#[test]
fn test_empty_string_signature_normalized_to_none() {
let mut mapper = GoogleEventMapper::new();
let response = GenerateContentResponse {
candidates: Some(vec![GenerateContentCandidate {
index: Some(0),
content: Content {
parts: vec![Part::FunctionCallPart(FunctionCallPart {
function_call: FunctionCall {
name: "test_function".to_string(),
args: json!({"arg": "value"}),
},
thought_signature: Some("".to_string()),
})],
role: GoogleRole::Model,
},
finish_reason: None,
finish_message: None,
safety_ratings: None,
citation_metadata: None,
}]),
prompt_feedback: None,
usage_metadata: None,
};
let events = mapper.map_event(response);
if let Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) = &events[0] {
assert!(tool_use.thought_signature.is_none());
} else {
panic!("Expected ToolUse event");
}
}
}

View File

@@ -0,0 +1,688 @@
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};
pub use language_model_core::ModelMode as GoogleModelMode;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub mod completion;
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>>> {
let api_key = api_key.trim();
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 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)
&& 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 = "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,
/// Thought signature returned by the model for function calls.
/// Only present on the first function call in parallel call scenarios.
#[serde(skip_serializing_if = "Option::is_none")]
pub thought_signature: Option<String>,
}
#[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: Option<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<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cached_content_token_count: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub candidates_token_count: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_use_prompt_token_count: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thoughts_token_count: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub total_token_count: Option<u64>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ThinkingConfig {
pub thinking_budget: 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)]
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 {
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-2.5-flash-lite",
alias = "gemini-2.5-flash-lite-preview-06-17",
alias = "gemini-2.0-flash-lite-preview"
)]
Gemini25FlashLite,
#[serde(
rename = "gemini-2.5-flash",
alias = "gemini-2.0-flash-thinking-exp",
alias = "gemini-2.5-flash-preview-04-17",
alias = "gemini-2.5-flash-preview-05-20",
alias = "gemini-2.5-flash-preview-latest",
alias = "gemini-2.0-flash"
)]
#[default]
Gemini25Flash,
#[serde(
rename = "gemini-2.5-pro",
alias = "gemini-2.0-pro-exp",
alias = "gemini-2.5-pro-preview-latest",
alias = "gemini-2.5-pro-exp-03-25",
alias = "gemini-2.5-pro-preview-03-25",
alias = "gemini-2.5-pro-preview-05-06",
alias = "gemini-2.5-pro-preview-06-05"
)]
Gemini25Pro,
#[serde(rename = "gemini-3.1-flash-lite")]
Gemini31FlashLite,
#[serde(rename = "gemini-3-flash-preview")]
Gemini3Flash,
#[serde(rename = "gemini-3.1-pro-preview", alias = "gemini-3-pro-preview")]
Gemini31Pro,
#[serde(rename = "custom")]
Custom {
name: String,
/// The name displayed in the UI, such as in the agent panel model dropdown menu.
display_name: Option<String>,
max_tokens: u64,
#[serde(default)]
mode: GoogleModelMode,
},
}
impl Model {
pub fn default_fast() -> Self {
Self::Gemini31FlashLite
}
pub fn id(&self) -> &str {
match self {
Self::Gemini25FlashLite => "gemini-2.5-flash-lite",
Self::Gemini25Flash => "gemini-2.5-flash",
Self::Gemini25Pro => "gemini-2.5-pro",
Self::Gemini31FlashLite => "gemini-3.1-flash-lite",
Self::Gemini3Flash => "gemini-3-flash-preview",
Self::Gemini31Pro => "gemini-3.1-pro-preview",
Self::Custom { name, .. } => name,
}
}
pub fn request_id(&self) -> &str {
match self {
Self::Gemini25FlashLite => "gemini-2.5-flash-lite",
Self::Gemini25Flash => "gemini-2.5-flash",
Self::Gemini25Pro => "gemini-2.5-pro",
Self::Gemini31FlashLite => "gemini-3.1-flash-lite",
Self::Gemini3Flash => "gemini-3-flash-preview",
Self::Gemini31Pro => "gemini-3.1-pro-preview",
Self::Custom { name, .. } => name,
}
}
pub fn display_name(&self) -> &str {
match self {
Self::Gemini25FlashLite => "Gemini 2.5 Flash-Lite",
Self::Gemini25Flash => "Gemini 2.5 Flash",
Self::Gemini25Pro => "Gemini 2.5 Pro",
Self::Gemini31FlashLite => "Gemini 3.1 Flash Lite",
Self::Gemini3Flash => "Gemini 3 Flash",
Self::Gemini31Pro => "Gemini 3.1 Pro",
Self::Custom {
name, display_name, ..
} => display_name.as_ref().unwrap_or(name),
}
}
pub fn max_token_count(&self) -> u64 {
match self {
Self::Gemini25FlashLite
| Self::Gemini25Flash
| Self::Gemini25Pro
| Self::Gemini31FlashLite
| Self::Gemini3Flash
| Self::Gemini31Pro => 1_048_576,
Self::Custom { max_tokens, .. } => *max_tokens,
}
}
pub fn max_output_tokens(&self) -> Option<u64> {
match self {
Model::Gemini25FlashLite
| Model::Gemini25Flash
| Model::Gemini25Pro
| Model::Gemini31FlashLite
| Model::Gemini3Flash
| Model::Gemini31Pro => Some(65_536),
Model::Custom { .. } => None,
}
}
pub fn supports_tools(&self) -> bool {
true
}
pub fn supports_images(&self) -> bool {
true
}
pub fn mode(&self) -> GoogleModelMode {
match self {
Self::Gemini25FlashLite | Self::Gemini25Flash | Self::Gemini25Pro => {
GoogleModelMode::Thinking {
// By default these models are set to "auto", so we preserve that behavior
// but indicate they are capable of thinking mode
budget_tokens: None,
}
}
Self::Gemini3Flash => GoogleModelMode::Default,
Self::Gemini31FlashLite => GoogleModelMode::Default,
Self::Gemini31Pro => GoogleModelMode::Thinking {
budget_tokens: None,
},
Self::Custom { mode, .. } => *mode,
}
}
}
impl std::fmt::Display for Model {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.id())
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_function_call_part_with_signature_serializes_correctly() {
let part = FunctionCallPart {
function_call: FunctionCall {
name: "test_function".to_string(),
args: json!({"arg": "value"}),
},
thought_signature: Some("test_signature".to_string()),
};
let serialized = serde_json::to_value(&part).unwrap();
assert_eq!(serialized["functionCall"]["name"], "test_function");
assert_eq!(serialized["functionCall"]["args"]["arg"], "value");
assert_eq!(serialized["thoughtSignature"], "test_signature");
}
#[test]
fn test_function_call_part_without_signature_omits_field() {
let part = FunctionCallPart {
function_call: FunctionCall {
name: "test_function".to_string(),
args: json!({"arg": "value"}),
},
thought_signature: None,
};
let serialized = serde_json::to_value(&part).unwrap();
assert_eq!(serialized["functionCall"]["name"], "test_function");
assert_eq!(serialized["functionCall"]["args"]["arg"], "value");
// thoughtSignature field should not be present when None
assert!(serialized.get("thoughtSignature").is_none());
}
#[test]
fn test_function_call_part_deserializes_with_signature() {
let json = json!({
"functionCall": {
"name": "test_function",
"args": {"arg": "value"}
},
"thoughtSignature": "test_signature"
});
let part: FunctionCallPart = serde_json::from_value(json).unwrap();
assert_eq!(part.function_call.name, "test_function");
assert_eq!(part.thought_signature, Some("test_signature".to_string()));
}
#[test]
fn test_function_call_part_deserializes_without_signature() {
let json = json!({
"functionCall": {
"name": "test_function",
"args": {"arg": "value"}
}
});
let part: FunctionCallPart = serde_json::from_value(json).unwrap();
assert_eq!(part.function_call.name, "test_function");
assert_eq!(part.thought_signature, None);
}
#[test]
fn test_function_call_part_round_trip() {
let original = FunctionCallPart {
function_call: FunctionCall {
name: "test_function".to_string(),
args: json!({"arg": "value", "nested": {"key": "val"}}),
},
thought_signature: Some("round_trip_signature".to_string()),
};
let serialized = serde_json::to_value(&original).unwrap();
let deserialized: FunctionCallPart = serde_json::from_value(serialized).unwrap();
assert_eq!(deserialized.function_call.name, original.function_call.name);
assert_eq!(deserialized.function_call.args, original.function_call.args);
assert_eq!(deserialized.thought_signature, original.thought_signature);
}
#[test]
fn test_function_call_part_with_empty_signature_serializes() {
let part = FunctionCallPart {
function_call: FunctionCall {
name: "test_function".to_string(),
args: json!({"arg": "value"}),
},
thought_signature: Some("".to_string()),
};
let serialized = serde_json::to_value(&part).unwrap();
// Empty string should still be serialized (normalization happens at a higher level)
assert_eq!(serialized["thoughtSignature"], "");
}
}