Skip to main content

semantic_operators/
each.rs

1//! Bounded independent calls, preserving input order and per-item failures.
2use crate::{
3    Answer, AnswerValue, Error, ErrorKind, Provider, Question, Request, Result, State, ask,
4};
5use futures_util::{StreamExt, TryStreamExt, stream};
6use indexmap::IndexMap;
7use serde_json::Value;
8use std::{collections::HashSet, num::NonZeroUsize};
9
10/// Options shared by per-item filtering and reranking.
11#[derive(Clone)]
12pub struct BatchOptions {
13    /// Maximum concurrently polled provider calls; defaults to eight.
14    pub concurrency: NonZeroUsize,
15    /// Shared background included in each item's independent request.
16    pub context: Option<Value>,
17    /// Explicit images available to each independent request.
18    #[cfg(feature = "images")]
19    #[cfg_attr(docsrs, doc(cfg(feature = "images")))]
20    pub assets: crate::image::Assets,
21}
22impl Default for BatchOptions {
23    fn default() -> Self {
24        Self {
25            concurrency: NonZeroUsize::new(8).expect("eight is positive"),
26            context: None,
27            #[cfg(feature = "images")]
28            assets: crate::image::Assets::default(),
29        }
30    }
31}
32impl std::fmt::Debug for BatchOptions {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        f.debug_struct("BatchOptions")
35            .field("concurrency", &self.concurrency)
36            .field("has_context", &self.context.is_some())
37            .finish_non_exhaustive()
38    }
39}
40/// One independently scored input. Identifiers are local metadata, never added to state.
41#[derive(Clone)]
42pub struct Item {
43    id: String,
44    state: State,
45}
46impl Item {
47    /// Validate an identifier and retain a previously validated state.
48    ///
49    /// # Errors
50    /// Returns `InvalidArgument` if `id` is empty. Batch functions separately
51    /// check identifiers for uniqueness.
52    pub fn new(id: impl Into<String>, state: State) -> Result<Self> {
53        let id = id.into();
54        validate_ids([id.as_str()])?;
55        Ok(Self { id, state })
56    }
57    /// Local identifier.
58    pub fn id(&self) -> &str {
59        &self.id
60    }
61    /// State sent to the provider.
62    pub fn state(&self) -> &State {
63        &self.state
64    }
65}
66impl std::fmt::Debug for Item {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        f.debug_struct("Item").finish_non_exhaustive()
69    }
70}
71/// Evidence or a provider failure for one input; an error is never an abstention.
72#[derive(Clone, Debug)]
73pub enum ItemResult {
74    /// A validated answer, possibly undecided.
75    Answer(Answer<AnswerValue>),
76    /// A provider, timeout, or malformed-response failure.
77    Error(Error),
78}
79impl ItemResult {
80    /// Answer evidence, if this item completed successfully.
81    pub fn answer(&self) -> Option<&Answer<AnswerValue>> {
82        match self {
83            Self::Answer(a) => Some(a),
84            Self::Error(_) => None,
85        }
86    }
87    /// Failure evidence, if this item failed.
88    pub fn error(&self) -> Option<&Error> {
89        match self {
90            Self::Answer(_) => None,
91            Self::Error(e) => Some(e),
92        }
93    }
94}
95
96pub(crate) fn validate_ids<'a>(ids: impl IntoIterator<Item = &'a str>) -> Result<()> {
97    let mut seen = HashSet::new();
98    for id in ids {
99        if id.is_empty() || !seen.insert(id) {
100            return Err(Error::new(
101                ErrorKind::InvalidArgument,
102                "item identifiers must be nonempty and unique",
103            ));
104        }
105    }
106    Ok(())
107}
108/// Ask one question independently for every state with bounded concurrency.
109///
110/// Output order matches input order. Provider, timeout, and malformed-response
111/// errors remain per-item results. Other errors abort and drop pending futures;
112/// dropping a future does not establish remote cancellation. Empty inputs do no I/O.
113/// This primitive sends its states unchanged; recipes apply `context` beforehand.
114///
115/// # Errors
116/// Propagates local request/image validation errors and provider errors
117/// other than `Provider`, `Timeout`, or `InvalidResponse`. Those three kinds
118/// are retained as per-item [`ItemResult::Error`] values instead of failing
119/// the batch. An outer error drops pending local futures.
120pub async fn ask_each(
121    provider: &dyn Provider,
122    question: &Question,
123    states: &[State],
124    options: &BatchOptions,
125) -> Result<Vec<ItemResult>> {
126    let questions = IndexMap::from([("q".to_owned(), question.clone())]);
127    let mut requests = Vec::with_capacity(states.len());
128    // Validate all locally available configuration before polling any provider.
129    for state in states {
130        let request = Request::new(state.clone(), questions.clone())?;
131        #[cfg(feature = "images")]
132        let request = {
133            options.assets.validate(state, &questions)?;
134            request.with_assets(options.assets.clone())
135        };
136        requests.push(request);
137    }
138    let mut results: Vec<(usize, ItemResult)> = stream::iter(requests.into_iter().enumerate().map(
139        |(position, request)| async move {
140            let result = match ask(provider, &request).await {
141                Ok(mut answers) => match answers.shift_remove("q") {
142                    Some(answer) => ItemResult::Answer(answer),
143                    None => ItemResult::Error(Error::new(
144                        ErrorKind::InvalidResponse,
145                        "provider omitted item answer",
146                    )),
147                },
148                Err(error) => match error.kind() {
149                    ErrorKind::Provider | ErrorKind::Timeout | ErrorKind::InvalidResponse => {
150                        ItemResult::Error(error)
151                    }
152                    _ => return Err(error),
153                },
154            };
155            Ok((position, result))
156        },
157    ))
158    .buffer_unordered(options.concurrency.get())
159    .try_collect()
160    .await?;
161    results.sort_unstable_by_key(|(position, _)| *position);
162    Ok(results.into_iter().map(|(_, result)| result).collect())
163}