1use 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#[derive(Debug, Clone)]
25#[non_exhaustive]
26pub struct RunRequest {
27 pub input: Content,
29 pub history: Vec<Content>,
31 pub response_schema: Option<serde_json::Value>,
34}
35
36impl RunRequest {
37 pub fn new(text: impl Into<String>) -> Self {
39 Self::from_content(Content::user(text.into()))
40 }
41
42 pub fn from_content(input: Content) -> Self {
44 Self {
45 input,
46 history: Vec::new(),
47 response_schema: None,
48 }
49 }
50
51 pub fn history(mut self, history: Vec<Content>) -> Self {
53 self.history = history;
54 self
55 }
56
57 pub fn response_schema(mut self, schema: serde_json::Value) -> Self {
59 self.response_schema = Some(schema);
60 self
61 }
62
63 pub fn input_text(&self) -> String {
65 text_of(&self.input)
66 }
67}
68
69#[derive(Debug, Clone)]
71#[non_exhaustive]
72pub struct ToolCallRecord {
73 pub name: String,
75 pub args: serde_json::Value,
77 pub outcome: Result<serde_json::Value, ToolError>,
79}
80
81impl ToolCallRecord {
82 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#[derive(Debug, Clone, Default)]
98#[non_exhaustive]
99pub struct RunResult {
100 pub text: String,
102 pub messages: Vec<Content>,
106 pub usage: TokenUsage,
109 pub tool_calls: Vec<ToolCallRecord>,
111 pub model_calls: u32,
113}
114
115impl RunResult {
116 pub fn from_text(text: impl Into<String>) -> Self {
118 Self {
119 text: text.into(),
120 ..Self::default()
121 }
122 }
123
124 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#[derive(Debug, Clone)]
145#[non_exhaustive]
146pub enum RunEvent {
147 TextDelta(String),
149 ToolCall {
151 name: String,
153 args: serde_json::Value,
155 },
156 ToolResult(ToolCallRecord),
158 Finished(RunResult),
160}
161
162pub struct Chat<A> {
189 agent: A,
190 history: Vec<Content>,
191 state: State,
192 usage: TokenUsage,
193}
194
195impl<A: TextAgent> Chat<A> {
196 pub fn new(agent: A) -> Self {
198 Self::with_state(agent, State::new())
199 }
200
201 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 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 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 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 pub fn history(&self) -> &[Content] {
253 &self.history
254 }
255
256 pub fn usage(&self) -> TokenUsage {
258 self.usage
259 }
260
261 pub fn state(&self) -> &State {
263 &self.state
264 }
265
266 pub fn clear(&mut self) {
268 self.history.clear();
269 }
270}
271
272pub(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
284pub(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
292fn 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}