Skip to main content

pidgin_lang/
context.rs

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    // Pick retrieval methods used across all refs
72    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}