gemini_adk_rs/tool/
context.rs

1//! What a tool knows about the call it serves.
2//!
3//! A tool is usually a pure function of its arguments. Some need more: the
4//! session [`State`] (the caller's account, a slot captured earlier), the
5//! call's id (for logs and idempotency), or a way to notice the user barged
6//! in and stop early. The runtime hands every call a [`ToolContext`]. A tool
7//! that wants it asks for it:
8//!
9//! - `#[tool]`: add a `ctx: ToolContext` parameter. It is filled by the
10//!   runtime and never shown to the model.
11//! - A closure: [`ContextTool`], or `T::contextual` in the fluent crate.
12//! - A hand-written [`ToolFunction`]: override
13//!   [`call_with_context`](ToolFunction::call_with_context).
14//!
15//! ```
16//! use gemini_adk_rs::tool::{ContextTool, ToolContext, ToolFunction};
17//! use gemini_adk_rs::State;
18//! use serde_json::json;
19//!
20//! # tokio_test::block_on(async {
21//! let balance = ContextTool::new("balance", "The caller's balance", None, |_args, ctx: ToolContext| async move {
22//!     let account: String = ctx.state.get("account_id").unwrap_or_default();
23//!     Ok(json!({ "account": account, "cents": 1200 }))
24//! });
25//!
26//! let state = State::new();
27//! state.set("account_id", "A-17").unwrap();
28//! let out = balance
29//!     .call_with_context(json!({}), ToolContext::new(state).with_call_id("call-1"))
30//!     .await
31//!     .unwrap();
32//! assert_eq!(out["account"], "A-17");
33//! # });
34//! ```
35
36use std::future::Future;
37use std::pin::Pin;
38use std::sync::Arc;
39
40use async_trait::async_trait;
41use serde_json::Value;
42use tokio_util::sync::CancellationToken;
43
44use super::ToolFunction;
45use crate::error::ToolError;
46use crate::state::State;
47
48/// The session a tool call runs in.
49#[derive(Debug, Clone)]
50#[non_exhaustive]
51pub struct ToolContext {
52    /// The session state: read what the conversation captured, write what
53    /// the tool learned.
54    pub state: State,
55    /// The model's id for this call, when it gave one.
56    pub call_id: Option<String>,
57    /// Cancelled when the call should stop: the user barged in, or the
58    /// session is ending. A long tool can check it between steps; the
59    /// runtime also drops the call's future when it fires.
60    pub cancel: CancellationToken,
61}
62
63impl ToolContext {
64    /// A context over `state`, with no call id and a token nothing cancels.
65    pub fn new(state: State) -> Self {
66        Self {
67            state,
68            call_id: None,
69            cancel: CancellationToken::new(),
70        }
71    }
72
73    /// A context for a call made outside any session (a direct
74    /// [`ToolFunction::call`], a unit test): fresh, empty state.
75    pub fn detached() -> Self {
76        Self::new(State::new())
77    }
78
79    /// With the model's call id.
80    pub fn with_call_id(mut self, id: impl Into<String>) -> Self {
81        self.call_id = Some(id.into());
82        self
83    }
84
85    /// With a cancellation token.
86    pub fn with_cancel(mut self, cancel: CancellationToken) -> Self {
87        self.cancel = cancel;
88        self
89    }
90}
91
92type ContextFn = Arc<
93    dyn Fn(Value, ToolContext) -> Pin<Box<dyn Future<Output = Result<Value, ToolError>> + Send>>
94        + Send
95        + Sync,
96>;
97
98/// A tool from a closure that receives the [`ToolContext`].
99pub struct ContextTool {
100    name: String,
101    description: String,
102    parameters: Option<Value>,
103    run: ContextFn,
104}
105
106impl ContextTool {
107    /// A tool named `name` running `f(args, ctx)`. `parameters` is the JSON
108    /// Schema of its arguments (`None` for none).
109    pub fn new<F, Fut>(
110        name: impl Into<String>,
111        description: impl Into<String>,
112        parameters: Option<Value>,
113        f: F,
114    ) -> Self
115    where
116        F: Fn(Value, ToolContext) -> Fut + Send + Sync + 'static,
117        Fut: Future<Output = Result<Value, ToolError>> + Send + 'static,
118    {
119        let f = Arc::new(f);
120        Self {
121            name: name.into(),
122            description: description.into(),
123            parameters,
124            run: Arc::new(move |args, ctx| {
125                let f = f.clone();
126                Box::pin(async move { f(args, ctx).await })
127            }),
128        }
129    }
130}
131
132#[async_trait]
133impl ToolFunction for ContextTool {
134    fn name(&self) -> &str {
135        &self.name
136    }
137
138    fn description(&self) -> &str {
139        &self.description
140    }
141
142    fn parameters(&self) -> Option<Value> {
143        self.parameters.clone()
144    }
145
146    async fn call(&self, args: Value) -> Result<Value, ToolError> {
147        self.call_with_context(args, ToolContext::detached()).await
148    }
149
150    async fn call_with_context(&self, args: Value, ctx: ToolContext) -> Result<Value, ToolError> {
151        (self.run)(args, ctx).await
152    }
153}
gemini_adk_rs/tool/context.rs

gemini_adk_rs/tool/
context.rs
