PicoBot/src/agent/media_handler.rs

102 lines
2.7 KiB
Rust

use std::collections::HashMap;
use std::fmt;
use std::io::Read;
use crate::bus::message::ContentBlock;
pub trait MediaHandler: Send + Sync {
fn media_type(&self) -> &str;
fn handle(&self, path: &str) -> Result<Vec<ContentBlock>, MediaHandlerError>;
}
#[derive(Debug)]
pub enum MediaHandlerError {
Io(std::io::Error),
UnsupportedFormat(String),
}
impl fmt::Display for MediaHandlerError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
MediaHandlerError::Io(e) => write!(f, "I/O error: {}", e),
MediaHandlerError::UnsupportedFormat(msg) => write!(f, "Unsupported format: {}", msg),
}
}
}
impl std::error::Error for MediaHandlerError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
MediaHandlerError::Io(e) => Some(e),
MediaHandlerError::UnsupportedFormat(_) => None,
}
}
}
pub struct ImageHandler;
impl MediaHandler for ImageHandler {
fn media_type(&self) -> &str {
"image"
}
fn handle(&self, path: &str) -> Result<Vec<ContentBlock>, MediaHandlerError> {
let (mime_type, base64_data) =
encode_image_to_base64(path).map_err(MediaHandlerError::Io)?;
let url = format!("data:{};base64,{}", mime_type, base64_data);
Ok(vec![ContentBlock::image_url(url)])
}
}
fn encode_image_to_base64(path: &str) -> Result<(String, String), std::io::Error> {
use base64::{engine::general_purpose::STANDARD, Engine as _};
let mut file = std::fs::File::open(path)?;
let mut buffer = Vec::new();
file.read_to_end(&mut buffer)?;
let mime = mime_guess::from_path(path)
.first_or_octet_stream()
.to_string();
let encoded = STANDARD.encode(&buffer);
Ok((mime, encoded))
}
pub struct MediaHandlerRegistry {
handlers: HashMap<String, Box<dyn MediaHandler>>,
}
impl MediaHandlerRegistry {
pub fn new() -> Self {
Self {
handlers: HashMap::new(),
}
}
pub fn register(&mut self, handler: Box<dyn MediaHandler>) {
self.handlers
.insert(handler.media_type().to_string(), handler);
}
pub fn handle(
&self,
media_type: &str,
path: &str,
) -> Result<Vec<ContentBlock>, MediaHandlerError> {
match self.handlers.get(media_type) {
Some(handler) => handler.handle(path),
None => Err(MediaHandlerError::UnsupportedFormat(format!(
"no handler for type: {}",
media_type
))),
}
}
pub fn with_defaults() -> Self {
let mut reg = Self::new();
reg.register(Box::new(ImageHandler));
reg
}
}