gzdaniel commited on
Commit
d3eb0c1
·
1 Parent(s): 442f485

Fix linting

Browse files
Files changed (1) hide show
  1. lightrag/llm/openai.py +12 -8
lightrag/llm/openai.py CHANGED
@@ -211,15 +211,17 @@ async def openai_complete_if_cache(
211
  # Track if we've started iterating
212
  iteration_started = False
213
  final_chunk_usage = None
214
-
215
  try:
216
  iteration_started = True
217
  async for chunk in response:
218
  # Check if this chunk has usage information (final chunk)
219
  if hasattr(chunk, "usage") and chunk.usage:
220
  final_chunk_usage = chunk.usage
221
- logger.debug(f"Received usage info in streaming chunk: {chunk.usage}")
222
-
 
 
223
  # Check if choices exists and is not empty
224
  if not hasattr(chunk, "choices") or not chunk.choices:
225
  logger.warning(f"Received chunk without choices: {chunk}")
@@ -231,21 +233,23 @@ async def openai_complete_if_cache(
231
  ):
232
  # This might be the final chunk, continue to check for usage
233
  continue
234
-
235
  content = chunk.choices[0].delta.content
236
  if content is None:
237
  continue
238
  if r"\u" in content:
239
  content = safe_unicode_decode(content.encode("utf-8"))
240
-
241
  yield content
242
-
243
  # After streaming is complete, track token usage
244
  if token_tracker and final_chunk_usage:
245
  # Use actual usage from the API
246
  token_counts = {
247
  "prompt_tokens": getattr(final_chunk_usage, "prompt_tokens", 0),
248
- "completion_tokens": getattr(final_chunk_usage, "completion_tokens", 0),
 
 
249
  "total_tokens": getattr(final_chunk_usage, "total_tokens", 0),
250
  }
251
  token_tracker.add_usage(token_counts)
@@ -471,4 +475,4 @@ async def openai_embed(
471
  response = await openai_async_client.embeddings.create(
472
  model=model, input=texts, encoding_format="float"
473
  )
474
- return np.array([dp.embedding for dp in response.data])
 
211
  # Track if we've started iterating
212
  iteration_started = False
213
  final_chunk_usage = None
214
+
215
  try:
216
  iteration_started = True
217
  async for chunk in response:
218
  # Check if this chunk has usage information (final chunk)
219
  if hasattr(chunk, "usage") and chunk.usage:
220
  final_chunk_usage = chunk.usage
221
+ logger.debug(
222
+ f"Received usage info in streaming chunk: {chunk.usage}"
223
+ )
224
+
225
  # Check if choices exists and is not empty
226
  if not hasattr(chunk, "choices") or not chunk.choices:
227
  logger.warning(f"Received chunk without choices: {chunk}")
 
233
  ):
234
  # This might be the final chunk, continue to check for usage
235
  continue
236
+
237
  content = chunk.choices[0].delta.content
238
  if content is None:
239
  continue
240
  if r"\u" in content:
241
  content = safe_unicode_decode(content.encode("utf-8"))
242
+
243
  yield content
244
+
245
  # After streaming is complete, track token usage
246
  if token_tracker and final_chunk_usage:
247
  # Use actual usage from the API
248
  token_counts = {
249
  "prompt_tokens": getattr(final_chunk_usage, "prompt_tokens", 0),
250
+ "completion_tokens": getattr(
251
+ final_chunk_usage, "completion_tokens", 0
252
+ ),
253
  "total_tokens": getattr(final_chunk_usage, "total_tokens", 0),
254
  }
255
  token_tracker.add_usage(token_counts)
 
475
  response = await openai_async_client.embeddings.create(
476
  model=model, input=texts, encoding_format="float"
477
  )
478
+ return np.array([dp.embedding for dp in response.data])