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

Merge pull request #608 from ParisNeo/main

Browse Source 
Temporary fix to introduced problem in ollama and lollm
This commit is contained in:
zrguo
2025-01-20 17:23:28 +08:00
committed by GitHub
parent a2993aeec2 501610c603
commit d9a22850c5
1 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
lightrag/llm.py
View File

@@ -350,7 +350,10 @@ 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 = {"Authorization": f"Bearer {api_key}"} if api_key else None headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}"
} if api_key else {"Content-Type": "application/json"}
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:
@@ -383,7 +386,10 @@ async def lollms_model_if_cache(
stream = True if kwargs.get("stream") else False stream = True if kwargs.get("stream") else False
api_key = kwargs.pop("api_key", None) api_key = kwargs.pop("api_key", None)
headers = {"Authorization": f"Bearer {api_key}"} if api_key else None headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}"
} if api_key else {"Content-Type": "application/json"}
# Extract lollms specific parameters # Extract lollms specific parameters
request_data = { request_data = {
@@ -1154,9 +1160,9 @@ 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 = (
{"Authorization": api_key, "Content-Type": "application/json"} {"Content-Type": "application/json", "Authorization": api_key}
if api_key if api_key
else None else {"Content-Type": "application/json"}
) )
kwargs["headers"] = headers kwargs["headers"] = headers
ollama_client = ollama.Client(**kwargs) ollama_client = ollama.Client(**kwargs)
@@ -1181,9 +1187,9 @@ async def lollms_embed(
""" """
api_key = kwargs.pop("api_key", None) api_key = kwargs.pop("api_key", None)
headers = ( headers = (
{"Authorization": api_key, "Content-Type": "application/json"} {"Content-Type": "application/json", "Authorization": api_key}
if api_key if api_key
else None else {"Content-Type": "application/json"}
) )
async with aiohttp.ClientSession(headers=headers) as session: async with aiohttp.ClientSession(headers=headers) as session:
embeddings = [] embeddings = []