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

Fix linting

Browse source
This commit is contained in:
yangdx 2025-02-06 23:12:35 +08:00
parent c838229f05
commit 8fdba5d4db
2 changed files with 11 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
lightrag/llm/ollama.py
View file

@ -94,7 +94,7 @@ async def ollama_model_if_cache(
api_key = kwargs.pop("api_key", None)
headers = {
"Content-Type": "application/json",
"User-Agent": f"LightRAG/{__api_version__}"
"User-Agent": f"LightRAG/{__api_version__}",
}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
@ -151,7 +151,7 @@ async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
api_key = kwargs.pop("api_key", None)
headers = {
"Content-Type": "application/json",
"User-Agent": f"LightRAG/{__api_version__}"
"User-Agent": f"LightRAG/{__api_version__}",
}
if api_key:
headers["Authorization"] = api_key

13
lightrag/llm/openai.py
View file

@ -78,10 +78,13 @@ from lightrag.api import __api_version__
import numpy as np
from typing import Union
class InvalidResponseError(Exception):
"""Custom exception class for triggering retry mechanism"""
pass
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
@ -105,10 +108,11 @@ async def openai_complete_if_cache(
default_headers = {
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
"Content-Type": "application/json"
"Content-Type": "application/json",
}
openai_async_client = (
AsyncOpenAI(default_headers=default_headers) if base_url is None
AsyncOpenAI(default_headers=default_headers)
if base_url is None
else AsyncOpenAI(base_url=base_url, default_headers=default_headers)
)
kwargs.pop("hashing_kv", None)
@ -295,10 +299,11 @@ async def openai_embed(
default_headers = {
"User-Agent": f"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
"Content-Type": "application/json"
"Content-Type": "application/json",
}
openai_async_client = (
AsyncOpenAI(default_headers=default_headers) if base_url is None
AsyncOpenAI(default_headers=default_headers)
if base_url is None
else AsyncOpenAI(base_url=base_url, default_headers=default_headers)
)
response = await openai_async_client.embeddings.create(