ferrin_core/
modality_hooks.rs

1//! Operation hooks for embedding and reranking.
2//!
3//! Event contracts are derived from the Vercel AI SDK (Apache-2.0,
4//! Copyright 2023 Vercel, Inc.), translated to Rust and modified; see NOTICE.
5
6use std::fmt;
7
8use ferrin_spec::Headers;
9use ferrin_spec::JsonValue;
10use ferrin_spec::ProviderMetadata;
11use ferrin_spec::ProviderOptions;
12use ferrin_spec::ResponseMetadata;
13use ferrin_spec::Warning;
14use serde_json::json;
15
16use crate::embed::Embedding;
17use crate::embed::EmbeddingUsage;
18use crate::hooks::HookList;
19use crate::rerank::Ranked;
20use crate::rerank::RerankDocument;
21use crate::telemetry::ModelIdentity;
22
23/// Input shape of an embedding operation.
24#[derive(Debug, Clone, PartialEq)]
25#[non_exhaustive]
26pub enum EmbeddingInput {
27    /// One value supplied to `embed`.
28    Single(String),
29    /// Values supplied to `embed_many`, in input order.
30    Many(Vec<String>),
31}
32
33/// Output shape of an embedding operation.
34#[derive(Debug, Clone, PartialEq)]
35#[non_exhaustive]
36pub enum EmbeddingOutput {
37    /// One vector returned by `embed`.
38    Single(Embedding),
39    /// Vectors returned by `embed_many`, in input order.
40    Many(Vec<Embedding>),
41}
42
43/// Response metadata shape of an embedding operation.
44#[derive(Debug, Clone, PartialEq)]
45#[non_exhaustive]
46pub enum EmbeddingResponse {
47    /// The response to `embed`.
48    Single(Box<ResponseMetadata>),
49    /// Responses to `embed_many`, in input chunk order.
50    Many(Vec<ResponseMetadata>),
51}
52
53/// Event before any provider attempt of an embedding operation.
54#[derive(Debug, Clone, PartialEq)]
55pub struct EmbedCallStartEvent {
56    /// Application context, absent only in restricted telemetry copies.
57    pub runtime_context: Option<JsonValue>,
58    /// Identifier shared by the operation and its model attempts.
59    pub call_id: String,
60    /// `ai.embed` or `ai.embedMany`.
61    pub operation_id: &'static str,
62    /// Provider and model identity.
63    pub model: ModelIdentity,
64    /// Original input; omitted when telemetry does not record inputs.
65    pub value: Option<EmbeddingInput>,
66    /// Maximum retries for each provider call.
67    pub max_retries: u32,
68    /// Request headers, including the core User-Agent suffix.
69    pub headers: Headers,
70    /// Provider-specific options.
71    pub provider_options: ProviderOptions,
72}
73
74/// Event after all chunks and retries of an embedding operation succeed.
75#[derive(Debug, Clone, PartialEq)]
76pub struct EmbedCallEndEvent {
77    /// Application context, absent only in restricted telemetry copies.
78    pub runtime_context: Option<JsonValue>,
79    /// Identifier shared with the start event.
80    pub call_id: String,
81    /// `ai.embed` or `ai.embedMany`.
82    pub operation_id: &'static str,
83    /// Provider and model identity.
84    pub model: ModelIdentity,
85    /// Original input; omitted when telemetry does not record inputs.
86    pub value: Option<EmbeddingInput>,
87    /// Generated vectors; omitted when telemetry does not record outputs.
88    pub embedding: Option<EmbeddingOutput>,
89    /// Aggregated token usage.
90    pub usage: EmbeddingUsage,
91    /// Warnings in input chunk order.
92    pub warnings: Vec<Warning>,
93    /// Provider metadata merged in input chunk order.
94    pub provider_metadata: Option<ProviderMetadata>,
95    /// Provider response metadata.
96    pub response: EmbeddingResponse,
97}
98
99/// Event before any provider attempt of a reranking operation.
100#[derive(Debug, Clone, PartialEq)]
101pub struct RerankCallStartEvent {
102    /// Application context, absent only in restricted telemetry copies.
103    pub runtime_context: Option<JsonValue>,
104    /// Identifier shared by the operation and its model attempts.
105    pub call_id: String,
106    /// `ai.rerank`.
107    pub operation_id: &'static str,
108    /// Provider and model identity.
109    pub model: ModelIdentity,
110    /// Original documents; omitted when telemetry does not record inputs.
111    pub documents: Option<Vec<RerankDocument>>,
112    /// Original query; omitted when telemetry does not record inputs.
113    pub query: Option<String>,
114    /// Requested maximum number of ranked documents.
115    pub top_n: Option<usize>,
116    /// Maximum retries for the provider call.
117    pub max_retries: u32,
118    /// Caller-provided request headers.
119    pub headers: Headers,
120    /// Provider-specific options.
121    pub provider_options: ProviderOptions,
122}
123
124/// Event after a successful reranking operation, including an empty input.
125#[derive(Debug, Clone, PartialEq)]
126pub struct RerankCallEndEvent {
127    /// Application context, absent only in restricted telemetry copies.
128    pub runtime_context: Option<JsonValue>,
129    /// Identifier shared with the start event.
130    pub call_id: String,
131    /// `ai.rerank`.
132    pub operation_id: &'static str,
133    /// Provider and model identity.
134    pub model: ModelIdentity,
135    /// Original documents; omitted when telemetry does not record inputs.
136    pub documents: Option<Vec<RerankDocument>>,
137    /// Original query; omitted when telemetry does not record inputs.
138    pub query: Option<String>,
139    /// Ranked documents; omitted when telemetry does not record outputs.
140    pub ranking: Option<Vec<Ranked<RerankDocument>>>,
141    /// Adapter warnings.
142    pub warnings: Vec<Warning>,
143    /// Provider-specific metadata.
144    pub provider_metadata: Option<ProviderMetadata>,
145    /// Response metadata, including resolved timestamp and model ID.
146    pub response: ResponseMetadata,
147}
148
149pub(crate) struct ModalityHooks<S, E> {
150    pub(crate) runtime_context: JsonValue,
151    pub(crate) on_start: HookList<S>,
152    pub(crate) on_end: HookList<E>,
153}
154
155impl<S, E> Default for ModalityHooks<S, E> {
156    fn default() -> Self {
157        Self {
158            runtime_context: json!({}),
159            on_start: Vec::new(),
160            on_end: Vec::new(),
161        }
162    }
163}
164
165impl<S, E> fmt::Debug for ModalityHooks<S, E> {
166    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
167        f.debug_struct("ModalityHooks")
168            .field("on_start", &self.on_start.len())
169            .field("on_end", &self.on_end.len())
170            .finish_non_exhaustive()
171    }
172}
173
174macro_rules! impl_modality_hooks {
175    ($ty:ident $(<$generic:ident>)?, $start:ty, $end:ty) => {
176        impl$(<$generic>)? $ty$(<$generic>)? {
177            /// Sets the application context passed to operation callbacks.
178            #[must_use]
179            pub fn runtime_context(mut self, context: ::ferrin_spec::JsonValue) -> Self {
180                self.hooks.runtime_context = context;
181                self
182            }
183
184            /// Adds an awaited callback before the first provider attempt.
185            #[must_use]
186            pub fn on_start(mut self, hook: impl $crate::hooks::HookFn<$start>) -> Self {
187                self.hooks.on_start.push(::std::sync::Arc::new(hook));
188                self
189            }
190
191            /// Adds an awaited callback after the entire operation succeeds.
192            #[must_use]
193            pub fn on_end(mut self, hook: impl $crate::hooks::HookFn<$end>) -> Self {
194                self.hooks.on_end.push(::std::sync::Arc::new(hook));
195                self
196            }
197        }
198    };
199}
200pub(crate) use impl_modality_hooks;
ferrin_core/modality_hooks.rs

ferrin_core/
modality_hooks.rs
