Skip to content

Commit 2e6f22e

Browse files
committed
Added robustness to rejected images, with tests
1 parent 37adc11 commit 2e6f22e

File tree

3 files changed

+434
-8
lines changed

3 files changed

+434
-8
lines changed

src/paperqa/settings.py

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
import asyncio
22
import importlib.resources
3+
import logging
34
import os
45
import pathlib
6+
import re
57
import warnings
68
from collections import defaultdict
79
from collections.abc import Awaitable, Callable, Mapping, Sequence
@@ -11,6 +13,7 @@
1113
from typing import Any, ClassVar, Protocol, Self, TypeAlias, cast, runtime_checkable
1214

1315
import anyio
16+
import litellm
1417
from aviary.core import Message, Tool, ToolSelector
1518
from lmi import (
1619
CommonLLMNames,
@@ -62,6 +65,8 @@
6265
from paperqa.types import Context, ParsedMedia, ParsedText
6366
from paperqa.utils import hexdigest, pqa_directory
6467

68+
logger = logging.getLogger(__name__)
69+
6570
# TODO: move to actual EnvironmentState
6671
# when we can do so without a circular import
6772
_EnvironmentState: TypeAlias = Any
@@ -1179,15 +1184,42 @@ async def enrich_single_media(
11791184
else ""
11801185
)
11811186
)
1182-
result = await llm.call_single(
1183-
messages=[
1184-
Message.create_message(
1185-
text=prompt, images=[media.to_image_url()]
1187+
try:
1188+
result = await llm.call_single(
1189+
messages=[
1190+
Message.create_message(
1191+
text=prompt, images=[media.to_image_url()]
1192+
)
1193+
]
1194+
)
1195+
if result.text:
1196+
media.info["enriched_description"] = result.text.strip()
1197+
except (litellm.InternalServerError, litellm.BadRequestError) as exc:
1198+
# Handle image > 5-MB failure mode 1:
1199+
# > litellm.BadRequestError: AnthropicException -
1200+
# > {"type":"error","error":{"type":"invalid_request_error",
1201+
# > "message":"messages.0.content.0.image.source.base64: image exceeds 5 MB maximum: 6229564 bytes > 5242880 bytes"}, # noqa: E501, W505
1202+
# > "request_id":"req_abc123"}.
1203+
# and image > 5-MB failure mode 2:
1204+
# > litellm.InternalServerError: AnthropicException -
1205+
# > {"type":"error","error":{"type":"invalid_request_error",
1206+
# > "message":"messages.0.content.0.image.source.base64: image exceeds 5 MB maximum: 5690780 bytes > 5242880 bytes"}, # noqa: E501, W505
1207+
# > "request_id":"req_abc123"}.
1208+
# And the image being corrupt (but a reasonable size)
1209+
if (
1210+
isinstance(exc, litellm.InternalServerError)
1211+
and re.search(
1212+
r"image exceeds .+ maximum", str(exc), re.IGNORECASE
11861213
)
1187-
]
1188-
)
1189-
if result.text:
1190-
media.info["enriched_description"] = result.text.strip()
1214+
) or isinstance(exc, litellm.BadRequestError):
1215+
logger.warning(
1216+
f"Skipping enrichment for media index {media.index}"
1217+
f" on page {page_num} with metadata {media.info} because"
1218+
f" it was rejected by the LLM provider for {llm.name!r}."
1219+
f" Full error message: {exc!r}"
1220+
)
1221+
else:
1222+
raise
11911223

11921224
await asyncio.gather(*list(starmap(enrich_single_media, media_to_enrich)))
11931225
count_enriched = sum(

0 commit comments

Comments
 (0)