feat: Add openai support for semantic parse_pdf (#253)

### TL;DR

Added PDF file support for OpenAI models with proper token counting and
estimation.

### What changed?

- Made `max_completion_tokens` optional in OpenAI chat completions
requests
- Implemented PDF file token counting for OpenAI models:
    - Added methods to count tokens for PDF files in input and output
- Updated token counter to handle PDF files with proper token estimation
- Added support for PDF parsing to various OpenAI models in the model
catalog
- Refactored openai token counting logic
- Mini refactor - separate `_max_output_tokens` user limit concept from
`_estimate_output_tokens` for cost estimation and throttling
    - added to openai and updated gemini
    - Max tokens (put in request)
        - Use max tokens provided by semantic_operator if exists
**- OR, for page parsing specifically, use an upper limit based on
output limit of our smallest VLM supported (8000 tokens)**
        - Add expected reasoning effort
    - Estimate output tokens (for cost estimate and throttling)
        - Use max tokens provided by semantic_operator if exists
        - OR estimate file output tokens
        - Add expected reasoning effort
- Added openai models to semantic_parse_pdf tests

### Out of scope

The token estimation should happen at the semantic operator level, since
it has the context of what its expecting from the model. Currently,
semantic operator only passes 'max token' limit to the client and we use
that upper limit in our estimates. As a future improvement we should
refactor and have the semantic operator decide on the output token limit
for the request

### How to test?

1. Run the new token counter tests: `pytest
tests/_inference/test_openai_token_counter.py`
2. Test PDF parsing with OpenAI models: `pytest
tests/_backends/local/functions/test_semantic_parse_pdf.py`
3. Verify that token estimation works correctly with PDF files by using
a model that supports PDF parsing

Author: YoungVor

src/fenic/_backends/local/semantic_operators/parse_pdf.py

@@ -48,11 +48,13 @@ def __init__(
         page_separator: Optional[str] = None,
         describe_images: bool = False,
         model_alias: Optional[ResolvedModelAlias] = None,
+        max_output_tokens: Optional[int] = None,
     ):
         self.page_separator = page_separator
         self.describe_images = describe_images
         self.model = model
         self.model_alias = model_alias
+        self.max_output_tokens = max_output_tokens
 
         DocFolderLoader.check_file_extensions(input.to_list(), "pdf")
 
@@ -62,7 +64,7 @@ def __init__(
                 model=model,
                 operator_name="semantic.parse_pdf",
                 inference_config=InferenceConfiguration(
-                    max_output_tokens=None,
+                    max_output_tokens=max_output_tokens,
                     temperature=1.0,  # Use a higher temperature so gemini flash models can handle complex table formatting.  For more info see the conversation here: https://discuss.ai.google.dev/t/gemini-2-0-flash-has-a-weird-bug/65119/26
                     model_profile=model_alias.profile if model_alias else None,
                 ),

src/fenic/_backends/local/transpiler/expr_converter.py

@@ -712,8 +712,9 @@ def parse_pdf_fn(batch: pl.Series) -> pl.Series:
                 page_separator=logical.page_separator,
                 describe_images=logical.describe_images,
                 model_alias=logical.model_alias,
+                max_output_tokens=logical.max_output_tokens,
             ).execute()
-        
+
         return self._convert_expr(logical.expr).map_batches(
             parse_pdf_fn, return_dtype=pl.Utf8
         )

src/fenic/_inference/anthropic/anthropic_batch_chat_completions_client.py

@@ -275,7 +275,7 @@ def _estimate_structured_output_overhead(self, response_format) -> int:
         """
         return self.estimate_response_format_tokens(response_format)
 
-    def _get_max_output_tokens(self, request: FenicCompletionsRequest) -> int:
+    def _get_max_output_token_request_limit(self, request: FenicCompletionsRequest) -> int:
         """Get maximum output tokens including thinking budget.
 
         Args:
@@ -329,7 +329,7 @@ def estimate_tokens_for_request(self, request: FenicCompletionsRequest):
         input_tokens += self._count_auxiliary_input_tokens(request)
 
         # Estimate output tokens
-        output_tokens = self._get_max_output_tokens(request)
+        output_tokens = self._get_max_output_token_request_limit(request)
 
         return TokenEstimate(
             input_tokens=input_tokens,

src/fenic/_inference/cohere/cohere_batch_embeddings_client.py

@@ -171,7 +171,7 @@ def estimate_tokens_for_request(self, request: FenicEmbeddingsRequest) -> TokenE
             output_tokens=0
         )
 
-    def _get_max_output_tokens(self, request: FenicEmbeddingsRequest) -> int:
+    def _get_max_output_token_request_limit(self, request: FenicEmbeddingsRequest) -> int:
         """Get maximum output tokens (always 0 for embeddings).
         
         Returns:

src/fenic/_inference/common_openai/openai_chat_completions_core.py

@@ -90,9 +90,13 @@ async def make_single_request(
             common_params: dict[str, Any] = {
                 "model": self._model,
                 "messages": convert_messages(request.messages),
-                "max_completion_tokens": request.max_completion_tokens + profile_configuration.expected_additional_reasoning_tokens,
                 "n": 1,
             }
+
+            max_completion_tokens = self.get_max_output_token_request_limit(request, profile_configuration)
+            if max_completion_tokens is not None:
+                common_params["max_completion_tokens"] = max_completion_tokens
+
             if request.temperature:
                 common_params.update({"temperature": request.temperature})
 
@@ -213,3 +217,13 @@ def get_request_key(self, request: FenicCompletionsRequest) -> str:
             A unique key for the request
         """
         return generate_completion_request_key(request)
+
+    def get_max_output_token_request_limit(self, request: FenicCompletionsRequest, profile_config:OpenAICompletionProfileConfiguration) -> Optional[int]:
+        """Return the maximum output token limit for a request.
+
+        Returns None if max_completion_tokens is not provided (no limit should be set).
+        If max_completion_tokens is provided, includes the thinking token budget with a safety margin.
+        """
+        if request.max_completion_tokens is None:
+            return None
+        return request.max_completion_tokens + profile_config.expected_additional_reasoning_tokens

src/fenic/_inference/google/gemini_batch_embeddings_client.py

@@ -121,7 +121,7 @@ def estimate_tokens_for_request(self, request: FenicEmbeddingsRequest) -> TokenE
             input_tokens=self.token_counter.count_tokens(request.doc), output_tokens=0
         )
 
