| 479 | |
| 480 | impl ApiClient for ScriptedApiClient { |
| 481 | fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> { |
| 482 | self.call_count += 1; |
| 483 | match self.call_count { |
| 484 | 1 => { |
| 485 | assert!(request |
| 486 | .messages |
| 487 | .iter() |
| 488 | .any(|message| message.role == MessageRole::User)); |
| 489 | Ok(vec![ |
| 490 | AssistantEvent::TextDelta("Let me calculate that.".to_string()), |
| 491 | AssistantEvent::ToolUse { |
| 492 | id: "tool-1".to_string(), |
| 493 | name: "add".to_string(), |
| 494 | input: "2,2".to_string(), |
| 495 | }, |
| 496 | AssistantEvent::Usage(TokenUsage { |
| 497 | input_tokens: 20, |
| 498 | output_tokens: 6, |
| 499 | cache_creation_input_tokens: 1, |
| 500 | cache_read_input_tokens: 2, |
| 501 | }), |
| 502 | AssistantEvent::MessageStop, |
| 503 | ]) |
| 504 | } |
| 505 | 2 => { |
| 506 | let last_message = request |
| 507 | .messages |
| 508 | .last() |
| 509 | .expect("tool result should be present"); |
| 510 | assert_eq!(last_message.role, MessageRole::Tool); |
| 511 | Ok(vec![ |
| 512 | AssistantEvent::TextDelta("The answer is 4.".to_string()), |
| 513 | AssistantEvent::Usage(TokenUsage { |
| 514 | input_tokens: 24, |
| 515 | output_tokens: 4, |
| 516 | cache_creation_input_tokens: 1, |
| 517 | cache_read_input_tokens: 3, |
| 518 | }), |
| 519 | AssistantEvent::MessageStop, |
| 520 | ]) |
| 521 | } |
| 522 | _ => Err(RuntimeError::new("unexpected extra API call")), |
| 523 | } |
| 524 | } |
| 525 | } |
| 526 | |
| 527 | struct PromptAllowOnce; |