1use serde::Serialize;
2use std::path::Path;
3
4use crate::ast::PgnPacket;
5use crate::resolver::{ResolutionStatus, ResolvedRef};
6
7#[derive(Debug, Serialize)]
8pub struct ContextPlan {
9 pub run_id: String,
10 pub primary_refs: Vec<ContextRef>,
11 pub retrieval_methods: Vec<String>,
12 pub token_budget: usize,
13 pub compression_allowed: bool,
14 pub total_bytes: usize,
15 pub estimated_tokens: usize,
16}
17
18#[derive(Debug, Serialize)]
19pub struct ContextRef {
20 pub reference: String,
21 pub resolved: bool,
22 pub path: Option<String>,
23 pub bytes: Option<usize>,
24 pub retrieval_method: String,
25}
26
27fn pick_retrieval_method(path: &Path) -> String {
28 match path.extension().and_then(|e| e.to_str()) {
29 Some("rs" | "py" | "js" | "ts" | "go" | "java" | "c" | "cpp" | "h") => {
30 "full_text".to_string()
31 }
32 Some("yaml" | "yml" | "json" | "toml") => "full_text".to_string(),
33 Some("md" | "txt" | "rst") => "full_text".to_string(),
34 Some("png" | "jpg" | "jpeg" | "gif" | "svg") => "metadata_only".to_string(),
35 Some("pdf") => "metadata_only".to_string(),
36 _ => "full_text".to_string(),
37 }
38}
39
40pub fn build_context_plan(packet: &PgnPacket, resolved: &[ResolvedRef]) -> ContextPlan {
41 let mut total_bytes = 0usize;
42 let primary_refs: Vec<ContextRef> = resolved
43 .iter()
44 .map(|r| {
45 let (bytes, method) = match (&r.resolved_path, r.status) {
46 (Some(path), ResolutionStatus::Resolved) => {
47 let size = std::fs::metadata(path)
48 .map(|m| m.len() as usize)
49 .unwrap_or(0);
50 let method = pick_retrieval_method(path);
51 (Some(size), method)
52 }
53 _ => (None, "full_text".to_string()),
54 };
55 if let Some(b) = bytes {
56 total_bytes += b;
57 }
58 ContextRef {
59 reference: r.original.clone(),
60 resolved: matches!(r.status, ResolutionStatus::Resolved),
61 path: r.resolved_path.as_ref().map(|p| p.display().to_string()),
62 bytes,
63 retrieval_method: method,
64 }
65 })
66 .collect();
67
68 let estimated_tokens = total_bytes.div_ceil(4);
69 let token_budget = estimated_tokens.max(1000);
70
71 let mut methods: Vec<String> = primary_refs
73 .iter()
74 .map(|r| r.retrieval_method.clone())
75 .collect::<std::collections::HashSet<_>>()
76 .into_iter()
77 .collect();
78 methods.sort();
79
80 ContextPlan {
81 run_id: packet.run_id.clone(),
82 primary_refs,
83 retrieval_methods: methods,
84 token_budget,
85 compression_allowed: total_bytes > 50_000,
86 total_bytes,
87 estimated_tokens,
88 }
89}