semantic_operators/
each.rs1use 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#[derive(Clone)]
12pub struct BatchOptions {
13 pub concurrency: NonZeroUsize,
15 pub context: Option<Value>,
17 #[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#[derive(Clone)]
42pub struct Item {
43 id: String,
44 state: State,
45}
46impl Item {
47 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 pub fn id(&self) -> &str {
59 &self.id
60 }
61 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#[derive(Clone, Debug)]
73pub enum ItemResult {
74 Answer(Answer<AnswerValue>),
76 Error(Error),
78}
79impl ItemResult {
80 pub fn answer(&self) -> Option<&Answer<AnswerValue>> {
82 match self {
83 Self::Answer(a) => Some(a),
84 Self::Error(_) => None,
85 }
86 }
87 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}
108pub 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 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}