""" Gemini LLM binding for LightRAG. This module provides asynchronous helpers that adapt Google's Gemini models to the same interface used by the rest of the LightRAG LLM bindings. The implementation mirrors the OpenAI helpers while relying on the official ``google-genai`` client under the hood. """ from __future__ import annotations import os import warnings from collections.abc import AsyncIterator from functools import lru_cache from typing import Any import numpy as np from tenacity import ( retry, stop_after_attempt, wait_exponential, retry_if_exception_type, ) from lightrag.utils import ( logger, remove_think_tags, safe_unicode_decode, wrap_embedding_func_with_attrs, ) import pipmaster as pm # 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 class InvalidResponseError(Exception): """Custom exception class for triggering retry mechanism when Gemini returns empty responses""" pass _DEFAULT_GEMINI_BASE_URLS = { "https://generativelanguage.googleapis.com", "https://generativelanguage.googleapis.com/", "https://generativelanguage.googleapis.com/v1beta", "https://generativelanguage.googleapis.com/v1beta/", "https://generativelanguage.googleapis.com/v1", "https://generativelanguage.googleapis.com/v1/", } def _normalize_gemini_base_url(base_url: str | None) -> str | None: """Treat Google's default Gemini API service roots as SDK defaults.""" if not base_url: return None normalized = base_url.strip() if not normalized or normalized == "DEFAULT_GEMINI_ENDPOINT": return None if normalized.rstrip("/") in { service_root.rstrip("/") for service_root in _DEFAULT_GEMINI_BASE_URLS }: return None return normalized @lru_cache(maxsize=8) def _get_gemini_client( api_key: str, base_url: str | None, timeout: int | None = None ) -> genai.Client: """ Create (or fetch cached) Gemini client. Args: api_key: Google Gemini API key (not used in Vertex AI mode). base_url: Optional custom API endpoint. timeout: Optional request timeout in milliseconds. Returns: genai.Client: Configured Gemini client instance. """ client_kwargs: dict[str, Any] = {} normalized_base_url = _normalize_gemini_base_url(base_url) # Add Vertex AI support use_vertexai = os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true" if use_vertexai: # Vertex AI mode: use project/location, NOT api_key client_kwargs["vertexai"] = True project = os.getenv("GOOGLE_CLOUD_PROJECT") if project: location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1") client_kwargs["project"] = project if location: client_kwargs["location"] = location else: raise ValueError( "GOOGLE_CLOUD_PROJECT must be set when using Vertex AI mode" ) else: # Standard Gemini API mode: use api_key client_kwargs["api_key"] = api_key if normalized_base_url is not None or timeout is not None: try: http_options_kwargs = {} if normalized_base_url is not None: http_options_kwargs["base_url"] = normalized_base_url if timeout is not None: http_options_kwargs["timeout"] = timeout client_kwargs["http_options"] = types.HttpOptions(**http_options_kwargs) except Exception as e: logger.error("Failed to apply custom Gemini http_options: %s", e) raise e return genai.Client(**client_kwargs) def _ensure_api_key(api_key: str | None) -> str: # In Vertex AI mode, API key is not required use_vertexai = os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true" if use_vertexai: # Return empty string for Vertex AI mode (not used) return "" key = api_key or os.getenv("LLM_BINDING_API_KEY") or os.getenv("GEMINI_API_KEY") if not key: raise ValueError( "Gemini API key not provided. " "Set LLM_BINDING_API_KEY or GEMINI_API_KEY in the environment." ) return key def _build_generation_config( base_config: dict[str, Any] | None, system_prompt: str | None, response_format: Any | None, ) -> types.GenerateContentConfig | None: config_data = dict(base_config or {}) if system_prompt: if config_data.get("system_instruction"): config_data["system_instruction"] = ( f"{config_data['system_instruction']}\n{system_prompt}" ) else: config_data["system_instruction"] = system_prompt # Translate response_format to Gemini's native generation config fields. if response_format is not None: config_data.setdefault("response_mime_type", "application/json") schema = _normalize_gemini_response_schema(response_format) if schema is not None and "response_json_schema" not in config_data: config_data["response_json_schema"] = schema # Remove entries that are explicitly set to None to avoid type errors sanitized = { key: value for key, value in config_data.items() if value is not None and value != "" } if not sanitized: return None return types.GenerateContentConfig(**sanitized) def _normalize_gemini_response_schema(response_format: Any) -> Any | None: """Extract a Gemini-compatible JSON schema from LightRAG/OpenAI inputs.""" if response_format is None: return None if isinstance(response_format, dict): if response_format.get("type") == "json_object": return None if response_format.get("type") == "json_schema": json_schema = response_format.get("json_schema") if isinstance(json_schema, dict): schema = json_schema.get("schema") if isinstance(schema, dict): return schema return json_schema return response_format return response_format def _validate_gemini_response_format(response_format: Any | None) -> None: """Reject typed structured-output helpers; only dict payloads are supported.""" if response_format is None or isinstance(response_format, dict): return raise TypeError( "gemini_complete_if_cache only supports dict response_format payloads; " "typed/Pydantic response_format values are not supported." ) def _format_history_messages(history_messages: list[dict[str, Any]] | None) -> str: if not history_messages: return "" history_lines: list[str] = [] for message in history_messages: role = message.get("role", "user") content = message.get("content", "") history_lines.append(f"[{role}] {content}") return "\n".join(history_lines) def _extract_response_text( response: Any, extract_thoughts: bool = False ) -> tuple[str, str]: """ Extract text content from Gemini response, separating regular content from thoughts. Args: response: Gemini API response object extract_thoughts: Whether to extract thought content separately Returns: Tuple of (regular_text, thought_text) """ candidates = getattr(response, "candidates", None) if not candidates: return ("", "") regular_parts: list[str] = [] thought_parts: list[str] = [] for candidate in candidates: if not getattr(candidate, "content", None): continue # Use 'or []' to handle None values from parts attribute for part in getattr(candidate.content, "parts", None) or []: text = getattr(part, "text", None) if not text: continue # Check if this part is thought content using the 'thought' attribute is_thought = getattr(part, "thought", False) if is_thought and extract_thoughts: thought_parts.append(text) elif not is_thought: regular_parts.append(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, system_prompt: str | None = None, history_messages: list[dict[str, Any]] | None = None, enable_cot: bool = False, base_url: str | None = None, api_key: str | None = None, token_tracker: Any | None = None, stream: bool | None = None, response_format: Any | None = None, keyword_extraction: bool = False, entity_extraction: bool = False, generation_config: dict[str, Any] | None = None, timeout: int | None = None, image_inputs: list[Any] | None = None, **_: Any, ) -> str | AsyncIterator[str]: """ Complete a prompt using Gemini's API with Chain of Thought (COT) support. This function supports automatic integration of reasoning content from Gemini models that provide Chain of Thought capabilities via the thinking_config API feature. Structured output note: - This adapter accepts OpenAI-style ``response_format`` and translates it to Gemini's native generation config fields. - ``response_format={"type": "json_object"}`` maps to ``response_mime_type="application/json"``. - Dict-form ``json_schema`` payloads map to ``response_mime_type="application/json"`` plus ``response_json_schema=``. - Typed/Pydantic ``response_format`` helpers are rejected explicitly. - Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are compatibility shims; when no explicit ``response_format`` is supplied, they are mapped to ``{"type": "json_object"}``. COT Integration: - When enable_cot=True: Thought content is wrapped in ... tags - When enable_cot=False: Thought content is filtered out, only regular content returned - Thought content is identified by the 'thought' attribute on response parts - Requires thinking_config to be enabled in generation_config for API to return thoughts Args: model: The Gemini model to use. prompt: The prompt to complete. system_prompt: Optional system prompt to include. history_messages: Optional list of previous messages in the conversation. api_key: Optional Gemini API key. If None, uses environment variable. base_url: Optional custom API endpoint. generation_config: Optional generation configuration dict. response_format: OpenAI-style structured output control translated to Gemini generation config. ``{"type": "json_object"}`` maps to ``response_mime_type="application/json"``; dict-form ``json_schema`` payloads map to ``response_json_schema``. Typed/Pydantic response_format values are rejected. token_tracker: Optional token usage tracker for monitoring API usage. stream: Whether to stream the response. hashing_kv: Storage interface (for interface parity with other bindings). enable_cot: Whether to include Chain of Thought content in the response. timeout: Request timeout in seconds (will be converted to milliseconds for Gemini API). **_: Additional keyword arguments (ignored). Returns: The completed text (with COT content if enable_cot=True) or an async iterator of text chunks if streaming. COT content is wrapped in ... tags. Raises: RuntimeError: If the response from Gemini is empty. ValueError: If API key is not provided or configured. """ key = _ensure_api_key(api_key) # Convert timeout from seconds to milliseconds for Gemini API timeout_ms = timeout * 1000 if timeout else None client = _get_gemini_client(key, base_url, timeout_ms) # Deprecation shims: map legacy boolean flags to response_format only when # an explicit response_format was not supplied. if response_format is None: if entity_extraction: warnings.warn( "gemini_complete_if_cache(entity_extraction=True) is deprecated; " "pass response_format={'type': 'json_object'} instead.", DeprecationWarning, stacklevel=2, ) response_format = {"type": "json_object"} elif keyword_extraction: warnings.warn( "gemini_complete_if_cache(keyword_extraction=True) is deprecated; " "pass response_format={'type': 'json_object'} instead.", DeprecationWarning, stacklevel=2, ) response_format = {"type": "json_object"} _validate_gemini_response_format(response_format) if response_format is not None: enable_cot = False history_block = _format_history_messages(history_messages) prompt_sections = [] if history_block: prompt_sections.append(history_block) prompt_sections.append(f"[user] {prompt}") combined_prompt = "\n".join(prompt_sections) config_obj = _build_generation_config( generation_config, system_prompt=system_prompt, response_format=response_format, ) if image_inputs: from lightrag.llm._vision_utils import normalize_image_inputs normalized_images = normalize_image_inputs(image_inputs) parts: list[Any] = [combined_prompt] parts.extend( types.Part.from_bytes(data=img.raw_bytes, mime_type=img.mime_type) for img in normalized_images ) contents: list[Any] = [parts] else: contents = [combined_prompt] request_kwargs: dict[str, Any] = { "model": model, "contents": contents, } if config_obj is not None: request_kwargs["config"] = config_obj if stream: async def _async_stream() -> AsyncIterator[str]: # COT state tracking for streaming cot_active = False cot_started = False initial_content_seen = False usage_metadata = None try: # Use native async streaming from genai SDK # Note: generate_content_stream returns Awaitable[AsyncIterator], need to await first stream_iter = await client.aio.models.generate_content_stream( **request_kwargs ) async for chunk in stream_iter: usage = getattr(chunk, "usage_metadata", None) if usage is not None: usage_metadata = usage # Extract both regular and thought content regular_text, thought_text = _extract_response_text( chunk, extract_thoughts=True ) if enable_cot: # Process regular content if regular_text: if not initial_content_seen: initial_content_seen = True # Close COT section if it was active if cot_active: yield "" cot_active = False # Process and yield regular content if "\\u" in regular_text: regular_text = safe_unicode_decode( regular_text.encode("utf-8") ) yield regular_text # Process thought content if thought_text: if not initial_content_seen and not cot_started: # Start COT section yield "" cot_active = True cot_started = True # Yield thought content if COT is active if cot_active: if "\\u" in thought_text: thought_text = safe_unicode_decode( thought_text.encode("utf-8") ) yield thought_text else: # COT disabled - only yield regular content if regular_text: if "\\u" in regular_text: regular_text = safe_unicode_decode( regular_text.encode("utf-8") ) yield regular_text # Ensure COT is properly closed if still active if cot_active: yield "" cot_active = False except Exception: # Try to close COT tag before re-raising if cot_active: try: yield "" except Exception: pass raise finally: # Track token usage after streaming completes if token_tracker and usage_metadata: token_tracker.add_usage( { "prompt_tokens": getattr( usage_metadata, "prompt_token_count", 0 ), "completion_tokens": getattr( usage_metadata, "candidates_token_count", 0 ), "total_tokens": getattr( usage_metadata, "total_token_count", 0 ), } ) return _async_stream() # Non-streaming: use native async client response = await client.aio.models.generate_content(**request_kwargs) # Extract both regular text and thought text regular_text, thought_text = _extract_response_text(response, extract_thoughts=True) # Apply COT filtering logic based on enable_cot parameter if enable_cot: # Include thought content wrapped in tags if thought_text and thought_text.strip(): if not regular_text or regular_text.strip() == "": # Only thought content available final_text = f"{thought_text}" else: # Both content types present: prepend thought to regular content final_text = f"{thought_text}{regular_text}" else: # No thought content, use regular content only final_text = regular_text or "" else: # Filter out thought content, return only regular content final_text = regular_text or "" if not final_text: 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")) final_text = remove_think_tags(final_text) usage = getattr(response, "usage_metadata", None) if token_tracker and usage: token_tracker.add_usage( { "prompt_tokens": getattr(usage, "prompt_token_count", 0), "completion_tokens": getattr(usage, "candidates_token_count", 0), "total_tokens": getattr(usage, "total_token_count", 0), } ) logger.debug("Gemini response length: %s", len(final_text)) return final_text async def gemini_model_complete( prompt: str, system_prompt: str | None = None, history_messages: list[dict[str, Any]] | None = None, response_format: Any | None = None, keyword_extraction: bool = False, entity_extraction: bool = False, **kwargs: Any, ) -> str | AsyncIterator[str]: # Accept legacy keyword if passed via kwargs to preserve backwards compat. entity_extraction = kwargs.pop("entity_extraction", entity_extraction) hashing_kv = kwargs.get("hashing_kv") model_name = None if hashing_kv is not None: model_name = hashing_kv.global_config.get("llm_model_name") if model_name is None: model_name = kwargs.pop("model_name", None) if model_name is None: raise ValueError("Gemini model name not provided in configuration.") return await gemini_complete_if_cache( model_name, prompt, system_prompt=system_prompt, history_messages=history_messages, response_format=response_format, keyword_extraction=keyword_extraction, entity_extraction=entity_extraction, **kwargs, ) @wrap_embedding_func_with_attrs( embedding_dim=1536, max_token_size=2048, model_name="gemini-embedding-001", supports_asymmetric=True, ) @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) ), ) async def gemini_embed( texts: list[str], model: str = "gemini-embedding-001", base_url: str | None = None, api_key: str | None = None, embedding_dim: int | None = None, max_token_size: int | None = None, task_type: str | None = None, timeout: int | None = None, token_tracker: Any | None = None, context: str = "document", ) -> np.ndarray: """Generate embeddings for a list of texts using Gemini's API. This function uses Google's Gemini embedding model to generate text embeddings. It supports dynamic dimension control and automatic normalization for dimensions less than 3072. Args: texts: List of texts to embed. model: The Gemini embedding model to use. Default is "gemini-embedding-001". base_url: Optional custom API endpoint. api_key: Optional Gemini API key. If None, uses environment variables. embedding_dim: Optional embedding dimension for dynamic dimension reduction. **IMPORTANT**: This parameter is automatically injected by the EmbeddingFunc wrapper. Do NOT manually pass this parameter when calling the function directly. The dimension is controlled by the @wrap_embedding_func_with_attrs decorator or the EMBEDDING_DIM environment variable. Supported range: 128-3072. Recommended values: 768, 1536, 3072. max_token_size: Maximum tokens per text. This parameter is automatically injected by the EmbeddingFunc wrapper when the underlying function signature supports it (via inspect.signature check). Gemini API will automatically truncate texts exceeding this limit (autoTruncate=True by default), so no client-side truncation is needed. task_type: Task type for embedding optimization. Default is "RETRIEVAL_DOCUMENT". Supported types: SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING, RETRIEVAL_DOCUMENT, RETRIEVAL_QUERY, CODE_RETRIEVAL_QUERY, QUESTION_ANSWERING, FACT_VERIFICATION. timeout: Request timeout in seconds (will be converted to milliseconds for Gemini API). token_tracker: Optional token usage tracker for monitoring API usage. context: The embedding context - "query" for search queries, "document" for indexed content. **IMPORTANT**: This parameter is automatically injected by the EmbeddingFunc wrapper when supports_asymmetric=True. Default is "document". Returns: A numpy array of embeddings, one per input text. For dimensions < 3072, the embeddings are L2-normalized to ensure optimal semantic similarity performance. Raises: ValueError: If API key is not provided or configured. RuntimeError: If the response from Gemini is invalid or empty. Note: - For dimension 3072: Embeddings are already normalized by the API - For dimensions < 3072: Embeddings are L2-normalized after retrieval - Normalization ensures accurate semantic similarity via cosine distance - Gemini API automatically truncates texts exceeding max_token_size (autoTruncate=True) """ # Note: max_token_size is received but not used for client-side truncation. # Gemini API handles truncation automatically with autoTruncate=True (default). _ = max_token_size # Acknowledge parameter to avoid unused variable warning key = _ensure_api_key(api_key) # Convert timeout from seconds to milliseconds for Gemini API timeout_ms = timeout * 1000 if timeout else None client = _get_gemini_client(key, base_url, timeout_ms) # Prepare embedding configuration config_kwargs: dict[str, Any] = {} # Add task_type to config if task_type is None: if context == "query": task_type = "RETRIEVAL_QUERY" elif context == "document": task_type = "RETRIEVAL_DOCUMENT" else: task_type = "RETRIEVAL_DOCUMENT" # Default for backward compatibility config_kwargs["task_type"] = task_type # Add output_dimensionality if embedding_dim is provided if embedding_dim is not None: config_kwargs["output_dimensionality"] = embedding_dim # Create config object if we have parameters config_obj = types.EmbedContentConfig(**config_kwargs) if config_kwargs else None request_kwargs: dict[str, Any] = { "model": model, "contents": texts, } if config_obj is not None: request_kwargs["config"] = config_obj # Use native async client for embedding response = await client.aio.models.embed_content(**request_kwargs) # Extract embeddings from response if not hasattr(response, "embeddings") or not response.embeddings: raise RuntimeError("Gemini response did not contain embeddings.") # Convert embeddings to numpy array embeddings = np.array( [np.array(e.values, dtype=np.float32) for e in response.embeddings] ) # Apply L2 normalization for dimensions < 3072 # The 3072 dimension embedding is already normalized by Gemini API if embedding_dim and embedding_dim < 3072: # Normalize each embedding vector to unit length norms = np.linalg.norm(embeddings, axis=1, keepdims=True) # Avoid division by zero norms = np.where(norms == 0, 1, norms) embeddings = embeddings / norms logger.debug( f"Applied L2 normalization to {len(embeddings)} embeddings of dimension {embedding_dim}" ) # Track token usage if tracker is provided # Note: Gemini embedding API may not provide usage metadata if token_tracker and hasattr(response, "usage_metadata"): usage = response.usage_metadata token_counts = { "prompt_tokens": getattr(usage, "prompt_token_count", 0), "total_tokens": getattr(usage, "total_token_count", 0), } token_tracker.add_usage(token_counts) logger.debug( f"Generated {len(embeddings)} Gemini embeddings with dimension {embeddings.shape[1]}" ) return embeddings __all__ = [ "gemini_complete_if_cache", "gemini_model_complete", "gemini_embed", ]