feat: Add pdf_parsing to openrouter

Author: YoungVor

.github/workflows/test.yaml

@@ -48,6 +48,7 @@ jobs:
     env:
       GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
       OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }}
     strategy:
       matrix: ${{ fromJson(needs.matrix.outputs.matrix) }}
     steps:

src/fenic/_inference/google/gemini_token_counter.py

@@ -40,7 +40,7 @@ def __init__(self, model_name: str, fallback_encoding: str = "gemini-2.5-flash")
         except ValueError:
             self.google_tokenizer = LocalTokenizer(model_name=fallback_encoding)
 
-    def count_tokens(self, messages: Tokenizable) -> int:
+    def count_tokens(self, messages: Tokenizable, ignore_file: bool = False) -> int:
         """Count tokens for a string, message list, or `LMRequestMessages`.
 
         Args:
@@ -53,7 +53,7 @@ def count_tokens(self, messages: Tokenizable) -> int:
         if isinstance(messages, str):
             return self._count_text_tokens(messages)
         elif isinstance(messages, LMRequestMessages):
-            return self._count_request_tokens(messages)
+            return self._count_request_tokens(messages, ignore_file)
 
     def count_file_input_tokens(self, messages: LMRequestMessages) -> int:
         # Gemini 2.0 charges 258 tokens per page for all PDF inputs.  For more detail, see https://gemini-api.apidog.io/doc-965859#technical-details
@@ -68,7 +68,7 @@ def count_file_output_tokens(self, messages: LMRequestMessages) -> int:
         # In our estimates we add buffer, both for markdown structure and in case we ask the model to describe images.
         return self.google_tokenizer.count_tokens(text).total_tokens
 
-    def _count_request_tokens(self, messages: LMRequestMessages) -> int:
+    def _count_request_tokens(self, messages: LMRequestMessages, ignore_file: bool = False) -> int:
         """Count tokens for an `LMRequestMessages` object."""
         contents = convert_text_messages(messages)
         tokens = 0
@@ -79,7 +79,7 @@ def _count_request_tokens(self, messages: LMRequestMessages) -> int:
             ).total_tokens
             tokens += count_tokens
 
-        if messages.user_file:
+        if messages.user_file and not ignore_file:
             tokens += self.count_file_input_tokens(messages)
         return tokens
 

src/fenic/_inference/openrouter/openrouter_batch_chat_completions_client.py

@@ -10,6 +10,7 @@
 
 from fenic._inference.common_openai.openai_utils import convert_messages
 from fenic._inference.common_openai.utils import handle_openai_compatible_response
+from fenic._inference.google.gemini_token_counter import GeminiLocalTokenCounter
 from fenic._inference.model_client import (
     FatalException,
     ModelClient,
@@ -65,16 +66,23 @@ def __init__(
         profiles: Optional[dict[str, object]] = None,
         default_profile_name: Optional[str] = None,
     ):
+        # Choose token counter based on the model's provider
+        token_counter = None
+        provider_and_model = model.split("/")
+        if provider_and_model[0] == "google":
+            token_counter = GeminiLocalTokenCounter(model_name=provider_and_model[1])
+        else:
+            token_counter = TiktokenTokenCounter(
+                model_name=provider_and_model[1], fallback_encoding="o200k_base"
+            )
         super().__init__(
             model=model,
             model_provider=ModelProvider.OPENROUTER,
             model_provider_class=OpenRouterModelProvider(),
             rate_limit_strategy=rate_limit_strategy,
             queue_size=queue_size,
             max_backoffs=max_backoffs,
-            token_counter=TiktokenTokenCounter(
-                model_name=model, fallback_encoding="o200k_base"
-            ),
+            token_counter=token_counter,
         )
         self._model_parameters = model_catalog.get_completion_model_parameters(
             ModelProvider.OPENROUTER, model
@@ -87,17 +95,22 @@ def __init__(
         self._aio_client = OpenRouterModelProvider().aio_client
         self._metrics = LMMetrics()
 
+
+
     async def make_single_request(
         self, request: FenicCompletionsRequest
     ) -> Union[None, FenicCompletionsResponse, TransientException, FatalException]:
         profile = self._profile_manager.get_profile_by_name(request.model_profile)
         common_params = {
                 "model": self.model,
                 "messages": convert_messages(request.messages),
-                "max_completion_tokens": self._get_max_output_token_request_limit(request),
                 "n": 1,
             }
 
+        max_completion_tokens = self._get_max_output_token_request_limit(request)
+        if max_completion_tokens is not None:
+            common_params["max_completion_tokens"] = max_completion_tokens
+
         if request.top_logprobs:
             common_params.update(
                 {"logprobs": True, "top_logprobs": request.top_logprobs}
@@ -238,8 +251,8 @@ def estimate_tokens_for_request(
         self, request: FenicCompletionsRequest
     ) -> TokenEstimate:
         return TokenEstimate(
-            input_tokens=self.token_counter.count_tokens(request.messages),
-            output_tokens=self.token_counter.count_tokens(request.messages) + self._get_expected_additional_reasoning_tokens(request),
+            input_tokens=self._estimate_input_tokens(request),
+            output_tokens=self._estimate_output_tokens(request),
         )
 
     def reset_metrics(self):
@@ -248,16 +261,39 @@ def reset_metrics(self):
     def get_metrics(self) -> LMMetrics:
         return self._metrics
 
-    def _get_max_output_token_request_limit(self, request: FenicCompletionsRequest) -> int:
-        """Get the upper limit of output tokens for a request.
+    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)
+        return base_tokens + self._get_expected_additional_reasoning_tokens(request)
 
