Fix linting
Browse files- 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(
|
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(
|
|
|
|
|
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])
|