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