kreuzberg 4.0.8 → 4.1.0

data/vendor/kreuzberg/src/api/startup.rs

@@ -0,0 +1,243 @@
+//! API server startup functions.
+
+use std::net::{IpAddr, SocketAddr};
+
+use crate::{ExtractionConfig, Result, core::ServerConfig};
+
+use super::{config::load_server_config, router::create_router_with_limits_and_server_config, types::ApiSizeLimits};
+
+/// Start the API server with config file discovery.
+///
+/// Searches for kreuzberg.toml/yaml/json in current and parent directories.
+/// If no config file is found, uses default configuration.
+///
+/// # Arguments
+///
+/// * `host` - IP address to bind to (e.g., "127.0.0.1" or "0.0.0.0")
+/// * `port` - Port number to bind to (e.g., 8000)
+///
+/// # Examples
+///
+/// ```no_run
+/// use kreuzberg::api::serve;
+///
+/// #[tokio::main]
+/// async fn main() -> kreuzberg::Result<()> {
+///     // Local development
+///     serve("127.0.0.1", 8000).await?;
+///     Ok(())
+/// }
+/// ```
+///
+/// ```no_run
+/// use kreuzberg::api::serve;
+///
+/// #[tokio::main]
+/// async fn main() -> kreuzberg::Result<()> {
+///     // Docker/production (listen on all interfaces)
+///     serve("0.0.0.0", 8000).await?;
+///     Ok(())
+/// }
+/// ```
+///
+/// # Environment Variables
+///
+/// ```bash
+/// # Python/Docker usage
+/// export KREUZBERG_HOST=0.0.0.0
+/// export KREUZBERG_PORT=8000
+///
+/// # CORS configuration (IMPORTANT for production security)
+/// # Default: allows all origins (permits CSRF attacks)
+/// # Production: set to comma-separated list of allowed origins
+/// export KREUZBERG_CORS_ORIGINS="https://app.example.com,https://api.example.com"
+///
+/// # Upload size limits (default: 100 MB)
+/// # Modern approach (in bytes):
+/// export KREUZBERG_MAX_REQUEST_BODY_BYTES=104857600       # 100 MB
+/// export KREUZBERG_MAX_MULTIPART_FIELD_BYTES=104857600    # 100 MB per file
+///
+/// # Legacy approach (in MB, applies to both limits):
+/// export KREUZBERG_MAX_UPLOAD_SIZE_MB=100  # 100 MB
+///
+/// python -m kreuzberg.api
+/// ```
+pub async fn serve(host: impl AsRef<str>, port: u16) -> Result<()> {
+    let extraction_config = match ExtractionConfig::discover()? {
+        Some(config) => {
+            tracing::info!("Loaded extraction config from discovered file");
+            config
+        }
+        None => {
+            tracing::info!("No config file found, using default configuration");
+            ExtractionConfig::default()
+        }
+    };
+
+    let server_config = load_server_config(None)?;
+    let limits = ApiSizeLimits::new(
+        server_config.max_request_body_bytes,
+        server_config.max_multipart_field_bytes,
+    );
+
+    serve_with_config_and_limits(host, port, extraction_config, limits).await
+}
+
+/// Start the API server with explicit config.
+///
+/// Uses default size limits (100 MB). For custom limits, use `serve_with_config_and_limits`.
+///
+/// # Arguments
+///
+/// * `host` - IP address to bind to (e.g., "127.0.0.1" or "0.0.0.0")
+/// * `port` - Port number to bind to (e.g., 8000)
+/// * `config` - Default extraction configuration for all requests
+///
+/// # Examples
+///
+/// ```no_run
+/// use kreuzberg::{ExtractionConfig, api::serve_with_config};
+///
+/// #[tokio::main]
+/// async fn main() -> kreuzberg::Result<()> {
+///     let config = ExtractionConfig::from_toml_file("config/kreuzberg.toml")?;
+///     serve_with_config("127.0.0.1", 8000, config).await?;
+///     Ok(())
+/// }
+/// ```
+pub async fn serve_with_config(host: impl AsRef<str>, port: u16, config: ExtractionConfig) -> Result<()> {
+    let limits = ApiSizeLimits::default();
+    tracing::info!(
+        "Upload size limit: 100 MB (default, {} bytes)",
+        limits.max_request_body_bytes
+    );
+    serve_with_config_and_limits(host, port, config, limits).await
+}
+
+/// Start the API server with explicit config and size limits.
+///
+/// # Arguments
+///
+/// * `host` - IP address to bind to (e.g., "127.0.0.1" or "0.0.0.0")
+/// * `port` - Port number to bind to (e.g., 8000)
+/// * `config` - Default extraction configuration for all requests
+/// * `limits` - Size limits for request bodies and multipart uploads
+///
+/// # Examples
+///
+/// ```no_run
+/// use kreuzberg::{ExtractionConfig, api::{serve_with_config_and_limits, ApiSizeLimits}};
+///
+/// #[tokio::main]
+/// async fn main() -> kreuzberg::Result<()> {
+///     let config = ExtractionConfig::from_toml_file("config/kreuzberg.toml")?;
+///     let limits = ApiSizeLimits::from_mb(200, 200);
+///     serve_with_config_and_limits("127.0.0.1", 8000, config, limits).await?;
+///     Ok(())
+/// }
+/// ```
+pub async fn serve_with_config_and_limits(
+    host: impl AsRef<str>,
+    port: u16,
+    config: ExtractionConfig,
+    limits: ApiSizeLimits,
+) -> Result<()> {
+    let ip: IpAddr = host
+        .as_ref()
+        .parse()
+        .map_err(|e| crate::error::KreuzbergError::validation(format!("Invalid host address: {}", e)))?;
+
+    let server_config = ServerConfig {
+        host: host.as_ref().to_string(),
+        port,
+        max_request_body_bytes: limits.max_request_body_bytes,
+        max_multipart_field_bytes: limits.max_multipart_field_bytes,
+        ..Default::default()
+    };
+
+    let addr = SocketAddr::new(ip, port);
+    let app = create_router_with_limits_and_server_config(config, limits, server_config);
+
+    tracing::info!("Starting Kreuzberg API server on http://{}:{}", ip, port);
+
+    let listener = tokio::net::TcpListener::bind(addr)
+        .await
+        .map_err(crate::error::KreuzbergError::Io)?;
+
+    axum::serve(listener, app)
+        .await
+        .map_err(|e| crate::error::KreuzbergError::Other(e.to_string()))?;
+
+    Ok(())
+}
+
+/// Start the API server with explicit extraction config and server config.
+///
+/// This function accepts a fully-configured ServerConfig, including CORS origins,
+/// size limits, host, and port. It respects all ServerConfig fields without
+/// re-parsing environment variables, making it ideal for CLI usage where
+/// configuration precedence has already been applied.
+///
+/// # Arguments
+///
+/// * `extraction_config` - Default extraction configuration for all requests
+/// * `server_config` - Server configuration including host, port, CORS, and size limits
+///
+/// # Examples
+///
+/// ```no_run
+/// use kreuzberg::{ExtractionConfig, api::serve_with_server_config, core::ServerConfig};
+///
+/// #[tokio::main]
+/// async fn main() -> kreuzberg::Result<()> {
+///     let extraction_config = ExtractionConfig::default();
+///     let mut server_config = ServerConfig::default();
+///     server_config.host = "0.0.0.0".to_string();
+///     server_config.port = 3000;
+///     server_config.cors_origins = vec!["https://example.com".to_string()];
+///
+///     serve_with_server_config(extraction_config, server_config).await?;
+///     Ok(())
+/// }
+/// ```
+pub async fn serve_with_server_config(extraction_config: ExtractionConfig, server_config: ServerConfig) -> Result<()> {
+    let ip: IpAddr = server_config
+        .host
+        .parse()
+        .map_err(|e| crate::error::KreuzbergError::validation(format!("Invalid host address: {}", e)))?;
+
+    let limits = ApiSizeLimits::new(
+        server_config.max_request_body_bytes,
+        server_config.max_multipart_field_bytes,
+    );
+
+    let addr = SocketAddr::new(ip, server_config.port);
+    let app = create_router_with_limits_and_server_config(extraction_config, limits, server_config.clone());
+
+    tracing::info!(
+        "Starting Kreuzberg API server on http://{}:{} (request_body_limit={} MB, multipart_field_limit={} MB)",
+        ip,
+        server_config.port,
+        server_config.max_request_body_mb(),
+        server_config.max_multipart_field_mb()
+    );
+
+    let listener = tokio::net::TcpListener::bind(addr)
+        .await
+        .map_err(crate::error::KreuzbergError::Io)?;
+
+    axum::serve(listener, app)
+        .await
+        .map_err(|e| crate::error::KreuzbergError::Other(e.to_string()))?;
+
+    Ok(())
+}
+
+/// Start the API server with default host and port.
+///
+/// Defaults: host = "127.0.0.1", port = 8000
+///
+/// Uses config file discovery (searches current/parent directories for kreuzberg.toml/yaml/json).
+pub async fn serve_default() -> Result<()> {
+    serve("127.0.0.1", 8000).await
+}

