1use qql_core::ast::{
2 EmbedKind, EmbeddingSpec, PointEntry, PointVectors, Stmt, UpsertStmt, VectorValue,
3};
4use qql_core::error::QqlError;
5
6use crate::embedder::Embedder;
7
8pub(crate) use super::resolve_query::ensure_batch_len;
9use super::resolve_query::resolve_query_embeddings;
10pub use super::resolve_query::{DENSE_VECTOR_NAME, SPARSE_VECTOR_NAME};
11
12pub async fn resolve_embeddings(stmt: &mut Stmt, embedder: &dyn Embedder) -> Result<(), QqlError> {
21 match stmt {
22 Stmt::Query(query) => resolve_query_embeddings(query, embedder).await?,
23 Stmt::Upsert(upsert) => resolve_upsert_embeddings(upsert, embedder).await?,
24 _ => {}
25 }
26 Ok(())
27}
28
29async fn resolve_upsert_embeddings(
30 upsert: &mut UpsertStmt,
31 embedder: &dyn Embedder,
32) -> Result<(), QqlError> {
33 if upsert.embedding.is_none() && upsert.embed.is_empty() {
34 let mut targets = Vec::new();
35 for (idx, point) in upsert.points.iter().enumerate() {
36 let PointEntry::Inline(inline) = point else {
37 continue;
38 };
39 if inline.vectors.is_none()
40 && let Some((_, qql_core::ast::Value::Str(text))) =
41 inline.payload.iter().find(|(k, _)| {
42 eq_lowered(k, "text") || eq_lowered(k, "body") || eq_lowered(k, "content")
43 })
44 && !text.is_empty()
45 {
46 targets.push((idx, text.clone()));
47 }
48 }
49 if !targets.is_empty() {
50 let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
55 let dense_vecs = embedder.embed_dense_batch(&texts, "default").await?;
56 ensure_batch_len(dense_vecs.len(), indices.len(), "default")?;
57 for (idx, d_vec) in indices.into_iter().zip(dense_vecs) {
58 let point = &mut upsert.points[idx];
59 add_point_vector(point, DENSE_VECTOR_NAME, VectorValue::Dense(d_vec))?;
60 }
61 }
62 }
63
64 if let Some(spec) = upsert.embedding.clone() {
65 let mut seen_vectors = Vec::new();
68 resolve_single_embedding_spec(upsert, &spec, embedder, &mut seen_vectors).await?;
69 }
70
71 for directive in &upsert.embed {
72 let field_name = &directive.source_field;
73 let target_vec_name = &directive.target_vector;
74 let mut targets = Vec::new();
75 for (idx, point) in upsert.points.iter().enumerate() {
76 let PointEntry::Inline(inline) = point else {
77 continue;
78 };
79 if let Some((_, qql_core::ast::Value::Str(text))) = inline
80 .payload
81 .iter()
82 .find(|(k, _)| k.eq_ignore_ascii_case(field_name))
83 && !text.is_empty()
84 {
85 targets.push((idx, text.clone()));
86 }
87 }
88
89 if !targets.is_empty() {
90 match &directive.kind {
91 EmbedKind::Dense { model } => {
92 let m_name = model.as_deref().unwrap_or("default");
93 let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
94 let vecs = embedder.embed_dense_batch(&texts, m_name).await?;
95 ensure_batch_len(vecs.len(), indices.len(), m_name)?;
96 for (idx, vec) in indices.into_iter().zip(vecs) {
97 let point = &mut upsert.points[idx];
98 add_point_vector(point, target_vec_name, VectorValue::Dense(vec))?;
99 }
100 }
101 EmbedKind::Sparse { model } => {
102 let m = model.as_deref().unwrap_or("default");
103 let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
104 let vecs = embedder.embed_sparse_document_batch(&texts, m).await?;
105 ensure_batch_len(vecs.len(), indices.len(), m)?;
106 for (idx, s_vec) in indices.into_iter().zip(vecs) {
107 let point = &mut upsert.points[idx];
108 add_point_vector(
109 point,
110 target_vec_name,
111 VectorValue::Sparse {
112 indices: s_vec.indices,
113 values: s_vec.values,
114 },
115 )?;
116 }
117 }
118 EmbedKind::Multi { model } => {
119 let m_name = model.as_deref().unwrap_or("default");
120 let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
121 let bags = embedder.embed_multi_batch(&texts, m_name).await?;
122 if bags.len() != indices.len() {
123 return Err(QqlError::execution(
124 "QQL-EMBEDDING-MULTI",
125 format!(
126 "embed_multi_batch returned {} bags for {} texts (model={m_name})",
127 bags.len(),
128 indices.len()
129 ),
130 None,
131 ));
132 }
133 for (idx, rows) in indices.into_iter().zip(bags) {
134 if rows.is_empty() {
135 return Err(QqlError::execution(
136 "QQL-EMBEDDING-MULTI",
137 "embed_multi returned an empty multivector",
138 None,
139 ));
140 }
141 let point = &mut upsert.points[idx];
142 add_point_vector(point, target_vec_name, VectorValue::MultiDense(rows))?;
143 }
144 }
145 EmbedKind::Image { model } => {
146 let m_name = model.as_deref().unwrap_or("default");
147 let (indices, sources): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
148 let vecs = embedder.embed_image_batch(&sources, m_name).await?;
149 if vecs.len() != indices.len() {
150 return Err(QqlError::execution(
151 "QQL-EMBEDDING-IMAGE",
152 format!(
153 "embed_image_batch returned {} vectors for {} sources (model={m_name})",
154 vecs.len(),
155 indices.len()
156 ),
157 None,
158 ));
159 }
160 for (idx, vec) in indices.into_iter().zip(vecs) {
161 let point = &mut upsert.points[idx];
162 add_point_vector(point, target_vec_name, VectorValue::Dense(vec))?;
163 }
164 }
165 }
166 }
167 }
168
169 Ok(())
170}
171
172async fn resolve_single_embedding_spec(
173 upsert: &mut UpsertStmt,
174 spec: &EmbeddingSpec,
175 embedder: &dyn Embedder,
176 seen_vectors: &mut Vec<String>,
177) -> Result<(), QqlError> {
178 match spec {
179 EmbeddingSpec::Multi(specs) => {
180 for sub_spec in specs {
181 Box::pin(resolve_single_embedding_spec(
182 upsert,
183 sub_spec,
184 embedder,
185 seen_vectors,
186 ))
187 .await?;
188 }
189 }
190 EmbeddingSpec::Dense {
191 model,
192 vector,
193 field,
194 } => {
195 let model_name = model.as_deref().unwrap_or("default");
196 let vector_name = vector.as_deref().unwrap_or(DENSE_VECTOR_NAME);
197 check_and_insert_vector_name(seen_vectors, vector_name)?;
198
199 let targets = collect_text_targets(&upsert.points, field.as_deref());
200 validate_non_empty_targets(upsert, &targets, "DENSE", field.as_deref())?;
201
202 let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
203 let vecs = embedder.embed_dense_batch(&texts, model_name).await?;
204 ensure_batch_len(vecs.len(), indices.len(), model_name)?;
205 for (idx, vec) in indices.into_iter().zip(vecs) {
206 let point = &mut upsert.points[idx];
207 add_point_vector(point, vector_name, VectorValue::Dense(vec))?;
208 }
209 }
210 EmbeddingSpec::Sparse {
211 model,
212 vector,
213 field,
214 } => {
215 let model_name = model.as_deref().unwrap_or("default");
216 let vector_name = vector.as_deref().unwrap_or(SPARSE_VECTOR_NAME);
217 check_and_insert_vector_name(seen_vectors, vector_name)?;
218
219 let targets = collect_text_targets(&upsert.points, field.as_deref());
220 validate_non_empty_targets(upsert, &targets, "SPARSE", field.as_deref())?;
221
222 let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
223 let vecs = embedder
224 .embed_sparse_document_batch(&texts, model_name)
225 .await?;
226 ensure_batch_len(vecs.len(), indices.len(), model_name)?;
227 for (idx, sparse_vec) in indices.into_iter().zip(vecs) {
228 add_point_vector(
229 &mut upsert.points[idx],
230 vector_name,
231 VectorValue::Sparse {
232 indices: sparse_vec.indices,
233 values: sparse_vec.values,
234 },
235 )?;
236 }
237 }
238 EmbeddingSpec::Hybrid {
239 dense_model,
240 dense_vector,
241 dense_field,
242 sparse_model,
243 sparse_vector,
244 sparse_field,
245 } => {
246 let d_model = dense_model.as_deref().unwrap_or("default");
247 let s_model = sparse_model.as_deref().unwrap_or("default");
248 let d_vec_name = dense_vector.as_deref().unwrap_or(DENSE_VECTOR_NAME);
249 let s_vec_name = sparse_vector.as_deref().unwrap_or(SPARSE_VECTOR_NAME);
250
251 check_and_insert_vector_name(seen_vectors, d_vec_name)?;
252 check_and_insert_vector_name(seen_vectors, s_vec_name)?;
253
254 let dense_targets = collect_text_targets(&upsert.points, dense_field.as_deref());
255 let sparse_targets = collect_text_targets(&upsert.points, sparse_field.as_deref());
256
257 validate_non_empty_targets(upsert, &dense_targets, "DENSE", dense_field.as_deref())?;
258 validate_non_empty_targets(upsert, &sparse_targets, "SPARSE", sparse_field.as_deref())?;
259
260 let (indices, texts): (Vec<usize>, Vec<String>) = dense_targets.into_iter().unzip();
261 let dense_vecs = embedder.embed_dense_batch(&texts, d_model).await?;
262 ensure_batch_len(dense_vecs.len(), indices.len(), d_model)?;
263 for (idx, d_vec) in indices.into_iter().zip(dense_vecs) {
264 let point = &mut upsert.points[idx];
265 add_point_vector(point, d_vec_name, VectorValue::Dense(d_vec))?;
266 }
267
268 let (sparse_indices, sparse_texts): (Vec<usize>, Vec<String>) =
269 sparse_targets.into_iter().unzip();
270 let sparse_vecs = embedder
271 .embed_sparse_document_batch(&sparse_texts, s_model)
272 .await?;
273 ensure_batch_len(sparse_vecs.len(), sparse_indices.len(), s_model)?;
274 for (idx, sparse_vec) in sparse_indices.into_iter().zip(sparse_vecs) {
275 let point = &mut upsert.points[idx];
276 add_point_vector(
277 point,
278 s_vec_name,
279 VectorValue::Sparse {
280 indices: sparse_vec.indices,
281 values: sparse_vec.values,
282 },
283 )?;
284 }
285 }
286 EmbeddingSpec::MultiVector {
287 model,
288 vector,
289 field,
290 } => {
291 let model_name = model.as_deref().unwrap_or("default");
292 let vector_name = vector.as_deref().unwrap_or("colbert");
293 check_and_insert_vector_name(seen_vectors, vector_name)?;
294
295 let targets = collect_text_targets(&upsert.points, field.as_deref());
296 validate_non_empty_targets(upsert, &targets, "MULTI", field.as_deref())?;
297
298 let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
299 let bags = embedder.embed_multi_batch(&texts, model_name).await?;
300 if bags.len() != indices.len() {
301 return Err(QqlError::execution(
302 "QQL-EMBEDDING-MULTI",
303 format!(
304 "embed_multi_batch returned {} bags for {} texts (model={model_name})",
305 bags.len(),
306 indices.len()
307 ),
308 None,
309 ));
310 }
311 for (idx, rows) in indices.into_iter().zip(bags) {
312 if rows.is_empty() {
313 return Err(QqlError::execution(
314 "QQL-EMBEDDING-MULTI",
315 "embed_multi returned an empty multivector",
316 None,
317 ));
318 }
319 add_point_vector(
320 &mut upsert.points[idx],
321 vector_name,
322 VectorValue::MultiDense(rows),
323 )?;
324 }
325 }
326 EmbeddingSpec::Image {
327 model,
328 vector,
329 field,
330 } => {
331 let model_name = model.as_deref().unwrap_or("default");
332 let vector_name = vector.as_deref().unwrap_or("image");
333 check_and_insert_vector_name(seen_vectors, vector_name)?;
334
335 let targets = collect_image_targets(&upsert.points, field.as_deref());
336 validate_non_empty_targets(upsert, &targets, "IMAGE", field.as_deref())?;
337
338 let (indices, sources): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
339 let vecs = embedder.embed_image_batch(&sources, model_name).await?;
340 if vecs.len() != indices.len() {
341 return Err(QqlError::execution(
342 "QQL-EMBEDDING-IMAGE",
343 format!(
344 "embed_image_batch returned {} vectors for {} sources (model={model_name})",
345 vecs.len(),
346 indices.len()
347 ),
348 None,
349 ));
350 }
351 for (idx, vec) in indices.into_iter().zip(vecs) {
352 add_point_vector(
353 &mut upsert.points[idx],
354 vector_name,
355 VectorValue::Dense(vec),
356 )?;
357 }
358 }
359 }
360 Ok(())
361}
362
363fn validate_non_empty_targets(
364 upsert: &UpsertStmt,
365 targets: &[(usize, String)],
366 kind: &str,
367 field: Option<&str>,
368) -> Result<(), QqlError> {
369 if targets.is_empty() {
370 let actual_fields = upsert
371 .points
372 .first()
373 .and_then(|p| match p {
374 PointEntry::Inline(inline) => Some(inline),
375 PointEntry::Param(..) | PointEntry::PositionalParam(..) => None,
376 })
377 .map(|p| {
378 p.payload
379 .iter()
380 .map(|(k, _)| k.as_str())
381 .collect::<Vec<_>>()
382 .join(", ")
383 })
384 .unwrap_or_default();
385
386 let err_msg = if let Some(f) = field {
387 format!(
388 "USING {kind} MODEL specified with ON FIELD '{f}' but no matching text payload field found. Found fields: {actual_fields}"
389 )
390 } else {
391 format!(
392 "USING {kind} MODEL specified but no text payload field found. Expected one of: {}. Found fields: {actual_fields}",
393 DEFAULT_TEXT_FIELDS_ORDERED.join(", ")
394 )
395 };
396
397 return Err(QqlError::execution("QQL-EMBEDDING", err_msg, None));
398 }
399 Ok(())
400}
401
402fn check_and_insert_vector_name(
403 seen_vectors: &mut Vec<String>,
404 vector_name: &str,
405) -> Result<(), QqlError> {
406 if seen_vectors.iter().any(|name| name == vector_name) {
407 return Err(QqlError::execution(
408 "QQL-EMBEDDING",
409 format!("duplicate target vector '{vector_name}' in multi-spec embedding clause"),
410 None,
411 ));
412 }
413 seen_vectors.push(vector_name.to_string());
414 Ok(())
415}
416
417const DEFAULT_TEXT_FIELDS_ORDERED: &[&str] = &[
418 "text",
419 "body",
420 "content",
421 "title",
422 "description",
423 "name",
424 "summary",
425 "document",
426];
427
428const DEFAULT_IMAGE_FIELDS_ORDERED: &[&str] = &[
429 "image",
430 "image_path",
431 "image_url",
432 "photo",
433 "picture",
434 "img",
435 "path",
436 "url",
437];
438
439fn eq_lowered(haystack: &str, needle_lower: &str) -> bool {
446 if haystack == needle_lower {
447 return true;
448 }
449 haystack.len() == needle_lower.len()
452 && haystack
453 .as_bytes()
454 .iter()
455 .zip(needle_lower.as_bytes())
456 .all(|(a, b)| a.to_ascii_lowercase() == *b)
457}
458
459fn collect_image_targets(
461 points: &[PointEntry],
462 field_override: Option<&str>,
463) -> Vec<(usize, String)> {
464 if let Some(target_field) = field_override {
465 let target_lower = target_field.to_ascii_lowercase();
466 points
467 .iter()
468 .enumerate()
469 .filter_map(|(idx, point)| {
470 let PointEntry::Inline(inline) = point else {
471 return None;
472 };
473 inline.payload.iter().find_map(|(key, value)| {
474 if eq_lowered(key, &target_lower)
475 && let qql_core::ast::Value::Str(source) = value
476 && !source.is_empty()
477 {
478 return Some((idx, source.clone()));
479 }
480 None
481 })
482 })
483 .collect()
484 } else {
485 points
486 .iter()
487 .enumerate()
488 .filter_map(|(idx, point)| {
489 let PointEntry::Inline(inline) = point else {
490 return None;
491 };
492 for &candidate in DEFAULT_IMAGE_FIELDS_ORDERED {
493 if let Some((_, qql_core::ast::Value::Str(source))) = inline
494 .payload
495 .iter()
496 .find(|(key, _)| eq_lowered(key, candidate))
497 && !source.is_empty()
498 {
499 return Some((idx, source.clone()));
500 }
501 }
502 None
503 })
504 .collect()
505 }
506}
507
508fn collect_text_targets(
509 points: &[PointEntry],
510 field_override: Option<&str>,
511) -> Vec<(usize, String)> {
512 if let Some(target_field) = field_override {
513 let target_lower = target_field.to_ascii_lowercase();
514 points
515 .iter()
516 .enumerate()
517 .filter_map(|(idx, point)| {
518 let PointEntry::Inline(inline) = point else {
519 return None;
520 };
521 inline.payload.iter().find_map(|(key, value)| {
522 if eq_lowered(key, &target_lower)
523 && let qql_core::ast::Value::Str(text) = value
524 && !text.is_empty()
525 {
526 return Some((idx, text.clone()));
527 }
528 None
529 })
530 })
531 .collect()
532 } else {
533 collect_default_text_targets(points)
534 }
535}
536
537fn collect_default_text_targets(points: &[PointEntry]) -> Vec<(usize, String)> {
538 points
539 .iter()
540 .enumerate()
541 .filter_map(|(idx, point)| {
542 let PointEntry::Inline(inline) = point else {
543 return None;
544 };
545 for &candidate in DEFAULT_TEXT_FIELDS_ORDERED {
546 if let Some((_, qql_core::ast::Value::Str(text))) = inline
547 .payload
548 .iter()
549 .find(|(key, _)| eq_lowered(key, candidate))
550 && !text.is_empty()
551 {
552 return Some((idx, text.clone()));
553 }
554 }
555 None
556 })
557 .collect()
558}
559
560fn add_point_vector(
561 point: &mut PointEntry,
562 name: &str,
563 vector: VectorValue,
564) -> Result<(), QqlError> {
565 let point = match point {
566 PointEntry::Inline(inline) => inline,
567 PointEntry::Param(name, span) => {
570 return Err(QqlError::execution(
571 "QQL-EMBEDDING",
572 format!("cannot embed into unbound point parameter ':{name}'"),
573 span.as_deref().copied(),
574 ));
575 }
576 PointEntry::PositionalParam(idx, span) => {
577 return Err(QqlError::execution(
578 "QQL-EMBEDDING",
579 format!("cannot embed into unbound point parameter '?{}'", *idx + 1),
580 span.as_deref().copied(),
581 ));
582 }
583 };
584 if name.is_empty() {
585 return match &mut point.vectors {
586 Some(PointVectors::Unnamed(existing)) => {
587 *existing = vector;
588 Ok(())
589 }
590 Some(PointVectors::Named(list)) => {
591 if let Some(existing) = list.iter_mut().find(|(key, _)| key.is_empty()) {
592 existing.1 = vector;
593 } else {
594 list.push((String::new(), vector));
595 }
596 Ok(())
597 }
598 Some(PointVectors::Param(..)) | Some(PointVectors::PositionalParam(..)) => {
599 point.vectors = Some(PointVectors::Unnamed(vector));
600 Ok(())
601 }
602 None => {
603 point.vectors = Some(PointVectors::Unnamed(vector));
604 Ok(())
605 }
606 };
607 }
608 match &mut point.vectors {
609 Some(PointVectors::Named(list)) => {
610 if let Some(existing) = list.iter_mut().find(|(k, _)| k == name) {
611 existing.1 = vector;
612 } else {
613 list.push((name.to_string(), vector));
614 }
615 Ok(())
616 }
617 Some(PointVectors::Unnamed(_))
618 | Some(PointVectors::Param(..))
619 | Some(PointVectors::PositionalParam(..)) => Err(QqlError::execution(
620 "QQL-EMBEDDING",
621 format!(
622 "cannot add named vector '{name}' to a point that already has an unnamed vector; \
623 provide an explicit named-vector topology or omit EMBED for this point"
624 ),
625 None,
626 )),
627 None => {
628 point.vectors = Some(PointVectors::Named(vec![(name.to_string(), vector)]));
629 Ok(())
630 }
631 }
632}