feat(server): add LLM processing and error handling strategies to chunkr

Author: skeptrunedev

pdf2md/server/src/operators/chunkr.rs

@@ -324,6 +324,39 @@ pub enum OcrStrategy {
     Auto,
 }
 
+#[derive(Debug, Serialize, Deserialize, Clone, ToSchema, Default)]
+/// Specifies the fallback strategy for LLM processing
+///
+/// This can be:
+/// 1. None - No fallback will be used
+/// 2. Default - The system default fallback model will be used
+/// 3. Model - A specific model ID will be used as fallback (check the documentation for the models.)
+pub enum FallbackStrategy {
+    /// No fallback will be used
+    None,
+    /// Use the system default fallback model
+    #[default]
+    Default,
+    /// Use a specific model as fallback
+    Model(String),
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone, ToSchema)]
+/// Controls the LLM used for the task.
+pub struct LlmProcessing {
+    /// The ID of the model to use for the task. If not provided, the default model will be used.
+    /// Please check the documentation for the model you want to use.
+    pub model_id: Option<String>,
+    /// The fallback strategy to use for the LLMs in the task.
+    #[serde(default)]
+    pub fallback_strategy: FallbackStrategy,
+    /// The maximum number of tokens to generate.
+    pub max_completion_tokens: Option<u32>,
+    /// The temperature to use for the LLM.
+    #[serde(default)]
+    pub temperature: f32,
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, Display, Eq, PartialEq, ToSchema, Default)]
 /// Controls the segmentation strategy:
 /// - `LayoutAnalysis`: Analyzes pages for layout elements (e.g., `Table`, `Picture`, `Formula`, etc.) using bounding boxes. Provides fine-grained segmentation and better chunking. (Latency penalty: ~TBD seconds per page).
@@ -334,6 +367,16 @@ pub enum SegmentationStrategy {
     Page,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone, Display, Eq, PartialEq, ToSchema, Default)]
+/// Controls how errors are handled during processing:
+/// - `Fail`: Stops processing and fails the task when any error occurs
+/// - `Continue`: Attempts to continue processing despite non-critical errors (eg. LLM refusals etc.)
+pub enum ErrorHandlingStrategy {
+    #[default]
+    Fail,
+    Continue,
+}
+
 #[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Display, Default)]
 pub enum PipelineType {
     Azure,
@@ -376,6 +419,8 @@ pub struct CreateFormWithoutFile {
     pub segment_processing: Option<SegmentProcessing>,
     #[schema(default = "LayoutAnalysis")]
     pub segmentation_strategy: Option<SegmentationStrategy>,
+    pub error_handling: Option<ErrorHandlingStrategy>,
+    pub llm_processing: Option<LlmProcessing>,
 }
 
 #[derive(Debug, Serialize, Clone, Deserialize, ToSchema, IntoParams)]
@@ -400,6 +445,8 @@ pub struct CreateForm {
     pub segment_processing: Option<SegmentProcessing>,
     #[schema(default = "LayoutAnalysis")]
     pub segmentation_strategy: Option<SegmentationStrategy>,
+    pub error_handling: Option<ErrorHandlingStrategy>,
+    pub llm_processing: Option<LlmProcessing>,
 }
 
 fn get_chunkr_credentials(api_key: Option<&str>) -> Result<(String, String), ServiceError> {
@@ -433,6 +480,13 @@ pub async fn create_chunkr_task(
             pipeline: Some(PipelineType::Chunkr),
             segment_processing: None,
             segmentation_strategy: Some(SegmentationStrategy::LayoutAnalysis),
+            error_handling: Some(ErrorHandlingStrategy::default()),
+            llm_processing: Some(LlmProcessing {
+                model_id: None,
+                fallback_strategy: FallbackStrategy::default(),
+                max_completion_tokens: None,
+                temperature: 0.0,
+            }),
         },
         |payload| CreateForm {
             file: file_base64.to_string(),
@@ -444,6 +498,8 @@ pub async fn create_chunkr_task(
             pipeline: payload.pipeline,
             segment_processing: payload.segment_processing,
             segmentation_strategy: payload.segmentation_strategy,
+            error_handling: payload.error_handling,
+            llm_processing: payload.llm_processing,
         },
     );