Skip to main content

laminar_db/ai/
provider.rs

1//! [`InferenceProvider`] trait and request/response types.
2//!
3//! Providers are I/O only: inputs in, outputs + usage out. Task framing and
4//! response parsing belong to the adapter.
5
6use async_trait::async_trait;
7use thiserror::Error;
8
9use crate::ai::registry::Task;
10
11/// One homogeneous batch of inputs for a single task and model.
12#[derive(Debug, Clone, PartialEq)]
13pub struct InferenceRequest {
14    /// Task to perform.
15    pub task: Task,
16    /// Vendor model id for remote backends; weight source for local.
17    pub model: String,
18    /// One input string per row, in row order.
19    pub inputs: Vec<String>,
20    /// Task-shaping parameters; also versions the result cache.
21    pub params: InferenceParams,
22}
23
24/// Parameters that shape a request and version the result cache.
25#[derive(Debug, Clone, PartialEq, Default)]
26pub struct InferenceParams {
27    /// Candidate labels for classification. Required for remote classify;
28    /// must match or be a subset of a local classifier's intrinsic labels.
29    pub labels: Option<Vec<String>>,
30}
31
32/// Per-row outputs of a batch. Shape is homogeneous within a request.
33#[derive(Debug, Clone, PartialEq)]
34pub enum InferenceOutputs {
35    /// Text outputs (classify, complete, …).
36    Text(Vec<String>),
37    /// Numeric vectors (embeddings, or classifier logits before adapter post-processing).
38    Vectors(Vec<Vec<f32>>),
39    /// Scalar sentiment scores in `[-1, 1]`. Produced by the adapter; providers
40    /// never return this shape directly.
41    Scores(Vec<f64>),
42}
43
44impl InferenceOutputs {
45    /// Number of output rows.
46    #[must_use]
47    pub fn len(&self) -> usize {
48        match self {
49            InferenceOutputs::Text(v) => v.len(),
50            InferenceOutputs::Vectors(v) => v.len(),
51            InferenceOutputs::Scores(v) => v.len(),
52        }
53    }
54
55    /// Whether no rows are present.
56    #[must_use]
57    pub fn is_empty(&self) -> bool {
58        self.len() == 0
59    }
60}
61
62/// Token and cost accounting for a batch call. Local backends report [`Usage::ZERO`].
63#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
64pub struct Usage {
65    /// Prompt/input tokens billed.
66    pub input_tokens: u64,
67    /// Completion/output tokens billed.
68    pub output_tokens: u64,
69    /// Micro-USD (millionths); 0 for local.
70    pub cost_micros: u64,
71}
72
73impl Usage {
74    /// Zero usage for local backends.
75    pub const ZERO: Usage = Usage {
76        input_tokens: 0,
77        output_tokens: 0,
78        cost_micros: 0,
79    };
80}
81
82/// Result of a batch inference call.
83#[derive(Debug, Clone, PartialEq)]
84pub struct InferenceResponse {
85    /// Per-row outputs, in input order.
86    pub outputs: InferenceOutputs,
87    /// Token/cost accounting.
88    pub usage: Usage,
89}
90
91/// I/O transport over a model backend. Shared as `Arc<dyn InferenceProvider>`;
92/// driven from Ring 1, never Ring 0.
93#[async_trait]
94pub trait InferenceProvider: Send + Sync {
95    /// Run one batch of inputs through the model.
96    ///
97    /// # Errors
98    ///
99    /// Returns [`ProviderError`] on transport failure, timeout, rate limiting,
100    /// a malformed response, or an unsupported task.
101    async fn infer_batch(
102        &self,
103        request: InferenceRequest,
104    ) -> Result<InferenceResponse, ProviderError>;
105
106    /// Backend name for logging (e.g. `anthropic`, `openai`, `local`).
107    fn name(&self) -> &'static str;
108
109    /// Labels intrinsic to a local classifier (`config.json` `id2label`).
110    /// Returns `None` for remote providers and embedding models. Lets a lazily
111    /// downloaded classifier return labels without them being known at startup.
112    fn intrinsic_labels(&self, _model: &str) -> Option<Vec<String>> {
113        None
114    }
115}
116
117/// Errors a provider can return for a batch call.
118#[derive(Debug, Error)]
119pub enum ProviderError {
120    /// Network or connection failure.
121    #[error("transport error: {0}")]
122    Transport(String),
123
124    /// The call exceeded its deadline.
125    #[error("request timed out after {0} ms")]
126    Timeout(u64),
127
128    /// The provider signalled rate limiting.
129    #[error("rate limited by provider")]
130    RateLimited,
131
132    /// The response could not be parsed into the expected shape.
133    #[error("malformed response: {0}")]
134    BadResponse(String),
135
136    /// The provider cannot perform the requested task.
137    #[error("provider does not support task '{0}'")]
138    UnsupportedTask(Task),
139}