| 143 | } |
| 144 | |
| 145 | pub fn assemble(&self, results: &[ContextResult]) -> ContextAssembly { |
| 146 | let mut deduped: HashMap<String, ContextItem> = HashMap::new(); |
| 147 | let mut source_count = 0usize; |
| 148 | |
| 149 | for result in results { |
| 150 | for item in &result.items { |
| 151 | source_count += 1; |
| 152 | let key = dedupe_key(item); |
| 153 | match deduped.get(&key) { |
| 154 | Some(existing) |
| 155 | if ranking_score(existing) |
| 156 | .total_cmp(&ranking_score(item)) |
| 157 | .then_with(|| existing.relevance.total_cmp(&item.relevance)) |
| 158 | .is_ge() => {} |
| 159 | _ => { |
| 160 | deduped.insert(key, item.clone()); |
| 161 | } |
| 162 | } |
| 163 | } |
| 164 | } |
| 165 | |
| 166 | let mut items = deduped.into_values().collect::<Vec<_>>(); |
| 167 | items.sort_by(|a, b| { |
| 168 | ranking_score(b) |
| 169 | .total_cmp(&ranking_score(a)) |
| 170 | .then_with(|| b.relevance.total_cmp(&a.relevance)) |
| 171 | .then_with(|| estimated_tokens(a).cmp(&estimated_tokens(b))) |
| 172 | .then_with(|| a.id.cmp(&b.id)) |
| 173 | }); |
| 174 | |
| 175 | let mut selected = Vec::new(); |
| 176 | let mut total_tokens = 0usize; |
| 177 | let mut truncated = source_count > items.len(); |
| 178 | let mut source_item_counts: HashMap<String, usize> = HashMap::new(); |
| 179 | let mut source_token_counts: HashMap<String, usize> = HashMap::new(); |
| 180 | |
| 181 | for item in items { |
| 182 | if selected.len() >= self.budget.max_items { |
| 183 | truncated = true; |
| 184 | break; |
| 185 | } |
| 186 | |
| 187 | let item_tokens = estimated_tokens(&item); |
| 188 | if total_tokens + item_tokens > self.budget.max_tokens { |
| 189 | truncated = true; |
| 190 | continue; |
| 191 | } |
| 192 | |
| 193 | let source_key = source_policy_key(&item); |
| 194 | if let Some(max_items) = self.source_policy.max_items_per_source { |
| 195 | let count = source_item_counts.get(&source_key).copied().unwrap_or(0); |
| 196 | if count >= max_items { |
| 197 | truncated = true; |
| 198 | continue; |
| 199 | } |
| 200 | } |
| 201 | if let Some(max_tokens) = self.source_policy.max_tokens_per_source { |
| 202 | let source_tokens = source_token_counts.get(&source_key).copied().unwrap_or(0); |