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    /// Task activation ownership, captured at admission when using skills.
58    pub task: Option<crate::tasks::TaskOwner>,
59    /// Cancelled when the call should stop: the user barged in, or the
60    /// session is ending. A long tool can check it between steps; the
61    /// runtime also drops the call's future when it fires.
62    pub cancel: CancellationToken,
63}
64
65impl ToolContext {
66    /// A context over `state`, with no call id and a token nothing cancels.
67    pub fn new(state: State) -> Self {
68        Self {
69            state,
70            call_id: None,
71            task: None,
72            cancel: CancellationToken::new(),
73        }
74    }
75
76    /// A context for a call made outside any session (a direct
77    /// [`ToolFunction::call`], a unit test): fresh, empty state.
78    pub fn detached() -> Self {
79        Self::new(State::new())
80    }
81
82    /// With the model's call id.
83    pub fn with_call_id(mut self, id: impl Into<String>) -> Self {
84        self.call_id = Some(id.into());
85        self
86    }
87
88    /// Bind execution and tracing to its original task activation.
89    pub fn with_task(mut self, owner: crate::tasks::TaskOwner) -> Self {
90        self.task = Some(owner);
91        self
92    }
93
94    /// With a cancellation token.
95    pub fn with_cancel(mut self, cancel: CancellationToken) -> Self {
96        self.cancel = cancel;
97        self
98    }
99}
100
101type ContextFn = Arc<
102    dyn Fn(Value, ToolContext) -> Pin<Box<dyn Future<Output = Result<Value, ToolError>> + Send>>
103        + Send
104        + Sync,
105>;
106
107/// A tool from a closure that receives the [`ToolContext`].
108pub struct ContextTool {
109    name: String,
110    description: String,
111    parameters: Option<Value>,
112    run: ContextFn,
113}
114
115impl ContextTool {
116    /// A tool named `name` running `f(args, ctx)`. `parameters` is the JSON
117    /// Schema of its arguments (`None` for none).
118    pub fn new<F, Fut>(
119        name: impl Into<String>,
120        description: impl Into<String>,
121        parameters: Option<Value>,
122        f: F,
123    ) -> Self
124    where
125        F: Fn(Value, ToolContext) -> Fut + Send + Sync + 'static,
126        Fut: Future<Output = Result<Value, ToolError>> + Send + 'static,
127    {
128        let f = Arc::new(f);
129        Self {
130            name: name.into(),
131            description: description.into(),
132            parameters,
133            run: Arc::new(move |args, ctx| {
134                let f = f.clone();
135                Box::pin(async move { f(args, ctx).await })
136            }),
137        }
138    }
139}
140
141#[async_trait]
142impl ToolFunction for ContextTool {
143    fn name(&self) -> &str {
144        &self.name
145    }
146
147    fn description(&self) -> &str {
148        &self.description
149    }
150
151    fn parameters(&self) -> Option<Value> {
152        self.parameters.clone()
153    }
154
155    async fn call(&self, args: Value) -> Result<Value, ToolError> {
156        self.call_with_context(args, ToolContext::detached()).await
157    }
158
159    async fn call_with_context(&self, args: Value, ctx: ToolContext) -> Result<Value, ToolError> {
160        (self.run)(args, ctx).await
161    }
162}
gemini_adk_rs/tool/context.rs

gemini_adk_rs/tool/
context.rs
