Add Rust OCR hidden params marker

litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs

@@ -367,6 +367,7 @@ mod tests {
             end_time: 1.0,
             stream: false,
             metadata: StandardLoggingMetadata::default(),
+            hidden_params: Default::default(),
             messages: None,
         });
 

litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs

@@ -211,6 +211,7 @@ mod tests {
                 user_api_key_team_id: Some("team".to_string()),
                 ..Default::default()
             },
+            hidden_params: Default::default(),
             messages: Some(json!([{"role": "user", "content": "read this"}])),
         }
     }

litellm-rust/crates/ai-gateway/src/integrations/types.rs

@@ -58,6 +58,8 @@ pub struct StandardLoggingPayload {
 
     pub metadata: StandardLoggingMetadata,
 
+    pub hidden_params: HashMap<String, Value>,
+
     /// Optional; stored as request input on the spend log row.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub messages: Option<Value>,

litellm-rust/crates/ai-gateway/src/ocr/hooks.rs

@@ -5,6 +5,7 @@ use std::pin::Pin;
 use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
 use litellm_core::error::CoreError;
 use litellm_core::ocr::transformation::OcrAuthStrategy;
+use litellm_core::ocr::types::litellm_rust_hidden_params;
 use litellm_core::CoreResult;
 use serde_json::{json, Map, Value};
 
@@ -168,6 +169,7 @@ impl OcrLifecycleHooks {
                 user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
                 ..Default::default()
             },
+            hidden_params: litellm_rust_hidden_params().into_iter().collect(),
             messages: None,
         }
     }

litellm-rust/crates/ai-gateway/src/ocr/tests.rs

@@ -351,6 +351,7 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
     .expect("ocr request succeeds");
 
     assert_eq!(response["pages"][0]["markdown"], "ok");
+    assert_eq!(response["_hidden_params"]["litellm_rust"], true);
     assert_eq!(
         guardrail.events(),
         vec!["async_pre_call_hook", "async_moderation_hook"]

litellm-rust/crates/ai-gateway/src/realtime/streaming.rs

@@ -178,6 +178,7 @@ impl RealTimeStreaming {
                 user_api_key_team_id: self.metadata.user_api_key_team_id.clone(),
                 ..Default::default()
             },
+            hidden_params: Default::default(),
             messages: None,
         }
     }

litellm-rust/crates/core/src/ocr/types.rs

@@ -1,5 +1,5 @@
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
+use serde_json::{json, Map, Value};
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct OcrRequestData {
@@ -24,6 +24,32 @@ impl OcrResponseData {
             "document_annotation": self.document_annotation,
             "usage_info": self.usage_info,
             "object": self.object,
+            "_hidden_params": Value::Object(litellm_rust_hidden_params()),
         })
     }
 }
+
+pub fn litellm_rust_hidden_params() -> Map<String, Value> {
+    let mut hidden_params = Map::new();
+    hidden_params.insert("litellm_rust".to_string(), json!(true));
+    hidden_params
+}
+
+#[cfg(test)]
+mod tests {
+    use super::OcrResponseData;
+
+    #[test]
+    fn ocr_response_data_marks_litellm_rust_hidden_params() {
+        let response = OcrResponseData {
+            pages: vec![],
+            model: "mistral-ocr-latest".to_string(),
+            document_annotation: None,
+            usage_info: None,
+            object: "ocr".to_string(),
+        }
+        .into_json();
+
+        assert_eq!(response["_hidden_params"]["litellm_rust"], true);
+    }
+}

litellm/litellm_core_utils/litellm_logging.py

@@ -5342,6 +5342,7 @@ class StandardLoggingPayloadSetup:
             batch_models=None,
             litellm_model_name=None,
             usage_object=None,
+            litellm_rust=None,
         )
         if hidden_params is not None:
             for key in StandardLoggingHiddenParams.__annotations__.keys():

litellm/llms/base_llm/ocr/transformation.py

@@ -77,6 +77,13 @@ class OCRResponse(LiteLLMPydanticObjectBase):
     # Define private attributes using PrivateAttr
     _hidden_params: dict = PrivateAttr(default_factory=dict)
 
+    def model_post_init(self, __context: Any) -> None:
+        extra_fields = getattr(self, "__pydantic_extra__", None)
+        if isinstance(extra_fields, dict):
+            hidden_params = extra_fields.pop("_hidden_params", None)
+            if isinstance(hidden_params, dict):
+                self._hidden_params.update(hidden_params)
+
 
 class OCRRequestData(LiteLLMPydanticObjectBase):
     """OCR request data structure."""

