1use super::response_handler::{
2 handle_final_answer_messages, handle_final_answer_steps, handle_structured_response_messages,
3 handle_structured_response_steps, ErrorSink, FinalAnswerContext, FinalAnswerStepsContext,
4 HandlerOutcome, StructuredResponseContext, StructuredResponseStepsContext,
5};
6use crate::{
7 core::{agent::Agent, memory::AgentMemory, steps::AgentStep},
8 error::{AgentError, Result},
9 schemas::validation::{
10 final_answer_tool_definition, inject_schema_instructions,
11 structured_response_tool_definition, structured_response_tool_name,
12 },
13 services::{
14 openai_client::ChatCompletionRequest,
15 tool_call_utils::{
16 extract_arguments_str, extract_function_info, extract_tool_call_id,
17 parse_function_arguments,
18 },
19 },
20 types::result::{RunResult, TokenUsage},
21};
22use serde_json::{json, Value};
23use std::time::Instant;
24use tokio::time::timeout;
25
26struct MemorySink<'a> {
28 memory: &'a mut AgentMemory,
29}
30
31impl<'a> ErrorSink for MemorySink<'a> {
32 fn report_error(&mut self, tool_call_id: &str, error_message: String) {
33 self.memory.add_step(AgentStep::Observation {
34 tool_call_id: tool_call_id.to_string(),
35 result: error_message,
36 is_error: true,
37 });
38 }
39
40 fn report_observation(&mut self, tool_call_id: &str, result: String, is_error: bool) {
41 self.memory.add_step(AgentStep::Observation {
42 tool_call_id: tool_call_id.to_string(),
43 result,
44 is_error,
45 });
46 }
47}
48
49struct MessagesSink<'a> {
51 messages: &'a mut Vec<Value>,
52}
53
54impl<'a> ErrorSink for MessagesSink<'a> {
55 fn report_error(&mut self, tool_call_id: &str, error_message: String) {
56 self.messages.push(json!({
57 "role": "tool",
58 "tool_call_id": tool_call_id,
59 "content": error_message
60 }));
61 }
62
63 fn report_observation(&mut self, tool_call_id: &str, result: String, _is_error: bool) {
64 self.messages.push(json!({
65 "role": "tool",
66 "tool_call_id": tool_call_id,
67 "content": result
68 }));
69 }
70}
71
72impl Agent {
73 pub async fn run_with_steps(&self, prompt: &str) -> Result<RunResult> {
74 let start_time = Instant::now();
75 let mut memory = AgentMemory::with_default_system();
76
77 memory.add_step(AgentStep::Task {
78 content: prompt.to_string(),
79 });
80
81 let mut iteration = 0;
82 let mut has_final_answer = false;
83 let mut final_answer_value: Option<String> = None;
84
85 while iteration < self.max_iterations() {
86 iteration += 1;
87
88 let mut messages = memory.as_messages();
89 if let Some(schema) = self.completion_schema() {
90 inject_schema_instructions(&mut messages, schema);
91 }
92
93 let mut tools = self.function_factory().get_openai_tools();
94 if let Some(schema) = self.completion_schema() {
95 tools.push(structured_response_tool_definition(schema));
96 } else {
97 tools.push(final_answer_tool_definition());
98 }
99
100 let mut chat_request =
101 ChatCompletionRequest::new(self.model().to_owned(), messages.clone())
102 .with_max_tokens(self.max_tokens());
103
104 if !tools.is_empty() {
105 chat_request = chat_request
106 .with_tools(tools)
107 .with_tool_choice(json!("auto"));
108 }
109
110 let request_body = chat_request.into_value();
111
112 let response = timeout(self.timeout(), self.make_raw_request(&request_body))
113 .await
114 .map_err(|_| AgentError::Timeout("OpenAI API call timed out".to_string()))??;
115
116 let choices = response
117 .get("choices")
118 .and_then(|value| value.as_array())
119 .ok_or_else(|| {
120 AgentError::Unknown(
121 "Missing 'choices' array in completion response".to_string(),
122 )
123 })?;
124
125 let first_choice = choices.first().ok_or_else(|| {
126 AgentError::Unknown("Completion response contained no choices".to_string())
127 })?;
128
129 let assistant_message = first_choice.get("message").cloned().ok_or_else(|| {
130 AgentError::Unknown("Completion response missing assistant message".to_string())
131 })?;
132
133 let token_usage = response.get("usage").and_then(|usage| {
134 Some(TokenUsage {
135 prompt_tokens: usage.get("prompt_tokens")?.as_u64()? as u32,
136 completion_tokens: usage.get("completion_tokens")?.as_u64()? as u32,
137 total_tokens: usage.get("total_tokens")?.as_u64()? as u32,
138 })
139 });
140
141 if let Some(tool_calls) = assistant_message.get("tool_calls") {
142 if let Some(tool_calls_array) = tool_calls.as_array() {
143 let turn_has_final_answer = tool_calls_array.iter().any(|tool_call| {
144 tool_call
145 .get("function")
146 .and_then(|f| f.get("name"))
147 .and_then(|name| name.as_str())
148 .map(|name| name == "final_answer")
149 .unwrap_or(false)
150 });
151
152 if turn_has_final_answer && tool_calls_array.len() > 1 {
153 memory.add_step(AgentStep::Observation {
154 tool_call_id: "final_answer".to_string(),
155 result: AgentError::InvalidFunctionCall(
156 "`final_answer` must be the only tool call in a single turn"
157 .to_string(),
158 )
159 .to_error_payload()
160 .to_string(),
161 is_error: true,
162 });
163 continue;
164 }
165
166 for tool_call in tool_calls_array {
167 let tool_call_id = extract_tool_call_id(tool_call);
168
169 let (function, function_name_opt) = match extract_function_info(tool_call) {
170 Some(info) => info,
171 None => {
172 memory.add_step(AgentStep::Observation {
173 tool_call_id: tool_call_id.to_string(),
174 result: "Tool call missing function".to_string(),
175 is_error: true,
176 });
177 continue;
178 }
179 };
180
181 let function_name = match function_name_opt {
182 Some(name) if !name.is_empty() => name,
183 _ => {
184 memory.add_step(AgentStep::Observation {
185 tool_call_id: tool_call_id.to_string(),
186 result: "Tool call missing function name".to_string(),
187 is_error: true,
188 });
189 continue;
190 }
191 };
192
193 let arguments_str = extract_arguments_str(&function);
194 let parsed_arguments =
195 parse_function_arguments(arguments_str, &function_name);
196
197 match parsed_arguments {
198 Ok(arguments_json) => {
199 if function_name == "final_answer" {
200 let steps = memory.steps().to_vec();
201 let mut sink = MemorySink {
202 memory: &mut memory,
203 };
204 let ctx = FinalAnswerStepsContext {
205 base: FinalAnswerContext {
206 tool_call_id,
207 arguments_json,
208 completion_schema: self.completion_schema(),
209 has_final_answer: &mut has_final_answer,
210 final_answer_value: &mut final_answer_value,
211 },
212 steps: &steps,
213 token_usage: token_usage.clone(),
214 start_duration: start_time.elapsed(),
215 iteration,
216 };
217
218 match handle_final_answer_steps(ctx, &mut sink)? {
219 HandlerOutcome::Continue => continue,
220 HandlerOutcome::ReturnResult(result) => return Ok(result),
221 HandlerOutcome::ReturnAnswer(_) => unreachable!(),
222 }
223 }
224
225 if function_name == structured_response_tool_name() {
226 let schema = match self.completion_schema() {
227 Some(schema) => schema.clone(),
228 None => {
229 let payload = AgentError::InvalidFunctionCall(
230 "No completion schema is active for structured response".to_string(),
231 )
232 .to_error_payload();
233 memory.add_step(AgentStep::Observation {
234 tool_call_id: tool_call_id.to_string(),
235 result: payload.to_string(),
236 is_error: true,
237 });
238 continue;
239 }
240 };
241
242 let steps = memory.steps().to_vec();
243 let mut sink = MemorySink {
244 memory: &mut memory,
245 };
246 let ctx = StructuredResponseStepsContext {
247 base: StructuredResponseContext {
248 tool_call_id,
249 arguments_json,
250 schema: &schema,
251 final_answer_value: final_answer_value.clone(),
252 },
253 steps: &steps,
254 token_usage: token_usage.clone(),
255 start_duration: start_time.elapsed(),
256 iteration,
257 };
258
259 match handle_structured_response_steps(ctx, &mut sink)? {
260 HandlerOutcome::Continue => continue,
261 HandlerOutcome::ReturnResult(result) => return Ok(result),
262 HandlerOutcome::ReturnAnswer(_) => unreachable!(),
263 }
264 }
265
266 memory.add_step(AgentStep::Action {
268 tool_name: function_name.to_string(),
269 tool_call_id: tool_call_id.to_string(),
270 arguments: arguments_json.clone(),
271 });
272
273 match self
274 .function_factory()
275 .execute_function(&function_name, arguments_json)
276 .await
277 {
278 Ok(result) => {
279 memory.add_step(AgentStep::Observation {
280 tool_call_id: tool_call_id.to_string(),
281 result: result.to_string(),
282 is_error: false,
283 });
284 }
285 Err(e) => {
286 let error_payload = e.to_error_payload();
287 memory.add_step(AgentStep::Observation {
288 tool_call_id: tool_call_id.to_string(),
289 result: error_payload.to_string(),
290 is_error: true,
291 });
292 }
293 };
294 }
295 Err(error) => {
296 memory.add_step(AgentStep::Observation {
297 tool_call_id: tool_call_id.to_string(),
298 result: error.to_error_payload().to_string(),
299 is_error: true,
300 });
301 }
302 }
303 }
304 }
305 } else {
306 let answer = assistant_message
307 .get("content")
308 .and_then(|value| value.as_str())
309 .unwrap_or("")
310 .trim()
311 .to_string();
312
313 let message = if !has_final_answer {
314 if answer.is_empty() {
315 "Assistant must call the `final_answer` tool to conclude the task, but returned no content.".to_string()
316 } else {
317 format!(
318 "Assistant must call the `final_answer` tool to conclude the task. Received plain response: {}",
319 answer
320 )
321 }
322 } else if answer.is_empty() {
323 format!(
324 "Assistant must call the `{}` tool with the structured schema payload, but returned no content.",
325 structured_response_tool_name()
326 )
327 } else {
328 format!(
329 "Assistant must call the `{}` tool with the structured schema payload instead of responding directly: {}",
330 structured_response_tool_name(),
331 answer
332 )
333 };
334
335 let tool_name = if self.completion_schema().is_some() {
336 structured_response_tool_name()
337 } else {
338 "final_answer"
339 };
340
341 memory.add_step(AgentStep::Observation {
342 tool_call_id: tool_name.to_string(),
343 result: message,
344 is_error: true,
345 });
346
347 continue;
348 }
349 }
350
351 Err(AgentError::MaxIterations(self.max_iterations()))
352 }
353
354 pub async fn run_with_messages(&self, mut messages: Vec<Value>) -> Result<String> {
355 let mut iteration = 0;
356 let mut has_final_answer = false;
357 let mut final_answer_value: Option<String> = None;
358
359 while iteration < self.max_iterations() {
360 iteration += 1;
361
362 if let Some(schema) = self.completion_schema() {
363 inject_schema_instructions(&mut messages, schema);
364 }
365
366 let mut tools = self.function_factory().get_openai_tools();
367 if let Some(schema) = self.completion_schema() {
368 tools.push(structured_response_tool_definition(schema));
369 } else {
370 tools.push(final_answer_tool_definition());
371 }
372
373 let mut chat_request =
374 ChatCompletionRequest::new(self.model().to_owned(), messages.clone())
375 .with_max_tokens(self.max_tokens());
376
377 if !tools.is_empty() {
378 chat_request = chat_request
379 .with_tools(tools)
380 .with_tool_choice(json!("auto"));
381 }
382
383 let request_body = chat_request.into_value();
384
385 let response = timeout(self.timeout(), self.make_raw_request(&request_body))
386 .await
387 .map_err(|_| AgentError::Timeout("OpenAI API call timed out".to_string()))??;
388
389 let choices = response
390 .get("choices")
391 .and_then(|value| value.as_array())
392 .ok_or_else(|| {
393 AgentError::Unknown(
394 "Missing 'choices' array in completion response".to_string(),
395 )
396 })?;
397
398 let first_choice = choices.first().ok_or_else(|| {
399 AgentError::Unknown("Completion response contained no choices".to_string())
400 })?;
401
402 let assistant_message = first_choice.get("message").cloned().ok_or_else(|| {
403 AgentError::Unknown("Completion response missing assistant message".to_string())
404 })?;
405
406 if let Some(tool_calls) = assistant_message.get("tool_calls") {
407 if let Some(tool_calls_array) = tool_calls.as_array() {
408 messages.push(json!({
409 "role": "assistant",
410 "content": assistant_message.get("content").unwrap_or(&json!("")),
411 "tool_calls": tool_calls
412 }));
413
414 let turn_has_final_answer = tool_calls_array.iter().any(|tool_call| {
415 tool_call
416 .get("function")
417 .and_then(|f| f.get("name"))
418 .and_then(|name| name.as_str())
419 .map(|name| name == "final_answer")
420 .unwrap_or(false)
421 });
422
423 if turn_has_final_answer && tool_calls_array.len() > 1 {
424 messages.push(json!({
425 "role": "tool",
426 "tool_call_id": "final_answer",
427 "content": AgentError::InvalidFunctionCall(
428 "`final_answer` must be the only tool call in a single turn".to_string()
429 )
430 .to_error_payload()
431 .to_string()
432 }));
433 continue;
434 }
435
436 for tool_call in tool_calls_array {
437 let tool_call_id = extract_tool_call_id(tool_call);
438
439 let (function, function_name_opt) = match extract_function_info(tool_call) {
440 Some(info) => info,
441 None => {
442 messages.push(json!({
443 "role": "tool",
444 "tool_call_id": tool_call_id,
445 "content": AgentError::InvalidFunctionCall(
446 "Tool call missing function".to_string()
447 )
448 .to_error_payload()
449 .to_string()
450 }));
451 continue;
452 }
453 };
454
455 let function_name = match function_name_opt {
456 Some(name) if !name.is_empty() => name,
457 _ => {
458 messages.push(json!({
459 "role": "tool",
460 "tool_call_id": tool_call_id,
461 "content": AgentError::InvalidFunctionCall(
462 "Tool call missing function name".to_string()
463 )
464 .to_error_payload()
465 .to_string()
466 }));
467 continue;
468 }
469 };
470
471 let arguments_str = extract_arguments_str(&function);
472 let parsed_arguments =
473 parse_function_arguments(arguments_str, &function_name);
474
475 if function_name == "final_answer" {
476 let arguments_json = match parsed_arguments {
477 Ok(val) => val,
478 Err(err) => {
479 messages.push(json!({
480 "role": "tool",
481 "tool_call_id": tool_call_id,
482 "content": err.to_error_payload().to_string()
483 }));
484 continue;
485 }
486 };
487
488 let mut sink = MessagesSink {
489 messages: &mut messages,
490 };
491 let ctx = FinalAnswerContext {
492 tool_call_id,
493 arguments_json,
494 completion_schema: self.completion_schema(),
495 has_final_answer: &mut has_final_answer,
496 final_answer_value: &mut final_answer_value,
497 };
498
499 match handle_final_answer_messages(ctx, &mut sink)? {
500 HandlerOutcome::Continue => continue,
501 HandlerOutcome::ReturnAnswer(answer) => return Ok(answer),
502 HandlerOutcome::ReturnResult(_) => unreachable!(),
503 }
504 }
505
506 if function_name == structured_response_tool_name() {
507 let schema = match self.completion_schema() {
508 Some(schema) => schema.clone(),
509 None => {
510 messages.push(json!({
511 "role": "tool",
512 "tool_call_id": tool_call_id,
513 "content": AgentError::InvalidFunctionCall(
514 "No completion schema is active for structured response".to_string()
515 )
516 .to_error_payload()
517 .to_string()
518 }));
519 continue;
520 }
521 };
522
523 let arguments_json = match parsed_arguments {
524 Ok(val) => val,
525 Err(err) => {
526 messages.push(json!({
527 "role": "tool",
528 "tool_call_id": tool_call_id,
529 "content": err.to_error_payload().to_string()
530 }));
531 continue;
532 }
533 };
534
535 let mut sink = MessagesSink {
536 messages: &mut messages,
537 };
538 let ctx = StructuredResponseContext {
539 tool_call_id,
540 arguments_json,
541 schema: &schema,
542 final_answer_value: final_answer_value.clone(),
543 };
544
545 match handle_structured_response_messages(ctx, &mut sink)? {
546 HandlerOutcome::Continue => continue,
547 HandlerOutcome::ReturnAnswer(answer) => return Ok(answer),
548 HandlerOutcome::ReturnResult(_) => unreachable!(),
549 }
550 }
551
552 let result = match parsed_arguments {
554 Ok(arguments_json) => match self
555 .function_factory()
556 .execute_function(&function_name, arguments_json)
557 .await
558 {
559 Ok(result) => result,
560 Err(e) => e.to_error_payload(),
561 },
562 Err(error) => error.to_error_payload(),
563 };
564
565 messages.push(json!({
566 "role": "tool",
567 "tool_call_id": tool_call_id,
568 "content": result.to_string()
569 }));
570 }
571 }
572 } else {
573 let answer = assistant_message
574 .get("content")
575 .and_then(|value| value.as_str())
576 .unwrap_or("")
577 .trim()
578 .to_string();
579
580 let content = if self.completion_schema().is_some() {
581 if answer.is_empty() {
582 format!(
583 "Reminder: Call the `{}` tool with the structured schema payload to complete the task.",
584 structured_response_tool_name()
585 )
586 } else {
587 format!(
588 "Reminder: Provide the structured schema by calling the `{}` tool instead of responding directly: {}",
589 structured_response_tool_name(),
590 answer
591 )
592 }
593 } else if answer.is_empty() {
594 "Reminder: You must call the `final_answer` tool with the completed answer to finish.".to_string()
595 } else {
596 format!(
597 "Reminder: Do not respond directly. Call the `final_answer` tool with the final answer instead of: {}",
598 answer
599 )
600 };
601
602 messages.push(json!({
603 "role": "system",
604 "content": content
605 }));
606
607 continue;
608 }
609 }
610
611 Err(AgentError::MaxIterations(self.max_iterations()))
612 }
613}