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}