public_02_context_budget_and_sections/
02_context_budget_and_sections.rs1#[path = "../support/utils.rs"]
2mod _utils;
3
4use _utils::{boxed_error, cleanup_run, create_client, new_run_id, print_summary, require};
5use mubit_sdk::{GetContextOptions, RememberOptions, TransportMode};
6use serde_json::json;
7use std::error::Error;
8use std::time::Instant;
9
10#[tokio::main(flavor = "current_thread")]
11async fn main() -> Result<(), Box<dyn Error>> {
12 let name = "public_02_context_budget_and_sections";
13 let started = Instant::now();
14 let client = create_client().await?;
15 let run_id = new_run_id("public_context_budget");
16 client.set_run_id(Some(run_id.clone()));
17 client.set_transport(TransportMode::Http);
18
19 let mut passed = true;
20 let mut detail = "validated context budgeting and sectioned disclosure".to_string();
21 let mut metrics = json!({});
22
23 let scenario = async {
24 let mut rule = RememberOptions::new("Always lead with policy facts before escalation notes.");
25 rule.run_id = Some(run_id.clone());
26 rule.agent_id = Some("planner".to_string());
27 rule.intent = Some("rule".to_string());
28 rule.metadata = Some(json!({"source": "public-example", "section": "rules"}));
29 rule.importance = Some("critical".to_string());
30 client.remember(rule).await?;
31
32 let mut fact = RememberOptions::new("Claim CLM-4242 has a policy limit of 5000 GBP and a deductible of 250 GBP.");
33 fact.run_id = Some(run_id.clone());
34 fact.agent_id = Some("planner".to_string());
35 fact.intent = Some("fact".to_string());
36 fact.metadata = Some(json!({"source": "public-example", "section": "facts"}));
37 fact.importance = Some("high".to_string());
38 client.remember(fact).await?;
39
40 let mut trace = RememberOptions::new("Planner noted that the claimant is waiting on a same-day callback.");
41 trace.run_id = Some(run_id.clone());
42 trace.agent_id = Some("planner".to_string());
43 trace.intent = Some("trace".to_string());
44 trace.metadata = Some(json!({"source": "public-example", "section": "traces"}));
45 trace.importance = Some("medium".to_string());
46 client.remember(trace).await?;
47
48 let mut mental = RememberOptions::new("Claim CLM-4242 overview: 5000 GBP policy limit, 250 GBP deductible, claimant awaiting same-day callback. Priority: high.");
49 mental.run_id = Some(run_id.clone());
50 mental.agent_id = Some("planner".to_string());
51 mental.intent = Some("mental_model".to_string());
52 mental.metadata = Some(json!({"source": "public-example", "section": "mental_models"}));
53 mental.importance = Some("critical".to_string());
54 client.remember(mental).await?;
55
56 let mut summary_opts = GetContextOptions::default();
57 summary_opts.run_id = Some(run_id.clone());
58 summary_opts.query = Some("Prepare a claims summary.".to_string());
59 summary_opts.mode = Some("summary".to_string());
60 summary_opts.entry_types = vec!["rule".to_string(), "fact".to_string(), "trace".to_string()];
61 summary_opts.max_token_budget = Some(220);
62 summary_opts.limit = Some(6);
63 let summary = client.get_context(summary_opts).await?;
64
65 let mut sections_opts = GetContextOptions::default();
66 sections_opts.run_id = Some(run_id.clone());
67 sections_opts.query = Some("Prepare a claims summary.".to_string());
68 sections_opts.mode = Some("sections".to_string());
69 sections_opts.sections = vec!["mental_models".to_string(), "rules".to_string(), "facts".to_string()];
70 sections_opts.entry_types = vec!["rule".to_string(), "fact".to_string(), "trace".to_string()];
71 sections_opts.max_token_budget = Some(220);
72 sections_opts.limit = Some(6);
73 let sections = client.get_context(sections_opts).await?;
74
75 require(summary.get("context_block").and_then(|v| v.as_str()).unwrap_or("").is_empty(), format!("summary mode should suppress context_block: {summary}"))?;
76 let summary_sections = summary.get("section_summaries").and_then(|v| v.as_array()).ok_or_else(|| boxed_error(format!("summary mode missing section summaries: {summary}")))?;
77 let section_sources = sections.get("sources").and_then(|v| v.as_array()).ok_or_else(|| boxed_error(format!("sections mode missing sources: {sections}")))?;
78 require(!summary_sections.is_empty(), format!("summary mode missing section summaries: {summary}"))?;
79 require(!section_sources.is_empty(), format!("sections mode missing sources: {sections}"))?;
80 require(sections.get("budget_used").and_then(|v| v.as_u64()).unwrap_or(0) >= 1, format!("budget_used missing: {sections}"))?;
81 require(sections.get("budget_remaining").and_then(|v| v.as_u64()).unwrap_or(0) <= 220, format!("budget_remaining missing: {sections}"))?;
82 require(section_sources.iter().any(|s| s.get("entry_type").and_then(|v| v.as_str()) == Some("rule")), format!("expected rule source in sections mode: {sections}"))?;
83
84 metrics = json!({
85 "run_id": run_id,
86 "summary_section_count": summary_sections.len(),
87 "section_source_count": section_sources.len(),
88 "budget_used": sections.get("budget_used").cloned().unwrap_or(serde_json::Value::Null),
89 "budget_remaining": sections.get("budget_remaining").cloned().unwrap_or(serde_json::Value::Null),
90 });
91 Ok::<(), Box<dyn Error>>(())
92 }
93 .await;
94
95 if let Err(err) = scenario {
96 passed = false;
97 detail = err.to_string();
98 }
99
100 let cleanup_ok = cleanup_run(&client, &run_id).await;
101 if !cleanup_ok {
102 passed = false;
103 detail = format!("{detail} | cleanup failures");
104 }
105
106 print_summary(name, passed, &detail, &metrics, started.elapsed().as_secs_f64(), cleanup_ok);
107 if passed { Ok(()) } else { Err(boxed_error(detail)) }
108}