1use std::collections::HashMap;
13
14use omgbase_search::{EmbeddingProvider, f32_to_blob};
15use omgbase_store::Store;
16use oqx::ast::{Expr, Query, SelectItem};
17use oqx::walk::{Clause, Node, VisitContext, Visitor, transform, visit};
18use oqx::{Consumer, Engine, InMemoryEngine, Value, build, resolve_aliases};
19use serde_json::{Map, Value as Json};
20
21use crate::context::{SemanticVec, StoreContext, render_row_values};
22use crate::cursor::{decode_cursor, encode_cursor};
23use crate::error::{Result, SurfaceError};
24use crate::planner::SqlitePlanner;
25
26pub const DEFAULT_LIMIT: usize = 50;
28
29const ROW_FNS: [&str; 12] = [
32 "text",
33 "semantic",
34 "under",
35 "under_heading",
36 "within",
37 "under_kind",
38 "yaml_path",
39 "json_pointer",
40 "has_edge",
41 "has_anchor",
42 "child_count",
43 "parent_type",
44];
45
46const ID_KEY: &str = "__oqx_id";
47const PATH_KEY: &str = "__oqx_path";
48const VALUE_KEY: &str = "__oqx_value";
52
53#[derive(Clone, Copy, Default)]
55pub struct QueryOptions<'a> {
56 pub limit: Option<usize>,
58 pub cursor: Option<&'a str>,
60 pub provider: Option<&'a dyn EmbeddingProvider>,
63 pub in_memory: bool,
66}
67
68#[derive(Clone, Debug, PartialEq)]
70pub struct OqxResult {
71 pub hits: Vec<Json>,
72 pub truncated: bool,
73 pub cursor: Option<String>,
74 pub consumer: Consumer,
75 pub count: Option<f64>,
76 pub exists: Option<bool>,
77 pub none: Option<bool>,
78 pub values: Option<Vec<Json>>,
80}
81
82impl OqxResult {
83 fn scalar(consumer: Consumer) -> Self {
84 Self {
85 hits: Vec::new(),
86 truncated: false,
87 cursor: None,
88 consumer,
89 count: None,
90 exists: None,
91 none: None,
92 values: None,
93 }
94 }
95
96 #[must_use]
99 pub fn to_json(&self) -> Json {
100 let mut m = Map::new();
101 m.insert("hits".to_owned(), Json::Array(self.hits.clone()));
102 m.insert("truncated".to_owned(), Json::Bool(self.truncated));
103 m.insert(
104 "cursor".to_owned(),
105 self.cursor.clone().map_or(Json::Null, Json::String),
106 );
107 m.insert(
108 "consumer".to_owned(),
109 Json::String(self.consumer.as_str().to_owned()),
110 );
111 if let Some(n) = self.count {
112 m.insert("count".to_owned(), Value::Number(n).to_canonical_json());
113 }
114 if let Some(b) = self.exists {
115 m.insert("exists".to_owned(), Json::Bool(b));
116 }
117 if let Some(b) = self.none {
118 m.insert("none".to_owned(), Json::Bool(b));
119 }
120 if let Some(v) = &self.values {
121 m.insert("values".to_owned(), Json::Array(v.clone()));
122 }
123 Json::Object(m)
124 }
125}
126
127#[must_use]
132pub fn rewrite_query(q: &Query) -> Query {
133 transform(q, &mut |e, _| match e {
134 Expr::Call {
135 recv: None,
136 name,
137 args,
138 span,
139 } if ROW_FNS.contains(&name.as_str()) => Expr::Call {
140 recv: Some(Box::new(build::ident("$self"))),
141 name,
142 args,
143 span,
144 },
145 other => other,
146 })
147}
148
149struct Phrases(Vec<String>);
153
154impl Visitor for Phrases {
155 fn enter(&mut self, node: Node<'_>, _ctx: &VisitContext<'_>) -> bool {
156 if let Node::Expr(Expr::Call {
157 recv: None,
158 name,
159 args,
160 ..
161 }) = node
162 && name == "semantic"
163 && let Some(Expr::Lit {
164 value: Value::Str(s),
165 ..
166 }) = args.first()
167 && !self.0.contains(s)
168 {
169 self.0.push(s.clone());
170 }
171 true
172 }
173}
174
175#[must_use]
178pub fn collect_semantic_phrases(source: &str) -> Vec<String> {
179 let Ok(q) = oqx::parse_string(source) else {
180 return Vec::new();
181 };
182 let mut phrases = Phrases(Vec::new());
183 visit(Node::Query(&q), &mut phrases);
184 phrases.0
185}
186
187fn mentions_outside_source(q: &Query, name: &str) -> bool {
193 struct Mentions<'n> {
194 name: &'n str,
195 found: bool,
196 }
197 impl Visitor for Mentions<'_> {
198 fn enter(&mut self, node: Node<'_>, ctx: &VisitContext<'_>) -> bool {
199 if self.found {
200 return false;
201 }
202 if ctx.clause == Some(Clause::Source) && ctx.path.len() == 1 {
204 return false;
205 }
206 if let Node::Expr(Expr::Ident { name, .. } | Expr::Outer { name, .. }) = node
207 && name == self.name
208 {
209 self.found = true;
210 }
211 !self.found
212 }
213 }
214 let mut m = Mentions { name, found: false };
215 visit(Node::Query(q), &mut m);
216 m.found
217}
218
219fn to_hit(row: Value) -> Value {
227 let Value::Object(o) = render_row_values(row) else {
228 return Value::Object(oqx::Object::new());
229 };
230 let mut id = Value::Undefined;
231 let mut path = Value::Undefined;
232 let mut rest = Vec::new();
233 for (k, v) in o {
234 match k.as_str() {
235 ID_KEY => id = v,
236 PATH_KEY => path = v,
237 _ => rest.push((k, v)),
238 }
239 }
240 let mut hit = oqx::Object::with_capacity(rest.len() + 2);
241 hit.insert("id", Value::Str(id.to_string()));
242 hit.insert(
243 "path",
244 Value::Str(if path.is_absent() {
245 String::new()
246 } else {
247 path.to_string()
248 }),
249 );
250 for (k, v) in rest {
251 hit.insert(k, v);
252 }
253 Value::Object(hit)
254}
255
256fn hit_str(hit: &Value, key: &str) -> String {
257 hit.as_object()
258 .and_then(|o| o.get(key))
259 .map(|v| v.to_string())
260 .unwrap_or_default()
261}
262
263fn dedup_hits_by_projection(hits: Vec<Value>) -> Vec<Value> {
267 let mut seen: Vec<String> = Vec::new();
268 let mut out = Vec::new();
269 for h in hits {
270 let mut pairs: Vec<(String, Value)> = h
271 .as_object()
272 .map(|o| {
273 o.iter()
274 .filter(|(k, _)| *k != "id" && *k != "path")
275 .map(|(k, v)| (k.to_owned(), v.clone()))
276 .collect()
277 })
278 .unwrap_or_default();
279 pairs.sort_by(|a, b| a.0.cmp(&b.0));
280 let key = Value::Array(
281 pairs
282 .into_iter()
283 .map(|(k, v)| Value::Array(vec![Value::Str(k), v]))
284 .collect(),
285 )
286 .to_canonical_json()
287 .to_string();
288 if seen.contains(&key) {
289 continue;
290 }
291 seen.push(key);
292 out.push(h);
293 }
294 out
295}
296
297fn const_bound(e: Option<&Expr>, word: &str) -> Result<Option<usize>> {
300 match e {
301 None => Ok(None),
302 Some(Expr::Lit {
303 value: Value::Number(n),
304 ..
305 }) if n.fract() == 0.0 && *n >= 0.0 && n.is_finite() => Ok(Some(*n as usize)),
306 Some(_) => Err(SurfaceError::filter_invalid(
307 format!("top-level {word} must be a non-negative integer literal"),
308 "OQX",
309 )),
310 }
311}
312
313fn value_of(hit: &Value) -> Value {
314 hit.as_object()
315 .and_then(|o| o.get(VALUE_KEY))
316 .cloned()
317 .unwrap_or(Value::Undefined)
318}
319
320fn without_value_key(hit: Value) -> Json {
321 hit.to_canonical_json()
322}
323
324pub fn query(
328 store: &Store,
329 repo_id: &str,
330 source: &str,
331 opts: QueryOptions<'_>,
332) -> Result<OqxResult> {
333 let phrases = collect_semantic_phrases(source);
338 let mut semantic: HashMap<String, SemanticVec> = HashMap::new();
339 if let Some(provider) = opts.provider.filter(|_| !phrases.is_empty()) {
340 for phrase in phrases {
341 let vec = provider
342 .embed_query(&phrase)
343 .map_err(|e| SurfaceError::new(e.code(), e.to_string()))?;
344 semantic.insert(
345 phrase,
346 SemanticVec {
347 model: provider.model().to_owned(),
348 vec: f32_to_blob(&vec),
349 },
350 );
351 }
352 }
353 let runner = Runner {
354 store,
355 repo_id,
356 semantic,
357 planned: !opts.in_memory,
358 };
359 run_inner(&runner, source, opts)
360}
361
362struct Runner<'a> {
368 store: &'a Store,
369 repo_id: &'a str,
370 semantic: HashMap<String, SemanticVec>,
371 planned: bool,
372}
373
374impl Runner<'_> {
375 fn run(&self, q: &Query) -> Result<oqx::OqxResult> {
382 let conn = self.store.conn();
383 let ctx = StoreContext::new(conn, self.repo_id, self.semantic.clone());
384 let plan = if self.planned {
385 SqlitePlanner::new(conn, self.repo_id)
386 .try_plan(q, &[])
387 .map_err(|e| SurfaceError::other(format!("sqlite: {e}")))?
388 } else {
389 None
390 };
391 let (ctx, residual) = match plan {
392 Some(plan) => {
397 let once = !mentions_outside_source(&plan.residual, oqx::ROWS_ROOT);
398 let ctx = if once {
399 ctx.with_rows_root_once(plan.rows)
400 } else {
401 ctx.with_rows_root(plan.rows)
402 };
403 (ctx, Some(plan.residual))
404 }
405 None => (ctx, None),
406 };
407 let engine = InMemoryEngine::new(ctx);
408 let out = engine.run(residual.as_ref().unwrap_or(q), &[]);
409 if let Some(failed) = engine.context().take_root_failure() {
412 return Err(failed.into());
413 }
414 Ok(out?)
417 }
418}
419
420fn run_inner(engine: &Runner<'_>, source: &str, opts: QueryOptions<'_>) -> Result<OqxResult> {
421 let parsed = rewrite_query(&resolve_aliases(&oqx::parse_string(source)?)?);
424 let consumer = parsed.consumer;
425
426 match consumer {
427 Consumer::Exists => {
428 let res = engine.run(&parsed)?;
429 let mut r = OqxResult::scalar(consumer);
430 r.exists = Some(matches!(res, oqx::OqxResult::Exists(true)));
431 return Ok(r);
432 }
433 Consumer::Count => {
434 let res = engine.run(&parsed)?;
435 let mut r = OqxResult::scalar(consumer);
436 r.count = Some(match res {
437 oqx::OqxResult::Count(n) => n,
438 _ => 0.0,
439 });
440 return Ok(r);
441 }
442 Consumer::None => {
443 let res = engine.run(&parsed)?;
444 let mut r = OqxResult::scalar(consumer);
445 r.none = Some(match res {
446 oqx::OqxResult::None(b) => b,
447 _ => true,
448 });
449 return Ok(r);
450 }
451 Consumer::Collect | Consumer::First | Consumer::Single => {}
452 }
453
454 let top_distinct = parsed.distinct;
460 let top_values = parsed.values;
461 let user_select: Vec<SelectItem> = if top_values {
462 parsed
463 .select
464 .first()
465 .map(|it| match it {
466 SelectItem::Field {
467 expr, lift, span, ..
468 } => SelectItem::Field {
469 name: VALUE_KEY.to_owned(),
470 expr: expr.clone(),
471 lift: *lift,
472 span: *span,
473 },
474 SelectItem::Collect { op, span, .. } => SelectItem::Collect {
475 name: VALUE_KEY.to_owned(),
476 op: op.clone(),
477 span: *span,
478 },
479 })
480 .into_iter()
481 .collect()
482 } else {
483 parsed.select.clone()
484 };
485 let id_item = build::field(ID_KEY, build::ident("$id"));
486 let path_item = build::field(PATH_KEY, build::ident("$path"));
487 let mut select = vec![id_item, path_item];
488 select.extend(user_select);
489 let (top_limit, top_offset) = (parsed.limit.clone(), parsed.offset.clone());
493 let q = Query {
494 distinct: false,
495 values: false,
496 select,
497 limit: if consumer == Consumer::Collect {
498 None
499 } else {
500 parsed.limit.clone()
501 },
502 offset: if consumer == Consumer::Collect {
503 None
504 } else {
505 parsed.offset.clone()
506 },
507 ..parsed.clone()
508 };
509 let res = engine.run(&q)?;
510
511 if matches!(consumer, Consumer::First | Consumer::Single) {
512 let row = match res {
513 oqx::OqxResult::First(r) | oqx::OqxResult::Single(r) => r,
514 _ => None,
515 };
516 let mut out = OqxResult::scalar(consumer);
517 match row {
518 None => {
519 if top_values {
520 out.values = Some(Vec::new());
521 }
522 }
523 Some(r) => {
524 let hit = to_hit(r);
525 if top_values {
526 out.values = Some(vec![value_of(&hit).to_canonical_json()]);
527 } else {
528 out.hits = vec![without_value_key(hit)];
529 }
530 }
531 }
532 return Ok(out);
533 }
534
535 let mut rows: Vec<Value> = match res {
537 oqx::OqxResult::Collect(rows) => rows,
538 _ => Vec::new(),
539 };
540 let custom = parsed.order_by.as_ref().is_some_and(|o| !o.is_empty());
541 let cursor = opts.cursor.filter(|c| !c.is_empty() && !custom);
542 let eager = top_distinct || cursor.is_some();
547 if eager {
548 rows = rows.into_iter().map(to_hit).collect();
549 }
550 if top_distinct {
551 rows = dedup_hits_by_projection(rows);
552 }
553 let offset = const_bound(top_offset.as_ref(), "offset")?.unwrap_or(0);
554 let limit = const_bound(top_limit.as_ref(), "limit")?;
555 if offset > 0 || limit.is_some() {
556 let end = limit.map_or(rows.len(), |l| (offset + l).min(rows.len()));
557 if offset >= rows.len() {
558 rows.clear();
559 } else {
560 rows.truncate(end);
561 rows.drain(..offset);
562 }
563 }
564 let cap = opts.limit.unwrap_or(DEFAULT_LIMIT);
565 let mut page = rows;
566 if let Some(cursor) = cursor {
567 let parts = decode_cursor(cursor, "query", 2)?;
568 let (path, id) = (&parts[0], &parts[1]);
569 page.retain(|h| {
570 let hp = hit_str(h, "path");
571 let hi = hit_str(h, "id");
572 hp > *path || (hp == *path && hi > *id)
573 });
574 }
575 let truncated = page.len() > cap;
576 page.truncate(cap);
577 if !eager {
578 page = page.into_iter().map(to_hit).collect();
579 }
580 let cursor = if truncated && !custom {
581 page.last()
582 .map(|last| encode_cursor(&[&hit_str(last, "path"), &hit_str(last, "id")]))
583 } else {
584 None
585 };
586 let mut out = OqxResult::scalar(Consumer::Collect);
587 out.truncated = truncated;
588 out.cursor = cursor;
589 if top_values {
590 out.values = Some(
591 page.iter()
592 .map(|h| value_of(h).to_canonical_json())
593 .collect(),
594 );
595 } else {
596 out.hits = page.into_iter().map(without_value_key).collect();
597 }
598 Ok(out)
599}
600
601#[cfg(test)]
602mod tests {
603 use super::*;
604 use oqx::ast::Where;
605
606 #[test]
607 fn row_functions_become_self_methods() {
608 let q = oqx::parse_string(
609 "from blocks where text(\"x\") && under_heading(\"h\") && size(attrs) > 0 && doc.$path.startsWith(\"a\")",
610 )
611 .unwrap();
612 let r = rewrite_query(&q);
613 let Some(Where::And { parts, .. }) = &r.r#where else {
614 panic!("and")
615 };
616 let Where::Scalar { expr, .. } = &parts[0] else {
617 panic!("scalar")
618 };
619 assert!(
620 matches!(expr, Expr::Call { recv: Some(r), name, .. } if name == "text" && **r == build::ident("$self"))
621 );
622 let Where::Scalar { expr, .. } = &parts[2] else {
623 panic!("scalar")
624 };
625 assert!(
626 matches!(expr, Expr::Binary { left, .. } if matches!(&**left, Expr::Call { recv: None, name, .. } if name == "size"))
627 );
628 }
629
630 #[test]
631 fn semantic_phrases_are_collected_distinct() {
632 let phrases = collect_semantic_phrases(
633 "select s: semantic(\"alpha\") from docs where semantic(\"alpha\") > 0.5 || nodes exists { where semantic(\"beta\") > 0 } order by semantic(\"gamma\") desc",
634 );
635 assert_eq!(phrases, ["alpha", "beta", "gamma"]);
636 assert!(collect_semantic_phrases("not a query {{").is_empty());
637 assert!(collect_semantic_phrases("from docs").is_empty());
638 }
639
640 #[test]
641 fn hits_peel_the_injected_columns() {
642 let mut o = oqx::Object::new();
643 o.insert(ID_KEY, Value::Str("d_1".into()));
644 o.insert(PATH_KEY, Value::Null);
645 o.insert("layer", Value::Str("canon".into()));
646 let hit = to_hit(Value::Object(o));
647 let ho = hit.as_object().unwrap();
648 assert_eq!(ho.keys().collect::<Vec<_>>(), ["id", "path", "layer"]);
649 assert_eq!(ho.get("path"), Some(&Value::Str(String::new())));
650 let mut o = oqx::Object::new();
652 o.insert(ID_KEY, Value::Str("d_1".into()));
653 o.insert(PATH_KEY, Value::Str("a.md".into()));
654 o.insert("id", Value::Number(7.0));
655 let hit = to_hit(Value::Object(o));
656 let ho = hit.as_object().unwrap();
657 assert_eq!(ho.keys().collect::<Vec<_>>(), ["id", "path"]);
658 assert_eq!(ho.get("id"), Some(&Value::Number(7.0)));
659 }
660
661 #[test]
662 fn distinct_dedups_by_user_projection_first_wins() {
663 let mk = |id: &str, t: &str| {
664 let mut o = oqx::Object::new();
665 o.insert("id", Value::Str(id.into()));
666 o.insert("path", Value::Str("p".into()));
667 o.insert("type", Value::Str(t.into()));
668 Value::Object(o)
669 };
670 let out = dedup_hits_by_projection(vec![mk("1", "a"), mk("2", "b"), mk("3", "a")]);
671 assert_eq!(out.len(), 2);
672 assert_eq!(hit_str(&out[0], "id"), "1");
673 assert_eq!(hit_str(&out[1], "id"), "2");
674 }
675
676 #[test]
677 fn top_level_bounds_must_be_literals() {
678 assert_eq!(const_bound(None, "limit").unwrap(), None);
679 assert_eq!(
680 const_bound(
681 Some(&Expr::Lit {
682 value: Value::Number(3.0),
683 span: oqx::Span::EMPTY
684 }),
685 "limit"
686 )
687 .unwrap(),
688 Some(3)
689 );
690 let e = const_bound(
691 Some(&Expr::Lit {
692 value: Value::Number(-1.0),
693 span: oqx::Span::EMPTY,
694 }),
695 "offset",
696 )
697 .unwrap_err();
698 assert_eq!(e.code, "filter_invalid");
699 assert!(e.message.contains("top-level offset"));
700 }
701}