litellm/proxy/_types.py

@@ -3197,6 +3197,7 @@ class SpendLogsMetadata(TypedDict):
         str
     ]  # S3/GCS object key for cold storage retrieval
     litellm_overhead_time_ms: Optional[float]  # LiteLLM overhead time in milliseconds
+    litellm_rust: Optional[bool]  # True when the request used the LiteLLM Rust path
     attempted_retries: Optional[
         int
     ]  # Number of retries attempted (0 = first attempt succeeded)

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -78,6 +78,7 @@ def _get_spend_logs_metadata(
     model_map_information: Optional[StandardLoggingModelInformation] = None,
     cold_storage_object_key: Optional[str] = None,
     litellm_overhead_time_ms: Optional[float] = None,
+    litellm_rust: Optional[bool] = None,
     cost_breakdown: Optional[CostBreakdown] = None,
     litellm_call_id: Optional[str] = None,
 ) -> SpendLogsMetadata:
@@ -107,6 +108,7 @@ def _get_spend_logs_metadata(
             eval_information=None,
             cold_storage_object_key=cold_storage_object_key,
             litellm_overhead_time_ms=None,
+            litellm_rust=litellm_rust,
             attempted_retries=None,
             max_retries=None,
             cost_breakdown=None,
@@ -134,6 +136,7 @@ def _get_spend_logs_metadata(
     clean_metadata["model_map_information"] = model_map_information
     clean_metadata["cold_storage_object_key"] = cold_storage_object_key
     clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms
+    clean_metadata["litellm_rust"] = litellm_rust
     clean_metadata["cost_breakdown"] = cost_breakdown
     clean_metadata["litellm_call_id"] = litellm_call_id
 
@@ -381,6 +384,11 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
             else None
         ),
         litellm_overhead_time_ms=litellm_overhead_time_ms,
+        litellm_rust=(
+            standard_logging_payload.get("hidden_params", {}).get("litellm_rust", None)
+            if standard_logging_payload is not None
+            else None
+        ),
         cost_breakdown=(
             standard_logging_payload.get("cost_breakdown", None)
             if standard_logging_payload is not None

litellm/types/utils.py

@@ -2728,6 +2728,7 @@ class StandardLoggingHiddenParams(TypedDict):
     batch_models: Optional[List[str]]
     litellm_model_name: Optional[str]  # the model name sent to the provider by litellm
     usage_object: Optional[dict]
+    litellm_rust: Optional[bool]
 
 
 class StandardLoggingModelInformation(TypedDict):

tests/ocr_tests/test_ocr_rust_bridge_callbacks.py

@@ -60,14 +60,11 @@ class OCRCustomGuardrail(CustomGuardrail):
         self, user_api_key_dict, cache, data: dict[str, Any], call_type: str
     ) -> None:
         self.calls.append({"data": data, "call_type": call_type})
-        data["document"] = {
-            **data["document"],
-            "guardrail_executed": True,
-        }
         data["optional_params"] = {
             **data["optional_params"],
             "include_image_base64": True,
         }
+        data["guardrail_executed"] = True
 
     async def async_log_success_event(
         self,
@@ -114,6 +111,7 @@ async def test_rust_ocr_executes_custom_logger_from_callback_manager():
     )
 
     assert len(response.pages) > 0
+    assert response._hidden_params["litellm_rust"] is True
     assert custom_logger.response_obj["object"] == "ocr"
     assert len(custom_logger.response_obj["value"]["pages"]) > 0
     assert isinstance(custom_logger.start_time, datetime)
@@ -128,6 +126,7 @@ async def test_rust_ocr_executes_custom_logger_from_callback_manager():
     assert logged_payload["custom_llm_provider"] == "mistral"
     assert logged_payload["call_type"] == "ocr"
     assert logged_payload["response_cost"] == 0.0
+    assert logged_payload["hidden_params"]["litellm_rust"] is True
 
 
 @pytest.mark.asyncio
@@ -144,8 +143,9 @@ async def test_rust_ocr_executes_custom_guardrail_from_callback_manager():
     )
 
     assert len(response.pages) > 0
+    assert response._hidden_params["litellm_rust"] is True
     assert custom_guardrail.calls[0]["call_type"] == "ocr"
-    assert custom_guardrail.calls[0]["data"]["document"]["guardrail_executed"] is True
+    assert custom_guardrail.calls[0]["data"]["guardrail_executed"] is True
     assert (
         custom_guardrail.calls[0]["data"]["optional_params"]["include_image_base64"]
         is True