yangdx commited on
Commit
8f6d644
·
1 Parent(s): 2dd7565

Add LightRAG version to User-Agent header for better request tracking

Browse files

• Add User-Agent header with version info
• Update header creation in Ollama client
• Update header creation in OpenAI client
• Ensure consistent header format
• Include Mozilla UA string for OpenAI

Files changed (2) hide show
  1. lightrag/llm/ollama.py +13 -10
  2. lightrag/llm/openai.py +13 -2
lightrag/llm/ollama.py CHANGED
@@ -66,6 +66,7 @@ from lightrag.exceptions import (
66
  RateLimitError,
67
  APITimeoutError,
68
  )
 
69
  import numpy as np
70
  from typing import Union
71
 
@@ -91,11 +92,12 @@ async def ollama_model_if_cache(
91
  timeout = kwargs.pop("timeout", None)
92
  kwargs.pop("hashing_kv", None)
93
  api_key = kwargs.pop("api_key", None)
94
- headers = (
95
- {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
96
- if api_key
97
- else {"Content-Type": "application/json"}
98
- )
 
99
  ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
100
  messages = []
101
  if system_prompt:
@@ -147,11 +149,12 @@ async def ollama_embedding(texts: list[str], embed_model, **kwargs) -> np.ndarra
147
 
148
  async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
149
  api_key = kwargs.pop("api_key", None)
150
- headers = (
151
- {"Content-Type": "application/json", "Authorization": api_key}
152
- if api_key
153
- else {"Content-Type": "application/json"}
154
- )
 
155
  kwargs["headers"] = headers
156
  ollama_client = ollama.Client(**kwargs)
157
  data = ollama_client.embed(model=embed_model, input=texts)
 
66
  RateLimitError,
67
  APITimeoutError,
68
  )
69
+ from lightrag.api import __api_version__
70
  import numpy as np
71
  from typing import Union
72
 
 
92
  timeout = kwargs.pop("timeout", None)
93
  kwargs.pop("hashing_kv", None)
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}"
101
  ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
102
  messages = []
103
  if system_prompt:
 
149
 
150
  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
158
  kwargs["headers"] = headers
159
  ollama_client = ollama.Client(**kwargs)
160
  data = ollama_client.embed(model=embed_model, input=texts)
lightrag/llm/openai.py CHANGED
@@ -73,6 +73,7 @@ from lightrag.utils import (
73
  logger,
74
  )
75
  from lightrag.types import GPTKeywordExtractionFormat
 
76
 
77
  import numpy as np
78
  from typing import Union
@@ -102,8 +103,13 @@ async def openai_complete_if_cache(
102
  if api_key:
103
  os.environ["OPENAI_API_KEY"] = api_key
104
 
 
 
 
 
105
  openai_async_client = (
106
- AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url)
 
107
  )
108
  kwargs.pop("hashing_kv", None)
109
  kwargs.pop("keyword_extraction", None)
@@ -287,8 +293,13 @@ async def openai_embed(
287
  if api_key:
288
  os.environ["OPENAI_API_KEY"] = api_key
289
 
 
 
 
 
290
  openai_async_client = (
291
- AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url)
 
292
  )
293
  response = await openai_async_client.embeddings.create(
294
  model=model, input=texts, encoding_format="float"
 
73
  logger,
74
  )
75
  from lightrag.types import GPTKeywordExtractionFormat
76
+ from lightrag.api import __api_version__
77
 
78
  import numpy as np
79
  from typing import Union
 
103
  if api_key:
104
  os.environ["OPENAI_API_KEY"] = api_key
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)
115
  kwargs.pop("keyword_extraction", None)
 
293
  if api_key:
294
  os.environ["OPENAI_API_KEY"] = api_key
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(
305
  model=model, input=texts, encoding_format="float"