qql_embed/
resolve.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4
5use qql_core::ast::{
6    EmbedKind, EmbeddingSpec, PointVectors, Prefetch, PrefetchSource, QueryExpr, QueryInput,
7    QueryStmt, Stmt, UpsertPoint, UpsertStmt, VectorKind, VectorTarget, VectorValue,
8};
9use qql_core::error::QqlError;
10
11use crate::embedder::Embedder;
12
13/// Default named dense vector for auto-embedding.
14pub const DENSE_VECTOR_NAME: &str = "dense";
15/// Default named sparse vector for auto-embedding.
16pub const SPARSE_VECTOR_NAME: &str = "sparse";
17
18#[cfg(not(target_arch = "wasm32"))]
19type BoxFut<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
20#[cfg(target_arch = "wasm32")]
21type BoxFut<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
22
23/// Dense vector iterator passed through recursive apply.
24#[cfg(not(target_arch = "wasm32"))]
25type DenseIter<'a> = &'a mut (dyn Iterator<Item = Vec<f32>> + Send);
26#[cfg(target_arch = "wasm32")]
27type DenseIter<'a> = &'a mut dyn Iterator<Item = Vec<f32>>;
28
29/// Resolve text → vectors on a statement before routing/execution.
30///
31/// Dense jobs are collected and sent through `embed_dense_batch` (grouped by
32/// model). Sparse stays local BM25 via the embedder.
33pub async fn resolve_embeddings(stmt: &mut Stmt, embedder: &dyn Embedder) -> Result<(), QqlError> {
34    match stmt {
35        Stmt::Query(query) => resolve_query_embeddings(query, embedder).await?,
36        Stmt::Upsert(upsert) => resolve_upsert_embeddings(upsert, embedder).await?,
37        _ => {}
38    }
39    Ok(())
40}
41
42async fn resolve_query_embeddings(
43    query: &mut QueryStmt,
44    embedder: &dyn Embedder,
45) -> Result<(), QqlError> {
46    let mut dense_jobs: Vec<(String, String)> = Vec::new();
47    collect_query_dense_jobs(query, &mut dense_jobs)?;
48
49    let dense_vecs = batch_dense_by_model(embedder, &dense_jobs).await?;
50    let mut dense_iter = dense_vecs.into_iter();
51    apply_query_embeddings(query, embedder, &mut dense_iter).await?;
52
53    if dense_iter.next().is_some() {
54        return Err(QqlError::execution(
55            "QQL-EMBEDDING",
56            "internal error: unused dense embeddings after apply",
57            None,
58        ));
59    }
60    Ok(())
61}
62
63async fn resolve_upsert_embeddings(
64    upsert: &mut UpsertStmt,
65    embedder: &dyn Embedder,
66) -> Result<(), QqlError> {
67    if upsert.embedding.is_none() && upsert.embed.is_empty() {
68        let mut targets = Vec::new();
69        for (idx, point) in upsert.points.iter().enumerate() {
70            if point.vectors.is_none() {
71                if let Some((_, qql_core::ast::Value::Str(text))) =
72                    point.payload.iter().find(|(k, _)| {
73                        k.eq_ignore_ascii_case("text")
74                            || k.eq_ignore_ascii_case("body")
75                            || k.eq_ignore_ascii_case("content")
76                    })
77                {
78                    if !text.is_empty() {
79                        targets.push((idx, text.clone()));
80                    }
81                }
82            }
83        }
84        if !targets.is_empty() {
85            // Topology-unaware fallback: dense only. Hybrid/sparse targets must
86            // be set by the executor (configure_upsert_embeddings) or explicit
87            // USING / EMBED directives before calling resolve_embeddings — so
88            // dense-only collections never receive orphan sparse vectors.
89            let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
90            let dense_vecs = embedder.embed_dense_batch(&texts, "default").await?;
91            ensure_batch_len(dense_vecs.len(), indices.len(), "default")?;
92            for (idx, d_vec) in indices.into_iter().zip(dense_vecs) {
93                let point = &mut upsert.points[idx];
94                add_point_vector(point, DENSE_VECTOR_NAME, VectorValue::Dense(d_vec))?;
95            }
96        }
97    }
98
99    if let Some(spec) = upsert.embedding.clone() {
100        let mut seen_vectors = std::collections::HashSet::new();
101        resolve_single_embedding_spec(upsert, &spec, embedder, &mut seen_vectors).await?;
102    }
103
104    for directive in &upsert.embed {
105        let field_name = &directive.source_field;
106        let target_vec_name = &directive.target_vector;
107        let mut targets = Vec::new();
108        for (idx, point) in upsert.points.iter().enumerate() {
109            if let Some((_, qql_core::ast::Value::Str(text))) = point
110                .payload
111                .iter()
112                .find(|(k, _)| k.eq_ignore_ascii_case(field_name))
113            {
114                if !text.is_empty() {
115                    targets.push((idx, text.clone()));
116                }
117            }
118        }
119
120        if !targets.is_empty() {
121            match &directive.kind {
122                EmbedKind::Dense { model } => {
123                    let m_name = model.as_deref().unwrap_or("default");
124                    let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
125                    let vecs = embedder.embed_dense_batch(&texts, m_name).await?;
126                    ensure_batch_len(vecs.len(), indices.len(), m_name)?;
127                    for (idx, vec) in indices.into_iter().zip(vecs) {
128                        let point = &mut upsert.points[idx];
129                        add_point_vector(point, target_vec_name, VectorValue::Dense(vec))?;
130                    }
131                }
132                EmbedKind::Sparse { model } => {
133                    let m = model.as_deref().unwrap_or("default");
134                    for (idx, text) in targets {
135                        let s_vec = embedder.embed_sparse(&text, m).await?;
136                        let point = &mut upsert.points[idx];
137                        add_point_vector(
138                            point,
139                            target_vec_name,
140                            VectorValue::Sparse {
141                                indices: s_vec.indices,
142                                values: s_vec.values,
143                            },
144                        )?;
145                    }
146                }
147                EmbedKind::Multi { model } => {
148                    let m_name = model.as_deref().unwrap_or("default");
149                    let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
150                    let bags = embedder.embed_multi_batch(&texts, m_name).await?;
151                    if bags.len() != indices.len() {
152                        return Err(QqlError::execution(
153                            "QQL-EMBEDDING-MULTI",
154                            format!(
155                                "embed_multi_batch returned {} bags for {} texts (model={m_name})",
156                                bags.len(),
157                                indices.len()
158                            ),
159                            None,
160                        ));
161                    }
162                    for (idx, rows) in indices.into_iter().zip(bags) {
163                        if rows.is_empty() {
164                            return Err(QqlError::execution(
165                                "QQL-EMBEDDING-MULTI",
166                                "embed_multi returned an empty multivector",
167                                None,
168                            ));
169                        }
170                        let point = &mut upsert.points[idx];
171                        add_point_vector(point, target_vec_name, VectorValue::MultiDense(rows))?;
172                    }
173                }
174                EmbedKind::Image { model } => {
175                    let m_name = model.as_deref().unwrap_or("default");
176                    let (indices, sources): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
177                    let vecs = embedder.embed_image_batch(&sources, m_name).await?;
178                    if vecs.len() != indices.len() {
179                        return Err(QqlError::execution(
180                            "QQL-EMBEDDING-IMAGE",
181                            format!(
182                                "embed_image_batch returned {} vectors for {} sources (model={m_name})",
183                                vecs.len(),
184                                indices.len()
185                            ),
186                            None,
187                        ));
188                    }
189                    for (idx, vec) in indices.into_iter().zip(vecs) {
190                        let point = &mut upsert.points[idx];
191                        add_point_vector(point, target_vec_name, VectorValue::Dense(vec))?;
192                    }
193                }
194            }
195        }
196    }
197
198    Ok(())
199}
200
201// ── Collect dense text jobs (model, text) in walk order ─────────────
202
203fn collect_query_dense_jobs(
204    query: &QueryStmt,
205    jobs: &mut Vec<(String, String)>,
206) -> Result<(), QqlError> {
207    for cte in &query.ctes {
208        collect_expr_dense_jobs(&cte.query.expression, jobs)?;
209    }
210    collect_expr_dense_jobs(&query.expression, jobs)
211}
212
213fn collect_prefetches_dense_jobs(
214    prefetches: &[Prefetch],
215    jobs: &mut Vec<(String, String)>,
216) -> Result<(), QqlError> {
217    for pref in prefetches {
218        if let PrefetchSource::Query(sub) = &pref.source {
219            collect_query_dense_jobs(sub, jobs)?;
220        }
221    }
222    Ok(())
223}
224
225fn collect_expr_dense_jobs(
226    expr: &QueryExpr,
227    jobs: &mut Vec<(String, String)>,
228) -> Result<(), QqlError> {
229    match expr {
230        QueryExpr::Nearest {
231            input,
232            using,
233            prefetch,
234            ..
235        } => {
236            collect_input_dense_job(input, require_embed_target(using)?, "default", jobs);
237            collect_prefetches_dense_jobs(prefetch, jobs)?;
238        }
239        QueryExpr::Recommend {
240            positive,
241            negative,
242            using,
243            prefetch,
244            ..
245        } => {
246            let target = require_embed_target(using)?;
247            for input in positive.iter().chain(negative.iter()) {
248                collect_input_dense_job(input, target, "default", jobs);
249            }
250            collect_prefetches_dense_jobs(prefetch, jobs)?;
251        }
252        QueryExpr::Context {
253            pairs,
254            using,
255            prefetch,
256            ..
257        } => {
258            let target = require_embed_target(using)?;
259            for pair in pairs {
260                collect_input_dense_job(&pair.positive, target, "default", jobs);
261                collect_input_dense_job(&pair.negative, target, "default", jobs);
262            }
263            collect_prefetches_dense_jobs(prefetch, jobs)?;
264        }
265        QueryExpr::Discover {
266            target,
267            context,
268            using,
269            prefetch,
270            ..
271        } => {
272            let emb = require_embed_target(using)?;
273            collect_input_dense_job(target, emb, "default", jobs);
274            for pair in context {
275                collect_input_dense_job(&pair.positive, emb, "default", jobs);
276                collect_input_dense_job(&pair.negative, emb, "default", jobs);
277            }
278            collect_prefetches_dense_jobs(prefetch, jobs)?;
279        }
280        QueryExpr::Fusion { prefetch, .. } | QueryExpr::Formula { prefetch, .. } => {
281            collect_prefetches_dense_jobs(prefetch, jobs)?;
282        }
283        QueryExpr::RelevanceFeedback {
284            target,
285            feedback,
286            using,
287            prefetch,
288            ..
289        } => {
290            let emb = require_embed_target(using)?;
291            collect_input_dense_job(target, emb, "default", jobs);
292            for fb in feedback {
293                collect_input_dense_job(&fb.example, emb, "default", jobs);
294            }
295            collect_prefetches_dense_jobs(prefetch, jobs)?;
296        }
297        QueryExpr::Hybrid { text, model, .. } => {
298            let m = model.as_deref().unwrap_or("default").to_string();
299            jobs.push((m, text.clone()));
300        }
301        QueryExpr::CrossRerank { prefetch, .. } => {
302            // Query string is scored by the pair model, not embedded.
303            collect_prefetches_dense_jobs(prefetch, jobs)?;
304        }
305        QueryExpr::Rerank {
306            input,
307            model,
308            using,
309            prefetch,
310            ..
311        } => {
312            // RERANK uses the MODEL string for dense/multi, not "default".
313            let mut emb = require_embed_target(using)?;
314            // RERANK is always dense-family; multi comes from USING / schema.
315            if emb.kind == VectorKind::Sparse {
316                return Err(QqlError::execution(
317                    "QQL-VECTOR-KIND",
318                    "RERANK requires a dense (or multivector) target, not sparse",
319                    None,
320                ));
321            }
322            emb.kind = VectorKind::Dense;
323            collect_input_dense_job(input, emb, model.as_str(), jobs);
324            collect_prefetches_dense_jobs(prefetch, jobs)?;
325        }
326        _ => {}
327    }
328    Ok(())
329}
330
331/// Only single-vector dense TEXT inputs join the dense batch; sparse and multi
332/// are applied one-by-one later.
333fn collect_input_dense_job(
334    input: &QueryInput,
335    target: EmbedTarget,
336    default_model: &str,
337    jobs: &mut Vec<(String, String)>,
338) {
339    if let QueryInput::Text { text, model } = input {
340        if target.kind == VectorKind::Dense && !target.multi {
341            let m = model.as_deref().unwrap_or(default_model).to_string();
342            jobs.push((m, text.clone()));
343        }
344    }
345}
346
347/// Group jobs by model, call `embed_dense_batch` once per model, restore walk order.
348async fn batch_dense_by_model(
349    embedder: &dyn Embedder,
350    jobs: &[(String, String)],
351) -> Result<Vec<Vec<f32>>, QqlError> {
352    if jobs.is_empty() {
353        return Ok(Vec::new());
354    }
355
356    let mut by_model: HashMap<String, Vec<usize>> = HashMap::new();
357    for (i, (model, _)) in jobs.iter().enumerate() {
358        by_model.entry(model.clone()).or_default().push(i);
359    }
360
361    let mut out: Vec<Option<Vec<f32>>> = vec![None; jobs.len()];
362    for (model, indices) in by_model {
363        let texts: Vec<String> = indices.iter().map(|&i| jobs[i].1.clone()).collect();
364        let vecs = embedder.embed_dense_batch(&texts, &model).await?;
365        if vecs.len() != indices.len() {
366            return Err(QqlError::execution(
367                "QQL-EMBEDDING",
368                format!(
369                    "embed_dense_batch returned {} vectors for {} texts (model={model})",
370                    vecs.len(),
371                    indices.len()
372                ),
373                None,
374            ));
375        }
376        for (idx, vec) in indices.into_iter().zip(vecs) {
377            out[idx] = Some(vec);
378        }
379    }
380
381    out.into_iter()
382        .enumerate()
383        .map(|(i, v)| {
384            v.ok_or_else(|| {
385                QqlError::execution(
386                    "QQL-EMBEDDING",
387                    format!("missing dense embedding at job index {i}"),
388                    None,
389                )
390            })
391        })
392        .collect()
393}
394
395// ── Apply dense vectors (in collect order) + resolve sparse ─────────
396
397fn apply_query_embeddings<'a>(
398    query: &'a mut QueryStmt,
399    embedder: &'a dyn Embedder,
400    dense: DenseIter<'a>,
401) -> BoxFut<'a, Result<(), QqlError>> {
402    Box::pin(async move {
403        for cte in &mut query.ctes {
404            apply_expr_embeddings(&mut cte.query.expression, embedder, dense).await?;
405        }
406        apply_expr_embeddings(&mut query.expression, embedder, dense).await?;
407        Ok(())
408    })
409}
410
411fn apply_prefetches_embeddings<'a>(
412    prefetches: &'a mut [Prefetch],
413    embedder: &'a dyn Embedder,
414    dense: DenseIter<'a>,
415) -> BoxFut<'a, Result<(), QqlError>> {
416    Box::pin(async move {
417        for pref in prefetches {
418            if let PrefetchSource::Query(sub) = &mut pref.source {
419                apply_query_embeddings(sub, embedder, dense).await?;
420            }
421        }
422        Ok(())
423    })
424}
425
426fn apply_expr_embeddings<'a>(
427    expr: &'a mut QueryExpr,
428    embedder: &'a dyn Embedder,
429    dense: DenseIter<'a>,
430) -> BoxFut<'a, Result<(), QqlError>> {
431    Box::pin(async move {
432        match expr {
433            QueryExpr::Nearest {
434                input,
435                using,
436                prefetch,
437                ..
438            } => {
439                apply_input(
440                    input,
441                    require_embed_target(using)?,
442                    "default",
443                    embedder,
444                    dense,
445                )
446                .await?;
447                apply_prefetches_embeddings(prefetch, embedder, dense).await?;
448            }
449            QueryExpr::Recommend {
450                positive,
451                negative,
452                using,
453                prefetch,
454                ..
455            } => {
456                let target = require_embed_target(using)?;
457                for input in positive.iter_mut().chain(negative.iter_mut()) {
458                    apply_input(input, target, "default", embedder, dense).await?;
459                }
460                apply_prefetches_embeddings(prefetch, embedder, dense).await?;
461            }
462            QueryExpr::Context {
463                pairs,
464                using,
465                prefetch,
466                ..
467            } => {
468                let target = require_embed_target(using)?;
469                for pair in pairs {
470                    apply_input(&mut pair.positive, target, "default", embedder, dense).await?;
471                    apply_input(&mut pair.negative, target, "default", embedder, dense).await?;
472                }
473                apply_prefetches_embeddings(prefetch, embedder, dense).await?;
474            }
475            QueryExpr::Discover {
476                target,
477                context,
478                using,
479                prefetch,
480                ..
481            } => {
482                let emb = require_embed_target(using)?;
483                apply_input(target, emb, "default", embedder, dense).await?;
484                for pair in context {
485                    apply_input(&mut pair.positive, emb, "default", embedder, dense).await?;
486                    apply_input(&mut pair.negative, emb, "default", embedder, dense).await?;
487                }
488                apply_prefetches_embeddings(prefetch, embedder, dense).await?;
489            }
490            QueryExpr::Fusion { prefetch, .. } | QueryExpr::Formula { prefetch, .. } => {
491                apply_prefetches_embeddings(prefetch, embedder, dense).await?;
492            }
493            QueryExpr::RelevanceFeedback {
494                target,
495                feedback,
496                using,
497                prefetch,
498                ..
499            } => {
500                let emb = require_embed_target(using)?;
501                apply_input(target, emb, "default", embedder, dense).await?;
502                for fb in feedback {
503                    apply_input(&mut fb.example, emb, "default", embedder, dense).await?;
504                }
505                apply_prefetches_embeddings(prefetch, embedder, dense).await?;
506            }
507            QueryExpr::Hybrid {
508                text,
509                dense_vector,
510                sparse_vector,
511                fusion,
512                ..
513            } => {
514                let d_vec = dense.next().ok_or_else(|| {
515                    QqlError::execution(
516                        "QQL-EMBEDDING",
517                        "internal error: ran out of dense embeddings for HYBRID",
518                        None,
519                    )
520                })?;
521                let s_vec = embedder.embed_sparse(text, "default").await?;
522                let d_vec_name = dense_vector.as_deref().unwrap_or(DENSE_VECTOR_NAME);
523                let s_vec_name = sparse_vector.as_deref().unwrap_or(SPARSE_VECTOR_NAME);
524
525                let dense_sub = QueryStmt {
526                    ctes: Vec::new(),
527                    collection: qql_core::ast::QueryCollection::Inherited,
528                    expression: QueryExpr::Nearest {
529                        input: QueryInput::Vector(VectorValue::Dense(d_vec)),
530                        using: Some(VectorTarget {
531                            name: d_vec_name.to_string(),
532                            kind: Some(VectorKind::Dense),
533                            multi: false,
534                        }),
535                        prefetch: Vec::new(),
536                        mmr: None,
537                    },
538                    filter: None,
539                    params: None,
540                    score_threshold: None,
541                    group: None,
542                    output: qql_core::ast::QueryOutput::default(),
543                    page: qql_core::ast::PageSpec::default(),
544                    shard_key: None,
545                };
546                let sparse_sub = QueryStmt {
547                    ctes: Vec::new(),
548                    collection: qql_core::ast::QueryCollection::Inherited,
549                    expression: QueryExpr::Nearest {
550                        input: QueryInput::Vector(VectorValue::Sparse {
551                            indices: s_vec.indices,
552                            values: s_vec.values,
553                        }),
554                        using: Some(VectorTarget {
555                            name: s_vec_name.to_string(),
556                            kind: Some(VectorKind::Sparse),
557                            multi: false,
558                        }),
559                        prefetch: Vec::new(),
560                        mmr: None,
561                    },
562                    filter: None,
563                    params: None,
564                    score_threshold: None,
565                    group: None,
566                    output: qql_core::ast::QueryOutput::default(),
567                    page: qql_core::ast::PageSpec::default(),
568                    shard_key: None,
569                };
570
571                *expr = QueryExpr::Fusion {
572                    method: *fusion,
573                    prefetch: vec![
574                        Prefetch {
575                            source: PrefetchSource::Query(Box::new(dense_sub)),
576                            filter: None,
577                            score_threshold: None,
578                            lookup: None,
579                        },
580                        Prefetch {
581                            source: PrefetchSource::Query(Box::new(sparse_sub)),
582                            filter: None,
583                            score_threshold: None,
584                            lookup: None,
585                        },
586                    ],
587                };
588            }
589            QueryExpr::Rerank {
590                input,
591                model,
592                using,
593                prefetch,
594                ..
595            } => {
596                let mut emb = require_embed_target(using)?;
597                emb.kind = VectorKind::Dense;
598                apply_input(input, emb, model.as_str(), embedder, dense).await?;
599                apply_prefetches_embeddings(prefetch, embedder, dense).await?;
600            }
601            QueryExpr::CrossRerank { prefetch, .. } => {
602                apply_prefetches_embeddings(prefetch, embedder, dense).await?;
603            }
604            _ => {}
605        }
606        Ok(())
607    })
608}
609
610async fn apply_input(
611    input: &mut QueryInput,
612    target: EmbedTarget,
613    default_model: &str,
614    embedder: &dyn Embedder,
615    dense: DenseIter<'_>,
616) -> Result<(), QqlError> {
617    match input {
618        QueryInput::Image { source, model } => {
619            // Images always produce single-vector dense (CLIP vision, etc.).
620            if target.kind == VectorKind::Sparse {
621                return Err(QqlError::execution(
622                    "QQL-VECTOR-KIND",
623                    "IMAGE input requires a dense target, not sparse",
624                    None,
625                ));
626            }
627            if target.multi {
628                return Err(QqlError::execution(
629                    "QQL-VECTOR-KIND",
630                    "IMAGE input produces single-vector dense, not multivector; use TEXT with AS MULTI for ColBERT",
631                    None,
632                ));
633            }
634            let model_name = model.as_deref().unwrap_or(default_model);
635            let vec = embedder.embed_image(source, model_name).await?;
636            if vec.is_empty() {
637                return Err(QqlError::execution(
638                    "QQL-EMBEDDING-IMAGE",
639                    "embed_image returned an empty vector",
640                    None,
641                ));
642            }
643            *input = QueryInput::Vector(VectorValue::Dense(vec));
644            Ok(())
645        }
646        QueryInput::Text { text, model } => {
647            let model_name = model.as_deref().unwrap_or(default_model);
648            if target.kind == VectorKind::Sparse {
649                let s_vec = embedder.embed_sparse(text, model_name).await?;
650                *input = QueryInput::Vector(VectorValue::Sparse {
651                    indices: s_vec.indices,
652                    values: s_vec.values,
653                });
654                return Ok(());
655            }
656            if target.multi {
657                let rows = embedder.embed_multi(text, model_name).await?;
658                if rows.is_empty() {
659                    return Err(QqlError::execution(
660                        "QQL-EMBEDDING-MULTI",
661                        "embed_multi returned an empty multivector",
662                        None,
663                    ));
664                }
665                *input = QueryInput::Vector(VectorValue::MultiDense(rows));
666                return Ok(());
667            }
668            let vec = dense.next().ok_or_else(|| {
669                QqlError::execution(
670                    "QQL-EMBEDDING",
671                    "internal error: ran out of dense embeddings",
672                    None,
673                )
674            })?;
675            *input = QueryInput::Vector(VectorValue::Dense(vec));
676            Ok(())
677        }
678        QueryInput::Vector(_) | QueryInput::Point(_) => Ok(()),
679    }
680}
681
682#[derive(Debug, Clone, Copy)]
683struct EmbedTarget {
684    kind: VectorKind,
685    multi: bool,
686}
687
688/// Resolve embed target for a `USING` clause.
689///
690/// - No `USING` → single dense.
691/// - `USING name AS …` / schema-filled kind → that kind; `multi` from AS MULTI or schema.
692/// - `USING name` with `kind: None` → error.
693fn require_embed_target(target: &Option<VectorTarget>) -> Result<EmbedTarget, QqlError> {
694    match target {
695        None => Ok(EmbedTarget {
696            kind: VectorKind::Dense,
697            multi: false,
698        }),
699        Some(t) => match t.kind {
700            Some(kind) => Ok(EmbedTarget {
701                kind,
702                multi: t.multi,
703            }),
704            None => Err(crate::topology::unknown_using_kind_error(&t.name)),
705        },
706    }
707}
708
709fn ensure_batch_len(got: usize, expected: usize, model: &str) -> Result<(), QqlError> {
710    if got != expected {
711        return Err(QqlError::execution(
712            "QQL-EMBEDDING",
713            format!(
714                "embed_dense_batch returned {got} vectors for {expected} texts (model={model})"
715            ),
716            None,
717        ));
718    }
719    Ok(())
720}
721
722async fn resolve_single_embedding_spec(
723    upsert: &mut UpsertStmt,
724    spec: &EmbeddingSpec,
725    embedder: &dyn Embedder,
726    seen_vectors: &mut std::collections::HashSet<String>,
727) -> Result<(), QqlError> {
728    match spec {
729        EmbeddingSpec::Multi(specs) => {
730            for sub_spec in specs {
731                Box::pin(resolve_single_embedding_spec(
732                    upsert,
733                    sub_spec,
734                    embedder,
735                    seen_vectors,
736                ))
737                .await?;
738            }
739        }
740        EmbeddingSpec::Dense {
741            model,
742            vector,
743            field,
744        } => {
745            let model_name = model.as_deref().unwrap_or("default");
746            let vector_name = vector.as_deref().unwrap_or(DENSE_VECTOR_NAME);
747            check_and_insert_vector_name(seen_vectors, vector_name)?;
748
749            let targets = collect_text_targets(&upsert.points, field.as_deref());
750            validate_non_empty_targets(upsert, &targets, "DENSE", field.as_deref())?;
751
752            let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
753            let vecs = embedder.embed_dense_batch(&texts, model_name).await?;
754            ensure_batch_len(vecs.len(), indices.len(), model_name)?;
755            for (idx, vec) in indices.into_iter().zip(vecs) {
756                let point = &mut upsert.points[idx];
757                add_point_vector(point, vector_name, VectorValue::Dense(vec))?;
758            }
759        }
760        EmbeddingSpec::Sparse {
761            model,
762            vector,
763            field,
764        } => {
765            let model_name = model.as_deref().unwrap_or("default");
766            let vector_name = vector.as_deref().unwrap_or(SPARSE_VECTOR_NAME);
767            check_and_insert_vector_name(seen_vectors, vector_name)?;
768
769            let targets = collect_text_targets(&upsert.points, field.as_deref());
770            validate_non_empty_targets(upsert, &targets, "SPARSE", field.as_deref())?;
771
772            for (idx, text) in targets {
773                let sparse_vec = embedder.embed_sparse(&text, model_name).await?;
774                add_point_vector(
775                    &mut upsert.points[idx],
776                    vector_name,
777                    VectorValue::Sparse {
778                        indices: sparse_vec.indices,
779                        values: sparse_vec.values,
780                    },
781                )?;
782            }
783        }
784        EmbeddingSpec::Hybrid {
785            dense_model,
786            dense_vector,
787            dense_field,
788            sparse_model,
789            sparse_vector,
790            sparse_field,
791        } => {
792            let d_model = dense_model.as_deref().unwrap_or("default");
793            let s_model = sparse_model.as_deref().unwrap_or("default");
794            let d_vec_name = dense_vector.as_deref().unwrap_or(DENSE_VECTOR_NAME);
795            let s_vec_name = sparse_vector.as_deref().unwrap_or(SPARSE_VECTOR_NAME);
796
797            check_and_insert_vector_name(seen_vectors, d_vec_name)?;
798            check_and_insert_vector_name(seen_vectors, s_vec_name)?;
799
800            let dense_targets = collect_text_targets(&upsert.points, dense_field.as_deref());
801            let sparse_targets = collect_text_targets(&upsert.points, sparse_field.as_deref());
802
803            validate_non_empty_targets(upsert, &dense_targets, "DENSE", dense_field.as_deref())?;
804            validate_non_empty_targets(upsert, &sparse_targets, "SPARSE", sparse_field.as_deref())?;
805
806            let (indices, texts): (Vec<usize>, Vec<String>) = dense_targets.into_iter().unzip();
807            let dense_vecs = embedder.embed_dense_batch(&texts, d_model).await?;
808            ensure_batch_len(dense_vecs.len(), indices.len(), d_model)?;
809            for (idx, d_vec) in indices.into_iter().zip(dense_vecs) {
810                let point = &mut upsert.points[idx];
811                add_point_vector(point, d_vec_name, VectorValue::Dense(d_vec))?;
812            }
813
814            for (idx, text) in sparse_targets {
815                let sparse_vec = embedder.embed_sparse(&text, s_model).await?;
816                let point = &mut upsert.points[idx];
817                add_point_vector(
818                    point,
819                    s_vec_name,
820                    VectorValue::Sparse {
821                        indices: sparse_vec.indices,
822                        values: sparse_vec.values,
823                    },
824                )?;
825            }
826        }
827        EmbeddingSpec::MultiVector {
828            model,
829            vector,
830            field,
831        } => {
832            let model_name = model.as_deref().unwrap_or("default");
833            let vector_name = vector.as_deref().unwrap_or("colbert");
834            check_and_insert_vector_name(seen_vectors, vector_name)?;
835
836            let targets = collect_text_targets(&upsert.points, field.as_deref());
837            validate_non_empty_targets(upsert, &targets, "MULTI", field.as_deref())?;
838
839            let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
840            let bags = embedder.embed_multi_batch(&texts, model_name).await?;
841            if bags.len() != indices.len() {
842                return Err(QqlError::execution(
843                    "QQL-EMBEDDING-MULTI",
844                    format!(
845                        "embed_multi_batch returned {} bags for {} texts (model={model_name})",
846                        bags.len(),
847                        indices.len()
848                    ),
849                    None,
850                ));
851            }
852            for (idx, rows) in indices.into_iter().zip(bags) {
853                if rows.is_empty() {
854                    return Err(QqlError::execution(
855                        "QQL-EMBEDDING-MULTI",
856                        "embed_multi returned an empty multivector",
857                        None,
858                    ));
859                }
860                add_point_vector(
861                    &mut upsert.points[idx],
862                    vector_name,
863                    VectorValue::MultiDense(rows),
864                )?;
865            }
866        }
867        EmbeddingSpec::Image {
868            model,
869            vector,
870            field,
871        } => {
872            let model_name = model.as_deref().unwrap_or("default");
873            let vector_name = vector.as_deref().unwrap_or("image");
874            check_and_insert_vector_name(seen_vectors, vector_name)?;
875
876            let targets = collect_image_targets(&upsert.points, field.as_deref());
877            validate_non_empty_targets(upsert, &targets, "IMAGE", field.as_deref())?;
878
879            let (indices, sources): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
880            let vecs = embedder.embed_image_batch(&sources, model_name).await?;
881            if vecs.len() != indices.len() {
882                return Err(QqlError::execution(
883                    "QQL-EMBEDDING-IMAGE",
884                    format!(
885                        "embed_image_batch returned {} vectors for {} sources (model={model_name})",
886                        vecs.len(),
887                        indices.len()
888                    ),
889                    None,
890                ));
891            }
892            for (idx, vec) in indices.into_iter().zip(vecs) {
893                add_point_vector(
894                    &mut upsert.points[idx],
895                    vector_name,
896                    VectorValue::Dense(vec),
897                )?;
898            }
899        }
900    }
901    Ok(())
902}
903
904fn validate_non_empty_targets(
905    upsert: &UpsertStmt,
906    targets: &[(usize, String)],
907    kind: &str,
908    field: Option<&str>,
909) -> Result<(), QqlError> {
910    if targets.is_empty() {
911        let actual_fields = upsert
912            .points
913            .first()
914            .map(|p| {
915                p.payload
916                    .iter()
917                    .map(|(k, _)| k.as_str())
918                    .collect::<Vec<_>>()
919                    .join(", ")
920            })
921            .unwrap_or_default();
922
923        let err_msg = if let Some(f) = field {
924            format!(
925                "USING {kind} MODEL specified with ON FIELD '{f}' but no matching text payload field found. Found fields: {actual_fields}"
926            )
927        } else {
928            format!(
929                "USING {kind} MODEL specified but no text payload field found. Expected one of: {}. Found fields: {actual_fields}",
930                DEFAULT_TEXT_FIELDS_ORDERED.join(", ")
931            )
932        };
933
934        return Err(QqlError::execution("QQL-EMBEDDING", err_msg, None));
935    }
936    Ok(())
937}
938
939fn check_and_insert_vector_name(
940    seen_vectors: &mut std::collections::HashSet<String>,
941    vector_name: &str,
942) -> Result<(), QqlError> {
943    if !seen_vectors.insert(vector_name.to_string()) {
944        return Err(QqlError::execution(
945            "QQL-EMBEDDING",
946            format!("duplicate target vector '{vector_name}' in multi-spec embedding clause"),
947            None,
948        ));
949    }
950    Ok(())
951}
952
953const DEFAULT_TEXT_FIELDS_ORDERED: &[&str] = &[
954    "text",
955    "body",
956    "content",
957    "title",
958    "description",
959    "name",
960    "summary",
961    "document",
962];
963
964const DEFAULT_IMAGE_FIELDS_ORDERED: &[&str] = &[
965    "image",
966    "image_path",
967    "image_url",
968    "photo",
969    "picture",
970    "img",
971    "path",
972    "url",
973];
974
975/// Collect image path/URL payload fields for IMAGE embedding specs.
976fn collect_image_targets(
977    points: &[UpsertPoint],
978    field_override: Option<&str>,
979) -> Vec<(usize, String)> {
980    if let Some(target_field) = field_override {
981        points
982            .iter()
983            .enumerate()
984            .filter_map(|(idx, point)| {
985                point.payload.iter().find_map(|(key, value)| {
986                    if key.eq_ignore_ascii_case(target_field) {
987                        if let qql_core::ast::Value::Str(source) = value {
988                            if !source.is_empty() {
989                                return Some((idx, source.clone()));
990                            }
991                        }
992                    }
993                    None
994                })
995            })
996            .collect()
997    } else {
998        points
999            .iter()
1000            .enumerate()
1001            .filter_map(|(idx, point)| {
1002                for &candidate in DEFAULT_IMAGE_FIELDS_ORDERED {
1003                    if let Some((_, qql_core::ast::Value::Str(source))) = point
1004                        .payload
1005                        .iter()
1006                        .find(|(key, _)| key.eq_ignore_ascii_case(candidate))
1007                    {
1008                        if !source.is_empty() {
1009                            return Some((idx, source.clone()));
1010                        }
1011                    }
1012                }
1013                None
1014            })
1015            .collect()
1016    }
1017}
1018
1019fn collect_text_targets(
1020    points: &[UpsertPoint],
1021    field_override: Option<&str>,
1022) -> Vec<(usize, String)> {
1023    if let Some(target_field) = field_override {
1024        points
1025            .iter()
1026            .enumerate()
1027            .filter_map(|(idx, point)| {
1028                point.payload.iter().find_map(|(key, value)| {
1029                    if key.eq_ignore_ascii_case(target_field) {
1030                        if let qql_core::ast::Value::Str(text) = value {
1031                            if !text.is_empty() {
1032                                return Some((idx, text.clone()));
1033                            }
1034                        }
1035                    }
1036                    None
1037                })
1038            })
1039            .collect()
1040    } else {
1041        collect_default_text_targets(points)
1042    }
1043}
1044
1045fn collect_default_text_targets(points: &[UpsertPoint]) -> Vec<(usize, String)> {
1046    points
1047        .iter()
1048        .enumerate()
1049        .filter_map(|(idx, point)| {
1050            for &candidate in DEFAULT_TEXT_FIELDS_ORDERED {
1051                if let Some((_, qql_core::ast::Value::Str(text))) = point
1052                    .payload
1053                    .iter()
1054                    .find(|(key, _)| key.eq_ignore_ascii_case(candidate))
1055                {
1056                    if !text.is_empty() {
1057                        return Some((idx, text.clone()));
1058                    }
1059                }
1060            }
1061            None
1062        })
1063        .collect()
1064}
1065
1066fn add_point_vector(
1067    point: &mut UpsertPoint,
1068    name: &str,
1069    vector: VectorValue,
1070) -> Result<(), QqlError> {
1071    if name.is_empty() {
1072        return match &mut point.vectors {
1073            Some(PointVectors::Unnamed(existing)) => {
1074                *existing = vector;
1075                Ok(())
1076            }
1077            Some(PointVectors::Named(list)) => {
1078                if let Some(existing) = list.iter_mut().find(|(key, _)| key.is_empty()) {
1079                    existing.1 = vector;
1080                } else {
1081                    list.push((String::new(), vector));
1082                }
1083                Ok(())
1084            }
1085            None => {
1086                point.vectors = Some(PointVectors::Unnamed(vector));
1087                Ok(())
1088            }
1089        };
1090    }
1091    match &mut point.vectors {
1092        Some(PointVectors::Named(list)) => {
1093            if let Some(existing) = list.iter_mut().find(|(k, _)| k == name) {
1094                existing.1 = vector;
1095            } else {
1096                list.push((name.to_string(), vector));
1097            }
1098            Ok(())
1099        }
1100        Some(PointVectors::Unnamed(_)) => Err(QqlError::execution(
1101            "QQL-EMBEDDING",
1102            format!(
1103                "cannot add named vector '{name}' to a point that already has an unnamed vector; \
1104                 provide an explicit named-vector topology or omit EMBED for this point"
1105            ),
1106            None,
1107        )),
1108        None => {
1109            point.vectors = Some(PointVectors::Named(vec![(name.to_string(), vector)]));
1110            Ok(())
1111        }
1112    }
1113}
qql_embed/resolve.rs

qql_embed/
resolve.rs
