use std::sync::atomic::{AtomicU64, Ordering};
use serde_json::{Map, Value, json};
use yapi_types::message::{
AssistantMessage, Content, ContentBlock, Message, StopReason, ThinkingLevel, ToolCall,
ToolDeclaration, ToolResultMessage, Usage, blocks_text,
};
use yapi_types::model::Model;
use crate::auth::google_adc;
use crate::cost::calculate_cost;
use crate::http::{self, Failure, Headers};
use crate::schema;
use crate::stream::{
EventSender, Request, StreamEvent, StreamOptions, ThinkingBudgets, check_complete, new_output,
now_ms, send_error,
};
use crate::thinking::{clamp_level, requested_max_tokens};
use crate::transcript::{
current_tools, initial_system_message, resolve_transcript, transform_messages,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Flavor {
Gemini,
Vertex,
}
const ABORTED: &str = http::ABORTED_READ;
#[derive(Clone, Debug, PartialEq, Eq)]
enum Thinking {
Disabled,
Level(&'static str),
Budget(i64),
}
fn gemini_major_version(id: &str) -> Option<u32> {
let id = id.to_lowercase();
let rest = id.strip_prefix("gemini-")?;
let rest = rest.strip_prefix("live-").unwrap_or(rest);
let digits: String = rest.chars().take_while(char::is_ascii_digit).collect();
digits.parse().ok()
}
pub fn requires_tool_call_id(id: &str) -> bool {
id.starts_with("claude-")
|| id.starts_with("gpt-oss-")
|| gemini_major_version(id).is_some_and(|major| major >= 3)
}
fn supports_multimodal_function_response(id: &str) -> bool {
gemini_major_version(id).is_none_or(|major| major >= 3)
}
fn supports_strict_tool_sampling(id: &str) -> bool {
gemini_major_version(id).is_some_and(|major| major >= 3)
}
fn uses_thinking_level(model: &Model) -> bool {
let id = model.id.to_lowercase();
let gemini_3 = id.find("gemini-3").is_some_and(|start| {
let rest = &id[start + "gemini-3".len()..];
let rest = match rest.strip_prefix('.') {
Some(minor) => {
let digits = minor.chars().take_while(char::is_ascii_digit).count();
if digits == 0 {
return false;
}
&minor[digits..]
}
None => rest,
};
rest.starts_with("-pro") || rest.starts_with("-flash")
});
gemini_3
|| id == "gemini-flash-latest"
|| id == "gemini-flash-lite-latest"
|| id.contains("gemma-4")
|| id.contains("gemma4")
}
fn resolve_level(model: &Model, level: ThinkingLevel) -> Result<ThinkingLevel, String> {
let mapped = model.thinking_level_value(level).flatten();
let resolved = mapped.map_or_else(|| level.as_str().to_owned(), str::to_lowercase);
match ThinkingLevel::parse(&resolved) {
Some(
level @ (ThinkingLevel::Minimal
| ThinkingLevel::Low
| ThinkingLevel::Medium
| ThinkingLevel::High),
) => Ok(level),
_ => Err(format!(
"Unsupported Google thinking level mapping for {}/{}: {} -> {}",
model.provider,
model.id,
level.as_str(),
match model.thinking_level_value(level) {
Some(Some(value)) => value.to_owned(),
Some(None) => "null".to_owned(),
None => "undefined".to_owned(),
}
)),
}
}
fn google_level(level: ThinkingLevel) -> &'static str {
match level {
ThinkingLevel::Minimal => "MINIMAL",
ThinkingLevel::Low => "LOW",
ThinkingLevel::Medium => "MEDIUM",
_ => "HIGH",
}
}
fn budget(model: &Model, level: ThinkingLevel, custom: &ThinkingBudgets, flavor: Flavor) -> i64 {
let custom = match level {
ThinkingLevel::Minimal => custom.minimal,
ThinkingLevel::Low => custom.low,
ThinkingLevel::Medium => custom.medium,
_ => custom.high,
};
if let Some(budget) = custom {
return budget as i64;
}
let table: [i64; 4] = if model.id.contains("2.5-pro") {
[128, 2048, 8192, 32768]
} else if flavor == Flavor::Gemini && model.id.contains("2.5-flash-lite") {
[512, 2048, 8192, 24576]
} else if model.id.contains("2.5-flash") {
[128, 2048, 8192, 24576]
} else {
return -1;
};
match level {
ThinkingLevel::Minimal => table[0],
ThinkingLevel::Low => table[1],
ThinkingLevel::Medium => table[2],
_ => table[3],
}
}
fn thinking(model: &Model, options: &StreamOptions, flavor: Flavor) -> Result<Thinking, String> {
let Some(level) = options.reasoning else {
return Ok(Thinking::Disabled);
};
let clamped = clamp_level(model, level);
if clamped == ThinkingLevel::Off {
return Ok(Thinking::Disabled);
}
let resolved = resolve_level(model, clamped)?;
Ok(if uses_thinking_level(model) {
Thinking::Level(google_level(resolved))
} else {
Thinking::Budget(budget(model, resolved, &options.thinking_budgets, flavor))
})
}
fn disabled_thinking_config(model: &Model) -> Result<Value, String> {
if !uses_thinking_level(model) {
return Ok(json!({"thinkingBudget": 0}));
}
let fallback = clamp_level(model, ThinkingLevel::Off);
if fallback == ThinkingLevel::Off {
return Ok(json!({"thinkingBudget": 0}));
}
Ok(json!({"thinkingLevel": google_level(resolve_level(model, fallback)?)}))
}
fn valid_signature(signature: Option<&str>) -> Option<&str> {
signature.filter(|signature| {
!signature.is_empty() && signature.len() % 4 == 0 && {
let body = signature.trim_end_matches('=');
signature.len() - body.len() <= 2
&& !body.is_empty()
&& body
.bytes()
.all(|b| b.is_ascii_alphanumeric() || b == b'+' || b == b'/')
}
})
}
fn inline_data(mime_type: &str, data: &str) -> Value {
json!({"inlineData": {"data": data, "mimeType": mime_type}})
}
fn part(text: Option<&str>, thought: bool, signature: Option<&str>) -> Value {
let mut part = Map::new();
if let Some(text) = text {
part.insert("text".into(), json!(text));
}
if thought {
part.insert("thought".into(), json!(true));
}
if let Some(signature) = signature {
part.insert("thoughtSignature".into(), json!(signature));
}
Value::Object(part)
}
fn convert_assistant(assistant: &AssistantMessage, model: &Model) -> Vec<Value> {
let same = assistant.provider == model.provider && assistant.model == model.id;
let signature = |value: Option<&str>| {
if same {
valid_signature(value).map(str::to_owned)
} else {
None
}
};
let mut parts = Vec::new();
for block in &assistant.content {
match block {
ContentBlock::Text(text) => {
let signature = signature(text.text_signature.as_deref());
if text.text.trim().is_empty() && signature.is_none() {
continue;
}
parts.push(part(Some(&text.text), false, signature.as_deref()));
}
ContentBlock::Thinking(thinking) => {
if same {
let signature = signature(thinking.thinking_signature.as_deref());
if thinking.thinking.trim().is_empty() && signature.is_none() {
continue;
}
parts.push(part(Some(&thinking.thinking), true, signature.as_deref()));
} else if !thinking.thinking.trim().is_empty() {
parts.push(part(Some(&thinking.thinking), false, None));
}
}
ContentBlock::ToolCall(call) => {
let mut function_call = Map::new();
function_call.insert("args".into(), Value::Object(call.arguments.clone()));
if requires_tool_call_id(&model.id) {
function_call.insert("id".into(), json!(call.id));
}
function_call.insert("name".into(), json!(call.name));
let mut value = Map::new();
value.insert("functionCall".into(), Value::Object(function_call));
if let Some(signature) = signature(call.thought_signature.as_deref()) {
value.insert("thoughtSignature".into(), json!(signature));
}
parts.push(Value::Object(value));
}
ContentBlock::Image(_) => {}
}
}
parts
}
fn convert_tool_result(result: &ToolResultMessage, model: &Model, contents: &mut Vec<Value>) {
let text = blocks_text(&result.content, "\n");
let images: Vec<Value> = if model.accepts_images() {
result
.content
.iter()
.filter_map(|block| match block {
ContentBlock::Image(image) => Some(inline_data(&image.mime_type, &image.data)),
_ => None,
})
.collect()
} else {
Vec::new()
};
let multimodal = supports_multimodal_function_response(&model.id);
let value = if !text.is_empty() {
text
} else if !images.is_empty() {
"(see attached image)".to_owned()
} else {
String::new()
};
let mut response = Map::new();
response.insert("name".into(), json!(result.tool_name));
response.insert(
"response".into(),
if result.is_error {
json!({"error": value})
} else {
json!({"output": value})
},
);
if !images.is_empty() && multimodal {
response.insert("parts".into(), Value::Array(images.clone()));
}
if requires_tool_call_id(&model.id) {
response.insert("id".into(), json!(result.tool_call_id));
}
let function_response = json!({"functionResponse": response});
let merged = contents.last_mut().is_some_and(|last| {
let is_responses = last["role"] == "user"
&& last["parts"]
.as_array()
.is_some_and(|parts| parts.iter().any(|p| p.get("functionResponse").is_some()));
if is_responses && let Some(parts) = last["parts"].as_array_mut() {
parts.push(function_response.clone());
}
is_responses
});
if !merged {
contents.push(json!({"parts": [function_response], "role": "user"}));
}
if !images.is_empty() && !multimodal {
let mut parts = vec![json!({"text": "Tool result image:"})];
parts.extend(images);
contents.push(json!({"parts": parts, "role": "user"}));
}
}
fn convert_messages(model: &Model, messages: &[Message]) -> Vec<Value> {
let collapsed = resolve_transcript(messages, false);
let conversation = match collapsed.first() {
Some(Message::System(_)) => &collapsed[1..],
_ => &collapsed[..],
};
let needs_ids = requires_tool_call_id(&model.id);
let normalize = |id: &str, _: &AssistantMessage| {
if needs_ids {
super::normalize_tool_call_id(id)
} else {
id.to_owned()
}
};
let transformed = transform_messages(conversation, model, Some(&normalize), now_ms());
let mut contents: Vec<Value> = Vec::new();
for message in &transformed {
match message {
Message::User(user) => {
let parts: Vec<Value> = match &user.content {
Content::Text(text) => vec![json!({"text": text})],
Content::Blocks(blocks) => blocks
.iter()
.filter_map(|block| match block {
ContentBlock::Text(text) => Some(json!({"text": text.text})),
ContentBlock::Image(image) => {
Some(inline_data(&image.mime_type, &image.data))
}
_ => None,
})
.collect(),
};
if !parts.is_empty() {
contents.push(json!({"parts": parts, "role": "user"}));
}
}
Message::Assistant(assistant) => {
let parts = convert_assistant(assistant, model);
if !parts.is_empty() {
contents.push(json!({"parts": parts, "role": "model"}));
}
}
Message::ToolResult(result) => convert_tool_result(result, model, &mut contents),
_ => {}
}
}
contents
}
fn convert_tools(tools: &[ToolDeclaration], strict_mode: bool) -> Result<Value, String> {
let declarations = tools
.iter()
.map(|tool| {
let strict = schema::strict_sampling(tool, strict_mode, None)?;
Ok(json!({
"name": tool.name,
"description": tool.description,
"parametersJsonSchema": schema::tool_parameters(tool, strict),
}))
})
.collect::<Result<Vec<Value>, String>>()?;
Ok(json!([{"functionDeclarations": declarations}]))
}
fn build_body(
model: &Model,
messages: &[Message],
options: &StreamOptions,
max_tokens: u64,
thinking: &Thinking,
) -> Result<Value, String> {
let contents = convert_messages(model, messages);
let collapsed = resolve_transcript(messages, false);
let tools = current_tools(&collapsed);
let strict_mode = supports_strict_tool_sampling(&model.id);
let mut body = Map::new();
body.insert("contents".into(), Value::Array(contents));
let system = initial_system_message(&collapsed)
.map(|system| system.text())
.unwrap_or_default();
if !system.is_empty() {
body.insert(
"systemInstruction".into(),
json!({"parts": [{"text": system}], "role": "user"}),
);
}
if !tools.is_empty() {
body.insert("tools".into(), convert_tools(&tools, strict_mode)?);
let validated = tools
.iter()
.map(|tool| schema::strict_sampling(tool, strict_mode, None))
.collect::<Result<Vec<bool>, String>>()?
.into_iter()
.any(|strict| strict);
if validated {
body.insert(
"toolConfig".into(),
json!({"functionCallingConfig": {"mode": "VALIDATED"}}),
);
}
}
let mut generation = Map::new();
if let Some(temperature) = options.temperature {
generation.insert("temperature".into(), json!(temperature));
}
generation.insert("maxOutputTokens".into(), json!(max_tokens));
if model.reasoning {
let config = match thinking {
Thinking::Level(level) => {
Some(json!({"includeThoughts": true, "thinkingLevel": level}))
}
Thinking::Budget(budget) => {
Some(json!({"includeThoughts": true, "thinkingBudget": budget}))
}
Thinking::Disabled => Some(disabled_thinking_config(model)?),
};
if let Some(config) = config {
generation.insert("thinkingConfig".into(), config);
}
}
body.insert("generationConfig".into(), Value::Object(generation));
Ok(Value::Object(body))
}
fn model_path(id: &str) -> String {
if id.starts_with("models/") || id.starts_with("tunedModels/") {
id.to_owned()
} else {
format!("models/{id}")
}
}
fn status_message(status: u16, body: &str) -> String {
match serde_json::from_str::<Value>(body) {
Ok(json) => yapi_types::json::stringify(&json),
Err(_) => {
let reason = reqwest::StatusCode::from_u16(status)
.ok()
.and_then(|code| code.canonical_reason())
.unwrap_or("");
let error = json!({"error": {"message": body, "code": status, "status": reason}});
yapi_types::json::stringify(&error)
}
}
}
fn failure_message(failure: Failure) -> String {
match failure {
Failure::Status { status, body } => status_message(status, &body),
Failure::Connection(_) => "fetch failed".to_owned(),
Failure::Aborted => ABORTED.to_owned(),
Failure::RetryDelay { delay, .. } => delay,
}
}
static TOOL_CALL_COUNTER: AtomicU64 = AtomicU64::new(0);
#[derive(Clone, Copy, PartialEq, Eq)]
enum Open {
Text,
Thinking,
}
struct State {
output: AssistantMessage,
open: Option<Open>,
}
fn map_stop_reason(reason: &str) -> Result<StopReason, String> {
match reason {
"STOP" => Ok(StopReason::Stop),
"MAX_TOKENS" => Ok(StopReason::Length),
"BLOCKLIST"
| "PROHIBITED_CONTENT"
| "SPII"
| "SAFETY"
| "IMAGE_SAFETY"
| "IMAGE_PROHIBITED_CONTENT"
| "IMAGE_RECITATION"
| "IMAGE_OTHER"
| "RECITATION"
| "FINISH_REASON_UNSPECIFIED"
| "OTHER"
| "LANGUAGE"
| "MALFORMED_FUNCTION_CALL"
| "UNEXPECTED_TOOL_CALL"
| "TOO_MANY_TOOL_CALLS"
| "NO_IMAGE" => Ok(StopReason::Error),
other => Err(format!("Unhandled stop reason: {other}")),
}
}
impl State {
fn last(&self) -> usize {
self.output.content.len().saturating_sub(1)
}
fn close(&mut self, sender: &EventSender) {
if self.open.take().is_some() {
sender.end(&self.output, self.last());
}
}
fn text_part(&mut self, part: &Value, text: &str, sender: &EventSender) {
let kind = if part["thought"] == true {
Open::Thinking
} else {
Open::Text
};
if self.open != Some(kind) {
self.close(sender);
let block = match kind {
Open::Thinking => ContentBlock::thinking("", None),
Open::Text => ContentBlock::text(""),
};
sender.start(&mut self.output, block);
self.open = Some(kind);
}
let incoming = part["thoughtSignature"]
.as_str()
.filter(|signature| !signature.is_empty())
.map(str::to_owned);
let index = self.last();
let signature = match &mut self.output.content[index] {
ContentBlock::Thinking(block) => &mut block.thinking_signature,
ContentBlock::Text(block) => &mut block.text_signature,
_ => return,
};
if incoming.is_some() {
*signature = incoming;
}
sender.delta(&mut self.output, index, text);
}
fn function_call(&mut self, part: &Value, sender: &EventSender) {
self.close(sender);
let call = &part["functionCall"];
let name = call["name"].as_str().unwrap_or_default().to_owned();
let provided = call["id"].as_str().filter(|id| !id.is_empty());
let duplicate = provided.is_some_and(|id| {
self.output
.content
.iter()
.any(|block| matches!(block, ContentBlock::ToolCall(c) if c.id == id))
});
let id = match provided {
Some(id) if !duplicate => id.to_owned(),
_ => format!(
"{name}_{}_{}",
now_ms(),
TOOL_CALL_COUNTER.fetch_add(1, Ordering::Relaxed) + 1
),
};
let arguments = call["args"].as_object().cloned().unwrap_or_default();
let delta = yapi_types::json::stringify(&arguments);
let tool_call = ContentBlock::ToolCall(ToolCall {
id,
name,
arguments,
thought_signature: part["thoughtSignature"]
.as_str()
.filter(|signature| !signature.is_empty())
.map(str::to_owned),
namespace: None,
});
let index = sender.start(&mut self.output, tool_call);
sender.delta(&mut self.output, index, &delta);
sender.end(&self.output, index);
}
fn handle(&mut self, chunk: &Value, model: &Model, sender: &EventSender) -> Result<(), String> {
if self.output.response_id.is_none()
&& let Some(id) = chunk["responseId"].as_str().filter(|id| !id.is_empty())
{
self.output.response_id = Some(id.to_owned());
}
let candidate = &chunk["candidates"][0];
for part in candidate["content"]["parts"]
.as_array()
.into_iter()
.flatten()
{
if let Some(text) = part.get("text").and_then(Value::as_str) {
self.text_part(part, text, sender);
}
if part["functionCall"].is_object() {
self.function_call(part, sender);
}
}
if let Some(reason) = candidate["finishReason"].as_str() {
self.output.raw_stop_reason = Some(reason.to_owned());
self.output.stop_reason = map_stop_reason(reason)?;
self.output.error_message = (self.output.stop_reason == StopReason::Error)
.then(|| format!("Provider stopped with: {reason}"));
if self.output.stop_reason == StopReason::Stop
&& self
.output
.content
.iter()
.any(|block| matches!(block, ContentBlock::ToolCall(_)))
{
self.output.stop_reason = StopReason::ToolUse;
}
}
let metadata = &chunk["usageMetadata"];
if metadata.is_object() {
let count = |key: &str| metadata[key].as_u64().unwrap_or(0);
let cached = count("cachedContentTokenCount");
let thoughts = count("thoughtsTokenCount");
let mut usage = Usage {
input: count("promptTokenCount").saturating_sub(cached),
output: count("candidatesTokenCount") + thoughts,
cache_read: cached,
cache_write: 0,
reasoning: Some(thoughts),
total_tokens: Some(count("totalTokenCount")),
..Usage::default()
};
calculate_cost(model, &mut usage);
self.output.usage = usage;
}
Ok(())
}
}
fn next_event(buffer: &mut String) -> Option<String> {
let (index, length) = ["\n\n", "\r\r", "\r\n\r\n"]
.iter()
.filter_map(|delimiter| buffer.find(delimiter).map(|index| (index, delimiter.len())))
.min_by_key(|(index, _)| *index)?;
let event = buffer[..index].to_owned();
buffer.drain(..index + length);
Some(event)
}
fn chunk_error(text: &str) -> Option<String> {
let json: Value = serde_json::from_str(text).ok()?;
let error = json.get("error")?;
let code = error["code"].as_f64()?;
(400.0..600.0).contains(&code).then(|| {
format!(
"got status: {}. {}",
match &error["status"] {
Value::String(status) => status.clone(),
Value::Null => "undefined".into(),
other => other.to_string(),
},
yapi_types::json::stringify(&json)
)
})
}
async fn consume(
state: &mut State,
mut response: reqwest::Response,
model: &Model,
sender: &EventSender,
options: &StreamOptions,
) -> Result<(), String> {
let mut pending: Vec<u8> = Vec::new();
let mut buffer = String::new();
loop {
let chunk = match http::read_chunk(&mut response, &options.cancel, ABORTED).await {
Ok(chunk) => chunk,
Err(_) if options.cancel.is_cancelled() => return Err(ABORTED.into()),
Err(message) => return Err(message),
};
let Some(bytes) = chunk else {
if !buffer.trim().is_empty() {
return Err("Incomplete JSON segment at the end".into());
}
break;
};
let text = yapi_types::js::decode_utf8_stream(&mut pending, &bytes);
if let Some(error) = chunk_error(&text) {
return Err(error);
}
buffer.push_str(&text);
while let Some(event) = next_event(&mut buffer) {
let Some(data) = event.trim().strip_prefix("data:") else {
continue;
};
let data = data.trim();
let json: Value = serde_json::from_str(data)
.map_err(|err| format!("exception parsing stream chunk {data}. {err}"))?;
options.hooks.stream_event(&json).await;
state.handle(&json, model, sender)?;
}
}
Ok(())
}
pub(super) async fn run(request: Request, sender: EventSender, flavor: Flavor) {
let Request {
model,
messages,
options,
} = request;
let output = new_output(&model, now_ms());
let response = match connect(&model, &messages, &options, flavor).await {
Ok(response) => response,
Err(message) => return send_error(&sender, output, &options.cancel, message),
};
sender.send(StreamEvent::Start(output.clone()));
let mut state = State { output, open: None };
let pending = match flavor {
Flavor::Gemini => "Google stream ended without a finish reason",
Flavor::Vertex => "Google Vertex stream ended without a finish reason",
};
let result = consume(&mut state, response, &model, &sender, &options)
.await
.and_then(|()| {
state.close(&sender);
check_complete(&state.output, &options.cancel, pending)
});
match result {
Ok(()) => sender.send(StreamEvent::Done(state.output)),
Err(message) => send_error(&sender, state.output, &options.cancel, message),
}
}
async fn connect(
model: &Model,
messages: &[Message],
options: &StreamOptions,
flavor: Flavor,
) -> Result<reqwest::Response, String> {
let api_key = options.api_key.clone().filter(|key| !key.is_empty());
if flavor == Flavor::Gemini && api_key.is_none() {
return Err(format!("No API key for provider: {}", model.provider));
}
let target = match flavor {
Flavor::Gemini => Target {
url: format!(
"{}/{}:streamGenerateContent?alt=sse",
model.base_url.trim_end_matches('/'),
model_path(&model.id)
),
auth: Auth::Key(api_key.unwrap_or_default()),
},
Flavor::Vertex => vertex::target(model, options)?,
};
let max_tokens = requested_max_tokens(model, messages, options);
let thinking = thinking(model, options, flavor)?;
let body = build_body(model, messages, options, max_tokens, &thinking)?;
let body = options.hooks.payload(body).await;
let body = yapi_types::json::stringify(&body);
let mut headers = Headers::default();
headers.set("content-type", Some("application/json"));
match target.auth {
Auth::Key(key) => headers.set("x-goog-api-key", Some(key)),
Auth::Credentials { credentials_file } => {
let token = google_adc::token(
credentials_file,
google_adc::OAUTH2_TOKEN_URL,
&options.cancel,
)
.await?;
headers.set(
"authorization",
Some(format!("Bearer {}", token.access_token)),
);
headers.set("x-goog-user-project", token.quota_project);
}
}
headers.extend_model(model.headers.as_ref());
headers.extend(&options.headers);
let build = || headers.apply(http::client().post(&target.url).body(body.clone()));
let options = StreamOptions {
timeout_ms: None,
..options.clone()
};
http::send_with(build, &options, false)
.await
.map_err(failure_message)
}
struct Target {
url: String,
auth: Auth,
}
enum Auth {
Key(String),
Credentials { credentials_file: Option<String> },
}
mod vertex {
use yapi_types::model::Model;
use super::{Auth, Target};
use crate::auth::google_adc;
use crate::credentials::provider_env_value;
use crate::stream::StreamOptions;
const CREDENTIALS_MARKER: &str = "gcp-vertex-credentials";
const API_VERSION: &str = "v1";
fn api_key(options: &StreamOptions) -> Option<String> {
let key = options.api_key.as_deref()?.trim();
let placeholder = key.len() > 2
&& key.starts_with('<')
&& key.ends_with('>')
&& !key[1..key.len() - 1].contains('>');
(!key.is_empty() && key != CREDENTIALS_MARKER && !placeholder).then(|| key.to_owned())
}
fn custom_base_url(model: &Model) -> Option<&str> {
let trimmed = model.base_url.trim();
(!trimmed.is_empty() && !trimmed.contains("{location}")).then_some(trimmed)
}
fn includes_api_version(base_url: &str) -> bool {
let is_version = |part: &str| {
part.strip_prefix('v').is_some_and(|rest| {
let digits = rest.chars().take_while(char::is_ascii_digit).count();
digits > 0 && {
let tail = &rest[digits..];
tail.is_empty()
|| tail
.strip_prefix("beta")
.is_some_and(|n| n.chars().all(|c| c.is_ascii_digit()))
}
})
};
match url::Url::parse(base_url) {
Ok(url) => url.path().split('/').any(is_version),
Err(_) => base_url.split('/').any(is_version),
}
}
fn model_path(id: &str) -> String {
if id.starts_with("publishers/") || id.starts_with("projects/") || id.starts_with("models/")
{
id.to_owned()
} else if let Some((publisher, model)) = id.split_once('/') {
let model = model.split('/').next().unwrap_or(model);
format!("publishers/{publisher}/models/{model}")
} else {
format!("publishers/google/models/{id}")
}
}
fn join(base: &str, version: bool, path: &str) -> String {
let base = base.strip_suffix('/').unwrap_or(base);
let mut url = base.to_owned();
if version {
url.push('/');
url.push_str(API_VERSION);
}
format!("{url}/{path}:streamGenerateContent?alt=sse")
}
pub(super) fn target(model: &Model, options: &StreamOptions) -> Result<Target, String> {
let env = |name: &str| provider_env_value(name, options.env.as_ref());
let custom = custom_base_url(model);
let path = model_path(&model.id);
if let Some(key) = api_key(options) {
let url = match custom {
Some(base) => join(base, !includes_api_version(base), &path),
None => join("https://aiplatform.googleapis.com", true, &path),
};
return Ok(Target {
url,
auth: Auth::Key(key),
});
}
let project = env("GOOGLE_CLOUD_PROJECT")
.or_else(|| env("GCLOUD_PROJECT"))
.ok_or_else(|| {
"Vertex AI requires a project ID. Set GOOGLE_CLOUD_PROJECT/GCLOUD_PROJECT or pass project in options.".to_owned()
})?;
let location = env("GOOGLE_CLOUD_LOCATION").ok_or_else(|| {
"Vertex AI requires a location. Set GOOGLE_CLOUD_LOCATION or pass location in options."
.to_owned()
})?;
let url = match custom {
Some(base) => join(base, !includes_api_version(base), &path),
None => {
let base = match location.as_str() {
"global" => "https://aiplatform.googleapis.com".to_owned(),
"us" | "eu" => format!("https://aiplatform.{location}.rep.googleapis.com"),
_ => format!("https://{location}-aiplatform.googleapis.com"),
};
let path = if path.starts_with("projects/") {
path
} else {
format!("projects/{project}/locations/{location}/{path}")
};
join(&base, true, &path)
}
};
Ok(Target {
url,
auth: Auth::Credentials {
credentials_file: google_adc::credentials_file(env),
},
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::credentials::ProviderEnv;
fn options(key: Option<&str>, vars: &[(&str, &str)]) -> StreamOptions {
let env: ProviderEnv = vars
.iter()
.map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
.collect();
StreamOptions {
api_key: key.map(str::to_owned),
env: Some(env),
..StreamOptions::default()
}
}
fn model(id: &str, base_url: &str) -> Model {
let mut model = crate::catalog::builtin_models("google-vertex")
.into_iter()
.next()
.unwrap();
model.id = id.into();
model.base_url = base_url.into();
model
}
#[test]
fn builds_sdk_urls() {
let template = "https://{location}-aiplatform.googleapis.com";
let url = |key: Option<&str>, location: &str, base: &str| {
target(
&model("gemini-2.5-flash", base),
&options(
key,
&[
("GOOGLE_CLOUD_PROJECT", "proj-1"),
("GOOGLE_CLOUD_LOCATION", location),
],
),
)
.map(|target| target.url)
};
assert_eq!(
url(Some("AIzaKEY"), "us-central1", template).unwrap(),
"https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
);
for (location, host) in [
(
"us-central1",
"https://us-central1-aiplatform.googleapis.com",
),
("global", "https://aiplatform.googleapis.com"),
("eu", "https://aiplatform.eu.rep.googleapis.com"),
] {
assert_eq!(
url(Some("gcp-vertex-credentials"), location, template).unwrap(),
format!(
"{host}/v1/projects/proj-1/locations/{location}/publishers/google/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
)
);
}
assert_eq!(
url(None, "us-central1", "http://proxy.local/vertex/v1beta1/").unwrap(),
"http://proxy.local/vertex/v1beta1/publishers/google/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
);
assert_eq!(
url(Some("<key>"), "us-central1", "http://proxy.local").unwrap(),
"http://proxy.local/v1/publishers/google/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
);
assert_eq!(model_path("meta/llama-4"), "publishers/meta/models/llama-4");
let missing = target(&model("gemini-2.5-flash", template), &options(None, &[]));
assert_eq!(
missing.err().as_deref(),
Some(
"Vertex AI requires a project ID. Set GOOGLE_CLOUD_PROJECT/GCLOUD_PROJECT or pass project in options."
)
);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn detects_model_families() {
assert_eq!(gemini_major_version("gemini-3.1-pro-preview"), Some(3));
assert_eq!(gemini_major_version("gemini-live-2.5-flash"), Some(2));
assert!(requires_tool_call_id("gemini-3-flash-preview"));
assert!(!requires_tool_call_id("gemini-2.5-flash"));
assert!(requires_tool_call_id("claude-sonnet-4"));
let model = |id: &str| -> Model {
serde_json::from_value(json!({
"id": id, "name": id, "api": "google-generative-ai", "provider": "google",
"baseUrl": "", "reasoning": true, "input": ["text"],
"cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0},
"contextWindow": 1, "maxTokens": 1,
}))
.unwrap()
};
assert!(uses_thinking_level(&model("gemini-3.8-flash")));
assert!(uses_thinking_level(&model("gemini-3-pro-preview")));
assert!(!uses_thinking_level(&model("gemini-2.5-pro")));
assert!(uses_thinking_level(&model("gemma-4-31b-it")));
}
#[test]
fn splits_events_and_signatures() {
let mut buffer = "data: {}\r\n\r\ndata: {\"a\":1}\n\nrest".to_owned();
assert_eq!(next_event(&mut buffer).as_deref(), Some("data: {}"));
assert_eq!(next_event(&mut buffer).as_deref(), Some("data: {\"a\":1}"));
assert_eq!(next_event(&mut buffer), None);
assert_eq!(buffer, "rest");
assert!(valid_signature(Some("c2ln")).is_some());
assert!(valid_signature(Some("c2lnbg==")).is_some());
assert!(valid_signature(Some("abc")).is_none());
assert!(valid_signature(Some("ab$=")).is_none());
}
}