gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix linting

Browse source
This commit is contained in:
yangdx 2025-07-08 18:17:21 +08:00
parent 8cbba6e9db
commit 2a0cff3ed6
1 changed files with 12 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
lightrag/llm/openai.py
View file

@ -211,15 +211,17 @@ async def openai_complete_if_cache(
# Track if we've started iterating # Track if we've started iterating
iteration_started = False iteration_started = False
final_chunk_usage = None final_chunk_usage = None
try: try:
iteration_started = True iteration_started = True
async for chunk in response: async for chunk in response:
# Check if this chunk has usage information (final chunk) # Check if this chunk has usage information (final chunk)
if hasattr(chunk, "usage") and chunk.usage: if hasattr(chunk, "usage") and chunk.usage:
final_chunk_usage = chunk.usage final_chunk_usage = chunk.usage
logger.debug(f"Received usage info in streaming chunk: {chunk.usage}") logger.debug(
f"Received usage info in streaming chunk: {chunk.usage}"
)
# Check if choices exists and is not empty # Check if choices exists and is not empty
if not hasattr(chunk, "choices") or not chunk.choices: if not hasattr(chunk, "choices") or not chunk.choices:
logger.warning(f"Received chunk without choices: {chunk}") logger.warning(f"Received chunk without choices: {chunk}")
@ -231,21 +233,23 @@ async def openai_complete_if_cache(
): ):
# This might be the final chunk, continue to check for usage # This might be the final chunk, continue to check for usage
continue continue
content = chunk.choices[0].delta.content content = chunk.choices[0].delta.content
if content is None: if content is None:
continue continue
if r"\u" in content: if r"\u" in content:
content = safe_unicode_decode(content.encode("utf-8")) content = safe_unicode_decode(content.encode("utf-8"))
yield content yield content
# After streaming is complete, track token usage # After streaming is complete, track token usage
if token_tracker and final_chunk_usage: if token_tracker and final_chunk_usage:
# Use actual usage from the API # Use actual usage from the API
token_counts = { token_counts = {
"prompt_tokens": getattr(final_chunk_usage, "prompt_tokens", 0), "prompt_tokens": getattr(final_chunk_usage, "prompt_tokens", 0),
"completion_tokens": getattr(final_chunk_usage, "completion_tokens", 0), "completion_tokens": getattr(
final_chunk_usage, "completion_tokens", 0
),
"total_tokens": getattr(final_chunk_usage, "total_tokens", 0), "total_tokens": getattr(final_chunk_usage, "total_tokens", 0),
} }
token_tracker.add_usage(token_counts) token_tracker.add_usage(token_counts)
@ -471,4 +475,4 @@ async def openai_embed(
response = await openai_async_client.embeddings.create( response = await openai_async_client.embeddings.create(
model=model, input=texts, encoding_format="float" model=model, input=texts, encoding_format="float"
) )
return np.array([dp.embedding for dp in response.data]) return np.array([dp.embedding for dp in response.data])