Add support for custom LLM error messages (#3501)

* Add support for custom LLM error messages * Fix mypy
onyx-dot-app · Dec 18, 2024 · 23ecf65 · 23ecf65
1 parent ddc6a6d
commit 23ecf65
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 1 deletion.
diff --git a/backend/onyx/configs/app_configs.py b/backend/onyx/configs/app_configs.py
@@ -1,6 +1,7 @@
 import json
 import os
 import urllib.parse
+from typing import cast
 
 from onyx.configs.constants import AuthType
 from onyx.configs.constants import DocumentIndexType
@@ -487,6 +488,21 @@
 
 PARSE_WITH_TRAFILATURA = os.environ.get("PARSE_WITH_TRAFILATURA", "").lower() == "true"
 
+# allow for custom error messages for different errors returned by litellm
+# for example, can specify: {"Violated content safety policy": "EVIL REQUEST!!!"}
+# to make it so that if an LLM call returns an error containing "Violated content safety policy"
+# the end user will see "EVIL REQUEST!!!" instead of the default error message.
+_LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS = os.environ.get(
+    "LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS", ""
+)
+LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS: dict[str, str] | None = None
+try:
+    LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS = cast(
+        dict[str, str], json.loads(_LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS)
+    )
+except json.JSONDecodeError:
+    pass
+
 #####
 # Enterprise Edition Configs
 #####

diff --git a/backend/onyx/llm/utils.py b/backend/onyx/llm/utils.py
@@ -28,6 +28,7 @@
 from litellm.exceptions import Timeout  # type: ignore
 from litellm.exceptions import UnprocessableEntityError  # type: ignore
 
+from onyx.configs.app_configs import LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS
 from onyx.configs.constants import MessageType
 from onyx.configs.model_configs import GEN_AI_MAX_TOKENS
 from onyx.configs.model_configs import GEN_AI_MODEL_FALLBACK_MAX_TOKENS
@@ -45,10 +46,19 @@
 
 
 def litellm_exception_to_error_msg(
-    e: Exception, llm: LLM, fallback_to_error_msg: bool = False
+    e: Exception,
+    llm: LLM,
+    fallback_to_error_msg: bool = False,
+    custom_error_msg_mappings: dict[str, str]
+    | None = LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS,
 ) -> str:
     error_msg = str(e)
 
+    if custom_error_msg_mappings:
+        for error_msg_pattern, custom_error_msg in custom_error_msg_mappings.items():
+            if error_msg_pattern in error_msg:
+                return custom_error_msg
+
     if isinstance(e, BadRequestError):
         error_msg = "Bad request: The server couldn't process your request. Please check your input."
     elif isinstance(e, AuthenticationError):

diff --git a/deployment/docker_compose/docker-compose.dev.yml b/deployment/docker_compose/docker-compose.dev.yml
@@ -92,6 +92,7 @@ services:
       - LOG_POSTGRES_LATENCY=${LOG_POSTGRES_LATENCY:-}
       - LOG_POSTGRES_CONN_COUNTS=${LOG_POSTGRES_CONN_COUNTS:-}
       - CELERY_BROKER_POOL_LIMIT=${CELERY_BROKER_POOL_LIMIT:-}
+      - LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS=${LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS:-}
 
       # Analytics Configs
       - SENTRY_DSN=${SENTRY_DSN:-}

diff --git a/deployment/docker_compose/docker-compose.gpu-dev.yml b/deployment/docker_compose/docker-compose.gpu-dev.yml
@@ -84,6 +84,7 @@ services:
       # (time spent on finding the right docs + time spent fetching summaries from disk)
       - LOG_VESPA_TIMING_INFORMATION=${LOG_VESPA_TIMING_INFORMATION:-}
       - CELERY_BROKER_POOL_LIMIT=${CELERY_BROKER_POOL_LIMIT:-}
+      - LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS=${LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS:-}
 
       # Chat Configs
       - HARD_DELETE_CHATS=${HARD_DELETE_CHATS:-}