-    def _get_max_output_tokens(self, request: FenicEmbeddingsRequest) -> int:
+    def _get_max_output_token_request_limit(self, request: FenicEmbeddingsRequest) -> int:
         return 0
 
     def reset_metrics(self):

src/fenic/_inference/google/gemini_native_chat_completions_client.py

@@ -132,56 +132,6 @@ def count_tokens(self, messages: Tokenizable) -> int:  # type: ignore[override]
         # Re-expose for mypy – same implementation as parent.
         return super().count_tokens(messages)
 
-    def _estimate_structured_output_overhead(self, response_format: ResolvedResponseFormat) -> int:
-        """Use Google-specific response schema token estimation.
-
-        Args:
-            response_format: Pydantic model class defining the response format
-
-        Returns:
-            Estimated token overhead for structured output
-        """
-        return self._estimate_response_schema_tokens(response_format)
-
-    def _get_max_output_tokens(self, request: FenicCompletionsRequest) -> Optional[int]:
-        """Get maximum output tokens including thinking budget.
-
-        If max_completion_tokens is not set, return None.
-
-        Conservative estimate that includes both completion tokens and
-        thinking token budget with a safety margin.
-
-        Args:
-            request: The completion request
-
-        Returns:
-            Maximum output tokens (completion + thinking budget with safety margin)
-        """
-        if request.max_completion_tokens is None:
-            return None
-        profile_config = self._profile_manager.get_profile_by_name(
-            request.model_profile
-        )
-        return request.max_completion_tokens + int(
-            1.5 * profile_config.thinking_token_budget
-        )
-
-    @cache  # noqa: B019 – builtin cache OK here.
-    def _estimate_response_schema_tokens(self, response_format: ResolvedResponseFormat) -> int:
-        """Estimate token count for a response format schema.
-
-        Uses Google's tokenizer to count tokens in a JSON schema representation
-        of the response format. Results are cached for performance.
-
-        Args:
-            response_format: Pydantic model class defining the response format
-
-        Returns:
-            Estimated token count for the response format
-        """
-        schema_str = response_format.schema_fingerprint
-        return self._token_counter.count_tokens(schema_str)
-
     def get_request_key(self, request: FenicCompletionsRequest) -> str:
         """Generate a unique key for the request.
 
@@ -196,19 +146,17 @@ def get_request_key(self, request: FenicCompletionsRequest) -> str:
     def estimate_tokens_for_request(self, request: FenicCompletionsRequest):
         """Estimate the number of tokens for a request.
 
