gemini_adk_rs/text/
run.rs

1//! One run of a text agent: what goes in ([`RunRequest`]), what comes out
2//! ([`RunResult`]), and a conversation that remembers ([`Chat`]).
3
4use futures_util::StreamExt;
5use futures_util::stream::BoxStream;
6use gemini_genai_rs::prelude::{Content, Part, Role};
7use serde::de::DeserializeOwned;
8
9use super::TextAgent;
10use crate::error::{AgentError, ToolError};
11use crate::llm::TokenUsage;
12use crate::state::State;
13
14/// What to run: the new user turn, the conversation before it, and optionally
15/// the JSON shape the reply must take.
16///
17/// ```
18/// use gemini_adk_rs::text::RunRequest;
19///
20/// let request = RunRequest::new("And in Celsius?")
21///     .history(vec![/* earlier turns, e.g. from `RunResult::messages` */]);
22/// assert_eq!(request.input_text(), "And in Celsius?");
23/// ```
24#[derive(Debug, Clone)]
25#[non_exhaustive]
26pub struct RunRequest {
27    /// The new user turn.
28    pub input: Content,
29    /// Earlier turns, oldest first, sent before `input`.
30    pub history: Vec<Content>,
31    /// A JSON Schema the reply must match; sent as the response schema by
32    /// agents that call a model.
33    pub response_schema: Option<serde_json::Value>,
34}
35
36impl RunRequest {
37    /// A request whose new turn is `text`.
38    pub fn new(text: impl Into<String>) -> Self {
39        Self::from_content(Content::user(text.into()))
40    }
41
42    /// A request whose new turn is `input` — text with images, audio or files.
43    pub fn from_content(input: Content) -> Self {
44        Self {
45            input,
46            history: Vec::new(),
47            response_schema: None,
48        }
49    }
50
51    /// Send `history` before the new turn.
52    pub fn history(mut self, history: Vec<Content>) -> Self {
53        self.history = history;
54        self
55    }
56
57    /// Require the reply to be JSON matching `schema`.
58    pub fn response_schema(mut self, schema: serde_json::Value) -> Self {
59        self.response_schema = Some(schema);
60        self
61    }
62
63    /// The text of the new turn, its text parts joined.
64    pub fn input_text(&self) -> String {
65        text_of(&self.input)
66    }
67}
68
69/// A tool call the model made during a run, and what it returned.
70#[derive(Debug, Clone)]
71#[non_exhaustive]
72pub struct ToolCallRecord {
73    /// The tool's name.
74    pub name: String,
75    /// The arguments the model supplied.
76    pub args: serde_json::Value,
77    /// What the tool returned, or why it did not run.
78    pub outcome: Result<serde_json::Value, ToolError>,
79}
80
81impl ToolCallRecord {
82    /// Record one call.
83    pub fn new(
84        name: impl Into<String>,
85        args: serde_json::Value,
86        outcome: Result<serde_json::Value, ToolError>,
87    ) -> Self {
88        Self {
89            name: name.into(),
90            args,
91            outcome,
92        }
93    }
94}
95
96/// The outcome of a run: the final text, and what it took to get there.
97#[derive(Debug, Clone, Default)]
98#[non_exhaustive]
99pub struct RunResult {
100    /// The final reply.
101    pub text: String,
102    /// The turns this run added, oldest first: the request's input, every
103    /// model turn and every tool response. Append them to the history to
104    /// continue the conversation.
105    pub messages: Vec<Content>,
106    /// Tokens used by this agent's own model calls. Composite agents
107    /// (pipelines, fan-outs) do not add up their children's usage.
108    pub usage: TokenUsage,
109    /// Every tool call, in order.
110    pub tool_calls: Vec<ToolCallRecord>,
111    /// How many times the model was called.
112    pub model_calls: u32,
113}
114
115impl RunResult {
116    /// A result carrying only its final text.
117    pub fn from_text(text: impl Into<String>) -> Self {
118        Self {
119            text: text.into(),
120            ..Self::default()
121        }
122    }
123
124    /// Parse the reply as JSON into `T`.
125    ///
126    /// A reply wrapped in a Markdown code fence is unwrapped first. Fails with
127    /// [`AgentError::InvalidOutput`], which names the type and carries the
128    /// reply.
129    pub fn parse<T: DeserializeOwned>(&self) -> Result<T, AgentError> {
130        serde_json::from_str(strip_code_fence(&self.text)).map_err(|e| AgentError::InvalidOutput {
131            expected: std::any::type_name::<T>(),
132            reason: e.to_string(),
133            text: self.text.clone(),
134        })
135    }
136}
137
138/// Something that happened during a streamed run; see
139/// [`TextAgent::run_stream`](super::TextAgent::run_stream).
140///
141/// Concatenating every [`TextDelta`](Self::TextDelta) gives the text of the
142/// model's turns; [`Finished`](Self::Finished) is always the last event of a
143/// successful run and carries the same [`RunResult`] `run_with` returns.
144#[derive(Debug, Clone)]
145#[non_exhaustive]
146pub enum RunEvent {
147    /// The next piece of the model's reply.
148    TextDelta(String),
149    /// The model called a tool; it runs next.
150    ToolCall {
151        /// The tool's name.
152        name: String,
153        /// The arguments the model supplied.
154        args: serde_json::Value,
155    },
156    /// A tool finished.
157    ToolResult(ToolCallRecord),
158    /// The run is over.
159    Finished(RunResult),
160}
161
162/// A conversation with an agent: every [`send`](Self::send) carries the turns
163/// before it, and the conversation keeps its own [`State`].
164///
165/// ```
166/// use gemini_adk_rs::llm::{LlmResponse, MockLlm};
167/// use gemini_adk_rs::text::{LlmTextAgent, TextAgent};
168///
169/// # tokio_test::block_on(async {
170/// let llm = MockLlm::script([
171///     LlmResponse::from_text("Nice to meet you, Ada."),
172///     LlmResponse::from_text("Your name is Ada."),
173/// ]);
174/// let agent = LlmTextAgent::new("assistant", llm.clone());
175///
176/// let mut chat = agent.chat();
177/// chat.send("Hi, I'm Ada.").await.unwrap();
178/// assert_eq!(chat.send("What's my name?").await.unwrap(), "Your name is Ada.");
179///
180/// // The second request carried the whole conversation.
181/// assert_eq!(llm.last_request().unwrap().contents.len(), 3);
182/// assert_eq!(chat.history().len(), 4);
183/// # });
184/// ```
185///
186/// `agent.chat()` borrows the agent; `Chat::new(agent)` takes it, for a
187/// conversation that outlives the scope that built the agent.
188pub struct Chat<A> {
189    agent: A,
190    history: Vec<Content>,
191    state: State,
192    usage: TokenUsage,
193}
194
195impl<A: TextAgent> Chat<A> {
196    /// A new conversation with `agent`, with empty history and fresh state.
197    pub fn new(agent: A) -> Self {
198        Self::with_state(agent, State::new())
199    }
200
201    /// A new conversation that reads and writes `state`.
202    pub fn with_state(agent: A, state: State) -> Self {
203        Self {
204            agent,
205            history: Vec::new(),
206            state,
207            usage: TokenUsage::default(),
208        }
209    }
210
211    /// Send a message and get the reply. The turn is added to the history
212    /// only when it succeeds.
213    pub async fn send(&mut self, message: impl Into<String>) -> Result<String, AgentError> {
214        Ok(self.send_request(RunRequest::new(message)).await?.text)
215    }
216
217    /// Send a request (media, a response schema) and get the full result.
218    /// The request's own history is replaced by the conversation's.
219    pub async fn send_request(&mut self, request: RunRequest) -> Result<RunResult, AgentError> {
220        let request = request.history(self.history.clone());
221        let result = self.agent.run_with(request, &self.state).await?;
222        self.history.extend(result.messages.iter().cloned());
223        self.usage += result.usage;
224        Ok(result)
225    }
226
227    /// Send a message and stream the reply; the turn joins the history when
228    /// the stream delivers [`RunEvent::Finished`].
229    pub fn send_stream(
230        &mut self,
231        message: impl Into<String>,
232    ) -> BoxStream<'_, Result<RunEvent, AgentError>> {
233        let request = RunRequest::new(message).history(self.history.clone());
234        let Chat {
235            agent,
236            history,
237            state,
238            usage,
239        } = self;
240        agent
241            .run_stream(request, state.clone())
242            .inspect(move |event| {
243                if let Ok(RunEvent::Finished(result)) = event {
244                    history.extend(result.messages.iter().cloned());
245                    *usage += result.usage;
246                }
247            })
248            .boxed()
249    }
250
251    /// Every turn so far, oldest first.
252    pub fn history(&self) -> &[Content] {
253        &self.history
254    }
255
256    /// Tokens used by the conversation so far.
257    pub fn usage(&self) -> TokenUsage {
258        self.usage
259    }
260
261    /// The conversation's state.
262    pub fn state(&self) -> &State {
263        &self.state
264    }
265
266    /// Forget the history (the state is kept).
267    pub fn clear(&mut self) {
268        self.history.clear();
269    }
270}
271
272/// The text parts of a turn, joined.
273pub(crate) fn text_of(content: &Content) -> String {
274    content
275        .parts
276        .iter()
277        .filter_map(|p| match p {
278            Part::Text { text } => Some(text.as_str()),
279            _ => None,
280        })
281        .collect()
282}
283
284/// A model turn saying `text`.
285pub(crate) fn model_turn(text: impl Into<String>) -> Content {
286    Content {
287        role: Some(Role::Model),
288        parts: vec![Part::Text { text: text.into() }],
289    }
290}
291
292/// `text` without a surrounding Markdown code fence (```` ```json … ``` ````).
293fn strip_code_fence(text: &str) -> &str {
294    let trimmed = text.trim();
295    let Some(body) = trimmed.strip_prefix("```") else {
296        return trimmed;
297    };
298    let body = body.split_once('\n').map_or("", |(_, rest)| rest);
299    body.trim_end().strip_suffix("```").unwrap_or(body).trim()
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305
306    #[test]
307    fn parse_reads_plain_and_fenced_json() {
308        #[derive(serde::Deserialize, Debug, PartialEq)]
309        struct City {
310            name: String,
311        }
312        let plain = RunResult::from_text(r#"{"name":"Paris"}"#);
313        assert_eq!(plain.parse::<City>().unwrap().name, "Paris");
314        let fenced = RunResult::from_text("```json\n{\"name\": \"Lyon\"}\n```");
315        assert_eq!(fenced.parse::<City>().unwrap().name, "Lyon");
316        let err = RunResult::from_text("Paris").parse::<City>().unwrap_err();
317        assert!(
318            matches!(&err, AgentError::InvalidOutput { text, expected, .. }
319                if text == "Paris" && expected.ends_with("City")),
320            "{err}"
321        );
322    }
323
324    #[test]
325    fn request_text_joins_text_parts() {
326        let request = RunRequest::from_content(Content {
327            role: Some(Role::User),
328            parts: vec![
329                Part::Text { text: "a".into() },
330                Part::inline_data("image/png", "AAAA"),
331                Part::Text { text: "b".into() },
332            ],
333        });
334        assert_eq!(request.input_text(), "ab");
335    }
336}