catch connection drops and reraise error in streaming

open-telemetry · Sep 6, 2024 · 885b7fd · 885b7fd
1 parent 8495a24
commit 885b7fd
Showing 1 changed file with 16 additions and 1 deletion.
diff --git a/...on/opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/patch.py b/...on/opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/patch.py
@@ -214,7 +214,15 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.cleanup()
+        try:
+            if exc_type is not None:
+                self.span.set_status(Status(StatusCode.ERROR, str(exc_val)))
+                self.span.set_attribute(
+                    ErrorAttributes.ERROR_TYPE, exc_type.__name__
+                )
+        finally:
+            self.cleanup()
+        return False  # Propagate the exception
 
     def __iter__(self):
         return self
@@ -227,6 +235,13 @@ def __next__(self):
         except StopIteration:
             self.cleanup()
             raise
+        except Exception as error:
+            self.span.set_status(Status(StatusCode.ERROR, str(error)))
+            self.span.set_attribute(
+                ErrorAttributes.ERROR_TYPE, type(error).__qualname__
+            )
+            self.cleanup()
+            raise
 
     def set_response_model(self, chunk):
         if self.response_model: