fix: check chunk has choises before printing

PostHog · Jan 11, 2025 · 43f39cf · 43f39cf
1 parent 4c41388
commit 43f39cf
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 8 deletions.
diff --git a/llm_observability_examples.py b/llm_observability_examples.py
@@ -29,8 +29,8 @@ def main_sync():
 
     try:
         basic_openai_call(distinct_id, trace_id, properties)
-        # streaming_openai_call(distinct_id, trace_id, properties)
-        # non_instrumented_openai_call()
+        streaming_openai_call(distinct_id, trace_id, properties)
+        non_instrumented_openai_call()
     except Exception as e:
         print("Error during OpenAI call:", str(e))
 
@@ -106,7 +106,8 @@ def streaming_openai_call(distinct_id, trace_id, properties):
     )
 
     for chunk in response:
-        print(chunk.choices[0].delta.content or "", end="")
+        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+            print(chunk.choices[0].delta.content or "", end="")
 
     return response
 
@@ -127,7 +128,8 @@ async def streaming_async_openai_call(distinct_id, trace_id, properties):
     )
 
     async for chunk in response:
-        print(chunk.choices[0].delta.content or "", end="")
+        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+            print(chunk.choices[0].delta.content or "", end="")
 
     return response
 

diff --git a/posthog/ai/providers/openai/openai.py b/posthog/ai/providers/openai/openai.py
@@ -100,9 +100,14 @@ def generator():
                                 "total_tokens",
                             ]
                         }
-                    if chunk.choices[0].delta.content:
-                        accumulated_content.append(chunk.choices[0].delta.content)
+
+                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            accumulated_content.append(content)
+
                     yield chunk
+
             finally:
                 end_time = time.time()
                 latency = end_time - start_time

diff --git a/posthog/ai/providers/openai/openai_async.py b/posthog/ai/providers/openai/openai_async.py
@@ -100,9 +100,13 @@ async def async_generator():
                                 "total_tokens",
                             ]
                         }
-                    if chunk.choices[0].delta.content:
-                        accumulated_content.append(chunk.choices[0].delta.content)
+                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            accumulated_content.append(content)
+
                     yield chunk
+
             finally:
                 end_time = time.time()
                 latency = end_time - start_time