1use std::cmp::Ordering;
10use std::collections::BTreeMap;
11
12use crate::engine::stable_id;
13use crate::event_log::EventLog;
14use crate::link_store::{LinkStore, LinkStoreError};
15use crate::links_format::format_lino_record;
16use crate::memory::MemoryEvent;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ProbabilityModel {
21 BayesianEvidence,
24 MarkovTransition,
27}
28
29impl ProbabilityModel {
30 #[must_use]
31 pub const fn slug(self) -> &'static str {
32 match self {
33 Self::BayesianEvidence => "bayesian_evidence",
34 Self::MarkovTransition => "markov_transition",
35 }
36 }
37}
38
39#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct ProbabilitySourceProvenance {
42 pub source_url: String,
43 pub fetched_at: String,
44 pub sha256: String,
45 pub cached: bool,
46}
47
48impl ProbabilitySourceProvenance {
49 #[must_use]
50 pub fn trace_payload(&self) -> String {
51 format!(
52 "{} fetched_at={} sha256={} cached={}",
53 self.source_url, self.fetched_at, self.sha256, self.cached
54 )
55 }
56}
57
58#[derive(Debug, Clone, PartialEq)]
60pub struct ProbabilityEvidence {
61 pub id: String,
62 pub target: String,
63 pub observation: String,
64 pub weight: f32,
65 pub model: ProbabilityModel,
66 pub provenance: String,
67 pub recorded_at: String,
68 pub source: Option<ProbabilitySourceProvenance>,
69 pub transition_from: Option<String>,
70}
71
72impl ProbabilityEvidence {
73 #[must_use]
74 pub fn symbolic(
75 target: impl Into<String>,
76 observation: impl Into<String>,
77 weight: f32,
78 provenance: impl Into<String>,
79 recorded_at: impl Into<String>,
80 ) -> Self {
81 let mut evidence = Self {
82 id: String::new(),
83 target: target.into(),
84 observation: observation.into(),
85 weight: finite_or_zero(weight),
86 model: ProbabilityModel::BayesianEvidence,
87 provenance: provenance.into(),
88 recorded_at: recorded_at.into(),
89 source: None,
90 transition_from: None,
91 };
92 evidence.id = evidence.stable_record_id();
93 evidence
94 }
95
96 #[must_use]
97 pub fn with_model(mut self, model: ProbabilityModel) -> Self {
98 self.model = model;
99 self.id = self.stable_record_id();
100 self
101 }
102
103 #[must_use]
104 pub fn with_source(mut self, source: ProbabilitySourceProvenance) -> Self {
105 self.source = Some(source);
106 self.id = self.stable_record_id();
107 self
108 }
109
110 #[must_use]
111 pub fn with_transition_from(mut self, transition_from: impl Into<String>) -> Self {
112 self.transition_from = Some(transition_from.into());
113 self.id = self.stable_record_id();
114 self
115 }
116
117 #[must_use]
118 pub fn trace_payload(&self) -> String {
119 let mut parts = vec![
120 format!("id={}", self.id),
121 format!("target={}", self.target),
122 format!("model={}", self.model.slug()),
123 format!("observation={}", self.observation),
124 format!("weight={:.6}", self.weight),
125 format!("provenance={}", self.provenance),
126 format!("recorded_at={}", self.recorded_at),
127 ];
128 if let Some(transition_from) = &self.transition_from {
129 parts.push(format!("transition_from={transition_from}"));
130 }
131 if let Some(source) = &self.source {
132 parts.push(format!("source_url={}", source.source_url));
133 parts.push(format!("fetched_at={}", source.fetched_at));
134 parts.push(format!("sha256={}", source.sha256));
135 parts.push(format!("cached={}", source.cached));
136 }
137 parts.join(" ")
138 }
139
140 #[must_use]
141 pub fn to_links_notation(&self) -> String {
142 let mut fields = vec![
143 ("id", self.id.clone()),
144 ("target", self.target.clone()),
145 ("observation", self.observation.clone()),
146 ("weight", format!("{:.6}", self.weight)),
147 ("model", self.model.slug().to_owned()),
148 ("provenance", self.provenance.clone()),
149 ("recorded_at", self.recorded_at.clone()),
150 ];
151 if let Some(transition_from) = &self.transition_from {
152 fields.push(("transition_from", transition_from.clone()));
153 }
154 if let Some(source) = &self.source {
155 fields.extend([
156 ("source_url", source.source_url.clone()),
157 ("fetched_at", source.fetched_at.clone()),
158 ("sha256", source.sha256.clone()),
159 ("cached", source.cached.to_string()),
160 ]);
161 }
162 format_lino_record("probability_evidence", &fields)
163 }
164
165 fn stable_record_id(&self) -> String {
166 let source_fingerprint = self.source.as_ref().map_or_else(String::new, |source| {
167 format!(
168 "{}:{}:{}:{}",
169 source.source_url, source.fetched_at, source.sha256, source.cached
170 )
171 });
172 stable_id(
173 "probability",
174 &format!(
175 "{}:{}:{:.6}:{}:{}:{}:{:?}:{}",
176 self.target,
177 self.observation,
178 self.weight,
179 self.model.slug(),
180 self.provenance,
181 self.recorded_at,
182 self.transition_from,
183 source_fingerprint
184 ),
185 )
186 }
187
188 fn usable_offline(&self, offline: bool) -> bool {
189 if !offline {
190 return true;
191 }
192 self.source.as_ref().is_none_or(|source| source.cached)
193 }
194
195 fn applies_to_markov_state(&self, markov_from: Option<&str>) -> bool {
196 match self.model {
197 ProbabilityModel::BayesianEvidence => true,
198 ProbabilityModel::MarkovTransition => self.transition_from.as_deref() == markov_from,
199 }
200 }
201}
202
203#[derive(Debug, Default, Clone, PartialEq)]
205pub struct ProbabilityStore {
206 records: Vec<ProbabilityEvidence>,
207}
208
209impl ProbabilityStore {
210 #[must_use]
211 pub const fn new() -> Self {
212 Self {
213 records: Vec::new(),
214 }
215 }
216
217 #[must_use]
218 pub const fn from_records(records: Vec<ProbabilityEvidence>) -> Self {
219 Self { records }
220 }
221
222 pub fn record(&mut self, evidence: ProbabilityEvidence) -> String {
223 let id = evidence.id.clone();
224 self.records.push(evidence);
225 id
226 }
227
228 pub fn update(
229 &mut self,
230 target: impl Into<String>,
231 observation: impl Into<String>,
232 weight: f32,
233 provenance: impl Into<String>,
234 recorded_at: impl Into<String>,
235 ) -> String {
236 self.record(ProbabilityEvidence::symbolic(
237 target,
238 observation,
239 weight,
240 provenance,
241 recorded_at,
242 ))
243 }
244
245 pub fn reinforce_transition_path<S: AsRef<str>>(
262 &mut self,
263 path: &[S],
264 reward: f32,
265 provenance: impl Into<String>,
266 recorded_at: impl Into<String>,
267 ) -> Vec<String> {
268 let provenance = provenance.into();
269 let recorded_at = recorded_at.into();
270 path.windows(2)
271 .map(|pair| {
272 let from = pair[0].as_ref();
273 let to = pair[1].as_ref();
274 self.record(
275 ProbabilityEvidence::symbolic(
276 to,
277 format!("episode_transition:{from}->{to}"),
278 reward,
279 provenance.clone(),
280 recorded_at.clone(),
281 )
282 .with_model(ProbabilityModel::MarkovTransition)
283 .with_transition_from(from),
284 )
285 })
286 .collect()
287 }
288
289 #[must_use]
290 pub fn records(&self) -> &[ProbabilityEvidence] {
291 &self.records
292 }
293
294 #[must_use]
295 pub fn target_weight(&self, target: &str, offline: bool, markov_from: Option<&str>) -> f32 {
296 self.records
297 .iter()
298 .filter(|evidence| evidence.target == target)
299 .filter(|evidence| evidence.usable_offline(offline))
300 .filter(|evidence| evidence.applies_to_markov_state(markov_from))
301 .map(|evidence| evidence.weight)
302 .sum()
303 }
304
305 #[must_use]
315 pub fn target_evidence_count(
316 &self,
317 target: &str,
318 offline: bool,
319 markov_from: Option<&str>,
320 ) -> usize {
321 self.records
322 .iter()
323 .filter(|evidence| evidence.target == target)
324 .filter(|evidence| evidence.usable_offline(offline))
325 .filter(|evidence| evidence.applies_to_markov_state(markov_from))
326 .count()
327 }
328
329 #[must_use]
339 pub fn nearest_similar_evidence(
340 &self,
341 target: &str,
342 offline: bool,
343 markov_from: Option<&str>,
344 threshold: f32,
345 ) -> Option<SimilarEvidence> {
346 let mut seen: Vec<&str> = Vec::new();
347 let mut best: Option<SimilarEvidence> = None;
348 for evidence in &self.records {
349 let other = evidence.target.as_str();
350 if other == target || seen.contains(&other) {
351 continue;
352 }
353 seen.push(other);
354 let count = self.target_evidence_count(other, offline, markov_from);
355 if count == 0 {
356 continue;
357 }
358 let similarity = symbolic_cosine_similarity(target, other);
359 if similarity < threshold {
360 continue;
361 }
362 let candidate = SimilarEvidence {
363 matched_target: other.to_owned(),
364 weight: self.target_weight(other, offline, markov_from),
365 count,
366 similarity,
367 };
368 let replace = best.as_ref().is_none_or(|current| {
369 match similarity.total_cmp(¤t.similarity) {
370 Ordering::Greater => true,
371 Ordering::Equal => candidate.matched_target < current.matched_target,
372 Ordering::Less => false,
373 }
374 });
375 if replace {
376 best = Some(candidate);
377 }
378 }
379 best
380 }
381
382 #[must_use]
383 pub fn to_links_notation(&self) -> String {
384 let mut blocks = vec![format_lino_record(
385 "probability_store",
386 &[("record_count", self.records.len().to_string())],
387 )];
388 blocks.extend(
389 self.records
390 .iter()
391 .map(ProbabilityEvidence::to_links_notation),
392 );
393 blocks.join("\n\n")
394 }
395
396 pub fn replay_into_event_log(&self, log: &mut EventLog, offline: bool) -> usize {
397 let mut replayed = 0;
398 for evidence in &self.records {
399 if !evidence.usable_offline(offline) {
400 if let Some(source) = &evidence.source {
401 log.append("policy:offline", source.trace_payload());
402 }
403 continue;
404 }
405 log.append("probability:evidence", evidence.trace_payload());
406 log.append("probability:model", evidence.model.slug().to_owned());
407 if let Some(source) = &evidence.source {
408 log.append("source:http", source.trace_payload());
409 if source.cached {
410 log.append("cache_hit", source.source_url.clone());
411 }
412 }
413 replayed += 1;
414 }
415 replayed
416 }
417
418 pub fn append_to_link_store<S: LinkStore>(
419 &self,
420 store: &mut S,
421 offline: bool,
422 ) -> Result<usize, LinkStoreError> {
423 let mut inserted = 0;
424 for evidence in &self.records {
425 if !evidence.usable_offline(offline) {
426 continue;
427 }
428 store.append_memory_event(MemoryEvent {
429 id: evidence.id.clone(),
430 kind: Some(String::from("probability:evidence")),
431 content: Some(evidence.to_links_notation()),
432 sent_at: Some(evidence.recorded_at.clone()),
433 evidence: vec![format!("probability:evidence:{}", evidence.id)],
434 ..MemoryEvent::default()
435 })?;
436 inserted += 1;
437 }
438 Ok(inserted)
439 }
440}
441
442#[derive(Debug, Clone, PartialEq)]
444pub struct SimilarEvidence {
445 pub matched_target: String,
447 pub weight: f32,
449 pub count: usize,
451 pub similarity: f32,
453}
454
455#[derive(Debug, Clone, PartialEq)]
457pub struct ProbabilityCandidate {
458 pub target: String,
459 pub prior_score: f32,
460}
461
462impl ProbabilityCandidate {
463 #[must_use]
464 pub fn new(target: impl Into<String>, prior_score: f32) -> Self {
465 Self {
466 target: target.into(),
467 prior_score: finite_or_zero(prior_score),
468 }
469 }
470}
471
472#[derive(Debug, Default, Clone, PartialEq)]
481pub struct ProbabilityRankingConfig {
482 pub temperature: f32,
483 pub offline: bool,
484 pub markov_from: Option<String>,
485 pub counted_utility: bool,
491 pub min_transition_utility: Option<f32>,
495 pub min_transition_count: Option<usize>,
499 pub similarity_threshold: Option<f32>,
505}
506
507impl ProbabilityRankingConfig {
508 #[must_use]
514 pub const fn with_decision_policy(mut self, policy: ProbabilityDecisionPolicy) -> Self {
515 self.counted_utility = policy.counted_utility;
516 self.min_transition_utility = policy.min_transition_utility;
517 self.min_transition_count = policy.min_transition_count;
518 self.similarity_threshold = policy.similarity_threshold;
519 self
520 }
521
522 #[must_use]
524 pub const fn decision_policy(&self) -> ProbabilityDecisionPolicy {
525 ProbabilityDecisionPolicy {
526 counted_utility: self.counted_utility,
527 min_transition_utility: self.min_transition_utility,
528 min_transition_count: self.min_transition_count,
529 similarity_threshold: self.similarity_threshold,
530 }
531 }
532}
533
534#[derive(Debug, Default, Clone, Copy, PartialEq)]
546pub struct ProbabilityDecisionPolicy {
547 pub counted_utility: bool,
549 pub min_transition_utility: Option<f32>,
551 pub min_transition_count: Option<usize>,
553 pub similarity_threshold: Option<f32>,
556}
557
558#[derive(Debug, Clone, PartialEq)]
566pub struct RankedProbabilityCandidate {
567 pub target: String,
568 pub prior_score: f32,
569 pub evidence_weight: f32,
570 pub evidence_count: usize,
571 pub similarity: f32,
576 pub posterior_score: f32,
577 pub probability: f32,
578}
579
580#[derive(Debug, Clone, PartialEq)]
582pub struct ProbabilityRanking {
583 pub ranked: Vec<RankedProbabilityCandidate>,
584 pub margin: f32,
585}
586
587impl ProbabilityRanking {
588 #[must_use]
589 pub fn probability_for(&self, target: &str) -> Option<f32> {
590 self.ranked
591 .iter()
592 .find(|candidate| candidate.target == target)
593 .map(|candidate| candidate.probability)
594 }
595
596 #[must_use]
597 pub fn trace_summary(&self) -> String {
598 self.ranked
599 .iter()
600 .map(|candidate| {
601 format!(
602 "{}:{:.6}:{:.6}",
603 candidate.target, candidate.posterior_score, candidate.probability
604 )
605 })
606 .collect::<Vec<_>>()
607 .join("|")
608 }
609}
610
611#[must_use]
612pub fn rank_probability_candidates(
613 candidates: &[ProbabilityCandidate],
614 store: &ProbabilityStore,
615 config: ProbabilityRankingConfig,
616) -> ProbabilityRanking {
617 if candidates.is_empty() {
618 return ProbabilityRanking {
619 ranked: Vec::new(),
620 margin: 0.0,
621 };
622 }
623
624 let ProbabilityRankingConfig {
625 temperature,
626 offline,
627 markov_from,
628 counted_utility,
629 min_transition_utility,
630 min_transition_count,
631 similarity_threshold,
632 } = config;
633 let markov_from = markov_from.as_deref();
634 let mut ranked = candidates
635 .iter()
636 .map(|candidate| {
637 let direct_weight = store.target_weight(&candidate.target, offline, markov_from);
638 let direct_count = store.target_evidence_count(&candidate.target, offline, markov_from);
639 let (raw_weight, raw_count, similarity) = if direct_count == 0 {
645 if let Some(found) = similarity_threshold.and_then(|threshold| {
646 store.nearest_similar_evidence(
647 &candidate.target,
648 offline,
649 markov_from,
650 threshold,
651 )
652 }) {
653 (
654 found.weight * found.similarity,
655 found.count,
656 found.similarity,
657 )
658 } else {
659 (direct_weight, direct_count, 1.0)
660 }
661 } else {
662 (direct_weight, direct_count, 1.0)
663 };
664 let below_utility =
668 min_transition_utility.is_some_and(|threshold| raw_weight < threshold);
669 let below_count = min_transition_count.is_some_and(|threshold| raw_count < threshold);
670 let (evidence_weight, evidence_count, similarity) = if below_utility || below_count {
671 (0.0, 0, 1.0)
672 } else {
673 (raw_weight, raw_count, similarity)
674 };
675 let contribution = if counted_utility {
678 evidence_weight * count_to_f32(evidence_count)
679 } else {
680 evidence_weight
681 };
682 let posterior_score = candidate.prior_score + contribution;
683 RankedProbabilityCandidate {
684 target: candidate.target.clone(),
685 prior_score: candidate.prior_score,
686 evidence_weight,
687 evidence_count,
688 similarity,
689 posterior_score,
690 probability: 0.0,
691 }
692 })
693 .collect::<Vec<_>>();
694
695 let probabilities = softmax_scores(
696 &ranked
697 .iter()
698 .map(|candidate| candidate.posterior_score)
699 .collect::<Vec<_>>(),
700 temperature,
701 );
702 for (candidate, probability) in ranked.iter_mut().zip(probabilities) {
703 candidate.probability = probability;
704 }
705
706 ranked.sort_by(|left, right| {
707 right
708 .probability
709 .total_cmp(&left.probability)
710 .then_with(|| right.posterior_score.total_cmp(&left.posterior_score))
711 .then_with(|| left.target.cmp(&right.target))
712 });
713 let margin = match ranked.as_slice() {
714 [first, second, ..] => first.probability - second.probability,
715 [_] => 1.0,
716 [] => 0.0,
717 };
718
719 ProbabilityRanking { ranked, margin }
720}
721
722fn softmax_scores(scores: &[f32], temperature: f32) -> Vec<f32> {
723 if scores.is_empty() {
724 return Vec::new();
725 }
726 let temperature = finite_clamped(temperature, 0.0, 1.0);
727 if temperature <= f32::EPSILON {
728 let mut probabilities = vec![0.0; scores.len()];
729 probabilities[highest_score_index(scores)] = 1.0;
730 return probabilities;
731 }
732
733 let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
734 let weights = scores
735 .iter()
736 .map(|score| ((*score - max_score) / temperature).exp())
737 .collect::<Vec<_>>();
738 let total = weights.iter().sum::<f32>();
739 if !total.is_finite() || total <= f32::EPSILON {
740 let uniform = 1.0 / usize_to_f32(scores.len());
741 return vec![uniform; scores.len()];
742 }
743 weights.iter().map(|weight| *weight / total).collect()
744}
745
746fn highest_score_index(scores: &[f32]) -> usize {
747 scores
748 .iter()
749 .enumerate()
750 .max_by(|(_, left), (_, right)| left.total_cmp(right))
751 .map_or(0, |(index, _)| index)
752}
753
754const fn finite_or_zero(value: f32) -> f32 {
755 if value.is_finite() {
756 value
757 } else {
758 0.0
759 }
760}
761
762const fn finite_clamped(value: f32, min: f32, max: f32) -> f32 {
763 if value.is_finite() {
764 value.clamp(min, max)
765 } else {
766 min
767 }
768}
769
770fn usize_to_f32(value: usize) -> f32 {
771 let bounded = u16::try_from(value).unwrap_or(u16::MAX);
772 f32::from(bounded.max(1))
773}
774
775fn count_to_f32(value: usize) -> f32 {
780 f32::from(u16::try_from(value).unwrap_or(u16::MAX))
781}
782
783#[must_use]
790pub fn symbolic_cosine_similarity(a: &str, b: &str) -> f32 {
791 let left = tokenize_symbolic(a);
792 let right = tokenize_symbolic(b);
793 if left.is_empty() || right.is_empty() {
794 return 0.0;
795 }
796 let left_counts = bag_of_words(&left);
797 let right_counts = bag_of_words(&right);
798 let mut dot = 0.0f32;
799 for (token, left_count) in &left_counts {
800 if let Some(right_count) = right_counts.get(token) {
801 dot = count_to_f32(*left_count).mul_add(count_to_f32(*right_count), dot);
802 }
803 }
804 let left_norm = vector_norm(&left_counts);
805 let right_norm = vector_norm(&right_counts);
806 if left_norm <= f32::EPSILON || right_norm <= f32::EPSILON {
807 return 0.0;
808 }
809 (dot / (left_norm * right_norm)).clamp(0.0, 1.0)
810}
811
812fn tokenize_symbolic(value: &str) -> Vec<String> {
813 value
814 .split(|character: char| !character.is_alphanumeric())
815 .filter(|token| !token.is_empty())
816 .map(str::to_lowercase)
817 .collect()
818}
819
820fn bag_of_words(tokens: &[String]) -> BTreeMap<String, usize> {
821 let mut counts = BTreeMap::new();
822 for token in tokens {
823 *counts.entry(token.clone()).or_insert(0) += 1;
824 }
825 counts
826}
827
828fn vector_norm(counts: &BTreeMap<String, usize>) -> f32 {
829 let sum_of_squares = counts
830 .values()
831 .map(|count| {
832 let value = count_to_f32(*count);
833 value * value
834 })
835 .sum::<f32>();
836 sum_of_squares.sqrt()
837}