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

Fix timing calculation logic in OllamaAPI stream generators

Browse source 
• Initialize first_chunk_time as None
• Set timing only when first chunk arrives
This commit is contained in:
yangdx 2025-02-06 04:53:05 +08:00
parent e49d3665aa
commit e124ad7f9c
1 changed files with 10 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
lightrag/api/ollama_api.py
View file

@ -203,14 +203,15 @@ class OllamaAPI:
)
async def stream_generator():
first_chunk_time = time.time_ns()
last_chunk_time = first_chunk_time
first_chunk_time = None
last_chunk_time = time.time_ns()
total_response = ""
try:
# Ensure response is an async generator
if isinstance(response, str):
# If it's a string, send in two parts
first_chunk_time = last_chunk_time
last_chunk_time = time.time_ns()
total_response = response
@ -282,7 +283,8 @@ class OllamaAPI:
}
yield f"{json.dumps(final_data, ensure_ascii=False)}\n"
return
if first_chunk_time is None:
first_chunk_time = last_chunk_time
completion_tokens = estimate_tokens(total_response)
total_time = last_chunk_time - start_time
prompt_eval_time = first_chunk_time - start_time
@ -407,14 +409,15 @@ class OllamaAPI:
)
async def stream_generator():
first_chunk_time = time.time_ns()
last_chunk_time = first_chunk_time
first_chunk_time = None
last_chunk_time = time.time_ns()
total_response = ""
try:
# Ensure response is an async generator
if isinstance(response, str):
# If it's a string, send in two parts
first_chunk_time = last_chunk_time
last_chunk_time = time.time_ns()
total_response = response
@ -499,6 +502,8 @@ class OllamaAPI:
yield f"{json.dumps(final_data, ensure_ascii=False)}\n"
return
if first_chunk_time is None:
first_chunk_time = last_chunk_time
completion_tokens = estimate_tokens(total_response)
total_time = last_chunk_time - start_time
prompt_eval_time = first_chunk_time - start_time