Skip to main content

topo_render/
jsonl.rs

1use serde::Serialize;
2use std::io::Write;
3use topo_core::ScoredFile;
4
5/// Writes scored files in JSONL v0.3 format.
6pub struct JsonlWriter {
7    query: String,
8    preset: String,
9    max_bytes: Option<u64>,
10    min_score: f64,
11    hint: Option<String>,
12}
13
14#[derive(Serialize)]
15#[serde(rename_all = "PascalCase")]
16struct Header {
17    version: String,
18    query: String,
19    preset: String,
20    budget: Budget,
21    min_score: f64,
22}
23
24#[derive(Serialize)]
25#[serde(rename_all = "PascalCase")]
26struct Budget {
27    #[serde(skip_serializing_if = "Option::is_none")]
28    max_bytes: Option<u64>,
29}
30
31#[derive(Serialize)]
32#[serde(rename_all = "PascalCase")]
33struct FileEntry {
34    path: String,
35    score: f64,
36    tokens: u64,
37    language: String,
38    role: String,
39}
40
41#[derive(Serialize)]
42#[serde(rename_all = "PascalCase")]
43struct Footer {
44    total_files: usize,
45    total_tokens: u64,
46    scanned_files: usize,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    hint: Option<String>,
49}
50
51impl JsonlWriter {
52    pub fn new(query: &str, preset: &str) -> Self {
53        Self {
54            query: query.to_string(),
55            preset: preset.to_string(),
56            max_bytes: None,
57            min_score: 0.0,
58            hint: None,
59        }
60    }
61
62    pub fn max_bytes(mut self, max_bytes: Option<u64>) -> Self {
63        self.max_bytes = max_bytes;
64        self
65    }
66
67    pub fn min_score(mut self, min_score: f64) -> Self {
68        self.min_score = min_score;
69        self
70    }
71
72    pub fn hint(mut self, hint: Option<String>) -> Self {
73        self.hint = hint;
74        self
75    }
76
77    /// Render scored files as JSONL v0.3 string.
78    pub fn render(&self, files: &[ScoredFile], scanned_count: usize) -> anyhow::Result<String> {
79        let mut buf = Vec::new();
80        self.write_to(&mut buf, files, scanned_count)?;
81        Ok(String::from_utf8(buf)?)
82    }
83
84    /// Write JSONL v0.3 output to a writer.
85    pub fn write_to(
86        &self,
87        writer: &mut dyn Write,
88        files: &[ScoredFile],
89        scanned_count: usize,
90    ) -> anyhow::Result<()> {
91        // Header
92        let header = Header {
93            version: "0.3".to_string(),
94            query: self.query.clone(),
95            preset: self.preset.clone(),
96            budget: Budget {
97                max_bytes: self.max_bytes,
98            },
99            min_score: self.min_score,
100        };
101        serde_json::to_writer(&mut *writer, &header)?;
102        writeln!(writer)?;
103
104        // File entries
105        let mut total_tokens = 0u64;
106        for file in files {
107            let entry = FileEntry {
108                path: file.path.clone(),
109                score: file.score,
110                tokens: file.tokens,
111                language: file.language.as_str().to_string(),
112                role: file.role.as_str().to_string(),
113            };
114            serde_json::to_writer(&mut *writer, &entry)?;
115            writeln!(writer)?;
116            total_tokens += file.tokens;
117        }
118
119        // Footer
120        let footer = Footer {
121            total_files: files.len(),
122            total_tokens,
123            scanned_files: scanned_count,
124            hint: self.hint.clone(),
125        };
126        serde_json::to_writer(&mut *writer, &footer)?;
127        writeln!(writer)?;
128
129        Ok(())
130    }
131}