gemini_adk_rs/text/
llm.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use gemini_genai_rs::prelude::{Content, FunctionCall, FunctionResponse, Part, Role};
5
6use super::TextAgent;
7use crate::context::AgentEvent;
8use crate::error::AgentError;
9use crate::llm::{BaseLlm, LlmRequest};
10use crate::middleware::MiddlewareChain;
11use crate::state::State;
12use crate::tool::ToolDispatcher;
13
14/// Maximum number of tool-dispatch round-trips before giving up.
15const MAX_TOOL_ROUNDS: usize = 10;
16
17/// A dynamic model source: state in, the model to use for this run out.
18type LlmProviderFn = Arc<dyn Fn(&State) -> Arc<dyn BaseLlm> + Send + Sync>;
19
20/// Core text agent — calls `BaseLlm::generate()`, dispatches tools, loops
21/// until the model produces a final text response.
22///
23/// Middleware hooks fire at each lifecycle point:
24///
25/// - `before_model` / `after_model` — wraps each `BaseLlm::generate()` call;
26///   `before_model` may return a cached response to skip the LLM entirely.
27/// - `before_tool` / `after_tool` / `on_tool_error` — wraps each tool dispatch.
28/// - `on_error` — called when `run()` is about to return an error.
29///
30/// Note: `before_agent`/`after_agent` are Live-session hooks that require an
31/// `InvocationContext` (a Live WebSocket concept) and are therefore not invoked
32/// by `LlmTextAgent`.  Use `before_model` or wrap in a custom `TextAgent` if you
33/// need entry/exit hooks for the text path.
34pub struct LlmTextAgent {
35    name: String,
36    llm: Arc<dyn BaseLlm>,
37    instruction: Option<String>,
38    /// Dynamic instruction source, resolved against state on every run;
39    /// wins over the static `instruction` when both are set.
40    instruction_provider: Option<Arc<dyn crate::instruction::InstructionProvider>>,
41    /// Dynamic model source, resolved against state on every run;
42    /// wins over the constructor's model when set. Risk-based escalation,
43    /// cost routing, per-tenant model selection without rebuilding the agent.
44    llm_provider: Option<LlmProviderFn>,
45    dispatcher: Option<Arc<ToolDispatcher>>,
46    temperature: Option<f32>,
47    max_output_tokens: Option<u32>,
48    middleware: MiddlewareChain,
49}
50
51impl LlmTextAgent {
52    /// Create a new LLM text agent.
53    pub fn new(name: impl Into<String>, llm: Arc<dyn BaseLlm>) -> Self {
54        Self {
55            name: name.into(),
56            llm,
57            instruction: None,
58            instruction_provider: None,
59            llm_provider: None,
60            dispatcher: None,
61            temperature: None,
62            max_output_tokens: None,
63            middleware: MiddlewareChain::new(),
64        }
65    }
66
67    /// Set the system instruction.
68    pub fn instruction(mut self, inst: impl Into<String>) -> Self {
69        self.instruction = Some(inst.into());
70        self
71    }
72
73    /// Set a dynamic instruction source — an
74    /// [`InstructionProvider`](crate::instruction::InstructionProvider)
75    /// (any `Fn(&State) -> String` closure, or a `TemplateInstruction`
76    /// under the `templates` feature) resolved against session state at
77    /// the start of every run. Wins over [`instruction`](Self::instruction).
78    pub fn instruction_provider(
79        mut self,
80        provider: impl crate::instruction::InstructionProvider + 'static,
81    ) -> Self {
82        self.instruction_provider = Some(Arc::new(provider));
83        self
84    }
85
86    /// Set a dynamic model source, resolved against session state at the
87    /// start of every run — risk-based escalation to a stronger model, cost
88    /// routing to a cheaper one, per-tenant model selection — without
89    /// rebuilding the agent. Wins over the constructor's model when set.
90    pub fn llm_provider<F>(mut self, provider: F) -> Self
91    where
92        F: Fn(&State) -> Arc<dyn BaseLlm> + Send + Sync + 'static,
93    {
94        self.llm_provider = Some(Arc::new(provider));
95        self
96    }
97
98    /// Set the tool dispatcher.
99    pub fn tools(mut self, dispatcher: Arc<ToolDispatcher>) -> Self {
100        self.dispatcher = Some(dispatcher);
101        self
102    }
103
104    /// Set temperature.
105    pub fn temperature(mut self, t: f32) -> Self {
106        self.temperature = Some(t);
107        self
108    }
109
110    /// Set max output tokens.
111    pub fn max_output_tokens(mut self, n: u32) -> Self {
112        self.max_output_tokens = Some(n);
113        self
114    }
115
116    /// Append a middleware layer to the chain.
117    ///
118    /// Layers are run in insertion order for `before_*` / `on_error` hooks
119    /// and in reverse insertion order for `after_*` hooks (outermost last).
120    pub fn add_middleware(mut self, mw: Arc<dyn crate::middleware::Middleware>) -> Self {
121        self.middleware.add(mw);
122        self
123    }
124
125    /// Replace the entire middleware chain (advanced — prefer `add_middleware`).
126    pub fn with_middleware_chain(mut self, chain: MiddlewareChain) -> Self {
127        self.middleware = chain;
128        self
129    }
130
131    /// Build an LlmRequest, taking ownership of contents to avoid cloning.
132    fn build_request(&self, contents: Vec<Content>, instruction: &Option<String>) -> LlmRequest {
133        let mut req = LlmRequest::from_contents(contents);
134        req.system_instruction = instruction.clone();
135        req.temperature = self.temperature;
136        req.max_output_tokens = self.max_output_tokens;
137
138        if let Some(dispatcher) = &self.dispatcher {
139            req.tools = dispatcher.to_tool_declarations();
140        }
141
142        req
143    }
144
145    /// Dispatch function calls and return function responses, firing middleware hooks.
146    async fn dispatch_tools(&self, calls: &[FunctionCall]) -> Vec<FunctionResponse> {
147        let dispatcher = match &self.dispatcher {
148            Some(d) => d,
149            None => return Vec::new(),
150        };
151
152        let mut responses = Vec::with_capacity(calls.len());
153        for call in calls {
154            // before_tool hook
155            if let Err(e) = self.middleware.run_before_tool(call).await {
156                // Hook error — record it and return an error response.
157                let _ = self
158                    .middleware
159                    .run_on_tool_error(
160                        call,
161                        &crate::error::ToolError::ExecutionFailed(e.to_string()),
162                    )
163                    .await;
164                responses.push(ToolDispatcher::build_response(
165                    call,
166                    Err(crate::error::ToolError::ExecutionFailed(e.to_string())),
167                ));
168                continue;
169            }
170
171            let result = dispatcher
172                .call_function(&call.name, call.args.clone())
173                .await;
174
175            match &result {
176                Ok(value) => {
177                    let _ = self.middleware.run_after_tool(call, value).await;
178                }
179                Err(e) => {
180                    let _ = self.middleware.run_on_tool_error(call, e).await;
181                }
182            }
183
184            responses.push(ToolDispatcher::build_response(call, result));
185        }
186        responses
187    }
188}
189
190#[async_trait]
191impl TextAgent for LlmTextAgent {
192    fn name(&self) -> &str {
193        &self.name
194    }
195
196    async fn run(&self, state: &State) -> Result<String, AgentError> {
197        // Build initial contents from state "input" key, or empty user message.
198        let input = state.get::<String>("input").unwrap_or_default();
199
200        let mut contents = vec![Content::user(&input)];
201
202        // Resolve the instruction for this run: provider (against live
203        // state) wins over the static string.
204        let instruction = match &self.instruction_provider {
205            Some(provider) => Some(provider.provide(state)),
206            None => self.instruction.clone(),
207        };
208
209        // Resolve the LLM for this run: provider (against live state) wins
210        // over the constructor's LLM.
211        let llm = self
212            .llm_provider
213            .as_ref()
214            .map(|p| p(state))
215            .unwrap_or_else(|| self.llm.clone());
216
217        // Lifecycle event — makes `on_event` (e.g. M::tap) observe agent start.
218        let _ = self
219            .middleware
220            .run_on_event(&AgentEvent::AgentStarted {
221                name: self.name.clone(),
222            })
223            .await;
224
225        // Enforce the tightest middleware timeout (M::timeout) over the whole run.
226        let result = match self.middleware.timeout() {
227            Some(limit) => {
228                match tokio::time::timeout(limit, self.run_inner(&mut contents, &instruction, &llm))
229                    .await
230                {
231                    Ok(r) => r,
232                    Err(_) => {
233                        let _ = self.middleware.run_on_event(&AgentEvent::Timeout).await;
234                        Err(AgentError::Other(format!(
235                            "agent '{}' timed out after {:?}",
236                            self.name, limit
237                        )))
238                    }
239                }
240            }
241            None => self.run_inner(&mut contents, &instruction, &llm).await,
242        };
243
244        if let Err(ref e) = result {
245            let _ = self.middleware.run_on_error(e).await;
246        } else if let Ok(ref text) = result {
247            let _ = state.set("output", text);
248            let _ = self
249                .middleware
250                .run_on_event(&AgentEvent::AgentCompleted {
251                    name: self.name.clone(),
252                })
253                .await;
254        }
255
256        result
257    }
258}
259
260impl LlmTextAgent {
261    /// Inner execution loop — separated so `on_error` fires exactly once.
262    async fn run_inner(
263        &self,
264        contents: &mut Vec<Content>,
265        instruction: &Option<String>,
266        llm: &Arc<dyn BaseLlm>,
267    ) -> Result<String, AgentError> {
268        for _round in 0..MAX_TOOL_ROUNDS {
269            let mut request = self.build_request(contents.clone(), instruction);
270
271            // transform_request hook — may rewrite the request (e.g. context
272            // policies trimming conversation history) before it is sent.
273            self.middleware.run_transform_request(&mut request).await?;
274
275            // before_model hook — may short-circuit with a cached response.
276            let response = match self.middleware.run_before_model(&request).await? {
277                Some(cached) => cached,
278                None => {
279                    let llm_response = llm
280                        .generate(request.clone())
281                        .await
282                        .map_err(|e| AgentError::Other(format!("LLM error: {e}")))?;
283
284                    // after_model hook — may replace the response.
285                    match self
286                        .middleware
287                        .run_after_model(&request, &llm_response)
288                        .await?
289                    {
290                        Some(replaced) => replaced,
291                        None => llm_response,
292                    }
293                }
294            };
295
296            let calls: Vec<FunctionCall> = response.function_calls().into_iter().cloned().collect();
297
298            if calls.is_empty() {
299                // No tool calls — we have a final text response.
300                return Ok(response.text());
301            }
302
303            // Move model response into conversation (no clone needed).
304            contents.push(response.content);
305
306            // Dispatch tools (middleware hooks inside). Media a tool
307            // attached under `_media` is lifted out of the JSON and
308            // delivered as inline_data parts in the same turn, so the
309            // model *sees* images rather than base64 noise.
310            let tool_responses = self.dispatch_tools(&calls).await;
311            let mut media_parts: Vec<Part> = Vec::new();
312            let mut response_parts: Vec<Part> = tool_responses
313                .into_iter()
314                .map(|mut fr| {
315                    for attachment in crate::tool::media::extract(&mut fr.response) {
316                        media_parts.push(Part::inline_data(
317                            attachment.mime_type,
318                            attachment.data_base64,
319                        ));
320                    }
321                    Part::FunctionResponse {
322                        function_response: fr,
323                    }
324                })
325                .collect();
326            response_parts.append(&mut media_parts);
327
328            contents.push(Content {
329                role: Some(Role::User),
330                parts: response_parts,
331            });
332        }
333
334        Err(AgentError::Other(format!(
335            "Agent '{}' exceeded max tool rounds ({})",
336            self.name, MAX_TOOL_ROUNDS
337        )))
338    }
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344    use crate::context::AgentEvent;
345    use crate::llm::{LlmError, LlmResponse};
346    use crate::middleware::Middleware;
347    use gemini_genai_rs::prelude::{Content, Part, Role};
348    use std::sync::atomic::{AtomicBool, Ordering};
349    use std::time::Duration;
350
351    fn text_response(t: &str) -> LlmResponse {
352        LlmResponse {
353            content: Content {
354                role: Some(Role::Model),
355                parts: vec![Part::Text { text: t.into() }],
356            },
357            finish_reason: Some("STOP".into()),
358            usage: None,
359        }
360    }
361
362    /// LLM that returns a function call on the first request, text on the
363    /// second, capturing every request it sees.
364    struct CapturingLlm {
365        requests: std::sync::Mutex<Vec<LlmRequest>>,
366    }
367    #[async_trait]
368    impl BaseLlm for CapturingLlm {
369        fn model_id(&self) -> &str {
370            "capturing"
371        }
372        async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
373            let mut requests = self.requests.lock().unwrap();
374            requests.push(req);
375            if requests.len() == 1 {
376                Ok(LlmResponse {
377                    content: Content {
378                        role: Some(Role::Model),
379                        parts: vec![Part::FunctionCall {
380                            function_call: gemini_genai_rs::prelude::FunctionCall {
381                                name: "snap".into(),
382                                args: serde_json::json!({}),
383                                id: None,
384                            },
385                        }],
386                    },
387                    finish_reason: None,
388                    usage: None,
389                })
390            } else {
391                Ok(text_response("described"))
392            }
393        }
394    }
395
396    #[tokio::test]
397    async fn tool_media_reaches_the_model_as_inline_data() {
398        use crate::tool::{SimpleTool, ToolDispatcher, media};
399        let llm = Arc::new(CapturingLlm {
400            requests: std::sync::Mutex::new(Vec::new()),
401        });
402        let mut dispatcher = ToolDispatcher::new();
403        dispatcher.register(SimpleTool::new(
404            "snap",
405            "Take a snapshot",
406            None,
407            |_args| async move {
408                let mut result = serde_json::json!({"took": true});
409                media::attach(&mut result, "image/png", b"fakepng");
410                Ok(result)
411            },
412        ));
413        let agent = LlmTextAgent::new("vision", llm.clone()).tools(Arc::new(dispatcher));
414        let state = State::new();
415        let _ = state.set("input", "what do you see?");
416        assert_eq!(agent.run(&state).await.unwrap(), "described");
417
418        let requests = llm.requests.lock().unwrap();
419        assert_eq!(requests.len(), 2);
420        // The second request's tool-response turn carries the image part and
421        // the function response JSON no longer contains the base64 blob.
422        let turn = requests[1].contents.last().unwrap();
423        let has_inline = turn
424            .parts
425            .iter()
426            .any(|p| matches!(p, Part::InlineData { .. }));
427        assert!(has_inline, "expected an inline_data part, got {turn:?}");
428        let fr_clean = turn.parts.iter().all(|p| match p {
429            Part::FunctionResponse { function_response } => {
430                function_response.response.get(media::MEDIA_KEY).is_none()
431            }
432            _ => true,
433        });
434        assert!(
435            fr_clean,
436            "media key should be stripped from the response JSON"
437        );
438    }
439
440    #[tokio::test]
441    async fn instruction_provider_resolves_against_state_each_run() {
442        let llm = Arc::new(CapturingLlm {
443            requests: std::sync::Mutex::new(Vec::new()),
444        });
445        let agent = LlmTextAgent::new("persona", llm.clone()).instruction_provider(|s: &State| {
446            format!(
447                "You are {}.",
448                s.get::<String>("persona").unwrap_or_default()
449            )
450        });
451        let state = State::new();
452        let _ = state.set("input", "hi");
453        let _ = state.set("persona", "a pirate");
454        // CapturingLlm returns a function call first; with no dispatcher the
455        // loop sends an empty tool-response turn and the second reply ends
456        // the run — both requests must carry the resolved instruction.
457        let _ = agent.run(&state).await.unwrap();
458        let requests = llm.requests.lock().unwrap();
459        assert!(
460            requests
461                .iter()
462                .all(|r| r.system_instruction.as_deref() == Some("You are a pirate."))
463        );
464    }
465
466    struct SlowLlm;
467    #[async_trait]
468    impl BaseLlm for SlowLlm {
469        fn model_id(&self) -> &str {
470            "slow"
471        }
472        async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
473            tokio::time::sleep(Duration::from_millis(500)).await;
474            Ok(text_response("done"))
475        }
476    }
477
478    struct FastLlm;
479    #[async_trait]
480    impl BaseLlm for FastLlm {
481        fn model_id(&self) -> &str {
482            "fast"
483        }
484        async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
485            Ok(text_response("hi"))
486        }
487    }
488
489    struct ShortTimeout;
490    #[async_trait]
491    impl Middleware for ShortTimeout {
492        fn name(&self) -> &str {
493            "short-timeout"
494        }
495        fn timeout(&self) -> Option<Duration> {
496            Some(Duration::from_millis(20))
497        }
498    }
499
500    struct EventFlag(Arc<AtomicBool>);
501    #[async_trait]
502    impl Middleware for EventFlag {
503        fn name(&self) -> &str {
504            "event-flag"
505        }
506        async fn on_event(&self, event: &AgentEvent) -> Result<(), AgentError> {
507            if matches!(event, AgentEvent::AgentStarted { .. }) {
508                self.0.store(true, Ordering::SeqCst);
509            }
510            Ok(())
511        }
512    }
513
514    #[tokio::test]
515    async fn timeout_aborts_slow_run() {
516        let agent =
517            LlmTextAgent::new("slowpoke", Arc::new(SlowLlm)).add_middleware(Arc::new(ShortTimeout));
518        let state = State::new();
519        let _ = state.set("input", "hi");
520        let err = agent.run(&state).await.expect_err("expected timeout");
521        assert!(format!("{err:?}").contains("timed out"), "got: {err:?}");
522    }
523
524    #[tokio::test]
525    async fn on_event_fires_for_agent_lifecycle() {
526        let flag = Arc::new(AtomicBool::new(false));
527        let agent = LlmTextAgent::new("a", Arc::new(FastLlm))
528            .add_middleware(Arc::new(EventFlag(flag.clone())));
529        let state = State::new();
530        let _ = state.set("input", "hi");
531        let _ = agent.run(&state).await;
532        assert!(
533            flag.load(Ordering::SeqCst),
534            "on_event(AgentStarted) should fire"
535        );
536    }
537
538    #[tokio::test]
539    async fn llm_provider_switches_model_per_run() {
540        // Two mock LLMs with distinguishable responses.
541        struct MockLlmA;
542        #[async_trait]
543        impl BaseLlm for MockLlmA {
544            fn model_id(&self) -> &str {
545                "mock-a"
546            }
547            async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
548                Ok(text_response("from-a"))
549            }
550        }
551
552        struct MockLlmB;
553        #[async_trait]
554        impl BaseLlm for MockLlmB {
555            fn model_id(&self) -> &str {
556                "mock-b"
557            }
558            async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
559                Ok(text_response("from-b"))
560            }
561        }
562
563        let llm_a = Arc::new(MockLlmA);
564        let llm_b = Arc::new(MockLlmB);
565
566        // Agent with provider that switches based on escalate flag.
567        let agent = LlmTextAgent::new("switcher", llm_a.clone()).llm_provider(move |state| {
568            if state.get::<bool>("escalate").unwrap_or(false) {
569                llm_b.clone()
570            } else {
571                llm_a.clone()
572            }
573        });
574
575        // Run 1: without escalate flag -> should use model A
576        let state = State::new();
577        let _ = state.set("input", "hi");
578        let result = agent.run(&state).await.unwrap();
579        assert_eq!(result, "from-a", "without escalate, should use model A");
580
581        // Run 2: with escalate flag -> should use model B
582        let state2 = State::new();
583        let _ = state2.set("input", "hi");
584        let _ = state2.set("escalate", true);
585        let result2 = agent.run(&state2).await.unwrap();
586        assert_eq!(result2, "from-b", "with escalate, should use model B");
587    }
588}
gemini_adk_rs/text/llm.rs

gemini_adk_rs/text/
llm.rs
