use anyhow::Result;
use std::fs;
use std::path::Path;
use super::modifier::Modifier;
use super::modifier::ModifierKey;
use crate::models::context::ContextModel;
use crate::utils;
pub struct TextModifier;
impl TextModifier {
fn html_to_text(value: &str) -> Result<Option<String>> {
let lower = value.to_lowercase();
if !(lower.contains("<html") || lower.contains("<!doctype") || lower.contains("</div>")) {
return Ok(None);
}
Ok(Some(html2text::from_read(value.as_bytes(), 120)?))
}
fn pdf_to_text(bytes: &[u8]) -> Result<Option<String>> {
if bytes.starts_with(b"%PDF") {
let text = pdf_extract::extract_text_from_mem(bytes)
.map_err(|e| anyhow::anyhow!("PDF text extraction failed: {}", e))?;
return Ok(Some(text));
}
Ok(None)
}
fn img_to_text(bytes: &[u8]) -> Result<Option<String>> {
let is_webp = bytes.starts_with(b"RIFF") && bytes.len() > 12 && &bytes[8..12] == b"WEBP";
if bytes.starts_with(b"\x89PNG")
|| bytes.starts_with(b"\xFF\xD8\xFF")
|| bytes.starts_with(b"GIF8")
|| is_webp
{
use base64::Engine;
let base64 = base64::engine::general_purpose::STANDARD.encode(bytes);
return Ok(Some(base64));
}
Ok(None)
}
}
impl Modifier for TextModifier {
fn key(&self) -> ModifierKey {
ModifierKey::Text
}
fn apply(&self, value: &ContextModel, _arg: &str) -> Result<ContextModel> {
match value {
ContextModel::String(s) => {
if crate::utils::image::is_image(s) {
return Ok(ContextModel::String(s.clone()));
}
if let Some(text) = Self::html_to_text(s)? {
return Ok(ContextModel::String(text));
}
if let Ok(path) = utils::path::resolve(Path::new(s)) {
if path.exists() && !path.is_dir() {
let bytes = fs::read(&path)?;
if let Some(text) = Self::img_to_text(&bytes)? {
return Ok(ContextModel::String(text));
}
if let Some(text) = Self::pdf_to_text(&bytes)? {
return Ok(ContextModel::String(text));
}
let content = String::from_utf8_lossy(&bytes).to_string();
if let Some(text) = Self::html_to_text(&content)? {
return Ok(ContextModel::String(text));
}
return Ok(ContextModel::String(content));
}
}
anyhow::bail!("File not found or unsupported content: '{}'", s)
}
ContextModel::List(items) => {
let results: Vec<String> = items
.iter()
.map(
|i| match self.apply(&ContextModel::String(i.clone()), "")? {
ContextModel::String(s) => Ok(s),
_ => unreachable!(),
},
)
.collect::<Result<Vec<_>>>()?;
Ok(ContextModel::List(results))
}
}
}
}