llm_bucket/
preprocess.rs

1use crate::code_to_pdf::{code_file_to_pdf, CodeToPdfError};
2use crate::contract::{
3    ExternalItemInput, ExternalSourceInput, ProcessConfig, ProcessError, ProcessInput,
4    ProcessorKind,
5};
6use tempfile;
7use tracing::{debug, error, info};
8
9/// Main processor struct for CLI usage: implements Preprocessor.
10pub struct Processor {
11    config: ProcessConfig,
12}
13
14impl Processor {
15    pub fn new(config: ProcessConfig) -> Self {
16        Self { config }
17    }
18}
19
20#[async_trait::async_trait]
21impl crate::contract::Preprocessor for Processor {
22    async fn process(&self, input: ProcessInput) -> Result<ExternalSourceInput, ProcessError> {
23        self.process_sync(input)
24    }
25}
26
27impl Processor {
28    /// Synchronous process logic for unit tests and internal use.
29    pub fn process_sync(&self, input: ProcessInput) -> Result<ExternalSourceInput, ProcessError> {
30        let config = &self.config;
31        info!(processor = ?config.kind, name = input.name, "Starting processing for source");
32        let result = match config.kind {
33            ProcessorKind::ReadmeToPDF => process_readme_to_pdf(input),
34            ProcessorKind::FlattenFiles => process_flatten_files(input),
35            // Add more processor kinds as needed
36        };
37        match &result {
38            Ok(ext) => info!(
39                items = ext.external_items.len(),
40                "Processing completed successfully"
41            ),
42            Err(e) => error!(error = ?e, "Processing failed"),
43        };
44        result
45    }
46}
47
48/// (No longer needed free function for process -- use Processor::process_sync or the trait.)
49
50fn process_readme_to_pdf(input: ProcessInput) -> Result<ExternalSourceInput, ProcessError> {
51    let readme_path = input.repo_path.join("README.md");
52    debug!(repo_path = %input.repo_path.display(), "Looking for README.md in repo path");
53
54    if !readme_path.exists() {
55        error!(path = %readme_path.display(), "No README.md found in repository");
56        return Err(ProcessError::NoReadme);
57    }
58
59    // Prepare a temp output file path for pdf generation
60    let tmp_pdf = tempfile::NamedTempFile::new().map_err(|e| {
61        error!(error = ?e, "Failed to create temp file for PDF output");
62        ProcessError::Io(e)
63    })?;
64    let tmp_pdf_path = tmp_pdf.path();
65
66    // Call the code_to_pdf module (on-disk)
67    code_file_to_pdf(&readme_path, tmp_pdf_path)
68        .map_err(|e| {
69            match &e {
70                CodeToPdfError::Io(err) => error!(path = %readme_path.display(), error = ?err, "IO error during PDF generation"),
71                CodeToPdfError::EmptyInput => error!("Attempted PDF generation with empty input"),
72                CodeToPdfError::Font(desc) => error!(desc = *desc, "Font error during PDF generation"),
73            }
74            match e {
75                CodeToPdfError::Io(e) => ProcessError::Io(e),
76                CodeToPdfError::EmptyInput => ProcessError::Other("PDF: Empty input".into()),
77                CodeToPdfError::Font(_) => ProcessError::Other("PDF: font error".into()),
78            }
79        })?;
80
81    // Read PDF as bytes
82    let content = std::fs::read(tmp_pdf_path).map_err(|e| {
83        error!(error = ?e, path = %tmp_pdf_path.display(), "Failed to read generated PDF from disk");
84        ProcessError::Io(e)
85    })?;
86
87    // Prepare the result structures
88    let ext_item = ExternalItemInput {
89        filename: "README.pdf".to_string(),
90        content,
91    };
92
93    info!(
94        filename = "README.pdf",
95        size = ext_item.content.len(),
96        "Generated README.pdf from README.md"
97    );
98    Ok(ExternalSourceInput {
99        name: input.name,
100        external_items: vec![ext_item],
101    })
102}
103
104/// Recursively flatten all files and output as items with "__" as directory separator.
105fn process_flatten_files(input: ProcessInput) -> Result<ExternalSourceInput, ProcessError> {
106    info!(path = %input.repo_path.display(), "Flattening files in repository");
107    let mut external_items = Vec::new();
108    let repo_path = &input.repo_path;
109    let _base_len = repo_path.components().count();
110
111    fn visit_dir(
112        dir: &std::path::Path,
113        repo_path: &std::path::Path,
114        results: &mut Vec<ExternalItemInput>,
115    ) -> Result<(), ProcessError> {
116        for entry_res in std::fs::read_dir(dir)? {
117            let entry = entry_res?;
118            let path = entry.path();
119            if path.is_dir() {
120                // Skip .git and target directories
121                let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
122                if file_name == ".git" || file_name == "target" {
123                    debug!(path = %path.display(), "Skipping directory");
124                    continue;
125                }
126                visit_dir(&path, repo_path, results)?;
127            } else if path.is_file() {
128                // compute a flat filename with "__" as a separator, with truncation logic
129                let rel_path = path.strip_prefix(repo_path).unwrap();
130                let mut segments: Vec<String> = Vec::new();
131                for comp in rel_path.components() {
132                    segments.push(comp.as_os_str().to_string_lossy().into_owned());
133                }
134                if segments.is_empty() {
135                    continue;
136                }
137                let basename = segments.pop().unwrap();
138                let mut joined: String;
139                let max_len = 180;
140                // Try to include as many trailing segments as possible
141                let mut from = 0;
142                loop {
143                    joined = if segments.len() > from {
144                        segments[from..].join("__") + "__" + &basename
145                    } else {
146                        basename.clone()
147                    };
148                    if joined.len() <= max_len || from >= segments.len() {
149                        break;
150                    }
151                    from += 1;
152                }
153                let flat_name = joined;
154                match std::fs::read(&path) {
155                    Ok(content) => {
156                        debug!(filename = %flat_name, size = content.len(), "Flattened file");
157                        results.push(ExternalItemInput {
158                            filename: flat_name,
159                            content,
160                        });
161                    }
162                    Err(e) => {
163                        error!(error = ?e, path = %path.display(), "Failed to read file while flattening");
164                        return Err(ProcessError::Io(e));
165                    }
166                }
167            }
168        }
169        Ok(())
170    }
171    if let Err(e) = visit_dir(repo_path, repo_path, &mut external_items) {
172        error!(error = ?e, "Error occurred during directory flattening");
173        return Err(e);
174    }
175
176    info!(
177        count = external_items.len(),
178        "Completed flattening files in repository"
179    );
180    Ok(ExternalSourceInput {
181        name: input.name,
182        external_items,
183    })
184}
llm_bucket/preprocess.rs

llm_bucket/
preprocess.rs