-        If max_completion_tokens is not set, don't apply a limit and return None.
+    def _get_max_output_token_request_limit(self, request: FenicCompletionsRequest) -> Optional[int]:
+        """Return the maximum output token limit for a request.
 
-        Include the thinking token budget with a safety margin."""
+        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 _estimate_input_tokens(self, request: FenicCompletionsRequest) -> int:
+        """Estimate the number of input tokens for a request."""
+        input_tokens = self.token_counter.count_tokens(request.messages, ignore_file=True)
+        if request.messages.user_file:
+            input_tokens += self._estimate_file_input_tokens(request)
+        return input_tokens
+
+    def _estimate_file_input_tokens(self, request: FenicCompletionsRequest) -> int:
+        """Estimate the number of input tokens from a file in a request."""
+        profile_config = self._profile_manager.get_profile_by_name(request.model_profile)
+        if profile_config.parsing_engine and profile_config.parsing_engine == "native":
+            return self.token_counter.count_file_input_tokens(messages=request.messages)
+        # OpenRouter's engine tool processes the file first and passes annotated text to the model.
+        # We can estimate by extracting the text and tokenizing it (which is what count_file_output_tokens does)
+        return self.token_counter.count_file_output_tokens(messages=request.messages)
+
     # This is a slightly less conservative estimate than the OpenRouter documentation on how reasoning_effort is used to
     # generate a reasoning.max_tokens for models that only support reasoning.max_tokens.
     # These percentages are slightly lower, since our use-cases generally require fewer reasoning tokens.

src/fenic/_inference/openrouter/openrouter_profile_manager.py

@@ -20,6 +20,7 @@
     ResolvedOpenRouterProviderRouting,
 )
 from fenic.core.types.provider_routing import StructuredOutputStrategy
+from fenic.core.types.semantic import ParsingEngine
 
 
 @dataclass
@@ -32,6 +33,7 @@ class OpenRouterCompletionProfileConfiguration(BaseProfileConfiguration):
     models: Optional[list[str]] = None
     provider: Optional[ResolvedOpenRouterProviderRouting] = None
     structured_output_strategy: Optional[StructuredOutputStrategy] = None
+    parsing_engine: Optional[ParsingEngine] = None
 
     @property
     def extra_body(self) -> dict[str, Any]:
@@ -70,6 +72,14 @@ def extra_body(self) -> dict[str, Any]:
         if reasoning_obj:
             reasoning_obj["exclude"] = True
             params["reasoning"] = reasoning_obj
+        params["plugins"] = [
+            {
+                "id": "file-parser",
+                "pdf": {
+                    "engine": self.parsing_engine or "native"
+                }
+            }
+        ]
         return params
 
 
@@ -116,13 +126,21 @@ def _process_profile(
                 and profile.reasoning_max_tokens is None
             ):
                 profile.reasoning_effort = "low"
