Skip to content

Commit 65e536e

Browse files
rishsrivClaude
and
Claude
committed
fix: handle missing prompt_tokens_details in OpenAI-compatible APIs
Add checks for the existence of prompt_tokens_details in response objects before attempting to access it. This fixes an issue when using third-party OpenAI-compatible APIs that don't include cache-related information. Fixes #82. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
1 parent a93cf55 commit 65e536e

File tree

1 file changed

+18
-11
lines changed

1 file changed

+18
-11
lines changed

defog/llm/utils.py

+18-11
Original file line numberDiff line numberDiff line change
@@ -753,13 +753,15 @@ async def _process_openai_response(
753753
if tools and len(tools) > 0:
754754
consecutive_exceptions = 0
755755
while True:
756-
total_input_tokens += (
757-
response.usage.prompt_tokens
758-
- response.usage.prompt_tokens_details.cached_tokens
759-
)
760-
total_cached_input_tokens += (
761-
response.usage.prompt_tokens_details.cached_tokens
762-
)
756+
# Check if prompt_tokens_details exists in the response
757+
if hasattr(response.usage, 'prompt_tokens_details') and response.usage.prompt_tokens_details is not None:
758+
cached_tokens = response.usage.prompt_tokens_details.cached_tokens
759+
total_input_tokens += response.usage.prompt_tokens - cached_tokens
760+
total_cached_input_tokens += cached_tokens
761+
else:
762+
# If prompt_tokens_details doesn't exist, assume all tokens are uncached
763+
total_input_tokens += response.usage.prompt_tokens
764+
763765
total_output_tokens += response.usage.completion_tokens
764766
message = response.choices[0].message
765767
if message.tool_calls:
@@ -872,10 +874,15 @@ async def _process_openai_response(
872874
content = response.choices[0].message.content
873875

874876
usage = response.usage
875-
total_cached_input_tokens += usage.prompt_tokens_details.cached_tokens
876-
total_input_tokens += (
877-
usage.prompt_tokens - usage.prompt_tokens_details.cached_tokens
878-
)
877+
# Check if prompt_tokens_details exists in the response
878+
if hasattr(usage, 'prompt_tokens_details') and usage.prompt_tokens_details is not None:
879+
cached_tokens = usage.prompt_tokens_details.cached_tokens
880+
total_cached_input_tokens += cached_tokens
881+
total_input_tokens += usage.prompt_tokens - cached_tokens
882+
else:
883+
# If prompt_tokens_details doesn't exist, assume all tokens are uncached
884+
total_input_tokens += usage.prompt_tokens
885+
879886
total_output_tokens += usage.completion_tokens
880887
return (
881888
content,

0 commit comments

Comments
 (0)