Skip to main content

kcode_gemini_api/
agent.rs

1use serde_json::{Value, json};
2
3use crate::{Error, Result, TokenUsage};
4
5const MAX_MODEL_CHARACTERS: usize = 128;
6const MAX_AGENT_INPUT_CHARACTERS: usize = 2_000_000;
7const MAX_TOOLS: usize = 64;
8
9/// One function available to a Gemini text turn.
10#[derive(Clone, Debug, Eq, PartialEq)]
11pub struct AgentTool {
12    /// Exact function name.
13    pub name: String,
14    /// Human-readable function description.
15    pub description: String,
16    /// JSON Schema object accepted by the function.
17    pub input_schema: Value,
18}
19
20/// One stateless Gemini text/tool interaction.
21#[derive(Clone, Debug, Eq, PartialEq)]
22pub struct AgentTurnRequest {
23    /// Exact Gemini model identifier.
24    pub model: String,
25    /// Complete input for this interaction.
26    pub input: String,
27    /// Provider-neutral reasoning effort.
28    pub reasoning_effort: String,
29    /// Functions available to the model.
30    pub tools: Vec<AgentTool>,
31}
32
33impl AgentTurnRequest {
34    /// Constructs a tool-free request with medium reasoning effort.
35    pub fn new(model: impl Into<String>, input: impl Into<String>) -> Self {
36        Self {
37            model: model.into(),
38            input: input.into(),
39            reasoning_effort: "medium".into(),
40            tools: Vec::new(),
41        }
42    }
43
44    pub(crate) fn validate(&self) -> Result<()> {
45        validate_model(&self.model)?;
46        if self.input.trim().is_empty() || self.input.chars().count() > MAX_AGENT_INPUT_CHARACTERS {
47            return Err(Error::InvalidInput(format!(
48                "agent input must contain 1 through {MAX_AGENT_INPUT_CHARACTERS} characters"
49            )));
50        }
51        if !matches!(
52            self.reasoning_effort.as_str(),
53            "none" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max"
54        ) {
55            return Err(Error::InvalidInput(
56                "reasoning effort is not supported".into(),
57            ));
58        }
59        if self.tools.len() > MAX_TOOLS {
60            return Err(Error::InvalidInput(format!(
61                "agent turn may expose at most {MAX_TOOLS} tools"
62            )));
63        }
64        for tool in &self.tools {
65            if tool.name.is_empty()
66                || tool.name.chars().count() > 100
67                || !tool
68                    .name
69                    .bytes()
70                    .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
71            {
72                return Err(Error::InvalidInput(
73                    "agent tool name must be a safe non-empty identifier".into(),
74                ));
75            }
76            if tool.description.trim().is_empty()
77                || tool.description.chars().count() > 4_000
78                || !tool.input_schema.is_object()
79            {
80                return Err(Error::InvalidInput(
81                    "agent tools require a bounded description and object JSON Schema".into(),
82                ));
83            }
84        }
85        Ok(())
86    }
87
88    pub(crate) fn payload(&self) -> Value {
89        let mut payload = json!({
90            "model": self.model,
91            "input": self.input,
92            "generation_config": {
93                "thinking_level": thinking_level(&self.reasoning_effort)
94            },
95            "service_tier": "standard",
96            "store": false
97        });
98        if !self.tools.is_empty() {
99            payload["tools"] = json!(
100                self.tools
101                    .iter()
102                    .map(|tool| json!({
103                        "type": "function",
104                        "name": tool.name,
105                        "description": tool.description,
106                        "parameters": tool.input_schema
107                    }))
108                    .collect::<Vec<_>>()
109            );
110        }
111        payload
112    }
113}
114
115/// One normalized function call returned by Gemini.
116#[derive(Clone, Debug, Eq, PartialEq)]
117pub struct AgentToolCall {
118    /// Provider call identifier.
119    pub call_id: String,
120    /// Exact function name.
121    pub name: String,
122    /// Parsed JSON arguments.
123    pub arguments: Value,
124}
125
126/// Normalized result of one Gemini text/tool interaction.
127#[derive(Clone, Debug, Eq, PartialEq)]
128pub struct AgentTurnResponse {
129    /// Gemini interaction identifier.
130    pub interaction_id: String,
131    /// Actual model identifier returned by Gemini.
132    pub model: String,
133    /// Ordered assistant text.
134    pub text: String,
135    /// At most one function call.
136    pub tool_call: Option<AgentToolCall>,
137    /// Provider token usage.
138    pub usage: TokenUsage,
139}
140
141/// Metadata advertised for one exact Gemini model.
142#[derive(Clone, Debug, Eq, PartialEq)]
143pub struct ModelMetadata {
144    /// Exact model identifier.
145    pub id: String,
146    /// Advertised context window, when returned.
147    pub context_window_tokens: Option<u64>,
148    /// Advertised maximum input, when returned.
149    pub max_input_tokens: Option<u64>,
150}
151
152pub(crate) fn validate_model(model: &str) -> Result<()> {
153    if model.trim().is_empty()
154        || model.chars().count() > MAX_MODEL_CHARACTERS
155        || !model.bytes().all(|byte| {
156            byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'-' | b'_' | b':' | b'/')
157        })
158    {
159        return Err(Error::InvalidInput(
160            "model must be an exact safe Gemini model identifier".into(),
161        ));
162    }
163    Ok(())
164}
165
166pub(crate) fn parse_tool_calls(value: &Value) -> Result<Option<AgentToolCall>> {
167    let mut calls = Vec::new();
168    for step in value
169        .get("steps")
170        .and_then(Value::as_array)
171        .into_iter()
172        .flatten()
173    {
174        if step.get("type").and_then(Value::as_str) == Some("function_call") {
175            calls.push(parse_call(step, calls.len())?);
176        }
177        for item in step
178            .get("content")
179            .and_then(Value::as_array)
180            .into_iter()
181            .flatten()
182        {
183            if item.get("type").and_then(Value::as_str) == Some("function_call") {
184                calls.push(parse_call(item, calls.len())?);
185            }
186        }
187    }
188    if calls.len() > 1 {
189        return Err(Error::Protocol(
190            "Gemini returned multiple function calls; agent tools must run serially".into(),
191        ));
192    }
193    Ok(calls.pop())
194}
195
196pub(crate) fn parse_model_metadata(value: &Value, requested: &str) -> ModelMetadata {
197    ModelMetadata {
198        id: value
199            .get("name")
200            .or_else(|| value.get("id"))
201            .and_then(Value::as_str)
202            .and_then(|value| value.strip_prefix("models/").or(Some(value)))
203            .unwrap_or(requested)
204            .to_owned(),
205        context_window_tokens: u64_field(
206            value,
207            &[
208                "contextWindowTokens",
209                "context_window_tokens",
210                "contextWindow",
211            ],
212        ),
213        max_input_tokens: u64_field(
214            value,
215            &[
216                "inputTokenLimit",
217                "input_token_limit",
218                "maxInputTokens",
219                "max_input_tokens",
220            ],
221        ),
222    }
223}
224
225fn parse_call(value: &Value, index: usize) -> Result<AgentToolCall> {
226    let arguments = match value.get("arguments") {
227        Some(Value::String(arguments)) => serde_json::from_str(arguments).map_err(|_| {
228            Error::Protocol("function call contained invalid JSON arguments".into())
229        })?,
230        Some(Value::Object(_)) => value["arguments"].clone(),
231        _ => {
232            return Err(Error::Protocol(
233                "function call omitted object arguments".into(),
234            ));
235        }
236    };
237    let name = value
238        .get("name")
239        .and_then(Value::as_str)
240        .filter(|name| !name.is_empty())
241        .ok_or_else(|| Error::Protocol("function call omitted its name".into()))?
242        .to_owned();
243    Ok(AgentToolCall {
244        call_id: value
245            .get("call_id")
246            .or_else(|| value.get("id"))
247            .and_then(Value::as_str)
248            .map(str::to_owned)
249            .unwrap_or_else(|| format!("gemini-call-{index}")),
250        name,
251        arguments,
252    })
253}
254
255fn thinking_level(effort: &str) -> &'static str {
256    match effort {
257        "none" | "minimal" => "minimal",
258        "low" => "low",
259        "medium" => "medium",
260        "high" | "xhigh" | "max" => "high",
261        _ => "medium",
262    }
263}
264
265fn u64_field(value: &Value, fields: &[&str]) -> Option<u64> {
266    fields
267        .iter()
268        .find_map(|field| value.get(*field).and_then(Value::as_u64))
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    #[test]
276    fn request_and_function_response_are_normalized() {
277        let mut request = AgentTurnRequest::new("gemini-3.1-pro-preview", "Do the task.");
278        request.tools.push(AgentTool {
279            name: "call_ktool".into(),
280            description: "Call one tool.".into(),
281            input_schema: json!({"type": "object"}),
282        });
283        assert_eq!(
284            request.payload()["generation_config"]["thinking_level"],
285            "medium"
286        );
287        let call = parse_tool_calls(&json!({
288            "steps": [{
289                "type": "model_output",
290                "content": [{
291                    "type": "function_call",
292                    "id": "call-1",
293                    "name": "call_ktool",
294                    "arguments": {"name": "Read"}
295                }]
296            }]
297        }))
298        .unwrap()
299        .unwrap();
300        assert_eq!(call.arguments["name"], "Read");
301    }
302}