+        
+        pdf_page_processing_cost = None
+        if profile.parsing_engine and profile.parsing_engine == "mistral-ocr":
+            pdf_page_processing_cost = 2/1000
+        elif profile.parsing_engine and profile.parsing_engine == "pdf-text":
+            pdf_page_processing_cost = 0
 
         return OpenRouterCompletionProfileConfiguration(
             models=profile.models,
             provider=profile.provider,
             reasoning_effort=profile.reasoning_effort,
             reasoning_max_tokens=profile.reasoning_max_tokens,
             structured_output_strategy=profile.structured_output_strategy,
+            parsing_engine=profile.parsing_engine,
+            pdf_page_processing_cost=pdf_page_processing_cost,
         )
 
     def get_default_profile(self) -> OpenRouterCompletionProfileConfiguration:

src/fenic/_inference/openrouter/openrouter_provider.py

@@ -123,6 +123,7 @@ def _translate_model(
             supports_reasoning=supports_reasoning,
             supports_custom_temperature=supports_custom_temperature,
             supports_verbosity=supports_verbosity,
+            supports_pdf_parsing=True, # Even if the model doesn't support pdf file processing, OpenRouter can use its separate processing engines
             supported_parameters=supported_params,
         )
 

src/fenic/_inference/token_counter.py

@@ -10,7 +10,7 @@
 Tokenizable = Union[str | LMRequestMessages]
 
 class TokenCounter(Protocol):
-    def count_tokens(self, messages: Tokenizable) -> int: ...
+    def count_tokens(self, messages: Tokenizable, ignore_file: bool = False) -> int: ...
     def count_file_input_tokens(self, messages: LMRequestMessages) -> int: ...
     def count_file_output_tokens(self, messages: LMRequestMessages) -> int: ...
 
@@ -22,11 +22,11 @@ def __init__(self, model_name: str, fallback_encoding: str = "o200k_base"):
         except KeyError:
             self.tokenizer = tiktoken.get_encoding(fallback_encoding)
 
-    def count_tokens(self, messages: Tokenizable) -> int:
+    def count_tokens(self, messages: Tokenizable, ignore_file: bool = False) -> int:
         if isinstance(messages, str):
             return len(self.tokenizer.encode(messages))
         elif isinstance(messages, LMRequestMessages):
-            return self._count_message_tokens(messages)
+            return self._count_message_tokens(messages, ignore_file)
         else:
             raise TypeError(f"Expected str or LMRequestMessages, got {type(messages)}")
 
@@ -55,7 +55,7 @@ def count_file_output_tokens(self, messages: LMRequestMessages) -> int:
         else:
             raise InternalError(f"File{messages.user_file.path}'s extension is not supported for llm completions.")
 
-    def _count_message_tokens(self, messages: LMRequestMessages) -> int:
+    def _count_message_tokens(self, messages: LMRequestMessages, ignore_file: bool = False) -> int:
         num_tokens = 0
         message_count = 2 # system message and user parent message
         num_tokens += self.count_tokens(messages.system)
@@ -66,7 +66,7 @@ def _count_message_tokens(self, messages: LMRequestMessages) -> int:
             num_tokens += self.count_tokens(example.user)
             num_tokens += self.count_tokens(example.assistant)
             message_count += 2
-        if messages.user_file:
+        if messages.user_file and not ignore_file:
             num_tokens += self.count_file_input_tokens(messages)
             message_count += 1
         num_tokens += message_count * PREFIX_TOKENS_PER_MESSAGE

src/fenic/api/session/config.py

@@ -644,6 +644,7 @@ class Profile(BaseModel):
             ge=1024,
         )
 
+ParsingEngine = Literal["mistral-ocr", "pdf-text", "native"]
 
 class OpenRouterLanguageModel(BaseModel):
     """Configuration for OpenRouter language models.
@@ -788,6 +789,8 @@ class Profile(BaseModel):
                 If the model does support reasoning, but not `reasoning_max_tokens`, a `reasoning_effort_ will be automatically
                 calculated based on `reasoning_max_tokens` as a percentage of the model's maximum output size
                 ([OpenRouter Documentation](https://openrouter.ai/docs/use-cases/reasoning-tokens#max-tokens-for-reasoning))
+            parsing_engine: The parsing engine to use for processing PDF files.  By default, the model's native parsing engine will be used.  If the model doesn't support PDF processing and the parsing engine is not provided, an error will be raised.  Note: 'mistral-ocr' incurs additional costs.
+                ([OpenRouter Documentation](https://openrouter.ai/docs/features/multimodal/pdfs))
         """
         model_config = ConfigDict(extra="forbid")
 