data/vendor/kreuzberg/src/api/types.rs

@@ -204,3 +204,81 @@ pub struct EmbedResponse {
     /// Number of embeddings generated
     pub count: usize,
 }
+
+/// Default chunker type.
+fn default_chunker_type() -> String {
+    "text".to_string()
+}
+
+/// Chunk request with text and configuration.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChunkRequest {
+    /// Text to chunk
+    pub text: String,
+    /// Optional chunking configuration
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub config: Option<ChunkingConfigRequest>,
+    /// Chunker type (text or markdown)
+    #[serde(default = "default_chunker_type")]
+    pub chunker_type: String,
+}
+
+/// Chunking configuration request.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ChunkingConfigRequest {
+    /// Maximum characters per chunk
+    pub max_characters: Option<usize>,
+    /// Overlap between chunks in characters
+    pub overlap: Option<usize>,
+    /// Whether to trim whitespace
+    pub trim: Option<bool>,
+}
+
+/// Chunk response with chunks and metadata.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChunkResponse {
+    /// List of chunks
+    pub chunks: Vec<ChunkItem>,
+    /// Total number of chunks
+    pub chunk_count: usize,
+    /// Configuration used for chunking
+    pub config: ChunkingConfigResponse,
+    /// Input text size in bytes
+    pub input_size_bytes: usize,
+    /// Chunker type used for chunking
+    pub chunker_type: String,
+}
+
+/// Individual chunk item with metadata.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChunkItem {
+    /// Chunk content
+    pub content: String,
+    /// Byte offset start position
+    pub byte_start: usize,
+    /// Byte offset end position
+    pub byte_end: usize,
+    /// Index of this chunk (0-based)
+    pub chunk_index: usize,
+    /// Total number of chunks
+    pub total_chunks: usize,
+    /// First page number (optional, for PDF chunking)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub first_page: Option<usize>,
+    /// Last page number (optional, for PDF chunking)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub last_page: Option<usize>,
+}
+
+/// Chunking configuration response.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChunkingConfigResponse {
+    /// Maximum characters per chunk
+    pub max_characters: usize,
+    /// Overlap between chunks in characters
+    pub overlap: usize,
+    /// Whether whitespace was trimmed
+    pub trim: bool,
+    /// Type of chunker used
+    pub chunker_type: String,
+}

