From d271a84effff556b28fc8c68e8d28342c7622fc1 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Sun, 2 Mar 2025 23:27:20 -0500 Subject: [PATCH] Add error message for contextexceedederror (#1866) --- skyvern/exceptions.py | 6 ++++++ skyvern/forge/sdk/api/llm/api_handler_factory.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py index a009cd20..0cf4a90b 100644 --- a/skyvern/exceptions.py +++ b/skyvern/exceptions.py @@ -655,3 +655,9 @@ class NoTOTPVerificationCodeFound(SkyvernHTTPException): if totp_identifier: msg += f" totp_identifier={totp_identifier}" super().__init__(msg) + + +class SkyvernContextWindowExceededError(SkyvernException): + def __init__(self) -> None: + message = "Context window exceeded. Please contact support@skyvern.com for help." + super().__init__(message) diff --git a/skyvern/forge/sdk/api/llm/api_handler_factory.py b/skyvern/forge/sdk/api/llm/api_handler_factory.py index 435a753a..2617b3a7 100644 --- a/skyvern/forge/sdk/api/llm/api_handler_factory.py +++ b/skyvern/forge/sdk/api/llm/api_handler_factory.py @@ -9,6 +9,7 @@ import structlog from jinja2 import Template from skyvern.config import settings +from skyvern.exceptions import SkyvernContextWindowExceededError from skyvern.forge import app from skyvern.forge.sdk.api.llm.config_registry import LLMConfigRegistry from skyvern.forge.sdk.api.llm.exceptions import ( @@ -125,6 +126,13 @@ class LLMAPIHandlerFactory: response = await router.acompletion(model=main_model_group, messages=messages, **parameters) except litellm.exceptions.APIError as e: raise LLMProviderErrorRetryableTask(llm_key) from e + except litellm.exceptions.ContextWindowExceededError as e: + LOG.exception( + "Context window exceeded", + llm_key=llm_key, + model=main_model_group, + ) + raise SkyvernContextWindowExceededError() from e except ValueError as e: LOG.exception( "LLM token limit exceeded", @@ -300,6 +308,13 @@ class LLMAPIHandlerFactory: ) except litellm.exceptions.APIError as e: raise LLMProviderErrorRetryableTask(llm_key) from e + except litellm.exceptions.ContextWindowExceededError as e: + LOG.exception( + "Context window exceeded", + llm_key=llm_key, + model=llm_config.model_name, + ) + raise SkyvernContextWindowExceededError() from e except CancelledError: t_llm_cancelled = time.perf_counter() LOG.error(