ferrin_core/generate_text/
prepare_step.rs

1//! Per-step overrides.
2
3use ferrin_message::Message;
4use ferrin_spec::BoxFuture;
5use ferrin_spec::JsonValue;
6use ferrin_spec::LanguageModelRef;
7use ferrin_spec::ToolChoice;
8use ferrin_spec::ToolName;
9
10use super::StepResult;
11use crate::error::Error;
12use crate::prompt::CallSettings;
13use crate::prompt::Instructions;
14use crate::telemetry::ModelIdentity;
15
16/// Information available to a [`PrepareStep`] callback.
17#[derive(Debug)]
18pub struct PrepareStepContext<'a> {
19    /// Steps completed so far.
20    pub steps: &'a [StepResult],
21    /// Zero-based index of the step about to run.
22    pub step_number: u32,
23    /// The model configured for the call.
24    pub model: &'a ModelIdentity,
25    /// Instructions as configured on the call.
26    pub instructions: Option<&'a Instructions>,
27    /// Messages that will be sent (initial messages plus response messages).
28    pub messages: &'a [Message],
29    /// The initial messages of the call.
30    pub initial_messages: &'a [Message],
31    /// Response messages accumulated so far.
32    pub response_messages: &'a [Message],
33    /// The tools context of the call.
34    pub tools_context: Option<&'a JsonValue>,
35}
36
37/// Overrides returned by a [`PrepareStep`] callback. Unset fields keep the
38/// call-level configuration.
39#[derive(Debug, Default)]
40pub struct StepOverrides {
41    /// Model for this step.
42    pub model: Option<LanguageModelRef>,
43    /// Tool choice for this step.
44    pub tool_choice: Option<ToolChoice>,
45    /// Active tools for this step.
46    pub active_tools: Option<Vec<ToolName>>,
47    /// Tool order for this step.
48    pub tool_order: Option<Vec<ToolName>>,
49    /// Instructions for this step.
50    pub instructions: Option<Instructions>,
51    /// Messages for this step (replaces initial plus response messages).
52    pub messages: Option<Vec<Message>>,
53    /// Tools context for this step.
54    pub tools_context: Option<JsonValue>,
55    /// Sampling settings overlaid on the call settings.
56    pub settings: Option<CallSettings>,
57}
58
59impl StepOverrides {
60    /// No overrides.
61    #[must_use]
62    pub fn none() -> Self {
63        Self::default()
64    }
65
66    /// Overrides the model.
67    #[must_use]
68    pub fn with_model(mut self, model: impl Into<LanguageModelRef>) -> Self {
69        self.model = Some(model.into());
70        self
71    }
72
73    /// Overrides the tool choice.
74    #[must_use]
75    pub fn with_tool_choice(mut self, tool_choice: ToolChoice) -> Self {
76        self.tool_choice = Some(tool_choice);
77        self
78    }
79
80    /// Overrides the active tools.
81    #[must_use]
82    pub fn with_active_tools(
83        mut self,
84        names: impl IntoIterator<Item = impl Into<ToolName>>,
85    ) -> Self {
86        self.active_tools = Some(names.into_iter().map(Into::into).collect());
87        self
88    }
89
90    /// Overrides the tool order.
91    #[must_use]
92    pub fn with_tool_order(mut self, names: impl IntoIterator<Item = impl Into<ToolName>>) -> Self {
93        self.tool_order = Some(names.into_iter().map(Into::into).collect());
94        self
95    }
96
97    /// Overrides the instructions.
98    #[must_use]
99    pub fn with_instructions(mut self, instructions: impl Into<Instructions>) -> Self {
100        self.instructions = Some(instructions.into());
101        self
102    }
103
104    /// Overrides the messages.
105    #[must_use]
106    pub fn with_messages(mut self, messages: impl IntoIterator<Item = Message>) -> Self {
107        self.messages = Some(messages.into_iter().collect());
108        self
109    }
110
111    /// Overrides the tools context.
112    #[must_use]
113    pub fn with_tools_context(mut self, context: JsonValue) -> Self {
114        self.tools_context = Some(context);
115        self
116    }
117
118    /// Overlays sampling settings.
119    #[must_use]
120    pub fn with_settings(mut self, settings: CallSettings) -> Self {
121        self.settings = Some(settings);
122        self
123    }
124}
125
126/// Computes overrides before each step.
127///
128/// Implemented for every `Fn(&PrepareStepContext<'_>) -> StepOverrides`
129/// closure; implement the trait directly when the decision is asynchronous.
130pub trait PrepareStep: Send + Sync {
131    /// Returns the overrides for the step described by `ctx`.
132    fn prepare_step<'a>(
133        &'a self,
134        ctx: PrepareStepContext<'a>,
135    ) -> BoxFuture<'a, Result<StepOverrides, Error>>;
136}
137
138impl<F> PrepareStep for F
139where
140    F: Fn(&PrepareStepContext<'_>) -> StepOverrides + Send + Sync,
141{
142    fn prepare_step<'a>(
143        &'a self,
144        ctx: PrepareStepContext<'a>,
145    ) -> BoxFuture<'a, Result<StepOverrides, Error>> {
146        let overrides = self(&ctx);
147        Box::pin(async move { Ok(overrides) })
148    }
149}
ferrin_core/generate_text/prepare_step.rs

ferrin_core/generate_text/
prepare_step.rs
