From 459a1e99354f70cb53e044fec8e72c886e71f657 Mon Sep 17 00:00:00 2001 From: ammar siddiqui Date: Wed, 27 May 2026 17:48:03 -0400 Subject: [PATCH] feat(agent): GeminiProvider over reqwest with function calling (#9) Implements LlmProvider for Gemini against :generateContent (no SDK, same posture as chat/anthropic.rs): - Pure, unit-tested conversions to_wire_request / from_wire_response; complete() is the thin HTTP shell around them. - System turns lifted to systemInstruction; user/assistant -> user/model; assistant tool_calls -> functionCall parts; tool results -> functionResponse parts (role user), routed by name. - API key sent via x-goog-api-key header (never in URL/query/logs). - Token usage reported; finish_reason forced to ToolCalls when the model emits function calls. - HTTP errors classified (400 InvalidRequest, 401/403 Auth, else Provider). - with_base_url() for mock-server tests. Adds ChatMessage.name (Option) so tool results carry the tool name for name-routed providers; tool_result() now takes (id, name, content). Tests: 7 non-live (wire mapping + parsing + error classification) plus an #[ignore] live smoke test. Verified live against gemini-2.5-flash with a real key: text completion returned usage, tool calling returned apply_seed{kind,n} with finish=ToolCalls. Closes #9. --- core/src/agent/provider/gemini.rs | 558 ++++++++++++++++++++++++++++++ core/src/agent/provider/mod.rs | 29 +- 2 files changed, 583 insertions(+), 4 deletions(-) create mode 100644 core/src/agent/provider/gemini.rs diff --git a/core/src/agent/provider/gemini.rs b/core/src/agent/provider/gemini.rs new file mode 100644 index 0000000..f2a8461 --- /dev/null +++ b/core/src/agent/provider/gemini.rs @@ -0,0 +1,558 @@ +//! Gemini provider (`#9`) — implements [`LlmProvider`] over `reqwest`. +//! +//! Talks to the Generative Language API `:generateContent` endpoint, no SDK +//! (same posture as `chat/anthropic.rs`). The conversion between our +//! provider-neutral types and Gemini's wire format lives in the free +//! functions [`to_wire_request`] / [`from_wire_response`], which are pure and +//! unit-tested without a network; [`GeminiProvider::complete`] is the thin +//! HTTP shell around them. +//! +//! Mapping notes: +//! - `System` turns are lifted into `system_instruction`. +//! - `User`/`Assistant` map to Gemini roles `user`/`model`. +//! - Assistant `tool_calls` become `functionCall` parts; `Tool` results +//! become `functionResponse` parts (role `user`), routed **by name** since +//! Gemini has no tool-call ids — hence [`ChatMessage::name`]. +//! - Gemini doesn't return tool-call ids, so we synthesize `call_`. +//! - Streaming is `#10`; this is the single-shot path. + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use super::{ + ChatMessage, ChatRequest, ChatResponse, FinishReason, LlmProvider, ProviderError, Role, + TokenUsage, ToolCall, +}; + +const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta"; + +/// Gemini provider. Holds a reusable `reqwest::Client`; the API key is passed +/// per call (never stored here) per the keychain design. +pub struct GeminiProvider { + client: reqwest::Client, + base_url: String, +} + +impl GeminiProvider { + pub fn new() -> Self { + Self { client: reqwest::Client::new(), base_url: DEFAULT_BASE_URL.to_string() } + } + + /// Override the base URL (used by tests to point at a local mock server). + pub fn with_base_url(mut self, base_url: impl Into) -> Self { + self.base_url = base_url.into(); + self + } +} + +impl Default for GeminiProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl LlmProvider for GeminiProvider { + fn id(&self) -> &'static str { + "gemini" + } + + async fn complete( + &self, + request: &ChatRequest, + api_key: &str, + ) -> Result { + let url = format!("{}/models/{}:generateContent", self.base_url, request.model); + let body = to_wire_request(request); + + let resp = self + .client + .post(&url) + // Key as a header (never in the URL/query) so it can't leak into + // request logs or proxies that record URLs. + .header("x-goog-api-key", api_key) + .json(&body) + .send() + .await + .map_err(|e| ProviderError::Transport(e.to_string()))?; + + let status = resp.status(); + let text = resp + .text() + .await + .map_err(|e| ProviderError::Transport(e.to_string()))?; + + if !status.is_success() { + return Err(classify_http_error(status.as_u16(), &text)); + } + + let wire: GeminiResponseBody = serde_json::from_str(&text) + .map_err(|e| ProviderError::Decode(format!("{e}: {text}")))?; + from_wire_response(wire) + } +} + +/// Map a non-2xx HTTP response to a typed error, pulling Gemini's +/// `error.message` when present. +fn classify_http_error(status: u16, body: &str) -> ProviderError { + let message = serde_json::from_str::(body) + .ok() + .map(|e| e.error.message) + .unwrap_or_else(|| body.to_string()); + match status { + 400 => ProviderError::InvalidRequest(message), + 401 | 403 => ProviderError::Auth(message), + _ => ProviderError::Provider(format!("HTTP {status}: {message}")), + } +} + +// --------------------------------------------------------------------------- +// Neutral -> Gemini wire +// --------------------------------------------------------------------------- + +/// Build the Gemini request body from a neutral [`ChatRequest`]. +fn to_wire_request(req: &ChatRequest) -> GeminiRequestBody { + let mut system_parts: Vec = Vec::new(); + let mut contents: Vec = Vec::new(); + + for msg in &req.messages { + match msg.role { + Role::System => { + if !msg.content.is_empty() { + system_parts.push(GeminiPart::text(&msg.content)); + } + } + Role::User => { + contents.push(GeminiContent::new("user", vec![GeminiPart::text(&msg.content)])); + } + Role::Assistant => { + let mut parts = Vec::new(); + if !msg.content.is_empty() { + parts.push(GeminiPart::text(&msg.content)); + } + for call in &msg.tool_calls { + parts.push(GeminiPart::function_call(&call.name, call.arguments.clone())); + } + contents.push(GeminiContent::new("model", parts)); + } + Role::Tool => { + // Gemini correlates by function name. Prefer the explicit + // name; fall back to the id if a caller omitted it. + let name = msg + .name + .clone() + .or_else(|| msg.tool_call_id.clone()) + .unwrap_or_default(); + contents.push(GeminiContent::new( + "user", + vec![GeminiPart::function_response(&name, tool_response_object(&msg.content))], + )); + } + } + } + + let system_instruction = + (!system_parts.is_empty()).then(|| GeminiContent { role: None, parts: system_parts }); + + let tools = (!req.tools.is_empty()).then(|| { + vec![GeminiTool { + function_declarations: req + .tools + .iter() + .map(|t| GeminiFunctionDeclaration { + name: t.name.clone(), + description: t.description.clone(), + parameters: t.parameters.clone(), + }) + .collect(), + }] + }); + + let generation_config = GeminiGenerationConfig { + temperature: req.config.temperature, + max_output_tokens: req.config.max_output_tokens, + top_p: req.config.top_p, + }; + + GeminiRequestBody { contents, system_instruction, tools, generation_config } +} + +/// Gemini's `functionResponse.response` must be a JSON object. If the tool +/// result text is already a JSON object, pass it through; otherwise wrap it. +fn tool_response_object(content: &str) -> serde_json::Value { + match serde_json::from_str::(content) { + Ok(v @ serde_json::Value::Object(_)) => v, + Ok(other) => serde_json::json!({ "result": other }), + Err(_) => serde_json::json!({ "result": content }), + } +} + +// --------------------------------------------------------------------------- +// Gemini wire -> neutral +// --------------------------------------------------------------------------- + +/// Convert a parsed Gemini response into a neutral [`ChatResponse`]. +fn from_wire_response(body: GeminiResponseBody) -> Result { + let candidate = body.candidates.into_iter().next().ok_or_else(|| { + // No candidate usually means the prompt was blocked. + let reason = body + .prompt_feedback + .and_then(|f| f.block_reason) + .unwrap_or_else(|| "no candidates returned".to_string()); + ProviderError::Provider(reason) + })?; + + let mut text = String::new(); + let mut tool_calls = Vec::new(); + for part in candidate.content.map(|c| c.parts).unwrap_or_default() { + if let Some(t) = part.text { + text.push_str(&t); + } + if let Some(fc) = part.function_call { + tool_calls.push(ToolCall { + id: format!("call_{}", tool_calls.len()), + name: fc.name, + arguments: fc.args.unwrap_or_else(|| serde_json::json!({})), + }); + } + } + + let usage = body + .usage_metadata + .map(|u| TokenUsage::new(u.prompt_token_count, u.candidates_token_count)) + .unwrap_or_default(); + + // If the model emitted tool calls, that's the branch the loop cares about, + // regardless of the textual finish reason Gemini reports. + let finish_reason = if !tool_calls.is_empty() { + FinishReason::ToolCalls + } else { + map_finish_reason(candidate.finish_reason.as_deref()) + }; + + Ok(ChatResponse { text, tool_calls, usage, finish_reason }) +} + +fn map_finish_reason(reason: Option<&str>) -> FinishReason { + match reason { + Some("STOP") => FinishReason::Stop, + Some("MAX_TOKENS") => FinishReason::Length, + Some("SAFETY") | Some("RECITATION") | Some("BLOCKLIST") | Some("PROHIBITED_CONTENT") => { + FinishReason::ContentFilter + } + _ => FinishReason::Other, + } +} + +// --------------------------------------------------------------------------- +// Wire structs +// --------------------------------------------------------------------------- + +#[derive(Debug, Serialize)] +struct GeminiRequestBody { + contents: Vec, + #[serde(rename = "systemInstruction", skip_serializing_if = "Option::is_none")] + system_instruction: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tools: Option>, + #[serde(rename = "generationConfig")] + generation_config: GeminiGenerationConfig, +} + +#[derive(Debug, Serialize, Deserialize)] +struct GeminiContent { + #[serde(skip_serializing_if = "Option::is_none")] + role: Option, + parts: Vec, +} + +impl GeminiContent { + fn new(role: &str, parts: Vec) -> Self { + Self { role: Some(role.to_string()), parts } + } +} + +#[derive(Debug, Default, Serialize, Deserialize)] +struct GeminiPart { + #[serde(skip_serializing_if = "Option::is_none")] + text: Option, + #[serde(rename = "functionCall", skip_serializing_if = "Option::is_none")] + function_call: Option, + #[serde(rename = "functionResponse", skip_serializing_if = "Option::is_none")] + function_response: Option, +} + +impl GeminiPart { + fn text(s: &str) -> Self { + Self { text: Some(s.to_string()), ..Default::default() } + } + fn function_call(name: &str, args: serde_json::Value) -> Self { + Self { + function_call: Some(GeminiFunctionCall { name: name.to_string(), args: Some(args) }), + ..Default::default() + } + } + fn function_response(name: &str, response: serde_json::Value) -> Self { + Self { + function_response: Some(GeminiFunctionResponse { name: name.to_string(), response }), + ..Default::default() + } + } +} + +#[derive(Debug, Serialize, Deserialize)] +struct GeminiFunctionCall { + name: String, + #[serde(skip_serializing_if = "Option::is_none")] + args: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +struct GeminiFunctionResponse { + name: String, + response: serde_json::Value, +} + +#[derive(Debug, Serialize)] +struct GeminiTool { + #[serde(rename = "functionDeclarations")] + function_declarations: Vec, +} + +#[derive(Debug, Serialize)] +struct GeminiFunctionDeclaration { + name: String, + description: String, + parameters: serde_json::Value, +} + +#[derive(Debug, Serialize)] +struct GeminiGenerationConfig { + #[serde(skip_serializing_if = "Option::is_none")] + temperature: Option, + #[serde(rename = "maxOutputTokens", skip_serializing_if = "Option::is_none")] + max_output_tokens: Option, + #[serde(rename = "topP", skip_serializing_if = "Option::is_none")] + top_p: Option, +} + +#[derive(Debug, Deserialize)] +struct GeminiResponseBody { + #[serde(default)] + candidates: Vec, + #[serde(rename = "usageMetadata")] + usage_metadata: Option, + #[serde(rename = "promptFeedback")] + prompt_feedback: Option, +} + +#[derive(Debug, Deserialize)] +struct GeminiCandidate { + content: Option, + #[serde(rename = "finishReason")] + finish_reason: Option, +} + +#[derive(Debug, Deserialize)] +struct GeminiUsage { + #[serde(rename = "promptTokenCount", default)] + prompt_token_count: u32, + #[serde(rename = "candidatesTokenCount", default)] + candidates_token_count: u32, +} + +#[derive(Debug, Deserialize)] +struct GeminiPromptFeedback { + #[serde(rename = "blockReason")] + block_reason: Option, +} + +#[derive(Debug, Deserialize)] +struct GeminiErrorEnvelope { + error: GeminiErrorBody, +} + +#[derive(Debug, Deserialize)] +struct GeminiErrorBody { + message: String, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::provider::{GenerationConfig, ToolDeclaration}; + + fn sample_tools() -> Vec { + vec![ToolDeclaration { + name: "apply_seed".into(), + description: "seed the cortex".into(), + parameters: serde_json::json!({ + "type": "object", + "properties": {"kind": {"type": "string"}, "n": {"type": "integer"}}, + "required": ["kind", "n"] + }), + }] + } + + #[test] + fn request_lifts_system_and_maps_roles_and_tools() { + let req = ChatRequest::new( + "gemini-2.0-flash", + vec![ + ChatMessage::system("you build cortexes"), + ChatMessage::user("seed 200 neurons"), + ], + ) + .with_tools(sample_tools()) + .with_config(GenerationConfig { temperature: Some(0.1), max_output_tokens: Some(256), top_p: None }); + + let body = to_wire_request(&req); + let json = serde_json::to_value(&body).unwrap(); + + // system lifted out of contents + assert_eq!(json["systemInstruction"]["parts"][0]["text"], "you build cortexes"); + assert_eq!(json["contents"].as_array().unwrap().len(), 1); + assert_eq!(json["contents"][0]["role"], "user"); + assert_eq!(json["contents"][0]["parts"][0]["text"], "seed 200 neurons"); + // tool declaration present + assert_eq!(json["tools"][0]["functionDeclarations"][0]["name"], "apply_seed"); + // generation config camelCased (f32 -> JSON, compare with tolerance) + assert!((json["generationConfig"]["temperature"].as_f64().unwrap() - 0.1).abs() < 1e-6); + assert_eq!(json["generationConfig"]["maxOutputTokens"], 256); + assert!(json["generationConfig"].get("topP").is_none()); + } + + #[test] + fn assistant_tool_call_and_tool_result_round_trip_to_wire() { + let req = ChatRequest::new( + "gemini-2.0-flash", + vec![ + ChatMessage::user("seed it"), + ChatMessage { + role: Role::Assistant, + content: String::new(), + tool_calls: vec![ToolCall { + id: "call_0".into(), + name: "apply_seed".into(), + arguments: serde_json::json!({"kind": "small_world", "n": 200}), + }], + tool_call_id: None, + name: None, + }, + ChatMessage::tool_result("call_0", "apply_seed", r#"{"added_nodes":200}"#), + ], + ); + let json = serde_json::to_value(&to_wire_request(&req)).unwrap(); + let contents = json["contents"].as_array().unwrap(); + // user, model(functionCall), user(functionResponse) + assert_eq!(contents.len(), 3); + assert_eq!(contents[1]["role"], "model"); + assert_eq!(contents[1]["parts"][0]["functionCall"]["name"], "apply_seed"); + assert_eq!(contents[1]["parts"][0]["functionCall"]["args"]["n"], 200); + assert_eq!(contents[2]["role"], "user"); + assert_eq!(contents[2]["parts"][0]["functionResponse"]["name"], "apply_seed"); + assert_eq!(contents[2]["parts"][0]["functionResponse"]["response"]["added_nodes"], 200); + } + + #[test] + fn tool_response_object_wraps_non_objects() { + assert_eq!(tool_response_object(r#"{"a":1}"#), serde_json::json!({"a": 1})); + assert_eq!(tool_response_object("plain text"), serde_json::json!({"result": "plain text"})); + assert_eq!(tool_response_object("42"), serde_json::json!({"result": 42})); + } + + #[test] + fn parses_text_response_with_usage() { + let wire = r#"{ + "candidates": [{ + "content": {"role": "model", "parts": [{"text": "200 neurons fired"}]}, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 120, "candidatesTokenCount": 18, "totalTokenCount": 138} + }"#; + let body: GeminiResponseBody = serde_json::from_str(wire).unwrap(); + let resp = from_wire_response(body).unwrap(); + assert_eq!(resp.text, "200 neurons fired"); + assert!(resp.tool_calls.is_empty()); + assert_eq!(resp.finish_reason, FinishReason::Stop); + assert_eq!(resp.usage.prompt_tokens, 120); + assert_eq!(resp.usage.completion_tokens, 18); + assert_eq!(resp.usage.total_tokens, 138); + } + + #[test] + fn parses_function_call_response_as_tool_calls() { + let wire = r#"{ + "candidates": [{ + "content": {"role": "model", "parts": [ + {"functionCall": {"name": "apply_seed", "args": {"kind": "ring", "n": 64}}} + ]}, + "finishReason": "STOP" + }], + "usageMetadata": {"promptTokenCount": 50, "candidatesTokenCount": 12} + }"#; + let body: GeminiResponseBody = serde_json::from_str(wire).unwrap(); + let resp = from_wire_response(body).unwrap(); + assert_eq!(resp.finish_reason, FinishReason::ToolCalls); + assert_eq!(resp.tool_calls.len(), 1); + assert_eq!(resp.tool_calls[0].name, "apply_seed"); + assert_eq!(resp.tool_calls[0].id, "call_0"); + assert_eq!(resp.tool_calls[0].arguments["n"], 64); + } + + #[test] + fn no_candidates_is_provider_error() { + let wire = r#"{"promptFeedback": {"blockReason": "SAFETY"}}"#; + let body: GeminiResponseBody = serde_json::from_str(wire).unwrap(); + let err = from_wire_response(body).unwrap_err(); + assert!(matches!(err, ProviderError::Provider(m) if m == "SAFETY")); + } + + #[test] + fn http_errors_are_classified() { + let body = r#"{"error": {"code": 403, "message": "API key not valid", "status": "PERMISSION_DENIED"}}"#; + assert!(matches!(classify_http_error(403, body), ProviderError::Auth(m) if m.contains("not valid"))); + assert!(matches!(classify_http_error(400, body), ProviderError::InvalidRequest(_))); + assert!(matches!(classify_http_error(503, body), ProviderError::Provider(_))); + } + + /// Live smoke test against the real Gemini API. Ignored by default; run with + /// `GOOGLE_API_KEY=... cargo test -p core gemini_live -- --ignored --nocapture`. + /// Optionally set `GEMINI_TEST_MODEL` (defaults to gemini-2.0-flash). + #[tokio::test] + #[ignore] + async fn gemini_live_smoke() { + let _ = dotenvy::dotenv(); + let key = match std::env::var("GOOGLE_API_KEY").or_else(|_| std::env::var("GEMINI_API_KEY")) { + Ok(k) => k, + Err(_) => { + eprintln!("skipping: GOOGLE_API_KEY not set"); + return; + } + }; + // gemini-2.0-flash has no free-tier quota on some keys; 2.5-flash does. + let model = std::env::var("GEMINI_TEST_MODEL").unwrap_or_else(|_| "gemini-2.5-flash".into()); + let provider = GeminiProvider::new(); + + // 1) plain text + let req = ChatRequest::new( + &model, + vec![ChatMessage::user("Reply with exactly the word: pong")], + ); + let resp = provider.complete(&req, &key).await.expect("text completion"); + eprintln!("TEXT -> {:?} | usage={:?}", resp.text, resp.usage); + assert!(!resp.text.is_empty()); + assert!(resp.usage.total_tokens > 0); + + // 2) tool calling + let req = ChatRequest::new( + &model, + vec![ChatMessage::user("Seed a small-world cortex with 200 neurons.")], + ) + .with_tools(sample_tools()); + let resp = provider.complete(&req, &key).await.expect("tool completion"); + eprintln!("TOOLCALLS -> {:?} | finish={:?}", resp.tool_calls, resp.finish_reason); + assert_eq!(resp.finish_reason, FinishReason::ToolCalls); + assert_eq!(resp.tool_calls[0].name, "apply_seed"); + } +} diff --git a/core/src/agent/provider/mod.rs b/core/src/agent/provider/mod.rs index 8efe885..df0fee0 100644 --- a/core/src/agent/provider/mod.rs +++ b/core/src/agent/provider/mod.rs @@ -22,6 +22,9 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; +pub mod gemini; +pub use gemini::GeminiProvider; + /// Who authored a message in the conversation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -57,12 +60,23 @@ pub struct ChatMessage { /// For `Role::Tool` messages: the id of the [`ToolCall`] this answers. #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_call_id: Option, + /// For `Role::Tool` messages: the tool's name. Providers route tool + /// results differently — Gemini matches `functionResponse` by name, + /// OpenAI carries a `name` field — so we keep both id and name. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub name: Option, } impl ChatMessage { /// A plain text message in the given role (no tool calls). pub fn text(role: Role, content: impl Into) -> Self { - Self { role, content: content.into(), tool_calls: Vec::new(), tool_call_id: None } + Self { + role, + content: content.into(), + tool_calls: Vec::new(), + tool_call_id: None, + name: None, + } } /// A `system` instruction message. @@ -75,13 +89,19 @@ impl ChatMessage { Self::text(Role::User, content) } - /// A `tool` result message answering a specific tool call. - pub fn tool_result(tool_call_id: impl Into, content: impl Into) -> Self { + /// A `tool` result message answering a specific tool call, carrying the + /// tool's name so name-routed providers (Gemini) can correlate it. + pub fn tool_result( + tool_call_id: impl Into, + name: impl Into, + content: impl Into, + ) -> Self { Self { role: Role::Tool, content: content.into(), tool_calls: Vec::new(), tool_call_id: Some(tool_call_id.into()), + name: Some(name.into()), } } } @@ -281,9 +301,10 @@ mod tests { fn message_constructors_set_roles() { assert_eq!(ChatMessage::system("hi").role, Role::System); assert_eq!(ChatMessage::user("hi").role, Role::User); - let t = ChatMessage::tool_result("call_1", "{}"); + let t = ChatMessage::tool_result("call_1", "graph_snapshot", "{}"); assert_eq!(t.role, Role::Tool); assert_eq!(t.tool_call_id.as_deref(), Some("call_1")); + assert_eq!(t.name.as_deref(), Some("graph_snapshot")); } #[test]