Fix simple_poem and use Sonnet 4.6 for multimodal testing (#714)

Author: mahesh-bespoke

examples/poem-generation/simple_poem.py

@@ -8,12 +8,12 @@
 # Use GPT-4o-mini for this example.
 llm = curator.LLM(model_name="gpt-4o-mini")
 poem = llm(["Write a poem about the importance of data in AI."])
-print(poem.to_pandas().iloc[0]["response"])
+print(poem.dataset.to_pandas().iloc[0]["response"])
 
-# Use Claude 3.5 Sonnet for this example.
-llm = curator.LLM(model_name="claude-3-5-sonnet-20240620", backend="litellm")
+# Use Claude Sonnet 4.6 for this example.
+llm = curator.LLM(model_name="claude-sonnet-4-6", backend="anthropic")
 poem = llm(["Write a poem about the importance of data in AI."])
-print(poem.to_pandas().iloc[0]["response"])
+print(poem.dataset.to_pandas().iloc[0]["response"])
 
 # Note that we can also pass a list of prompts to generate multiple responses.
 poems = llm(

src/bespokelabs/curator/request_processor/online/base_online_request_processor.py

@@ -273,10 +273,13 @@ def requests_to_responses(
         Args:
             generic_request_files: List of request files to process
         """
+        max_requests_per_minute = self.max_requests_per_minute
+        max_tokens_per_minute = self.max_tokens_per_minute
+
         self.tracker = OnlineStatusTracker(
             token_limit_strategy=self.token_limit_strategy,
-            max_requests_per_minute=self.max_requests_per_minute,
-            max_tokens_per_minute=self.max_tokens_per_minute,
+            max_requests_per_minute=max_requests_per_minute,
+            max_tokens_per_minute=max_tokens_per_minute,
             compatible_provider=self.compatible_provider,
             viewer_client=self._viewer_client,
             model=self.prompt_formatter.model_name,

src/bespokelabs/curator/request_processor/online/litellm_online_request_processor.py

@@ -22,6 +22,15 @@
 _OTPM_LIMIT["anthropic"] = "max_tokens"
 _CONCURRENT_ONLY_RATELIMIT_PROVIDERS = {"deepinfra"}
 _FILE_UPLOAD_LIMIT_PROVIDERS = {"openai": 20}  # MB
+_ANTHROPIC_MULTIMODAL_SUPPORTED_PREFIXES = (
+    "claude-3",
+    "claude-sonnet-4",
+    "claude-haiku-4",
+    "claude-opus-4",
+)
+_ANTHROPIC_PDF_BETA_HEADER = {"anthropic-beta": "pdfs-2024-09-25"}
+_ANTHROPIC_IMAGE_TOKEN_ESTIMATE = 1024
+_ANTHROPIC_DOCUMENT_TOKEN_ESTIMATE = 2048
 
 
 class LiteLLMOnlineRequestProcessor(BaseOnlineRequestProcessor):
@@ -48,7 +57,7 @@ def __init__(self, config: OnlineRequestProcessorConfig):
         # NOTE: Removed global litellm.api_base assignment - now passed per-request
         # to allow multiple LLM instances with different base_urls
         self.client = instructor.from_litellm(litellm.acompletion)
-        self.header_based_max_requests_per_minute, self.header_based_max_tokens_per_minute = self.get_header_based_rate_limits()
+        self._rate_limits_initialized = False
 
         self._validate_config(config)
         self._set_manual_tpm(config)
@@ -66,6 +75,12 @@ def _validate_config(self, config):
             assert config.max_input_tokens_per_minute is None, "`max_input_tokens_per_minute` cannot be used with `combined` token strategy"
             assert config.max_output_tokens_per_minute is None, "`max_output_tokens_per_minute` cannot be used with `combined` token strategy"
 
+    def _ensure_rate_limits(self):
+        """Lazily initialize rate limits from API headers on first access."""
+        if not self._rate_limits_initialized:
+            self.header_based_max_requests_per_minute, self.header_based_max_tokens_per_minute = self.get_header_based_rate_limits()
+            self._rate_limits_initialized = True
+
     @property
     def _provider(self):
         return self.config.model.split("/")[0]
@@ -96,6 +111,7 @@ def max_requests_per_minute(self) -> int | None:
         or uses default value as last resort.
         """
         if not self._concurrency_only_rate_limited:
+            self._ensure_rate_limits()
             return super().max_requests_per_minute
 
     @property
@@ -106,6 +122,7 @@ def max_tokens_per_minute(self) -> int | None:
         or uses default value as last resort.
         """
         if not self._concurrency_only_rate_limited:
+            self._ensure_rate_limits()
             return super().max_tokens_per_minute
 
     @property
@@ -129,7 +146,54 @@ def file_upload_limit_check(self, base64_image: str) -> None:
 
     @property
     def _multimodal_prompt_supported(self) -> bool:
-        return litellm.supports_vision(self.config.model)
+        if litellm.supports_vision(self.config.model):
+            return True
+
+        # LiteLLM's local model map can lag Anthropic releases and return a
+        # false negative for newer or provider-prefixed Claude aliases.
+        model_name = self.config.model.split("/", 1)[-1]
+        if model_name.startswith("claude-"):
+            return any(prefix in model_name for prefix in _ANTHROPIC_MULTIMODAL_SUPPORTED_PREFIXES)
+
+        return False
+
+    def _uses_anthropic_multimodal_format(self, model: str | None = None) -> bool:
+        model_name = (model or self.config.model).split("/", 1)[-1]
+        return model_name.startswith("claude-")
+
+    def _format_multimodal(self, data, mime_type="image/png"):
+        mime_type = mime_type or "image/png"
+        if not self._uses_anthropic_multimodal_format():
+            return super()._format_multimodal(data, mime_type=mime_type)
+
+        block_type = "image" if "image" in mime_type else "document"
+        if data.url and not data.is_local:
+            return {"type": block_type, "source": {"type": "url", "url": data.url}}
+
+        base64_content = data.serialize()
+        self.file_upload_limit_check(base64_content)
+        return {
+            "type": block_type,
+            "source": {
+                "type": "base64",
+                "media_type": mime_type,
+                "data": base64_content,
+            },
+        }
+
+    @staticmethod
+    def _has_anthropic_pdf_document(messages: list[dict]) -> bool:
+        for message in messages:
+            content = message.get("content", [])
+            if not isinstance(content, list):
+                continue
+            for block in content:
+                if block.get("type") != "document":
+                    continue
+                source = block.get("source", {})
+                if source.get("media_type") == "application/pdf" or source.get("type") == "url":
+                    return True
+        return False
 
     def check_structured_output_support(self):
         """Verify if the model supports structured output via litellm.
@@ -176,10 +240,42 @@ def estimate_total_tokens(self, messages: list) -> _TokenUsage:
         Returns:
             _TokenUsage: Total estimated tokens (input and output)
         """
-        input_tokens = litellm.token_counter(model=self.config.model, messages=messages)
+        try:
+            input_tokens = litellm.token_counter(model=self.config.model, messages=messages)
+        except ValueError:
+            if self._uses_anthropic_multimodal_format() and self._has_anthropic_pdf_document(messages):
+                input_tokens = self._estimate_anthropic_document_tokens(messages)
+            else:
+                raise
         output_tokens = int(self.estimate_output_tokens())
         return _TokenUsage(input=input_tokens, output=output_tokens)
 
+    def _estimate_anthropic_document_tokens(self, messages: list[dict]) -> int:
+        input_tokens = 0
+        for message in messages:
+            content = message.get("content", "")
+            role = message.get("role", "user")
+            if isinstance(content, str):
+                input_tokens += litellm.token_counter(model=self.config.model, messages=[{"role": role, "content": content}])
+                continue
+
+            if not isinstance(content, list):
+                continue
+
+            for block in content:
+                block_type = block.get("type")
+                if block_type == "text":
+                    input_tokens += litellm.token_counter(
+                        model=self.config.model,
+                        messages=[{"role": role, "content": block.get("text", "")}],
+                    )
+                elif block_type == "image":
+                    input_tokens += _ANTHROPIC_IMAGE_TOKEN_ESTIMATE
+                elif block_type == "document":
+                    input_tokens += _ANTHROPIC_DOCUMENT_TOKEN_ESTIMATE
+
+        return input_tokens
+
     def test_call(self):
         """Test call to get rate limits."""
         completion = litellm.completion(
@@ -282,6 +378,12 @@ def create_api_specific_request_online(self, generic_request: GenericRequest) ->
                 },
             ]
 
+        if self._uses_anthropic_multimodal_format(generic_request.model) and self._has_anthropic_pdf_document(generic_request.messages):
+            request["extra_headers"] = {
+                **request.get("extra_headers", {}),
+                **_ANTHROPIC_PDF_BETA_HEADER,
+            }
+
         return request
 
     async def call_single_request(

tests/integrations/helper.py

@@ -170,7 +170,7 @@ def create_camel(temp_working_dir, batch=False, batch_check_interval=1):
 
 _DEFAULT_MODEL_MAP = {
     "openai": "gpt-3.5-turbo",
-    "anthropic": "claude-3-5-sonnet-20241022",
+    "anthropic": "claude-sonnet-4-6",
     "litellm": "gpt-3.5-turbo",
     "vllm": "Qwen/Qwen2.5-1.5B-Instruct",
     "gemini": "gemini-1.5-flash-002",