@@ -156,29 +156,31 @@ def calculate_cost_and_tokens(result):
156156 # Extract 'extraction_result' from the result
157157 extraction_result = result .get ("extraction_result" , [])
158158
159- if extraction_result :
160- extraction_data = extraction_result [0 ].get ("result" , "" )
161-
162- # If extraction_data is a string, attempt to parse it as JSON
163- if isinstance (extraction_data , str ):
164- try :
165- extraction_data = json .loads (extraction_data ) if extraction_data else {}
166- except json .JSONDecodeError :
167- logger .warning ("Failed to decode JSON for extraction data; defaulting to empty dictionary." )
168- extraction_data = {}
169-
170-
171- metadata = extraction_data .get ("metadata" , {})
172- embedding_llm = metadata .get ("embedding" , [])
173- extraction_llm = metadata .get ("extraction_llm" , [])
174-
175- # Calculate total cost
176- total_embedding_cost += sum (float (item .get ("cost_in_dollars" , "0" )) for item in embedding_llm )
177- total_llm_cost += sum (float (item .get ("cost_in_dollars" , "0" )) for item in extraction_llm )
178-
179- # Calculate total tokens
180- total_embedding_tokens += sum (item .get ("embedding_tokens" , 0 ) for item in embedding_llm )
181- total_llm_tokens += sum (item .get ("total_tokens" , 0 ) for item in extraction_llm )
159+ if not extraction_result :
160+ return total_embedding_cost , total_llm_cost , total_embedding_tokens , total_llm_tokens
161+
162+ extraction_data = extraction_result [0 ].get ("result" , "" )
163+
164+ # If extraction_data is a string, attempt to parse it as JSON
165+ if isinstance (extraction_data , str ):
166+ try :
167+ extraction_data = json .loads (extraction_data ) if extraction_data else {}
168+ except json .JSONDecodeError :
169+ logger .warning ("Failed to decode JSON for extraction data; defaulting to empty dictionary." )
170+ extraction_data = {}
171+
172+
173+ metadata = extraction_data .get ("metadata" , {})
174+ embedding_llm = metadata .get ("embedding" , [])
175+ extraction_llm = metadata .get ("extraction_llm" , [])
176+
177+ # Calculate total cost
178+ total_embedding_cost += sum (float (item .get ("cost_in_dollars" , "0" )) for item in embedding_llm )
179+ total_llm_cost += sum (float (item .get ("cost_in_dollars" , "0" )) for item in extraction_llm )
180+
181+ # Calculate total tokens
182+ total_embedding_tokens += sum (item .get ("embedding_tokens" , 0 ) for item in embedding_llm )
183+ total_llm_tokens += sum (item .get ("total_tokens" , 0 ) for item in extraction_llm )
182184
183185 return total_embedding_cost , total_llm_cost , total_embedding_tokens , total_llm_tokens
184186
0 commit comments