|
| 1 | +//! Mock model types with streaming support for running agents without a real LLM provider. |
| 2 | +//! |
| 3 | +//! Wraps upstream [`serdes_ai_models`] mock types so they work with |
| 4 | +//! [`Agent::run_stream`][`serdes_ai::Agent::run_stream`]. |
| 5 | +//! |
| 6 | +//! # Quick start |
| 7 | +//! |
| 8 | +//! ```text |
| 9 | +//! use reloaded_code_serdesai::mock::{Streamed, tool_then_text}; |
| 10 | +//! use serde_json::json; |
| 11 | +//! |
| 12 | +//! let model = tool_then_text("glob", json!({"pattern": "*.rs"}), "Done."); |
| 13 | +//! let stream = agent.run_stream("prompt", ()).await?; // OK |
| 14 | +//! ``` |
| 15 | +//! |
| 16 | +//! When using [`crate::AgentBuildContext`], call |
| 17 | +//! [`with_model_override`](crate::AgentBuildContext::with_model_override) |
| 18 | +//! to inject the mock model before calling [`build()`](crate::AgentBuildContext::build). |
| 19 | +
|
| 20 | +// Re-export upstream mock types so users can still access the raw variants when needed. |
| 21 | +pub use serdes_ai_models::{FunctionModel, MockModel, TestModel}; |
| 22 | + |
| 23 | +use async_trait::async_trait; |
| 24 | +use futures::stream; |
| 25 | +use serdes_ai::core::{ |
| 26 | + FinishReason, ModelRequest, ModelResponse, ModelResponsePart, ModelResponseStreamEvent, |
| 27 | +}; |
| 28 | +use serdes_ai_models::Model as ModelTrait; |
| 29 | +// Re-export the types from where serdes-ai-models exposes them. |
| 30 | +use serdes_ai::core::ModelSettings; |
| 31 | +use serdes_ai_models::{ |
| 32 | + ModelCapability, ModelError, ModelProfile, ModelRequestParameters, StreamedResponse, |
| 33 | +}; |
| 34 | + |
| 35 | +// ============================================================================ |
| 36 | +// Streamed - wrapper that adds streaming support to any Model |
| 37 | +// ============================================================================ |
| 38 | + |
| 39 | +/// Wrapper adding [`request_stream`](ModelTrait::request_stream) support to any [`ModelTrait`] implementation. |
| 40 | +/// |
| 41 | +/// Delegates [`request`](ModelTrait::request) directly to the inner model and converts the non-streaming |
| 42 | +/// response into a sequence of [`ModelResponseStreamEvent`]s for streaming callers. |
| 43 | +/// |
| 44 | +/// # Example |
| 45 | +/// |
| 46 | +/// ```rust,no_run |
| 47 | +/// use reloaded_code_serdesai::mock::{FunctionModel, Streamed}; |
| 48 | +/// use serde_json::json; |
| 49 | +/// |
| 50 | +/// let model = Streamed::new(FunctionModel::tool_call("glob", json!({"pattern": "*.rs"}))); |
| 51 | +/// ``` |
| 52 | +#[derive(Clone, Debug)] |
| 53 | +pub struct Streamed<T> { |
| 54 | + inner: T, |
| 55 | + name: String, |
| 56 | +} |
| 57 | + |
| 58 | +impl<T> Streamed<T> { |
| 59 | + /// Wrap a model to add streaming support. |
| 60 | + /// |
| 61 | + /// The `name` defaults to the inner model's [`name()`](ModelTrait::name). |
| 62 | + pub fn new(inner: T) -> Self |
| 63 | + where |
| 64 | + T: ModelTrait, |
| 65 | + { |
| 66 | + let name = inner.name().to_string(); |
| 67 | + Self { inner, name } |
| 68 | + } |
| 69 | + |
| 70 | + /// Set a custom name for the wrapped model. |
| 71 | + pub fn with_name(mut self, name: impl Into<String>) -> Self { |
| 72 | + self.name = name.into(); |
| 73 | + self |
| 74 | + } |
| 75 | +} |
| 76 | + |
| 77 | +#[async_trait] |
| 78 | +impl<T: ModelTrait + Send + Sync> ModelTrait for Streamed<T> { |
| 79 | + fn name(&self) -> &str { |
| 80 | + &self.name |
| 81 | + } |
| 82 | + |
| 83 | + fn system(&self) -> &str { |
| 84 | + self.inner.system() |
| 85 | + } |
| 86 | + |
| 87 | + fn profile(&self) -> &ModelProfile { |
| 88 | + self.inner.profile() |
| 89 | + } |
| 90 | + |
| 91 | + async fn request( |
| 92 | + &self, |
| 93 | + messages: &[ModelRequest], |
| 94 | + settings: &ModelSettings, |
| 95 | + params: &ModelRequestParameters, |
| 96 | + ) -> Result<ModelResponse, ModelError> { |
| 97 | + self.inner.request(messages, settings, params).await |
| 98 | + } |
| 99 | + |
| 100 | + async fn request_stream( |
| 101 | + &self, |
| 102 | + messages: &[ModelRequest], |
| 103 | + settings: &ModelSettings, |
| 104 | + params: &ModelRequestParameters, |
| 105 | + ) -> Result<StreamedResponse, ModelError> { |
| 106 | + let response = self.inner.request(messages, settings, params).await?; |
| 107 | + let events = response_to_stream_events(response); |
| 108 | + Ok(Box::pin(stream::iter(events.into_iter().map(Ok)))) |
| 109 | + } |
| 110 | + |
| 111 | + fn supports(&self, capability: ModelCapability) -> bool { |
| 112 | + self.inner.supports(capability) |
| 113 | + } |
| 114 | +} |
| 115 | + |
| 116 | +// ============================================================================ |
| 117 | +// Convenience helpers |
| 118 | +// ============================================================================ |
| 119 | + |
| 120 | +/// Build a mock model that calls `tool_name` with `args` on the **first** turn, |
| 121 | +/// then returns text that incorporates the real tool return on the **second** turn. |
| 122 | +/// |
| 123 | +/// This prevents infinite loops when running agent examples that stream, |
| 124 | +/// because after the tool result is fed back the model answers with text. |
| 125 | +/// |
| 126 | +/// The second-turn response includes whatever the real tool returned, so |
| 127 | +/// the output reflects actual tool execution rather than a canned message. |
| 128 | +/// |
| 129 | +/// # Example |
| 130 | +/// |
| 131 | +/// ```rust,no_run |
| 132 | +/// use reloaded_code_serdesai::mock::tool_then_text; |
| 133 | +/// use serde_json::json; |
| 134 | +/// |
| 135 | +/// let model = tool_then_text("glob", json!({"pattern": "*.rs"}), "Done."); |
| 136 | +/// ``` |
| 137 | +pub fn tool_then_text( |
| 138 | + tool_name: impl Into<String>, |
| 139 | + args: serde_json::Value, |
| 140 | + fallback_text: impl Into<String>, |
| 141 | +) -> Streamed<FunctionModel> { |
| 142 | + let tool_name = tool_name.into(); |
| 143 | + let fallback_text = fallback_text.into(); |
| 144 | + let tool_name_clone = tool_name.clone(); |
| 145 | + |
| 146 | + let model = FunctionModel::new(move |messages, _settings| { |
| 147 | + // Check whether the conversation already contains a tool return from a |
| 148 | + // previous turn. If it does, we are on the second call and should |
| 149 | + // produce a text response incorporating the real result. |
| 150 | + let has_tool_return = messages.iter().any(|m| { |
| 151 | + m.parts |
| 152 | + .iter() |
| 153 | + .any(|p| matches!(p, serdes_ai::core::ModelRequestPart::ToolReturn(_))) |
| 154 | + }); |
| 155 | + |
| 156 | + if has_tool_return { |
| 157 | + // Collect tool return content from the message history. |
| 158 | + let tool_results: String = messages |
| 159 | + .iter() |
| 160 | + .flat_map(|m| m.tool_returns()) |
| 161 | + .map(extract_tool_return_text) |
| 162 | + .collect::<Vec<_>>() |
| 163 | + .join("\n"); |
| 164 | + |
| 165 | + let text = if tool_results.is_empty() { |
| 166 | + fallback_text.clone() |
| 167 | + } else { |
| 168 | + format!("{fallback_text}\n\n{tool_results}") |
| 169 | + }; |
| 170 | + |
| 171 | + ModelResponse::text(text) |
| 172 | + } else { |
| 173 | + // First call: emit a tool call so the agent executes the real tool. |
| 174 | + ModelResponse::with_parts(vec![ |
| 175 | + ModelResponsePart::text(format!("Calling {tool_name}...")), |
| 176 | + ModelResponsePart::tool_call(tool_name_clone.clone(), args.clone()), |
| 177 | + ]) |
| 178 | + .with_finish_reason(FinishReason::ToolCall) |
| 179 | + } |
| 180 | + }); |
| 181 | + |
| 182 | + Streamed::new(model) |
| 183 | +} |
| 184 | + |
| 185 | +// ============================================================================ |
| 186 | +// Private helpers |
| 187 | +// ============================================================================ |
| 188 | + |
| 189 | +fn response_to_stream_events(response: ModelResponse) -> Vec<ModelResponseStreamEvent> { |
| 190 | + let mut events = Vec::with_capacity(response.parts.len() * 2 + 1); |
| 191 | + |
| 192 | + for (index, part) in response.parts.into_iter().enumerate() { |
| 193 | + events.push(ModelResponseStreamEvent::part_start(index, part)); |
| 194 | + events.push(ModelResponseStreamEvent::part_end(index)); |
| 195 | + } |
| 196 | + |
| 197 | + events |
| 198 | +} |
| 199 | + |
| 200 | +/// Extract human-readable text from a [`ToolReturnPart`]. |
| 201 | +/// |
| 202 | +/// Uses serde JSON round-tripping to avoid depending on the |
| 203 | +/// non-public `ToolReturnContent` enum variants directly. |
| 204 | +fn extract_tool_return_text(tr: &serdes_ai::core::ToolReturnPart) -> String { |
| 205 | + // Serialize the content field to JSON, then extract readable text. |
| 206 | + // ToolReturnContent variants produce: |
| 207 | + // Text -> {"type":"text","content":"..."} |
| 208 | + // Json -> {"type":"json","content":{...}} |
| 209 | + // Error -> {"type":"error","message":"..."} |
| 210 | + // Multiple-> {"type":"multiple","items":[...]} |
| 211 | + // Image -> {"type":"image","image":{...}} |
| 212 | + let Ok(val) = serde_json::to_value(&tr.content) else { |
| 213 | + return format!("{:?}", tr.content); |
| 214 | + }; |
| 215 | + |
| 216 | + // Try text content field first (most common case). |
| 217 | + if let Some(text) = val.get("content").and_then(|v| v.as_str()) { |
| 218 | + return text.to_string(); |
| 219 | + } |
| 220 | + |
| 221 | + // JSON content field. |
| 222 | + if let Some(json_val) = val.get("content") |
| 223 | + && let Ok(pretty) = serde_json::to_string_pretty(json_val) |
| 224 | + { |
| 225 | + return pretty; |
| 226 | + } |
| 227 | + |
| 228 | + // Error message field. |
| 229 | + if let Some(msg) = val.get("message").and_then(|v| v.as_str()) { |
| 230 | + return format!("[error] {msg}"); |
| 231 | + } |
| 232 | + |
| 233 | + // Fallback: pretty-print the whole thing. |
| 234 | + serde_json::to_string_pretty(&val).unwrap_or_else(|_| format!("{:?}", tr.content)) |
| 235 | +} |
0 commit comments