yangdx commited on
Commit
8662c14
·
1 Parent(s): b802259

Fix linting

Browse files
Files changed (2) hide show
  1. lightrag/llm/ollama.py +2 -2
  2. lightrag/llm/openai.py +9 -4
lightrag/llm/ollama.py CHANGED
@@ -94,7 +94,7 @@ async def ollama_model_if_cache(
94
  api_key = kwargs.pop("api_key", None)
95
  headers = {
96
  "Content-Type": "application/json",
97
- "User-Agent": f"LightRAG/{__api_version__}"
98
  }
99
  if api_key:
100
  headers["Authorization"] = f"Bearer {api_key}"
@@ -151,7 +151,7 @@ async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
151
  api_key = kwargs.pop("api_key", None)
152
  headers = {
153
  "Content-Type": "application/json",
154
- "User-Agent": f"LightRAG/{__api_version__}"
155
  }
156
  if api_key:
157
  headers["Authorization"] = api_key
 
94
  api_key = kwargs.pop("api_key", None)
95
  headers = {
96
  "Content-Type": "application/json",
97
+ "User-Agent": f"LightRAG/{__api_version__}",
98
  }
99
  if api_key:
100
  headers["Authorization"] = f"Bearer {api_key}"
 
151
  api_key = kwargs.pop("api_key", None)
152
  headers = {
153
  "Content-Type": "application/json",
154
+ "User-Agent": f"LightRAG/{__api_version__}",
155
  }
156
  if api_key:
157
  headers["Authorization"] = api_key
lightrag/llm/openai.py CHANGED
@@ -78,10 +78,13 @@ from lightrag.api import __api_version__
78
  import numpy as np
79
  from typing import Union
80
 
 
81
  class InvalidResponseError(Exception):
82
  """Custom exception class for triggering retry mechanism"""
 
83
  pass
84
 
 
85
  @retry(
86
  stop=stop_after_attempt(3),
87
  wait=wait_exponential(multiplier=1, min=4, max=10),
@@ -105,10 +108,11 @@ async def openai_complete_if_cache(
105
 
106
  default_headers = {
107
  "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
108
- "Content-Type": "application/json"
109
  }
110
  openai_async_client = (
111
- AsyncOpenAI(default_headers=default_headers) if base_url is None
 
112
  else AsyncOpenAI(base_url=base_url, default_headers=default_headers)
113
  )
114
  kwargs.pop("hashing_kv", None)
@@ -295,10 +299,11 @@ async def openai_embed(
295
 
296
  default_headers = {
297
  "User-Agent": f"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
298
- "Content-Type": "application/json"
299
  }
300
  openai_async_client = (
301
- AsyncOpenAI(default_headers=default_headers) if base_url is None
 
302
  else AsyncOpenAI(base_url=base_url, default_headers=default_headers)
303
  )
304
  response = await openai_async_client.embeddings.create(
 
78
  import numpy as np
79
  from typing import Union
80
 
81
+
82
  class InvalidResponseError(Exception):
83
  """Custom exception class for triggering retry mechanism"""
84
+
85
  pass
86
 
87
+
88
  @retry(
89
  stop=stop_after_attempt(3),
90
  wait=wait_exponential(multiplier=1, min=4, max=10),
 
108
 
109
  default_headers = {
110
  "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
111
+ "Content-Type": "application/json",
112
  }
113
  openai_async_client = (
114
+ AsyncOpenAI(default_headers=default_headers)
115
+ if base_url is None
116
  else AsyncOpenAI(base_url=base_url, default_headers=default_headers)
117
  )
118
  kwargs.pop("hashing_kv", None)
 
299
 
300
  default_headers = {
301
  "User-Agent": f"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
302
+ "Content-Type": "application/json",
303
  }
304
  openai_async_client = (
305
+ AsyncOpenAI(default_headers=default_headers)
306
+ if base_url is None
307
  else AsyncOpenAI(base_url=base_url, default_headers=default_headers)
308
  )
309
  response = await openai_async_client.embeddings.create(