Catch LLM Generation Failure (#1712)

backend/danswer/llm/chat_llm.py

@@ -5,6 +5,7 @@ from typing import Any
 from typing import cast
 
 import litellm  # type: ignore
+from httpx import RemoteProtocolError
 from langchain.schema.language_model import LanguageModelInput
 from langchain_core.messages import AIMessage
 from langchain_core.messages import AIMessageChunk
@@ -338,6 +339,7 @@ class DefaultMultiLLM(LLM):
 
         output = None
         response = self._completion(prompt, tools, tool_choice, True)
+        try:
             for part in response:
                 if len(part["choices"]) == 0:
                     continue
@@ -350,6 +352,11 @@ class DefaultMultiLLM(LLM):
 
                 yield message_chunk
 
+        except RemoteProtocolError:
+            raise RuntimeError(
+                "The AI model failed partway through generation, please try again."
+            )
+
         if LOG_ALL_MODEL_INTERACTIONS and output:
             content = output.content or ""
             if isinstance(output, AIMessage):