@@ -807,6 +810,10 @@ class Profile(BaseModel):
         provider: Optional[OpenRouterLanguageModel.Provider] = Field(
             default=None, description="Provider routing configuration"
         )
+        parsing_engine: Optional[ParsingEngine] = Field(
+            default=None,
+            description="The parsing engine to use for processing PDF files. By default, the model's native parsing engine will be used."
+        )
 
 
 CohereEmbeddingTaskType = Literal[
@@ -1341,6 +1348,7 @@ def resolve_model(model: ModelConfig) -> ResolvedModelConfig:
                             reasoning_effort=profile.reasoning_effort,
                             reasoning_max_tokens=profile.reasoning_max_tokens,
                             models=profile.models,
+                            parsing_engine=profile.parsing_engine,
                             provider=(
                                 ResolvedOpenRouterProviderRouting(
                                     **(profile.provider.model_dump())

src/fenic/core/_inference/model_catalog.py

@@ -32,7 +32,6 @@ def __init__(
         self.cached_input_token_write_cost = cached_input_token_write_cost
         self.output_token_cost = output_token_cost
 
-
 class CompletionModelParameters:
     """Parameters for completion models including costs and context window size.
 
@@ -42,6 +41,7 @@ class CompletionModelParameters:
         cached_input_token_write_cost: Cost per cached input token write in USD
         output_token_cost: Cost per output token in USD
         context_window_length: Maximum number of tokens in the context window
+        tiered_token_costs: Tiered token costs for the model.
         max_output_tokens: Maximum number of tokens the model can generate in a single request.
         max_temperature: Maximum temperature for the model.
         supports_profiles: Whether the model supports parameter profiles.
@@ -897,7 +897,6 @@ def _initialize_google_gla_models(self):
                 context_window_length=1_048_576,
                 max_output_tokens=8_192,
                 max_temperature=2.0,
-                supports_profiles=False,
                 supports_pdf_parsing=True,
             ),
             snapshots=["gemini-2.0-flash-lite-001"],
@@ -914,7 +913,6 @@ def _initialize_google_gla_models(self):
                 context_window_length=1_048_576,
                 max_output_tokens=8_192,
                 max_temperature=2.0,
-                supports_profiles=False,
                 supports_pdf_parsing=True,
             ),
             snapshots=["gemini-2.0-flash-001", "gemini-2.0-flash-exp"],

src/fenic/core/_resolved_session_config.py

@@ -19,6 +19,7 @@
     ProviderSort,
     StructuredOutputStrategy,
 )
+from fenic.core.types.semantic import ParsingEngine
 
 ReasoningEffort = Literal["minimal", "low", "medium", "high"]
 Verbosity = Literal["low", "medium", "high"]
@@ -79,6 +80,7 @@ class ResolvedOpenRouterModelProfile:
     models: Optional[list[str]] = None
     provider: Optional[ResolvedOpenRouterProviderRouting] = None
     structured_output_strategy: Optional[StructuredOutputStrategy] = None
+    parsing_engine: Optional[ParsingEngine] = None
 
 
 @dataclass
@@ -128,6 +130,7 @@ class ResolvedOpenRouterModelConfig:
     profiles: Optional[dict[str, ResolvedOpenRouterModelProfile]] = None
     model_provider: ModelProvider = ModelProvider.OPENROUTER
     default_profile: Optional[str] = None
+    parsing_engine: Optional[ParsingEngine] = None
 
 
 ResolvedModelConfig = Union[

src/fenic/core/types/semantic.py

@@ -1,12 +1,13 @@
 """Types used to configure model selection for semantic functions."""
 from __future__ import annotations
 
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 from pydantic import BaseModel
 
 from fenic.core._logical_plan.resolved_types import ResolvedModelAlias
 
+ParsingEngine = Literal["mistral-ocr", "pdf-text", "native"]
 
 class ModelAlias(BaseModel):
     """A combination of a model name and a required profile for that model.