data/vendor/kreuzberg/src/cache/cleanup.rs

@@ -0,0 +1,277 @@
+//! Cache cleanup operations for managing cache size and age.
+
+use crate::error::{KreuzbergError, Result};
+use std::fs;
+use std::path::Path;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use super::core::{CacheEntry, CacheScanResult, CacheStats};
+use super::utilities::get_available_disk_space;
+
+pub(super) fn scan_cache_directory(cache_dir: &str) -> Result<CacheScanResult> {
+    let dir_path = Path::new(cache_dir);
+
+    if !dir_path.exists() {
+        return Ok(CacheScanResult {
+            stats: CacheStats {
+                total_files: 0,
+                total_size_mb: 0.0,
+                available_space_mb: get_available_disk_space(cache_dir)?,
+                oldest_file_age_days: 0.0,
+                newest_file_age_days: 0.0,
+            },
+            entries: Vec::new(),
+        });
+    }
+
+    let current_time = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs() as f64;
+
+    let read_dir =
+        fs::read_dir(dir_path).map_err(|e| KreuzbergError::cache(format!("Failed to read cache directory: {}", e)))?;
+
+    let mut total_size = 0u64;
+    let mut oldest_age = 0.0f64;
+    let mut newest_age = f64::INFINITY;
+    let mut entries = Vec::new();
+
+    for entry in read_dir {
+        let entry = match entry {
+            Ok(e) => e,
+            Err(e) => {
+                tracing::debug!("Error reading cache entry: {}", e);
+                continue;
+            }
+        };
+
+        let metadata = match entry.metadata() {
+            Ok(m) if m.is_file() => m,
+            _ => continue,
+        };
+
+        let path = entry.path();
+        if path.extension().and_then(|s| s.to_str()) != Some("msgpack") {
+            continue;
+        }
+
+        let modified = match metadata.modified() {
+            Ok(m) => m,
+            Err(e) => {
+                tracing::debug!("Error getting modification time for {:?}: {}", path, e);
+                continue;
+            }
+        };
+
+        let size = metadata.len();
+        total_size += size;
+
+        if let Ok(duration) = modified.duration_since(UNIX_EPOCH) {
+            let age_days = (current_time - duration.as_secs() as f64) / (24.0 * 3600.0);
+            oldest_age = oldest_age.max(age_days);
+            newest_age = newest_age.min(age_days);
+        }
+
+        entries.push(CacheEntry { path, size, modified });
+    }
+
+    if entries.is_empty() {
+        oldest_age = 0.0;
+        newest_age = 0.0;
+    }
+
+    Ok(CacheScanResult {
+        stats: CacheStats {
+            total_files: entries.len(),
+            total_size_mb: total_size as f64 / (1024.0 * 1024.0),
+            available_space_mb: get_available_disk_space(cache_dir)?,
+            oldest_file_age_days: oldest_age,
+            newest_file_age_days: newest_age,
+        },
+        entries,
+    })
+}
+
+pub fn get_cache_metadata(cache_dir: &str) -> Result<CacheStats> {
+    let scan_result = scan_cache_directory(cache_dir)?;
+    Ok(scan_result.stats)
+}
+
+pub fn cleanup_cache(
+    cache_dir: &str,
+    max_age_days: f64,
+    max_size_mb: f64,
+    target_size_ratio: f64,
+) -> Result<(usize, f64)> {
+    let scan_result = scan_cache_directory(cache_dir)?;
+
+    if scan_result.entries.is_empty() {
+        return Ok((0, 0.0));
+    }
+
+    let current_time = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs() as f64;
+    let max_age_seconds = max_age_days * 24.0 * 3600.0;
+
+    let mut removed_count = 0;
+    let mut removed_size = 0.0;
+    let mut remaining_entries = Vec::new();
+    let mut total_remaining_size = 0u64;
+
+    for entry in scan_result.entries {
+        if let Ok(age) = entry.modified.duration_since(UNIX_EPOCH) {
+            let age_seconds = current_time - age.as_secs() as f64;
+            if age_seconds > max_age_seconds {
+                match fs::remove_file(&entry.path) {
+                    Ok(_) => {
+                        removed_count += 1;
+                        removed_size += entry.size as f64 / (1024.0 * 1024.0);
+                    }
+                    Err(e) => {
+                        tracing::debug!("Failed to remove {:?}: {}", entry.path, e);
+                    }
+                }
+            } else {
+                total_remaining_size += entry.size;
+                remaining_entries.push(entry);
+            }
+        }
+    }
+
+    let mut total_size_mb = total_remaining_size as f64 / (1024.0 * 1024.0);
+
+    if total_size_mb > max_size_mb {
+        remaining_entries.sort_by_key(|e| e.modified);
+
+        let target_size = max_size_mb * target_size_ratio;
+
+        for entry in remaining_entries {
+            if total_size_mb <= target_size {
+                break;
+            }
+
+            match fs::remove_file(&entry.path) {
+                Ok(_) => {
+                    let size_mb = entry.size as f64 / (1024.0 * 1024.0);
+                    removed_count += 1;
+                    removed_size += size_mb;
+                    total_size_mb -= size_mb;
+                }
+                Err(e) => {
+                    tracing::debug!("Failed to remove {:?}: {}", entry.path, e);
+                }
+            }
+        }
+    }
+
+    Ok((removed_count, removed_size))
+}
+
+pub fn smart_cleanup_cache(
+    cache_dir: &str,
+    max_age_days: f64,
+    max_size_mb: f64,
+    min_free_space_mb: f64,
+) -> Result<(usize, f64)> {
+    let stats = get_cache_metadata(cache_dir)?;
+
+    let needs_cleanup = stats.available_space_mb < min_free_space_mb
+        || stats.total_size_mb > max_size_mb
+        || stats.oldest_file_age_days > max_age_days;
+
+    if !needs_cleanup {
+        return Ok((0, 0.0));
+    }
+
+    let target_ratio = if stats.available_space_mb < min_free_space_mb {
+        0.5
+    } else {
+        0.8
+    };
+
+    cleanup_cache(cache_dir, max_age_days, max_size_mb, target_ratio)
+}
+
+pub fn is_cache_valid(cache_path: &str, max_age_days: f64) -> bool {
+    let path = Path::new(cache_path);
+
+    if !path.exists() {
+        return false;
+    }
+
+    match fs::metadata(path) {
+        Ok(metadata) => match metadata.modified() {
+            Ok(modified) => match SystemTime::now().duration_since(modified) {
+                Ok(elapsed) => {
+                    let age_days = elapsed.as_secs() as f64 / (24.0 * 3600.0);
+                    age_days <= max_age_days
+                }
+                Err(_) => false,
+            },
+            Err(_) => false,
+        },
+        Err(_) => false,
+    }
+}
+
+pub fn clear_cache_directory(cache_dir: &str) -> Result<(usize, f64)> {
+    let dir_path = Path::new(cache_dir);
+
+    if !dir_path.exists() {
+        return Ok((0, 0.0));
+    }
+
+    let mut removed_count = 0;
+    let mut removed_size = 0.0;
+
+    let read_dir =
+        fs::read_dir(dir_path).map_err(|e| KreuzbergError::cache(format!("Failed to read cache directory: {}", e)))?;
+
+    for entry in read_dir {
+        let entry = match entry {
+            Ok(e) => e,
+            Err(e) => {
+                tracing::debug!("Error reading entry: {}", e);
+                continue;
+            }
+        };
+
+        let metadata = match entry.metadata() {
+            Ok(m) if m.is_file() => m,
+            _ => continue,
+        };
+
+        let path = entry.path();
+        if path.extension().and_then(|s| s.to_str()) != Some("msgpack") {
+            continue;
+        }
+
+        let size_mb = metadata.len() as f64 / (1024.0 * 1024.0);
+        match fs::remove_file(&path) {
+            Ok(_) => {
+                removed_count += 1;
+                removed_size += size_mb;
+            }
+            Err(e) => {
+                tracing::debug!("Failed to remove {:?}: {}", path, e);
+            }
+        }
+    }
+
+    Ok((removed_count, removed_size))
+}
+
+pub fn batch_cleanup_caches(
+    cache_dirs: &[&str],
+    max_age_days: f64,
+    max_size_mb: f64,
+    min_free_space_mb: f64,
+) -> Result<Vec<(usize, f64)>> {
+    cache_dirs
+        .iter()
+        .map(|dir| smart_cleanup_cache(dir, max_age_days, max_size_mb, min_free_space_mb))
+        .collect()
+}