|
28 | 28 | SERVICE_PROVIDERS, |
29 | 29 | ) |
30 | 30 | from langtrace_python_sdk.constants.instrumentation.openai import APIS |
31 | | -from langtrace_python_sdk.utils.llm import calculate_prompt_tokens, estimate_tokens |
| 31 | +from langtrace_python_sdk.utils.llm import ( |
| 32 | + calculate_prompt_tokens, |
| 33 | + estimate_tokens, |
| 34 | + get_tool_calls, |
| 35 | +) |
32 | 36 | from openai._types import NOT_GIVEN |
33 | 37 |
|
34 | 38 |
|
@@ -430,9 +434,10 @@ def traced_method(wrapped, instance, args, kwargs): |
430 | 434 | # handle tool calls in the kwargs |
431 | 435 | llm_prompts = [] |
432 | 436 | for item in kwargs.get("messages", []): |
433 | | - if hasattr(item, "tool_calls") and item.tool_calls is not None: |
| 437 | + tools = get_tool_calls(item) |
| 438 | + if tools is not None: |
434 | 439 | tool_calls = [] |
435 | | - for tool_call in item.tool_calls: |
| 440 | + for tool_call in tools: |
436 | 441 | tool_call_dict = { |
437 | 442 | "id": tool_call.id if hasattr(tool_call, "id") else "", |
438 | 443 | "type": tool_call.type if hasattr(tool_call, "type") else "", |
@@ -611,9 +616,10 @@ async def traced_method(wrapped, instance, args, kwargs): |
611 | 616 | # handle tool calls in the kwargs |
612 | 617 | llm_prompts = [] |
613 | 618 | for item in kwargs.get("messages", []): |
614 | | - if hasattr(item, "tool_calls") and item.tool_calls is not None: |
| 619 | + tools = get_tool_calls(item) |
| 620 | + if tools is not None: |
615 | 621 | tool_calls = [] |
616 | | - for tool_call in item.tool_calls: |
| 622 | + for tool_call in tools: |
617 | 623 | tool_call_dict = { |
618 | 624 | "id": tool_call.id if hasattr(tool_call, "id") else "", |
619 | 625 | "type": tool_call.type if hasattr(tool_call, "type") else "", |
|
0 commit comments