gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
• 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
This commit is contained in:
yangdx
2025-02-06 22:55:22 +08:00
committed by ultrageopro
parent cec8da7f91
commit b90f3f14be
2 changed files with 26 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
lightrag/llm/ollama.py
View File

@@ -66,6 +66,7 @@ from lightrag.exceptions import (
RateLimitError, RateLimitError,
APITimeoutError, APITimeoutError,
) )
from lightrag.api import __api_version__
from lightrag.utils import extract_reasoning from lightrag.utils import extract_reasoning
import numpy as np import numpy as np
from typing import Union from typing import Union
@@ -93,11 +94,12 @@ async def ollama_model_if_cache(
timeout = kwargs.pop("timeout", None) timeout = kwargs.pop("timeout", None)
kwargs.pop("hashing_kv", None) kwargs.pop("hashing_kv", None)
api_key = kwargs.pop("api_key", None) api_key = kwargs.pop("api_key", None)
headers = ( headers = {
{"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"} "Content-Type": "application/json",
if api_key "User-Agent": f"LightRAG/{__api_version__}",
else {"Content-Type": "application/json"} }
) if api_key:
headers["Authorization"] = f"Bearer {api_key}"
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers) ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
messages = [] messages = []
if system_prompt: if system_prompt:
@@ -161,11 +163,12 @@ async def ollama_embedding(texts: list[str], embed_model, **kwargs) -> np.ndarra
async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray: async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
api_key = kwargs.pop("api_key", None) api_key = kwargs.pop("api_key", None)
headers = ( headers = {
{"Content-Type": "application/json", "Authorization": api_key} "Content-Type": "application/json",
if api_key "User-Agent": f"LightRAG/{__api_version__}",
else {"Content-Type": "application/json"} }
) if api_key:
headers["Authorization"] = api_key
kwargs["headers"] = headers kwargs["headers"] = headers
ollama_client = ollama.Client(**kwargs) ollama_client = ollama.Client(**kwargs)
data = ollama_client.embed(model=embed_model, input=texts) data = ollama_client.embed(model=embed_model, input=texts)

15
lightrag/llm/openai.py
View File

@@ -73,6 +73,7 @@ from lightrag.utils import (
logger, logger,
) )
from lightrag.types import GPTKeywordExtractionFormat from lightrag.types import GPTKeywordExtractionFormat
from lightrag.api import __api_version__
import numpy as np import numpy as np
from typing import Union from typing import Union
@@ -102,8 +103,13 @@ async def openai_complete_if_cache(
if api_key: if api_key:
os.environ["OPENAI_API_KEY"] = api_key os.environ["OPENAI_API_KEY"] = api_key
default_headers = {
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
"Content-Type": "application/json"
}
openai_async_client = ( openai_async_client = (
AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url) 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) kwargs.pop("hashing_kv", None)
kwargs.pop("keyword_extraction", None) kwargs.pop("keyword_extraction", None)
@@ -287,8 +293,13 @@ async def openai_embed(
if api_key: if api_key:
os.environ["OPENAI_API_KEY"] = api_key os.environ["OPENAI_API_KEY"] = api_key
default_headers = {
"User-Agent": f"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
"Content-Type": "application/json"
}
openai_async_client = ( openai_async_client = (
AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url) 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( response = await openai_async_client.embeddings.create(
model=model, input=texts, encoding_format="float" model=model, input=texts, encoding_format="float"