Add error message for contextexceedederror (#1866)

This commit is contained in:
Shuchang Zheng 2025-03-02 23:27:20 -05:00 committed by GitHub
parent b3b548ed38
commit d271a84eff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 0 deletions

View file

@ -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)

View file

@ -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(