opendev_context/
context_picker.rs

1//! Dynamic context selection for LLM calls.
2//!
3//! Provides the data models and types for context picking. The actual
4//! picking logic (file injection, playbook strategies, etc.) depends on
5//! higher-level crates and is wired up at the application layer.
6//!
7//! All decisions are logged as `ContextReason` objects for full traceability.
8
9use std::collections::HashMap;
10use std::fmt;
11use std::path::Path;
12
13use serde::{Deserialize, Serialize};
14
15/// Category of context piece for organization and filtering.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
17#[serde(rename_all = "snake_case")]
18pub enum ContextCategory {
19    SystemPrompt,
20    FileReference,
21    DirectoryListing,
22    ConversationHistory,
23    PlaybookStrategy,
24    ImageContent,
25    PdfContent,
26    ToolResult,
27    UserQuery,
28}
29
30impl ContextCategory {
31    pub fn as_str(&self) -> &'static str {
32        match self {
33            Self::SystemPrompt => "system_prompt",
34            Self::FileReference => "file_reference",
35            Self::DirectoryListing => "directory_listing",
36            Self::ConversationHistory => "conversation_history",
37            Self::PlaybookStrategy => "playbook_strategy",
38            Self::ImageContent => "image_content",
39            Self::PdfContent => "pdf_content",
40            Self::ToolResult => "tool_result",
41            Self::UserQuery => "user_query",
42        }
43    }
44}
45
46impl fmt::Display for ContextCategory {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        f.write_str(self.as_str())
49    }
50}
51
52/// Documents why a context piece was included.
53///
54/// This is the key to traceability -- every piece of context should have
55/// a clear reason for inclusion that can be logged and debugged.
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct ContextReason {
58    pub source: String,
59    pub reason: String,
60    #[serde(default = "default_relevance")]
61    pub relevance_score: f64,
62    #[serde(default)]
63    pub tokens_estimate: usize,
64    #[serde(default)]
65    pub metadata: HashMap<String, serde_json::Value>,
66}
67
68fn default_relevance() -> f64 {
69    1.0
70}
71
72impl ContextReason {
73    pub fn new(source: &str, reason: &str) -> Self {
74        Self {
75            source: source.to_string(),
76            reason: reason.to_string(),
77            relevance_score: 1.0,
78            tokens_estimate: 0,
79            metadata: HashMap::new(),
80        }
81    }
82
83    pub fn with_tokens(mut self, tokens: usize) -> Self {
84        self.tokens_estimate = tokens;
85        self
86    }
87
88    pub fn with_score(mut self, score: f64) -> Self {
89        self.relevance_score = score;
90        self
91    }
92}
93
94impl fmt::Display for ContextReason {
95    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96        let score_str = if self.relevance_score < 1.0 {
97            format!(" (score={:.2})", self.relevance_score)
98        } else {
99            String::new()
100        };
101        let tokens_str = if self.tokens_estimate > 0 {
102            format!(" [{} tokens]", self.tokens_estimate)
103        } else {
104            String::new()
105        };
106        write!(
107            f,
108            "[{}]{}{}: {}",
109            self.source, score_str, tokens_str, self.reason
110        )
111    }
112}
113
114/// A single piece of context to include in the LLM call.
115#[derive(Debug, Clone, Serialize, Deserialize)]
116pub struct ContextPiece {
117    pub content: String,
118    pub reason: ContextReason,
119    pub category: ContextCategory,
120    /// Ordering hint (lower = earlier in context).
121    #[serde(default = "default_order")]
122    pub order: i32,
123}
124
125fn default_order() -> i32 {
126    100
127}
128
129impl ContextPiece {
130    pub fn new(content: String, reason: ContextReason, category: ContextCategory) -> Self {
131        Self {
132            content,
133            reason,
134            category,
135            order: 100,
136        }
137    }
138
139    pub fn with_order(mut self, order: i32) -> Self {
140        self.order = order;
141        self
142    }
143
144    /// Estimated token count (from reason or calculated).
145    pub fn tokens_estimate(&self) -> usize {
146        if self.reason.tokens_estimate > 0 {
147            return self.reason.tokens_estimate;
148        }
149        self.content.len() / 4
150    }
151}
152
153impl fmt::Display for ContextPiece {
154    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
155        let preview: String = self.content.chars().take(50).collect();
156        let preview = preview.replace('\n', "\\n");
157        let ellipsis = if self.content.len() > 50 { "..." } else { "" };
158        write!(f, "{}: {}{}", self.category, preview, ellipsis)
159    }
160}
161
162/// Final assembled context ready for LLM call.
163///
164/// Contains everything needed for an LLM call plus traceability information.
165#[derive(Debug, Clone, Default, Serialize, Deserialize)]
166pub struct AssembledContext {
167    pub system_prompt: String,
168    pub messages: Vec<serde_json::Map<String, serde_json::Value>>,
169    #[serde(default)]
170    pub pieces: Vec<ContextPiece>,
171    #[serde(default)]
172    pub image_blocks: Vec<serde_json::Value>,
173    #[serde(default)]
174    pub total_tokens_estimate: usize,
175}
176
177impl AssembledContext {
178    /// Return concise summary of context for display.
179    pub fn summary(&self) -> String {
180        let mut by_category: HashMap<ContextCategory, Vec<&ContextPiece>> = HashMap::new();
181        for piece in &self.pieces {
182            by_category.entry(piece.category).or_default().push(piece);
183        }
184
185        let mut parts = Vec::new();
186        for (category, pieces) in &by_category {
187            let total_tokens: usize = pieces.iter().map(|p| p.tokens_estimate()).sum();
188            parts.push(format!("{}: ~{} tokens", category, total_tokens));
189        }
190
191        let mut summary = format!("Context: {} tokens", self.total_tokens_estimate);
192        if !parts.is_empty() {
193            summary = format!("{} ({})", summary, parts.join(", "));
194        }
195        summary
196    }
197}
198
199/// Simple tracer for context selection decisions.
200pub struct ContextTracer;
201
202impl ContextTracer {
203    pub fn new() -> Self {
204        Self
205    }
206
207    pub fn trace(&self, context: &AssembledContext) {
208        tracing::debug!("[ContextPicker] {}", context.summary());
209    }
210
211    /// Export trace to a JSON file for debugging.
212    pub fn export_trace(&self, context: &AssembledContext, path: &Path) -> std::io::Result<()> {
213        let trace_data = serde_json::json!({
214            "timestamp": chrono::Utc::now().to_rfc3339(),
215            "total_tokens_estimate": context.total_tokens_estimate,
216            "message_count": context.messages.len(),
217            "piece_count": context.pieces.len(),
218            "image_count": context.image_blocks.len(),
219            "pieces": context.pieces.iter().map(|p| {
220                serde_json::json!({
221                    "category": p.category.as_str(),
222                    "source": p.reason.source,
223                    "tokens_estimate": p.tokens_estimate(),
224                })
225            }).collect::<Vec<_>>(),
226        });
227
228        if let Some(parent) = path.parent() {
229            std::fs::create_dir_all(parent)?;
230        }
231        let file = std::fs::File::create(path)?;
232        serde_json::to_writer_pretty(file, &trace_data)?;
233        tracing::debug!("Context trace exported to {}", path.display());
234        Ok(())
235    }
236}
237
238impl Default for ContextTracer {
239    fn default() -> Self {
240        Self::new()
241    }
242}
243
244#[cfg(test)]
245#[path = "context_picker_tests.rs"]
246mod tests;
opendev_context/context_picker.rs

opendev_context/
context_picker.rs
