lix/mnote
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: init monorepo snapshot

Browse Source
This commit is contained in:
liaibo
2025-11-23 10:55:04 +08:00
commit c70ff52869
941 changed files with 246586 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
LightRAG/lightrag/llm/__init__.py
View File
LightRAG/lightrag/llm/anthropic.py
+333
View File
@@ -0,0 +1,333 @@
from ..utils import verbose_debug, VERBOSE_DEBUG
import sys
import os
import logging
import numpy as np
from typing import Any, Union, AsyncIterator
import pipmaster as pm # Pipmaster for dynamic library install
if sys.version_info < (3, 9):
from typing import AsyncIterator
else:
from collections.abc import AsyncIterator
# Install Anthropic SDK if not present
if not pm.is_installed("anthropic"):
pm.install("anthropic")
# Add Voyage AI import
if not pm.is_installed("voyageai"):
pm.install("voyageai")
import voyageai
from anthropic import (
AsyncAnthropic,
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.utils import (
safe_unicode_decode,
logger,
)
from lightrag.api import __api_version__
# Custom exception for retry mechanism
class InvalidResponseError(Exception):
"""Custom exception class for triggering retry mechanism"""
pass
# Core Anthropic completion function with retry
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError, InvalidResponseError)
),
)
async def anthropic_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,
**kwargs: Any,
) -> Union[str, AsyncIterator[str]]:
if history_messages is None:
history_messages = []
if enable_cot:
logger.debug(
"enable_cot=True is not supported for the Anthropic API and will be ignored."
)
if not api_key:
api_key = os.environ.get("ANTHROPIC_API_KEY")
default_headers = {
"User-Agent": f"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}",
"Content-Type": "application/json",
}
# Set logger level to INFO when VERBOSE_DEBUG is off
if not VERBOSE_DEBUG and logger.level == logging.DEBUG:
logging.getLogger("anthropic").setLevel(logging.INFO)
kwargs.pop("hashing_kv", None)
kwargs.pop("keyword_extraction", None)
timeout = kwargs.pop("timeout", None)
anthropic_async_client = (
AsyncAnthropic(
default_headers=default_headers, api_key=api_key, timeout=timeout
)
if base_url is None
else AsyncAnthropic(
base_url=base_url,
default_headers=default_headers,
api_key=api_key,
timeout=timeout,
)
)
messages: list[dict[str, Any]] = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.extend(history_messages)
messages.append({"role": "user", "content": prompt})
logger.debug("===== Sending Query to Anthropic LLM =====")
logger.debug(f"Model: {model} Base URL: {base_url}")
logger.debug(f"Additional kwargs: {kwargs}")
verbose_debug(f"Query: {prompt}")
verbose_debug(f"System prompt: {system_prompt}")
try:
response = await anthropic_async_client.messages.create(
model=model, messages=messages, stream=True, **kwargs
)
except APIConnectionError as e:
logger.error(f"Anthropic API Connection Error: {e}")
raise
except RateLimitError as e:
logger.error(f"Anthropic API Rate Limit Error: {e}")
raise
except APITimeoutError as e:
logger.error(f"Anthropic API Timeout Error: {e}")
raise
except Exception as e:
logger.error(
f"Anthropic API Call Failed,\nModel: {model},\nParams: {kwargs}, Got: {e}"
)
raise
async def stream_response():
try:
async for event in response:
content = (
event.delta.text
if hasattr(event, "delta") and event.delta.text
else None
)
if content is None:
continue
if r"\u" in content:
content = safe_unicode_decode(content.encode("utf-8"))
yield content
except Exception as e:
logger.error(f"Error in stream response: {str(e)}")
raise
return stream_response()
# Generic Anthropic completion function
async def anthropic_complete(
prompt: str,
system_prompt: str | None = None,
history_messages: list[dict[str, Any]] | None = None,
enable_cot: bool = False,
**kwargs: Any,
) -> Union[str, AsyncIterator[str]]:
if history_messages is None:
history_messages = []
model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
return await anthropic_complete_if_cache(
model_name,
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
# Claude 3 Opus specific completion
async def claude_3_opus_complete(
prompt: str,
system_prompt: str | None = None,
history_messages: list[dict[str, Any]] | None = None,
enable_cot: bool = False,
**kwargs: Any,
) -> Union[str, AsyncIterator[str]]:
if history_messages is None:
history_messages = []
return await anthropic_complete_if_cache(
"claude-3-opus-20240229",
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
# Claude 3 Sonnet specific completion
async def claude_3_sonnet_complete(
prompt: str,
system_prompt: str | None = None,
history_messages: list[dict[str, Any]] | None = None,
enable_cot: bool = False,
**kwargs: Any,
) -> Union[str, AsyncIterator[str]]:
if history_messages is None:
history_messages = []
return await anthropic_complete_if_cache(
"claude-3-sonnet-20240229",
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
# Claude 3 Haiku specific completion
async def claude_3_haiku_complete(
prompt: str,
system_prompt: str | None = None,
history_messages: list[dict[str, Any]] | None = None,
enable_cot: bool = False,
**kwargs: Any,
) -> Union[str, AsyncIterator[str]]:
if history_messages is None:
history_messages = []
return await anthropic_complete_if_cache(
"claude-3-haiku-20240307",
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
# Embedding function (placeholder, as Anthropic does not provide embeddings)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def anthropic_embed(
texts: list[str],
model: str = "voyage-3", # Default to voyage-3 as a good general-purpose model
base_url: str = None,
api_key: str = None,
) -> np.ndarray:
"""
Generate embeddings using Voyage AI since Anthropic doesn't provide native embedding support.
Args:
texts: List of text strings to embed
model: Voyage AI model name (e.g., "voyage-3", "voyage-3-large", "voyage-code-3")
base_url: Optional custom base URL (not used for Voyage AI)
api_key: API key for Voyage AI (defaults to VOYAGE_API_KEY environment variable)
Returns:
numpy array of shape (len(texts), embedding_dimension) containing the embeddings
"""
if not api_key:
api_key = os.environ.get("VOYAGE_API_KEY")
if not api_key:
logger.error("VOYAGE_API_KEY environment variable not set")
raise ValueError(
"VOYAGE_API_KEY environment variable is required for embeddings"
)
try:
# Initialize Voyage AI client
voyage_client = voyageai.Client(api_key=api_key)
# Get embeddings
result = voyage_client.embed(
texts,
model=model,
input_type="document", # Assuming document context; could be made configurable
)
# Convert list of embeddings to numpy array
embeddings = np.array(result.embeddings, dtype=np.float32)
logger.debug(f"Generated embeddings for {len(texts)} texts using {model}")
verbose_debug(f"Embedding shape: {embeddings.shape}")
return embeddings
except Exception as e:
logger.error(f"Voyage AI embedding failed: {str(e)}")
raise
# Optional: a helper function to get available embedding models
def get_available_embedding_models() -> dict[str, dict]:
"""
Returns a dictionary of available Voyage AI embedding models and their properties.
"""
return {
"voyage-3-large": {
"context_length": 32000,
"dimension": 1024,
"description": "Best general-purpose and multilingual",
},
"voyage-3": {
"context_length": 32000,
"dimension": 1024,
"description": "General-purpose and multilingual",
},
"voyage-3-lite": {
"context_length": 32000,
"dimension": 512,
"description": "Optimized for latency and cost",
},
"voyage-code-3": {
"context_length": 32000,
"dimension": 1024,
"description": "Optimized for code",
},
"voyage-finance-2": {
"context_length": 32000,
"dimension": 1024,
"description": "Optimized for finance",
},
"voyage-law-2": {
"context_length": 16000,
"dimension": 1024,
"description": "Optimized for legal",
},
"voyage-multimodal-3": {
"context_length": 32000,
"dimension": 1024,
"description": "Multimodal text and images",
},
}
LightRAG/lightrag/llm/azure_openai.py
+177
View File
@@ -0,0 +1,177 @@
from collections.abc import Iterable
import os
import pipmaster as pm # Pipmaster for dynamic library install
# install specific modules
if not pm.is_installed("openai"):
pm.install("openai")
from openai import (
AsyncAzureOpenAI,
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from openai.types.chat import ChatCompletionMessageParam
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.utils import (
wrap_embedding_func_with_attrs,
safe_unicode_decode,
logger,
)
import numpy as np
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APIConnectionError)
),
)
async def azure_openai_complete_if_cache(
model,
prompt,
system_prompt: str | None = None,
history_messages: Iterable[ChatCompletionMessageParam] | None = None,
enable_cot: bool = False,
base_url: str | None = None,
api_key: str | None = None,
api_version: str | None = None,
**kwargs,
):
if enable_cot:
logger.debug(
"enable_cot=True is not supported for the Azure OpenAI API and will be ignored."
)
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")
)
kwargs.pop("hashing_kv", None)
kwargs.pop("keyword_extraction", None)
timeout = kwargs.pop("timeout", None)
openai_async_client = AsyncAzureOpenAI(
azure_endpoint=base_url,
azure_deployment=deployment,
api_key=api_key,
api_version=api_version,
timeout=timeout,
)
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
if history_messages:
messages.extend(history_messages)
if prompt is not None:
messages.append({"role": "user", "content": prompt})
if "response_format" in kwargs:
response = await openai_async_client.beta.chat.completions.parse(
model=model, messages=messages, **kwargs
)
else:
response = await openai_async_client.chat.completions.create(
model=model, messages=messages, **kwargs
)
if hasattr(response, "__aiter__"):
async def inner():
async for chunk in response:
if len(chunk.choices) == 0:
continue
content = chunk.choices[0].delta.content
if content is None:
continue
if r"\u" in content:
content = safe_unicode_decode(content.encode("utf-8"))
yield content
return inner()
else:
content = response.choices[0].message.content
if r"\u" in content:
content = safe_unicode_decode(content.encode("utf-8"))
return content
async def azure_openai_complete(
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
kwargs.pop("keyword_extraction", None)
result = await azure_openai_complete_if_cache(
os.getenv("LLM_MODEL", "gpt-4o-mini"),
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
**kwargs,
)
return result
@wrap_embedding_func_with_attrs(embedding_dim=1536)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
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:
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")
)
openai_async_client = AsyncAzureOpenAI(
azure_endpoint=base_url,
azure_deployment=deployment,
api_key=api_key,
api_version=api_version,
)
response = await openai_async_client.embeddings.create(
model=model, input=texts, encoding_format="float"
)
return np.array([dp.embedding for dp in response.data])
LightRAG/lightrag/llm/bedrock.py
+483
View File
@@ -0,0 +1,483 @@
import copy
import os
import json
import logging
import pipmaster as pm # Pipmaster for dynamic library install
if not pm.is_installed("aioboto3"):
pm.install("aioboto3")
import aioboto3
import numpy as np
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
import sys
from lightrag.utils import wrap_embedding_func_with_attrs
if sys.version_info < (3, 9):
from typing import AsyncIterator
else:
from collections.abc import AsyncIterator
from typing import Union
# Import botocore exceptions for proper exception handling
try:
from botocore.exceptions import (
ClientError,
ConnectionError as BotocoreConnectionError,
ReadTimeoutError,
)
except ImportError:
# If botocore is not installed, define placeholders
ClientError = Exception
BotocoreConnectionError = Exception
ReadTimeoutError = Exception
class BedrockError(Exception):
"""Generic error for issues related to Amazon Bedrock"""
class BedrockRateLimitError(BedrockError):
"""Error for rate limiting and throttling issues"""
class BedrockConnectionError(BedrockError):
"""Error for network and connection issues"""
class BedrockTimeoutError(BedrockError):
"""Error for timeout issues"""
def _set_env_if_present(key: str, value):
"""Set environment variable only if a non-empty value is provided."""
if value is not None and value != "":
os.environ[key] = value
def _handle_bedrock_exception(e: Exception, operation: str = "Bedrock API") -> None:
"""Convert AWS Bedrock exceptions to appropriate custom exceptions.
Args:
e: The exception to handle
operation: Description of the operation for error messages
Raises:
BedrockRateLimitError: For rate limiting and throttling issues (retryable)
BedrockConnectionError: For network and server issues (retryable)
BedrockTimeoutError: For timeout issues (retryable)
BedrockError: For validation and other non-retryable errors
"""
error_message = str(e)
# Handle botocore ClientError with specific error codes
if isinstance(e, ClientError):
error_code = e.response.get("Error", {}).get("Code", "")
error_msg = e.response.get("Error", {}).get("Message", error_message)
# Rate limiting and throttling errors (retryable)
if error_code in [
"ThrottlingException",
"ProvisionedThroughputExceededException",
]:
logging.error(f"{operation} rate limit error: {error_msg}")
raise BedrockRateLimitError(f"Rate limit error: {error_msg}")
# Server errors (retryable)
elif error_code in ["ServiceUnavailableException", "InternalServerException"]:
logging.error(f"{operation} connection error: {error_msg}")
raise BedrockConnectionError(f"Service error: {error_msg}")
# Check for 5xx HTTP status codes (retryable)
elif e.response.get("ResponseMetadata", {}).get("HTTPStatusCode", 0) >= 500:
logging.error(f"{operation} server error: {error_msg}")
raise BedrockConnectionError(f"Server error: {error_msg}")
# Validation and other client errors (non-retryable)
else:
logging.error(f"{operation} client error: {error_msg}")
raise BedrockError(f"Client error: {error_msg}")
# Connection errors (retryable)
elif isinstance(e, BotocoreConnectionError):
logging.error(f"{operation} connection error: {error_message}")
raise BedrockConnectionError(f"Connection error: {error_message}")
# Timeout errors (retryable)
elif isinstance(e, (ReadTimeoutError, TimeoutError)):
logging.error(f"{operation} timeout error: {error_message}")
raise BedrockTimeoutError(f"Timeout error: {error_message}")
# Custom Bedrock errors (already properly typed)
elif isinstance(
e,
(
BedrockRateLimitError,
BedrockConnectionError,
BedrockTimeoutError,
BedrockError,
),
):
raise
# Unknown errors (non-retryable)
else:
logging.error(f"{operation} unexpected error: {error_message}")
raise BedrockError(f"Unexpected error: {error_message}")
@retry(
stop=stop_after_attempt(5),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=(
retry_if_exception_type(BedrockRateLimitError)
| retry_if_exception_type(BedrockConnectionError)
| retry_if_exception_type(BedrockTimeoutError)
),
)
async def bedrock_complete_if_cache(
model,
prompt,
system_prompt=None,
history_messages=[],
enable_cot: bool = False,
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
**kwargs,
) -> Union[str, AsyncIterator[str]]:
if enable_cot:
import logging
logging.debug(
"enable_cot=True is not supported for Bedrock and will be ignored."
)
# Respect existing env; only set if a non-empty value is available
access_key = os.environ.get("AWS_ACCESS_KEY_ID") or aws_access_key_id
secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY") or aws_secret_access_key
session_token = os.environ.get("AWS_SESSION_TOKEN") or aws_session_token
_set_env_if_present("AWS_ACCESS_KEY_ID", access_key)
_set_env_if_present("AWS_SECRET_ACCESS_KEY", secret_key)
_set_env_if_present("AWS_SESSION_TOKEN", session_token)
# Region handling: prefer env, else kwarg (optional)
region = os.environ.get("AWS_REGION") or kwargs.pop("aws_region", None)
kwargs.pop("hashing_kv", None)
# Capture stream flag (if provided) and remove from kwargs since it's not a Bedrock API parameter
# We'll use this to determine whether to call converse_stream or converse
stream = bool(kwargs.pop("stream", False))
# Remove unsupported args for Bedrock Converse API
for k in [
"response_format",
"tools",
"tool_choice",
"seed",
"presence_penalty",
"frequency_penalty",
"n",
"logprobs",
"top_logprobs",
"max_completion_tokens",
"response_format",
]:
kwargs.pop(k, None)
# Fix message history format
messages = []
for history_message in history_messages:
message = copy.copy(history_message)
message["content"] = [{"text": message["content"]}]
messages.append(message)
# Add user prompt
messages.append({"role": "user", "content": [{"text": prompt}]})
# Initialize Converse API arguments
args = {"modelId": model, "messages": messages}
# Define system prompt
if system_prompt:
args["system"] = [{"text": system_prompt}]
# Map and set up inference parameters
inference_params_map = {
"max_tokens": "maxTokens",
"top_p": "topP",
"stop_sequences": "stopSequences",
}
if inference_params := list(
set(kwargs) & set(["max_tokens", "temperature", "top_p", "stop_sequences"])
):
args["inferenceConfig"] = {}
for param in inference_params:
args["inferenceConfig"][inference_params_map.get(param, param)] = (
kwargs.pop(param)
)
# Import logging for error handling
import logging
# For streaming responses, we need a different approach to keep the connection open
if stream:
# Create a session that will be used throughout the streaming process
session = aioboto3.Session()
client = None
# Define the generator function that will manage the client lifecycle
async def stream_generator():
nonlocal client
# Create the client outside the generator to ensure it stays open
client = await session.client(
"bedrock-runtime", region_name=region
).__aenter__()
event_stream = None
iteration_started = False
try:
# Make the API call
response = await client.converse_stream(**args, **kwargs)
event_stream = response.get("stream")
iteration_started = True
# Process the stream
async for event in event_stream:
# Validate event structure
if not event or not isinstance(event, dict):
continue
if "contentBlockDelta" in event:
delta = event["contentBlockDelta"].get("delta", {})
text = delta.get("text")
if text:
yield text
# Handle other event types that might indicate stream end
elif "messageStop" in event:
break
except Exception as e:
# Try to clean up resources if possible
if (
iteration_started
and event_stream
and hasattr(event_stream, "aclose")
and callable(getattr(event_stream, "aclose", None))
):
try:
await event_stream.aclose()
except Exception as close_error:
logging.warning(
f"Failed to close Bedrock event stream: {close_error}"
)
# Convert to appropriate exception type
_handle_bedrock_exception(e, "Bedrock streaming")
finally:
# Clean up the event stream
if (
iteration_started
and event_stream
and hasattr(event_stream, "aclose")
and callable(getattr(event_stream, "aclose", None))
):
try:
await event_stream.aclose()
except Exception as close_error:
logging.warning(
f"Failed to close Bedrock event stream in finally block: {close_error}"
)
# Clean up the client
if client:
try:
await client.__aexit__(None, None, None)
except Exception as client_close_error:
logging.warning(
f"Failed to close Bedrock client: {client_close_error}"
)
# Return the generator that manages its own lifecycle
return stream_generator()
# For non-streaming responses, use the standard async context manager pattern
session = aioboto3.Session()
async with session.client(
"bedrock-runtime", region_name=region
) as bedrock_async_client:
try:
# Use converse for non-streaming responses
response = await bedrock_async_client.converse(**args, **kwargs)
# Validate response structure
if (
not response
or "output" not in response
or "message" not in response["output"]
or "content" not in response["output"]["message"]
or not response["output"]["message"]["content"]
):
raise BedrockError("Invalid response structure from Bedrock API")
content = response["output"]["message"]["content"][0]["text"]
if not content or content.strip() == "":
raise BedrockError("Received empty content from Bedrock API")
return content
except Exception as e:
# Convert to appropriate exception type
_handle_bedrock_exception(e, "Bedrock converse")
# Generic Bedrock completion function
async def bedrock_complete(
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> Union[str, AsyncIterator[str]]:
kwargs.pop("keyword_extraction", None)
model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
result = await bedrock_complete_if_cache(
model_name,
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
**kwargs,
)
return result
@wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
@retry(
stop=stop_after_attempt(5),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=(
retry_if_exception_type(BedrockRateLimitError)
| retry_if_exception_type(BedrockConnectionError)
| retry_if_exception_type(BedrockTimeoutError)
),
)
async def bedrock_embed(
texts: list[str],
model: str = "amazon.titan-embed-text-v2:0",
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
) -> np.ndarray:
# Respect existing env; only set if a non-empty value is available
access_key = os.environ.get("AWS_ACCESS_KEY_ID") or aws_access_key_id
secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY") or aws_secret_access_key
session_token = os.environ.get("AWS_SESSION_TOKEN") or aws_session_token
_set_env_if_present("AWS_ACCESS_KEY_ID", access_key)
_set_env_if_present("AWS_SECRET_ACCESS_KEY", secret_key)
_set_env_if_present("AWS_SESSION_TOKEN", session_token)
# Region handling: prefer env
region = os.environ.get("AWS_REGION")
session = aioboto3.Session()
async with session.client(
"bedrock-runtime", region_name=region
) as bedrock_async_client:
try:
if (model_provider := model.split(".")[0]) == "amazon":
embed_texts = []
for text in texts:
try:
if "v2" in model:
body = json.dumps(
{
"inputText": text,
# 'dimensions': embedding_dim,
"embeddingTypes": ["float"],
}
)
elif "v1" in model:
body = json.dumps({"inputText": text})
else:
raise BedrockError(f"Model {model} is not supported!")
response = await bedrock_async_client.invoke_model(
modelId=model,
body=body,
accept="application/json",
contentType="application/json",
)
response_body = await response.get("body").json()
# Validate response structure
if not response_body or "embedding" not in response_body:
raise BedrockError(
f"Invalid embedding response structure for text: {text[:50]}..."
)
embedding = response_body["embedding"]
if not embedding:
raise BedrockError(
f"Received empty embedding for text: {text[:50]}..."
)
embed_texts.append(embedding)
except Exception as e:
# Convert to appropriate exception type
_handle_bedrock_exception(
e, "Bedrock embedding (amazon, text chunk)"
)
elif model_provider == "cohere":
try:
body = json.dumps(
{
"texts": texts,
"input_type": "search_document",
"truncate": "NONE",
}
)
response = await bedrock_async_client.invoke_model(
model=model,
body=body,
accept="application/json",
contentType="application/json",
)
response_body = json.loads(response.get("body").read())
# Validate response structure
if not response_body or "embeddings" not in response_body:
raise BedrockError(
"Invalid embedding response structure from Cohere"
)
embeddings = response_body["embeddings"]
if not embeddings or len(embeddings) != len(texts):
raise BedrockError(
f"Invalid embeddings count: expected {len(texts)}, got {len(embeddings) if embeddings else 0}"
)
embed_texts = embeddings
except Exception as e:
# Convert to appropriate exception type
_handle_bedrock_exception(e, "Bedrock embedding (cohere)")
else:
raise BedrockError(
f"Model provider '{model_provider}' is not supported!"
)
# Final validation
if not embed_texts:
raise BedrockError("No embeddings generated")
return np.array(embed_texts)
except Exception as e:
# Convert to appropriate exception type
_handle_bedrock_exception(e, "Bedrock embedding")
LightRAG/lightrag/llm/binding_options.py
+740
View File
@@ -0,0 +1,740 @@
"""
Module that implements containers for specific LLM bindings.
This module provides container implementations for various Large Language Model
bindings and integrations.
"""
from argparse import ArgumentParser, Namespace
import argparse
import json
from dataclasses import asdict, dataclass, field
from typing import Any, ClassVar, List, get_args, get_origin
from lightrag.utils import get_env_value
from lightrag.constants import DEFAULT_TEMPERATURE
def _resolve_optional_type(field_type: Any) -> Any:
"""Return the concrete type for Optional/Union annotations."""
origin = get_origin(field_type)
if origin in (list, dict, tuple):
return field_type
args = get_args(field_type)
if args:
non_none_args = [arg for arg in args if arg is not type(None)]
if len(non_none_args) == 1:
return non_none_args[0]
return field_type
# =============================================================================
# BindingOptions Base Class
# =============================================================================
#
# The BindingOptions class serves as the foundation for all LLM provider bindings
# in LightRAG. It provides a standardized framework for:
#
# 1. Configuration Management:
# - Defines how each LLM provider's configuration parameters are structured
# - Handles default values and type information for each parameter
# - Maps configuration options to command-line arguments and environment variables
#
# 2. Environment Integration:
# - Automatically generates environment variable names from binding parameters
# - Provides methods to create sample .env files for easy configuration
# - Supports configuration via environment variables with fallback to defaults
#
# 3. Command-Line Interface:
# - Dynamically generates command-line arguments for all registered bindings
# - Maintains consistent naming conventions across different LLM providers
# - Provides help text and type validation for each configuration option
#
# 4. Extensibility:
# - Uses class introspection to automatically discover all binding subclasses
# - Requires minimal boilerplate code when adding new LLM provider bindings
# - Maintains separation of concerns between different provider configurations
#
# This design pattern ensures that adding support for a new LLM provider requires
# only defining the provider-specific parameters and help text, while the base
# class handles all the common functionality for argument parsing, environment
# variable handling, and configuration management.
#
# Instances of a derived class of BindingOptions can be used to store multiple
# runtime configurations of options for a single LLM provider. using the
# asdict() method to convert the options to a dictionary.
#
# =============================================================================
@dataclass
class BindingOptions:
"""Base class for binding options."""
# mandatory name of binding
_binding_name: ClassVar[str]
# optional help message for each option
_help: ClassVar[dict[str, str]]
@staticmethod
def _all_class_vars(klass: type, include_inherited=True) -> dict[str, Any]:
"""Print class variables, optionally including inherited ones"""
if include_inherited:
# Get all class variables from MRO
vars_dict = {}
for base in reversed(klass.__mro__[:-1]): # Exclude 'object'
vars_dict.update(
{
k: v
for k, v in base.__dict__.items()
if (
not k.startswith("_")
and not callable(v)
and not isinstance(v, classmethod)
)
}
)
else:
# Only direct class variables
vars_dict = {
k: v
for k, v in klass.__dict__.items()
if (
not k.startswith("_")
and not callable(v)
and not isinstance(v, classmethod)
)
}
return vars_dict
@classmethod
def add_args(cls, parser: ArgumentParser):
group = parser.add_argument_group(f"{cls._binding_name} binding options")
for arg_item in cls.args_env_name_type_value():
# Handle JSON parsing for list types
if arg_item["type"] is List[str]:
def json_list_parser(value):
try:
parsed = json.loads(value)
if not isinstance(parsed, list):
raise argparse.ArgumentTypeError(
f"Expected JSON array, got {type(parsed).__name__}"
)
return parsed
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"Invalid JSON: {e}")
# Get environment variable with JSON parsing
env_value = get_env_value(f"{arg_item['env_name']}", argparse.SUPPRESS)
if env_value is not argparse.SUPPRESS:
try:
env_value = json_list_parser(env_value)
except argparse.ArgumentTypeError:
env_value = argparse.SUPPRESS
group.add_argument(
f"--{arg_item['argname']}",
type=json_list_parser,
default=env_value,
help=arg_item["help"],
)
# Handle JSON parsing for dict types
elif arg_item["type"] is dict:
def json_dict_parser(value):
try:
parsed = json.loads(value)
if not isinstance(parsed, dict):
raise argparse.ArgumentTypeError(
f"Expected JSON object, got {type(parsed).__name__}"
)
return parsed
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"Invalid JSON: {e}")
# Get environment variable with JSON parsing
env_value = get_env_value(f"{arg_item['env_name']}", argparse.SUPPRESS)
if env_value is not argparse.SUPPRESS:
try:
env_value = json_dict_parser(env_value)
except argparse.ArgumentTypeError:
env_value = argparse.SUPPRESS
group.add_argument(
f"--{arg_item['argname']}",
type=json_dict_parser,
default=env_value,
help=arg_item["help"],
)
# Handle boolean types specially to avoid argparse bool() constructor issues
elif arg_item["type"] is bool:
def bool_parser(value):
"""Custom boolean parser that handles string representations correctly"""
if isinstance(value, bool):
return value
if isinstance(value, str):
return value.lower() in ("true", "1", "yes", "t", "on")
return bool(value)
# Get environment variable with proper type conversion
env_value = get_env_value(
f"{arg_item['env_name']}", argparse.SUPPRESS, bool
)
group.add_argument(
f"--{arg_item['argname']}",
type=bool_parser,
default=env_value,
help=arg_item["help"],
)
else:
resolved_type = arg_item["type"]
if resolved_type is not None:
resolved_type = _resolve_optional_type(resolved_type)
group.add_argument(
f"--{arg_item['argname']}",
type=resolved_type,
default=get_env_value(f"{arg_item['env_name']}", argparse.SUPPRESS),
help=arg_item["help"],
)
@classmethod
def args_env_name_type_value(cls):
import dataclasses
args_prefix = f"{cls._binding_name}".replace("_", "-")
env_var_prefix = f"{cls._binding_name}_".upper()
help = cls._help
# Check if this is a dataclass and use dataclass fields
if dataclasses.is_dataclass(cls):
for field in dataclasses.fields(cls):
# Skip private fields
if field.name.startswith("_"):
continue
# Get default value
if field.default is not dataclasses.MISSING:
default_value = field.default
elif field.default_factory is not dataclasses.MISSING:
default_value = field.default_factory()
else:
default_value = None
argdef = {
"argname": f"{args_prefix}-{field.name}",
"env_name": f"{env_var_prefix}{field.name.upper()}",
"type": _resolve_optional_type(field.type),
"default": default_value,
"help": f"{cls._binding_name} -- " + help.get(field.name, ""),
}
yield argdef
else:
# Fallback to old method for non-dataclass classes
class_vars = {
key: value
for key, value in cls._all_class_vars(cls).items()
if not callable(value) and not key.startswith("_")
}
# Get type hints to properly detect List[str] types
type_hints = {}
for base in cls.__mro__:
if hasattr(base, "__annotations__"):
type_hints.update(base.__annotations__)
for class_var in class_vars:
# Use type hint if available, otherwise fall back to type of value
var_type = type_hints.get(class_var, type(class_vars[class_var]))
argdef = {
"argname": f"{args_prefix}-{class_var}",
"env_name": f"{env_var_prefix}{class_var.upper()}",
"type": var_type,
"default": class_vars[class_var],
"help": f"{cls._binding_name} -- " + help.get(class_var, ""),
}
yield argdef
@classmethod
def generate_dot_env_sample(cls):
"""
Generate a sample .env file for all LightRAG binding options.
This method creates a .env file that includes all the binding options
defined by the subclasses of BindingOptions. It uses the args_env_name_type_value()
method to get the list of all options and their default values.
Returns:
str: A string containing the contents of the sample .env file.
"""
from io import StringIO
sample_top = (
"#" * 80
+ "\n"
+ (
"# Autogenerated .env entries list for LightRAG binding options\n"
"#\n"
"# To generate run:\n"
"# $ python -m lightrag.llm.binding_options\n"
)
+ "#" * 80
+ "\n"
)
sample_bottom = (
("#\n# End of .env entries for LightRAG binding options\n")
+ "#" * 80
+ "\n"
)
sample_stream = StringIO()
sample_stream.write(sample_top)
for klass in cls.__subclasses__():
for arg_item in klass.args_env_name_type_value():
if arg_item["help"]:
sample_stream.write(f"# {arg_item['help']}\n")
# Handle JSON formatting for list and dict types
if arg_item["type"] is List[str] or arg_item["type"] is dict:
default_value = json.dumps(arg_item["default"])
else:
default_value = arg_item["default"]
sample_stream.write(f"# {arg_item['env_name']}={default_value}\n\n")
sample_stream.write(sample_bottom)
return sample_stream.getvalue()
@classmethod
def options_dict(cls, args: Namespace) -> dict[str, Any]:
"""
Extract options dictionary for a specific binding from parsed arguments.
This method filters the parsed command-line arguments to return only those
that belong to the specific binding class. It removes the binding prefix
from argument names to create a clean options dictionary.
Args:
args (Namespace): Parsed command-line arguments containing all binding options
Returns:
dict[str, Any]: Dictionary mapping option names (without prefix) to their values
Example:
If args contains {'ollama_num_ctx': 512, 'other_option': 'value'}
and this is called on OllamaOptions, it returns {'num_ctx': 512}
"""
prefix = cls._binding_name + "_"
skipchars = len(prefix)
options = {
key[skipchars:]: value
for key, value in vars(args).items()
if key.startswith(prefix)
}
return options
def asdict(self) -> dict[str, Any]:
"""
Convert an instance of binding options to a dictionary.
This method uses dataclasses.asdict() to convert the dataclass instance
into a dictionary representation, including all its fields and values.
Returns:
dict[str, Any]: Dictionary representation of the binding options instance
"""
return asdict(self)
# =============================================================================
# Binding Options for Ollama
# =============================================================================
#
# Ollama binding options provide configuration for the Ollama local LLM server.
# These options control model behavior, sampling parameters, hardware utilization,
# and performance settings. The parameters are based on Ollama's API specification
# and provide fine-grained control over model inference and generation.
#
# The _OllamaOptionsMixin defines the complete set of available options, while
# OllamaEmbeddingOptions and OllamaLLMOptions provide specialized configurations
# for embedding and language model tasks respectively.
# =============================================================================
@dataclass
class _OllamaOptionsMixin:
"""Options for Ollama bindings."""
# Core context and generation parameters
num_ctx: int = 32768 # Context window size (number of tokens)
num_predict: int = 128 # Maximum number of tokens to predict
num_keep: int = 0 # Number of tokens to keep from the initial prompt
seed: int = -1 # Random seed for generation (-1 for random)
# Sampling parameters
temperature: float = DEFAULT_TEMPERATURE # Controls randomness (0.0-2.0)
top_k: int = 40 # Top-k sampling parameter
top_p: float = 0.9 # Top-p (nucleus) sampling parameter
tfs_z: float = 1.0 # Tail free sampling parameter
typical_p: float = 1.0 # Typical probability mass
min_p: float = 0.0 # Minimum probability threshold
# Repetition control
repeat_last_n: int = 64 # Number of tokens to consider for repetition penalty
repeat_penalty: float = 1.1 # Penalty for repetition
presence_penalty: float = 0.0 # Penalty for token presence
frequency_penalty: float = 0.0 # Penalty for token frequency
# Mirostat sampling
mirostat: int = (
# Mirostat sampling algorithm (0=disabled, 1=Mirostat 1.0, 2=Mirostat 2.0)
0
)
mirostat_tau: float = 5.0 # Mirostat target entropy
mirostat_eta: float = 0.1 # Mirostat learning rate
# Hardware and performance parameters
numa: bool = False # Enable NUMA optimization
num_batch: int = 512 # Batch size for processing
num_gpu: int = -1 # Number of GPUs to use (-1 for auto)
main_gpu: int = 0 # Main GPU index
low_vram: bool = False # Optimize for low VRAM
num_thread: int = 0 # Number of CPU threads (0 for auto)
# Memory and model parameters
f16_kv: bool = True # Use half-precision for key/value cache
logits_all: bool = False # Return logits for all tokens
vocab_only: bool = False # Only load vocabulary
use_mmap: bool = True # Use memory mapping for model files
use_mlock: bool = False # Lock model in memory
embedding_only: bool = False # Only use for embeddings
# Output control
penalize_newline: bool = True # Penalize newline tokens
stop: List[str] = field(default_factory=list) # Stop sequences
# optional help strings
_help: ClassVar[dict[str, str]] = {
"num_ctx": "Context window size (number of tokens)",
"num_predict": "Maximum number of tokens to predict",
"num_keep": "Number of tokens to keep from the initial prompt",
"seed": "Random seed for generation (-1 for random)",
"temperature": "Controls randomness (0.0-2.0, higher = more creative)",
"top_k": "Top-k sampling parameter (0 = disabled)",
"top_p": "Top-p (nucleus) sampling parameter (0.0-1.0)",
"tfs_z": "Tail free sampling parameter (1.0 = disabled)",
"typical_p": "Typical probability mass (1.0 = disabled)",
"min_p": "Minimum probability threshold (0.0 = disabled)",
"repeat_last_n": "Number of tokens to consider for repetition penalty",
"repeat_penalty": "Penalty for repetition (1.0 = no penalty)",
"presence_penalty": "Penalty for token presence (-2.0 to 2.0)",
"frequency_penalty": "Penalty for token frequency (-2.0 to 2.0)",
"mirostat": "Mirostat sampling algorithm (0=disabled, 1=Mirostat 1.0, 2=Mirostat 2.0)",
"mirostat_tau": "Mirostat target entropy",
"mirostat_eta": "Mirostat learning rate",
"numa": "Enable NUMA optimization",
"num_batch": "Batch size for processing",
"num_gpu": "Number of GPUs to use (-1 for auto)",
"main_gpu": "Main GPU index",
"low_vram": "Optimize for low VRAM",
"num_thread": "Number of CPU threads (0 for auto)",
"f16_kv": "Use half-precision for key/value cache",
"logits_all": "Return logits for all tokens",
"vocab_only": "Only load vocabulary",
"use_mmap": "Use memory mapping for model files",
"use_mlock": "Lock model in memory",
"embedding_only": "Only use for embeddings",
"penalize_newline": "Penalize newline tokens",
"stop": 'Stop sequences (JSON array of strings, e.g., \'["</s>", "\\n\\n"]\')',
}
@dataclass
class OllamaEmbeddingOptions(_OllamaOptionsMixin, BindingOptions):
"""Options for Ollama embeddings with specialized configuration for embedding tasks."""
# mandatory name of binding
_binding_name: ClassVar[str] = "ollama_embedding"
@dataclass
class OllamaLLMOptions(_OllamaOptionsMixin, BindingOptions):
"""Options for Ollama LLM with specialized configuration for LLM tasks."""
# mandatory name of binding
_binding_name: ClassVar[str] = "ollama_llm"
# =============================================================================
# Binding Options for Gemini
# =============================================================================
@dataclass
class GeminiLLMOptions(BindingOptions):
"""Options for Google Gemini models."""
_binding_name: ClassVar[str] = "gemini_llm"
temperature: float = DEFAULT_TEMPERATURE
top_p: float = 0.95
top_k: int = 40
max_output_tokens: int | None = None
candidate_count: int = 1
presence_penalty: float = 0.0
frequency_penalty: float = 0.0
stop_sequences: List[str] = field(default_factory=list)
seed: int | None = None
thinking_config: dict | None = None
safety_settings: dict | None = None
_help: ClassVar[dict[str, str]] = {
"temperature": "Controls randomness (0.0-2.0, higher = more creative)",
"top_p": "Nucleus sampling parameter (0.0-1.0)",
"top_k": "Limits sampling to the top K tokens (1 disables the limit)",
"max_output_tokens": "Maximum tokens generated in the response",
"candidate_count": "Number of candidates returned per request",
"presence_penalty": "Penalty for token presence (-2.0 to 2.0)",
"frequency_penalty": "Penalty for token frequency (-2.0 to 2.0)",
"stop_sequences": "Stop sequences (JSON array of strings, e.g., '[\"END\"]')",
"seed": "Random seed for reproducible generation (leave empty for random)",
"thinking_config": "Thinking configuration (JSON dict, e.g., '{\"thinking_budget\": 1024}' or '{\"include_thoughts\": true}')",
"safety_settings": "JSON object with Gemini safety settings overrides",
}
@dataclass
class GeminiEmbeddingOptions(BindingOptions):
"""Options for Google Gemini embedding models."""
_binding_name: ClassVar[str] = "gemini_embedding"
task_type: str = "RETRIEVAL_DOCUMENT"
_help: ClassVar[dict[str, str]] = {
"task_type": "Task type for embedding optimization (RETRIEVAL_DOCUMENT, RETRIEVAL_QUERY, SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING, CODE_RETRIEVAL_QUERY, QUESTION_ANSWERING, FACT_VERIFICATION)",
}
# =============================================================================
# Binding Options for OpenAI
# =============================================================================
#
# OpenAI binding options provide configuration for OpenAI's API and Azure OpenAI.
# These options control model behavior, sampling parameters, and generation settings.
# The parameters are based on OpenAI's API specification and provide fine-grained
# control over model inference and generation.
#
# =============================================================================
@dataclass
class OpenAILLMOptions(BindingOptions):
"""Options for OpenAI LLM with configuration for OpenAI and Azure OpenAI API calls."""
# mandatory name of binding
_binding_name: ClassVar[str] = "openai_llm"
# Sampling and generation parameters
frequency_penalty: float = 0.0 # Penalty for token frequency (-2.0 to 2.0)
max_completion_tokens: int = None # Maximum number of tokens to generate
presence_penalty: float = 0.0 # Penalty for token presence (-2.0 to 2.0)
reasoning_effort: str = "medium" # Reasoning effort level (low, medium, high)
safety_identifier: str = "" # Safety identifier for content filtering
service_tier: str = "" # Service tier for API usage
stop: List[str] = field(default_factory=list) # Stop sequences
temperature: float = DEFAULT_TEMPERATURE # Controls randomness (0.0 to 2.0)
top_p: float = 1.0 # Nucleus sampling parameter (0.0 to 1.0)
max_tokens: int = None # Maximum number of tokens to generate(deprecated, use max_completion_tokens instead)
extra_body: dict = None # Extra body parameters for OpenRouter of vLLM
# Help descriptions
_help: ClassVar[dict[str, str]] = {
"frequency_penalty": "Penalty for token frequency (-2.0 to 2.0, positive values discourage repetition)",
"max_completion_tokens": "Maximum number of tokens to generate (optional, leave empty for model default)",
"presence_penalty": "Penalty for token presence (-2.0 to 2.0, positive values encourage new topics)",
"reasoning_effort": "Reasoning effort level for o1 models (low, medium, high)",
"safety_identifier": "Safety identifier for content filtering (optional)",
"service_tier": "Service tier for API usage (optional)",
"stop": 'Stop sequences (JSON array of strings, e.g., \'["</s>", "\\n\\n"]\')',
"temperature": "Controls randomness (0.0-2.0, higher = more creative)",
"top_p": "Nucleus sampling parameter (0.0-1.0, lower = more focused)",
"max_tokens": "Maximum number of tokens to generate (deprecated, use max_completion_tokens instead)",
"extra_body": 'Extra body parameters for OpenRouter of vLLM (JSON dict, e.g., \'"reasoning": {"reasoning": {"enabled": false}}\')',
}
# =============================================================================
# Main Section - For Testing and Sample Generation
# =============================================================================
#
# When run as a script, this module:
# 1. Generates and prints a sample .env file with all binding options
# 2. If "test" argument is provided, demonstrates argument parsing with Ollama binding
#
# Usage:
# python -m lightrag.llm.binding_options # Generate .env sample
# python -m lightrag.llm.binding_options test # Test argument parsing
#
# =============================================================================
if __name__ == "__main__":
import sys
import dotenv
# from io import StringIO
dotenv.load_dotenv(dotenv_path=".env", override=False)
# env_strstream = StringIO(
# ("OLLAMA_LLM_TEMPERATURE=0.1\nOLLAMA_EMBEDDING_TEMPERATURE=0.2\n")
# )
# # Load environment variables from .env file
# dotenv.load_dotenv(stream=env_strstream)
if len(sys.argv) > 1 and sys.argv[1] == "test":
# Add arguments for OllamaEmbeddingOptions, OllamaLLMOptions, and OpenAILLMOptions
parser = ArgumentParser(description="Test binding options")
OllamaEmbeddingOptions.add_args(parser)
OllamaLLMOptions.add_args(parser)
OpenAILLMOptions.add_args(parser)
# Parse arguments test
args = parser.parse_args(
[
"--ollama-embedding-num_ctx",
"1024",
"--ollama-llm-num_ctx",
"2048",
"--openai-llm-temperature",
"0.7",
"--openai-llm-max_completion_tokens",
"1000",
"--openai-llm-stop",
'["</s>", "\\n\\n"]',
"--openai-llm-reasoning",
'{"effort": "high", "max_tokens": 2000, "exclude": false, "enabled": true}',
]
)
print("Final args for LLM and Embedding:")
print(f"{args}\n")
print("Ollama LLM options:")
print(OllamaLLMOptions.options_dict(args))
print("\nOllama Embedding options:")
print(OllamaEmbeddingOptions.options_dict(args))
print("\nOpenAI LLM options:")
print(OpenAILLMOptions.options_dict(args))
# Test creating OpenAI options instance
openai_options = OpenAILLMOptions(
temperature=0.8,
max_completion_tokens=1500,
frequency_penalty=0.1,
presence_penalty=0.2,
stop=["<|end|>", "\n\n"],
)
print("\nOpenAI LLM options instance:")
print(openai_options.asdict())
# Test creating OpenAI options instance with reasoning parameter
openai_options_with_reasoning = OpenAILLMOptions(
temperature=0.9,
max_completion_tokens=2000,
reasoning={
"effort": "medium",
"max_tokens": 1500,
"exclude": True,
"enabled": True,
},
)
print("\nOpenAI LLM options instance with reasoning:")
print(openai_options_with_reasoning.asdict())
# Test dict parsing functionality
print("\n" + "=" * 50)
print("TESTING DICT PARSING FUNCTIONALITY")
print("=" * 50)
# Test valid JSON dict parsing
test_parser = ArgumentParser(description="Test dict parsing")
OpenAILLMOptions.add_args(test_parser)
try:
test_args = test_parser.parse_args(
["--openai-llm-reasoning", '{"effort": "low", "max_tokens": 1000}']
)
print("✓ Valid JSON dict parsing successful:")
print(
f" Parsed reasoning: {OpenAILLMOptions.options_dict(test_args)['reasoning']}"
)
except Exception as e:
print(f"✗ Valid JSON dict parsing failed: {e}")
# Test invalid JSON dict parsing
try:
test_args = test_parser.parse_args(
[
"--openai-llm-reasoning",
'{"effort": "low", "max_tokens": 1000', # Missing closing brace
]
)
print("✗ Invalid JSON should have failed but didn't")
except SystemExit:
print("✓ Invalid JSON dict parsing correctly rejected")
except Exception as e:
print(f"✓ Invalid JSON dict parsing correctly rejected: {e}")
# Test non-dict JSON parsing
try:
test_args = test_parser.parse_args(
[
"--openai-llm-reasoning",
'["not", "a", "dict"]', # Array instead of dict
]
)
print("✗ Non-dict JSON should have failed but didn't")
except SystemExit:
print("✓ Non-dict JSON parsing correctly rejected")
except Exception as e:
print(f"✓ Non-dict JSON parsing correctly rejected: {e}")
print("\n" + "=" * 50)
print("TESTING ENVIRONMENT VARIABLE SUPPORT")
print("=" * 50)
# Test environment variable support for dict
import os
os.environ["OPENAI_LLM_REASONING"] = (
'{"effort": "high", "max_tokens": 3000, "exclude": false}'
)
env_parser = ArgumentParser(description="Test env var dict parsing")
OpenAILLMOptions.add_args(env_parser)
try:
env_args = env_parser.parse_args(
[]
) # No command line args, should use env var
reasoning_from_env = OpenAILLMOptions.options_dict(env_args).get(
"reasoning"
)
if reasoning_from_env:
print("✓ Environment variable dict parsing successful:")
print(f" Parsed reasoning from env: {reasoning_from_env}")
else:
print("✗ Environment variable dict parsing failed: No reasoning found")
except Exception as e:
print(f"✗ Environment variable dict parsing failed: {e}")
finally:
# Clean up environment variable
if "OPENAI_LLM_REASONING" in os.environ:
del os.environ["OPENAI_LLM_REASONING"]
else:
print(BindingOptions.generate_dot_env_sample())
LightRAG/lightrag/llm/deprecated/siliconcloud.py
+69
View File
@@ -0,0 +1,69 @@
import sys
if sys.version_info < (3, 9):
pass
else:
pass
import pipmaster as pm # Pipmaster for dynamic library install
# install specific modules
if not pm.is_installed("lmdeploy"):
pm.install("lmdeploy")
from openai import (
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
import numpy as np
import aiohttp
import base64
import struct
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def siliconcloud_embedding(
texts: list[str],
model: str = "netease-youdao/bce-embedding-base_v1",
base_url: str = "https://api.siliconflow.cn/v1/embeddings",
max_token_size: int = 8192,
api_key: str = None,
) -> np.ndarray:
if api_key and not api_key.startswith("Bearer "):
api_key = "Bearer " + api_key
headers = {"Authorization": api_key, "Content-Type": "application/json"}
truncate_texts = [text[0:max_token_size] for text in texts]
payload = {"model": model, "input": truncate_texts, "encoding_format": "base64"}
base64_strings = []
async with aiohttp.ClientSession() as session:
async with session.post(base_url, headers=headers, json=payload) as response:
content = await response.json()
if "code" in content:
raise ValueError(content)
base64_strings = [item["embedding"] for item in content["data"]]
embeddings = []
for string in base64_strings:
decode_bytes = base64.b64decode(string)
n = len(decode_bytes) // 4
float_array = struct.unpack("<" + "f" * n, decode_bytes)
embeddings.append(float_array)
return np.array(embeddings)
LightRAG/lightrag/llm/gemini.py
+595
View File
@@ -0,0 +1,595 @@
"""
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 asyncio
import logging
import os
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
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
) -> genai.Client:
"""
Create (or fetch cached) Gemini client.
Args:
api_key: Google Gemini API key.
base_url: Optional custom API endpoint.
timeout: Optional request timeout in milliseconds.
Returns:
genai.Client: Configured Gemini client instance.
"""
client_kwargs: dict[str, Any] = {"api_key": api_key}
if base_url and base_url != DEFAULT_GEMINI_ENDPOINT or timeout is not None:
try:
http_options_kwargs = {}
if base_url and base_url != DEFAULT_GEMINI_ENDPOINT:
http_options_kwargs["api_endpoint"] = base_url
if timeout is not None:
http_options_kwargs["timeout"] = timeout
client_kwargs["http_options"] = types.HttpOptions(**http_options_kwargs)
except Exception as exc: # pragma: no cover - defensive
LOG.warning("Failed to apply custom Gemini http_options: %s", exc)
try:
return genai.Client(**client_kwargs)
except TypeError:
# Older google-genai releases don't accept http_options; retry without it.
client_kwargs.pop("http_options", None)
return genai.Client(**client_kwargs)
def _ensure_api_key(api_key: str | None) -> str:
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,
keyword_extraction: bool,
) -> 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
if keyword_extraction and not config_data.get("response_mime_type"):
config_data["response_mime_type"] = "application/json"
# 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 _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,
keyword_extraction: bool = False,
generation_config: dict[str, Any] | None = None,
timeout: int | 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.
COT Integration:
- When enable_cot=True: Thought content is wrapped in <think>...</think> 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.
keyword_extraction: Whether to use JSON response format.
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 <think>...</think> tags.
Raises:
RuntimeError: If the response from Gemini is empty.
ValueError: If API key is not provided or configured.
"""
loop = asyncio.get_running_loop()
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)
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,
keyword_extraction=keyword_extraction,
)
request_kwargs: dict[str, Any] = {
"model": model,
"contents": [combined_prompt],
}
if config_obj is not None:
request_kwargs["config"] = config_obj
def _call_model():
return client.models.generate_content(**request_kwargs)
if stream:
queue: asyncio.Queue[Any] = asyncio.Queue()
usage_container: dict[str, Any] = {}
def _stream_model() -> None:
# COT state tracking for streaming
cot_active = False
cot_started = False
initial_content_seen = False
try:
stream_kwargs = dict(request_kwargs)
stream_iterator = client.models.generate_content_stream(**stream_kwargs)
for chunk in stream_iterator:
usage = getattr(chunk, "usage_metadata", None)
if usage is not None:
usage_container["usage"] = 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:
loop.call_soon_threadsafe(queue.put_nowait, "</think>")
cot_active = False
# Send regular content
loop.call_soon_threadsafe(queue.put_nowait, regular_text)
# Process thought content
if thought_text:
if not initial_content_seen and not cot_started:
# Start COT section
loop.call_soon_threadsafe(queue.put_nowait, "<think>")
cot_active = True
cot_started = True
# Send thought content if COT is active
if cot_active:
loop.call_soon_threadsafe(
queue.put_nowait, thought_text
)
else:
# COT disabled - only send regular content
if regular_text:
loop.call_soon_threadsafe(queue.put_nowait, regular_text)
# Ensure COT is properly closed if still active
if cot_active:
loop.call_soon_threadsafe(queue.put_nowait, "</think>")
loop.call_soon_threadsafe(queue.put_nowait, None)
except Exception as exc: # pragma: no cover - surface runtime issues
# Try to close COT tag before reporting error
if cot_active:
try:
loop.call_soon_threadsafe(queue.put_nowait, "</think>")
except Exception:
pass
loop.call_soon_threadsafe(queue.put_nowait, exc)
loop.run_in_executor(None, _stream_model)
async def _async_stream() -> AsyncIterator[str]:
try:
while True:
item = await queue.get()
if item is None:
break
if isinstance(item, Exception):
raise item
chunk_text = str(item)
if "\\u" in chunk_text:
chunk_text = safe_unicode_decode(chunk_text.encode("utf-8"))
# Yield the chunk directly without filtering
# COT filtering is already handled in _stream_model()
yield chunk_text
finally:
usage = usage_container.get("usage")
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),
}
)
return _async_stream()
response = await asyncio.to_thread(_call_model)
# 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 <think> tags
if thought_text and thought_text.strip():
if not regular_text or regular_text.strip() == "":
# Only thought content available
final_text = f"<think>{thought_text}</think>"
else:
# Both content types present: prepend thought to regular content
final_text = f"<think>{thought_text}</think>{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,
keyword_extraction: bool = False,
**kwargs: Any,
) -> str | AsyncIterator[str]:
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,
keyword_extraction=keyword_extraction,
**kwargs,
)
@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=2048)
@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,
task_type: str = "RETRIEVAL_DOCUMENT",
timeout: int | None = None,
token_tracker: Any | None = None,
) -> 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.
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.
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
"""
loop = asyncio.get_running_loop()
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:
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
def _call_embed() -> Any:
"""Call Gemini embedding API in executor thread."""
request_kwargs: dict[str, Any] = {
"model": model,
"contents": texts,
}
if config_obj is not None:
request_kwargs["config"] = config_obj
return client.models.embed_content(**request_kwargs)
# Execute API call in thread pool
response = await loop.run_in_executor(None, _call_embed)
# 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",
]
LightRAG/lightrag/llm/hf.py
+175
View File
@@ -0,0 +1,175 @@
import copy
import os
from functools import lru_cache
import pipmaster as pm # Pipmaster for dynamic library install
# install specific modules
if not pm.is_installed("transformers"):
pm.install("transformers")
if not pm.is_installed("torch"):
pm.install("torch")
if not pm.is_installed("numpy"):
pm.install("numpy")
from transformers import AutoTokenizer, AutoModelForCausalLM
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.exceptions import (
APIConnectionError,
RateLimitError,
APITimeoutError,
)
import torch
import numpy as np
from lightrag.utils import wrap_embedding_func_with_attrs
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@lru_cache(maxsize=1)
def initialize_hf_model(model_name):
hf_tokenizer = AutoTokenizer.from_pretrained(
model_name, device_map="auto", trust_remote_code=True
)
hf_model = AutoModelForCausalLM.from_pretrained(
model_name, device_map="auto", trust_remote_code=True
)
if hf_tokenizer.pad_token is None:
hf_tokenizer.pad_token = hf_tokenizer.eos_token
return hf_model, hf_tokenizer
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def hf_model_if_cache(
model,
prompt,
system_prompt=None,
history_messages=[],
enable_cot: bool = False,
**kwargs,
) -> str:
if enable_cot:
from lightrag.utils import logger
logger.debug(
"enable_cot=True is not supported for Hugging Face local models and will be ignored."
)
model_name = model
hf_model, hf_tokenizer = initialize_hf_model(model_name)
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.extend(history_messages)
messages.append({"role": "user", "content": prompt})
kwargs.pop("hashing_kv", None)
input_prompt = ""
try:
input_prompt = hf_tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
except Exception:
try:
ori_message = copy.deepcopy(messages)
if messages[0]["role"] == "system":
messages[1]["content"] = (
"<system>"
+ messages[0]["content"]
+ "</system>\n"
+ messages[1]["content"]
)
messages = messages[1:]
input_prompt = hf_tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
except Exception:
len_message = len(ori_message)
for msgid in range(len_message):
input_prompt = (
input_prompt
+ "<"
+ ori_message[msgid]["role"]
+ ">"
+ ori_message[msgid]["content"]
+ "</"
+ ori_message[msgid]["role"]
+ ">\n"
)
input_ids = hf_tokenizer(
input_prompt, return_tensors="pt", padding=True, truncation=True
).to("cuda")
inputs = {k: v.to(hf_model.device) for k, v in input_ids.items()}
output = hf_model.generate(
**input_ids, max_new_tokens=512, num_return_sequences=1, early_stopping=True
)
response_text = hf_tokenizer.decode(
output[0][len(inputs["input_ids"][0]) :], skip_special_tokens=True
)
return response_text
async def hf_model_complete(
prompt,
system_prompt=None,
history_messages=[],
keyword_extraction=False,
enable_cot: bool = False,
**kwargs,
) -> str:
kwargs.pop("keyword_extraction", None)
model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
result = await hf_model_if_cache(
model_name,
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
return result
@wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
async def hf_embed(texts: list[str], tokenizer, embed_model) -> np.ndarray:
# Detect the appropriate device
if torch.cuda.is_available():
device = next(embed_model.parameters()).device # Use CUDA if available
elif torch.backends.mps.is_available():
device = torch.device("mps") # Use MPS for Apple Silicon
else:
device = torch.device("cpu") # Fallback to CPU
# Move the model to the detected device
embed_model = embed_model.to(device)
# Tokenize the input texts and move them to the same device
encoded_texts = tokenizer(
texts, return_tensors="pt", padding=True, truncation=True
).to(device)
# Perform inference
with torch.no_grad():
outputs = embed_model(
input_ids=encoded_texts["input_ids"],
attention_mask=encoded_texts["attention_mask"],
)
embeddings = outputs.last_hidden_state.mean(dim=1)
# Convert embeddings to NumPy
if embeddings.dtype == torch.bfloat16:
return embeddings.detach().to(torch.float32).cpu().numpy()
else:
return embeddings.detach().cpu().numpy()
LightRAG/lightrag/llm/jina.py
+152
View File
@@ -0,0 +1,152 @@
import os
import pipmaster as pm # Pipmaster for dynamic library install
# install specific modules
if not pm.is_installed("aiohttp"):
pm.install("aiohttp")
if not pm.is_installed("tenacity"):
pm.install("tenacity")
import numpy as np
import base64
import aiohttp
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.utils import wrap_embedding_func_with_attrs, logger
async def fetch_data(url, headers, data):
async with aiohttp.ClientSession() as session:
async with session.post(url, headers=headers, json=data) as response:
if response.status != 200:
error_text = await response.text()
# Check if the error response is HTML (common for 502, 503, etc.)
content_type = response.headers.get("content-type", "").lower()
is_html_error = (
error_text.strip().startswith("<!DOCTYPE html>")
or "text/html" in content_type
)
if is_html_error:
# Provide clean, user-friendly error messages for HTML error pages
if response.status == 502:
clean_error = "Bad Gateway (502) - Jina AI service temporarily unavailable. Please try again in a few minutes."
elif response.status == 503:
clean_error = "Service Unavailable (503) - Jina AI service is temporarily overloaded. Please try again later."
elif response.status == 504:
clean_error = "Gateway Timeout (504) - Jina AI service request timed out. Please try again."
else:
clean_error = f"HTTP {response.status} - Jina AI service error. Please try again later."
else:
# Use original error text if it's not HTML
clean_error = error_text
logger.error(f"Jina API error {response.status}: {clean_error}")
raise aiohttp.ClientResponseError(
request_info=response.request_info,
history=response.history,
status=response.status,
message=f"Jina API error: {clean_error}",
)
response_json = await response.json()
data_list = response_json.get("data", [])
return data_list
@wrap_embedding_func_with_attrs(embedding_dim=2048, max_token_size=8192)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=(
retry_if_exception_type(aiohttp.ClientError)
| retry_if_exception_type(aiohttp.ClientResponseError)
),
)
async def jina_embed(
texts: list[str],
embedding_dim: int = 2048,
late_chunking: bool = False,
base_url: str = None,
api_key: str = None,
) -> np.ndarray:
"""Generate embeddings for a list of texts using Jina AI's API.
Args:
texts: List of texts to embed.
embedding_dim: The embedding dimensions (default: 2048 for jina-embeddings-v4).
**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 Jina API for dimension reduction.
late_chunking: Whether to use late chunking.
base_url: Optional base URL for the Jina API.
api_key: Optional Jina API key. If None, uses the JINA_API_KEY environment variable.
Returns:
A numpy array of embeddings, one per input text.
Raises:
aiohttp.ClientError: If there is a connection error with the Jina API.
aiohttp.ClientResponseError: If the Jina API returns an error response.
"""
if api_key:
os.environ["JINA_API_KEY"] = api_key
if "JINA_API_KEY" not in os.environ:
raise ValueError("JINA_API_KEY environment variable is required")
url = base_url or "https://api.jina.ai/v1/embeddings"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {os.environ['JINA_API_KEY']}",
}
data = {
"model": "jina-embeddings-v4",
"task": "text-matching",
"dimensions": embedding_dim,
"embedding_type": "base64",
"input": texts,
}
# Only add optional parameters if they have non-default values
if late_chunking:
data["late_chunking"] = late_chunking
logger.debug(
f"Jina embedding request: {len(texts)} texts, dimensions: {embedding_dim}"
)
try:
data_list = await fetch_data(url, headers, data)
if not data_list:
logger.error("Jina API returned empty data list")
raise ValueError("Jina API returned empty data list")
if len(data_list) != len(texts):
logger.error(
f"Jina API returned {len(data_list)} embeddings for {len(texts)} texts"
)
raise ValueError(
f"Jina API returned {len(data_list)} embeddings for {len(texts)} texts"
)
embeddings = np.array(
[
np.frombuffer(base64.b64decode(dp["embedding"]), dtype=np.float32)
for dp in data_list
]
)
logger.debug(f"Jina embeddings generated: shape {embeddings.shape}")
return embeddings
except Exception as e:
logger.error(f"Jina embedding error: {e}")
raise
LightRAG/lightrag/llm/llama_index_impl.py
+208
View File
@@ -0,0 +1,208 @@
import pipmaster as pm
from llama_index.core.llms import (
ChatMessage,
MessageRole,
ChatResponse,
)
from typing import List, Optional
from lightrag.utils import logger
# Install required dependencies
if not pm.is_installed("llama-index"):
pm.install("llama-index")
from llama_index.core.embeddings import BaseEmbedding
from llama_index.core.settings import Settings as LlamaIndexSettings
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.utils import (
wrap_embedding_func_with_attrs,
)
from lightrag.exceptions import (
APIConnectionError,
RateLimitError,
APITimeoutError,
)
import numpy as np
def configure_llama_index(settings: LlamaIndexSettings = None, **kwargs):
"""
Configure LlamaIndex settings.
Args:
settings: LlamaIndex Settings instance. If None, uses default settings.
**kwargs: Additional settings to override/configure
"""
if settings is None:
settings = LlamaIndexSettings()
# Update settings with any provided kwargs
for key, value in kwargs.items():
if hasattr(settings, key):
setattr(settings, key, value)
else:
logger.warning(f"Unknown LlamaIndex setting: {key}")
# Set as global settings
LlamaIndexSettings.set_global(settings)
return settings
def format_chat_messages(messages):
"""Format chat messages into LlamaIndex format."""
formatted_messages = []
for msg in messages:
role = msg.get("role", "user")
content = msg.get("content", "")
if role == "system":
formatted_messages.append(
ChatMessage(role=MessageRole.SYSTEM, content=content)
)
elif role == "assistant":
formatted_messages.append(
ChatMessage(role=MessageRole.ASSISTANT, content=content)
)
elif role == "user":
formatted_messages.append(
ChatMessage(role=MessageRole.USER, content=content)
)
else:
logger.warning(f"Unknown role {role}, treating as user message")
formatted_messages.append(
ChatMessage(role=MessageRole.USER, content=content)
)
return formatted_messages
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def llama_index_complete_if_cache(
model: str,
prompt: str,
system_prompt: Optional[str] = None,
history_messages: List[dict] = [],
enable_cot: bool = False,
chat_kwargs={},
) -> str:
"""Complete the prompt using LlamaIndex."""
if enable_cot:
logger.debug(
"enable_cot=True is not supported for LlamaIndex implementation and will be ignored."
)
try:
# Format messages for chat
formatted_messages = []
# Add system message if provided
if system_prompt:
formatted_messages.append(
ChatMessage(role=MessageRole.SYSTEM, content=system_prompt)
)
# Add history messages
for msg in history_messages:
formatted_messages.append(
ChatMessage(
role=MessageRole.USER
if msg["role"] == "user"
else MessageRole.ASSISTANT,
content=msg["content"],
)
)
# Add current prompt
formatted_messages.append(ChatMessage(role=MessageRole.USER, content=prompt))
response: ChatResponse = await model.achat(
messages=formatted_messages, **chat_kwargs
)
# In newer versions, the response is in message.content
content = response.message.content
return content
except Exception as e:
logger.error(f"Error in llama_index_complete_if_cache: {str(e)}")
raise
async def llama_index_complete(
prompt,
system_prompt=None,
history_messages=None,
enable_cot: bool = False,
keyword_extraction=False,
settings: LlamaIndexSettings = None,
**kwargs,
) -> str:
"""
Main completion function for LlamaIndex
Args:
prompt: Input prompt
system_prompt: Optional system prompt
history_messages: Optional chat history
keyword_extraction: Whether to extract keywords from response
settings: Optional LlamaIndex settings
**kwargs: Additional arguments
"""
if history_messages is None:
history_messages = []
kwargs.pop("keyword_extraction", None)
result = await llama_index_complete_if_cache(
kwargs.get("llm_instance"),
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
return result
@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def llama_index_embed(
texts: list[str],
embed_model: BaseEmbedding = None,
settings: LlamaIndexSettings = None,
**kwargs,
) -> np.ndarray:
"""
Generate embeddings using LlamaIndex
Args:
texts: List of texts to embed
embed_model: LlamaIndex embedding model
settings: Optional LlamaIndex settings
**kwargs: Additional arguments
"""
if settings:
configure_llama_index(settings)
if embed_model is None:
raise ValueError("embed_model must be provided")
# Use _get_text_embeddings for batch processing
embeddings = embed_model._get_text_embeddings(texts)
return np.array(embeddings)
LightRAG/lightrag/llm/lmdeploy.py
+154
View File
@@ -0,0 +1,154 @@
import pipmaster as pm # Pipmaster for dynamic library install
# install specific modules
if not pm.is_installed("lmdeploy"):
pm.install("lmdeploy[all]")
from lightrag.exceptions import (
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from functools import lru_cache
@lru_cache(maxsize=1)
def initialize_lmdeploy_pipeline(
model,
tp=1,
chat_template=None,
log_level="WARNING",
model_format="hf",
quant_policy=0,
):
from lmdeploy import pipeline, ChatTemplateConfig, TurbomindEngineConfig
lmdeploy_pipe = pipeline(
model_path=model,
backend_config=TurbomindEngineConfig(
tp=tp, model_format=model_format, quant_policy=quant_policy
),
chat_template_config=(
ChatTemplateConfig(model_name=chat_template) if chat_template else None
),
log_level="WARNING",
)
return lmdeploy_pipe
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def lmdeploy_model_if_cache(
model,
prompt,
system_prompt=None,
history_messages=[],
enable_cot: bool = False,
chat_template=None,
model_format="hf",
quant_policy=0,
**kwargs,
) -> str:
"""
Args:
model (str): The path to the model.
It could be one of the following options:
- i) A local directory path of a turbomind model which is
converted by `lmdeploy convert` command or download
from ii) and iii).
- ii) The model_id of a lmdeploy-quantized model hosted
inside a model repo on huggingface.co, such as
"InternLM/internlm-chat-20b-4bit",
"lmdeploy/llama2-chat-70b-4bit", etc.
- iii) The model_id of a model hosted inside a model repo
on huggingface.co, such as "internlm/internlm-chat-7b",
"Qwen/Qwen-7B-Chat ", "baichuan-inc/Baichuan2-7B-Chat"
and so on.
chat_template (str): needed when model is a pytorch model on
huggingface.co, such as "internlm-chat-7b",
"Qwen-7B-Chat ", "Baichuan2-7B-Chat" and so on,
and when the model name of local path did not match the original model name in HF.
tp (int): tensor parallel
prompt (Union[str, List[str]]): input texts to be completed.
do_preprocess (bool): whether pre-process the messages. Default to
True, which means chat_template will be applied.
skip_special_tokens (bool): Whether or not to remove special tokens
in the decoding. Default to be True.
do_sample (bool): Whether or not to use sampling, use greedy decoding otherwise.
Default to be False, which means greedy decoding will be applied.
"""
if enable_cot:
from lightrag.utils import logger
logger.debug(
"enable_cot=True is not supported for lmdeploy and will be ignored."
)
try:
import lmdeploy
from lmdeploy import version_info, GenerationConfig
except Exception:
raise ImportError("Please install lmdeploy before initialize lmdeploy backend.")
kwargs.pop("hashing_kv", None)
kwargs.pop("response_format", None)
max_new_tokens = kwargs.pop("max_tokens", 512)
tp = kwargs.pop("tp", 1)
skip_special_tokens = kwargs.pop("skip_special_tokens", True)
do_preprocess = kwargs.pop("do_preprocess", True)
do_sample = kwargs.pop("do_sample", False)
gen_params = kwargs
version = version_info
if do_sample is not None and version < (0, 6, 0):
raise RuntimeError(
"`do_sample` parameter is not supported by lmdeploy until "
f"v0.6.0, but currently using lmdeloy {lmdeploy.__version__}"
)
else:
do_sample = True
gen_params.update(do_sample=do_sample)
lmdeploy_pipe = initialize_lmdeploy_pipeline(
model=model,
tp=tp,
chat_template=chat_template,
model_format=model_format,
quant_policy=quant_policy,
log_level="WARNING",
)
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.extend(history_messages)
messages.append({"role": "user", "content": prompt})
gen_config = GenerationConfig(
skip_special_tokens=skip_special_tokens,
max_new_tokens=max_new_tokens,
**gen_params,
)
response = ""
async for res in lmdeploy_pipe.generate(
messages,
gen_config=gen_config,
do_preprocess=do_preprocess,
stream_response=False,
session_id=1,
):
response += res.response
return response
LightRAG/lightrag/llm/lollms.py
+175
View File
@@ -0,0 +1,175 @@
import sys
if sys.version_info < (3, 9):
from typing import AsyncIterator
else:
from collections.abc import AsyncIterator
import pipmaster as pm # Pipmaster for dynamic library install
if not pm.is_installed("aiohttp"):
pm.install("aiohttp")
import aiohttp
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.exceptions import (
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from typing import Union, List
import numpy as np
from lightrag.utils import (
wrap_embedding_func_with_attrs,
)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def lollms_model_if_cache(
model,
prompt,
system_prompt=None,
history_messages=[],
enable_cot: bool = False,
base_url="http://localhost:9600",
**kwargs,
) -> Union[str, AsyncIterator[str]]:
"""Client implementation for lollms generation."""
if enable_cot:
from lightrag.utils import logger
logger.debug("enable_cot=True is not supported for lollms and will be ignored.")
stream = True if kwargs.get("stream") else False
api_key = kwargs.pop("api_key", None)
headers = (
{"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
if api_key
else {"Content-Type": "application/json"}
)
# Extract lollms specific parameters
request_data = {
"prompt": prompt,
"model_name": model,
"personality": kwargs.get("personality", -1),
"n_predict": kwargs.get("n_predict", None),
"stream": stream,
"temperature": kwargs.get("temperature", 1.0),
"top_k": kwargs.get("top_k", 50),
"top_p": kwargs.get("top_p", 0.95),
"repeat_penalty": kwargs.get("repeat_penalty", 0.8),
"repeat_last_n": kwargs.get("repeat_last_n", 40),
"seed": kwargs.get("seed", None),
"n_threads": kwargs.get("n_threads", 8),
}
# Prepare the full prompt including history
full_prompt = ""
if system_prompt:
full_prompt += f"{system_prompt}\n"
for msg in history_messages:
full_prompt += f"{msg['role']}: {msg['content']}\n"
full_prompt += prompt
request_data["prompt"] = full_prompt
timeout = aiohttp.ClientTimeout(total=kwargs.get("timeout", None))
async with aiohttp.ClientSession(timeout=timeout, headers=headers) as session:
if stream:
async def inner():
async with session.post(
f"{base_url}/lollms_generate", json=request_data
) as response:
async for line in response.content:
yield line.decode().strip()
return inner()
else:
async with session.post(
f"{base_url}/lollms_generate", json=request_data
) as response:
return await response.text()
async def lollms_model_complete(
prompt,
system_prompt=None,
history_messages=[],
enable_cot: bool = False,
keyword_extraction=False,
**kwargs,
) -> Union[str, AsyncIterator[str]]:
"""Complete function for lollms model generation."""
# Extract and remove keyword_extraction from kwargs if present
keyword_extraction = kwargs.pop("keyword_extraction", None)
# Get model name from config
model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
# If keyword extraction is needed, we might need to modify the prompt
# or add specific parameters for JSON output (if lollms supports it)
if keyword_extraction:
# Note: You might need to adjust this based on how lollms handles structured output
pass
return await lollms_model_if_cache(
model_name,
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
@wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
async def lollms_embed(
texts: List[str], embed_model=None, base_url="http://localhost:9600", **kwargs
) -> np.ndarray:
"""
Generate embeddings for a list of texts using lollms server.
Args:
texts: List of strings to embed
embed_model: Model name (not used directly as lollms uses configured vectorizer)
base_url: URL of the lollms server
**kwargs: Additional arguments passed to the request
Returns:
np.ndarray: Array of embeddings
"""
api_key = kwargs.pop("api_key", None)
headers = (
{"Content-Type": "application/json", "Authorization": api_key}
if api_key
else {"Content-Type": "application/json"}
)
async with aiohttp.ClientSession(headers=headers) as session:
embeddings = []
for text in texts:
request_data = {"text": text}
async with session.post(
f"{base_url}/lollms_embed",
json=request_data,
) as response:
result = await response.json()
embeddings.append(result["vector"])
return np.array(embeddings)
LightRAG/lightrag/llm/nvidia_openai.py
+66
View File
@@ -0,0 +1,66 @@
import sys
import os
if sys.version_info < (3, 9):
pass
else:
pass
import pipmaster as pm # Pipmaster for dynamic library install
# install specific modules
if not pm.is_installed("openai"):
pm.install("openai")
from openai import (
AsyncOpenAI,
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.utils import (
wrap_embedding_func_with_attrs,
)
import numpy as np
@wrap_embedding_func_with_attrs(embedding_dim=2048, max_token_size=8192)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def nvidia_openai_embed(
texts: list[str],
model: str = "nvidia/llama-3.2-nv-embedqa-1b-v1",
# refer to https://build.nvidia.com/nim?filters=usecase%3Ausecase_text_to_embedding
base_url: str = "https://integrate.api.nvidia.com/v1",
api_key: str = None,
input_type: str = "passage", # query for retrieval, passage for embedding
trunc: str = "NONE", # NONE or START or END
encode: str = "float", # float or base64
) -> np.ndarray:
if api_key:
os.environ["OPENAI_API_KEY"] = api_key
openai_async_client = (
AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url)
)
response = await openai_async_client.embeddings.create(
model=model,
input=texts,
encoding_format=encode,
extra_body={"input_type": input_type, "truncate": trunc},
)
return np.array([dp.embedding for dp in response.data])
LightRAG/lightrag/llm/ollama.py
+214
View File
@@ -0,0 +1,214 @@
from collections.abc import AsyncIterator
import os
import re
import pipmaster as pm
# install specific modules
if not pm.is_installed("ollama"):
pm.install("ollama")
import ollama
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.exceptions import (
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from lightrag.api import __api_version__
import numpy as np
from typing import Optional, Union
from lightrag.utils import (
wrap_embedding_func_with_attrs,
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 (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
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def _ollama_model_if_cache(
model,
prompt,
system_prompt=None,
history_messages=[],
enable_cot: bool = False,
**kwargs,
) -> Union[str, AsyncIterator[str]]:
if enable_cot:
logger.debug("enable_cot=True is not supported for ollama and will be ignored.")
stream = True if kwargs.get("stream") else False
kwargs.pop("max_tokens", None)
# kwargs.pop("response_format", None) # allow json
host = kwargs.pop("host", None)
timeout = kwargs.pop("timeout", None)
if timeout == 0:
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__}",
}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
host = _coerce_host_for_cloud_model(host, model)
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
try:
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.extend(history_messages)
messages.append({"role": "user", "content": prompt})
response = await ollama_client.chat(model=model, messages=messages, **kwargs)
if stream:
"""cannot cache stream response and process reasoning"""
async def inner():
try:
async for chunk in response:
yield chunk["message"]["content"]
except Exception as e:
logger.error(f"Error in stream response: {str(e)}")
raise
finally:
try:
await ollama_client._client.aclose()
logger.debug("Successfully closed Ollama client for streaming")
except Exception as close_error:
logger.warning(f"Failed to close Ollama client: {close_error}")
return inner()
else:
model_response = response["message"]["content"]
"""
If the model also wraps its thoughts in a specific tag,
this information is not needed for the final
response and can simply be trimmed.
"""
return model_response
except Exception as e:
try:
await ollama_client._client.aclose()
logger.debug("Successfully closed Ollama client after exception")
except Exception as close_error:
logger.warning(
f"Failed to close Ollama client after exception: {close_error}"
)
raise e
finally:
if not stream:
try:
await ollama_client._client.aclose()
logger.debug(
"Successfully closed Ollama client for non-streaming response"
)
except Exception as close_error:
logger.warning(
f"Failed to close Ollama client in finally block: {close_error}"
)
async def ollama_model_complete(
prompt,
system_prompt=None,
history_messages=[],
enable_cot: bool = False,
keyword_extraction=False,
**kwargs,
) -> Union[str, AsyncIterator[str]]:
keyword_extraction = kwargs.pop("keyword_extraction", None)
if keyword_extraction:
kwargs["format"] = "json"
model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
return await _ollama_model_if_cache(
model_name,
prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
@wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
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__}",
}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
host = kwargs.pop("host", None)
timeout = kwargs.pop("timeout", None)
host = _coerce_host_for_cloud_model(host, embed_model)
ollama_client = ollama.AsyncClient(host=host, timeout=timeout, headers=headers)
try:
options = kwargs.pop("options", {})
data = await ollama_client.embed(
model=embed_model, input=texts, options=options
)
return np.array(data["embeddings"])
except Exception as e:
logger.error(f"Error in ollama_embed: {str(e)}")
try:
await ollama_client._client.aclose()
logger.debug("Successfully closed Ollama client after exception in embed")
except Exception as close_error:
logger.warning(
f"Failed to close Ollama client after exception in embed: {close_error}"
)
raise e
finally:
try:
await ollama_client._client.aclose()
logger.debug("Successfully closed Ollama client after embed")
except Exception as close_error:
logger.warning(f"Failed to close Ollama client after embed: {close_error}")
LightRAG/lightrag/llm/openai.py
+686
View File
@@ -0,0 +1,686 @@
from ..utils import verbose_debug, VERBOSE_DEBUG
import os
import logging
from collections.abc import AsyncIterator
import pipmaster as pm
# install specific modules
if not pm.is_installed("openai"):
pm.install("openai")
from openai import (
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.utils import (
wrap_embedding_func_with_attrs,
safe_unicode_decode,
logger,
)
from lightrag.types import GPTKeywordExtractionFormat
from lightrag.api import __api_version__
import numpy as np
import base64
from typing import Any, Union
from dotenv import load_dotenv
# Try to import Langfuse for LLM observability (optional)
# Falls back to standard OpenAI client if not available
# Langfuse requires proper configuration to work correctly
LANGFUSE_ENABLED = False
try:
# Check if required Langfuse environment variables are set
langfuse_public_key = os.environ.get("LANGFUSE_PUBLIC_KEY")
langfuse_secret_key = os.environ.get("LANGFUSE_SECRET_KEY")
# Only enable Langfuse if both keys are configured
if langfuse_public_key and langfuse_secret_key:
from langfuse.openai import AsyncOpenAI # type: ignore[import-untyped]
LANGFUSE_ENABLED = True
logger.info("Langfuse observability enabled for OpenAI client")
else:
from openai import AsyncOpenAI
logger.debug(
"Langfuse environment variables not configured, using standard OpenAI client"
)
except ImportError:
from openai import AsyncOpenAI
logger.debug("Langfuse not available, using standard OpenAI client")
# use the .env that is inside the current folder
# allows to use different .env file for each lightrag instance
# the OS environment variables take precedence over the .env file
load_dotenv(dotenv_path=".env", override=False)
class InvalidResponseError(Exception):
"""Custom exception class for triggering retry mechanism"""
pass
def create_openai_async_client(
api_key: str | None = None,
base_url: str | None = None,
client_configs: dict[str, Any] | None = None,
) -> AsyncOpenAI:
"""Create an AsyncOpenAI 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.
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.
"""
if not api_key:
api_key = os.environ["OPENAI_API_KEY"]
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"
)
return AsyncOpenAI(**merged_configs)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=(
retry_if_exception_type(RateLimitError)
| retry_if_exception_type(APIConnectionError)
| retry_if_exception_type(APITimeoutError)
| retry_if_exception_type(InvalidResponseError)
),
)
async def openai_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,
timeout: int | None = None,
keyword_extraction: bool = False,
**kwargs: Any,
) -> str:
"""Complete a prompt using OpenAI's API with caching support and Chain of Thought (COT) integration.
This function supports automatic integration of reasoning content from models that provide
Chain of Thought capabilities. The reasoning content is seamlessly integrated into the response
using <think>...</think> tags.
Note on `reasoning_content`: This feature relies on a Deepseek Style `reasoning_content`
in the API response, which may be provided by OpenAI-compatible endpoints that support
Chain of Thought.
COT Integration Rules:
1. COT content is accepted only when regular content is empty and `reasoning_content` has content.
2. COT processing stops when regular content becomes available.
3. If both `content` and `reasoning_content` are present simultaneously, reasoning is ignored.
4. If both fields have content from the start, COT is never activated.
5. For streaming: COT content is inserted into the content stream with <think> tags.
6. For non-streaming: COT content is prepended to regular content with <think> tags.
Args:
model: The OpenAI 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.
base_url: Optional base URL for the OpenAI API.
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).
Returns:
The completed text (with integrated COT content if available) or an async iterator
of text chunks if streaming. COT content is wrapped in <think>...</think> tags.
Raises:
InvalidResponseError: If the response from OpenAI is invalid or empty.
APIConnectionError: If there is a connection error with the OpenAI API.
RateLimitError: If the OpenAI API rate limit is exceeded.
APITimeoutError: If the OpenAI API request times out.
"""
if history_messages is None:
history_messages = []
# Set openai logger level to INFO when VERBOSE_DEBUG is off
if not VERBOSE_DEBUG and logger.level == logging.DEBUG:
logging.getLogger("openai").setLevel(logging.INFO)
# Remove special kwargs that shouldn't be passed to OpenAI
kwargs.pop("hashing_kv", None)
# Extract client configuration options
client_configs = kwargs.pop("openai_client_configs", {})
# Create the OpenAI client
openai_async_client = create_openai_async_client(
api_key=api_key,
base_url=base_url,
client_configs=client_configs,
)
# Prepare messages
messages: list[dict[str, Any]] = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.extend(history_messages)
messages.append({"role": "user", "content": prompt})
logger.debug("===== Entering func of LLM =====")
logger.debug(f"Model: {model} Base URL: {base_url}")
logger.debug(f"Client Configs: {client_configs}")
logger.debug(f"Additional kwargs: {kwargs}")
logger.debug(f"Num of history messages: {len(history_messages)}")
verbose_debug(f"System prompt: {system_prompt}")
verbose_debug(f"Query: {prompt}")
logger.debug("===== Sending Query to LLM =====")
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(
model=model, messages=messages, **kwargs
)
else:
response = await openai_async_client.chat.completions.create(
model=model, messages=messages, **kwargs
)
except APIConnectionError as e:
logger.error(f"OpenAI API Connection Error: {e}")
await openai_async_client.close() # Ensure client is closed
raise
except RateLimitError as e:
logger.error(f"OpenAI API Rate Limit Error: {e}")
await openai_async_client.close() # Ensure client is closed
raise
except APITimeoutError as e:
logger.error(f"OpenAI API Timeout Error: {e}")
await openai_async_client.close() # Ensure client is closed
raise
except Exception as e:
logger.error(
f"OpenAI API Call Failed,\nModel: {model},\nParams: {kwargs}, Got: {e}"
)
await openai_async_client.close() # Ensure client is closed
raise
if hasattr(response, "__aiter__"):
async def inner():
# Track if we've started iterating
iteration_started = False
final_chunk_usage = None
# COT (Chain of Thought) state tracking
cot_active = False
cot_started = False
initial_content_seen = False
try:
iteration_started = True
async for chunk in response:
# Check if this chunk has usage information (final chunk)
if hasattr(chunk, "usage") and chunk.usage:
final_chunk_usage = chunk.usage
logger.debug(
f"Received usage info in streaming chunk: {chunk.usage}"
)
# Check if choices exists and is not empty
if not hasattr(chunk, "choices") or not chunk.choices:
logger.warning(f"Received chunk without choices: {chunk}")
continue
# Check if delta exists
if not hasattr(chunk.choices[0], "delta"):
# This might be the final chunk, continue to check for usage
continue
delta = chunk.choices[0].delta
content = getattr(delta, "content", None)
reasoning_content = getattr(delta, "reasoning_content", "")
# Handle COT logic for streaming (only if enabled)
if enable_cot:
if content:
# Regular content is present
if not initial_content_seen:
initial_content_seen = True
# If both content and reasoning_content are present initially, don't start COT
if reasoning_content:
cot_active = False
cot_started = False
# If COT was active, end it
if cot_active:
yield "</think>"
cot_active = False
# Process regular content
if r"\u" in content:
content = safe_unicode_decode(content.encode("utf-8"))
yield content
elif reasoning_content:
# Only reasoning content is present
if not initial_content_seen and not cot_started:
# Start COT if we haven't seen initial content yet
if not cot_active:
yield "<think>"
cot_active = True
cot_started = True
# Process reasoning content if COT is active
if cot_active:
if r"\u" in reasoning_content:
reasoning_content = safe_unicode_decode(
reasoning_content.encode("utf-8")
)
yield reasoning_content
else:
# COT disabled, only process regular content
if content:
if r"\u" in content:
content = safe_unicode_decode(content.encode("utf-8"))
yield content
# If neither content nor reasoning_content, continue to next chunk
if content is None and reasoning_content is None:
continue
# Ensure COT is properly closed if still active after stream ends
if enable_cot and cot_active:
yield "</think>"
cot_active = False
# After streaming is complete, track token usage
if token_tracker and final_chunk_usage:
# Use actual usage from the API
token_counts = {
"prompt_tokens": getattr(final_chunk_usage, "prompt_tokens", 0),
"completion_tokens": getattr(
final_chunk_usage, "completion_tokens", 0
),
"total_tokens": getattr(final_chunk_usage, "total_tokens", 0),
}
token_tracker.add_usage(token_counts)
logger.debug(f"Streaming token usage (from API): {token_counts}")
elif token_tracker:
logger.debug("No usage information available in streaming response")
except Exception as e:
# Ensure COT is properly closed before handling exception
if enable_cot and cot_active:
try:
yield "</think>"
cot_active = False
except Exception as close_error:
logger.warning(
f"Failed to close COT tag during exception handling: {close_error}"
)
logger.error(f"Error in stream response: {str(e)}")
# Try to clean up resources if possible
if (
iteration_started
and hasattr(response, "aclose")
and callable(getattr(response, "aclose", None))
):
try:
await response.aclose()
logger.debug("Successfully closed stream response after error")
except Exception as close_error:
logger.warning(
f"Failed to close stream response: {close_error}"
)
# Ensure client is closed in case of exception
await openai_async_client.close()
raise
finally:
# Final safety check for unclosed COT tags
if enable_cot and cot_active:
try:
yield "</think>"
cot_active = False
except Exception as final_close_error:
logger.warning(
f"Failed to close COT tag in finally block: {final_close_error}"
)
# Ensure resources are released even if no exception occurs
# Note: Some wrapped clients (e.g., Langfuse) may not implement aclose() properly
if iteration_started and hasattr(response, "aclose"):
aclose_method = getattr(response, "aclose", None)
if callable(aclose_method):
try:
await response.aclose()
logger.debug("Successfully closed stream response")
except (AttributeError, TypeError) as close_error:
# Some wrapper objects may report hasattr(aclose) but fail when called
# This is expected behavior for certain client wrappers
logger.debug(
f"Stream response cleanup not supported by client wrapper: {close_error}"
)
except Exception as close_error:
logger.warning(
f"Unexpected error during stream response cleanup: {close_error}"
)
# This prevents resource leaks since the caller doesn't handle closing
try:
await openai_async_client.close()
logger.debug(
"Successfully closed OpenAI client for streaming response"
)
except Exception as client_close_error:
logger.warning(
f"Failed to close OpenAI client in streaming finally block: {client_close_error}"
)
return inner()
else:
try:
if (
not response
or not response.choices
or not hasattr(response.choices[0], "message")
):
logger.error("Invalid response from OpenAI API")
await openai_async_client.close() # Ensure client is closed
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 = ""
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:
# 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"<think>{reasoning_content}</think>{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")
# Apply Unicode decoding to final content if needed
if r"\u" in final_content:
final_content = safe_unicode_decode(final_content.encode("utf-8"))
if token_tracker and hasattr(response, "usage"):
token_counts = {
"prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
"completion_tokens": getattr(
response.usage, "completion_tokens", 0
),
"total_tokens": getattr(response.usage, "total_tokens", 0),
}
token_tracker.add_usage(token_counts)
logger.debug(f"Response content len: {len(final_content)}")
verbose_debug(f"Response: {response}")
return final_content
finally:
# Ensure client is closed in all cases for non-streaming responses
await openai_async_client.close()
async def openai_complete(
prompt,
system_prompt=None,
history_messages=None,
keyword_extraction=False,
**kwargs,
) -> Union[str, AsyncIterator[str]]:
if history_messages is None:
history_messages = []
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,
)
async def gpt_4o_complete(
prompt,
system_prompt=None,
history_messages=None,
enable_cot: bool = False,
keyword_extraction=False,
**kwargs,
) -> str:
if history_messages is None:
history_messages = []
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,
)
async def gpt_4o_mini_complete(
prompt,
system_prompt=None,
history_messages=None,
enable_cot: bool = False,
keyword_extraction=False,
**kwargs,
) -> str:
if history_messages is None:
history_messages = []
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,
)
async def nvidia_openai_complete(
prompt,
system_prompt=None,
history_messages=None,
enable_cot: bool = False,
keyword_extraction=False,
**kwargs,
) -> str:
if history_messages is None:
history_messages = []
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,
)
return result
@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=(
retry_if_exception_type(RateLimitError)
| retry_if_exception_type(APIConnectionError)
| retry_if_exception_type(APITimeoutError)
),
)
async def openai_embed(
texts: list[str],
model: str = "text-embedding-3-small",
base_url: str | None = None,
api_key: str | None = None,
embedding_dim: int | None = None,
client_configs: dict[str, Any] | None = None,
token_tracker: Any | None = None,
) -> np.ndarray:
"""Generate embeddings for a list of texts using OpenAI's API.
Args:
texts: List of texts to 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 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).
token_tracker: Optional token usage tracker for monitoring API usage.
Returns:
A numpy array of embeddings, one per input text.
Raises:
APIConnectionError: If there is a connection error with the OpenAI API.
RateLimitError: If the OpenAI API rate limit is exceeded.
APITimeoutError: If the OpenAI API request times out.
"""
# Create the OpenAI client
openai_async_client = create_openai_async_client(
api_key=api_key, base_url=base_url, client_configs=client_configs
)
async with openai_async_client:
# Prepare API call parameters
api_params = {
"model": model,
"input": texts,
"encoding_format": "base64",
}
# Add dimensions parameter only if embedding_dim is provided
if embedding_dim is not None:
api_params["dimensions"] = embedding_dim
# Make API call
response = await openai_async_client.embeddings.create(**api_params)
if token_tracker and hasattr(response, "usage"):
token_counts = {
"prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
"total_tokens": getattr(response.usage, "total_tokens", 0),
}
token_tracker.add_usage(token_counts)
return np.array(
[
np.array(dp.embedding, dtype=np.float32)
if isinstance(dp.embedding, list)
else np.frombuffer(base64.b64decode(dp.embedding), dtype=np.float32)
for dp in response.data
]
)
LightRAG/lightrag/llm/zhipu.py
+217
View File
@@ -0,0 +1,217 @@
import sys
import re
import json
from ..utils import verbose_debug
if sys.version_info < (3, 9):
pass
else:
pass
import pipmaster as pm # Pipmaster for dynamic library install
# install specific modules
if not pm.is_installed("zhipuai"):
pm.install("zhipuai")
from openai import (
APIConnectionError,
RateLimitError,
APITimeoutError,
)
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
from lightrag.utils import (
wrap_embedding_func_with_attrs,
logger,
)
from lightrag.types import GPTKeywordExtractionFormat
import numpy as np
from typing import Union, List, Optional, Dict
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def zhipu_complete_if_cache(
prompt: Union[str, List[Dict[str, str]]],
model: str = "glm-4-flashx", # The most cost/performance balance model in glm-4 series
api_key: Optional[str] = None,
system_prompt: Optional[str] = None,
history_messages: List[Dict[str, str]] = [],
enable_cot: bool = False,
**kwargs,
) -> str:
if enable_cot:
logger.debug(
"enable_cot=True is not supported for ZhipuAI and will be ignored."
)
# dynamically load ZhipuAI
try:
from zhipuai import ZhipuAI
except ImportError:
raise ImportError("Please install zhipuai before initialize zhipuai backend.")
if api_key:
client = ZhipuAI(api_key=api_key)
else:
# please set ZHIPUAI_API_KEY in your environment
# os.environ["ZHIPUAI_API_KEY"]
client = ZhipuAI()
messages = []
if not system_prompt:
system_prompt = "You are a helpful assistant. Note that sensitive words in the content should be replaced with ***"
# Add system prompt if provided
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.extend(history_messages)
messages.append({"role": "user", "content": prompt})
# Add debug logging
logger.debug("===== Query Input to LLM =====")
logger.debug(f"Query: {prompt}")
verbose_debug(f"System prompt: {system_prompt}")
# Remove unsupported kwargs
kwargs = {
k: v for k, v in kwargs.items() if k not in ["hashing_kv", "keyword_extraction"]
}
response = client.chat.completions.create(model=model, messages=messages, **kwargs)
return response.choices[0].message.content
async def zhipu_complete(
prompt,
system_prompt=None,
history_messages=[],
keyword_extraction=False,
enable_cot: bool = False,
**kwargs,
):
# Pop keyword_extraction from kwargs to avoid passing it to zhipu_complete_if_cache
keyword_extraction = kwargs.pop("keyword_extraction", None)
if keyword_extraction:
# Add a system prompt to guide the model to return JSON format
extraction_prompt = """You are a helpful assistant that extracts keywords from text.
Please analyze the content and extract two types of keywords:
1. High-level keywords: Important concepts and main themes
2. Low-level keywords: Specific details and supporting elements
Return your response in this exact JSON format:
{
"high_level_keywords": ["keyword1", "keyword2"],
"low_level_keywords": ["keyword1", "keyword2", "keyword3"]
}
Only return the JSON, no other text."""
# Combine with existing system prompt if any
if system_prompt:
system_prompt = f"{system_prompt}\n\n{extraction_prompt}"
else:
system_prompt = extraction_prompt
try:
response = await zhipu_complete_if_cache(
prompt=prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
# Try to parse as JSON
try:
data = json.loads(response)
return GPTKeywordExtractionFormat(
high_level_keywords=data.get("high_level_keywords", []),
low_level_keywords=data.get("low_level_keywords", []),
)
except json.JSONDecodeError:
# If direct JSON parsing fails, try to extract JSON from text
match = re.search(r"\{[\s\S]*\}", response)
if match:
try:
data = json.loads(match.group())
return GPTKeywordExtractionFormat(
high_level_keywords=data.get("high_level_keywords", []),
low_level_keywords=data.get("low_level_keywords", []),
)
except json.JSONDecodeError:
pass
# If all parsing fails, log warning and return empty format
logger.warning(
f"Failed to parse keyword extraction response: {response}"
)
return GPTKeywordExtractionFormat(
high_level_keywords=[], low_level_keywords=[]
)
except Exception as e:
logger.error(f"Error during keyword extraction: {str(e)}")
return GPTKeywordExtractionFormat(
high_level_keywords=[], low_level_keywords=[]
)
else:
# For non-keyword-extraction, just return the raw response string
return await zhipu_complete_if_cache(
prompt=prompt,
system_prompt=system_prompt,
history_messages=history_messages,
enable_cot=enable_cot,
**kwargs,
)
@wrap_embedding_func_with_attrs(embedding_dim=1024)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=60),
retry=retry_if_exception_type(
(RateLimitError, APIConnectionError, APITimeoutError)
),
)
async def zhipu_embedding(
texts: list[str], model: str = "embedding-3", api_key: str = None, **kwargs
) -> np.ndarray:
# dynamically load ZhipuAI
try:
from zhipuai import ZhipuAI
except ImportError:
raise ImportError("Please install zhipuai before initialize zhipuai backend.")
if api_key:
client = ZhipuAI(api_key=api_key)
else:
# please set ZHIPUAI_API_KEY in your environment
# os.environ["ZHIPUAI_API_KEY"]
client = ZhipuAI()
# Convert single text to list if needed
if isinstance(texts, str):
texts = [texts]
embeddings = []
for text in texts:
try:
response = client.embeddings.create(model=model, input=[text], **kwargs)
embeddings.append(response.data[0].embedding)
except Exception as e:
raise Exception(f"Error calling ChatGLM Embedding API: {str(e)}")
return np.array(embeddings)