From 9103e7f46339c10a0e7d66f477300d26fbf133b9 Mon Sep 17 00:00:00 2001 From: yangdx Date: Thu, 6 Feb 2025 10:42:49 +0800 Subject: [PATCH] fix: improve timing accuracy and variable scoping in OllamaAPI --- lightrag/api/ollama_api.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/lightrag/api/ollama_api.py b/lightrag/api/ollama_api.py index 132601c3..edf97993 100644 --- a/lightrag/api/ollama_api.py +++ b/lightrag/api/ollama_api.py @@ -203,15 +203,15 @@ class OllamaAPI: ) async def stream_generator(): - first_chunk_time = None - last_chunk_time = time.time_ns() - total_response = "" - try: + first_chunk_time = None + last_chunk_time = time.time_ns() + total_response = "" + # 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 + first_chunk_time = start_time last_chunk_time = time.time_ns() total_response = response @@ -284,7 +284,7 @@ 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 + first_chunk_time = start_time completion_tokens = estimate_tokens(total_response) total_time = last_chunk_time - start_time prompt_eval_time = first_chunk_time - start_time @@ -409,15 +409,15 @@ class OllamaAPI: ) async def stream_generator(): - first_chunk_time = None - last_chunk_time = time.time_ns() - total_response = "" - try: + first_chunk_time = None + last_chunk_time = time.time_ns() + total_response = "" + # 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 + first_chunk_time = start_time last_chunk_time = time.time_ns() total_response = response @@ -503,7 +503,7 @@ class OllamaAPI: return if first_chunk_time is None: - first_chunk_time = last_chunk_time + first_chunk_time = start_time completion_tokens = estimate_tokens(total_response) total_time = last_chunk_time - start_time prompt_eval_time = first_chunk_time - start_time