pre-commit run --all-files

This commit is contained in:
yangdx
2025-01-19 08:07:26 +08:00
parent ea88981146
commit a78be2ab17

View File

@@ -33,6 +33,7 @@ from dotenv import load_dotenv
load_dotenv() load_dotenv()
def estimate_tokens(text: str) -> int: def estimate_tokens(text: str) -> int:
"""Estimate the number of tokens in text """Estimate the number of tokens in text
Chinese characters: approximately 1.5 tokens per character Chinese characters: approximately 1.5 tokens per character
@@ -587,9 +588,12 @@ def create_app(args):
# Initialize document manager # Initialize document manager
doc_manager = DocumentManager(args.input_dir) doc_manager = DocumentManager(args.input_dir)
async def openai_alike_model_complete( async def openai_alike_model_complete(
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs prompt,
system_prompt=None,
history_messages=[],
keyword_extraction=False,
**kwargs,
) -> str: ) -> str:
return await openai_complete_if_cache( return await openai_complete_if_cache(
args.llm_model, args.llm_model,
@@ -602,7 +606,11 @@ def create_app(args):
) )
async def azure_openai_model_complete( async def azure_openai_model_complete(
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs prompt,
system_prompt=None,
history_messages=[],
keyword_extraction=False,
**kwargs,
) -> str: ) -> str:
return await azure_openai_complete_if_cache( return await azure_openai_complete_if_cache(
args.llm_model, args.llm_model,
@@ -642,7 +650,7 @@ def create_app(args):
) )
# Initialize RAG # Initialize RAG
if args.llm_binding in ["lollms", "ollama"] : if args.llm_binding in ["lollms", "ollama"]:
rag = LightRAG( rag = LightRAG(
working_dir=args.working_dir, working_dir=args.working_dir,
llm_model_func=lollms_model_complete llm_model_func=lollms_model_complete
@@ -658,7 +666,7 @@ def create_app(args):
}, },
embedding_func=embedding_func, embedding_func=embedding_func,
) )
else : else:
rag = LightRAG( rag = LightRAG(
working_dir=args.working_dir, working_dir=args.working_dir,
llm_model_func=azure_openai_model_complete llm_model_func=azure_openai_model_complete