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
- lightrag/llm/ollama.py +13 -10
- 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 |
-
|
96 |
-
|
97 |
-
|
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 |
-
|
152 |
-
|
153 |
-
|
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
|
|
|
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
|
|
|
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"
|