From 70425b0357cdca8767b8546c2d3d9eab7284f693 Mon Sep 17 00:00:00 2001 From: Saifeddine ALOUI Date: Mon, 20 Jan 2025 00:26:28 +0100 Subject: [PATCH] fixed linting --- extra/OpenWebuiTool/openwebui_tool.py | 6 +---- lightrag/api/lightrag_server.py | 28 +++++++++----------- lightrag/llm.py | 38 ++++++++++++++++++--------- 3 files changed, 39 insertions(+), 33 deletions(-) diff --git a/extra/OpenWebuiTool/openwebui_tool.py b/extra/OpenWebuiTool/openwebui_tool.py index 260b779c..1dabc68d 100644 --- a/extra/OpenWebuiTool/openwebui_tool.py +++ b/extra/OpenWebuiTool/openwebui_tool.py @@ -23,12 +23,8 @@ Requirements: - Lightrag """ -from pathlib import Path -from typing import Optional, List, Dict, Union, Any -from datetime import datetime - # Tool version __version__ = "1.0.0" __author__ = "ParisNeo" __author_email__ = "parisneoai@gmail.com" -__description__ = "Lightrag integration for OpenWebui" \ No newline at end of file +__description__ = "Lightrag integration for OpenWebui" diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py index e4cbac57..21fc68ab 100644 --- a/lightrag/api/lightrag_server.py +++ b/lightrag/api/lightrag_server.py @@ -297,15 +297,13 @@ def parse_args() -> argparse.Namespace: default=default_llm_host, help=f"llm server host URL (default: from env or {default_llm_host})", ) - - default_llm_api_key = get_env_value( - "LLM_BINDING_API_KEY", None - ) - + + default_llm_api_key = get_env_value("LLM_BINDING_API_KEY", None) + parser.add_argument( "--llm-binding-api-key", default=default_llm_api_key, - help=f"llm server API key (default: from env or empty string)", + help="llm server API key (default: from env or empty string)", ) parser.add_argument( @@ -323,14 +321,12 @@ def parse_args() -> argparse.Namespace: default=default_embedding_host, help=f"embedding server host URL (default: from env or {default_embedding_host})", ) - - default_embedding_api_key = get_env_value( - "EMBEDDING_BINDING_API_KEY", "" - ) + + default_embedding_api_key = get_env_value("EMBEDDING_BINDING_API_KEY", "") parser.add_argument( "--embedding-binding-api-key", default=default_embedding_api_key, - help=f"embedding server API key (default: from env or empty string)", + help="embedding server API key (default: from env or empty string)", ) parser.add_argument( @@ -649,26 +645,26 @@ def create_app(args): texts, embed_model=args.embedding_model, host=args.embedding_binding_host, - api_key = args.embedding_binding_api_key + api_key=args.embedding_binding_api_key, ) if args.embedding_binding == "lollms" else ollama_embed( texts, embed_model=args.embedding_model, host=args.embedding_binding_host, - api_key = args.embedding_binding_api_key + api_key=args.embedding_binding_api_key, ) if args.embedding_binding == "ollama" else azure_openai_embedding( texts, model=args.embedding_model, # no host is used for openai, - api_key = args.embedding_binding_api_key + api_key=args.embedding_binding_api_key, ) if args.embedding_binding == "azure_openai" else openai_embedding( texts, model=args.embedding_model, # no host is used for openai, - api_key = args.embedding_binding_api_key + api_key=args.embedding_binding_api_key, ), ) @@ -686,7 +682,7 @@ def create_app(args): "host": args.llm_binding_host, "timeout": args.timeout, "options": {"num_ctx": args.max_tokens}, - "api_key": args.llm_binding_api_key + "api_key": args.llm_binding_api_key, }, embedding_func=embedding_func, ) diff --git a/lightrag/llm.py b/lightrag/llm.py index 02fe3961..c3e395c9 100644 --- a/lightrag/llm.py +++ b/lightrag/llm.py @@ -349,8 +349,8 @@ async def ollama_model_if_cache( host = kwargs.pop("host", None) timeout = kwargs.pop("timeout", None) kwargs.pop("hashing_kv", None) - api_key = kwargs.pop("api_key", None) - headers={'Authorization': f'Bearer {api_key}'} if api_key else None + api_key = kwargs.pop("api_key", None) + headers = {"Authorization": f"Bearer {api_key}"} if api_key else None ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers) messages = [] if system_prompt: @@ -382,8 +382,8 @@ async def lollms_model_if_cache( """Client implementation for lollms generation.""" stream = True if kwargs.get("stream") else False - api_key = kwargs.pop("api_key", None) - headers={'Authorization': f'Bearer {api_key}'} if api_key else None + api_key = kwargs.pop("api_key", None) + headers = {"Authorization": f"Bearer {api_key}"} if api_key else None # Extract lollms specific parameters request_data = { @@ -412,7 +412,7 @@ async def lollms_model_if_cache( request_data["prompt"] = full_prompt timeout = aiohttp.ClientTimeout(total=kwargs.get("timeout", None)) - async with aiohttp.ClientSession(timeout=timeout,headers=headers) as session: + async with aiohttp.ClientSession(timeout=timeout, headers=headers) as session: if stream: async def inner(): @@ -626,7 +626,12 @@ async def nvidia_openai_complete( async def azure_openai_complete( - model: str = "gpt-4o-mini", prompt="", system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs + model: str = "gpt-4o-mini", + prompt="", + system_prompt=None, + history_messages=[], + keyword_extraction=False, + **kwargs, ) -> str: keyword_extraction = kwargs.pop("keyword_extraction", None) result = await azure_openai_complete_if_cache( @@ -1152,9 +1157,13 @@ async def ollama_embedding(texts: list[str], embed_model, **kwargs) -> np.ndarra async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray: - api_key = kwargs.pop("api_key",None) - headers = {"Authorization": api_key, "Content-Type": "application/json"} if api_key else None - kwargs["headers"]=headers + api_key = kwargs.pop("api_key", None) + headers = ( + {"Authorization": api_key, "Content-Type": "application/json"} + if api_key + else None + ) + kwargs["headers"] = headers ollama_client = ollama.Client(**kwargs) data = ollama_client.embed(model=embed_model, input=texts) return data["embeddings"] @@ -1175,15 +1184,20 @@ async def lollms_embed( Returns: np.ndarray: Array of embeddings """ - api_key = kwargs.pop("api_key",None) - headers = {"Authorization": api_key, "Content-Type": "application/json"} if api_key else None + api_key = kwargs.pop("api_key", None) + headers = ( + {"Authorization": api_key, "Content-Type": "application/json"} + if api_key + else None + ) async with aiohttp.ClientSession(headers=headers) as session: embeddings = [] for text in texts: request_data = {"text": text} async with session.post( - f"{base_url}/lollms_embed", json=request_data, + f"{base_url}/lollms_embed", + json=request_data, ) as response: result = await response.json() embeddings.append(result["vector"])