1use std::collections::HashMap;
39use std::sync::Arc;
40use std::time::{Duration, Instant};
41
42use futures_util::future::join_all;
43
44use crate::budget::{apply_budget, PruneStrategy, WeightedContribution};
45use crate::error::{BriefError, Result};
46use crate::governance::{Governance, GovernanceVerdict, NoOpGovernance};
47use crate::receipt::{BriefReceipt, DroppedContribution, SourceStats, Timings};
48use crate::source::{Contribution, ContributionContent, Source, SourceError};
49use crate::tokenizer::{CharApproxTokenizer, Tokenizer};
50use crate::types::{
51 Brief, BriefContext, BriefMessage, Priority, SourceId, TokenBudget, ToolSchema,
52};
53
54pub struct BriefBuilder {
61 sources: Vec<Arc<dyn Source>>,
62 governance: Arc<dyn Governance>,
63 tokenizer: Arc<dyn Tokenizer>,
64 budget: TokenBudget,
65 strategy: PruneStrategy,
66}
67
68impl Default for BriefBuilder {
69 fn default() -> Self {
70 BriefBuilder::new()
71 }
72}
73
74impl std::fmt::Debug for BriefBuilder {
75 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76 f.debug_struct("BriefBuilder")
77 .field("sources", &self.sources.len())
78 .field("budget", &self.budget)
79 .field("strategy", &self.strategy)
80 .finish_non_exhaustive()
81 }
82}
83
84impl BriefBuilder {
85 pub fn new() -> Self {
90 BriefBuilder {
91 sources: Vec::new(),
92 governance: Arc::new(NoOpGovernance),
93 tokenizer: Arc::new(CharApproxTokenizer),
94 budget: TokenBudget::default(),
95 strategy: PruneStrategy::default(),
96 }
97 }
98
99 pub fn source(mut self, source: Arc<dyn Source>) -> Self {
102 self.sources.push(source);
103 self
104 }
105
106 pub fn governance(mut self, governance: Arc<dyn Governance>) -> Self {
108 self.governance = governance;
109 self
110 }
111
112 pub fn tokenizer(mut self, tokenizer: Arc<dyn Tokenizer>) -> Self {
114 self.tokenizer = tokenizer;
115 self
116 }
117
118 pub fn budget(mut self, budget: TokenBudget) -> Self {
120 self.budget = budget;
121 self
122 }
123
124 pub fn strategy(mut self, strategy: PruneStrategy) -> Self {
127 self.strategy = strategy;
128 self
129 }
130
131 pub fn sources(&self) -> &[Arc<dyn Source>] {
133 &self.sources
134 }
135
136 pub fn current_budget(&self) -> &TokenBudget {
138 &self.budget
139 }
140
141 pub async fn build(&self, ctx: &BriefContext) -> Result<Brief> {
164 let start = Instant::now();
165
166 self.check_unique_ids()?;
168
169 let (per_source_results, fanout_elapsed) = self.fan_out(ctx).await;
171
172 for (sid, res) in &per_source_results {
176 if let Err(err) = res {
177 match err {
178 SourceError::Skipped(_) => {} other => {
180 return Err(BriefError::Source {
181 source_id: sid.to_string(),
182 message: other.to_string(),
183 });
184 }
185 }
186 }
187 }
188
189 let mut by_source_priority: HashMap<SourceId, Priority> = HashMap::new();
193 let mut by_source_contribs: HashMap<SourceId, usize> = HashMap::new();
194 let mut flat: Vec<WeightedContribution> = Vec::new();
195
196 for (sid, res) in per_source_results.into_iter() {
197 let priority = self
198 .sources
199 .iter()
200 .find(|s| s.id() == sid)
201 .map(|s| s.priority())
202 .unwrap_or(Priority::Normal);
203 by_source_priority.insert(sid.clone(), priority);
204
205 let contributions = match res {
206 Ok(contribs) => contribs,
207 Err(SourceError::Skipped(_)) => Vec::new(),
208 Err(_) => unreachable!("fatal source errors handled above"),
210 };
211
212 by_source_contribs.insert(sid.clone(), contributions.len());
213
214 let tokenize_start = Instant::now();
216 for (idx, contribution) in contributions.into_iter().enumerate() {
217 let actual_tokens = measure_contribution(&*self.tokenizer, &contribution.content);
218 flat.push(WeightedContribution {
219 source: sid.clone(),
220 priority,
221 actual_tokens,
222 source_index: idx,
223 contribution,
224 });
225 }
226 let _ = tokenize_start.elapsed();
230 }
231
232 let tokenize_total = compute_tokenize_total(&self.tokenizer, &flat);
238
239 let prune_start = Instant::now();
241 let (kept, dropped) = apply_budget(flat, &self.budget, self.strategy);
242 let prune_elapsed = prune_start.elapsed();
243
244 let (mut draft, per_source_stats) =
246 self.assemble_draft(kept, &dropped, &by_source_priority, &by_source_contribs);
247
248 let governance_start = Instant::now();
250 let verdict = self
251 .governance
252 .review(&mut draft, ctx)
253 .await
254 .map_err(|e| BriefError::Rejected(e.to_string()))?;
255 let governance_elapsed = governance_start.elapsed();
256
257 match verdict {
258 GovernanceVerdict::Allow => {}
259 GovernanceVerdict::Redacted(records) => {
260 draft.receipt.redactions = records;
261 }
262 GovernanceVerdict::Rejected(reason) => {
263 return Err(BriefError::Rejected(reason));
264 }
265 }
266
267 draft.receipt.dropped = dropped;
269 draft.receipt.by_source = per_source_stats;
270 draft.receipt.total_tokens = draft.receipt.by_source.values().map(|s| s.tokens).sum();
271 draft.receipt.timings = Timings {
272 fanout: fanout_elapsed,
273 tokenize: tokenize_total,
274 prune: prune_elapsed,
275 governance: governance_elapsed,
276 total: start.elapsed(),
277 };
278 draft.receipt.built_at = std::time::SystemTime::now();
279
280 let prompt_budget = self.budget.prompt_budget();
285 if draft.receipt.total_tokens > prompt_budget {
286 return Err(BriefError::BudgetUnsatisfiable {
287 needed: draft.receipt.total_tokens,
288 available: prompt_budget,
289 });
290 }
291
292 Ok(draft)
293 }
294
295 fn check_unique_ids(&self) -> Result<()> {
296 let mut seen = std::collections::HashSet::new();
297 for source in &self.sources {
298 let id = source.id();
299 if !seen.insert(id.clone()) {
300 return Err(BriefError::Source {
301 source_id: id.to_string(),
302 message: "duplicate source id registered on BriefBuilder".into(),
303 });
304 }
305 }
306 Ok(())
307 }
308
309 async fn fan_out(
310 &self,
311 ctx: &BriefContext,
312 ) -> (
313 Vec<(
314 SourceId,
315 std::result::Result<Vec<Contribution>, SourceError>,
316 )>,
317 Duration,
318 ) {
319 let start = Instant::now();
320 let futures = self.sources.iter().map(|source| {
321 let source = source.clone();
322 let ctx = ctx.clone();
323 async move {
324 let id = source.id();
325 let res = source.contribute(&ctx).await;
326 (id, res)
327 }
328 });
329
330 let results = join_all(futures).await;
331 let elapsed = start.elapsed();
335 (results, elapsed)
336 }
337
338 fn assemble_draft(
339 &self,
340 kept: Vec<WeightedContribution>,
341 dropped: &[DroppedContribution],
342 by_source_priority: &HashMap<SourceId, Priority>,
343 by_source_contribs: &HashMap<SourceId, usize>,
344 ) -> (Brief, HashMap<SourceId, SourceStats>) {
345 let mut system_chunks: Vec<String> = Vec::new();
346 let mut messages: Vec<BriefMessage> = Vec::new();
347 let mut tools: Vec<ToolSchema> = Vec::new();
348
349 let mut kept_count: HashMap<SourceId, usize> = HashMap::new();
351 let mut kept_tokens: HashMap<SourceId, usize> = HashMap::new();
352
353 for w in &kept {
354 *kept_count.entry(w.source.clone()).or_insert(0) += 1;
355 *kept_tokens.entry(w.source.clone()).or_insert(0) += w.actual_tokens;
356 }
357
358 for w in kept {
360 match w.contribution.content {
361 ContributionContent::System { text } => {
362 system_chunks.push(text);
363 }
364 ContributionContent::Text { role, content } => {
365 messages.push(BriefMessage::Text {
366 role,
367 content,
368 source: w.source,
369 });
370 }
371 ContributionContent::Image { role, data, alt } => {
372 messages.push(BriefMessage::Image {
373 role,
374 data,
375 alt,
376 source: w.source,
377 });
378 }
379 ContributionContent::ToolCall { id, name, args } => {
380 messages.push(BriefMessage::ToolCall {
381 id,
382 name,
383 args,
384 source: w.source,
385 });
386 }
387 ContributionContent::ToolResult { id, content } => {
388 messages.push(BriefMessage::ToolResult {
389 id,
390 content,
391 source: w.source,
392 });
393 }
394 ContributionContent::Tool { mut schema } => {
395 schema.source = w.source.clone();
398 tools.push(schema);
399 }
400 }
401 }
402
403 let system = if system_chunks.is_empty() {
404 None
405 } else {
406 Some(system_chunks.join("\n\n"))
407 };
408
409 let mut by_source_stats: HashMap<SourceId, SourceStats> = HashMap::new();
412 for source in &self.sources {
413 let sid = source.id();
414 let priority = by_source_priority
415 .get(&sid)
416 .copied()
417 .unwrap_or(Priority::Normal);
418 let contributions = by_source_contribs.get(&sid).copied().unwrap_or(0);
419 let kept = kept_count.get(&sid).copied().unwrap_or(0);
420 let tokens = kept_tokens.get(&sid).copied().unwrap_or(0);
421 by_source_stats.insert(
422 sid,
423 SourceStats {
424 contributions,
425 kept,
426 tokens,
427 priority,
428 },
429 );
430 }
431
432 let mut receipt = BriefReceipt::empty();
433 receipt.dropped = dropped.to_vec();
434
435 (
436 Brief {
437 system,
438 messages,
439 tools,
440 receipt,
441 },
442 by_source_stats,
443 )
444 }
445}
446
447fn measure_contribution(tokenizer: &dyn Tokenizer, content: &ContributionContent) -> usize {
448 match content {
449 ContributionContent::System { text } => tokenizer.count(text),
450 ContributionContent::Text { content, .. } => tokenizer.count(content),
451 ContributionContent::Image { alt, .. } => {
452 alt.as_deref().map(|a| tokenizer.count(a)).unwrap_or(0)
455 }
456 ContributionContent::ToolCall { name, args, .. } => {
457 let mut total = tokenizer.count(name);
458 if let Ok(s) = serde_json::to_string(args) {
460 total += tokenizer.count(&s);
461 }
462 total
463 }
464 ContributionContent::ToolResult { content, .. } => tokenizer.count(content),
465 ContributionContent::Tool { schema } => {
466 let mut total = tokenizer.count(&schema.name) + tokenizer.count(&schema.description);
467 if let Ok(s) = serde_json::to_string(&schema.input_schema) {
468 total += tokenizer.count(&s);
469 }
470 total
471 }
472 }
473}
474
475fn compute_tokenize_total(
478 tokenizer: &Arc<dyn Tokenizer>,
479 flat: &[WeightedContribution],
480) -> Duration {
481 let start = Instant::now();
482 for w in flat {
483 let _ = measure_contribution(&**tokenizer, &w.contribution.content);
484 }
485 start.elapsed()
486}
487
488#[cfg(test)]
489mod tests {
490 use super::*;
491
492 use async_trait::async_trait;
493
494 use crate::source::{Contribution, Source, SourceError};
495 use crate::types::{Role, SourceId};
496
497 struct FixedSource {
498 id: &'static str,
499 priority: Priority,
500 contributions: Vec<Contribution>,
501 }
502
503 #[async_trait]
504 impl Source for FixedSource {
505 fn id(&self) -> SourceId {
506 SourceId::new(self.id)
507 }
508
509 fn priority(&self) -> Priority {
510 self.priority
511 }
512
513 async fn contribute(
514 &self,
515 _ctx: &BriefContext,
516 ) -> std::result::Result<Vec<Contribution>, SourceError> {
517 Ok(self.contributions.clone())
518 }
519 }
520
521 struct ErrSource;
522
523 #[async_trait]
524 impl Source for ErrSource {
525 fn id(&self) -> SourceId {
526 SourceId::new("err")
527 }
528
529 fn priority(&self) -> Priority {
530 Priority::Normal
531 }
532
533 async fn contribute(
534 &self,
535 _ctx: &BriefContext,
536 ) -> std::result::Result<Vec<Contribution>, SourceError> {
537 Err(SourceError::Backend("boom".into()))
538 }
539 }
540
541 struct SkipSource;
542
543 #[async_trait]
544 impl Source for SkipSource {
545 fn id(&self) -> SourceId {
546 SourceId::new("skip")
547 }
548
549 fn priority(&self) -> Priority {
550 Priority::Normal
551 }
552
553 async fn contribute(
554 &self,
555 _ctx: &BriefContext,
556 ) -> std::result::Result<Vec<Contribution>, SourceError> {
557 Err(SourceError::Skipped("nothing to add".into()))
558 }
559 }
560
561 fn ctx() -> BriefContext {
562 BriefContext::new(TokenBudget::default())
563 }
564
565 #[tokio::test]
566 async fn builds_a_brief_from_a_single_system_source() {
567 let src = Arc::new(FixedSource {
568 id: "sys",
569 priority: Priority::Critical,
570 contributions: vec![Contribution::system("be helpful", 3)],
571 });
572 let builder = BriefBuilder::new().source(src);
573 let brief = builder.build(&ctx()).await.expect("build ok");
574 assert_eq!(brief.system.as_deref(), Some("be helpful"));
575 assert!(brief.messages.is_empty());
576 assert_eq!(brief.receipt.by_source.len(), 1);
577 let stats = brief
578 .receipt
579 .by_source
580 .get(&SourceId::new("sys"))
581 .expect("stats");
582 assert_eq!(stats.contributions, 1);
583 assert_eq!(stats.kept, 1);
584 assert!(stats.tokens > 0);
585 }
586
587 #[tokio::test]
588 async fn skipped_sources_do_not_fail_the_build() {
589 let sys = Arc::new(FixedSource {
590 id: "sys",
591 priority: Priority::Critical,
592 contributions: vec![Contribution::system("ok", 1)],
593 });
594 let skip = Arc::new(SkipSource);
595 let builder = BriefBuilder::new().source(sys).source(skip);
596 let brief = builder.build(&ctx()).await.expect("build ok");
597 let skip_stats = brief
598 .receipt
599 .by_source
600 .get(&SourceId::new("skip"))
601 .expect("skip recorded");
602 assert_eq!(skip_stats.contributions, 0);
603 assert_eq!(skip_stats.kept, 0);
604 }
605
606 #[tokio::test]
607 async fn fatal_source_error_propagates() {
608 let bad = Arc::new(ErrSource);
609 let builder = BriefBuilder::new().source(bad);
610 let err = builder.build(&ctx()).await.expect_err("should error");
611 match err {
612 BriefError::Source { source_id, .. } => assert_eq!(source_id, "err"),
613 other => panic!("expected Source error, got {other:?}"),
614 }
615 }
616
617 #[tokio::test]
618 async fn duplicate_source_ids_rejected() {
619 let a = Arc::new(FixedSource {
620 id: "same",
621 priority: Priority::Normal,
622 contributions: vec![],
623 });
624 let b = Arc::new(FixedSource {
625 id: "same",
626 priority: Priority::Normal,
627 contributions: vec![],
628 });
629 let builder = BriefBuilder::new().source(a).source(b);
630 let err = builder.build(&ctx()).await.expect_err("should error");
631 match err {
632 BriefError::Source { source_id, message } => {
633 assert_eq!(source_id, "same");
634 assert!(message.contains("duplicate"));
635 }
636 other => panic!("expected duplicate Source error, got {other:?}"),
637 }
638 }
639
640 #[tokio::test]
641 async fn budget_drops_lowest_importance_first() {
642 let src = Arc::new(FixedSource {
645 id: "mem",
646 priority: Priority::Normal,
647 contributions: vec![
648 Contribution::text(Role::User, "x".repeat(100), 25).with_importance(0.9),
649 Contribution::text(Role::User, "y".repeat(100), 25).with_importance(0.1),
650 ],
651 });
652 let budget = TokenBudget::new(30, 0);
655 let builder = BriefBuilder::new().source(src).budget(budget);
656 let brief = builder.build(&ctx()).await.expect("build ok");
657 assert_eq!(brief.messages.len(), 1);
658 assert_eq!(brief.receipt.dropped.len(), 1);
659 match &brief.messages[0] {
661 BriefMessage::Text { content, .. } => {
662 assert!(
663 content.starts_with('x'),
664 "expected x's first, got {content}"
665 );
666 }
667 other => panic!("expected Text message, got {other:?}"),
668 }
669 }
670
671 #[tokio::test]
672 async fn budget_unsatisfiable_when_critical_overflows() {
673 let huge = "x".repeat(5000);
674 let src = Arc::new(FixedSource {
675 id: "sys",
676 priority: Priority::Critical,
677 contributions: vec![Contribution::system(huge, 1250)],
678 });
679 let budget = TokenBudget::new(100, 0);
680 let builder = BriefBuilder::new().source(src).budget(budget);
681 let err = builder.build(&ctx()).await.expect_err("should error");
682 match err {
683 BriefError::BudgetUnsatisfiable { needed, available } => {
684 assert!(needed > available);
685 assert_eq!(available, 100);
686 }
687 other => panic!("expected BudgetUnsatisfiable, got {other:?}"),
688 }
689 }
690
691 #[tokio::test]
692 async fn tokenizer_swap_is_used_for_pruning() {
693 struct Fat;
696 impl Tokenizer for Fat {
697 fn count(&self, _text: &str) -> usize {
698 1000
699 }
700 }
701 let src = Arc::new(FixedSource {
702 id: "mem",
703 priority: Priority::Normal,
704 contributions: vec![Contribution::text(Role::User, "hi", 1)],
705 });
706 let budget = TokenBudget::new(30, 0);
707 let builder = BriefBuilder::new()
708 .source(src)
709 .budget(budget)
710 .tokenizer(Arc::new(Fat));
711 let brief = builder.build(&ctx()).await.expect("build ok");
712 assert!(
713 brief.messages.is_empty(),
714 "Fat tokenizer should drop everything"
715 );
716 assert_eq!(brief.receipt.dropped.len(), 1);
717 }
718
719 #[tokio::test]
720 async fn governance_rejection_surfaces_as_error() {
721 use crate::governance::{Governance, GovernanceError, GovernanceVerdict};
722
723 struct AlwaysReject;
724 #[async_trait]
725 impl Governance for AlwaysReject {
726 async fn review(
727 &self,
728 _draft: &mut Brief,
729 _ctx: &BriefContext,
730 ) -> std::result::Result<GovernanceVerdict, GovernanceError> {
731 Ok(GovernanceVerdict::Rejected("no".into()))
732 }
733 }
734
735 let src = Arc::new(FixedSource {
736 id: "sys",
737 priority: Priority::Critical,
738 contributions: vec![Contribution::system("hi", 1)],
739 });
740 let builder = BriefBuilder::new()
741 .source(src)
742 .governance(Arc::new(AlwaysReject));
743 let err = builder.build(&ctx()).await.expect_err("should reject");
744 match err {
745 BriefError::Rejected(reason) => assert_eq!(reason, "no"),
746 other => panic!("expected Rejected, got {other:?}"),
747 }
748 }
749
750 #[tokio::test]
751 async fn governance_redactions_land_on_receipt() {
752 use crate::governance::{Governance, GovernanceError, GovernanceVerdict};
753 use crate::receipt::RedactionRecord;
754
755 struct Redact;
756 #[async_trait]
757 impl Governance for Redact {
758 async fn review(
759 &self,
760 _draft: &mut Brief,
761 _ctx: &BriefContext,
762 ) -> std::result::Result<GovernanceVerdict, GovernanceError> {
763 Ok(GovernanceVerdict::Redacted(vec![RedactionRecord {
764 source: SourceId::new("sys"),
765 rule: "rule:test".into(),
766 }]))
767 }
768 }
769
770 let src = Arc::new(FixedSource {
771 id: "sys",
772 priority: Priority::Critical,
773 contributions: vec![Contribution::system("hi", 1)],
774 });
775 let builder = BriefBuilder::new().source(src).governance(Arc::new(Redact));
776 let brief = builder.build(&ctx()).await.expect("ok");
777 assert_eq!(brief.receipt.redactions.len(), 1);
778 assert_eq!(brief.receipt.redactions[0].rule, "rule:test");
779 }
780
781 #[tokio::test]
782 async fn priority_floor_protects_normal_bucket() {
783 let critical = Arc::new(FixedSource {
787 id: "sys",
788 priority: Priority::Critical,
789 contributions: vec![Contribution::system("c".repeat(100), 25)],
790 });
791 let normal = Arc::new(FixedSource {
792 id: "hist",
793 priority: Priority::Normal,
794 contributions: vec![
795 Contribution::text(Role::User, "n".repeat(100), 25).with_importance(0.5),
796 Contribution::text(Role::User, "x".repeat(100), 25).with_importance(0.4),
797 ],
798 });
799 let low = Arc::new(FixedSource {
800 id: "noise",
801 priority: Priority::Low,
802 contributions: vec![
803 Contribution::text(Role::User, "l".repeat(100), 25).with_importance(0.9)
804 ],
805 });
806
807 let budget = TokenBudget::new(80, 0).with_floor(Priority::Normal, 25);
810 let builder = BriefBuilder::new()
811 .source(critical)
812 .source(normal)
813 .source(low)
814 .budget(budget);
815 let brief = builder.build(&ctx()).await.expect("ok");
816 let dropped_sources: Vec<&str> = brief
817 .receipt
818 .dropped
819 .iter()
820 .map(|d| d.source.as_str())
821 .collect();
822 assert!(
823 dropped_sources.contains(&"noise"),
824 "expected noise to be dropped, got dropped={dropped_sources:?}"
825 );
826 assert!(!dropped_sources.contains(&"hist"));
827 }
828}