diff --git a/lightrag/llm/gemini.py b/lightrag/llm/gemini.py index 3954e814..983d6b9f 100644 --- a/lightrag/llm/gemini.py +++ b/lightrag/llm/gemini.py @@ -33,18 +33,27 @@ from lightrag.utils import ( import pipmaster as pm -# Install the Google Gemini client on demand +# Install the Google Gemini client and its dependencies on demand if not pm.is_installed("google-genai"): pm.install("google-genai") +if not pm.is_installed("google-api-core"): + pm.install("google-api-core") from google import genai # type: ignore from google.genai import types # type: ignore +from google.api_core import exceptions as google_api_exceptions # type: ignore DEFAULT_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com" LOG = logging.getLogger(__name__) +class InvalidResponseError(Exception): + """Custom exception class for triggering retry mechanism when Gemini returns empty responses""" + + pass + + @lru_cache(maxsize=8) def _get_gemini_client( api_key: str, base_url: str | None, timeout: int | None = None @@ -176,6 +185,21 @@ def _extract_response_text( return ("\n".join(regular_parts), "\n".join(thought_parts)) +@retry( + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=4, max=60), + retry=( + retry_if_exception_type(google_api_exceptions.InternalServerError) + | retry_if_exception_type(google_api_exceptions.ServiceUnavailable) + | retry_if_exception_type(google_api_exceptions.ResourceExhausted) + | retry_if_exception_type(google_api_exceptions.GatewayTimeout) + | retry_if_exception_type(google_api_exceptions.BadGateway) + | retry_if_exception_type(google_api_exceptions.DeadlineExceeded) + | retry_if_exception_type(google_api_exceptions.Aborted) + | retry_if_exception_type(google_api_exceptions.Unknown) + | retry_if_exception_type(InvalidResponseError) + ), +) async def gemini_complete_if_cache( model: str, prompt: str, @@ -382,7 +406,7 @@ async def gemini_complete_if_cache( final_text = regular_text or "" if not final_text: - raise RuntimeError("Gemini response did not contain any text content.") + raise InvalidResponseError("Gemini response did not contain any text content.") if "\\u" in final_text: final_text = safe_unicode_decode(final_text.encode("utf-8")) @@ -434,7 +458,14 @@ async def gemini_model_complete( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=60), retry=( - retry_if_exception_type(Exception) # Gemini uses generic exceptions + retry_if_exception_type(google_api_exceptions.InternalServerError) + | retry_if_exception_type(google_api_exceptions.ServiceUnavailable) + | retry_if_exception_type(google_api_exceptions.ResourceExhausted) + | retry_if_exception_type(google_api_exceptions.GatewayTimeout) + | retry_if_exception_type(google_api_exceptions.BadGateway) + | retry_if_exception_type(google_api_exceptions.DeadlineExceeded) + | retry_if_exception_type(google_api_exceptions.Aborted) + | retry_if_exception_type(google_api_exceptions.Unknown) ), ) async def gemini_embed( diff --git a/pyproject.toml b/pyproject.toml index aad54bc2..d0052b74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "aiohttp", "configparser", "future", + "google-api-core>=2.0.0,<3.0.0", "google-genai>=1.0.0,<2.0.0", "json_repair", "nano-vectordb", @@ -60,6 +61,7 @@ api = [ "tenacity", "tiktoken", "xlsxwriter>=3.1.0", + "google-api-core>=2.0.0,<3.0.0", "google-genai>=1.0.0,<2.0.0", # API-specific dependencies "aiofiles", @@ -108,6 +110,7 @@ offline-llm = [ "aioboto3>=12.0.0,<16.0.0", "voyageai>=0.2.0,<1.0.0", "llama-index>=0.9.0,<1.0.0", + "google-api-core>=2.0.0,<3.0.0", "google-genai>=1.0.0,<2.0.0", ] diff --git a/requirements-offline-llm.txt b/requirements-offline-llm.txt index 4e8b7168..1539552a 100644 --- a/requirements-offline-llm.txt +++ b/requirements-offline-llm.txt @@ -10,6 +10,7 @@ # LLM provider dependencies (with version constraints matching pyproject.toml) aioboto3>=12.0.0,<16.0.0 anthropic>=0.18.0,<1.0.0 +google-api-core>=2.0.0,<3.0.0 google-genai>=1.0.0,<2.0.0 llama-index>=0.9.0,<1.0.0 ollama>=0.1.0,<1.0.0