cupel 1.2.0

Context window management pipeline for LLM applications
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

use std::collections::HashMap;

use crate::CupelError;
use crate::model::{ContextBudget, ContextItem, OverflowStrategy, ScoredItem};
use crate::placer::Placer;

/// Result type for [`place_items`]: `(placed_items, truncated_items_with_scores)`.
type PlaceResult = Result<(Vec<ContextItem>, Vec<(ContextItem, f64)>), CupelError>;

/// Merges pinned items with sliced items, handles overflow, and delegates to
/// the placer for final ordering.
///
/// Returns `(placed, truncated)` where `truncated` is a vec of `(item, score)`
/// pairs for any items dropped by `Truncate` overflow handling. `truncated` is
/// always empty unless `OverflowStrategy::Truncate` fires.
///
/// Pinned items are assigned score 1.0. Overflow detection compares against the
/// ORIGINAL `targetTokens` (not the effective budget used for slicing).
pub(crate) fn place_items(
    pinned: &[ContextItem],
    sliced: &[ContextItem],
    sorted_scored: &[ScoredItem],
    budget: &ContextBudget,
    overflow_strategy: OverflowStrategy,
    placer: &dyn Placer,
) -> PlaceResult {
    // Build content -> score map (first occurrence = highest score since sorted desc)
    let mut score_map: HashMap<&str, f64> = HashMap::with_capacity(sorted_scored.len());
    for si in sorted_scored {
        score_map.entry(si.item.content()).or_insert(si.score);
    }

    // Step 1: Merge — pinned items with score 1.0, then sliced items with original scores
    let mut merged: Vec<ScoredItem> = Vec::with_capacity(pinned.len() + sliced.len());

    for item in pinned {
        merged.push(ScoredItem {
            item: item.clone(),
            score: 1.0,
        });
    }

    for item in sliced {
        let score = score_map.get(item.content()).copied().unwrap_or(0.0);

        merged.push(ScoredItem {
            item: item.clone(),
            score,
        });
    }

    // Step 2: Overflow detection — compare against ORIGINAL targetTokens
    let merged_tokens: i64 = merged.iter().map(|si| si.item.tokens()).sum();

    let truncated: Vec<ScoredItem>;
    if merged_tokens > budget.target_tokens() {
        let (kept, dropped) = handle_overflow(merged, budget.target_tokens(), overflow_strategy)?;
        merged = kept;
        truncated = dropped;
    } else {
        truncated = vec![];
    }

    // Step 3: Delegate to placer for final ordering
    let result = placer.place(&merged);
    let truncated_with_scores: Vec<(ContextItem, f64)> = truncated
        .into_iter()
        .map(|si| (si.item, si.score))
        .collect();

    Ok((result, truncated_with_scores))
}

fn handle_overflow(
    mut merged: Vec<ScoredItem>,
    target_tokens: i64,
    strategy: OverflowStrategy,
) -> Result<(Vec<ScoredItem>, Vec<ScoredItem>), CupelError> {
    match strategy {
        OverflowStrategy::Throw => {
            let merged_tokens: i64 = merged.iter().map(|si| si.item.tokens()).sum();
            Err(CupelError::Overflow {
                merged_tokens,
                target_tokens,
            })
        }
        OverflowStrategy::Truncate => {
            // Sort non-pinned items by score descending so truncation removes
            // lowest-priority items first (pinned items are always kept).
            merged.sort_by(|a, b| match (a.item.pinned(), b.item.pinned()) {
                (true, false) => std::cmp::Ordering::Less,
                (false, true) => std::cmp::Ordering::Greater,
                _ => b.score.total_cmp(&a.score),
            });

            let mut kept = Vec::new();
            let mut dropped = Vec::new();
            let mut current_tokens: i64 = 0;

            for si in merged {
                let fits = si.item.pinned() || current_tokens + si.item.tokens() <= target_tokens;
                if fits {
                    current_tokens += si.item.tokens();
                    kept.push(si);
                } else {
                    dropped.push(si);
                }
            }

            Ok((kept, dropped))
        }
        OverflowStrategy::Proceed => {
            // Accept over-budget selection
            Ok((merged, vec![]))
        }
    }
}