gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: improve timing accuracy and variable scoping in OllamaAPI

Browse Source
This commit is contained in:
yangdx
2025-02-06 10:42:49 +08:00
parent e124ad7f9c
commit 9103e7f463
1 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
lightrag/api/ollama_api.py
View File

@@ -203,15 +203,15 @@ class OllamaAPI:
) )
async def stream_generator(): async def stream_generator():
try:
first_chunk_time = None first_chunk_time = None
last_chunk_time = time.time_ns() last_chunk_time = time.time_ns()
total_response = "" total_response = ""
try:
# Ensure response is an async generator # Ensure response is an async generator
if isinstance(response, str): if isinstance(response, str):
# If it's a string, send in two parts # 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() last_chunk_time = time.time_ns()
total_response = response total_response = response
@@ -284,7 +284,7 @@ class OllamaAPI:
yield f"{json.dumps(final_data, ensure_ascii=False)}\n" yield f"{json.dumps(final_data, ensure_ascii=False)}\n"
return return
if first_chunk_time is None: if first_chunk_time is None:
first_chunk_time = last_chunk_time first_chunk_time = start_time
completion_tokens = estimate_tokens(total_response) completion_tokens = estimate_tokens(total_response)
total_time = last_chunk_time - start_time total_time = last_chunk_time - start_time
prompt_eval_time = first_chunk_time - start_time prompt_eval_time = first_chunk_time - start_time
@@ -409,15 +409,15 @@ class OllamaAPI:
) )
async def stream_generator(): async def stream_generator():
try:
first_chunk_time = None first_chunk_time = None
last_chunk_time = time.time_ns() last_chunk_time = time.time_ns()
total_response = "" total_response = ""
try:
# Ensure response is an async generator # Ensure response is an async generator
if isinstance(response, str): if isinstance(response, str):
# If it's a string, send in two parts # 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() last_chunk_time = time.time_ns()
total_response = response total_response = response
@@ -503,7 +503,7 @@ class OllamaAPI:
return return
if first_chunk_time is None: if first_chunk_time is None:
first_chunk_time = last_chunk_time first_chunk_time = start_time
completion_tokens = estimate_tokens(total_response) completion_tokens = estimate_tokens(total_response)
total_time = last_chunk_time - start_time total_time = last_chunk_time - start_time
prompt_eval_time = first_chunk_time - start_time prompt_eval_time = first_chunk_time - start_time