Skip to content

Commit 1bd1c19

Browse files
shah-sidddshah-siddd
committed
fix(open-9027): azure openai tracer fix for chunks
1 parent 4cd559e commit 1bd1c19

File tree

3 files changed

+69
-13
lines changed

3 files changed

+69
-13
lines changed

‎src/openlayer/lib/integrations/async_openai_tracer.py‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,11 @@ async def handle_async_streaming_create(
206206
num_of_completion_tokens = i + 1
207207
i += 1
208208

209+
choices = getattr(chunk, "choices", None)
210+
if not choices:
211+
yield chunk
212+
continue
213+
209214
delta = chunk.choices[0].delta
210215

211216
if delta.content:
@@ -235,7 +240,13 @@ async def handle_async_streaming_create(
235240
if collected_output_data:
236241
output_data = "".join(collected_output_data)
237242
else:
238-
collected_function_call["arguments"] = json.loads(collected_function_call["arguments"])
243+
if collected_function_call["arguments"]:
244+
try:
245+
collected_function_call["arguments"] = json.loads(
246+
collected_function_call["arguments"]
247+
)
248+
except json.JSONDecodeError:
249+
pass
239250
output_data = collected_function_call
240251

241252
trace_args = create_trace_args(
@@ -543,6 +554,12 @@ async def handle_async_streaming_parse(
543554
num_of_completion_tokens = i + 1
544555
i += 1
545556

557+
# Skip chunks with empty choices (e.g., Azure OpenAI heartbeat chunks)
558+
choices = getattr(chunk, "choices", None)
559+
if not choices:
560+
yield chunk
561+
continue
562+
546563
delta = chunk.choices[0].delta
547564

548565
if delta.content:
@@ -578,9 +595,13 @@ async def handle_async_streaming_parse(
578595
if collected_output_data:
579596
output_data = "".join(collected_output_data)
580597
else:
581-
collected_function_call["arguments"] = json.loads(
582-
collected_function_call["arguments"]
583-
)
598+
if collected_function_call["arguments"]:
599+
try:
600+
collected_function_call["arguments"] = json.loads(
601+
collected_function_call["arguments"]
602+
)
603+
except json.JSONDecodeError:
604+
pass
584605
output_data = collected_function_call
585606

586607
trace_args = create_trace_args(

‎src/openlayer/lib/integrations/groq_tracer.py‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,11 @@ def stream_chunks(
133133
if i > 0:
134134
num_of_completion_tokens = i + 1
135135

136+
choices = getattr(chunk, "choices", None)
137+
if not choices:
138+
yield chunk
139+
continue
140+
136141
delta = chunk.choices[0].delta
137142

138143
if delta.content:
@@ -161,7 +166,13 @@ def stream_chunks(
161166
if collected_output_data:
162167
output_data = "".join(collected_output_data)
163168
else:
164-
collected_function_call["arguments"] = json.loads(collected_function_call["arguments"])
169+
if collected_function_call["arguments"]:
170+
try:
171+
collected_function_call["arguments"] = json.loads(
172+
collected_function_call["arguments"]
173+
)
174+
except json.JSONDecodeError:
175+
pass
165176
output_data = collected_function_call
166177

167178
# Get usage data from the last chunk
@@ -321,14 +332,25 @@ def parse_non_streaming_output_data(
321332
output_data = output_content.strip()
322333
elif output_function_call or output_tool_calls:
323334
if output_function_call:
335+
args_str = getattr(output_function_call, "arguments", "") or ""
336+
try:
337+
arguments = json.loads(args_str) if args_str.strip() else {}
338+
except json.JSONDecodeError:
339+
arguments = args_str
324340
function_call = {
325341
"name": output_function_call.name,
326-
"arguments": json.loads(output_function_call.arguments),
342+
"arguments": arguments,
327343
}
328344
else:
345+
func = output_tool_calls[0].function
346+
args_str = getattr(func, "arguments", "") or ""
347+
try:
348+
arguments = json.loads(args_str) if args_str.strip() else {}
349+
except json.JSONDecodeError:
350+
arguments = args_str
329351
function_call = {
330-
"name": output_tool_calls[0].function.name,
331-
"arguments": json.loads(output_tool_calls[0].function.arguments),
352+
"name": func.name,
353+
"arguments": arguments,
332354
}
333355
output_data = function_call
334356
else:

‎src/openlayer/lib/integrations/openai_tracer.py‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,8 @@ def stream_chunks(
216216
num_of_completion_tokens = i + 1
217217

218218
# Skip chunks with empty choices (e.g., Azure OpenAI heartbeat chunks)
219-
if not chunk.choices:
219+
choices = getattr(chunk, "choices", None)
220+
if not choices:
220221
yield chunk
221222
continue
222223

@@ -1344,14 +1345,25 @@ def parse_non_streaming_output_data(
13441345
# Function/tool call response
13451346
if output_function_call or output_tool_calls:
13461347
if output_function_call:
1348+
args_str = getattr(output_function_call, "arguments", "") or ""
1349+
try:
1350+
arguments = json.loads(args_str) if args_str.strip() else {}
1351+
except json.JSONDecodeError:
1352+
arguments = args_str
13471353
return {
13481354
"name": output_function_call.name,
1349-
"arguments": json.loads(output_function_call.arguments),
1355+
"arguments": arguments,
13501356
}
13511357
else:
1358+
func = output_tool_calls[0].function
1359+
args_str = getattr(func, "arguments", "") or ""
1360+
try:
1361+
arguments = json.loads(args_str) if args_str.strip() else {}
1362+
except json.JSONDecodeError:
1363+
arguments = args_str
13521364
return {
1353-
"name": output_tool_calls[0].function.name,
1354-
"arguments": json.loads(output_tool_calls[0].function.arguments),
1365+
"name": func.name,
1366+
"arguments": arguments,
13551367
}
13561368

13571369
return None
@@ -1417,7 +1429,8 @@ def stream_parse_chunks(
14171429
num_of_completion_tokens = i + 1
14181430

14191431
# Skip chunks with empty choices (e.g., Azure OpenAI heartbeat chunks)
1420-
if not chunk.choices:
1432+
choices = getattr(chunk, "choices", None)
1433+
if not choices:
14211434
yield chunk
14221435
continue
14231436

0 commit comments

Comments
 (0)