+        If the request provides a max_completion_tokens value, use that.  Otherwise, estimate the output tokens based on the file size.
+
         Args:
             request: The request to estimate tokens for
 
         Returns:
             TokenEstimate: The estimated token usage
         """
-
-        # Count input tokens
         input_tokens = self.count_tokens(request.messages)
         input_tokens += self._count_auxiliary_input_tokens(request)
-
-        output_tokens = self._get_max_output_tokens(request) or self._model_parameters.max_output_tokens
-
+        output_tokens = self._estimate_output_tokens(request)
         return TokenEstimate(input_tokens=input_tokens, output_tokens=output_tokens)
 
     async def make_single_request(
@@ -228,16 +176,17 @@ async def make_single_request(
         """
 
         profile_config = self._profile_manager.get_profile_by_name(request.model_profile)
-        max_output_tokens = self._get_max_output_tokens(request)
-
         generation_config: GenerateContentConfigDict = {
             "temperature": request.temperature,
             "response_logprobs": request.top_logprobs is not None,
             "logprobs": request.top_logprobs,
             "system_instruction": request.messages.system,
         }
+
+        max_output_tokens = self._get_max_output_token_request_limit(request)
         if max_output_tokens is not None:
             generation_config["max_output_tokens"] = max_output_tokens
+
         generation_config.update(profile_config.additional_generation_config)
         if request.structured_output is not None:
             generation_config.update(
@@ -355,3 +304,54 @@ async def make_single_request(
         finally:
             if file_obj:
                 await delete_file(self._client, file_obj.name)
+
+    @cache  # noqa: B019 – builtin cache OK here.
+    def _estimate_response_schema_tokens(self, response_format: ResolvedResponseFormat) -> int:
+        """Estimate token count for a response format schema.
+
+        Uses Google's tokenizer to count tokens in a JSON schema representation
+        of the response format. Results are cached for performance.
+
+        Args:
+            response_format: Pydantic model class defining the response format
+
+        Returns:
+            Estimated token count for the response format
+        """
+        schema_str = response_format.schema_fingerprint
+        return self._token_counter.count_tokens(schema_str)
+
+    def _estimate_structured_output_overhead(self, response_format: ResolvedResponseFormat) -> int:
+        """Use Google-specific response schema token estimation.
+
+        Args:
+            response_format: Pydantic model class defining the response format
+
+        Returns:
+            Estimated token overhead for structured output
+        """
+        return self._estimate_response_schema_tokens(response_format)
+
+    def _estimate_output_tokens(self, request: FenicCompletionsRequest) -> int:
+        """Estimate the number of output tokens for a request."""
+        estimated_output_tokens = request.max_completion_tokens or 0
+        if request.max_completion_tokens is None and request.messages.user_file:
+            # TODO(DY): the semantic operator should dictate how the file affects the token estimate
+            estimated_output_tokens = self.token_counter.count_file_output_tokens(request.messages)
+        return estimated_output_tokens + self._get_expected_additional_reasoning_tokens(request)
+
+    def _get_max_output_token_request_limit(self, request: FenicCompletionsRequest) -> Optional[int]:
+        """Get the upper limit of output tokens for a request.
+
+        Returns None if max_completion_tokens is not provided (no limit should be set).
+        If max_completion_tokens is provided, includes the thinking token budget with a safety margin."""
+        if request.max_completion_tokens is None:
+            return None
+        return request.max_completion_tokens + self._get_expected_additional_reasoning_tokens(request)
+
+    def _get_expected_additional_reasoning_tokens(self, request: FenicCompletionsRequest) -> int:
+        """Get the expected additional reasoning tokens for a request.  Include a safety margin."""
+        profile_config = self._profile_manager.get_profile_by_name(request.model_profile)
+        return int(
+            1.5 * profile_config.thinking_token_budget
+        )

src/fenic/_inference/language_model.py

@@ -22,7 +22,7 @@
 
 @dataclass
 class InferenceConfiguration:
-    # If max_output_tokens is not provided, do not include it in the request.
+    # If max_output_tokens is not provided, model_client will add a guardrail based on the estimated output tokens.
     max_output_tokens: Optional[int]
     temperature: float
     top_logprobs: Optional[int] = None

src/fenic/_inference/model_client.py

@@ -245,8 +245,8 @@ def _estimate_structured_output_overhead(self, response_format: ResolvedResponse
 
 
     @abstractmethod
-    def _get_max_output_tokens(self, request: RequestT) -> int:
-        """Get conservative output token estimate. Override in subclasses for provider-specific logic."""
+    def _get_max_output_token_request_limit(self, request: RequestT) -> int:
+        """Get the upper limit of output tokens to set on a request."""
         pass
 
     #

src/fenic/_inference/openai/openai_batch_chat_completions_client.py

@@ -65,6 +65,7 @@ def __init__(
             profile_configurations=profiles,
             default_profile_name=default_profile_name,
         )
+
         self._core = OpenAIChatCompletionsCore(
             model=model,
             model_provider=ModelProvider.OPENAI,
@@ -108,7 +109,7 @@ def estimate_tokens_for_request(self, request: FenicCompletionsRequest) -> Token
         """
         return TokenEstimate(
             input_tokens=self.token_counter.count_tokens(request.messages),
-            output_tokens=self._get_max_output_tokens(request)
+            output_tokens=self._estimate_output_tokens(request)
         )
 
     def reset_metrics(self):
@@ -123,10 +124,21 @@ def get_metrics(self) -> LMMetrics:
         """
         return self._core.get_metrics()
 
-    def _get_max_output_tokens(self, request: FenicCompletionsRequest) -> int:
-        """Conservative estimate: max_completion_tokens + reasoning effort-based thinking tokens."""
-        base_tokens = request.max_completion_tokens
-
-        # Get profile-specific reasoning effort
+    def _estimate_output_tokens(self, request: FenicCompletionsRequest) -> int:
+        """Estimate the number of output tokens for a request."""
+        base_tokens = request.max_completion_tokens or 0
+        if request.max_completion_tokens is None and request.messages.user_file:
+            # TODO(DY): the semantic operator should dictate how the file affects the token estimate
+            base_tokens += self.token_counter.count_file_output_tokens(messages=request.messages)
         profile_config = self._profile_manager.get_profile_by_name(request.model_profile)
         return base_tokens + profile_config.expected_additional_reasoning_tokens
+
+    def _get_max_output_token_request_limit(self, request: FenicCompletionsRequest) -> int:
+        """Return the maximum output token limit for a request.
+
+        For file parsing requests, use a guardrail limit of 8192 tokens (the lowest output limit of a VLM model we support).
+
+        Include the thinking token budget with a safety margin.
+        """
+        profile_config = self._profile_manager.get_profile_by_name(request.model_profile)
+        return self._core.get_max_output_token_request_limit(request, profile_config)