vv_agent/
context_providers.rs

1use std::collections::BTreeMap;
2
3use serde_json::Value;
4use sha2::{Digest, Sha256};
5
6use crate::types::Metadata;
7
8#[derive(Debug, Clone, PartialEq)]
9pub struct ContextFragment {
10    pub id: String,
11    pub text: String,
12    pub stable: bool,
13    pub priority: i32,
14    pub source: Option<String>,
15    pub metadata: Metadata,
16}
17
18impl ContextFragment {
19    pub fn new(id: impl Into<String>, text: impl Into<String>) -> Self {
20        Self {
21            id: id.into(),
22            text: text.into(),
23            stable: false,
24            priority: 100,
25            source: None,
26            metadata: Metadata::new(),
27        }
28    }
29
30    pub fn stable(mut self, stable: bool) -> Self {
31        self.stable = stable;
32        self
33    }
34
35    pub fn priority(mut self, priority: i32) -> Self {
36        self.priority = priority;
37        self
38    }
39
40    pub fn source(mut self, source: impl Into<String>) -> Self {
41        self.source = Some(source.into());
42        self
43    }
44
45    pub fn cache_hint(mut self, cache_hint: impl Into<String>) -> Self {
46        self.metadata
47            .insert("cache_hint".to_string(), Value::String(cache_hint.into()));
48        self
49    }
50
51    pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
52        self.metadata.insert(key.into(), value);
53        self
54    }
55}
56
57pub struct ContextRequest<'a> {
58    pub agent_name: &'a str,
59    pub input: &'a str,
60    pub max_prompt_chars: Option<usize>,
61    pub metadata: Metadata,
62}
63
64impl<'a> ContextRequest<'a> {
65    pub fn new(agent_name: &'a str, input: &'a str) -> Self {
66        Self {
67            agent_name,
68            input,
69            max_prompt_chars: None,
70            metadata: Metadata::new(),
71        }
72    }
73
74    pub fn for_test(agent_name: &'a str, input: &'a str) -> Self {
75        Self::new(agent_name, input)
76    }
77
78    pub fn max_prompt_chars(mut self, max_prompt_chars: usize) -> Self {
79        self.max_prompt_chars = Some(max_prompt_chars);
80        self
81    }
82
83    pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
84        self.metadata.insert(key.into(), value);
85        self
86    }
87}
88
89#[derive(Debug, Clone, PartialEq)]
90pub struct ContextSection {
91    pub id: String,
92    pub text: String,
93    pub stable: bool,
94    pub metadata: Metadata,
95}
96
97#[derive(Debug, Clone, PartialEq)]
98pub struct ContextBundle {
99    pub prompt: String,
100    pub sections: Vec<ContextSection>,
101    pub stable_hash: String,
102    pub sources: BTreeMap<String, String>,
103    pub omitted_section_ids: Vec<String>,
104}
105
106pub trait ContextProvider: Send + Sync {
107    fn fragments(&self, request: &ContextRequest<'_>)
108        -> Result<Vec<ContextFragment>, ContextError>;
109}
110
111#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct ContextError {
113    message: String,
114}
115
116impl ContextError {
117    pub fn new(message: impl Into<String>) -> Self {
118        Self {
119            message: message.into(),
120        }
121    }
122}
123
124impl std::fmt::Display for ContextError {
125    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        formatter.write_str(&self.message)
127    }
128}
129
130impl std::error::Error for ContextError {}
131
132pub fn collect_context_fragments(
133    request: &ContextRequest<'_>,
134    providers: &[std::sync::Arc<dyn ContextProvider>],
135) -> Result<Vec<ContextFragment>, ContextError> {
136    let mut fragments = Vec::new();
137    for provider in providers {
138        fragments.extend(provider.fragments(request)?);
139    }
140    Ok(fragments)
141}
142
143pub fn assemble_context_fragments(
144    request: &ContextRequest<'_>,
145    mut fragments: Vec<ContextFragment>,
146) -> Result<ContextBundle, ContextError> {
147    fragments.sort_by(|left, right| {
148        left.priority
149            .cmp(&right.priority)
150            .then_with(|| right.stable.cmp(&left.stable))
151            .then_with(|| left.id.cmp(&right.id))
152    });
153
154    let mut prompt_parts = Vec::new();
155    let mut sections = Vec::new();
156    let mut stable_parts = Vec::new();
157    let mut sources = BTreeMap::new();
158    let mut omitted_section_ids = Vec::new();
159
160    for fragment in fragments {
161        let text = fragment.text.trim().to_string();
162        if text.is_empty() {
163            continue;
164        }
165        let candidate_len = if prompt_parts.is_empty() {
166            text.len()
167        } else {
168            prompt_parts.join("\n\n").len() + 2 + text.len()
169        };
170        if request
171            .max_prompt_chars
172            .is_some_and(|max_chars| candidate_len > max_chars)
173        {
174            omitted_section_ids.push(fragment.id);
175            continue;
176        }
177        if let Some(source) = fragment.source.as_ref() {
178            sources.insert(fragment.id.clone(), source.clone());
179        }
180        if fragment.stable {
181            stable_parts.push(text.clone());
182        }
183        prompt_parts.push(text.clone());
184        sections.push(ContextSection {
185            id: fragment.id,
186            text,
187            stable: fragment.stable,
188            metadata: fragment.metadata,
189        });
190    }
191
192    Ok(ContextBundle {
193        prompt: prompt_parts.join("\n\n"),
194        sections,
195        stable_hash: sha256_hex(stable_parts.join("").as_bytes()),
196        sources,
197        omitted_section_ids,
198    })
199}
200
201fn sha256_hex(bytes: &[u8]) -> String {
202    let mut hasher = Sha256::new();
203    hasher.update(bytes);
204    hex_lower(&hasher.finalize())
205}
206
207fn hex_lower(bytes: &[u8]) -> String {
208    bytes.iter().map(|byte| format!("{byte:02x}")).collect()
209}
vv_agent/context_providers.rs

vv_agent/
context_providers.rs
