import json import os from collections.abc import Mapping from dataclasses import dataclass, field from typing import Any, Literal, TypeAlias, TypeVar, overload import httpx from openai import APIConnectionError, APIStatusError, AsyncOpenAI, RateLimitError from openai.types.chat.chat_completion import ChatCompletion from openai.types.shared_params.response_format_json_schema import ( JSONSchema, ResponseFormatJSONSchema, ) from pydantic import BaseModel from browser_use.llm.base import BaseChatModel from browser_use.llm.exceptions import ModelProviderError, ModelRateLimitError from browser_use.llm.messages import BaseMessage, ContentPartTextParam, SystemMessage from browser_use.llm.schema import SchemaOptimizer from browser_use.llm.vercel.serializer import VercelMessageSerializer from browser_use.llm.views import ChatInvokeCompletion, ChatInvokeUsage T = TypeVar('T', bound=BaseModel) ChatVercelModel: TypeAlias = Literal[ 'alibaba/qwen-3-14b', 'alibaba/qwen-3-235b', 'alibaba/qwen-3-30b', 'alibaba/qwen-3-32b', 'alibaba/qwen3-235b-a22b-thinking', 'alibaba/qwen3-coder', 'alibaba/qwen3-coder-30b-a3b', 'alibaba/qwen3-coder-next', 'alibaba/qwen3-coder-plus', 'alibaba/qwen3-embedding-0.6b', 'alibaba/qwen3-embedding-4b', 'alibaba/qwen3-embedding-8b', 'alibaba/qwen3-max', 'alibaba/qwen3-max-preview', 'alibaba/qwen3-max-thinking', 'alibaba/qwen3-next-80b-a3b-instruct', 'alibaba/qwen3-next-80b-a3b-thinking', 'alibaba/qwen3-vl-instruct', 'alibaba/qwen3-vl-thinking', 'alibaba/qwen3.5-flash', 'alibaba/qwen3.5-plus', 'alibaba/wan-v2.5-t2v-preview', 'alibaba/wan-v2.6-i2v', 'alibaba/wan-v2.6-i2v-flash', 'alibaba/wan-v2.6-r2v', 'alibaba/wan-v2.6-r2v-flash', 'alibaba/wan-v2.6-t2v', 'amazon/nova-2-lite', 'amazon/nova-lite', 'amazon/nova-micro', 'amazon/nova-pro', 'amazon/titan-embed-text-v2', 'anthropic/claude-3-haiku', 'anthropic/claude-3-opus', 'anthropic/claude-3.5-haiku', 'anthropic/claude-3.5-sonnet', 'anthropic/claude-3.5-sonnet-20240620', 'anthropic/claude-3.7-sonnet', 'anthropic/claude-fable-5', 'anthropic/claude-haiku-4.5', 'anthropic/claude-opus-4', 'anthropic/claude-opus-4.1', 'anthropic/claude-opus-4.5', 'anthropic/claude-opus-4.6', 'anthropic/claude-sonnet-4', 'anthropic/claude-sonnet-4.5', 'anthropic/claude-sonnet-4.6', 'arcee-ai/trinity-large-preview', 'arcee-ai/trinity-mini', 'bfl/flux-kontext-max', 'bfl/flux-kontext-pro', 'bfl/flux-pro-1.0-fill', 'bfl/flux-pro-1.1', 'bfl/flux-pro-1.1-ultra', 'bytedance/seed-1.6', 'bytedance/seed-1.8', 'bytedance/seedance-v1.0-lite-i2v', 'bytedance/seedance-v1.0-lite-t2v', 'bytedance/seedance-v1.0-pro', 'bytedance/seedance-v1.0-pro-fast', 'bytedance/seedance-v1.5-pro', 'cohere/command-a', 'cohere/embed-v4.0', 'deepseek/deepseek-r1', 'deepseek/deepseek-v3', 'deepseek/deepseek-v3.1', 'deepseek/deepseek-v3.1-terminus', 'deepseek/deepseek-v3.2', 'deepseek/deepseek-v3.2-thinking', 'google/gemini-2.0-flash', 'google/gemini-2.0-flash-lite', 'google/gemini-2.5-flash', 'google/gemini-2.5-flash-image', 'google/gemini-2.5-flash-lite', 'google/gemini-2.5-flash-lite-preview-09-2025', 'google/gemini-2.5-flash-preview-09-2025', 'google/gemini-2.5-pro', 'google/gemini-3-flash', 'google/gemini-3-pro-image', 'google/gemini-3-pro-preview', 'google/gemini-3.1-flash-image-preview', 'google/gemini-3.1-flash-lite-preview', 'google/gemini-3.1-pro-preview', 'google/gemini-embedding-001', 'google/imagen-4.0-fast-generate-001', 'google/imagen-4.0-generate-001', 'google/imagen-4.0-ultra-generate-001', 'google/text-embedding-005', 'google/text-multilingual-embedding-002', 'google/veo-3.0-fast-generate-001', 'google/veo-3.0-generate-001', 'google/veo-3.1-fast-generate-001', 'google/veo-3.1-generate-001', 'inception/mercury-2', 'inception/mercury-coder-small', 'klingai/kling-v2.5-turbo-i2v', 'klingai/kling-v2.5-turbo-t2v', 'klingai/kling-v2.6-i2v', 'klingai/kling-v2.6-motion-control', 'klingai/kling-v2.6-t2v', 'klingai/kling-v3.0-i2v', 'klingai/kling-v3.0-t2v', 'kwaipilot/kat-coder-pro-v1', 'meituan/longcat-flash-chat', 'meituan/longcat-flash-thinking', 'meta/llama-3.1-70b', 'meta/llama-3.1-8b', 'meta/llama-3.2-11b', 'meta/llama-3.2-1b', 'meta/llama-3.2-3b', 'meta/llama-3.2-90b', 'meta/llama-3.3-70b', 'meta/llama-4-maverick', 'meta/llama-4-scout', 'minimax/minimax-m2', 'minimax/minimax-m2.1', 'minimax/minimax-m2.1-lightning', 'minimax/minimax-m2.5', 'minimax/minimax-m2.5-highspeed', 'mistral/codestral', 'mistral/codestral-embed', 'mistral/devstral-2', 'mistral/devstral-small', 'mistral/devstral-small-2', 'mistral/magistral-medium', 'mistral/magistral-small', 'mistral/ministral-14b', 'mistral/ministral-3b', 'mistral/ministral-8b', 'mistral/mistral-embed', 'mistral/mistral-large-3', 'mistral/mistral-medium', 'mistral/mistral-nemo', 'mistral/mistral-small', 'mistral/mixtral-8x22b-instruct', 'mistral/pixtral-12b', 'mistral/pixtral-large', 'moonshotai/kimi-k2', 'moonshotai/kimi-k2-0905', 'moonshotai/kimi-k2-thinking', 'moonshotai/kimi-k2-thinking-turbo', 'moonshotai/kimi-k2-turbo', 'moonshotai/kimi-k2.5', 'morph/morph-v3-fast', 'morph/morph-v3-large', 'nvidia/nemotron-3-nano-30b-a3b', 'nvidia/nemotron-nano-12b-v2-vl', 'nvidia/nemotron-nano-9b-v2', 'openai/gpt-3.5-turbo', 'openai/gpt-3.5-turbo-instruct', 'openai/gpt-4-turbo', 'openai/gpt-4.1', 'openai/gpt-4.1-mini', 'openai/gpt-4.1-nano', 'openai/gpt-4o', 'openai/gpt-4o-mini', 'openai/gpt-4o-mini-search-preview', 'openai/gpt-5', 'openai/gpt-5-chat', 'openai/gpt-5-codex', 'openai/gpt-5-mini', 'openai/gpt-5-nano', 'openai/gpt-5-pro', 'openai/gpt-5.1-codex', 'openai/gpt-5.1-codex-max', 'openai/gpt-5.1-codex-mini', 'openai/gpt-5.1-instant', 'openai/gpt-5.1-thinking', 'openai/gpt-5.2', 'openai/gpt-5.2-chat', 'openai/gpt-5.2-codex', 'openai/gpt-5.2-pro', 'openai/gpt-5.3-chat', 'openai/gpt-5.3-codex', 'openai/gpt-5.4', 'openai/gpt-5.4-pro', 'openai/gpt-image-1', 'openai/gpt-image-1-mini', 'openai/gpt-image-1.5', 'openai/gpt-oss-120b', 'openai/gpt-oss-20b', 'openai/gpt-oss-safeguard-20b', 'openai/o1', 'openai/o3', 'openai/o3-deep-research', 'openai/o3-mini', 'openai/o3-pro', 'openai/o4-mini', 'openai/text-embedding-3-large', 'openai/text-embedding-3-small', 'openai/text-embedding-ada-002', 'perplexity/sonar', 'perplexity/sonar-pro', 'perplexity/sonar-reasoning', 'perplexity/sonar-reasoning-pro', 'prime-intellect/intellect-3', 'recraft/recraft-v2', 'recraft/recraft-v3', 'recraft/recraft-v4', 'recraft/recraft-v4-pro', 'stealth/sonoma-dusk-alpha', 'stealth/sonoma-sky-alpha', 'vercel/v0-1.0-md', 'vercel/v0-1.5-md', 'voyage/voyage-3-large', 'voyage/voyage-3.5', 'voyage/voyage-3.5-lite', 'voyage/voyage-4', 'voyage/voyage-4-large', 'voyage/voyage-4-lite', 'voyage/voyage-code-2', 'voyage/voyage-code-3', 'voyage/voyage-finance-2', 'voyage/voyage-law-2', 'xai/grok-2-vision', 'xai/grok-3', 'xai/grok-3-fast', 'xai/grok-3-mini', 'xai/grok-3-mini-fast', 'xai/grok-4', 'xai/grok-4-fast-non-reasoning', 'xai/grok-4-fast-reasoning', 'xai/grok-4.1-fast-non-reasoning', 'xai/grok-4.1-fast-reasoning', 'xai/grok-4.20-multi-agent-beta', 'xai/grok-4.20-non-reasoning-beta', 'xai/grok-4.20-reasoning-beta', 'xai/grok-code-fast-1', 'xai/grok-imagine-image', 'xai/grok-imagine-image-pro', 'xai/grok-imagine-video', 'xiaomi/mimo-v2-flash', 'zai/glm-4.5', 'zai/glm-4.5-air', 'zai/glm-4.5v', 'zai/glm-4.6', 'zai/glm-4.6v', 'zai/glm-4.6v-flash', 'zai/glm-4.7', 'zai/glm-4.7-flashx', 'zai/glm-5', ] @dataclass class ChatVercel(BaseChatModel): """ A wrapper around Vercel AI Gateway's API, which provides OpenAI-compatible access to various LLM models with features like rate limiting, caching, and monitoring. Examples: ```python from browser_use import Agent, ChatVercel llm = ChatVercel(model='openai/gpt-4o', api_key='your_vercel_api_key') agent = Agent(task='Your task here', llm=llm) ``` Args: model: The model identifier api_key: Your Vercel AI Gateway API key. If not provided, falls back to AI_GATEWAY_API_KEY or VERCEL_OIDC_TOKEN environment variables. base_url: The Vercel AI Gateway endpoint (defaults to https://ai-gateway.vercel.sh/v1) temperature: Sampling temperature (0-2) max_tokens: Maximum tokens to generate reasoning_models: List of reasoning model patterns (e.g., 'o1', 'gpt-oss') that need prompt-based JSON extraction. Auto-detects common reasoning models by default. timeout: Request timeout in seconds max_retries: Maximum number of retries for failed requests provider_options: Provider routing options for the gateway. Use this to control which providers are used and in what order. Example: {'gateway': {'order': ['vertex', 'anthropic']}} reasoning: Optional provider-specific reasoning configuration. Merged into providerOptions under the appropriate provider key. Example for Anthropic: {'anthropic': {'thinking': {'type': 'adaptive'}}}. Example for OpenAI: {'openai': {'reasoningEffort': 'high', 'reasoningSummary': 'detailed'}}. model_fallbacks: Optional list of fallback model IDs tried in order if the primary model fails. Passed as providerOptions.gateway.models. caching: Optional caching mode for the gateway. Currently supports 'auto', which enables provider-specific prompt caching via providerOptions.gateway.caching. """ # Model configuration model: ChatVercelModel | str # Model params temperature: float | None = None max_tokens: int | None = None top_p: float | None = None reasoning_models: list[str] | None = field( default_factory=lambda: [ 'o1', 'o3', 'o4', 'gpt-oss', 'gpt-5.2-pro', 'gpt-5.4-pro', 'deepseek-r1', '-thinking', 'perplexity/sonar-reasoning', ] ) # Client initialization parameters api_key: str | None = None base_url: str | httpx.URL = 'https://ai-gateway.vercel.sh/v1' timeout: float | httpx.Timeout | None = None max_retries: int = 5 default_headers: Mapping[str, str] | None = None default_query: Mapping[str, object] | None = None http_client: httpx.AsyncClient | None = None _strict_response_validation: bool = False provider_options: dict[str, Any] | None = None reasoning: dict[str, dict[str, Any]] | None = None model_fallbacks: list[str] | None = None caching: Literal['auto'] | None = None # Static @property def provider(self) -> str: return 'vercel' def _get_client_params(self) -> dict[str, Any]: """Prepare client parameters dictionary.""" api_key = self.api_key or os.getenv('AI_GATEWAY_API_KEY') or os.getenv('VERCEL_OIDC_TOKEN') base_params = { 'api_key': api_key, 'base_url': self.base_url, 'timeout': self.timeout, 'max_retries': self.max_retries, 'default_headers': self.default_headers, 'default_query': self.default_query, '_strict_response_validation': self._strict_response_validation, } client_params = {k: v for k, v in base_params.items() if v is not None} if self.http_client is not None: client_params['http_client'] = self.http_client return client_params def get_client(self) -> AsyncOpenAI: """ Returns an AsyncOpenAI client configured for Vercel AI Gateway. Returns: AsyncOpenAI: An instance of the AsyncOpenAI client with Vercel base URL. """ if not hasattr(self, '_client'): client_params = self._get_client_params() self._client = AsyncOpenAI(**client_params) return self._client @property def name(self) -> str: return str(self.model) def _get_usage(self, response: ChatCompletion) -> ChatInvokeUsage | None: """Extract usage information from the Vercel response.""" if response.usage is None: return None prompt_details = getattr(response.usage, 'prompt_tokens_details', None) cached_tokens = prompt_details.cached_tokens if prompt_details else None return ChatInvokeUsage( prompt_tokens=response.usage.prompt_tokens, prompt_cached_tokens=cached_tokens, prompt_cache_creation_tokens=None, prompt_image_tokens=None, completion_tokens=response.usage.completion_tokens, total_tokens=response.usage.total_tokens, ) def _fix_gemini_schema(self, schema: dict[str, Any]) -> dict[str, Any]: """ Convert a Pydantic model to a Gemini-compatible schema. This function removes unsupported properties like 'additionalProperties' and resolves $ref references that Gemini doesn't support. """ # Handle $defs and $ref resolution if '$defs' in schema: defs = schema.pop('$defs') def resolve_refs(obj: Any) -> Any: if isinstance(obj, dict): if '$ref' in obj: ref = obj.pop('$ref') ref_name = ref.split('/')[-1] if ref_name in defs: # Replace the reference with the actual definition resolved = defs[ref_name].copy() # Merge any additional properties from the reference for key, value in obj.items(): if key != '$ref': resolved[key] = value return resolve_refs(resolved) return obj else: # Recursively process all dictionary values return {k: resolve_refs(v) for k, v in obj.items()} elif isinstance(obj, list): return [resolve_refs(item) for item in obj] return obj schema = resolve_refs(schema) # Remove unsupported properties def clean_schema(obj: Any) -> Any: if isinstance(obj, dict): # Remove unsupported properties cleaned = {} for key, value in obj.items(): if key not in ['additionalProperties', 'title', 'default']: cleaned_value = clean_schema(value) # Handle empty object properties - Gemini doesn't allow empty OBJECT types if ( key == 'properties' and isinstance(cleaned_value, dict) and len(cleaned_value) == 0 and isinstance(obj.get('type', ''), str) and obj.get('type', '').upper() == 'OBJECT' ): # Convert empty object to have at least one property cleaned['properties'] = {'_placeholder': {'type': 'string'}} else: cleaned[key] = cleaned_value # If this is an object type with empty properties, add a placeholder if ( isinstance(cleaned.get('type', ''), str) and cleaned.get('type', '').upper() == 'OBJECT' and 'properties' in cleaned and isinstance(cleaned['properties'], dict) and len(cleaned['properties']) == 0 ): cleaned['properties'] = {'_placeholder': {'type': 'string'}} # Also remove 'title' from the required list if it exists if 'required' in cleaned and isinstance(cleaned.get('required'), list): cleaned['required'] = [p for p in cleaned['required'] if p != 'title'] return cleaned elif isinstance(obj, list): return [clean_schema(item) for item in obj] return obj return clean_schema(schema) @overload async def ainvoke( self, messages: list[BaseMessage], output_format: None = None, **kwargs: Any ) -> ChatInvokeCompletion[str]: ... @overload async def ainvoke(self, messages: list[BaseMessage], output_format: type[T], **kwargs: Any) -> ChatInvokeCompletion[T]: ... async def ainvoke( self, messages: list[BaseMessage], output_format: type[T] | None = None, **kwargs: Any ) -> ChatInvokeCompletion[T] | ChatInvokeCompletion[str]: """ Invoke the model with the given messages through Vercel AI Gateway. Args: messages: List of chat messages output_format: Optional Pydantic model class for structured output Returns: Either a string response or an instance of output_format """ vercel_messages = VercelMessageSerializer.serialize_messages(messages) try: model_params: dict[str, Any] = {} if self.temperature is not None: model_params['temperature'] = self.temperature if self.max_tokens is not None: model_params['max_tokens'] = self.max_tokens if self.top_p is not None: model_params['top_p'] = self.top_p extra_body: dict[str, Any] = {} provider_opts: dict[str, Any] = {} if self.provider_options: provider_opts.update(self.provider_options) if self.reasoning: # Merge provider-specific reasoning options (ex: {'anthropic': {'thinking': ...}}) for provider_name, opts in self.reasoning.items(): existing = provider_opts.get(provider_name, {}) existing.update(opts) provider_opts[provider_name] = existing gateway_opts: dict[str, Any] = provider_opts.get('gateway', {}) if self.model_fallbacks: gateway_opts['models'] = self.model_fallbacks if self.caching: gateway_opts['caching'] = self.caching if gateway_opts: provider_opts['gateway'] = gateway_opts if provider_opts: extra_body['providerOptions'] = provider_opts if extra_body: model_params['extra_body'] = extra_body if output_format is None: # Return string response response = await self.get_client().chat.completions.create( model=self.model, messages=vercel_messages, **model_params, ) usage = self._get_usage(response) return ChatInvokeCompletion( completion=response.choices[0].message.content or '', usage=usage, stop_reason=response.choices[0].finish_reason if response.choices else None, ) else: is_google_model = self.model.startswith('google/') is_anthropic_model = self.model.startswith('anthropic/') is_reasoning_model = self.reasoning_models and any( str(pattern).lower() in str(self.model).lower() for pattern in self.reasoning_models ) if is_google_model or is_anthropic_model or is_reasoning_model: modified_messages = [m.model_copy(deep=True) for m in messages] schema = SchemaOptimizer.create_gemini_optimized_schema(output_format) json_instruction = f'\n\nIMPORTANT: You must respond with ONLY a valid JSON object (no markdown, no code blocks, no explanations) that exactly matches this schema:\n{json.dumps(schema, indent=2)}' instruction_added = False if modified_messages and modified_messages[0].role == 'system': if isinstance(modified_messages[0].content, str): modified_messages[0].content += json_instruction instruction_added = True elif isinstance(modified_messages[0].content, list): modified_messages[0].content.append(ContentPartTextParam(text=json_instruction)) instruction_added = True elif modified_messages and modified_messages[-1].role == 'user': if isinstance(modified_messages[-1].content, str): modified_messages[-1].content += json_instruction instruction_added = True elif isinstance(modified_messages[-1].content, list): modified_messages[-1].content.append(ContentPartTextParam(text=json_instruction)) instruction_added = True if not instruction_added: modified_messages.insert(0, SystemMessage(content=json_instruction)) vercel_messages = VercelMessageSerializer.serialize_messages(modified_messages) response = await self.get_client().chat.completions.create( model=self.model, messages=vercel_messages, **model_params, ) content = response.choices[0].message.content if response.choices else None if not content: raise ModelProviderError( message='No response from model', status_code=500, model=self.name, ) try: text = content.strip() if text.startswith('```json') and text.endswith('```'): text = text[7:-3].strip() elif text.startswith('```') and text.endswith('```'): text = text[3:-3].strip() parsed_data = json.loads(text) parsed = output_format.model_validate(parsed_data) usage = self._get_usage(response) return ChatInvokeCompletion( completion=parsed, usage=usage, stop_reason=response.choices[0].finish_reason if response.choices else None, ) except (json.JSONDecodeError, ValueError) as e: raise ModelProviderError( message=f'Failed to parse JSON response: {str(e)}. Raw response: {content[:200]}', status_code=500, model=self.name, ) from e else: schema = SchemaOptimizer.create_optimized_json_schema(output_format) response_format_schema: JSONSchema = { 'name': 'agent_output', 'strict': True, 'schema': schema, } response = await self.get_client().chat.completions.create( model=self.model, messages=vercel_messages, response_format=ResponseFormatJSONSchema( json_schema=response_format_schema, type='json_schema', ), **model_params, ) content = response.choices[0].message.content if response.choices else None if not content: raise ModelProviderError( message='Failed to parse structured output from model response - empty or null content', status_code=500, model=self.name, ) usage = self._get_usage(response) parsed = output_format.model_validate_json(content) return ChatInvokeCompletion( completion=parsed, usage=usage, stop_reason=response.choices[0].finish_reason if response.choices else None, ) except RateLimitError as e: raise ModelRateLimitError(message=e.message, model=self.name) from e except APIConnectionError as e: raise ModelProviderError(message=str(e), model=self.name) from e except APIStatusError as e: raise ModelProviderError(message=e.message, status_code=e.status_code, model=self.name) from e except Exception as e: raise ModelProviderError(message=str(e), model=self.name) from e