From f7432a260e0fd543dccbfb9606a8db7d316fed50 Mon Sep 17 00:00:00 2001 From: Louis Lacombe Date: Wed, 12 Nov 2025 16:11:05 +0000 Subject: [PATCH 1/3] Add support for environment variable fallback for API key and default host for cloud models --- lightrag/llm/ollama.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/lightrag/llm/ollama.py b/lightrag/llm/ollama.py index b013496e..28d91075 100644 --- a/lightrag/llm/ollama.py +++ b/lightrag/llm/ollama.py @@ -1,4 +1,5 @@ from collections.abc import AsyncIterator +import os import pipmaster as pm @@ -53,6 +54,9 @@ async def _ollama_model_if_cache( timeout = None kwargs.pop("hashing_kv", None) api_key = kwargs.pop("api_key", None) + # fallback to environment variable when not provided explicitly + if not api_key: + api_key = os.getenv("OLLAMA_API_KEY") headers = { "Content-Type": "application/json", "User-Agent": f"LightRAG/{__api_version__}", @@ -60,6 +64,16 @@ async def _ollama_model_if_cache( if api_key: headers["Authorization"] = f"Bearer {api_key}" + # If this is a cloud model (names include '-cloud' or ':cloud'), default + # the host to Ollama cloud when no explicit host was provided. + try: + model_name_str = str(model) if model is not None else "" + except Exception: + model_name_str = "" + + if host is None and ("-cloud" in model_name_str or ":cloud" in model_name_str): + host = "https://ollama.com" + ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers) try: @@ -144,6 +158,8 @@ async def ollama_model_complete( async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray: api_key = kwargs.pop("api_key", None) + if not api_key: + api_key = os.getenv("OLLAMA_API_KEY") headers = { "Content-Type": "application/json", "User-Agent": f"LightRAG/{__api_version__}", @@ -154,6 +170,15 @@ async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray: host = kwargs.pop("host", None) timeout = kwargs.pop("timeout", None) + # If embed_model targets Ollama cloud, default host when not provided + try: + embed_model_name = str(embed_model) if embed_model is not None else "" + except Exception: + embed_model_name = "" + + if host is None and ("-cloud" in embed_model_name or ":cloud" in embed_model_name): + host = "https://ollama.com" + ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers) try: options = kwargs.pop("options", {}) From 844537e378471b95a42705de8542edc82e7d304f Mon Sep 17 00:00:00 2001 From: LacombeLouis Date: Thu, 13 Nov 2025 12:17:51 +0100 Subject: [PATCH 2/3] Add a better regex --- lightrag/llm/ollama.py | 38 ++++++++++++++++++++------------------ 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/lightrag/llm/ollama.py b/lightrag/llm/ollama.py index 28d91075..9c2d17ee 100644 --- a/lightrag/llm/ollama.py +++ b/lightrag/llm/ollama.py @@ -1,5 +1,6 @@ from collections.abc import AsyncIterator import os +import re import pipmaster as pm @@ -23,10 +24,26 @@ from lightrag.exceptions import ( from lightrag.api import __api_version__ import numpy as np -from typing import Union +from typing import Optional, Union from lightrag.utils import logger +_OLLAMA_CLOUD_HOST = "https://ollama.com" +_CLOUD_MODEL_SUFFIX_PATTERN = re.compile(r"(?:-cloud|:cloud)$") + + +def _coerce_host_for_cloud_model(host: Optional[str], model: object) -> Optional[str]: + if host: + return host + try: + model_name_str = str(model) if model is not None else "" + except Exception: + model_name_str = "" + if _CLOUD_MODEL_SUFFIX_PATTERN.search(model_name_str): + return _OLLAMA_CLOUD_HOST + return host + + @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), @@ -64,15 +81,7 @@ async def _ollama_model_if_cache( if api_key: headers["Authorization"] = f"Bearer {api_key}" - # If this is a cloud model (names include '-cloud' or ':cloud'), default - # the host to Ollama cloud when no explicit host was provided. - try: - model_name_str = str(model) if model is not None else "" - except Exception: - model_name_str = "" - - if host is None and ("-cloud" in model_name_str or ":cloud" in model_name_str): - host = "https://ollama.com" + host = _coerce_host_for_cloud_model(host, model) ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers) @@ -170,14 +179,7 @@ async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray: host = kwargs.pop("host", None) timeout = kwargs.pop("timeout", None) - # If embed_model targets Ollama cloud, default host when not provided - try: - embed_model_name = str(embed_model) if embed_model is not None else "" - except Exception: - embed_model_name = "" - - if host is None and ("-cloud" in embed_model_name or ":cloud" in embed_model_name): - host = "https://ollama.com" + host = _coerce_host_for_cloud_model(host, embed_model) ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers) try: From 77ad906d3a4a6913f6d3f5e1d613f61ba27c0628 Mon Sep 17 00:00:00 2001 From: yangdx Date: Thu, 13 Nov 2025 20:41:44 +0800 Subject: [PATCH 3/3] Improve error handling and logging in cloud model detection --- lightrag/llm/ollama.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/lightrag/llm/ollama.py b/lightrag/llm/ollama.py index 9c2d17ee..670351bc 100644 --- a/lightrag/llm/ollama.py +++ b/lightrag/llm/ollama.py @@ -37,9 +37,13 @@ def _coerce_host_for_cloud_model(host: Optional[str], model: object) -> Optional return host try: model_name_str = str(model) if model is not None else "" - except Exception: + except (TypeError, ValueError, AttributeError) as e: + logger.warning(f"Failed to convert model to string: {e}, using empty string") model_name_str = "" if _CLOUD_MODEL_SUFFIX_PATTERN.search(model_name_str): + logger.debug( + f"Detected cloud model '{model_name_str}', using Ollama Cloud host" + ) return _OLLAMA_CLOUD_HOST return host