From 23b96159c080d6c9b413a69548d325aa40315b95 Mon Sep 17 00:00:00 2001
From: Robert Craigie <robert@craigie.dev>
Date: Tue, 10 Sep 2024 16:46:51 +0100
Subject: [PATCH] feat(errors): include completion in LengthFinishReasonError
 (#1701)

---
 src/openai/_exceptions.py                     | 23 +++++++++++++++----
 src/openai/lib/_parsing/_completions.py       |  2 +-
 src/openai/lib/streaming/chat/_completions.py |  4 +++-
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/openai/_exceptions.py b/src/openai/_exceptions.py
index f44f90b52f..e326ed9578 100644
--- a/src/openai/_exceptions.py
+++ b/src/openai/_exceptions.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 from typing_extensions import Literal
 
 import httpx
@@ -10,6 +10,9 @@
 from ._utils import is_dict
 from ._models import construct_type
 
+if TYPE_CHECKING:
+    from .types.chat import ChatCompletion
+
 __all__ = [
     "BadRequestError",
     "AuthenticationError",
@@ -130,10 +133,20 @@ class InternalServerError(APIStatusError):
 
 
 class LengthFinishReasonError(OpenAIError):
-    def __init__(self) -> None:
-        super().__init__(
-            f"Could not parse response content as the length limit was reached",
-        )
+    completion: ChatCompletion
+    """The completion that caused this error.
+
+    Note: this will *not* be a complete `ChatCompletion` object when streaming as `usage`
+          will not be included.
+    """
+
+    def __init__(self, *, completion: ChatCompletion) -> None:
+        msg = "Could not parse response content as the length limit was reached"
+        if completion.usage:
+            msg += f" - {completion.usage}"
+
+        super().__init__(msg)
+        self.completion = completion
 
 
 class ContentFilterFinishReasonError(OpenAIError):
diff --git a/src/openai/lib/_parsing/_completions.py b/src/openai/lib/_parsing/_completions.py
index 2ef1bf3553..f1fa9f2b55 100644
--- a/src/openai/lib/_parsing/_completions.py
+++ b/src/openai/lib/_parsing/_completions.py
@@ -69,7 +69,7 @@ def parse_chat_completion(
     choices: list[ParsedChoice[ResponseFormatT]] = []
     for choice in chat_completion.choices:
         if choice.finish_reason == "length":
-            raise LengthFinishReasonError()
+            raise LengthFinishReasonError(completion=chat_completion)
 
         if choice.finish_reason == "content_filter":
             raise ContentFilterFinishReasonError()
diff --git a/src/openai/lib/streaming/chat/_completions.py b/src/openai/lib/streaming/chat/_completions.py
index 342a5e2b95..a4b0f856f7 100644
--- a/src/openai/lib/streaming/chat/_completions.py
+++ b/src/openai/lib/streaming/chat/_completions.py
@@ -394,7 +394,9 @@ def _accumulate_chunk(self, chunk: ChatCompletionChunk) -> ParsedChatCompletionS
 
                 if has_parseable_input(response_format=self._response_format, input_tools=self._input_tools):
                     if choice.finish_reason == "length":
-                        raise LengthFinishReasonError()
+                        # at the time of writing, `.usage` will always be `None` but
+                        # we include it here in case that is changed in the future
+                        raise LengthFinishReasonError(completion=completion_snapshot)
 
                     if choice.finish_reason == "content_filter":
                         raise ContentFilterFinishReasonError()