Skip to main content

ferrin_core/generate_text/
prepare_step.rs

1//! Per-step overrides.
2
3use std::fmt;
4use std::sync::Arc;
5
6use ferrin_message::Message;
7use ferrin_spec::BoxFuture;
8use ferrin_spec::DynLanguageModel;
9use ferrin_spec::JsonValue;
10use ferrin_spec::LanguageModelRef;
11use ferrin_spec::ToolChoice;
12use ferrin_spec::ToolName;
13
14use super::StepResult;
15use crate::error::Error;
16use crate::prompt::CallSettings;
17use crate::prompt::Instructions;
18
19/// Information available to a [`PrepareStep`] callback.
20pub struct PrepareStepContext<'a> {
21    /// Steps completed so far.
22    pub steps: &'a [StepResult],
23    /// Zero-based index of the step about to run.
24    pub step_number: u32,
25    /// The model configured for the call.
26    pub model: &'a Arc<dyn DynLanguageModel>,
27    /// Instructions retained from the preceding step, initially configured on the call.
28    pub instructions: Option<&'a Instructions>,
29    /// Instructions originally configured for this invocation.
30    pub initial_instructions: Option<&'a Instructions>,
31    /// Current messages, including new responses since the latest message override.
32    pub messages: &'a [Message],
33    /// The initial messages of the call.
34    pub initial_messages: &'a [Message],
35    /// Response messages accumulated so far.
36    pub response_messages: &'a [Message],
37    /// The tools context retained from the preceding step.
38    pub tools_context: Option<&'a JsonValue>,
39    /// Application state retained from the preceding step.
40    pub runtime_context: Option<&'a JsonValue>,
41    /// The sandbox configured for the invocation before this step's override.
42    #[cfg(feature = "sandbox")]
43    pub sandbox: Option<&'a Arc<dyn ferrin_tool::Sandbox>>,
44}
45
46impl fmt::Debug for PrepareStepContext<'_> {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        f.debug_struct("PrepareStepContext")
49            .field("step_number", &self.step_number)
50            .field("provider", self.model.provider())
51            .field("model_id", self.model.model_id())
52            .field("steps", &self.steps.len())
53            .field("messages", &self.messages.len())
54            .finish_non_exhaustive()
55    }
56}
57
58/// Overrides returned by a [`PrepareStep`] callback. Unset fields keep the
59/// prior state for messages, instructions and contexts; other fields keep call defaults.
60#[derive(Default)]
61pub struct StepOverrides {
62    /// Model for this step.
63    pub model: Option<LanguageModelRef>,
64    /// Tool choice for this step.
65    pub tool_choice: Option<ToolChoice>,
66    /// Active tools for this step.
67    pub active_tools: Option<Vec<ToolName>>,
68    /// Tool order for this step.
69    pub tool_order: Option<Vec<ToolName>>,
70    /// Instructions for this and subsequent steps.
71    pub instructions: Option<Instructions>,
72    /// Messages for this and subsequent steps (new responses are appended).
73    pub messages: Option<Vec<Message>>,
74    /// Tools context for this and subsequent steps.
75    pub tools_context: Option<JsonValue>,
76    /// Application state for this and subsequent steps (`None` preserves the prior value).
77    pub runtime_context: Option<JsonValue>,
78    /// Sandbox for this step only; an unset value uses the invocation sandbox.
79    #[cfg(feature = "sandbox")]
80    pub sandbox: Option<Arc<dyn ferrin_tool::Sandbox>>,
81    /// Sampling settings overlaid on the call settings.
82    pub settings: Option<CallSettings>,
83}
84
85impl fmt::Debug for StepOverrides {
86    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
87        f.debug_struct("StepOverrides")
88            .field("model", &self.model)
89            .field("tool_choice", &self.tool_choice)
90            .field("active_tools", &self.active_tools)
91            .field("tool_order", &self.tool_order)
92            .field("instructions", &self.instructions)
93            .field("messages", &self.messages)
94            .field("has_tools_context", &self.tools_context.is_some())
95            .field("has_runtime_context", &self.runtime_context.is_some())
96            .field("settings", &self.settings)
97            .finish_non_exhaustive()
98    }
99}
100
101impl StepOverrides {
102    /// No overrides.
103    #[must_use]
104    pub fn none() -> Self {
105        Self::default()
106    }
107
108    /// Overrides the model.
109    #[must_use]
110    pub fn with_model(mut self, model: impl Into<LanguageModelRef>) -> Self {
111        self.model = Some(model.into());
112        self
113    }
114
115    /// Overrides the tool choice.
116    #[must_use]
117    pub fn with_tool_choice(mut self, tool_choice: ToolChoice) -> Self {
118        self.tool_choice = Some(tool_choice);
119        self
120    }
121
122    /// Overrides the active tools.
123    #[must_use]
124    pub fn with_active_tools(
125        mut self,
126        names: impl IntoIterator<Item = impl Into<ToolName>>,
127    ) -> Self {
128        self.active_tools = Some(names.into_iter().map(Into::into).collect());
129        self
130    }
131
132    /// Overrides the tool order.
133    #[must_use]
134    pub fn with_tool_order(mut self, names: impl IntoIterator<Item = impl Into<ToolName>>) -> Self {
135        self.tool_order = Some(names.into_iter().map(Into::into).collect());
136        self
137    }
138
139    /// Overrides the instructions.
140    #[must_use]
141    pub fn with_instructions(mut self, instructions: impl Into<Instructions>) -> Self {
142        self.instructions = Some(instructions.into());
143        self
144    }
145
146    /// Overrides the messages.
147    #[must_use]
148    pub fn with_messages(mut self, messages: impl IntoIterator<Item = Message>) -> Self {
149        self.messages = Some(messages.into_iter().collect());
150        self
151    }
152
153    /// Overrides the tools context.
154    #[must_use]
155    pub fn with_tools_context(mut self, context: JsonValue) -> Self {
156        self.tools_context = Some(context);
157        self
158    }
159
160    /// Replaces application state for this and subsequent steps.
161    #[must_use]
162    pub fn with_runtime_context(mut self, context: JsonValue) -> Self {
163        self.runtime_context = Some(context);
164        self
165    }
166
167    /// Uses a sandbox for this step only.
168    #[cfg(feature = "sandbox")]
169    #[must_use]
170    pub fn with_sandbox(mut self, sandbox: Arc<dyn ferrin_tool::Sandbox>) -> Self {
171        self.sandbox = Some(sandbox);
172        self
173    }
174
175    /// Overlays sampling settings.
176    #[must_use]
177    pub fn with_settings(mut self, settings: CallSettings) -> Self {
178        self.settings = Some(settings);
179        self
180    }
181}
182
183/// Computes overrides before each step.
184///
185/// Implemented for every `Fn(&PrepareStepContext<'_>) -> StepOverrides`
186/// closure; implement the trait directly when the decision is asynchronous.
187pub trait PrepareStep: Send + Sync {
188    /// Returns the overrides for the step described by `ctx`.
189    fn prepare_step<'a>(
190        &'a self,
191        ctx: PrepareStepContext<'a>,
192    ) -> BoxFuture<'a, Result<StepOverrides, Error>>;
193}
194
195impl<F> PrepareStep for F
196where
197    F: Fn(&PrepareStepContext<'_>) -> StepOverrides + Send + Sync,
198{
199    fn prepare_step<'a>(
200        &'a self,
201        ctx: PrepareStepContext<'a>,
202    ) -> BoxFuture<'a, Result<StepOverrides, Error>> {
203        let overrides = self(&ctx);
204        Box::pin(async move { Ok(overrides) })
205    }
206}