Catch LLM Generation Failure (#1712)

This commit is contained in:
Yuhong Sun 2024-06-26 10:44:22 -07:00 committed by GitHub
parent 0d814939ee
commit 20c4cdbdda
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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,17 +339,23 @@ class DefaultMultiLLM(LLM):
output = None
response = self._completion(prompt, tools, tool_choice, True)
for part in response:
if len(part["choices"]) == 0:
continue
delta = part["choices"][0]["delta"]
message_chunk = _convert_delta_to_message_chunk(delta, output)
if output is None:
output = message_chunk
else:
output += message_chunk
try:
for part in response:
if len(part["choices"]) == 0:
continue
delta = part["choices"][0]["delta"]
message_chunk = _convert_delta_to_message_chunk(delta, output)
if output is None:
output = message_chunk
else:
output += message_chunk
yield message_chunk
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 ""