Skip to main content

writerslogic_warm/core/engine/
multi_hop.rs

1// Copyright 2024-2026 WritersLogic Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::core::query_limits::{QueryBudget, QueryLimitExceeded, QueryLimits};
5use crate::core::rules::RuleStore;
6use crate::core::triple_store::TripleStore;
7
8use super::structural::{fuzzy_structural_query, MeaningContext};
9
10#[derive(Clone, Debug)]
11pub struct MultiHopResult {
12    pub entity_id: String,
13    pub confidence: f64,
14    pub method: MultiHopMethod,
15    pub hops: Vec<HopDetail>,
16}
17
18#[derive(Clone, Debug)]
19pub enum MultiHopMethod {
20    RuleRewrite { rule_name: String },
21    ChainedLookup,
22    SingleAlgebraic,
23}
24
25#[derive(Clone, Debug)]
26pub struct HopDetail {
27    pub from_entity: String,
28    pub relation: String,
29    pub to_entity: String,
30    pub confidence: f64,
31}
32
33pub fn multi_hop_query(
34    start_entity: &str,
35    relations: &[&str],
36    ctx: &MeaningContext<'_>,
37    rule_store: &RuleStore,
38    max_depth: usize,
39) -> Vec<MultiHopResult> {
40    if relations.len() > max_depth {
41        return Vec::new();
42    }
43    multi_hop_query_with_limits(
44        start_entity,
45        relations,
46        ctx,
47        rule_store,
48        max_depth,
49        QueryLimits::default(),
50    )
51    .expect("multi-hop query budget exhausted; use multi_hop_with_limits")
52}
53
54pub fn multi_hop_query_with_limits(
55    start_entity: &str,
56    relations: &[&str],
57    ctx: &MeaningContext<'_>,
58    rule_store: &RuleStore,
59    max_depth: usize,
60    limits: QueryLimits,
61) -> Result<Vec<MultiHopResult>, QueryLimitExceeded> {
62    if relations.is_empty() || relations.len() > max_depth {
63        return if relations.is_empty() {
64            Ok(Vec::new())
65        } else {
66            Err(QueryLimitExceeded {
67                resource: "hop depth",
68                limit: max_depth,
69            })
70        };
71    }
72    let mut budget = QueryBudget::new(limits);
73    if relations.len() == 1 {
74        preflight_algebraic(ctx, &mut budget)?;
75        return bound_algebraic(single_hop(start_entity, relations[0], ctx), &mut budget);
76    }
77
78    if relations.len() == 2 {
79        if let Some(rule) = rule_store.find_rule(relations[0], relations[1]) {
80            preflight_algebraic(ctx, &mut budget)?;
81            return bound_algebraic(
82                rule_rewrite(start_entity, &rule.output_relation, &rule.name, ctx),
83                &mut budget,
84            );
85        }
86    }
87
88    chained_lookup(start_entity, relations, ctx.triple_store, &mut budget)
89}
90
91impl super::WarmCore {
92    /// Runs the same multi-hop query with explicit work, frontier and byte limits. Budget
93    /// exhaustion is an error; distinct paths to an endpoint remain distinct results.
94    pub fn multi_hop_with_limits(
95        &self,
96        start: &str,
97        relations: &[&str],
98        limits: QueryLimits,
99    ) -> anyhow::Result<Vec<MultiHopResult>> {
100        let _transaction = self.mutation_gate.read();
101        let (Some(ctx), Some(rules)) = (self.meaning_ctx(), &self.rule_store) else {
102            return Ok(Vec::new());
103        };
104        Ok(multi_hop_query_with_limits(
105            start,
106            relations,
107            &ctx,
108            rules,
109            self.config.meaning.max_hop_depth,
110            limits,
111        )?)
112    }
113}
114
115fn preflight_algebraic(
116    ctx: &MeaningContext<'_>,
117    budget: &mut QueryBudget,
118) -> Result<(), QueryLimitExceeded> {
119    let atoms = ctx.atom_memory.inner().vectors_read();
120    let composites = ctx.composite_memory.inner().vectors_read();
121    budget.frontier(composites.len())?;
122    budget.work(
123        composites
124            .len()
125            .saturating_mul(atoms.len().saturating_add(1))
126            .saturating_mul(ctx.max_iter.saturating_add(1)),
127    )?;
128    for (id, vector) in atoms.iter().chain(composites.iter()) {
129        budget.reserve(
130            id.len()
131                .saturating_mul(2)
132                .saturating_add(vector.indices().len().saturating_mul(64))
133                .saturating_add(128),
134        )?;
135    }
136    Ok(())
137}
138
139fn bound_algebraic(
140    results: Vec<MultiHopResult>,
141    budget: &mut QueryBudget,
142) -> Result<Vec<MultiHopResult>, QueryLimitExceeded> {
143    budget.frontier(results.len())?;
144    for result in &results {
145        budget.reserve(path_bytes(&result.entity_id, &result.hops))?;
146    }
147    Ok(results)
148}
149
150fn single_hop(start_entity: &str, relation: &str, ctx: &MeaningContext<'_>) -> Vec<MultiHopResult> {
151    algebraic_hop(start_entity, relation, ctx, || {
152        MultiHopMethod::SingleAlgebraic
153    })
154}
155
156fn rule_rewrite(
157    start_entity: &str,
158    output_relation: &str,
159    rule_name: &str,
160    ctx: &MeaningContext<'_>,
161) -> Vec<MultiHopResult> {
162    algebraic_hop(start_entity, output_relation, ctx, || {
163        MultiHopMethod::RuleRewrite {
164            rule_name: rule_name.to_string(),
165        }
166    })
167}
168
169fn algebraic_hop(
170    start_entity: &str,
171    relation: &str,
172    ctx: &MeaningContext<'_>,
173    method: impl Fn() -> MultiHopMethod,
174) -> Vec<MultiHopResult> {
175    let Some(s_vec) = ctx.atom_memory.get(start_entity) else {
176        return Vec::new();
177    };
178    let Some(r_vec) = ctx.atom_memory.get(relation) else {
179        return Vec::new();
180    };
181    fuzzy_structural_query(ctx, &[("subject", &s_vec), ("relation", &r_vec)], "object")
182        .into_iter()
183        .map(|r| MultiHopResult {
184            entity_id: r.entity_id.clone(),
185            confidence: r.confidence,
186            method: method(),
187            hops: vec![HopDetail {
188                from_entity: start_entity.to_string(),
189                relation: relation.to_string(),
190                to_entity: r.entity_id,
191                confidence: r.confidence,
192            }],
193        })
194        .collect()
195}
196
197/// Evidence-weighted confidence for a single hop.
198/// When an entity has N matches for a relation, each match gets confidence
199/// 1/sqrt(N) — a unique match is high confidence, one of many is low.
200/// Base decay of 0.9 per hop is applied on top.
201const HOP_DECAY: f64 = 0.9;
202
203fn chained_lookup(
204    start_entity: &str,
205    relations: &[&str],
206    triple_store: &TripleStore,
207    budget: &mut QueryBudget,
208) -> Result<Vec<MultiHopResult>, QueryLimitExceeded> {
209    let bytes = path_bytes(start_entity, &[]);
210    budget.frontier(1)?;
211    budget.reserve(bytes)?;
212    let mut current = vec![(start_entity.to_owned(), Vec::<HopDetail>::new(), bytes)];
213
214    for &relation in relations {
215        let mut next = Vec::new();
216        for (entity, hops, _) in &current {
217            let mut fan_out = 0usize;
218            triple_store.visit_bounded(Some(entity), Some(relation), None, budget, |_, _| {
219                fan_out += 1;
220                Ok(())
221            })?;
222            if fan_out == 0 {
223                continue;
224            }
225            let evidence_weight = 1.0 / (fan_out as f64).sqrt();
226            let hop_confidence = HOP_DECAY * evidence_weight.min(1.0);
227            triple_store.visit_bounded(
228                Some(entity),
229                Some(relation),
230                None,
231                budget,
232                |t, budget| {
233                    budget.work(hops.len().saturating_add(1))?;
234                    budget.frontier(next.len().saturating_add(1))?;
235                    let bytes = path_bytes(&t.object_id, hops).saturating_add(hop_bytes(
236                        entity,
237                        relation,
238                        &t.object_id,
239                    ));
240                    budget.reserve(bytes)?;
241                    let mut path = hops.clone();
242                    path.push(HopDetail {
243                        from_entity: entity.clone(),
244                        relation: relation.to_owned(),
245                        to_entity: t.object_id.clone(),
246                        confidence: hop_confidence,
247                    });
248                    next.push((t.object_id.clone(), path, bytes));
249                    Ok(())
250                },
251            )?;
252        }
253        if next.is_empty() {
254            return Ok(Vec::new());
255        }
256        for (_, _, bytes) in &current {
257            budget.release(*bytes);
258        }
259        current = next;
260    }
261
262    let mut results: Vec<MultiHopResult> = current
263        .into_iter()
264        .map(|(entity, hops, _)| {
265            let chain_confidence = hops.iter().map(|h| h.confidence).product();
266            MultiHopResult {
267                entity_id: entity,
268                confidence: chain_confidence,
269                method: MultiHopMethod::ChainedLookup,
270                hops,
271            }
272        })
273        .collect();
274
275    results.sort_by(|a, b| {
276        b.confidence
277            .partial_cmp(&a.confidence)
278            .unwrap_or(std::cmp::Ordering::Equal)
279    });
280    Ok(results)
281}
282
283fn hop_bytes(from: &str, relation: &str, to: &str) -> usize {
284    2 * std::mem::size_of::<HopDetail>()
285        + from
286            .len()
287            .saturating_add(relation.len())
288            .saturating_add(to.len())
289            .saturating_mul(2)
290}
291
292fn path_bytes(entity: &str, hops: &[HopDetail]) -> usize {
293    hops.iter().fold(
294        entity
295            .len()
296            .saturating_mul(2)
297            .saturating_add(2 * std::mem::size_of::<MultiHopResult>()),
298        |bytes, hop| {
299            bytes.saturating_add(hop_bytes(&hop.from_entity, &hop.relation, &hop.to_entity))
300        },
301    )
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::core::engine::structural::tests::Fixture;
308    use crate::core::rules::CompositionRule;
309
310    #[test]
311    fn bounded_chains_preserve_distinct_paths_at_work_and_frontier_boundaries() {
312        let fixture = Fixture::new();
313        for (from, to, id) in [
314            ("a", "b", "ab"),
315            ("a", "c", "ac"),
316            ("b", "d", "bd"),
317            ("c", "d", "cd"),
318        ] {
319            fixture.triples.add(from, "r", to, id);
320        }
321        let rules = RuleStore::new();
322        for work in [13, 14, 15] {
323            let result = multi_hop_query_with_limits(
324                "a",
325                &["r", "r"],
326                &fixture.context(),
327                &rules,
328                10,
329                QueryLimits {
330                    max_work: work,
331                    ..Default::default()
332                },
333            );
334            if work == 13 {
335                assert_eq!(result.unwrap_err().resource, "work");
336            } else {
337                let paths = result.unwrap();
338                assert_eq!(paths.len(), 2);
339                assert!(paths
340                    .iter()
341                    .all(|path| path.entity_id == "d" && path.hops.len() == 2));
342                assert_ne!(paths[0].hops[0].to_entity, paths[1].hops[0].to_entity);
343            }
344        }
345        for frontier in [1, 2, 3] {
346            let result = multi_hop_query_with_limits(
347                "a",
348                &["r", "r"],
349                &fixture.context(),
350                &rules,
351                10,
352                QueryLimits {
353                    max_frontier: frontier,
354                    ..Default::default()
355                },
356            );
357            if frontier == 1 {
358                assert_eq!(result.unwrap_err().resource, "frontier");
359            } else {
360                assert_eq!(result.unwrap().len(), 2);
361            }
362        }
363        let error = multi_hop_query_with_limits(
364            "a",
365            &["r", "r"],
366            &fixture.context(),
367            &rules,
368            10,
369            QueryLimits {
370                max_bytes: 1,
371                ..Default::default()
372            },
373        )
374        .unwrap_err();
375        assert_eq!(error.resource, "bytes");
376        fixture.triples.add("a", "unrelated", "z", "az");
377        let error = multi_hop_query_with_limits(
378            "a",
379            &["r", "r"],
380            &fixture.context(),
381            &rules,
382            10,
383            QueryLimits {
384                max_work: 14,
385                ..Default::default()
386            },
387        )
388        .unwrap_err();
389        assert_eq!(error.resource, "work");
390    }
391
392    #[test]
393    fn cyclic_branching_queries_fail_explicitly_before_frontier_growth() {
394        let fixture = Fixture::new();
395        for from in ["a", "b", "c"] {
396            for to in ["a", "b", "c"] {
397                fixture.triples.add(from, "r", to, &format!("{from}{to}"));
398            }
399        }
400        let error = multi_hop_query_with_limits(
401            "a",
402            &["r"; 10],
403            &fixture.context(),
404            &RuleStore::new(),
405            10,
406            QueryLimits {
407                max_frontier: 8,
408                ..Default::default()
409            },
410        )
411        .unwrap_err();
412        assert_eq!(error.resource, "frontier");
413    }
414
415    #[test]
416    fn querying_an_unknown_relation_does_not_create_an_atom() {
417        let fixture = Fixture::new();
418        fixture.atoms.get_or_insert("a");
419        let count = fixture.atoms.count();
420        assert!(
421            multi_hop_query("a", &["unknown"], &fixture.context(), &RuleStore::new(), 10)
422                .is_empty()
423        );
424        assert_eq!(fixture.atoms.count(), count);
425    }
426
427    #[test]
428    fn chained_lookup_follows_each_hop() {
429        let fixture = Fixture::new();
430        for atom in ["john", "mark", "bob", "father"] {
431            fixture.atoms.get_or_insert(atom);
432        }
433        fixture.triples.add("john", "father", "mark", "c1");
434        fixture.triples.add("mark", "father", "bob", "c2");
435
436        let results = multi_hop_query(
437            "john",
438            &["father", "father"],
439            &fixture.context(),
440            &RuleStore::new(),
441            10,
442        );
443
444        assert_eq!(results[0].entity_id, "bob");
445        assert!(matches!(results[0].method, MultiHopMethod::ChainedLookup));
446        assert_eq!(results[0].hops.len(), 2);
447    }
448
449    #[test]
450    fn composition_rule_rewrites_two_hops_into_one_relation() {
451        let fixture = Fixture::new();
452        let rule_store = RuleStore::new();
453        rule_store.add_rule(CompositionRule {
454            name: "grandfather_rule".to_string(),
455            input_relations: vec!["father".to_string(), "father".to_string()],
456            output_relation: "grandfather".to_string(),
457        });
458        fixture.bind("john", "grandfather", "bob", "grandfather_triple");
459
460        let results = multi_hop_query(
461            "john",
462            &["father", "father"],
463            &fixture.context(),
464            &rule_store,
465            10,
466        );
467
468        assert_eq!(results[0].entity_id, "bob");
469        assert!(matches!(
470            results[0].method,
471            MultiHopMethod::RuleRewrite { .. }
472        ));
473    }
474
475    #[test]
476    fn fan_out_at_a_hop_lowers_chain_confidence() {
477        let fixture = Fixture::new();
478        let rule_store = RuleStore::new();
479        for atom in ["a", "b", "c", "d"] {
480            fixture.atoms.get_or_insert(atom);
481        }
482        fixture.triples.add("a", "r", "b", "c1");
483        fixture.triples.add("b", "r", "c", "c2");
484        let unique = multi_hop_query("a", &["r", "r"], &fixture.context(), &rule_store, 10);
485        let unique_confidence = unique[0].confidence;
486
487        fixture.triples.add("a", "r", "d", "c3");
488        let ambiguous = multi_hop_query("a", &["r", "r"], &fixture.context(), &rule_store, 10);
489        let through_b = ambiguous
490            .iter()
491            .find(|r| r.entity_id == "c")
492            .expect("the path through b still reaches c");
493        assert!(
494            through_b.confidence < unique_confidence,
495            "ambiguous {:.4} >= unique {unique_confidence:.4}",
496            through_b.confidence
497        );
498    }
499
500    #[test]
501    fn results_rank_by_confidence() {
502        let fixture = Fixture::new();
503        fixture.atoms.get_or_insert("start");
504        // start -> a -> end is unique at both hops; b also reaches end, among three objects.
505        for (subject, relation, object, id) in [
506            ("start", "r1", "a", "c1"),
507            ("start", "r1", "b", "c2"),
508            ("a", "r2", "end", "c3"),
509            ("b", "r2", "end", "c4"),
510            ("b", "r2", "noise1", "c5"),
511            ("b", "r2", "noise2", "c6"),
512        ] {
513            fixture.triples.add(subject, relation, object, id);
514        }
515
516        let results = multi_hop_query(
517            "start",
518            &["r1", "r2"],
519            &fixture.context(),
520            &RuleStore::new(),
521            10,
522        );
523
524        assert!(results.iter().any(|r| r.entity_id == "end"));
525        assert!(results
526            .windows(2)
527            .all(|w| w[0].confidence >= w[1].confidence));
528    }
529}