scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

fix(ocr): avoid duplicate post-call callback after Azure polling

Browse source
This commit is contained in:
Yujong Lee 2026-09-15 09:14:48 -07:00
parent 451e96661e
commit ff3185f385
2 changed files with 15 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs
View file

@ -530,7 +530,7 @@ mod polling {
crate::ocr::client::read_response_bytes(response, connection.max_response_bytes)
.await?;
crate::ocr::handler::post_call(hooks, &bytes).await?;
poll_operation(http_client, operation, headers, connection, native, hooks).await
poll_operation(http_client, operation, headers, connection, native).await
}
async fn poll_operation(
@ -539,7 +539,6 @@ mod polling {
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
hooks: &Arc<dyn OcrHooks>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
let deadline = Instant::now()
.checked_add(connection.poll_timeout)
@ -583,10 +582,7 @@ mod polling {
.await
.map_err(|_| OcrPollingError::PollTimeout)??;
match &decoded.data.status {
Some(OperationStatus::Succeeded) => {
crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?;
return Ok(decoded);
}
Some(OperationStatus::Succeeded) => return Ok(decoded),
Some(OperationStatus::Running | OperationStatus::NotStarted) => {
tokio::time::timeout_at(
deadline,

25
litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
View file

@ -180,6 +180,7 @@ async fn accepted_response_polls_to_success_with_only_credentials() {
struct SubmissionBoundary {
request_count: Arc<Mutex<Vec<String>>>,
post_calls: Arc<Mutex<Vec<Value>>>,
}
impl super::hooks::OcrHooks for SubmissionBoundary {
@ -188,24 +189,18 @@ impl super::hooks::OcrHooks for SubmissionBoundary {
request: super::hooks::OcrPostCallRequest,
) -> super::hooks::OcrHookFuture<'_, super::hooks::OcrPostCallRequest> {
Box::pin(async move {
match self.request_count.lock().unwrap().len() {
1 => assert_eq!(request.original_response, json!(r#"{"submitted":true}"#)),
2 => assert!(
request
.original_response
.as_str()
.unwrap()
.contains("succeeded")
),
count => panic!("unexpected callback after {count} requests"),
}
assert_eq!(self.request_count.lock().unwrap().len(), 1);
self.post_calls
.lock()
.unwrap()
.push(request.original_response.clone());
Ok(request)
})
}
}
#[tokio::test]
async fn accepted_response_runs_post_call_before_polling() {
async fn accepted_response_runs_post_call_once_before_polling() {
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
@ -215,9 +210,11 @@ async fn accepted_response_runs_post_call_before_polling() {
MockResponse::json(json!({"status":"succeeded"})),
])
.await;
let post_calls = Arc::new(Mutex::new(Vec::new()));
let request = super::LiteLLMOcrRequest {
hooks: Arc::new(SubmissionBoundary {
request_count: seen.clone(),
post_calls: post_calls.clone(),
}),
..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}))
};
@ -225,6 +222,10 @@ async fn accepted_response_runs_post_call_before_polling() {
perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 2);
assert_eq!(
*post_calls.lock().unwrap(),
[json!(r#"{"submitted":true}"#)]
);
}
#[tokio::test]