From a0514eec1a948c5d4400d73ca97fd34ac318215d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20MANSUY?= Date: Thu, 4 Dec 2025 19:14:28 +0800 Subject: [PATCH] cherry-pick 0c4cba38 --- lightrag/llm/openai.py | 374 ++++++++++++++++++++++++++++++++--------- lightrag/utils.py | 71 +++++++- 2 files changed, 367 insertions(+), 78 deletions(-) diff --git a/lightrag/llm/openai.py b/lightrag/llm/openai.py index 61db965f..4052727b 100644 --- a/lightrag/llm/openai.py +++ b/lightrag/llm/openai.py @@ -77,46 +77,73 @@ class InvalidResponseError(Exception): def create_openai_async_client( api_key: str | None = None, base_url: str | None = None, + use_azure: bool = False, + azure_deployment: str | None = None, + api_version: str | None = None, + timeout: int | None = None, client_configs: dict[str, Any] | None = None, ) -> AsyncOpenAI: - """Create an AsyncOpenAI client with the given configuration. + """Create an AsyncOpenAI or AsyncAzureOpenAI client with the given configuration. Args: api_key: OpenAI API key. If None, uses the OPENAI_API_KEY environment variable. base_url: Base URL for the OpenAI API. If None, uses the default OpenAI API URL. + use_azure: Whether to create an Azure OpenAI client. Default is False. + azure_deployment: Azure OpenAI deployment name (only used when use_azure=True). + api_version: Azure OpenAI API version (only used when use_azure=True). + timeout: Request timeout in seconds. client_configs: Additional configuration options for the AsyncOpenAI client. These will override any default configurations but will be overridden by explicit parameters (api_key, base_url). Returns: - An AsyncOpenAI client instance. + An AsyncOpenAI or AsyncAzureOpenAI client instance. """ - if not api_key: - api_key = os.environ["OPENAI_API_KEY"] + if use_azure: + from openai import AsyncAzureOpenAI - default_headers = { - "User-Agent": f"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}", - "Content-Type": "application/json", - } + if not api_key: + api_key = os.environ.get("AZURE_OPENAI_API_KEY") or os.environ.get( + "LLM_BINDING_API_KEY" + ) - if client_configs is None: - client_configs = {} - - # Create a merged config dict with precedence: explicit params > client_configs > defaults - merged_configs = { - **client_configs, - "default_headers": default_headers, - "api_key": api_key, - } - - if base_url is not None: - merged_configs["base_url"] = base_url - else: - merged_configs["base_url"] = os.environ.get( - "OPENAI_API_BASE", "https://api.openai.com/v1" + return AsyncAzureOpenAI( + azure_endpoint=base_url, + azure_deployment=azure_deployment, + api_key=api_key, + api_version=api_version, + timeout=timeout, ) + else: + if not api_key: + api_key = os.environ["OPENAI_API_KEY"] - return AsyncOpenAI(**merged_configs) + default_headers = { + "User-Agent": f"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}", + "Content-Type": "application/json", + } + + if client_configs is None: + client_configs = {} + + # Create a merged config dict with precedence: explicit params > client_configs > defaults + merged_configs = { + **client_configs, + "default_headers": default_headers, + "api_key": api_key, + } + + if base_url is not None: + merged_configs["base_url"] = base_url + else: + merged_configs["base_url"] = os.environ.get( + "OPENAI_API_BASE", "https://api.openai.com/v1" + ) + + if timeout is not None: + merged_configs["timeout"] = timeout + + return AsyncOpenAI(**merged_configs) @retry( @@ -138,6 +165,12 @@ async def openai_complete_if_cache( base_url: str | None = None, api_key: str | None = None, token_tracker: Any | None = None, + stream: bool | None = None, + timeout: int | None = None, + keyword_extraction: bool = False, + use_azure: bool = False, + azure_deployment: str | None = None, + api_version: str | None = None, **kwargs: Any, ) -> str: """Complete a prompt using OpenAI's API with caching support and Chain of Thought (COT) integration. @@ -167,13 +200,15 @@ async def openai_complete_if_cache( api_key: Optional OpenAI API key. If None, uses the OPENAI_API_KEY environment variable. token_tracker: Optional token usage tracker for monitoring API usage. enable_cot: Whether to enable Chain of Thought (COT) processing. Default is False. + stream: Whether to stream the response. Default is False. + timeout: Request timeout in seconds. Default is None. + keyword_extraction: Whether to enable keyword extraction mode. When True, triggers + special response formatting for keyword extraction. Default is False. **kwargs: Additional keyword arguments to pass to the OpenAI API. Special kwargs: - openai_client_configs: Dict of configuration options for the AsyncOpenAI client. These will be passed to the client constructor but will be overridden by explicit parameters (api_key, base_url). - - hashing_kv: Will be removed from kwargs before passing to OpenAI. - - keyword_extraction: Will be removed from kwargs before passing to OpenAI. Returns: The completed text (with integrated COT content if available) or an async iterator @@ -194,15 +229,22 @@ async def openai_complete_if_cache( # Remove special kwargs that shouldn't be passed to OpenAI kwargs.pop("hashing_kv", None) - kwargs.pop("keyword_extraction", None) # Extract client configuration options client_configs = kwargs.pop("openai_client_configs", {}) - # Create the OpenAI client + # Handle keyword extraction mode + if keyword_extraction: + kwargs["response_format"] = GPTKeywordExtractionFormat + + # Create the OpenAI client (supports both OpenAI and Azure) openai_async_client = create_openai_async_client( api_key=api_key, base_url=base_url, + use_azure=use_azure, + azure_deployment=azure_deployment, + api_version=api_version, + timeout=timeout, client_configs=client_configs, ) @@ -224,10 +266,16 @@ async def openai_complete_if_cache( messages = kwargs.pop("messages", messages) + # Add explicit parameters back to kwargs so they're passed to OpenAI API + if stream is not None: + kwargs["stream"] = stream + if timeout is not None: + kwargs["timeout"] = timeout + try: # Don't use async with context manager, use client directly if "response_format" in kwargs: - response = await openai_async_client.beta.chat.completions.parse( + response = await openai_async_client.chat.completions.parse( model=model, messages=messages, **kwargs ) else: @@ -439,46 +487,57 @@ async def openai_complete_if_cache( raise InvalidResponseError("Invalid response from OpenAI API") message = response.choices[0].message - content = getattr(message, "content", None) - reasoning_content = getattr(message, "reasoning_content", "") - # Handle COT logic for non-streaming responses (only if enabled) - final_content = "" + # Handle parsed responses (structured output via response_format) + # When using beta.chat.completions.parse(), the response is in message.parsed + if hasattr(message, "parsed") and message.parsed is not None: + # Serialize the parsed structured response to JSON + final_content = message.parsed.model_dump_json() + logger.debug("Using parsed structured response from API") + else: + # Handle regular content responses + content = getattr(message, "content", None) + reasoning_content = getattr(message, "reasoning_content", "") - if enable_cot: - # Check if we should include reasoning content - should_include_reasoning = False - if reasoning_content and reasoning_content.strip(): - if not content or content.strip() == "": - # Case 1: Only reasoning content, should include COT - should_include_reasoning = True - final_content = ( - content or "" - ) # Use empty string if content is None + # Handle COT logic for non-streaming responses (only if enabled) + final_content = "" + + if enable_cot: + # Check if we should include reasoning content + should_include_reasoning = False + if reasoning_content and reasoning_content.strip(): + if not content or content.strip() == "": + # Case 1: Only reasoning content, should include COT + should_include_reasoning = True + final_content = ( + content or "" + ) # Use empty string if content is None + else: + # Case 3: Both content and reasoning_content present, ignore reasoning + should_include_reasoning = False + final_content = content else: - # Case 3: Both content and reasoning_content present, ignore reasoning - should_include_reasoning = False - final_content = content + # No reasoning content, use regular content + final_content = content or "" + + # Apply COT wrapping if needed + if should_include_reasoning: + if r"\u" in reasoning_content: + reasoning_content = safe_unicode_decode( + reasoning_content.encode("utf-8") + ) + final_content = ( + f"{reasoning_content}{final_content}" + ) else: - # No reasoning content, use regular content + # COT disabled, only use regular content final_content = content or "" - # Apply COT wrapping if needed - if should_include_reasoning: - if r"\u" in reasoning_content: - reasoning_content = safe_unicode_decode( - reasoning_content.encode("utf-8") - ) - final_content = f"{reasoning_content}{final_content}" - else: - # COT disabled, only use regular content - final_content = content or "" - - # Validate final content - if not final_content or final_content.strip() == "": - logger.error("Received empty content from OpenAI API") - await openai_async_client.close() # Ensure client is closed - raise InvalidResponseError("Received empty content from OpenAI API") + # Validate final content + if not final_content or final_content.strip() == "": + logger.error("Received empty content from OpenAI API") + await openai_async_client.close() # Ensure client is closed + raise InvalidResponseError("Received empty content from OpenAI API") # Apply Unicode decoding to final content if needed if r"\u" in final_content: @@ -512,15 +571,13 @@ async def openai_complete( ) -> Union[str, AsyncIterator[str]]: if history_messages is None: history_messages = [] - keyword_extraction = kwargs.pop("keyword_extraction", None) - if keyword_extraction: - kwargs["response_format"] = "json" model_name = kwargs["hashing_kv"].global_config["llm_model_name"] return await openai_complete_if_cache( model_name, prompt, system_prompt=system_prompt, history_messages=history_messages, + keyword_extraction=keyword_extraction, **kwargs, ) @@ -535,15 +592,13 @@ async def gpt_4o_complete( ) -> str: if history_messages is None: history_messages = [] - keyword_extraction = kwargs.pop("keyword_extraction", None) - if keyword_extraction: - kwargs["response_format"] = GPTKeywordExtractionFormat return await openai_complete_if_cache( "gpt-4o", prompt, system_prompt=system_prompt, history_messages=history_messages, enable_cot=enable_cot, + keyword_extraction=keyword_extraction, **kwargs, ) @@ -558,15 +613,13 @@ async def gpt_4o_mini_complete( ) -> str: if history_messages is None: history_messages = [] - keyword_extraction = kwargs.pop("keyword_extraction", None) - if keyword_extraction: - kwargs["response_format"] = GPTKeywordExtractionFormat return await openai_complete_if_cache( "gpt-4o-mini", prompt, system_prompt=system_prompt, history_messages=history_messages, enable_cot=enable_cot, + keyword_extraction=keyword_extraction, **kwargs, ) @@ -581,13 +634,13 @@ async def nvidia_openai_complete( ) -> str: if history_messages is None: history_messages = [] - kwargs.pop("keyword_extraction", None) result = await openai_complete_if_cache( "nvidia/llama-3.1-nemotron-70b-instruct", # context length 128k prompt, system_prompt=system_prompt, history_messages=history_messages, enable_cot=enable_cot, + keyword_extraction=keyword_extraction, base_url="https://integrate.api.nvidia.com/v1", **kwargs, ) @@ -609,9 +662,12 @@ async def openai_embed( model: str = "text-embedding-3-small", base_url: str | None = None, api_key: str | None = None, - embedding_dim: int = None, + embedding_dim: int | None = None, client_configs: dict[str, Any] | None = None, token_tracker: Any | None = None, + use_azure: bool = False, + azure_deployment: str | None = None, + api_version: str | None = None, ) -> np.ndarray: """Generate embeddings for a list of texts using OpenAI's API. @@ -620,7 +676,12 @@ async def openai_embed( model: The OpenAI embedding model to use. base_url: Optional base URL for the OpenAI API. api_key: Optional OpenAI API key. If None, uses the OPENAI_API_KEY environment variable. - embedding_dim: Optional embedding dimension. If None, uses the default embedding dimension for the model. (will be passed to API for dimension reduction). + 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. + Manually passing a different value will trigger a warning and be ignored. + When provided (by EmbeddingFunc), it will be passed to the OpenAI API for dimension reduction. client_configs: Additional configuration options for the AsyncOpenAI client. These will override any default configurations but will be overridden by explicit parameters (api_key, base_url). @@ -634,9 +695,14 @@ async def openai_embed( RateLimitError: If the OpenAI API rate limit is exceeded. APITimeoutError: If the OpenAI API request times out. """ - # Create the OpenAI client + # Create the OpenAI client (supports both OpenAI and Azure) openai_async_client = create_openai_async_client( - api_key=api_key, base_url=base_url, client_configs=client_configs + api_key=api_key, + base_url=base_url, + use_azure=use_azure, + azure_deployment=azure_deployment, + api_version=api_version, + client_configs=client_configs, ) async with openai_async_client: @@ -669,3 +735,157 @@ async def openai_embed( for dp in response.data ] ) + + +# Azure OpenAI wrapper functions for backward compatibility +async def azure_openai_complete_if_cache( + model, + prompt, + 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, + api_version: str | None = None, + keyword_extraction: bool = False, + **kwargs, +): + """Azure OpenAI completion wrapper function. + + This function provides backward compatibility by wrapping the unified + openai_complete_if_cache implementation with Azure-specific parameter handling. + """ + # Handle Azure-specific environment variables and parameters + deployment = os.getenv("AZURE_OPENAI_DEPLOYMENT") or model or os.getenv("LLM_MODEL") + base_url = ( + base_url or os.getenv("AZURE_OPENAI_ENDPOINT") or os.getenv("LLM_BINDING_HOST") + ) + api_key = ( + api_key or os.getenv("AZURE_OPENAI_API_KEY") or os.getenv("LLM_BINDING_API_KEY") + ) + api_version = ( + api_version + or os.getenv("AZURE_OPENAI_API_VERSION") + or os.getenv("OPENAI_API_VERSION") + ) + + # Pop timeout from kwargs if present (will be handled by openai_complete_if_cache) + timeout = kwargs.pop("timeout", None) + + # Call the unified implementation with Azure-specific parameters + return await openai_complete_if_cache( + model=model, + prompt=prompt, + system_prompt=system_prompt, + history_messages=history_messages, + enable_cot=enable_cot, + base_url=base_url, + api_key=api_key, + timeout=timeout, + use_azure=True, + azure_deployment=deployment, + api_version=api_version, + keyword_extraction=keyword_extraction, + **kwargs, + ) + + +async def azure_openai_complete( + prompt, + system_prompt=None, + history_messages=None, + keyword_extraction=False, + **kwargs, +) -> str: + """Azure OpenAI complete wrapper function. + + Provides backward compatibility for azure_openai_complete calls. + """ + if history_messages is None: + history_messages = [] + result = await azure_openai_complete_if_cache( + os.getenv("LLM_MODEL", "gpt-4o-mini"), + prompt, + system_prompt=system_prompt, + history_messages=history_messages, + keyword_extraction=keyword_extraction, + **kwargs, + ) + return result + + +@wrap_embedding_func_with_attrs(embedding_dim=1536) +async def azure_openai_embed( + texts: list[str], + model: str | None = None, + base_url: str | None = None, + api_key: str | None = None, + api_version: str | None = None, +) -> np.ndarray: + """Azure OpenAI embedding wrapper function. + + This function provides backward compatibility by wrapping the unified + openai_embed implementation with Azure-specific parameter handling. + + IMPORTANT - Decorator Usage: + + 1. This function is decorated with @wrap_embedding_func_with_attrs to provide + the EmbeddingFunc interface for users who need to access embedding_dim + and other attributes. + + 2. This function does NOT use @retry decorator to avoid double-wrapping, + since the underlying openai_embed.func already has retry logic. + + 3. This function calls openai_embed.func (the unwrapped function) instead of + openai_embed (the EmbeddingFunc instance) to avoid double decoration issues: + + ✅ Correct: await openai_embed.func(...) # Calls unwrapped function with retry + ❌ Wrong: await openai_embed(...) # Would cause double EmbeddingFunc wrapping + + Double decoration causes: + - Double injection of embedding_dim parameter + - Incorrect parameter passing to the underlying implementation + - Runtime errors due to parameter conflicts + + The call chain with correct implementation: + azure_openai_embed(texts) + → EmbeddingFunc.__call__(texts) # azure's decorator + → azure_openai_embed_impl(texts, embedding_dim=1536) + → openai_embed.func(texts, ...) + → @retry_wrapper(texts, ...) # openai's retry (only one layer) + → openai_embed_impl(texts, ...) + → actual embedding computation + """ + # Handle Azure-specific environment variables and parameters + deployment = ( + os.getenv("AZURE_EMBEDDING_DEPLOYMENT") + or model + or os.getenv("EMBEDDING_MODEL", "text-embedding-3-small") + ) + base_url = ( + base_url + or os.getenv("AZURE_EMBEDDING_ENDPOINT") + or os.getenv("EMBEDDING_BINDING_HOST") + ) + api_key = ( + api_key + or os.getenv("AZURE_EMBEDDING_API_KEY") + or os.getenv("EMBEDDING_BINDING_API_KEY") + ) + api_version = ( + api_version + or os.getenv("AZURE_EMBEDDING_API_VERSION") + or os.getenv("OPENAI_API_VERSION") + ) + + # CRITICAL: Call openai_embed.func (unwrapped) to avoid double decoration + # openai_embed is an EmbeddingFunc instance, .func accesses the underlying function + return await openai_embed.func( + texts=texts, + model=model or deployment, + base_url=base_url, + api_key=api_key, + use_azure=True, + azure_deployment=deployment, + api_version=api_version, + ) diff --git a/lightrag/utils.py b/lightrag/utils.py index aa04f338..10b24da7 100644 --- a/lightrag/utils.py +++ b/lightrag/utils.py @@ -949,7 +949,76 @@ def priority_limit_async_func_call( def wrap_embedding_func_with_attrs(**kwargs): - """Wrap a function with attributes""" + """Decorator to add embedding dimension and token limit attributes to embedding functions. + + This decorator wraps an async embedding function and returns an EmbeddingFunc instance + that automatically handles dimension parameter injection and attribute management. + + WARNING: DO NOT apply this decorator to wrapper functions that call other + decorated embedding functions. This will cause double decoration and parameter + injection conflicts. + + Correct usage patterns: + + 1. Direct implementation (decorated): + ```python + @wrap_embedding_func_with_attrs(embedding_dim=1536) + async def my_embed(texts, embedding_dim=None): + # Direct implementation + return embeddings + ``` + + 2. Wrapper calling decorated function (DO NOT decorate wrapper): + ```python + # my_embed is already decorated above + + async def my_wrapper(texts, **kwargs): # ❌ DO NOT decorate this! + # Must call .func to access unwrapped implementation + return await my_embed.func(texts, **kwargs) + ``` + + 3. Wrapper calling decorated function (properly decorated): + ```python + @wrap_embedding_func_with_attrs(embedding_dim=1536) + async def my_wrapper(texts, **kwargs): # ✅ Can decorate if calling .func + # Calling .func avoids double decoration + return await my_embed.func(texts, **kwargs) + ``` + + The decorated function becomes an EmbeddingFunc instance with: + - embedding_dim: The embedding dimension + - max_token_size: Maximum token limit (optional) + - func: The original unwrapped function (access via .func) + - __call__: Wrapper that injects embedding_dim parameter + + Double decoration causes: + - Double injection of embedding_dim parameter + - Incorrect parameter passing to the underlying implementation + - Runtime errors due to parameter conflicts + + Args: + embedding_dim: The dimension of embedding vectors + max_token_size: Maximum number of tokens (optional) + send_dimensions: Whether to inject embedding_dim as a keyword argument (optional) + + Returns: + A decorator that wraps the function as an EmbeddingFunc instance + + Example of correct wrapper implementation: + ```python + @wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192) + @retry(...) + async def openai_embed(texts, ...): + # Base implementation + pass + + @wrap_embedding_func_with_attrs(embedding_dim=1536) # Note: No @retry here! + async def azure_openai_embed(texts, ...): + # CRITICAL: Call .func to access unwrapped function + return await openai_embed.func(texts, ...) # ✅ Correct + # return await openai_embed(texts, ...) # ❌ Wrong - double decoration! + ``` + """ def final_decro(func) -> EmbeddingFunc: new_func = EmbeddingFunc(**kwargs, func=func)