feat: full tool call support (OpenAI + Gemini endpoints)
- store.rs: Add tool context storage (active tools, tool config, pending tool results, call_id mapping, last function calls for history rewrite) - types.rs: Add tools/tool_choice fields to ResponsesRequest, add build_function_call_output helper for OpenAI function_call output items - modify.rs: Replace hardcoded get_weather with dynamic ToolContext injection. Add openai_tools_to_gemini and openai_tool_choice_to_gemini converters. Add conversation history rewriting for tool result turns (replaces fake 'Tool call completed' model turn with real functionCall, injects functionResponse before last user turn) - proxy.rs: Build ToolContext from MitmStore before calling modify_request. Save last_function_calls for history rewriting on subsequent turns - responses.rs: Store client tools in MitmStore before LS call. Detect function_call_output in input array for tool result submission. Return captured functionCalls as OpenAI function_call output items with generated call_ids and stringified arguments - gemini.rs: New Gemini-native endpoint (POST /v1/gemini) with zero format translation. Accepts functionDeclarations directly, returns functionCall in Gemini format directly - mod.rs: Wire /v1/gemini route, bump version to 3.3.0
This commit is contained in:
@@ -8,14 +8,29 @@ use regex::Regex;
|
||||
use serde_json::Value;
|
||||
use tracing::info;
|
||||
|
||||
use super::store::{CapturedFunctionCall, PendingToolResult};
|
||||
|
||||
/// Strip ALL tool definitions.
|
||||
/// Must be true: with tools present, the LS enters full agentic mode
|
||||
/// (multi-turn tool calls, file searches, etc.) burning quota.
|
||||
const STRIP_ALL_TOOLS: bool = true;
|
||||
|
||||
/// Context for tool injection during request modification.
|
||||
/// Built from MitmStore data before calling modify_request.
|
||||
pub struct ToolContext {
|
||||
/// Gemini-format tool declarations (functionDeclarations).
|
||||
pub tools: Option<Vec<Value>>,
|
||||
/// Gemini-format toolConfig.
|
||||
pub tool_config: Option<Value>,
|
||||
/// Pending tool results to inject as functionResponse.
|
||||
pub pending_results: Vec<PendingToolResult>,
|
||||
/// Last captured function calls for history rewriting.
|
||||
pub last_calls: Vec<CapturedFunctionCall>,
|
||||
}
|
||||
|
||||
/// Modify a streamGenerateContent request body in-place.
|
||||
/// Returns the modified JSON bytes, or None if modification wasn't possible.
|
||||
pub fn modify_request(body: &[u8]) -> Option<Vec<u8>> {
|
||||
pub fn modify_request(body: &[u8], tool_ctx: Option<&ToolContext>) -> Option<Vec<u8>> {
|
||||
let mut json: Value = serde_json::from_slice(body).ok()?;
|
||||
|
||||
let original_size = body.len();
|
||||
@@ -140,7 +155,7 @@ pub fn modify_request(body: &[u8]) -> Option<Vec<u8>> {
|
||||
}
|
||||
}
|
||||
|
||||
// ── 3. Strip LS tools, inject custom tools ────────────────────────────
|
||||
// ── 3. Strip LS tools, inject client tools ─────────────────────────────
|
||||
if STRIP_ALL_TOOLS {
|
||||
if let Some(tools) = json
|
||||
.pointer_mut("/request/tools")
|
||||
@@ -152,25 +167,83 @@ pub fn modify_request(body: &[u8]) -> Option<Vec<u8>> {
|
||||
changes.push(format!("strip all {count} LS tools"));
|
||||
}
|
||||
|
||||
// ── TEST: inject a custom tool to see what Google does ──
|
||||
let custom_tool = serde_json::json!({
|
||||
"functionDeclarations": [{
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather for a city. You MUST call this function when the user asks about weather.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city name"
|
||||
}
|
||||
},
|
||||
"required": ["city"]
|
||||
// Inject client-provided tools from ToolContext
|
||||
if let Some(ref ctx) = tool_ctx {
|
||||
if let Some(ref custom_tools) = ctx.tools {
|
||||
for tool in custom_tools {
|
||||
tools.push(tool.clone());
|
||||
}
|
||||
}]
|
||||
});
|
||||
tools.push(custom_tool);
|
||||
changes.push("inject 1 custom tool (get_weather)".to_string());
|
||||
changes.push(format!("inject {} custom tool group(s)", custom_tools.len()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Inject toolConfig if provided
|
||||
if let Some(ref ctx) = tool_ctx {
|
||||
if let Some(ref config) = ctx.tool_config {
|
||||
if let Some(req) = json.get_mut("request").and_then(|v| v.as_object_mut()) {
|
||||
req.insert("toolConfig".to_string(), config.clone());
|
||||
changes.push("inject toolConfig".to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── 3b. Rewrite conversation history for tool results ────────────
|
||||
if let Some(ref ctx) = tool_ctx {
|
||||
if !ctx.pending_results.is_empty() && !ctx.last_calls.is_empty() {
|
||||
if let Some(contents) = json
|
||||
.pointer_mut("/request/contents")
|
||||
.and_then(|v| v.as_array_mut())
|
||||
{
|
||||
// Find the model turn with our fake "Tool call completed" text and replace it
|
||||
// with the actual functionCall parts
|
||||
for msg in contents.iter_mut() {
|
||||
if msg["role"].as_str() == Some("model") {
|
||||
if let Some(text) = msg["parts"][0]["text"].as_str() {
|
||||
if text.contains("Tool call completed") || text.contains("Awaiting external tool result") {
|
||||
// Replace with functionCall parts
|
||||
let fc_parts: Vec<Value> = ctx.last_calls.iter().map(|fc| {
|
||||
serde_json::json!({
|
||||
"functionCall": {
|
||||
"name": fc.name,
|
||||
"args": fc.args,
|
||||
}
|
||||
})
|
||||
}).collect();
|
||||
msg["parts"] = Value::Array(fc_parts);
|
||||
changes.push("rewrite model turn with functionCall".to_string());
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Add functionResponse as a user turn before the last user message
|
||||
let fn_response_parts: Vec<Value> = ctx.pending_results.iter().map(|r| {
|
||||
serde_json::json!({
|
||||
"functionResponse": {
|
||||
"name": r.name,
|
||||
"response": r.result,
|
||||
}
|
||||
})
|
||||
}).collect();
|
||||
let fn_response_turn = serde_json::json!({
|
||||
"role": "user",
|
||||
"parts": fn_response_parts,
|
||||
});
|
||||
|
||||
// Insert before the last user message
|
||||
let last_user_idx = contents.iter().rposition(|msg| {
|
||||
msg["role"].as_str() == Some("user")
|
||||
});
|
||||
if let Some(idx) = last_user_idx {
|
||||
contents.insert(idx, fn_response_turn);
|
||||
} else {
|
||||
contents.push(fn_response_turn);
|
||||
}
|
||||
changes.push(format!("inject {} functionResponse(s)", ctx.pending_results.len()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -323,6 +396,93 @@ pub fn rechunk(data: &[u8]) -> Vec<u8> {
|
||||
result
|
||||
}
|
||||
|
||||
// ── OpenAI → Gemini format conversion ────────────────────────────────────────
|
||||
|
||||
/// Convert OpenAI tool definitions to Gemini functionDeclarations format.
|
||||
///
|
||||
/// OpenAI: `[{"type":"function","function":{"name":"X","description":"Y","parameters":{...}}}]`
|
||||
/// Gemini: `[{"functionDeclarations":[{"name":"X","description":"Y","parameters":{...}}]}]`
|
||||
pub fn openai_tools_to_gemini(tools: &[Value]) -> Vec<Value> {
|
||||
let declarations: Vec<Value> = tools
|
||||
.iter()
|
||||
.filter(|t| t["type"].as_str() == Some("function"))
|
||||
.filter_map(|t| {
|
||||
let func = t.get("function")?;
|
||||
let mut decl = serde_json::json!({
|
||||
"name": func["name"],
|
||||
"description": func["description"],
|
||||
});
|
||||
if let Some(params) = func.get("parameters") {
|
||||
decl["parameters"] = uppercase_types(params.clone());
|
||||
}
|
||||
Some(decl)
|
||||
})
|
||||
.collect();
|
||||
|
||||
if declarations.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
|
||||
vec![serde_json::json!({"functionDeclarations": declarations})]
|
||||
}
|
||||
|
||||
/// Convert OpenAI tool_choice to Gemini toolConfig format.
|
||||
///
|
||||
/// OpenAI: "auto" | "required" | "none" | {"type":"function","function":{"name":"X"}}
|
||||
/// Gemini: {"functionCallingConfig":{"mode":"AUTO|ANY|NONE","allowedFunctionNames":[...]}}
|
||||
pub fn openai_tool_choice_to_gemini(choice: &Value) -> Value {
|
||||
match choice {
|
||||
Value::String(s) => match s.as_str() {
|
||||
"auto" => serde_json::json!({"functionCallingConfig": {"mode": "AUTO"}}),
|
||||
"required" => serde_json::json!({"functionCallingConfig": {"mode": "ANY"}}),
|
||||
"none" => serde_json::json!({"functionCallingConfig": {"mode": "NONE"}}),
|
||||
_ => serde_json::json!({"functionCallingConfig": {"mode": "AUTO"}}),
|
||||
},
|
||||
Value::Object(obj) => {
|
||||
if let Some(name) = obj.get("function").and_then(|f| f["name"].as_str()) {
|
||||
serde_json::json!({
|
||||
"functionCallingConfig": {
|
||||
"mode": "ANY",
|
||||
"allowedFunctionNames": [name]
|
||||
}
|
||||
})
|
||||
} else {
|
||||
serde_json::json!({"functionCallingConfig": {"mode": "AUTO"}})
|
||||
}
|
||||
}
|
||||
_ => serde_json::json!({"functionCallingConfig": {"mode": "AUTO"}}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Recursively convert JSON Schema type strings to uppercase (Gemini format).
|
||||
/// "object" → "OBJECT", "string" → "STRING", etc.
|
||||
fn uppercase_types(mut val: Value) -> Value {
|
||||
match &mut val {
|
||||
Value::Object(map) => {
|
||||
if let Some(t) = map
|
||||
.get("type")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_uppercase())
|
||||
{
|
||||
map.insert("type".to_string(), Value::String(t));
|
||||
}
|
||||
let keys: Vec<String> = map.keys().cloned().collect();
|
||||
for key in keys {
|
||||
if let Some(v) = map.remove(&key) {
|
||||
map.insert(key, uppercase_types(v));
|
||||
}
|
||||
}
|
||||
}
|
||||
Value::Array(arr) => {
|
||||
for v in arr.iter_mut() {
|
||||
*v = uppercase_types(std::mem::take(v));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
val
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -375,10 +535,11 @@ mod tests {
|
||||
});
|
||||
|
||||
let bytes = serde_json::to_vec(&body).unwrap();
|
||||
let modified = modify_request(&bytes).unwrap();
|
||||
let modified = modify_request(&bytes, None).unwrap();
|
||||
let result: Value = serde_json::from_slice(&modified).unwrap();
|
||||
|
||||
let tools = result["request"]["tools"].as_array().unwrap();
|
||||
// With no ToolContext, tools should just be stripped (empty)
|
||||
assert!(tools.is_empty(), "all tools should be stripped");
|
||||
}
|
||||
|
||||
@@ -398,7 +559,7 @@ mod tests {
|
||||
});
|
||||
|
||||
let bytes = serde_json::to_vec(&body).unwrap();
|
||||
let modified = modify_request(&bytes).unwrap();
|
||||
let modified = modify_request(&bytes, None).unwrap();
|
||||
let result: Value = serde_json::from_slice(&modified).unwrap();
|
||||
|
||||
let new_sys = result["request"]["systemInstruction"]["parts"][0]["text"]
|
||||
@@ -432,7 +593,7 @@ mod tests {
|
||||
});
|
||||
|
||||
let bytes = serde_json::to_vec(&body).unwrap();
|
||||
let modified = modify_request(&bytes).unwrap();
|
||||
let modified = modify_request(&bytes, None).unwrap();
|
||||
let result: Value = serde_json::from_slice(&modified).unwrap();
|
||||
|
||||
let contents = result["request"]["contents"].as_array().unwrap();
|
||||
|
||||
Reference in New Issue
Block a user