zeph_memory/
embed_probe.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Shared embedding-dimension probe helper.
5//!
6//! Several ingest and memory subsystems need to learn an embedding provider's output
7//! dimension before creating a Qdrant collection: they call the provider once with a
8//! throwaway probe string and use the resulting vector's length as the collection's
9//! vector size. [`probe_vector_size`] centralizes that sequence (including an optional
10//! timeout) so callers stop re-implementing it inline.
11
12use std::future::Future;
13use std::time::Duration;
14
15/// Error surfaced by [`probe_vector_size`].
16#[derive(Debug, thiserror::Error)]
17pub enum ProbeError<E> {
18    /// The probe embedding call exceeded `timeout`.
19    #[error("embedding probe timed out after {0:?}")]
20    Timeout(Duration),
21    /// The probe embedding call itself returned an error.
22    #[error("{0}")]
23    Embed(E),
24}
25
26/// Learn an embedding provider's output dimension.
27///
28/// Awaits `embed_fut` — the caller's already-invoked embed call for a throwaway probe string —
29/// and returns the length of the resulting vector. When `timeout` is `Some`, the await is
30/// bounded by it; pass `None` to await unconditionally.
31///
32/// Taking the future itself (rather than a `Fn(&str) -> Fut` closure) sidesteps the higher-rank
33/// lifetime that a borrowing `embed(&self, text: &str)` call would otherwise require, and lets
34/// each caller keep its own probe string literal.
35///
36/// This function only performs the probe — pair the returned size with the caller's own
37/// `ensure_collection` (or `ensure_collection_with_quantization`) call, since collection setup
38/// differs across call sites (plain vs. quantized, fixed vs. caller-supplied collection name).
39///
40/// # Errors
41///
42/// Returns [`ProbeError::Timeout`] if `timeout` elapses, or [`ProbeError::Embed`] if the probe
43/// call fails.
44///
45/// # Examples
46///
47/// ```
48/// use std::time::Duration;
49/// use zeph_memory::embed_probe::probe_vector_size;
50///
51/// # async fn probe() -> Result<(), Box<dyn std::error::Error>> {
52/// let embed_fut = async { Ok::<_, std::convert::Infallible>(vec![0.0_f32; 384]) };
53/// let size = probe_vector_size(embed_fut, Some(Duration::from_secs(15))).await?;
54/// assert_eq!(size, 384);
55/// # Ok(())
56/// # }
57/// ```
58#[tracing::instrument(
59    name = "memory.embed_probe.probe_vector_size",
60    skip_all,
61    err,
62    level = "debug"
63)]
64pub async fn probe_vector_size<Fut, E>(
65    embed_fut: Fut,
66    timeout: Option<Duration>,
67) -> Result<u64, ProbeError<E>>
68where
69    Fut: Future<Output = Result<Vec<f32>, E>>,
70    E: std::fmt::Display,
71{
72    let vector = match timeout {
73        Some(d) => tokio::time::timeout(d, embed_fut)
74            .await
75            .map_err(|_| ProbeError::Timeout(d))?
76            .map_err(ProbeError::Embed)?,
77        None => embed_fut.await.map_err(ProbeError::Embed)?,
78    };
79    // Safe: a Vec<f32> with 4B+ elements is impossible in practice on any 64-bit platform.
80    Ok(vector.len() as u64)
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use std::convert::Infallible;
87
88    #[tokio::test]
89    async fn probe_vector_size_returns_length() {
90        let size = probe_vector_size(async { Ok::<_, Infallible>(vec![0.0_f32; 128]) }, None)
91            .await
92            .unwrap();
93        assert_eq!(size, 128);
94    }
95
96    #[tokio::test]
97    async fn probe_vector_size_propagates_embed_error() {
98        let result = probe_vector_size(async { Err::<Vec<f32>, _>("boom") }, None).await;
99        assert!(matches!(result, Err(ProbeError::Embed("boom"))));
100    }
101
102    #[tokio::test]
103    async fn probe_vector_size_times_out() {
104        let result = probe_vector_size(
105            async {
106                tokio::time::sleep(Duration::from_millis(50)).await;
107                Ok::<_, Infallible>(vec![0.0_f32; 4])
108            },
109            Some(Duration::from_millis(1)),
110        )
111        .await;
112        assert!(matches!(result, Err(ProbeError::Timeout(_))));
113    }
114
115    #[tokio::test]
116    async fn probe_vector_size_no_timeout_awaits_unconditionally() {
117        let size = probe_vector_size(
118            async {
119                tokio::time::sleep(Duration::from_millis(20)).await;
120                Ok::<_, Infallible>(vec![0.0_f32; 7])
121            },
122            None,
123        )
124        .await
125        .unwrap();
126        assert_eq!(size, 7);
127    }
128}
zeph_memory/embed_probe.rs

zeph_memory/
embed_probe.rs
