Repository navigation
fix(discovery): persist provider token ceilings and context windows #952
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -25,7 +25,7 @@ | |
| import urllib.request | ||
| import certifi | ||
| from dataclasses import dataclass, replace | ||
| from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence | ||
| from urllib.parse import quote, urlsplit, urlunsplit | ||
|
|
||
| from .chat_capability import ( | ||
|
|
@@ -273,6 +273,10 @@ class DiscoveredModel: | |
| capabilities: tuple[str, ...] = () | ||
| input_modalities: tuple[str, ...] = () | ||
| output_modalities: tuple[str, ...] = () | ||
| max_output_tokens: int | None = None | ||
| context_window: int | None = None | ||
| max_output_tokens_conflicted: bool = False | ||
| context_window_conflicted: bool = False | ||
| prompt_price_per_1k: float | None = None | ||
| completion_price_per_1k: float | None = None | ||
| currency_code: str = "USD" | ||
|
|
@@ -296,6 +300,29 @@ def __init__(self, provider_name: str, error_code: str) -> None: | |
| super().__init__(f"model discovery failed for provider {provider_name!r}: {error_code}") | ||
|
|
||
|
|
||
| def _positive_int_metadata(value: object) -> int | None: | ||
| """Return one exact positive integer metadata field or ``None`` when unknown.""" | ||
| if isinstance(value, bool): | ||
| return None | ||
| if isinstance(value, int): | ||
| return value if 0 < value <= 9_223_372_036_854_775_807 else None | ||
| if isinstance(value, float): | ||
| if not math.isfinite(value) or not value.is_integer(): | ||
| return None | ||
| integer = int(value) | ||
| return integer if 0 < integer <= 9_223_372_036_854_775_807 else None | ||
| if isinstance(value, str): | ||
| stripped = value.strip() | ||
| if not stripped.isdigit(): | ||
| return None | ||
| try: | ||
| integer = int(stripped) | ||
| except ValueError: | ||
| return None | ||
| return integer if 0 < integer <= 9_223_372_036_854_775_807 else None | ||
| return None | ||
|
seonghobae marked this conversation as resolved.
|
||
|
|
||
|
|
||
| def _fetch_json(url: str, *, api_key: str = "", auth_scheme: str = "Bearer", timeout: float) -> Any: | ||
| if not url.startswith("https://"): | ||
| # Every caller passes one of the hardcoded PROVIDER_SOURCES chat_base_url | ||
|
|
@@ -503,6 +530,8 @@ def _deduplicate_discovered_models( | |
| chosen = min((previous, model), key=_source_tiebreaker) | ||
| unique[identity] = replace( | ||
| chosen, | ||
| max_output_tokens=None, | ||
| context_window=None, | ||
|
seonghobae marked this conversation as resolved.
|
||
| prompt_price_per_1k=None, | ||
| completion_price_per_1k=None, | ||
| unit_prices=(), | ||
|
|
@@ -625,6 +654,7 @@ def _merge_models_dev_metadata(payload: Any, metadata: Any, provider: str) -> An | |
| if _valid_price_component(value): | ||
| pricing[target_key] = str(Decimal(str(value)) / Decimal(1_000_000)) | ||
| modalities = model.get("modalities") if isinstance(model.get("modalities"), dict) else {} | ||
| limits = model.get("limit") if isinstance(model.get("limit"), dict) else {} | ||
| enriched.append( | ||
| { | ||
| **row, | ||
|
|
@@ -633,6 +663,8 @@ def _merge_models_dev_metadata(payload: Any, metadata: Any, provider: str) -> An | |
| "input_modalities": modalities.get("input"), | ||
| "output_modalities": modalities.get("output"), | ||
| }, | ||
| "max_output_tokens": _positive_int_metadata(limits.get("output")), | ||
| "context_window": _positive_int_metadata(limits.get("context")), | ||
| "is_free": _models_dev_cost_is_free(cost), | ||
| } | ||
| ) | ||
|
|
@@ -676,6 +708,10 @@ def _merge_configured_gateway_metadata(payload: Any, metadata: Any) -> Any: | |
| row.pop("pricing", None) | ||
| row.pop("architecture", None) | ||
| row.pop("unit_pricing", None) | ||
| row.pop("max_output_tokens", None) | ||
| row.pop("context_window", None) | ||
|
seonghobae marked this conversation as resolved.
|
||
| row.pop("context_length", None) | ||
| row.pop("max_completion_tokens", None) | ||
|
seonghobae marked this conversation as resolved.
|
||
| for key in ( | ||
| "supports_zero_data_retention", | ||
| "supports_no_training", | ||
|
|
@@ -689,6 +725,8 @@ def _merge_configured_gateway_metadata(payload: Any, metadata: Any) -> Any: | |
| prices: set[tuple[object, object]] = set() | ||
| pricing_complete = bool(model_details) | ||
| unit_price_maps: list[tuple[tuple[str, object], ...]] = [] | ||
| completion_limits: list[int | None] = [] | ||
| context_windows: list[int | None] = [] | ||
| privacy_values = { | ||
| key: [] | ||
| for key in ( | ||
|
|
@@ -746,6 +784,22 @@ def _merge_configured_gateway_metadata(payload: Any, metadata: Any) -> Any: | |
| completion = info.get( | ||
| "output_cost_per_token", params.get("output_cost_per_token") | ||
| ) | ||
| completion_limits.append( | ||
| _positive_int_metadata( | ||
| info.get("max_output_tokens", params.get("max_output_tokens")) | ||
| ) | ||
| or _positive_int_metadata( | ||
| info.get("max_completion_tokens", params.get("max_completion_tokens")) | ||
| ) | ||
| ) | ||
| context_windows.append( | ||
| _positive_int_metadata( | ||
| info.get("context_window", params.get("context_window")) | ||
| ) | ||
| or _positive_int_metadata( | ||
| info.get("context_length", params.get("context_length")) | ||
| ) | ||
| ) | ||
| if _valid_price_component(prompt) and _valid_price_component(completion): | ||
| prices.add((prompt, completion)) | ||
| else: | ||
|
|
@@ -777,6 +831,18 @@ def _merge_configured_gateway_metadata(payload: Any, metadata: Any) -> Any: | |
| "input_modalities": list(deployment_inputs[0]), | ||
| "output_modalities": list(deployment_outputs[0]), | ||
| } | ||
| if completion_limits and all(value is not None for value in completion_limits): | ||
| unique_completion_limits = {value for value in completion_limits if value is not None} | ||
| if len(unique_completion_limits) == 1: | ||
| row["max_output_tokens"] = unique_completion_limits.pop() | ||
| else: | ||
| row["_max_output_tokens_conflicted"] = True | ||
| if context_windows and all(value is not None for value in context_windows): | ||
| unique_context_windows = {value for value in context_windows if value is not None} | ||
| if len(unique_context_windows) == 1: | ||
| row["context_window"] = unique_context_windows.pop() | ||
|
seonghobae marked this conversation as resolved.
|
||
| else: | ||
| row["_context_window_conflicted"] = True | ||
| if pricing_complete and len(prices) == 1: | ||
| prompt, completion = prices.pop() | ||
| if prompt is not None and completion is not None: | ||
|
|
@@ -941,10 +1007,29 @@ def _parse_openai_compatible(payload: Any, source: ProviderModelSource) -> list[ | |
| if isinstance(row.get("supported_parameters"), list) | ||
| else [] | ||
| ) | ||
| top_provider = row.get("top_provider") if isinstance(row.get("top_provider"), dict) else {} | ||
| raw_inputs = architecture.get("input_modalities") | ||
| raw_outputs = architecture.get("output_modalities") | ||
| inputs = tuple(value for value in raw_inputs if isinstance(value, str)) if isinstance(raw_inputs, list) else () | ||
| outputs = tuple(value for value in raw_outputs if isinstance(value, str)) if isinstance(raw_outputs, list) else () | ||
| raw_output_limits = [ | ||
| value | ||
| for container, key in ( | ||
| (row, "max_output_tokens"), | ||
| (top_provider, "max_completion_tokens"), | ||
| ) | ||
| if key in container | ||
| for value in (container[key],) | ||
| ] | ||
| raw_context_limits = [ | ||
| row[key] | ||
| for key in ("context_window", "context_length") | ||
| if key in row | ||
| ] | ||
| output_limits = [_positive_int_metadata(value) for value in raw_output_limits] | ||
| context_limits = [_positive_int_metadata(value) for value in raw_context_limits] | ||
| max_output_tokens = next((value for value in output_limits if value is not None), None) | ||
| context_window = next((value for value in context_limits if value is not None), None) | ||
| if ( | ||
| not outputs | ||
| and not any(capability != "chat" for capability in source.capabilities) | ||
|
|
@@ -1004,6 +1089,16 @@ def _parse_openai_compatible(payload: Any, source: ProviderModelSource) -> list[ | |
| capabilities=capabilities, | ||
| input_modalities=inputs, | ||
| output_modalities=outputs, | ||
| max_output_tokens=max_output_tokens, | ||
| context_window=context_window, | ||
| max_output_tokens_conflicted=( | ||
| row.get("_max_output_tokens_conflicted") is True | ||
| or any(value is None for value in output_limits) | ||
|
Contributor
There was a problem hiding this comment. 🟡 Missing enrichment clears saved limits When Models.dev omits a limit, Prompt for agents
Was this helpful? React with 👍 or 👎 to provide feedback. |
||
| ), | ||
| context_window_conflicted=( | ||
| row.get("_context_window_conflicted") is True | ||
| or any(value is None for value in context_limits) | ||
| ), | ||
| prompt_price_per_1k=prompt_price, | ||
| completion_price_per_1k=completion_price, | ||
| unit_prices=unit_prices, | ||
|
|
@@ -1484,6 +1579,8 @@ def agent_from_discovered(discovered: DiscoveredModel, *, priority: int = 0) -> | |
| ), | ||
| priority=priority, | ||
| disabled=True, | ||
| max_output_tokens=discovered.max_output_tokens, | ||
| context_window=discovered.context_window, | ||
|
seonghobae marked this conversation as resolved.
|
||
| ) | ||
|
|
||
|
|
||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.