diff --git a/docker/Dockerfile b/docker/Dockerfile index 5acee2902..115b0c48e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -28,7 +28,8 @@ RUN chmod 777 -R /tmp && \ g++ \ make \ git \ - wget && \ + wget \ + poppler-utils && \ rm -rf /var/lib/apt/lists/* RUN case ${TARGETPLATFORM} in \ diff --git a/lightllm/server/api_anthropic.py b/lightllm/server/api_anthropic.py index 60e7f6746..47730ed05 100644 --- a/lightllm/server/api_anthropic.py +++ b/lightllm/server/api_anthropic.py @@ -10,19 +10,48 @@ """ from __future__ import annotations +import asyncio +import base64 +import binascii +import hashlib +import os +import select +import shutil +import subprocess +import tempfile +import time import uuid import ujson as json +from collections import OrderedDict from http import HTTPStatus +from threading import Lock from typing import Any, Dict, Tuple from fastapi import Request from fastapi.responses import JSONResponse, Response +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) _cached_adapter: Any = None +_PDF_MAX_BYTES = 20 * 1024 * 1024 +_PDF_MAX_DOCUMENTS = 20 +_PDF_MAX_RENDER_PAGES = 20 +_PDF_MAX_TEXT_BYTES = 10 * 1024 * 1024 +_PDF_MAX_RENDER_DIMENSION = 2048 +_PDF_MAX_RENDER_PAGE_BYTES = 10 * 1024 * 1024 +_PDF_MAX_RENDERED_BYTES = 20 * 1024 * 1024 +_PDF_MAX_CONVERTED_BYTES = 28 * 1024 * 1024 +_PDF_SUBPROCESS_TIMEOUT_SECONDS = 30 +_PDF_PARSING_ENV = "LIGHTLLM_ANTHROPIC_ENABLE_PDF_PARSING" +_PDF_REQUIRED_TOOLS = ("pdftotext", "pdftoppm", "pdfinfo") +_PDF_CACHE_MAX_BYTES = 2 * 1024 * 1024 * 1024 +_PDF_CACHE: OrderedDict[tuple[str, bytes, int], tuple[int, Any]] = OrderedDict() +_PDF_CACHE_BYTES = 0 +_PDF_CACHE_LOCK = Lock() +_PDF_CACHE_MISS = object() def get_anthropic_messages_adapter() -> Any: @@ -66,6 +95,8 @@ def _anthropic_to_chat_request(anthropic_body: Dict[str, Any]) -> Tuple[Dict[str """ adapter = get_anthropic_messages_adapter() + _replace_anthropic_pdf_documents(anthropic_body) + tool_result_image_urls = _tool_result_image_urls(anthropic_body) openai_request, tool_name_mapping = adapter.translate_anthropic_to_openai(anthropic_body) if hasattr(openai_request, "model_dump"): @@ -77,6 +108,8 @@ def _anthropic_to_chat_request(anthropic_body: Dict[str, Any]) -> Tuple[Dict[str if "max_tokens" in anthropic_body: openai_dict["max_tokens"] = anthropic_body["max_tokens"] + _restore_tool_result_image_urls(openai_dict, tool_result_image_urls) + # Forward LightLLM-specific fields nested under ``extra_body`` (OpenAI SDK # convention) so clients hitting /v1/messages can reach ChatCompletionRequest # options Anthropic's own schema does not expose — notably chat_template_kwargs @@ -98,6 +131,530 @@ def _anthropic_to_chat_request(anthropic_body: Dict[str, Any]) -> Tuple[Dict[str return openai_dict, tool_name_mapping +def _restore_tool_result_image_urls(openai_dict: Dict[str, Any], tool_result_image_urls: Dict[str, set[str]]) -> None: + """LiteLLM flattens Anthropic tool_result image blocks into data URL text. + + LightLLM's multimodal path only sees OpenAI image_url content parts, so + restore image-only tool results to that shape before constructing the + ChatCompletionRequest. + """ + for msg in openai_dict.get("messages") or []: + if not isinstance(msg, dict) or msg.get("role") != "tool": + continue + content = msg.get("content") + if not isinstance(content, str): + continue + if content not in tool_result_image_urls.get(msg.get("tool_call_id"), set()): + continue + msg["content"] = [{"type": "image_url", "image_url": {"url": content}}] + + +def _pdf_data_url_to_anthropic_parts(data_url: str, deadline: float | None = None) -> list[Dict[str, Any]]: + _ensure_pdf_parsing_supported() + pdf_bytes = _decode_pdf_data_url(data_url) + if deadline is None: + deadline = time.monotonic() + _PDF_SUBPROCESS_TIMEOUT_SECONDS + page_count = _ensure_pdf_page_limit(pdf_bytes, deadline) + return _pdf_bytes_to_anthropic_parts(pdf_bytes, page_count, deadline) + + +def _pdf_bytes_to_anthropic_parts(pdf_bytes: bytes, page_count: int, deadline: float) -> list[Dict[str, Any]]: + if _is_vision_enabled(): + pages = _render_pdf_pages_to_png_b64(pdf_bytes, page_count, deadline) + if pages: + return [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": page}} + for page in pages + ] + return [_pdf_bytes_to_text_part(pdf_bytes, deadline)] + + +def check_pdf_parsing_supported_at_startup() -> None: + if _is_pdf_parsing_enabled(): + _ensure_pdf_tools_installed() + + +def _ensure_pdf_parsing_supported() -> None: + if not _is_pdf_parsing_enabled(): + raise RuntimeError( + f"PDF document parsing is disabled. Set {_PDF_PARSING_ENV}=1 to enable it; " + "requires Poppler tools: pdftotext, pdftoppm, pdfinfo." + ) + _ensure_pdf_tools_installed() + + +def _is_pdf_parsing_enabled() -> bool: + return os.getenv(_PDF_PARSING_ENV, "False").upper() in {"1", "TRUE", "ON"} + + +def _ensure_pdf_tools_installed() -> None: + missing = [tool for tool in _PDF_REQUIRED_TOOLS if shutil.which(tool) is None] + if missing: + raise RuntimeError( + f"PDF document parsing requires Poppler tools when {_PDF_PARSING_ENV}=1. " + f"Missing: {', '.join(missing)}. Required: {', '.join(_PDF_REQUIRED_TOOLS)}." + ) + + +def _ensure_pdf_page_limit(pdf_bytes: bytes, deadline: float | None = None) -> int: + page_count = _pdf_page_count(pdf_bytes, deadline) + if page_count is None: + raise ValueError("Unable to determine PDF page count") + if page_count > _PDF_MAX_RENDER_PAGES: + raise ValueError( + f"PDF document has {page_count} pages and exceeds configured page limit of {_PDF_MAX_RENDER_PAGES}" + ) + return page_count + + +def _pdf_page_count(pdf_bytes: bytes, deadline: float | None = None) -> int | None: + pdfinfo = shutil.which("pdfinfo") + if not pdfinfo: + return None + timeout = _remaining_pdf_time(deadline) + if timeout <= 0: + return None + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: + tmp.write(pdf_bytes) + tmp_path = tmp.name + try: + try: + proc = subprocess.run( + [pdfinfo, tmp_path], + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=timeout, + ) + except Exception: + return None + finally: + try: + os.unlink(tmp_path) + except OSError: + pass + if proc.returncode != 0: + return None + for line in proc.stdout.decode("utf-8", "replace").splitlines(): + key, _, value = line.partition(":") + if key == "Pages": + try: + return int(value.strip()) + except ValueError: + return None + return None + + +def _pdf_bytes_to_text_part(pdf_bytes: bytes, deadline: float | None = None) -> Dict[str, str]: + text = _extract_pdf_text(pdf_bytes, deadline).strip() + if not text: + text = ( + "[PDF document attached, but no extractable text was found. " + "This backend does not OCR scanned PDFs in the Anthropic adapter. " + "Ask the client to read specific pages or use an OCR/vision-enabled backend.]" + ) + return {"type": "text", "text": f"[PDF extracted text]\n{text}"} + + +def _is_vision_enabled() -> bool: + try: + args = get_env_start_args() + except Exception: + return False + return bool(getattr(args, "enable_multimodal", False) and not getattr(args, "disable_vision", True)) + + +def _decode_pdf_data_url(data_url: str) -> bytes: + try: + _, encoded = data_url.split("base64,", 1) + except ValueError as exc: + raise ValueError("Invalid base64 PDF document block") from exc + if len(encoded) > ((_PDF_MAX_BYTES + 2) // 3 * 4 + 4): + raise ValueError("PDF document block exceeds configured size limit") + try: + pdf_bytes = base64.b64decode(encoded, validate=True) + except binascii.Error as exc: + raise ValueError("Invalid base64 PDF document block") from exc + if len(pdf_bytes) > _PDF_MAX_BYTES: + raise ValueError("PDF document block exceeds configured size limit") + return pdf_bytes + + +def _remaining_pdf_time(deadline: float | None) -> float: + if deadline is None: + return _PDF_SUBPROCESS_TIMEOUT_SECONDS + return max(0.0, deadline - time.monotonic()) + + +def _kill_pdf_process(proc: subprocess.Popen) -> None: + try: + proc.kill() + except OSError: + pass + try: + proc.wait() + except OSError: + pass + + +def _extract_pdf_text(pdf_bytes: bytes, deadline: float | None = None) -> str: + key = _pdf_cache_key("text", pdf_bytes) + cached = _pdf_cache_get(key) + if cached is not _PDF_CACHE_MISS: + return cached + + pdftotext = shutil.which("pdftotext") + if not pdftotext: + text = "" + else: + try: + text = _extract_pdf_text_with_pdftotext(pdftotext, pdf_bytes, deadline) + except Exception: + text = "" + text = text if text.strip() else "" + if text: + _pdf_cache_set(key, text) + return text + + +def _extract_pdf_text_with_pdftotext(pdftotext: str, pdf_bytes: bytes, deadline: float | None = None) -> str: + if deadline is None: + deadline = time.monotonic() + _PDF_SUBPROCESS_TIMEOUT_SECONDS + if _remaining_pdf_time(deadline) <= 0: + return "" + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: + tmp.write(pdf_bytes) + tmp_path = tmp.name + try: + proc = subprocess.Popen( + [pdftotext, "-layout", tmp_path, "-"], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + ) + try: + output_parts = [] + output_size = 0 + while True: + timeout = _remaining_pdf_time(deadline) + if timeout <= 0: + _kill_pdf_process(proc) + return "" + readable, _, _ = select.select([proc.stdout], [], [], timeout) + if not readable: + _kill_pdf_process(proc) + return "" + chunk = os.read( + proc.stdout.fileno(), + min(64 * 1024, _PDF_MAX_TEXT_BYTES + 1 - output_size), + ) + if not chunk: + break + output_parts.append(chunk) + output_size += len(chunk) + if output_size > _PDF_MAX_TEXT_BYTES: + _kill_pdf_process(proc) + return "" + try: + proc.wait(timeout=_remaining_pdf_time(deadline)) + except subprocess.TimeoutExpired: + _kill_pdf_process(proc) + return "" + except Exception: + _kill_pdf_process(proc) + return "" + finally: + proc.stdout.close() + finally: + try: + os.unlink(tmp_path) + except OSError: + pass + if proc.returncode != 0: + return "" + output = b"".join(output_parts) + return output.decode("utf-8", "replace") + + +def _render_pdf_pages_to_png_b64( + pdf_bytes: bytes, + page_count: int = _PDF_MAX_RENDER_PAGES, + deadline: float | None = None, +) -> tuple[str, ...]: + key = _pdf_cache_key("vision", pdf_bytes) + cached = _pdf_cache_get(key) + if cached is not _PDF_CACHE_MISS: + return cached + + pdftoppm = shutil.which("pdftoppm") + if not pdftoppm: + return () + if deadline is None: + deadline = time.monotonic() + _PDF_SUBPROCESS_TIMEOUT_SECONDS + timeout = _remaining_pdf_time(deadline) + if timeout <= 0: + return () + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: + tmp.write(pdf_bytes) + tmp_path = tmp.name + try: + with tempfile.TemporaryDirectory() as out_dir: + prefix = os.path.join(out_dir, "page") + pages = [] + rendered_bytes = 0 + try: + proc = subprocess.run( + [ + pdftoppm, + "-png", + "-scale-to", + str(_PDF_MAX_RENDER_DIMENSION), + "-f", + "1", + "-l", + str(page_count), + tmp_path, + prefix, + ], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=timeout, + ) + except Exception: + return () + if proc.returncode != 0: + return () + for page_path in _rendered_pdf_page_paths(out_dir): + try: + page_size = os.path.getsize(page_path) + except OSError: + break + if page_size > _PDF_MAX_RENDER_PAGE_BYTES or rendered_bytes + page_size > _PDF_MAX_RENDERED_BYTES: + return () + with open(page_path, "rb") as f: + page = f.read() + if len(page) != page_size: + return () + rendered_bytes += page_size + pages.append(base64.b64encode(page).decode("ascii")) + pages = tuple(pages) + if pages: + _pdf_cache_set(key, pages) + return pages + finally: + try: + os.unlink(tmp_path) + except OSError: + pass + + +def _rendered_pdf_page_paths(out_dir: str) -> list[str]: + numbered_paths = [] + for name in os.listdir(out_dir): + if not name.startswith("page-") or not name.endswith(".png"): + continue + try: + page_number = int(name[len("page-") : -len(".png")]) + except ValueError: + continue + numbered_paths.append((page_number, os.path.join(out_dir, name))) + return [path for _, path in sorted(numbered_paths)] + + +def _pdf_cache_key(kind: str, pdf_bytes: bytes) -> tuple[str, bytes, int]: + return (kind, hashlib.sha256(pdf_bytes).digest(), len(pdf_bytes)) + + +def _pdf_cache_get(key: tuple[str, bytes, int]) -> Any: + with _PDF_CACHE_LOCK: + item = _PDF_CACHE.get(key) + if item is None: + return _PDF_CACHE_MISS + _PDF_CACHE.move_to_end(key) + return item[1] + + +def _pdf_cache_set(key: tuple[str, bytes, int], value: Any) -> None: + global _PDF_CACHE_BYTES + size = _pdf_cache_size(value) + if size == 0 or size > _PDF_CACHE_MAX_BYTES: + return + with _PDF_CACHE_LOCK: + old = _PDF_CACHE.pop(key, None) + if old is not None: + _PDF_CACHE_BYTES -= old[0] + while _PDF_CACHE_BYTES + size > _PDF_CACHE_MAX_BYTES and _PDF_CACHE: + _, (old_size, _) = _PDF_CACHE.popitem(last=False) + _PDF_CACHE_BYTES -= old_size + _PDF_CACHE[key] = (size, value) + _PDF_CACHE_BYTES += size + + +def _pdf_cache_size(value: Any) -> int: + if isinstance(value, str): + return len(value.encode("utf-8")) + if isinstance(value, tuple): + return sum(len(item.encode("utf-8")) for item in value if isinstance(item, str)) + return 0 + + +def _clear_pdf_cache() -> None: + global _PDF_CACHE_BYTES + with _PDF_CACHE_LOCK: + _PDF_CACHE.clear() + _PDF_CACHE_BYTES = 0 + + +_extract_pdf_text.cache_clear = _clear_pdf_cache +_render_pdf_pages_to_png_b64.cache_clear = _clear_pdf_cache + + +def _replace_anthropic_pdf_documents(anthropic_body: Dict[str, Any]) -> None: + pdf_blocks = list(_iter_anthropic_pdf_documents(anthropic_body)) + if not pdf_blocks: + return + if any(block["source"].get("type") == "url" for block in pdf_blocks): + raise ValueError("URL-backed PDF document blocks are not supported; send the PDF as base64 instead") + _ensure_pdf_parsing_supported() + if len(pdf_blocks) > _PDF_MAX_DOCUMENTS: + raise ValueError( + f"Request contains {len(pdf_blocks)} PDF documents and exceeds configured document limit " + f"of {_PDF_MAX_DOCUMENTS}" + ) + + pdf_bytes_list = [] + total_bytes = 0 + for block in pdf_blocks: + pdf_bytes = _decode_pdf_data_url(f"data:application/pdf;base64,{block['source']['data']}") + total_bytes += len(pdf_bytes) + if total_bytes > _PDF_MAX_BYTES: + raise ValueError("PDF documents exceed configured request size limit") + pdf_bytes_list.append(pdf_bytes) + + deadline = time.monotonic() + _PDF_SUBPROCESS_TIMEOUT_SECONDS + page_counts = [] + total_pages = 0 + for pdf_bytes in pdf_bytes_list: + page_count = _ensure_pdf_page_limit(pdf_bytes, deadline) + total_pages += page_count + if total_pages > _PDF_MAX_RENDER_PAGES: + raise ValueError( + f"PDF documents have {total_pages} pages and exceed configured request page limit " + f"of {_PDF_MAX_RENDER_PAGES}" + ) + page_counts.append(page_count) + + replacements = {} + converted_bytes = 0 + for block, pdf_bytes, page_count in zip(pdf_blocks, pdf_bytes_list, page_counts): + replacement = _pdf_bytes_to_anthropic_parts(pdf_bytes, page_count, deadline) + converted_bytes += _pdf_converted_size(replacement) + if converted_bytes > _PDF_MAX_CONVERTED_BYTES: + raise ValueError("PDF conversions exceed configured request output size limit") + replacements[id(block)] = replacement + for message in anthropic_body.get("messages") or []: + if isinstance(message, dict): + _replace_pdf_documents_in_content(message.get("content"), replacements) + + +def _pdf_converted_size(parts: list[Dict[str, Any]]) -> int: + size = 0 + for part in parts: + if part.get("type") == "text" and isinstance(part.get("text"), str): + size += len(part["text"].encode("utf-8")) + continue + source = part.get("source") + if ( + part.get("type") == "image" + and isinstance(source, dict) + and source.get("type") == "base64" + and isinstance(source.get("data"), str) + ): + size += len(source["data"]) + return size + + +def _replace_pdf_documents_in_content(value: Any, replacements: Dict[int, list[Dict[str, Any]]]) -> None: + if isinstance(value, list): + new_items = [] + changed = False + for item in value: + replacement = replacements.get(id(item)) + if replacement is not None: + new_items.extend(replacement) + changed = True + else: + if isinstance(item, dict) and item.get("type") == "tool_result": + _replace_pdf_documents_in_content(item.get("content"), replacements) + new_items.append(item) + if changed: + value[:] = new_items + + +def _iter_anthropic_pdf_documents( + anthropic_body: Dict[str, Any], +): + for message in anthropic_body.get("messages") or []: + if isinstance(message, dict): + yield from _iter_pdf_documents_in_content(message.get("content")) + + +def _iter_pdf_documents_in_content(value: Any): + if not isinstance(value, list): + return + for item in value: + if _is_anthropic_pdf_document(item): + yield item + elif isinstance(item, dict) and item.get("type") == "tool_result": + yield from _iter_pdf_documents_in_content(item.get("content")) + + +def _tool_result_image_urls(anthropic_body: Dict[str, Any]) -> Dict[str, set[str]]: + image_urls: Dict[str, set[str]] = {} + for message in anthropic_body.get("messages") or []: + if not isinstance(message, dict) or not isinstance(message.get("content"), list): + continue + for block in message["content"]: + if not isinstance(block, dict) or block.get("type") != "tool_result": + continue + tool_use_id = block.get("tool_use_id") + content = block.get("content") + if not isinstance(tool_use_id, str) or not isinstance(content, list): + continue + for item in content: + if not isinstance(item, dict) or item.get("type") != "image": + continue + source = item.get("source") + if not isinstance(source, dict): + continue + if source.get("type") == "url" and isinstance(source.get("url"), str) and source["url"]: + image_urls.setdefault(tool_use_id, set()).add(source["url"]) + elif ( + source.get("type") == "base64" + and isinstance(source.get("media_type"), str) + and source["media_type"].startswith("image/") + and isinstance(source.get("data"), str) + ): + image_urls.setdefault(tool_use_id, set()).add( + f"data:{source['media_type']};base64,{source['data']}" + ) + return image_urls + + +def _is_anthropic_pdf_document(value: Any) -> bool: + if not isinstance(value, dict) or value.get("type") != "document": + return False + source = value.get("source") + return isinstance(source, dict) and ( + ( + source.get("type") == "base64" + and source.get("media_type") == "application/pdf" + and isinstance(source.get("data"), str) + ) + or (source.get("type") == "url" and isinstance(source.get("url"), str) and bool(source["url"])) + ) + + # --------------------------------------------------------------------------- # Response translation # --------------------------------------------------------------------------- @@ -553,7 +1110,7 @@ async def anthropic_messages_impl(raw_request: Request) -> Response: is_stream = bool(raw_body.get("stream")) try: - chat_dict, tool_name_mapping = _anthropic_to_chat_request(raw_body) + chat_dict, tool_name_mapping = await asyncio.to_thread(_anthropic_to_chat_request, raw_body) except Exception as exc: logger.exception("Failed to translate Anthropic request") return _anthropic_error_response(HTTPStatus.BAD_REQUEST, f"Request translation failed: {exc}") diff --git a/test/test_api/test_anthropic_extra_body.py b/test/test_api/test_anthropic_extra_body.py index 99f08a01a..b2b272b6b 100644 --- a/test/test_api/test_anthropic_extra_body.py +++ b/test/test_api/test_anthropic_extra_body.py @@ -10,11 +10,13 @@ """ import asyncio +import base64 import pytest import ujson as json pytest.importorskip("litellm") +import lightllm.server.api_anthropic as api_anthropic from lightllm.server.api_anthropic import _anthropic_to_chat_request, _openai_sse_to_anthropic_events @@ -26,6 +28,91 @@ def _base_body(): } +def _pdf_document_block(): + return { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": base64.b64encode(b"%PDF-1.4\n").decode("ascii"), + }, + } + + +def _url_pdf_document_block(): + return { + "type": "document", + "source": { + "type": "url", + "url": "https://example.com/document.pdf", + }, + } + + +def _enable_pdf_parsing(monkeypatch): + monkeypatch.setenv("LIGHTLLM_ANTHROPIC_ENABLE_PDF_PARSING", "1") + monkeypatch.setattr(api_anthropic, "_ensure_pdf_tools_installed", lambda: None) + monkeypatch.setattr(api_anthropic, "_pdf_page_count", lambda _, deadline=None: 1) + + +def _user_pdf_body(): + return { + "model": "test-model", + "max_tokens": 32, + "messages": [ + { + "role": "user", + "content": [ + _pdf_document_block(), + {"type": "text", "text": "What is in the PDF?"}, + ], + } + ], + } + + +def _tool_result_body(content): + return { + "model": "test-model", + "max_tokens": 32, + "tools": [ + { + "name": "Read", + "description": "Read a file", + "input_schema": { + "type": "object", + "properties": {"file_path": {"type": "string"}}, + "required": ["file_path"], + }, + } + ], + "messages": [ + {"role": "user", "content": "read a file"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "call_read", + "name": "Read", + "input": {"file_path": "file"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_read", + "content": content, + } + ], + }, + ], + } + + def test_extra_body_chat_template_kwargs_forwarded(): body = _base_body() body["extra_body"] = {"chat_template_kwargs": {"enable_thinking": False}} @@ -77,6 +164,522 @@ def test_non_dict_extra_body_is_ignored(): assert "extra_body" not in chat_dict +def test_pdf_document_block_becomes_text_not_pdf_image_url(monkeypatch): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_is_vision_enabled", lambda: False) + monkeypatch.setattr(api_anthropic, "_extract_pdf_text", lambda _, deadline=None: "PDF_SENTINEL_DIRECT") + body = _user_pdf_body() + + chat_dict, _ = _anthropic_to_chat_request(body) + + content = chat_dict["messages"][0]["content"] + assert content[0]["type"] == "text" + assert "PDF_SENTINEL_DIRECT" in content[0]["text"] + assert "data:application/pdf" not in json.dumps(chat_dict) + + +def test_pdf_document_block_becomes_images_when_vision_enabled(monkeypatch): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_is_vision_enabled", lambda: True) + monkeypatch.setattr( + api_anthropic, + "_render_pdf_pages_to_png_b64", + lambda _, page_count=None, deadline=None: ["UE5HMQ==", "UE5HMg=="], + ) + body = _user_pdf_body() + + chat_dict, _ = _anthropic_to_chat_request(body) + + content = chat_dict["messages"][0]["content"] + assert [p["type"] for p in content[:2]] == ["image_url", "image_url"] + assert [p["image_url"]["url"] for p in content[:2]] == [ + "data:image/png;base64,UE5HMQ==", + "data:image/png;base64,UE5HMg==", + ] + assert "data:application/pdf" not in json.dumps(chat_dict) + assert "PDF extracted text" not in json.dumps(chat_dict) + + +def test_pdf_document_block_with_invalid_base64_fails_cleanly(monkeypatch): + _enable_pdf_parsing(monkeypatch) + body = { + "model": "test-model", + "max_tokens": 32, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "not-base64!", + }, + } + ], + } + ], + } + + with pytest.raises(ValueError, match="Invalid base64 PDF document block"): + _anthropic_to_chat_request(body) + + +def test_pdf_document_block_over_size_fails_cleanly(monkeypatch): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_PDF_MAX_BYTES", 4) + + with pytest.raises(ValueError, match="PDF document block exceeds configured size limit"): + _anthropic_to_chat_request(_user_pdf_body()) + + +def test_pdf_request_rejects_too_many_documents_before_page_checks(monkeypatch): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_PDF_MAX_DOCUMENTS", 1) + monkeypatch.setattr( + api_anthropic, + "_pdf_page_count", + lambda *_: pytest.fail("page count should not run before document-count validation"), + ) + body = _user_pdf_body() + body["messages"][0]["content"].insert(0, _pdf_document_block()) + + with pytest.raises(ValueError, match="document limit"): + _anthropic_to_chat_request(body) + + +def test_pdf_request_enforces_total_byte_limit(monkeypatch): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_PDF_MAX_BYTES", 10) + monkeypatch.setattr( + api_anthropic, + "_pdf_page_count", + lambda *_: pytest.fail("page count should not run before byte validation"), + ) + first = _pdf_document_block() + second = _pdf_document_block() + first["source"]["data"] = base64.b64encode(b"123456").decode("ascii") + second["source"]["data"] = base64.b64encode(b"abcdef").decode("ascii") + body = _user_pdf_body() + body["messages"][0]["content"] = [first, second] + + with pytest.raises(ValueError, match="request size limit"): + _anthropic_to_chat_request(body) + + +def test_pdf_request_validates_total_pages_before_conversion(monkeypatch): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_pdf_page_count", lambda *args: 11) + monkeypatch.setattr( + api_anthropic, + "_pdf_bytes_to_anthropic_parts", + lambda *_: pytest.fail("conversion should not run before page validation"), + ) + body = _user_pdf_body() + body["messages"][0]["content"].insert(0, _pdf_document_block()) + + with pytest.raises(ValueError, match="request page limit"): + _anthropic_to_chat_request(body) + + +def test_pdf_request_uses_one_subprocess_deadline(monkeypatch): + _enable_pdf_parsing(monkeypatch) + deadlines = [] + + def fake_page_count(_pdf_bytes, deadline): + deadlines.append(deadline) + return 1 + + def fake_convert(_pdf_bytes, _page_count, deadline): + deadlines.append(deadline) + return [{"type": "text", "text": "PDF"}] + + monkeypatch.setattr(api_anthropic, "_pdf_page_count", fake_page_count) + monkeypatch.setattr(api_anthropic, "_pdf_bytes_to_anthropic_parts", fake_convert) + body = _user_pdf_body() + body["messages"][0]["content"].insert(0, _pdf_document_block()) + + _anthropic_to_chat_request(body) + + assert len(deadlines) == 4 + assert len(set(deadlines)) == 1 + + +@pytest.mark.parametrize( + "replacement", + [ + [{"type": "text", "text": "1234"}], + [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "YWJj"}}], + ], + ids=["text", "rendered-image"], +) +def test_pdf_request_enforces_total_converted_output_limit(monkeypatch, replacement): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_PDF_MAX_CONVERTED_BYTES", 7) + calls = {"count": 0} + + def fake_convert(_pdf_bytes, _page_count, _deadline): + calls["count"] += 1 + return replacement + + monkeypatch.setattr(api_anthropic, "_pdf_bytes_to_anthropic_parts", fake_convert) + body = _user_pdf_body() + body["messages"][0]["content"] = [ + _pdf_document_block(), + _pdf_document_block(), + _pdf_document_block(), + ] + + with pytest.raises(ValueError, match="request output size limit"): + _anthropic_to_chat_request(body) + + assert calls["count"] == 2 + + +def test_url_pdf_document_block_is_rejected_explicitly(): + body = _user_pdf_body() + body["messages"][0]["content"][0] = _url_pdf_document_block() + + with pytest.raises(ValueError, match="URL-backed PDF document blocks are not supported"): + _anthropic_to_chat_request(body) + + +def test_tool_result_url_pdf_document_block_is_rejected_explicitly(): + body = _tool_result_body([_url_pdf_document_block()]) + + with pytest.raises(ValueError, match="URL-backed PDF document blocks are not supported"): + _anthropic_to_chat_request(body) + + +def test_tool_result_pdf_document_block_becomes_text_not_pdf_image_url(monkeypatch): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_is_vision_enabled", lambda: False) + monkeypatch.setattr(api_anthropic, "_extract_pdf_text", lambda _, deadline=None: "PDF_SENTINEL_TOOL") + body = _tool_result_body([_pdf_document_block()]) + + chat_dict, _ = _anthropic_to_chat_request(body) + + assert chat_dict["messages"][2]["role"] == "tool" + assert "PDF_SENTINEL_TOOL" in chat_dict["messages"][2]["content"] + assert "data:application/pdf" not in json.dumps(chat_dict) + + +def test_tool_result_pdf_document_block_becomes_images_when_vision_enabled(monkeypatch): + _enable_pdf_parsing(monkeypatch) + monkeypatch.setattr(api_anthropic, "_is_vision_enabled", lambda: True) + monkeypatch.setattr( + api_anthropic, + "_render_pdf_pages_to_png_b64", + lambda _, page_count=None, deadline=None: ["UE5HMQ=="], + ) + body = _tool_result_body([_pdf_document_block()]) + + chat_dict, _ = _anthropic_to_chat_request(body) + + assert chat_dict["messages"][2]["role"] == "tool" + assert chat_dict["messages"][2]["content"] == [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,UE5HMQ=="}, + } + ] + assert "data:application/pdf" not in json.dumps(chat_dict) + + +def test_pdf_vision_render_limits_pages(monkeypatch): + captured = [] + + api_anthropic._render_pdf_pages_to_png_b64.cache_clear() + + def fake_run(cmd, **kwargs): + captured.append(cmd) + with open(f"{cmd[-1]}-1.png", "wb") as f: + f.write(b"png") + return type("Proc", (), {"returncode": 0})() + + monkeypatch.setattr(api_anthropic.shutil, "which", lambda _: "/usr/bin/pdftoppm") + monkeypatch.setattr(api_anthropic.subprocess, "run", fake_run) + + assert api_anthropic._render_pdf_pages_to_png_b64(b"%PDF-1.4\n") == ("cG5n",) + assert len(captured) == 1 + for cmd in captured: + assert "-l" in cmd + assert cmd[cmd.index("-f") + 1] == "1" + assert cmd[cmd.index("-l") + 1] == str(api_anthropic._PDF_MAX_RENDER_PAGES) + assert "-scale-to" in cmd + assert "-scale-to-x" not in cmd + assert "-scale-to-y" not in cmd + assert str(api_anthropic._PDF_MAX_RENDER_DIMENSION) in cmd + + +def test_pdf_vision_reads_zero_padded_page_names(monkeypatch): + api_anthropic._render_pdf_pages_to_png_b64.cache_clear() + + def fake_run(cmd, **kwargs): + with open(f"{cmd[-1]}-01.png", "wb") as f: + f.write(b"one") + with open(f"{cmd[-1]}-02.png", "wb") as f: + f.write(b"two") + return type("Proc", (), {"returncode": 0})() + + monkeypatch.setattr(api_anthropic.shutil, "which", lambda _: "/usr/bin/pdftoppm") + monkeypatch.setattr(api_anthropic.subprocess, "run", fake_run) + + assert api_anthropic._render_pdf_pages_to_png_b64(b"%PDF-padded\n", page_count=2) == ( + "b25l", + "dHdv", + ) + + +def test_pdf_vision_rejects_large_rendered_page(monkeypatch): + api_anthropic._render_pdf_pages_to_png_b64.cache_clear() + monkeypatch.setattr(api_anthropic, "_PDF_MAX_RENDER_PAGE_BYTES", 2) + + def fake_run(cmd, **kwargs): + with open(f"{cmd[-1]}-1.png", "wb") as f: + f.write(b"png") + return type("Proc", (), {"returncode": 0})() + + monkeypatch.setattr(api_anthropic.shutil, "which", lambda _: "/usr/bin/pdftoppm") + monkeypatch.setattr(api_anthropic.subprocess, "run", fake_run) + + assert api_anthropic._render_pdf_pages_to_png_b64(b"%PDF-1.4\n") == () + + +def test_pdf_text_extraction_limits_output(monkeypatch): + class FakeStdout: + def fileno(self): + return 42 + + def close(self): + pass + + class FakeProcess: + stdout = FakeStdout() + returncode = -9 + killed = False + + def kill(self): + self.killed = True + + def wait(self, **kwargs): + pass + + process = FakeProcess() + monkeypatch.setattr(api_anthropic, "_PDF_MAX_TEXT_BYTES", 4) + monkeypatch.setattr(api_anthropic.subprocess, "Popen", lambda *args, **kwargs: process) + monkeypatch.setattr(api_anthropic.select, "select", lambda *args: ([process.stdout], [], [])) + monkeypatch.setattr(api_anthropic.os, "read", lambda _fd, size: b"x" * size) + + assert api_anthropic._extract_pdf_text_with_pdftotext("/usr/bin/pdftotext", b"%PDF-1.4\n") == "" + assert process.killed + + +def test_pdf_text_extraction_times_out_while_stdout_is_open(monkeypatch): + class FakeStdout: + def fileno(self): + return 42 + + def close(self): + pass + + class FakeProcess: + stdout = FakeStdout() + returncode = -9 + killed = False + + def kill(self): + self.killed = True + + def wait(self, **kwargs): + pass + + process = FakeProcess() + timeouts = [] + monkeypatch.setattr(api_anthropic.subprocess, "Popen", lambda *args, **kwargs: process) + + def fake_select(*args): + timeouts.append(args[-1]) + return [], [], [] + + monkeypatch.setattr(api_anthropic.select, "select", fake_select) + + assert api_anthropic._extract_pdf_text_with_pdftotext("/usr/bin/pdftotext", b"%PDF-1.4\n") == "" + assert process.killed + assert 0 < timeouts[0] <= api_anthropic._PDF_SUBPROCESS_TIMEOUT_SECONDS + + +def test_pdf_text_extraction_does_not_start_after_request_deadline(monkeypatch): + monkeypatch.setattr( + api_anthropic.subprocess, + "Popen", + lambda *args, **kwargs: pytest.fail("expired request must not start pdftotext"), + ) + + assert api_anthropic._extract_pdf_text_with_pdftotext("/usr/bin/pdftotext", b"%PDF-1.4\n", deadline=0) == "" + + +def test_pdf_text_extraction_is_cached(monkeypatch): + calls = {"count": 0} + + api_anthropic._extract_pdf_text.cache_clear() + + def fake_extract(_pdftotext, _pdf_bytes, _deadline=None): + calls["count"] += 1 + return "PDF_SENTINEL" + + monkeypatch.setattr(api_anthropic.shutil, "which", lambda _: "/usr/bin/pdftotext") + monkeypatch.setattr(api_anthropic, "_extract_pdf_text_with_pdftotext", fake_extract) + + assert api_anthropic._extract_pdf_text(b"%PDF-1.4\n") == "PDF_SENTINEL" + assert api_anthropic._extract_pdf_text(b"%PDF-1.4\n") == "PDF_SENTINEL" + assert calls["count"] == 1 + + +def test_pdf_cache_evicts_by_memory_budget(monkeypatch): + calls = {"count": 0} + + api_anthropic._extract_pdf_text.cache_clear() + monkeypatch.setattr(api_anthropic, "_PDF_CACHE_MAX_BYTES", 8) + + def fake_extract(_pdftotext, pdf_bytes, _deadline=None): + calls["count"] += 1 + return pdf_bytes.decode("ascii") + + monkeypatch.setattr(api_anthropic.shutil, "which", lambda _: "/usr/bin/pdftotext") + monkeypatch.setattr(api_anthropic, "_extract_pdf_text_with_pdftotext", fake_extract) + + assert api_anthropic._extract_pdf_text(b"aaaa") == "aaaa" + assert api_anthropic._extract_pdf_text(b"bbbbbbbb") == "bbbbbbbb" + assert api_anthropic._extract_pdf_text(b"aaaa") == "aaaa" + assert calls["count"] == 3 + + +def test_pdf_cache_does_not_store_empty_results(): + api_anthropic._clear_pdf_cache() + + api_anthropic._pdf_cache_set(("text", b"a", 1), "") + api_anthropic._pdf_cache_set(("vision", b"b", 1), ()) + + assert not api_anthropic._PDF_CACHE + assert api_anthropic._PDF_CACHE_BYTES == 0 + + +def test_anthropic_messages_impl_runs_translation_in_thread(monkeypatch): + called = {} + + async def fake_to_thread(fn, *args, **kwargs): + called["to_thread"] = True + return fn(*args, **kwargs) + + def fake_translate(_body): + return {"model": "test-model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1}, {} + + async def fake_chat_completions_impl(_request, _raw_request): + from fastapi.responses import Response + + return Response("ok") + + class FakeRequest: + async def json(self): + return {"model": "test-model", "max_tokens": 1, "messages": [{"role": "user", "content": "hi"}]} + + import lightllm.server.api_openai as api_openai + + monkeypatch.setattr(api_anthropic.asyncio, "to_thread", fake_to_thread) + monkeypatch.setattr(api_anthropic, "_anthropic_to_chat_request", fake_translate) + monkeypatch.setattr(api_openai, "chat_completions_impl", fake_chat_completions_impl) + + response = asyncio.run(api_anthropic.anthropic_messages_impl(FakeRequest())) + + assert called["to_thread"] + assert response.body == b"ok" + + +def test_tool_result_image_blocks_survive_as_image_url_parts(): + body = _tool_result_body( + [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "YWJjZA==", + }, + } + ] + ) + + chat_dict, _ = _anthropic_to_chat_request(body) + + assert chat_dict["messages"][2]["role"] == "tool" + assert chat_dict["messages"][2]["content"] == [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,YWJjZA=="}, + } + ] + + +def test_tool_result_url_image_block_survives_as_image_url_part(): + body = _tool_result_body( + [ + { + "type": "image", + "source": { + "type": "url", + "url": "https://example.com/image.png", + }, + } + ] + ) + + chat_dict, _ = _anthropic_to_chat_request(body) + + assert chat_dict["messages"][2]["role"] == "tool" + assert chat_dict["messages"][2]["content"] == [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + } + ] + + +def test_tool_result_textual_image_data_url_stays_text(): + data_url = "data:image/png;base64,YWJjZA==" + body = _tool_result_body(data_url) + + chat_dict, _ = _anthropic_to_chat_request(body) + + assert chat_dict["messages"][2]["role"] == "tool" + assert chat_dict["messages"][2]["content"] == data_url + + +def test_pdf_shaped_tool_schema_example_is_not_rewritten(monkeypatch): + body = _base_body() + body["tools"] = [ + { + "name": "inspect", + "description": "inspect", + "input_schema": { + "type": "object", + "examples": [_pdf_document_block()], + }, + } + ] + original = json.loads(json.dumps(body)) + monkeypatch.setattr( + api_anthropic, + "_ensure_pdf_parsing_supported", + lambda: pytest.fail("tool schemas are not message content"), + ) + + api_anthropic._replace_anthropic_pdf_documents(body) + + assert body == original + + # Helpers for streaming test def _chunk(delta, finish_reason=None, usage=None): obj = {"choices": [{"delta": delta, "finish_reason": finish_reason}]}