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

fix (token_tracker): prevented double token tracker

Browse source
This commit is contained in:
GGrassia 2025-11-28 15:49:04 +01:00
parent 49ce064a11
commit d971e372c9
1 changed files with 8 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
lightrag/utils.py
View file

@ -1694,10 +1694,13 @@ async def use_llm_func_with_cache(
if max_tokens is not None: if max_tokens is not None:
kwargs["max_tokens"] = max_tokens kwargs["max_tokens"] = max_tokens
# Only add token_tracker if not already in kwargs (to avoid duplicate argument error)
if token_tracker is not None and "token_tracker" not in kwargs:
kwargs["token_tracker"] = token_tracker
res: str = await use_llm_func( res: str = await use_llm_func(
safe_user_prompt, safe_user_prompt,
system_prompt=safe_system_prompt, system_prompt=safe_system_prompt,
token_tracker=token_tracker,
**kwargs, **kwargs,
) )
@ -1731,11 +1734,14 @@ async def use_llm_func_with_cache(
if max_tokens is not None: if max_tokens is not None:
kwargs["max_tokens"] = max_tokens kwargs["max_tokens"] = max_tokens
# Only add token_tracker if not already in kwargs (to avoid duplicate argument error)
if token_tracker is not None and "token_tracker" not in kwargs:
kwargs["token_tracker"] = token_tracker
try: try:
res = await use_llm_func( res = await use_llm_func(
safe_user_prompt, safe_user_prompt,
system_prompt=safe_system_prompt, system_prompt=safe_system_prompt,
token_tracker=token_tracker,
**kwargs, **kwargs,
) )
except Exception as e: except Exception as e: