chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,45 @@
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ._version import __version__ as __version__
|
||||
|
||||
__all__ = [
|
||||
"LightRAG",
|
||||
"QueryParam",
|
||||
"RoleLLMConfig",
|
||||
"RoleSpec",
|
||||
"ROLES",
|
||||
"__version__",
|
||||
]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .lightrag import (
|
||||
LightRAG as LightRAG,
|
||||
QueryParam as QueryParam,
|
||||
ROLES as ROLES,
|
||||
RoleLLMConfig as RoleLLMConfig,
|
||||
RoleSpec as RoleSpec,
|
||||
)
|
||||
|
||||
|
||||
_LAZY_EXPORTS = {"LightRAG", "QueryParam", "RoleLLMConfig", "RoleSpec", "ROLES"}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _LAZY_EXPORTS:
|
||||
from .lightrag import LightRAG, QueryParam, RoleLLMConfig, RoleSpec, ROLES
|
||||
|
||||
values = {
|
||||
"LightRAG": LightRAG,
|
||||
"QueryParam": QueryParam,
|
||||
"RoleLLMConfig": RoleLLMConfig,
|
||||
"RoleSpec": RoleSpec,
|
||||
"ROLES": ROLES,
|
||||
}
|
||||
value = values[name]
|
||||
globals()[name] = value
|
||||
return value
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__author__ = "Zirui Guo"
|
||||
__url__ = "https://github.com/HKUDS/LightRAG"
|
||||
@@ -0,0 +1,4 @@
|
||||
"""Lightweight version definitions shared by packaging and runtime code."""
|
||||
|
||||
__version__ = "1.5.5"
|
||||
__api_version__ = "0315"
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Addon parameters: observable mapping + normalization helper.
|
||||
|
||||
``addon_params`` is a free-form configuration dict on :class:`LightRAG` that
|
||||
controls things like summary language and entity-type prompt overrides. The
|
||||
module exposes:
|
||||
|
||||
- :class:`ObservableAddonParams` — a ``dict`` subclass that calls a callback
|
||||
whenever the contents change so the LightRAG runtime can invalidate cached
|
||||
derived state.
|
||||
- :func:`default_addon_params` — environment-driven defaults.
|
||||
- :func:`normalize_addon_params` — converts an arbitrary input into a plain
|
||||
``dict`` with the env-driven defaults backfilled.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable, Mapping
|
||||
|
||||
from lightrag.constants import DEFAULT_SUMMARY_LANGUAGE
|
||||
from lightrag.utils import get_env_value, logger
|
||||
|
||||
|
||||
# Keys that used to live in addon_params but have been superseded by
|
||||
# per-document ``process_options``. We log once when callers still pass them
|
||||
# so existing configs surface their drift without breaking.
|
||||
_DEPRECATED_ADDON_PARAM_KEYS: tuple[str, ...] = ("enable_multimodal_pipeline",)
|
||||
_warned_deprecated_keys: set[str] = set()
|
||||
|
||||
|
||||
def _emit_deprecated_addon_warnings(params: Mapping[str, Any]) -> None:
|
||||
for key in _DEPRECATED_ADDON_PARAM_KEYS:
|
||||
if key in params and key not in _warned_deprecated_keys:
|
||||
logger.warning(
|
||||
f"addon_params['{key}'] is deprecated and ignored; per-document "
|
||||
f"behaviour is now controlled by filename-hint process_options "
|
||||
f"(see docs/FileProcessingConfiguration-zh.md)."
|
||||
)
|
||||
_warned_deprecated_keys.add(key)
|
||||
|
||||
|
||||
def default_addon_params() -> dict[str, Any]:
|
||||
# Lazy import to avoid the parser_routing → utils → … cycle that
|
||||
# would otherwise form when parser_routing imports back into this
|
||||
# module via ``LightRAG`` construction paths.
|
||||
from lightrag.parser.routing import default_chunker_config
|
||||
|
||||
return {
|
||||
"language": get_env_value("SUMMARY_LANGUAGE", DEFAULT_SUMMARY_LANGUAGE, str),
|
||||
"entity_type_prompt_file": get_env_value("ENTITY_TYPE_PROMPT_FILE", "", str),
|
||||
# Per-strategy chunker parameters; mutate at runtime (e.g.
|
||||
# ``rag.addon_params["chunker"]["recursive_character"]["separators"]
|
||||
# = [...]``) to change defaults applied to subsequently
|
||||
# enqueued documents. Per-document snapshots are persisted to
|
||||
# ``full_docs[doc_id]["chunk_options"]`` at enqueue time and
|
||||
# are not affected by later runtime mutations.
|
||||
"chunker": default_chunker_config(),
|
||||
}
|
||||
|
||||
|
||||
def normalize_addon_params(addon_params: Mapping[str, Any] | None) -> dict[str, Any]:
|
||||
"""Coerce ``addon_params`` to a plain dict with env defaults backfilled."""
|
||||
from lightrag.parser.routing import default_chunker_config
|
||||
|
||||
if addon_params is None:
|
||||
normalized = default_addon_params()
|
||||
elif isinstance(addon_params, Mapping):
|
||||
_emit_deprecated_addon_warnings(addon_params)
|
||||
normalized = {
|
||||
k: v
|
||||
for k, v in addon_params.items()
|
||||
if k not in _DEPRECATED_ADDON_PARAM_KEYS
|
||||
}
|
||||
else:
|
||||
raise TypeError(
|
||||
f"addon_params must be a Mapping or None, got {type(addon_params).__name__}"
|
||||
)
|
||||
|
||||
# When the caller supplies addon_params explicitly, the dataclass
|
||||
# default_factory is skipped — fall back to environment variables so
|
||||
# ENTITY_TYPE_PROMPT_FILE / SUMMARY_LANGUAGE / chunker still apply.
|
||||
normalized.setdefault(
|
||||
"language", get_env_value("SUMMARY_LANGUAGE", DEFAULT_SUMMARY_LANGUAGE, str)
|
||||
)
|
||||
normalized.setdefault(
|
||||
"entity_type_prompt_file",
|
||||
get_env_value("ENTITY_TYPE_PROMPT_FILE", "", str),
|
||||
)
|
||||
# Build the chunker default lazily — `default_chunker_config()` reads env
|
||||
# vars (e.g. CHUNK_R_SEPARATORS via json.loads) and would raise on a
|
||||
# malformed value, which would prevent an explicit caller-supplied
|
||||
# `chunker` from bypassing a broken environment.
|
||||
if "chunker" not in normalized:
|
||||
normalized["chunker"] = default_chunker_config()
|
||||
return normalized
|
||||
|
||||
|
||||
class ObservableAddonParams(dict[str, Any]):
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
on_change: Callable[[], None] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._on_change = on_change
|
||||
|
||||
def _changed(self) -> None:
|
||||
if self._on_change is not None:
|
||||
self._on_change()
|
||||
|
||||
def __setitem__(self, key: str, value: Any) -> None:
|
||||
super().__setitem__(key, value)
|
||||
self._changed()
|
||||
|
||||
def __delitem__(self, key: str) -> None:
|
||||
super().__delitem__(key)
|
||||
self._changed()
|
||||
|
||||
def clear(self) -> None:
|
||||
if self:
|
||||
super().clear()
|
||||
self._changed()
|
||||
|
||||
def pop(self, key: str, default: Any = ...):
|
||||
existed = key in self
|
||||
if default is ...:
|
||||
value = super().pop(key)
|
||||
self._changed()
|
||||
else:
|
||||
value = super().pop(key, default)
|
||||
if existed:
|
||||
self._changed()
|
||||
return value
|
||||
|
||||
def popitem(self) -> tuple[str, Any]:
|
||||
item = super().popitem()
|
||||
self._changed()
|
||||
return item
|
||||
|
||||
def setdefault(self, key: str, default: Any = None) -> Any:
|
||||
if key in self:
|
||||
return self[key]
|
||||
value = super().setdefault(key, default)
|
||||
self._changed()
|
||||
return value
|
||||
|
||||
def update(self, *args: Any, **kwargs: Any) -> None:
|
||||
if not args and not kwargs:
|
||||
return
|
||||
super().update(*args, **kwargs)
|
||||
self._changed()
|
||||
|
||||
def __ior__(self, other: Mapping[str, Any]):
|
||||
super().__ior__(other)
|
||||
self._changed()
|
||||
return self
|
||||
@@ -0,0 +1,7 @@
|
||||
AZURE_OPENAI_API_VERSION=2024-08-01-preview
|
||||
AZURE_OPENAI_DEPLOYMENT=gpt-4o
|
||||
AZURE_OPENAI_API_KEY=myapikey
|
||||
AZURE_OPENAI_ENDPOINT=https://myendpoint.openai.azure.com
|
||||
|
||||
AZURE_EMBEDDING_DEPLOYMENT=text-embedding-3-large
|
||||
AZURE_EMBEDDING_API_VERSION=2023-05-15
|
||||
@@ -0,0 +1,2 @@
|
||||
inputs
|
||||
rag_storage
|
||||
@@ -0,0 +1 @@
|
||||
from .._version import __api_version__ as __api_version__
|
||||
@@ -0,0 +1,163 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import jwt
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import HTTPException, status
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ..utils import logger
|
||||
from .config import DEFAULT_TOKEN_SECRET, global_args
|
||||
from .passwords import verify_password
|
||||
|
||||
# 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 TokenPayload(BaseModel):
|
||||
sub: str # Username
|
||||
exp: datetime # Expiration time
|
||||
role: str = "user" # User role, default is regular user
|
||||
metadata: dict = {} # Additional metadata
|
||||
|
||||
|
||||
class AuthHandler:
|
||||
def __init__(self):
|
||||
auth_accounts = global_args.auth_accounts
|
||||
self.secret = global_args.token_secret
|
||||
if not self.secret:
|
||||
if auth_accounts:
|
||||
raise ValueError(
|
||||
"TOKEN_SECRET must be explicitly set to a non-default value when AUTH_ACCOUNTS is configured."
|
||||
)
|
||||
self.secret = DEFAULT_TOKEN_SECRET
|
||||
logger.warning(
|
||||
"TOKEN_SECRET not set and AUTH_ACCOUNTS is not configured. "
|
||||
"Falling back to the default guest-mode JWT secret. "
|
||||
)
|
||||
algorithm = global_args.jwt_algorithm
|
||||
if not algorithm or algorithm.lower() == "none":
|
||||
raise ValueError(
|
||||
"JWT_ALGORITHM must be set to a secure algorithm (e.g. HS256). "
|
||||
"The 'none' algorithm is not permitted."
|
||||
)
|
||||
self.algorithm = algorithm
|
||||
self.expire_hours = global_args.token_expire_hours
|
||||
self.guest_expire_hours = global_args.guest_token_expire_hours
|
||||
self.accounts = {}
|
||||
invalid_accounts = []
|
||||
if auth_accounts:
|
||||
for account in auth_accounts.split(","):
|
||||
try:
|
||||
username, password = account.split(":", 1)
|
||||
if not username or not password:
|
||||
raise ValueError
|
||||
self.accounts[username] = password
|
||||
except ValueError:
|
||||
invalid_accounts.append(account)
|
||||
if invalid_accounts:
|
||||
invalid_entries = ", ".join(invalid_accounts)
|
||||
logger.error(f"Invalid account format in AUTH_ACCOUNTS: {invalid_entries}")
|
||||
raise ValueError(
|
||||
"AUTH_ACCOUNTS must use comma-separated user:password pairs."
|
||||
)
|
||||
|
||||
def verify_password(self, username: str, plain_password: str) -> bool:
|
||||
"""
|
||||
Verify password for a user. Supports explicit bcrypt values and plaintext.
|
||||
|
||||
Args:
|
||||
username: Username to verify
|
||||
plain_password: Plaintext password to check
|
||||
|
||||
Returns:
|
||||
bool: True if password is correct, False otherwise
|
||||
"""
|
||||
if username not in self.accounts:
|
||||
return False
|
||||
|
||||
stored_password = self.accounts[username]
|
||||
return verify_password(plain_password, stored_password)
|
||||
|
||||
def create_token(
|
||||
self,
|
||||
username: str,
|
||||
role: str = "user",
|
||||
custom_expire_hours: int = None,
|
||||
metadata: dict = None,
|
||||
) -> str:
|
||||
"""
|
||||
Create JWT token
|
||||
|
||||
Args:
|
||||
username: Username
|
||||
role: User role, default is "user", guest is "guest"
|
||||
custom_expire_hours: Custom expiration time (hours), if None use default value
|
||||
metadata: Additional metadata
|
||||
|
||||
Returns:
|
||||
str: Encoded JWT token
|
||||
"""
|
||||
# Choose default expiration time based on role
|
||||
if custom_expire_hours is None:
|
||||
if role == "guest":
|
||||
expire_hours = self.guest_expire_hours
|
||||
else:
|
||||
expire_hours = self.expire_hours
|
||||
else:
|
||||
expire_hours = custom_expire_hours
|
||||
|
||||
expire = datetime.now(timezone.utc) + timedelta(hours=expire_hours)
|
||||
|
||||
# Create payload
|
||||
payload = TokenPayload(
|
||||
sub=username, exp=expire, role=role, metadata=metadata or {}
|
||||
)
|
||||
|
||||
return jwt.encode(payload.model_dump(), self.secret, algorithm=self.algorithm)
|
||||
|
||||
def validate_token(self, token: str) -> dict:
|
||||
"""
|
||||
Validate JWT token
|
||||
|
||||
Args:
|
||||
token: JWT token
|
||||
|
||||
Returns:
|
||||
dict: Dictionary containing user information
|
||||
|
||||
Raises:
|
||||
HTTPException: If token is invalid or expired
|
||||
"""
|
||||
try:
|
||||
# Explicitly exclude 'none' to prevent algorithm confusion attacks
|
||||
allowed_algorithms = [self.algorithm]
|
||||
if "none" in (a.lower() for a in allowed_algorithms):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Insecure JWT algorithm configuration",
|
||||
)
|
||||
payload = jwt.decode(token, self.secret, algorithms=allowed_algorithms)
|
||||
expire_timestamp = payload["exp"]
|
||||
expire_time = datetime.fromtimestamp(expire_timestamp, timezone.utc)
|
||||
|
||||
if datetime.now(timezone.utc) > expire_time:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail="Token expired"
|
||||
)
|
||||
|
||||
# Return complete payload instead of just username
|
||||
return {
|
||||
"username": payload["sub"],
|
||||
"role": payload.get("role", "user"),
|
||||
"metadata": payload.get("metadata", {}),
|
||||
"exp": expire_time,
|
||||
}
|
||||
except jwt.PyJWTError:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token"
|
||||
)
|
||||
|
||||
|
||||
auth_handler = AuthHandler()
|
||||
@@ -0,0 +1,873 @@
|
||||
"""
|
||||
Configs for the LightRAG API.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import argparse
|
||||
import logging
|
||||
from dotenv import load_dotenv
|
||||
from lightrag import ROLES
|
||||
from lightrag.utils import get_env_value, logger
|
||||
from lightrag.llm.binding_options import (
|
||||
BedrockLLMOptions,
|
||||
GeminiEmbeddingOptions,
|
||||
GeminiLLMOptions,
|
||||
OllamaEmbeddingOptions,
|
||||
OllamaLLMOptions,
|
||||
OpenAILLMOptions,
|
||||
)
|
||||
from lightrag.base import OllamaServerInfos
|
||||
import sys
|
||||
|
||||
from lightrag.constants import (
|
||||
DEFAULT_WOKERS,
|
||||
DEFAULT_TIMEOUT,
|
||||
DEFAULT_TOP_K,
|
||||
DEFAULT_CHUNK_TOP_K,
|
||||
DEFAULT_MAX_ENTITY_TOKENS,
|
||||
DEFAULT_MAX_RELATION_TOKENS,
|
||||
DEFAULT_MAX_TOTAL_TOKENS,
|
||||
DEFAULT_COSINE_THRESHOLD,
|
||||
DEFAULT_RELATED_CHUNK_NUMBER,
|
||||
DEFAULT_MIN_RERANK_SCORE,
|
||||
DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE,
|
||||
DEFAULT_MAX_ASYNC,
|
||||
DEFAULT_MAX_PARALLEL_INSERT,
|
||||
DEFAULT_SUMMARY_MAX_TOKENS,
|
||||
DEFAULT_SUMMARY_LENGTH_RECOMMENDED,
|
||||
DEFAULT_SUMMARY_CONTEXT_SIZE,
|
||||
DEFAULT_SUMMARY_LANGUAGE,
|
||||
DEFAULT_EMBEDDING_FUNC_MAX_ASYNC,
|
||||
DEFAULT_EMBEDDING_BATCH_NUM,
|
||||
DEFAULT_OLLAMA_MODEL_NAME,
|
||||
DEFAULT_OLLAMA_MODEL_TAG,
|
||||
DEFAULT_RERANK_BINDING,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
DEFAULT_EMBEDDING_TIMEOUT,
|
||||
DEFAULT_RERANK_TIMEOUT,
|
||||
)
|
||||
|
||||
# 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)
|
||||
|
||||
|
||||
ollama_server_infos = OllamaServerInfos()
|
||||
DEFAULT_TOKEN_SECRET = "lightrag-jwt-default-secret-key!"
|
||||
NO_PREFIX_SENTINEL = "NO_PREFIX"
|
||||
PROVIDER_ASYMMETRIC_EMBEDDING_BINDINGS = {"gemini", "jina", "voyageai"}
|
||||
PREFIX_ASYMMETRIC_EMBEDDING_BINDINGS = {"azure_openai", "ollama", "openai"}
|
||||
|
||||
|
||||
class DefaultRAGStorageConfig:
|
||||
KV_STORAGE = "JsonKVStorage"
|
||||
VECTOR_STORAGE = "NanoVectorDBStorage"
|
||||
GRAPH_STORAGE = "NetworkXStorage"
|
||||
DOC_STATUS_STORAGE = "JsonDocStatusStorage"
|
||||
|
||||
|
||||
def get_default_host(binding_type: str) -> str:
|
||||
default_hosts = {
|
||||
"ollama": os.getenv("LLM_BINDING_HOST", "http://localhost:11434"),
|
||||
"lollms": os.getenv("LLM_BINDING_HOST", "http://localhost:9600"),
|
||||
"azure_openai": os.getenv("AZURE_OPENAI_ENDPOINT", "https://api.openai.com/v1"),
|
||||
"openai": os.getenv("LLM_BINDING_HOST", "https://api.openai.com/v1"),
|
||||
# Let boto3 select the regional Bedrock endpoint unless the user
|
||||
# explicitly overrides LLM_BINDING_HOST / EMBEDDING_BINDING_HOST.
|
||||
"bedrock": os.getenv("LLM_BINDING_HOST", "DEFAULT_BEDROCK_ENDPOINT"),
|
||||
# Let google-genai pick the correct default endpoint/version unless the
|
||||
# user explicitly overrides LLM_BINDING_HOST / EMBEDDING_BINDING_HOST.
|
||||
"gemini": os.getenv("LLM_BINDING_HOST", "DEFAULT_GEMINI_ENDPOINT"),
|
||||
}
|
||||
return default_hosts.get(
|
||||
binding_type, os.getenv("LLM_BINDING_HOST", "http://localhost:11434")
|
||||
) # fallback to ollama if unknown
|
||||
|
||||
|
||||
def resolve_asymmetric_embedding_opt_in(
|
||||
*,
|
||||
binding: str,
|
||||
embedding_asymmetric: bool,
|
||||
embedding_asymmetric_configured: bool,
|
||||
query_prefix: str | None,
|
||||
document_prefix: str | None,
|
||||
query_prefix_configured: bool = False,
|
||||
document_prefix_configured: bool = False,
|
||||
) -> bool:
|
||||
"""Resolve whether query/document-aware embedding behavior should be enabled."""
|
||||
has_non_empty_prefix = bool(query_prefix or document_prefix)
|
||||
has_prefix_config = query_prefix_configured or document_prefix_configured
|
||||
|
||||
if not embedding_asymmetric:
|
||||
if has_prefix_config:
|
||||
state = "false" if embedding_asymmetric_configured else "unset"
|
||||
logger.warning(
|
||||
f"EMBEDDING_ASYMMETRIC is {state}; "
|
||||
"EMBEDDING_QUERY_PREFIX and EMBEDDING_DOCUMENT_PREFIX will be ignored."
|
||||
)
|
||||
return False
|
||||
|
||||
if binding in PROVIDER_ASYMMETRIC_EMBEDDING_BINDINGS:
|
||||
if has_prefix_config:
|
||||
logger.warning(
|
||||
f"{binding} embeddings use provider task parameters for asymmetric "
|
||||
"mode; EMBEDDING_QUERY_PREFIX and EMBEDDING_DOCUMENT_PREFIX will be ignored."
|
||||
)
|
||||
return True
|
||||
|
||||
if binding in PREFIX_ASYMMETRIC_EMBEDDING_BINDINGS:
|
||||
if not query_prefix_configured or not document_prefix_configured:
|
||||
raise ValueError(
|
||||
f"EMBEDDING_ASYMMETRIC=true for {binding} embeddings requires both "
|
||||
"EMBEDDING_QUERY_PREFIX and EMBEDDING_DOCUMENT_PREFIX. Use "
|
||||
f"{NO_PREFIX_SENTINEL} for a side that should intentionally have no prefix."
|
||||
)
|
||||
|
||||
if not has_non_empty_prefix:
|
||||
raise ValueError(
|
||||
"At least one of EMBEDDING_QUERY_PREFIX or EMBEDDING_DOCUMENT_PREFIX "
|
||||
f"must be non-empty. Use {NO_PREFIX_SENTINEL} only for the side that "
|
||||
"should intentionally have no prefix."
|
||||
)
|
||||
return True
|
||||
|
||||
raise ValueError(
|
||||
f"EMBEDDING_ASYMMETRIC=true is not supported for {binding} embeddings."
|
||||
)
|
||||
|
||||
|
||||
def get_embedding_prefix_config(env_key: str) -> tuple[str | None, bool]:
|
||||
"""Read an embedding prefix and whether it was explicitly configured."""
|
||||
if env_key not in os.environ:
|
||||
return None, False
|
||||
|
||||
value = os.environ[env_key]
|
||||
if value == "None":
|
||||
return None, False
|
||||
if value == NO_PREFIX_SENTINEL:
|
||||
return "", True
|
||||
if value == "":
|
||||
raise ValueError(
|
||||
f"{env_key} is empty. Use {NO_PREFIX_SENTINEL} to explicitly request "
|
||||
"no prefix, or remove the variable to leave it unconfigured."
|
||||
)
|
||||
return value, True
|
||||
|
||||
|
||||
def validate_auth_configuration(args: argparse.Namespace) -> None:
|
||||
"""Reject insecure JWT auth settings before the API starts."""
|
||||
auth_accounts = (getattr(args, "auth_accounts", "") or "").strip()
|
||||
token_secret = (getattr(args, "token_secret", "") or "").strip()
|
||||
|
||||
if auth_accounts and (not token_secret or token_secret == DEFAULT_TOKEN_SECRET):
|
||||
raise ValueError(
|
||||
"TOKEN_SECRET must be explicitly set to a non-default value when AUTH_ACCOUNTS is configured."
|
||||
)
|
||||
|
||||
|
||||
def _is_set(value: str | None) -> bool:
|
||||
return bool((value or "").strip())
|
||||
|
||||
|
||||
def validate_bedrock_auth_configuration(args: argparse.Namespace) -> None:
|
||||
"""Reject Bedrock configuration with no explicit supported auth source."""
|
||||
bearer_token = os.getenv("AWS_BEARER_TOKEN_BEDROCK")
|
||||
|
||||
def has_valid_auth(prefix: str | None = None) -> bool:
|
||||
if _is_set(bearer_token):
|
||||
return True
|
||||
|
||||
if prefix:
|
||||
role_access_key = getattr(args, f"{prefix}_aws_access_key_id", None)
|
||||
role_secret_key = getattr(args, f"{prefix}_aws_secret_access_key", None)
|
||||
if _is_set(role_access_key) or _is_set(role_secret_key):
|
||||
return _is_set(role_access_key) and _is_set(role_secret_key)
|
||||
|
||||
access_key = getattr(args, "aws_access_key_id", None)
|
||||
secret_key = getattr(args, "aws_secret_access_key", None)
|
||||
return _is_set(access_key) and _is_set(secret_key)
|
||||
|
||||
if getattr(args, "llm_binding", None) == "bedrock":
|
||||
if not has_valid_auth():
|
||||
raise ValueError(
|
||||
"Bedrock LLM binding requires AWS_ACCESS_KEY_ID and "
|
||||
"AWS_SECRET_ACCESS_KEY, or process-level AWS_BEARER_TOKEN_BEDROCK."
|
||||
)
|
||||
if _is_set(getattr(args, "llm_binding_api_key", None)):
|
||||
logging.warning(
|
||||
"LLM_BINDING_API_KEY is set but ignored for Bedrock LLM binding. "
|
||||
"Use SigV4 AWS_* variables or process-level AWS_BEARER_TOKEN_BEDROCK instead."
|
||||
)
|
||||
|
||||
if getattr(args, "embedding_binding", None) == "bedrock":
|
||||
if not has_valid_auth():
|
||||
raise ValueError(
|
||||
"Bedrock embedding binding requires AWS_ACCESS_KEY_ID and "
|
||||
"AWS_SECRET_ACCESS_KEY, or process-level AWS_BEARER_TOKEN_BEDROCK."
|
||||
)
|
||||
if _is_set(getattr(args, "embedding_binding_api_key", None)):
|
||||
logging.warning(
|
||||
"EMBEDDING_BINDING_API_KEY is set but ignored for Bedrock embedding binding. "
|
||||
"Use SigV4 AWS_* variables or process-level AWS_BEARER_TOKEN_BEDROCK instead."
|
||||
)
|
||||
|
||||
for spec in ROLES:
|
||||
role = spec.name
|
||||
if getattr(
|
||||
args, f"{role}_llm_binding", None
|
||||
) == "bedrock" and not has_valid_auth(role):
|
||||
raise ValueError(
|
||||
f"Bedrock role '{role}' requires {spec.env_prefix}_AWS_ACCESS_KEY_ID "
|
||||
f"and {spec.env_prefix}_AWS_SECRET_ACCESS_KEY, global "
|
||||
"AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY, or process-level "
|
||||
"AWS_BEARER_TOKEN_BEDROCK."
|
||||
)
|
||||
|
||||
|
||||
def normalize_binding_name(binding: str | None) -> str | None:
|
||||
"""Normalize environment-provided binding aliases to canonical names."""
|
||||
if binding == "aws_bedrock":
|
||||
return "bedrock"
|
||||
return binding
|
||||
|
||||
|
||||
def get_binding_env_value(env_key: str, default: str) -> str:
|
||||
"""Read a binding env var and normalize legacy aliases."""
|
||||
return normalize_binding_name(get_env_value(env_key, default)) or default
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
"""
|
||||
Parse command line arguments with environment variable fallback
|
||||
|
||||
Args:
|
||||
is_uvicorn_mode: Whether running under uvicorn mode
|
||||
|
||||
Returns:
|
||||
argparse.Namespace: Parsed arguments
|
||||
"""
|
||||
|
||||
parser = argparse.ArgumentParser(description="LightRAG API Server")
|
||||
|
||||
# Server configuration
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
default=get_env_value("HOST", "0.0.0.0"),
|
||||
help="Server host (default: from env or 0.0.0.0)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=get_env_value("PORT", 9621, int),
|
||||
help="Server port (default: from env or 9621)",
|
||||
)
|
||||
|
||||
# Directory configuration
|
||||
parser.add_argument(
|
||||
"--working-dir",
|
||||
default=get_env_value("WORKING_DIR", "./rag_storage"),
|
||||
help="Working directory for RAG storage (default: from env or ./rag_storage)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-dir",
|
||||
default=get_env_value("INPUT_DIR", "./inputs"),
|
||||
help="Directory containing input documents (default: from env or ./inputs)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--timeout",
|
||||
default=get_env_value("TIMEOUT", DEFAULT_TIMEOUT, int, special_none=True),
|
||||
type=int,
|
||||
help="Timeout in seconds (useful when using slow AI). Use None for infinite timeout",
|
||||
)
|
||||
|
||||
# RAG configuration
|
||||
parser.add_argument(
|
||||
"--max-async",
|
||||
type=int,
|
||||
default=get_env_value(
|
||||
"MAX_ASYNC_LLM", get_env_value("MAX_ASYNC", DEFAULT_MAX_ASYNC, int), int
|
||||
),
|
||||
help=f"Maximum async operations (default: from env or {DEFAULT_MAX_ASYNC})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--summary-max-tokens",
|
||||
type=int,
|
||||
default=get_env_value("SUMMARY_MAX_TOKENS", DEFAULT_SUMMARY_MAX_TOKENS, int),
|
||||
help=f"Maximum token size for entity/relation summary(default: from env or {DEFAULT_SUMMARY_MAX_TOKENS})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--summary-context-size",
|
||||
type=int,
|
||||
default=get_env_value(
|
||||
"SUMMARY_CONTEXT_SIZE", DEFAULT_SUMMARY_CONTEXT_SIZE, int
|
||||
),
|
||||
help=f"LLM Summary Context size (default: from env or {DEFAULT_SUMMARY_CONTEXT_SIZE})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--summary-length-recommended",
|
||||
type=int,
|
||||
default=get_env_value(
|
||||
"SUMMARY_LENGTH_RECOMMENDED", DEFAULT_SUMMARY_LENGTH_RECOMMENDED, int
|
||||
),
|
||||
help=f"LLM Summary Context size (default: from env or {DEFAULT_SUMMARY_LENGTH_RECOMMENDED})",
|
||||
)
|
||||
|
||||
# Logging configuration
|
||||
parser.add_argument(
|
||||
"--log-level",
|
||||
default=get_env_value("LOG_LEVEL", "INFO"),
|
||||
choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
|
||||
help="Logging level (default: from env or INFO)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
action="store_true",
|
||||
default=get_env_value("VERBOSE", False, bool),
|
||||
help="Enable verbose debug output(only valid for DEBUG log-level)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--key",
|
||||
type=str,
|
||||
default=get_env_value("LIGHTRAG_API_KEY", None),
|
||||
help="API key for authentication. This protects lightrag server against unauthorized access",
|
||||
)
|
||||
|
||||
# Optional https parameters
|
||||
parser.add_argument(
|
||||
"--ssl",
|
||||
action="store_true",
|
||||
default=get_env_value("SSL", False, bool),
|
||||
help="Enable HTTPS (default: from env or False)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ssl-certfile",
|
||||
default=get_env_value("SSL_CERTFILE", None),
|
||||
help="Path to SSL certificate file (required if --ssl is enabled)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ssl-keyfile",
|
||||
default=get_env_value("SSL_KEYFILE", None),
|
||||
help="Path to SSL private key file (required if --ssl is enabled)",
|
||||
)
|
||||
|
||||
# Ollama model configuration
|
||||
parser.add_argument(
|
||||
"--simulated-model-name",
|
||||
type=str,
|
||||
default=get_env_value("OLLAMA_EMULATING_MODEL_NAME", DEFAULT_OLLAMA_MODEL_NAME),
|
||||
help="Name for the simulated Ollama model (default: from env or lightrag)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--simulated-model-tag",
|
||||
type=str,
|
||||
default=get_env_value("OLLAMA_EMULATING_MODEL_TAG", DEFAULT_OLLAMA_MODEL_TAG),
|
||||
help="Tag for the simulated Ollama model (default: from env or latest)",
|
||||
)
|
||||
|
||||
# Namespace
|
||||
parser.add_argument(
|
||||
"--workspace",
|
||||
type=str,
|
||||
default=get_env_value("WORKSPACE", ""),
|
||||
help="Default workspace for all storage",
|
||||
)
|
||||
|
||||
# Path prefix configuration
|
||||
parser.add_argument(
|
||||
"--api-prefix",
|
||||
type=str,
|
||||
default=get_env_value("LIGHTRAG_API_PREFIX", ""),
|
||||
help="API path prefix (e.g., /api/v1). Prepended to all API routes. Default: none (root).",
|
||||
)
|
||||
|
||||
# Server workers configuration
|
||||
parser.add_argument(
|
||||
"--workers",
|
||||
type=int,
|
||||
default=get_env_value("WORKERS", DEFAULT_WOKERS, int),
|
||||
help="Number of worker processes (default: from env or 1)",
|
||||
)
|
||||
|
||||
# LLM and embedding bindings
|
||||
parser.add_argument(
|
||||
"--llm-binding",
|
||||
type=str,
|
||||
default=get_binding_env_value("LLM_BINDING", "ollama"),
|
||||
choices=[
|
||||
"lollms",
|
||||
"ollama",
|
||||
"openai",
|
||||
"openai-ollama",
|
||||
"azure_openai",
|
||||
"bedrock",
|
||||
"gemini",
|
||||
],
|
||||
help="LLM binding type (default: from env or ollama)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--embedding-binding",
|
||||
type=str,
|
||||
default=get_binding_env_value("EMBEDDING_BINDING", "ollama"),
|
||||
choices=[
|
||||
"lollms",
|
||||
"ollama",
|
||||
"openai",
|
||||
"azure_openai",
|
||||
"bedrock",
|
||||
"jina",
|
||||
"gemini",
|
||||
"voyageai",
|
||||
],
|
||||
help="Embedding binding type (default: from env or ollama)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rerank-binding",
|
||||
type=str,
|
||||
default=get_env_value("RERANK_BINDING", DEFAULT_RERANK_BINDING),
|
||||
choices=["null", "cohere", "jina", "aliyun"],
|
||||
help=f"Rerank binding type (default: from env or {DEFAULT_RERANK_BINDING})",
|
||||
)
|
||||
|
||||
# Conditionally add binding-specific options (Ollama, OpenAI, Azure OpenAI, Gemini)
|
||||
# This registers command line arguments (e.g., --openai-llm-temperature)
|
||||
# and reads corresponding environment variables (e.g., OPENAI_LLM_TEMPERATURE)
|
||||
|
||||
# Determine LLM binding value consistently from command line or environment
|
||||
llm_binding_value = None
|
||||
if "--llm-binding" in sys.argv:
|
||||
try:
|
||||
idx = sys.argv.index("--llm-binding")
|
||||
if idx + 1 < len(sys.argv) and not sys.argv[idx + 1].startswith("-"):
|
||||
llm_binding_value = sys.argv[idx + 1]
|
||||
except IndexError:
|
||||
pass
|
||||
|
||||
# Fall back to environment variable using same function as argparse default
|
||||
if llm_binding_value is None:
|
||||
llm_binding_value = get_binding_env_value("LLM_BINDING", "ollama")
|
||||
|
||||
# Add LLM binding options based on determined value
|
||||
if llm_binding_value == "ollama":
|
||||
OllamaLLMOptions.add_args(parser)
|
||||
elif llm_binding_value in ["openai", "azure_openai"]:
|
||||
OpenAILLMOptions.add_args(parser)
|
||||
elif llm_binding_value == "gemini":
|
||||
GeminiLLMOptions.add_args(parser)
|
||||
elif llm_binding_value == "bedrock":
|
||||
BedrockLLMOptions.add_args(parser)
|
||||
|
||||
# Determine embedding binding value consistently from command line or environment
|
||||
embedding_binding_value = None
|
||||
if "--embedding-binding" in sys.argv:
|
||||
try:
|
||||
idx = sys.argv.index("--embedding-binding")
|
||||
if idx + 1 < len(sys.argv) and not sys.argv[idx + 1].startswith("-"):
|
||||
embedding_binding_value = sys.argv[idx + 1]
|
||||
except IndexError:
|
||||
pass
|
||||
|
||||
# Fall back to environment variable using same function as argparse default
|
||||
if embedding_binding_value is None:
|
||||
embedding_binding_value = get_binding_env_value("EMBEDDING_BINDING", "ollama")
|
||||
|
||||
# Add embedding binding options based on determined value
|
||||
if embedding_binding_value == "ollama":
|
||||
OllamaEmbeddingOptions.add_args(parser)
|
||||
elif embedding_binding_value == "gemini":
|
||||
GeminiEmbeddingOptions.add_args(parser)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# convert relative path to absolute path
|
||||
args.working_dir = os.path.abspath(args.working_dir)
|
||||
args.input_dir = os.path.abspath(args.input_dir)
|
||||
|
||||
# Inject storage configuration from environment variables
|
||||
args.kv_storage = get_env_value(
|
||||
"LIGHTRAG_KV_STORAGE", DefaultRAGStorageConfig.KV_STORAGE
|
||||
)
|
||||
args.doc_status_storage = get_env_value(
|
||||
"LIGHTRAG_DOC_STATUS_STORAGE", DefaultRAGStorageConfig.DOC_STATUS_STORAGE
|
||||
)
|
||||
args.graph_storage = get_env_value(
|
||||
"LIGHTRAG_GRAPH_STORAGE", DefaultRAGStorageConfig.GRAPH_STORAGE
|
||||
)
|
||||
args.vector_storage = get_env_value(
|
||||
"LIGHTRAG_VECTOR_STORAGE", DefaultRAGStorageConfig.VECTOR_STORAGE
|
||||
)
|
||||
|
||||
# Get MAX_PARALLEL_INSERT from environment
|
||||
args.max_parallel_insert = get_env_value(
|
||||
"MAX_PARALLEL_INSERT", DEFAULT_MAX_PARALLEL_INSERT, int
|
||||
)
|
||||
|
||||
# Get MAX_GRAPH_NODES from environment
|
||||
args.max_graph_nodes = get_env_value("MAX_GRAPH_NODES", 1000, int)
|
||||
|
||||
# Handle openai-ollama special case
|
||||
if args.llm_binding == "openai-ollama":
|
||||
args.llm_binding = "openai"
|
||||
args.embedding_binding = "ollama"
|
||||
|
||||
args.llm_binding_host = get_env_value(
|
||||
"LLM_BINDING_HOST", get_default_host(args.llm_binding)
|
||||
)
|
||||
args.embedding_binding_host = get_env_value(
|
||||
"EMBEDDING_BINDING_HOST", get_default_host(args.embedding_binding)
|
||||
)
|
||||
args.llm_binding_api_key = get_env_value("LLM_BINDING_API_KEY", None)
|
||||
args.embedding_binding_api_key = get_env_value("EMBEDDING_BINDING_API_KEY", "")
|
||||
|
||||
args.aws_region = get_env_value("AWS_REGION", None, special_none=True)
|
||||
args.aws_access_key_id = get_env_value("AWS_ACCESS_KEY_ID", None, special_none=True)
|
||||
args.aws_secret_access_key = get_env_value(
|
||||
"AWS_SECRET_ACCESS_KEY", None, special_none=True
|
||||
)
|
||||
args.aws_session_token = get_env_value("AWS_SESSION_TOKEN", None, special_none=True)
|
||||
|
||||
# Inject model configuration
|
||||
args.llm_model = get_env_value("LLM_MODEL", "mistral-nemo:latest")
|
||||
# EMBEDDING_MODEL defaults to None - each binding will use its own default model
|
||||
# e.g., OpenAI uses "text-embedding-3-small", Jina uses "jina-embeddings-v4"
|
||||
args.embedding_model = get_env_value("EMBEDDING_MODEL", None, special_none=True)
|
||||
# EMBEDDING_DIM defaults to None - each binding will use its own default dimension
|
||||
# Value is inherited from provider defaults via wrap_embedding_func_with_attrs decorator
|
||||
args.embedding_dim = get_env_value("EMBEDDING_DIM", None, int, special_none=True)
|
||||
args.embedding_send_dim = get_env_value("EMBEDDING_SEND_DIM", False, bool)
|
||||
|
||||
# Inject chunk configuration
|
||||
args.chunk_size = get_env_value("CHUNK_SIZE", 1200, int)
|
||||
args.chunk_overlap_size = get_env_value("CHUNK_OVERLAP_SIZE", 100, int)
|
||||
|
||||
# Inject LLM cache configuration
|
||||
# Should not be disabled； LLM cache is required for entity/realtion rebuild after file deletion.
|
||||
args.enable_llm_cache_for_extract = get_env_value(
|
||||
"ENABLE_LLM_CACHE_FOR_EXTRACT", True, bool
|
||||
)
|
||||
args.enable_llm_cache = get_env_value("ENABLE_LLM_CACHE", True, bool)
|
||||
|
||||
# --- Per-role LLM configuration (driven by lightrag.ROLES registry) ---
|
||||
for spec in ROLES:
|
||||
prefix = spec.env_prefix
|
||||
attr_prefix = spec.name
|
||||
binding_key = f"{prefix}_LLM_BINDING"
|
||||
model_key = f"{prefix}_LLM_MODEL"
|
||||
host_key = f"{prefix}_LLM_BINDING_HOST"
|
||||
apikey_key = f"{prefix}_LLM_BINDING_API_KEY"
|
||||
max_async_key = f"{prefix}_MAX_ASYNC_LLM"
|
||||
timeout_key = f"{prefix}_LLM_TIMEOUT"
|
||||
|
||||
role_binding = normalize_binding_name(
|
||||
get_env_value(binding_key, None, special_none=True)
|
||||
)
|
||||
role_model = get_env_value(model_key, None, special_none=True)
|
||||
role_host = get_env_value(host_key, None, special_none=True)
|
||||
role_apikey = get_env_value(apikey_key, None, special_none=True)
|
||||
role_max_async = get_env_value(max_async_key, None, int, special_none=True)
|
||||
role_timeout = get_env_value(timeout_key, None, int, special_none=True)
|
||||
role_aws_region = get_env_value(f"{prefix}_AWS_REGION", None, special_none=True)
|
||||
role_aws_access_key_id = get_env_value(
|
||||
f"{prefix}_AWS_ACCESS_KEY_ID", None, special_none=True
|
||||
)
|
||||
role_aws_secret_access_key = get_env_value(
|
||||
f"{prefix}_AWS_SECRET_ACCESS_KEY", None, special_none=True
|
||||
)
|
||||
role_aws_session_token = get_env_value(
|
||||
f"{prefix}_AWS_SESSION_TOKEN", None, special_none=True
|
||||
)
|
||||
|
||||
setattr(args, f"{attr_prefix}_llm_binding", role_binding)
|
||||
setattr(args, f"{attr_prefix}_llm_model", role_model)
|
||||
setattr(args, f"{attr_prefix}_llm_binding_host", role_host)
|
||||
setattr(args, f"{attr_prefix}_llm_binding_api_key", role_apikey)
|
||||
setattr(args, f"{attr_prefix}_llm_max_async", role_max_async)
|
||||
setattr(args, f"{attr_prefix}_llm_timeout", role_timeout)
|
||||
setattr(args, f"{attr_prefix}_aws_region", role_aws_region)
|
||||
setattr(args, f"{attr_prefix}_aws_access_key_id", role_aws_access_key_id)
|
||||
setattr(
|
||||
args, f"{attr_prefix}_aws_secret_access_key", role_aws_secret_access_key
|
||||
)
|
||||
setattr(args, f"{attr_prefix}_aws_session_token", role_aws_session_token)
|
||||
|
||||
if role_binding == "bedrock" and role_apikey:
|
||||
raise SystemExit(
|
||||
f"Bedrock role '{spec.name}' does not support {apikey_key}; use "
|
||||
"role-specific SigV4 AWS_* variables or process-level "
|
||||
"AWS_BEARER_TOKEN_BEDROCK."
|
||||
)
|
||||
|
||||
# Cross-provider validation
|
||||
if role_binding and role_binding != args.llm_binding:
|
||||
missing = []
|
||||
if not role_model:
|
||||
missing.append(model_key)
|
||||
if not role_host:
|
||||
role_host = get_default_host(role_binding)
|
||||
setattr(args, f"{attr_prefix}_llm_binding_host", role_host)
|
||||
if role_binding != "bedrock" and not role_apikey:
|
||||
missing.append(apikey_key)
|
||||
if missing:
|
||||
raise SystemExit(
|
||||
f"Cross-provider error for role '{spec.name}': "
|
||||
f"binding={role_binding} differs from base={args.llm_binding}, "
|
||||
f"but required env vars are missing: {', '.join(missing)}"
|
||||
)
|
||||
|
||||
# VLM multimodal master switch — when off, the pipeline emits a warning
|
||||
# and skips every i/t/e item without touching the VLM. When on, the
|
||||
# effective VLM binding must support image inputs.
|
||||
args.vlm_process_enable = get_env_value("VLM_PROCESS_ENABLE", False, bool)
|
||||
if args.vlm_process_enable:
|
||||
effective_vlm_binding = (
|
||||
getattr(args, "vlm_llm_binding", None) or args.llm_binding
|
||||
)
|
||||
vlm_incompatible = {"lollms"}
|
||||
if effective_vlm_binding in vlm_incompatible:
|
||||
raise SystemExit(
|
||||
f"VLM_PROCESS_ENABLE=true but the effective VLM binding "
|
||||
f"'{effective_vlm_binding}' does not support image inputs. "
|
||||
"Configure VLM_LLM_BINDING (or LLM_BINDING) to one of: "
|
||||
"openai, azure_openai, gemini, bedrock, ollama."
|
||||
)
|
||||
|
||||
# Add environment variables that were previously read directly
|
||||
args.cors_origins = get_env_value("CORS_ORIGINS", "*")
|
||||
args.summary_language = get_env_value("SUMMARY_LANGUAGE", DEFAULT_SUMMARY_LANGUAGE)
|
||||
args.whitelist_paths = get_env_value("WHITELIST_PATHS", "/health,/api/*")
|
||||
|
||||
# For JWT Auth
|
||||
args.auth_accounts = get_env_value("AUTH_ACCOUNTS", "")
|
||||
args.token_secret = get_env_value("TOKEN_SECRET", None)
|
||||
args.token_expire_hours = get_env_value("TOKEN_EXPIRE_HOURS", 48, float)
|
||||
args.guest_token_expire_hours = get_env_value("GUEST_TOKEN_EXPIRE_HOURS", 24, float)
|
||||
args.jwt_algorithm = get_env_value("JWT_ALGORITHM", "HS256")
|
||||
|
||||
# Token auto-renewal configuration (sliding window expiration)
|
||||
args.token_auto_renew = get_env_value("TOKEN_AUTO_RENEW", True, bool)
|
||||
args.token_renew_threshold = get_env_value("TOKEN_RENEW_THRESHOLD", 0.5, float)
|
||||
|
||||
# Rerank model configuration
|
||||
args.rerank_model = get_env_value("RERANK_MODEL", None)
|
||||
args.rerank_binding_host = get_env_value("RERANK_BINDING_HOST", None)
|
||||
args.rerank_binding_api_key = get_env_value("RERANK_BINDING_API_KEY", None)
|
||||
# Note: rerank_binding is already set by argparse, no need to override from env
|
||||
|
||||
# Min rerank score configuration
|
||||
args.min_rerank_score = get_env_value(
|
||||
"MIN_RERANK_SCORE", DEFAULT_MIN_RERANK_SCORE, float
|
||||
)
|
||||
|
||||
# LLM / Embedding request timeouts
|
||||
args.llm_timeout = get_env_value("LLM_TIMEOUT", DEFAULT_LLM_TIMEOUT, int)
|
||||
args.embedding_timeout = get_env_value(
|
||||
"EMBEDDING_TIMEOUT", DEFAULT_EMBEDDING_TIMEOUT, int
|
||||
)
|
||||
|
||||
# Rerank async/timeout configuration (independent from base LLM)
|
||||
# rerank_max_async falls back to MAX_ASYNC_LLM; rerank_timeout has its own default.
|
||||
args.rerank_max_async = get_env_value("MAX_ASYNC_RERANK", args.max_async, int)
|
||||
args.rerank_timeout = get_env_value("RERANK_TIMEOUT", DEFAULT_RERANK_TIMEOUT, int)
|
||||
|
||||
# Query configuration
|
||||
args.top_k = get_env_value("TOP_K", DEFAULT_TOP_K, int)
|
||||
args.chunk_top_k = get_env_value("CHUNK_TOP_K", DEFAULT_CHUNK_TOP_K, int)
|
||||
args.max_entity_tokens = get_env_value(
|
||||
"MAX_ENTITY_TOKENS", DEFAULT_MAX_ENTITY_TOKENS, int
|
||||
)
|
||||
args.max_relation_tokens = get_env_value(
|
||||
"MAX_RELATION_TOKENS", DEFAULT_MAX_RELATION_TOKENS, int
|
||||
)
|
||||
args.max_total_tokens = get_env_value(
|
||||
"MAX_TOTAL_TOKENS", DEFAULT_MAX_TOTAL_TOKENS, int
|
||||
)
|
||||
args.cosine_threshold = get_env_value(
|
||||
"COSINE_THRESHOLD", DEFAULT_COSINE_THRESHOLD, float
|
||||
)
|
||||
args.related_chunk_number = get_env_value(
|
||||
"RELATED_CHUNK_NUMBER", DEFAULT_RELATED_CHUNK_NUMBER, int
|
||||
)
|
||||
|
||||
# Add missing environment variables for health endpoint
|
||||
args.force_llm_summary_on_merge = get_env_value(
|
||||
"FORCE_LLM_SUMMARY_ON_MERGE", DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE, int
|
||||
)
|
||||
args.embedding_func_max_async = get_env_value(
|
||||
"EMBEDDING_FUNC_MAX_ASYNC", DEFAULT_EMBEDDING_FUNC_MAX_ASYNC, int
|
||||
)
|
||||
args.embedding_batch_num = get_env_value(
|
||||
"EMBEDDING_BATCH_NUM", DEFAULT_EMBEDDING_BATCH_NUM, int
|
||||
)
|
||||
|
||||
# Embedding token limit configuration
|
||||
args.embedding_token_limit = get_env_value(
|
||||
"EMBEDDING_TOKEN_LIMIT", None, int, special_none=True
|
||||
)
|
||||
|
||||
# File upload size limit (in bytes, None for unlimited)
|
||||
# Default: 100MB (104857600 bytes)
|
||||
args.max_upload_size = get_env_value(
|
||||
"MAX_UPLOAD_SIZE", 104857600, int, special_none=True
|
||||
)
|
||||
|
||||
# Embedding prefix configuration for context-aware embeddings. Empty prefixes
|
||||
# must be explicit via NO_PREFIX so missing config is distinguishable.
|
||||
(
|
||||
args.embedding_document_prefix,
|
||||
args.embedding_document_prefix_configured,
|
||||
) = get_embedding_prefix_config("EMBEDDING_DOCUMENT_PREFIX")
|
||||
(
|
||||
args.embedding_query_prefix,
|
||||
args.embedding_query_prefix_configured,
|
||||
) = get_embedding_prefix_config("EMBEDDING_QUERY_PREFIX")
|
||||
args.embedding_prefix_no_prefix_sentinel = NO_PREFIX_SENTINEL
|
||||
args.embedding_prefixes_configured = (
|
||||
args.embedding_document_prefix_configured
|
||||
or args.embedding_query_prefix_configured
|
||||
)
|
||||
# Asymmetric embedding behavior toggle
|
||||
args.embedding_asymmetric_configured = "EMBEDDING_ASYMMETRIC" in os.environ
|
||||
args.embedding_asymmetric = get_env_value("EMBEDDING_ASYMMETRIC", False, bool)
|
||||
|
||||
ollama_server_infos.LIGHTRAG_NAME = args.simulated_model_name
|
||||
ollama_server_infos.LIGHTRAG_TAG = args.simulated_model_tag
|
||||
|
||||
# Sanitize workspace: only alphanumeric characters and underscores are allowed
|
||||
if args.workspace:
|
||||
sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", args.workspace)
|
||||
if sanitized != args.workspace:
|
||||
logging.warning(
|
||||
f"Workspace name '{args.workspace}' contains invalid characters. "
|
||||
f"It has been sanitized to '{sanitized}'. "
|
||||
"Only alphanumeric characters and underscores are allowed."
|
||||
)
|
||||
args.workspace = sanitized
|
||||
|
||||
validate_auth_configuration(args)
|
||||
validate_bedrock_auth_configuration(args)
|
||||
return args
|
||||
|
||||
|
||||
def update_uvicorn_mode_config():
|
||||
# If in uvicorn mode and workers > 1, force it to 1 and log warning
|
||||
if global_args.workers > 1:
|
||||
original_workers = global_args.workers
|
||||
global_args.workers = 1
|
||||
# Log warning directly here
|
||||
logging.debug(
|
||||
f">> Forcing workers=1 in uvicorn mode(Ignoring workers={original_workers})"
|
||||
)
|
||||
|
||||
|
||||
# Global configuration with lazy initialization
|
||||
_global_args = None
|
||||
_initialized = False
|
||||
|
||||
|
||||
def initialize_config(args=None, force=False):
|
||||
"""Initialize global configuration
|
||||
|
||||
This function allows explicit initialization of the configuration,
|
||||
which is useful for programmatic usage, testing, or embedding LightRAG
|
||||
in other applications.
|
||||
|
||||
Args:
|
||||
args: Pre-parsed argparse.Namespace or None to parse from sys.argv
|
||||
force: Force re-initialization even if already initialized
|
||||
|
||||
Returns:
|
||||
argparse.Namespace: The configured arguments
|
||||
|
||||
Example:
|
||||
# Use parsed command line arguments (default)
|
||||
initialize_config()
|
||||
|
||||
# Use custom configuration programmatically
|
||||
custom_args = argparse.Namespace(
|
||||
host='localhost',
|
||||
port=8080,
|
||||
working_dir='./custom_rag',
|
||||
# ... other config
|
||||
)
|
||||
initialize_config(custom_args)
|
||||
"""
|
||||
global _global_args, _initialized
|
||||
|
||||
if _initialized and not force:
|
||||
return _global_args
|
||||
|
||||
resolved_args = args if args is not None else parse_args()
|
||||
validate_auth_configuration(resolved_args)
|
||||
validate_bedrock_auth_configuration(resolved_args)
|
||||
_global_args = resolved_args
|
||||
_initialized = True
|
||||
return _global_args
|
||||
|
||||
|
||||
def get_config():
|
||||
"""Get global configuration, auto-initializing if needed
|
||||
|
||||
Returns:
|
||||
argparse.Namespace: The configured arguments
|
||||
"""
|
||||
if not _initialized:
|
||||
initialize_config()
|
||||
return _global_args
|
||||
|
||||
|
||||
class _GlobalArgsProxy:
|
||||
"""Proxy object that auto-initializes configuration on first access
|
||||
|
||||
This maintains backward compatibility with existing code while
|
||||
allowing programmatic control over initialization timing.
|
||||
|
||||
The proxy fully delegates to the underlying argparse.Namespace,
|
||||
including support for vars() calls which is used by binding_options
|
||||
to extract provider-specific configuration options.
|
||||
"""
|
||||
|
||||
def __getattribute__(self, name):
|
||||
"""Override attribute access to support vars() and regular attribute access.
|
||||
|
||||
This method intercepts __dict__ access (used by vars()) and delegates
|
||||
to the underlying _global_args namespace, ensuring binding options
|
||||
can be properly extracted.
|
||||
"""
|
||||
global _initialized, _global_args
|
||||
|
||||
# Handle __dict__ access for vars() support
|
||||
if name == "__dict__":
|
||||
if not _initialized:
|
||||
initialize_config()
|
||||
return vars(_global_args)
|
||||
|
||||
# Handle class-level attributes that should come from the proxy itself
|
||||
if name in ("__class__", "__repr__", "__getattribute__", "__setattr__"):
|
||||
return object.__getattribute__(self, name)
|
||||
|
||||
# Delegate all other attribute access to the underlying namespace
|
||||
if not _initialized:
|
||||
initialize_config()
|
||||
return getattr(_global_args, name)
|
||||
|
||||
def __setattr__(self, name, value):
|
||||
global _initialized, _global_args
|
||||
if not _initialized:
|
||||
initialize_config()
|
||||
setattr(_global_args, name, value)
|
||||
|
||||
def __repr__(self):
|
||||
global _initialized, _global_args
|
||||
if not _initialized:
|
||||
return "<GlobalArgsProxy: Not initialized>"
|
||||
return repr(_global_args)
|
||||
|
||||
|
||||
# Create proxy instance for backward compatibility
|
||||
# Existing code like `from config import global_args` continues to work
|
||||
# The proxy will auto-initialize on first attribute access
|
||||
global_args = _GlobalArgsProxy()
|
||||
@@ -0,0 +1,162 @@
|
||||
# gunicorn_config.py
|
||||
import os
|
||||
import logging
|
||||
from lightrag.kg.shared_storage import finalize_share_data
|
||||
from lightrag.utils import setup_logger, get_env_value
|
||||
from lightrag.constants import (
|
||||
DEFAULT_LOG_MAX_BYTES,
|
||||
DEFAULT_LOG_BACKUP_COUNT,
|
||||
DEFAULT_LOG_FILENAME,
|
||||
)
|
||||
|
||||
|
||||
# Get log directory path from environment variable
|
||||
log_dir = os.getenv("LOG_DIR", os.getcwd())
|
||||
log_file_path = os.path.abspath(os.path.join(log_dir, DEFAULT_LOG_FILENAME))
|
||||
|
||||
# Ensure log directory exists
|
||||
os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
|
||||
|
||||
# Get log file max size and backup count from environment variables
|
||||
log_max_bytes = get_env_value("LOG_MAX_BYTES", DEFAULT_LOG_MAX_BYTES, int)
|
||||
log_backup_count = get_env_value("LOG_BACKUP_COUNT", DEFAULT_LOG_BACKUP_COUNT, int)
|
||||
|
||||
# These variables will be set by run_with_gunicorn.py
|
||||
workers = None
|
||||
bind = None
|
||||
loglevel = None
|
||||
certfile = None
|
||||
keyfile = None
|
||||
|
||||
# Enable preload_app option
|
||||
preload_app = True
|
||||
|
||||
# Use Uvicorn worker
|
||||
worker_class = "uvicorn.workers.UvicornWorker"
|
||||
|
||||
# Other Gunicorn configurations
|
||||
|
||||
# Logging configuration
|
||||
errorlog = os.getenv("ERROR_LOG", log_file_path) # Default write to lightrag.log
|
||||
accesslog = os.getenv("ACCESS_LOG", log_file_path) # Default write to lightrag.log
|
||||
|
||||
logconfig_dict = {
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
"formatters": {
|
||||
"standard": {"format": "%(asctime)s [%(levelname)s] %(name)s: %(message)s"},
|
||||
},
|
||||
"handlers": {
|
||||
"console": {
|
||||
"class": "logging.StreamHandler",
|
||||
"formatter": "standard",
|
||||
"stream": "ext://sys.stdout",
|
||||
},
|
||||
"file": {
|
||||
"class": "logging.handlers.RotatingFileHandler",
|
||||
"formatter": "standard",
|
||||
"filename": log_file_path,
|
||||
"maxBytes": log_max_bytes,
|
||||
"backupCount": log_backup_count,
|
||||
"encoding": "utf8",
|
||||
},
|
||||
},
|
||||
"filters": {
|
||||
"path_filter": {
|
||||
"()": "lightrag.utils.LightragPathFilter",
|
||||
},
|
||||
},
|
||||
"loggers": {
|
||||
"lightrag": {
|
||||
"handlers": ["console", "file"],
|
||||
"level": loglevel.upper() if loglevel else "INFO",
|
||||
"propagate": False,
|
||||
},
|
||||
"gunicorn": {
|
||||
"handlers": ["console", "file"],
|
||||
"level": loglevel.upper() if loglevel else "INFO",
|
||||
"propagate": False,
|
||||
},
|
||||
"gunicorn.error": {
|
||||
"handlers": ["console", "file"],
|
||||
"level": loglevel.upper() if loglevel else "INFO",
|
||||
"propagate": False,
|
||||
},
|
||||
"gunicorn.access": {
|
||||
"handlers": ["console", "file"],
|
||||
"level": loglevel.upper() if loglevel else "INFO",
|
||||
"propagate": False,
|
||||
"filters": ["path_filter"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def on_starting(server):
|
||||
"""
|
||||
Executed when Gunicorn starts, before forking the first worker processes
|
||||
You can use this function to do more initialization tasks for all processes
|
||||
"""
|
||||
print("=" * 80)
|
||||
print(f"GUNICORN MASTER PROCESS: on_starting jobs for {workers} worker(s)")
|
||||
print(f"Process ID: {os.getpid()}")
|
||||
print("=" * 80)
|
||||
|
||||
# Memory usage monitoring
|
||||
try:
|
||||
import psutil
|
||||
|
||||
process = psutil.Process(os.getpid())
|
||||
memory_info = process.memory_info()
|
||||
msg = (
|
||||
f"Memory usage after initialization: {memory_info.rss / 1024 / 1024:.2f} MB"
|
||||
)
|
||||
print(msg)
|
||||
except ImportError:
|
||||
print("psutil not installed, skipping memory usage reporting")
|
||||
|
||||
# Log the location of the LightRAG log file
|
||||
print(f"LightRAG log file: {log_file_path}\n")
|
||||
|
||||
print("Gunicorn initialization complete, forking workers...\n")
|
||||
|
||||
|
||||
def on_exit(server):
|
||||
"""
|
||||
Executed when Gunicorn is shutting down.
|
||||
This is a good place to release shared resources.
|
||||
"""
|
||||
print("=" * 80)
|
||||
print("GUNICORN MASTER PROCESS: Shutting down")
|
||||
print(f"Process ID: {os.getpid()}")
|
||||
|
||||
print("Finalizing shared storage...")
|
||||
finalize_share_data()
|
||||
|
||||
print("Gunicorn shutdown complete")
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
def post_fork(server, worker):
|
||||
"""
|
||||
Executed after a worker has been forked.
|
||||
This is a good place to set up worker-specific configurations.
|
||||
"""
|
||||
# Set up main loggers
|
||||
log_level = loglevel.upper() if loglevel else "INFO"
|
||||
setup_logger("uvicorn", log_level, add_filter=False, log_file_path=log_file_path)
|
||||
setup_logger(
|
||||
"uvicorn.access", log_level, add_filter=True, log_file_path=log_file_path
|
||||
)
|
||||
setup_logger("lightrag", log_level, add_filter=True, log_file_path=log_file_path)
|
||||
|
||||
# Set up lightrag submodule loggers
|
||||
for name in logging.root.manager.loggerDict:
|
||||
if name.startswith("lightrag."):
|
||||
setup_logger(name, log_level, add_filter=True, log_file_path=log_file_path)
|
||||
|
||||
# Disable uvicorn.error logger
|
||||
uvicorn_error_logger = logging.getLogger("uvicorn.error")
|
||||
uvicorn_error_logger.handlers = []
|
||||
uvicorn_error_logger.setLevel(logging.CRITICAL)
|
||||
uvicorn_error_logger.propagate = False
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,26 @@
|
||||
import bcrypt
|
||||
|
||||
BCRYPT_PASSWORD_PREFIX = "{bcrypt}"
|
||||
|
||||
|
||||
def hash_password(password: str) -> str:
|
||||
"""Return an AUTH_ACCOUNTS-ready bcrypt password value."""
|
||||
salt = bcrypt.gensalt()
|
||||
hashed = bcrypt.hashpw(password.encode("utf-8"), salt).decode("utf-8")
|
||||
return f"{BCRYPT_PASSWORD_PREFIX}{hashed}"
|
||||
|
||||
|
||||
def verify_password(plain_password: str, stored_password: str) -> bool:
|
||||
"""Verify a plaintext password against a stored password spec."""
|
||||
if stored_password.startswith(BCRYPT_PASSWORD_PREFIX):
|
||||
hashed_password = stored_password[len(BCRYPT_PASSWORD_PREFIX) :]
|
||||
if not hashed_password:
|
||||
return False
|
||||
try:
|
||||
return bcrypt.checkpw(
|
||||
plain_password.encode("utf-8"), hashed_password.encode("utf-8")
|
||||
)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
return stored_password == plain_password
|
||||
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
This module contains all the routers for the LightRAG API.
|
||||
|
||||
The document/query/graph routers are intentionally NOT re-exported here:
|
||||
they are constructed per-app via the `create_*_routes` factory functions
|
||||
in their respective modules. A module-level singleton would accumulate
|
||||
duplicate routes if the factory is invoked more than once in the same
|
||||
process (e.g. across tests), which produced "Duplicate Operation ID"
|
||||
warnings before the factories were converted to local routers.
|
||||
"""
|
||||
|
||||
from .ollama_api import OllamaAPI
|
||||
|
||||
__all__ = ["OllamaAPI"]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,807 @@
|
||||
"""
|
||||
This module contains all graph-related routes for the LightRAG API.
|
||||
"""
|
||||
|
||||
from typing import Optional, Dict, Any
|
||||
import traceback
|
||||
from fastapi import APIRouter, Depends, Query, HTTPException
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
from lightrag.base import DeletionResult
|
||||
from lightrag.utils import logger
|
||||
from ..utils_api import get_combined_auth_dependency
|
||||
from .document_routes import check_pipeline_busy_or_raise
|
||||
|
||||
|
||||
class EntityUpdateRequest(BaseModel):
|
||||
entity_name: str
|
||||
updated_data: Dict[str, Any]
|
||||
allow_rename: bool = False
|
||||
allow_merge: bool = False
|
||||
|
||||
|
||||
class RelationUpdateRequest(BaseModel):
|
||||
source_id: str
|
||||
target_id: str
|
||||
updated_data: Dict[str, Any]
|
||||
|
||||
|
||||
class EntityMergeRequest(BaseModel):
|
||||
entities_to_change: list[str] = Field(
|
||||
...,
|
||||
description="List of entity names to be merged and deleted. These are typically duplicate or misspelled entities.",
|
||||
min_length=1,
|
||||
examples=[["Elon Msk", "Ellon Musk"]],
|
||||
)
|
||||
entity_to_change_into: str = Field(
|
||||
...,
|
||||
description="Target entity name that will receive all relationships from the source entities. This entity will be preserved.",
|
||||
min_length=1,
|
||||
examples=["Elon Musk"],
|
||||
)
|
||||
|
||||
|
||||
class EntityCreateRequest(BaseModel):
|
||||
entity_name: str = Field(
|
||||
...,
|
||||
description="Unique name for the new entity",
|
||||
min_length=1,
|
||||
examples=["Tesla"],
|
||||
)
|
||||
entity_data: Dict[str, Any] = Field(
|
||||
...,
|
||||
description="Dictionary containing entity properties. Common fields include 'description' and 'entity_type'.",
|
||||
examples=[
|
||||
{
|
||||
"description": "Electric vehicle manufacturer",
|
||||
"entity_type": "ORGANIZATION",
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class DeleteEntityRequest(BaseModel):
|
||||
entity_name: str = Field(..., description="The name of the entity to delete.")
|
||||
|
||||
@field_validator("entity_name", mode="after")
|
||||
@classmethod
|
||||
def validate_entity_name(cls, entity_name: str) -> str:
|
||||
if not entity_name or not entity_name.strip():
|
||||
raise ValueError("Entity name cannot be empty")
|
||||
return entity_name.strip()
|
||||
|
||||
|
||||
class DeleteRelationRequest(BaseModel):
|
||||
source_entity: str = Field(..., description="The name of the source entity.")
|
||||
target_entity: str = Field(..., description="The name of the target entity.")
|
||||
|
||||
@field_validator("source_entity", "target_entity", mode="after")
|
||||
@classmethod
|
||||
def validate_entity_names(cls, entity_name: str) -> str:
|
||||
if not entity_name or not entity_name.strip():
|
||||
raise ValueError("Entity name cannot be empty")
|
||||
return entity_name.strip()
|
||||
|
||||
|
||||
class RelationCreateRequest(BaseModel):
|
||||
source_entity: str = Field(
|
||||
...,
|
||||
description="Name of the source entity. This entity must already exist in the knowledge graph.",
|
||||
min_length=1,
|
||||
examples=["Elon Musk"],
|
||||
)
|
||||
target_entity: str = Field(
|
||||
...,
|
||||
description="Name of the target entity. This entity must already exist in the knowledge graph.",
|
||||
min_length=1,
|
||||
examples=["Tesla"],
|
||||
)
|
||||
relation_data: Dict[str, Any] = Field(
|
||||
...,
|
||||
description="Dictionary containing relationship properties. Common fields include 'description', 'keywords', and 'weight'.",
|
||||
examples=[
|
||||
{
|
||||
"description": "Elon Musk is the CEO of Tesla",
|
||||
"keywords": "CEO, founder",
|
||||
"weight": 1.0,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def create_graph_routes(rag, api_key: Optional[str] = None):
|
||||
# Fresh router per call. A module-level instance would accumulate
|
||||
# duplicate routes when the factory is invoked more than once in the
|
||||
# same process (e.g. across tests), which triggers FastAPI's
|
||||
# "Duplicate Operation ID" warnings.
|
||||
router = APIRouter(tags=["graph"])
|
||||
|
||||
combined_auth = get_combined_auth_dependency(api_key)
|
||||
|
||||
@router.get("/graph/label/list", dependencies=[Depends(combined_auth)])
|
||||
async def get_graph_labels():
|
||||
"""
|
||||
Get all graph labels
|
||||
|
||||
Returns:
|
||||
List[str]: List of graph labels
|
||||
"""
|
||||
try:
|
||||
return await rag.get_graph_labels()
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting graph labels: {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error getting graph labels: {str(e)}"
|
||||
)
|
||||
|
||||
@router.get("/graph/label/popular", dependencies=[Depends(combined_auth)])
|
||||
async def get_popular_labels(
|
||||
limit: int = Query(
|
||||
300, description="Maximum number of popular labels to return", ge=1, le=1000
|
||||
),
|
||||
):
|
||||
"""
|
||||
Get popular labels by node degree (most connected entities)
|
||||
|
||||
Args:
|
||||
limit (int): Maximum number of labels to return (default: 300, max: 1000)
|
||||
|
||||
Returns:
|
||||
List[str]: List of popular labels sorted by degree (highest first)
|
||||
"""
|
||||
try:
|
||||
return await rag.chunk_entity_relation_graph.get_popular_labels(limit)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting popular labels: {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error getting popular labels: {str(e)}"
|
||||
)
|
||||
|
||||
@router.get("/graph/label/search", dependencies=[Depends(combined_auth)])
|
||||
async def search_labels(
|
||||
q: str = Query(..., description="Search query string"),
|
||||
limit: int = Query(
|
||||
50, description="Maximum number of search results to return", ge=1, le=100
|
||||
),
|
||||
):
|
||||
"""
|
||||
Search labels with fuzzy matching
|
||||
|
||||
Args:
|
||||
q (str): Search query string
|
||||
limit (int): Maximum number of results to return (default: 50, max: 100)
|
||||
|
||||
Returns:
|
||||
List[str]: List of matching labels sorted by relevance
|
||||
"""
|
||||
try:
|
||||
return await rag.chunk_entity_relation_graph.search_labels(q, limit)
|
||||
except Exception as e:
|
||||
logger.error(f"Error searching labels with query '{q}': {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error searching labels: {str(e)}"
|
||||
)
|
||||
|
||||
@router.get("/graphs", dependencies=[Depends(combined_auth)])
|
||||
async def get_knowledge_graph(
|
||||
label: str = Query(..., description="Label to get knowledge graph for"),
|
||||
max_depth: int = Query(3, description="Maximum depth of graph", ge=1),
|
||||
max_nodes: int = Query(1000, description="Maximum nodes to return", ge=1),
|
||||
):
|
||||
"""
|
||||
Retrieve a connected subgraph of nodes where the label includes the specified label.
|
||||
When reducing the number of nodes, the prioritization criteria are as follows:
|
||||
1. Hops(path) to the staring node take precedence
|
||||
2. Followed by the degree of the nodes
|
||||
|
||||
Args:
|
||||
label (str): Label of the starting node
|
||||
max_depth (int, optional): Maximum depth of the subgraph,Defaults to 3
|
||||
max_nodes: Maxiumu nodes to return
|
||||
|
||||
Returns:
|
||||
Dict[str, List[str]]: Knowledge graph for label
|
||||
"""
|
||||
try:
|
||||
# Log the label parameter to check for leading spaces
|
||||
logger.debug(
|
||||
f"get_knowledge_graph called with label: '{label}' (length: {len(label)}, repr: {repr(label)})"
|
||||
)
|
||||
|
||||
return await rag.get_knowledge_graph(
|
||||
node_label=label,
|
||||
max_depth=max_depth,
|
||||
max_nodes=max_nodes,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting knowledge graph for label '{label}': {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error getting knowledge graph: {str(e)}"
|
||||
)
|
||||
|
||||
@router.get("/graph/entity/exists", dependencies=[Depends(combined_auth)])
|
||||
async def check_entity_exists(
|
||||
name: str = Query(..., description="Entity name to check"),
|
||||
):
|
||||
"""
|
||||
Check if an entity with the given name exists in the knowledge graph
|
||||
|
||||
Args:
|
||||
name (str): Name of the entity to check
|
||||
|
||||
Returns:
|
||||
Dict[str, bool]: Dictionary with 'exists' key indicating if entity exists
|
||||
"""
|
||||
try:
|
||||
exists = await rag.chunk_entity_relation_graph.has_node(name)
|
||||
return {"exists": exists}
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking entity existence for '{name}': {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error checking entity existence: {str(e)}"
|
||||
)
|
||||
|
||||
@router.post("/graph/entity/edit", dependencies=[Depends(combined_auth)])
|
||||
async def update_entity(request: EntityUpdateRequest):
|
||||
"""
|
||||
Update an entity's properties in the knowledge graph
|
||||
|
||||
This endpoint allows updating entity properties, including renaming entities.
|
||||
When renaming to an existing entity name, the behavior depends on allow_merge:
|
||||
|
||||
Args:
|
||||
request (EntityUpdateRequest): Request containing:
|
||||
- entity_name (str): Name of the entity to update
|
||||
- updated_data (Dict[str, Any]): Dictionary of properties to update
|
||||
- allow_rename (bool): Whether to allow entity renaming (default: False)
|
||||
- allow_merge (bool): Whether to merge into existing entity when renaming
|
||||
causes name conflict (default: False)
|
||||
|
||||
Returns:
|
||||
Dict with the following structure:
|
||||
{
|
||||
"status": "success",
|
||||
"message": "Entity updated successfully" | "Entity merged successfully into 'target_name'",
|
||||
"data": {
|
||||
"entity_name": str, # Final entity name
|
||||
"description": str, # Entity description
|
||||
"entity_type": str, # Entity type
|
||||
"source_id": str, # Source chunk IDs
|
||||
... # Other entity properties
|
||||
},
|
||||
"operation_summary": {
|
||||
"merged": bool, # Whether entity was merged into another
|
||||
"merge_status": str, # "success" | "failed" | "not_attempted"
|
||||
"merge_error": str | None, # Error message if merge failed
|
||||
"operation_status": str, # "success" | "partial_success" | "failure"
|
||||
"target_entity": str | None, # Target entity name if renaming/merging
|
||||
"final_entity": str, # Final entity name after operation
|
||||
"renamed": bool # Whether entity was renamed
|
||||
}
|
||||
}
|
||||
|
||||
operation_status values explained:
|
||||
- "success": All operations completed successfully
|
||||
* For simple updates: entity properties updated
|
||||
* For renames: entity renamed successfully
|
||||
* For merges: non-name updates applied AND merge completed
|
||||
|
||||
- "partial_success": Update succeeded but merge failed
|
||||
* Non-name property updates were applied successfully
|
||||
* Merge operation failed (entity not merged)
|
||||
* Original entity still exists with updated properties
|
||||
* Use merge_error for failure details
|
||||
|
||||
- "failure": Operation failed completely
|
||||
* If merge_status == "failed": Merge attempted but both update and merge failed
|
||||
* If merge_status == "not_attempted": Regular update failed
|
||||
* No changes were applied to the entity
|
||||
|
||||
merge_status values explained:
|
||||
- "success": Entity successfully merged into target entity
|
||||
- "failed": Merge operation was attempted but failed
|
||||
- "not_attempted": No merge was attempted (normal update/rename)
|
||||
|
||||
Behavior when renaming to an existing entity:
|
||||
- If allow_merge=False: Raises ValueError with 400 status (default behavior)
|
||||
- If allow_merge=True: Automatically merges the source entity into the existing target entity,
|
||||
preserving all relationships and applying non-name updates first
|
||||
|
||||
Example Request (simple update):
|
||||
POST /graph/entity/edit
|
||||
{
|
||||
"entity_name": "Tesla",
|
||||
"updated_data": {"description": "Updated description"},
|
||||
"allow_rename": false,
|
||||
"allow_merge": false
|
||||
}
|
||||
|
||||
Example Response (simple update success):
|
||||
{
|
||||
"status": "success",
|
||||
"message": "Entity updated successfully",
|
||||
"data": { ... },
|
||||
"operation_summary": {
|
||||
"merged": false,
|
||||
"merge_status": "not_attempted",
|
||||
"merge_error": null,
|
||||
"operation_status": "success",
|
||||
"target_entity": null,
|
||||
"final_entity": "Tesla",
|
||||
"renamed": false
|
||||
}
|
||||
}
|
||||
|
||||
Example Request (rename with auto-merge):
|
||||
POST /graph/entity/edit
|
||||
{
|
||||
"entity_name": "Elon Msk",
|
||||
"updated_data": {
|
||||
"entity_name": "Elon Musk",
|
||||
"description": "Corrected description"
|
||||
},
|
||||
"allow_rename": true,
|
||||
"allow_merge": true
|
||||
}
|
||||
|
||||
Example Response (merge success):
|
||||
{
|
||||
"status": "success",
|
||||
"message": "Entity merged successfully into 'Elon Musk'",
|
||||
"data": { ... },
|
||||
"operation_summary": {
|
||||
"merged": true,
|
||||
"merge_status": "success",
|
||||
"merge_error": null,
|
||||
"operation_status": "success",
|
||||
"target_entity": "Elon Musk",
|
||||
"final_entity": "Elon Musk",
|
||||
"renamed": true
|
||||
}
|
||||
}
|
||||
|
||||
Example Response (partial success - update succeeded but merge failed):
|
||||
{
|
||||
"status": "success",
|
||||
"message": "Entity updated successfully",
|
||||
"data": { ... }, # Data reflects updated "Elon Msk" entity
|
||||
"operation_summary": {
|
||||
"merged": false,
|
||||
"merge_status": "failed",
|
||||
"merge_error": "Target entity locked by another operation",
|
||||
"operation_status": "partial_success",
|
||||
"target_entity": "Elon Musk",
|
||||
"final_entity": "Elon Msk", # Original entity still exists
|
||||
"renamed": true
|
||||
}
|
||||
}
|
||||
"""
|
||||
try:
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
result = await rag.aedit_entity(
|
||||
entity_name=request.entity_name,
|
||||
updated_data=request.updated_data,
|
||||
allow_rename=request.allow_rename,
|
||||
allow_merge=request.allow_merge,
|
||||
)
|
||||
|
||||
# Extract operation_summary from result, with fallback for backward compatibility
|
||||
operation_summary = result.get(
|
||||
"operation_summary",
|
||||
{
|
||||
"merged": False,
|
||||
"merge_status": "not_attempted",
|
||||
"merge_error": None,
|
||||
"operation_status": "success",
|
||||
"target_entity": None,
|
||||
"final_entity": request.updated_data.get(
|
||||
"entity_name", request.entity_name
|
||||
),
|
||||
"renamed": request.updated_data.get(
|
||||
"entity_name", request.entity_name
|
||||
)
|
||||
!= request.entity_name,
|
||||
},
|
||||
)
|
||||
|
||||
# Separate entity data from operation_summary for clean response
|
||||
entity_data = dict(result)
|
||||
entity_data.pop("operation_summary", None)
|
||||
|
||||
# Generate appropriate response message based on merge status
|
||||
response_message = (
|
||||
f"Entity merged successfully into '{operation_summary['final_entity']}'"
|
||||
if operation_summary.get("merged")
|
||||
else "Entity updated successfully"
|
||||
)
|
||||
return {
|
||||
"status": "success",
|
||||
"message": response_message,
|
||||
"data": entity_data,
|
||||
"operation_summary": operation_summary,
|
||||
}
|
||||
except HTTPException:
|
||||
raise
|
||||
except ValueError as ve:
|
||||
logger.error(
|
||||
f"Validation error updating entity '{request.entity_name}': {str(ve)}"
|
||||
)
|
||||
raise HTTPException(status_code=400, detail=str(ve))
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating entity '{request.entity_name}': {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error updating entity: {str(e)}"
|
||||
)
|
||||
|
||||
@router.post("/graph/relation/edit", dependencies=[Depends(combined_auth)])
|
||||
async def update_relation(request: RelationUpdateRequest):
|
||||
"""Update a relation's properties in the knowledge graph
|
||||
|
||||
Args:
|
||||
request (RelationUpdateRequest): Request containing source ID, target ID and updated data
|
||||
|
||||
Returns:
|
||||
Dict: Updated relation information
|
||||
"""
|
||||
try:
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
result = await rag.aedit_relation(
|
||||
source_entity=request.source_id,
|
||||
target_entity=request.target_id,
|
||||
updated_data=request.updated_data,
|
||||
)
|
||||
return {
|
||||
"status": "success",
|
||||
"message": "Relation updated successfully",
|
||||
"data": result,
|
||||
}
|
||||
except HTTPException:
|
||||
raise
|
||||
except ValueError as ve:
|
||||
logger.error(
|
||||
f"Validation error updating relation between '{request.source_id}' and '{request.target_id}': {str(ve)}"
|
||||
)
|
||||
raise HTTPException(status_code=400, detail=str(ve))
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error updating relation between '{request.source_id}' and '{request.target_id}': {str(e)}"
|
||||
)
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error updating relation: {str(e)}"
|
||||
)
|
||||
|
||||
@router.post("/graph/entity/create", dependencies=[Depends(combined_auth)])
|
||||
async def create_entity(request: EntityCreateRequest):
|
||||
"""
|
||||
Create a new entity in the knowledge graph
|
||||
|
||||
This endpoint creates a new entity node in the knowledge graph with the specified
|
||||
properties. The system automatically generates vector embeddings for the entity
|
||||
to enable semantic search and retrieval.
|
||||
|
||||
Request Body:
|
||||
entity_name (str): Unique name identifier for the entity
|
||||
entity_data (dict): Entity properties including:
|
||||
- description (str): Textual description of the entity
|
||||
- entity_type (str): Category/type of the entity (e.g., PERSON, ORGANIZATION, LOCATION)
|
||||
- source_id (str): Related chunk_id from which the description originates
|
||||
- Additional custom properties as needed
|
||||
|
||||
Response Schema:
|
||||
{
|
||||
"status": "success",
|
||||
"message": "Entity 'Tesla' created successfully",
|
||||
"data": {
|
||||
"entity_name": "Tesla",
|
||||
"description": "Electric vehicle manufacturer",
|
||||
"entity_type": "ORGANIZATION",
|
||||
"source_id": "chunk-123<SEP>chunk-456"
|
||||
... (other entity properties)
|
||||
}
|
||||
}
|
||||
|
||||
HTTP Status Codes:
|
||||
200: Entity created successfully
|
||||
400: Invalid request (e.g., missing required fields, duplicate entity)
|
||||
500: Internal server error
|
||||
|
||||
Example Request:
|
||||
POST /graph/entity/create
|
||||
{
|
||||
"entity_name": "Tesla",
|
||||
"entity_data": {
|
||||
"description": "Electric vehicle manufacturer",
|
||||
"entity_type": "ORGANIZATION"
|
||||
}
|
||||
}
|
||||
"""
|
||||
try:
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
# Use the proper acreate_entity method which handles:
|
||||
# - Graph lock for concurrency
|
||||
# - Vector embedding creation in entities_vdb
|
||||
# - Metadata population and defaults
|
||||
# - Index consistency via _edit_entity_done
|
||||
result = await rag.acreate_entity(
|
||||
entity_name=request.entity_name,
|
||||
entity_data=request.entity_data,
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Entity '{request.entity_name}' created successfully",
|
||||
"data": result,
|
||||
}
|
||||
except HTTPException:
|
||||
raise
|
||||
except ValueError as ve:
|
||||
logger.error(
|
||||
f"Validation error creating entity '{request.entity_name}': {str(ve)}"
|
||||
)
|
||||
raise HTTPException(status_code=400, detail=str(ve))
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating entity '{request.entity_name}': {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error creating entity: {str(e)}"
|
||||
)
|
||||
|
||||
@router.post("/graph/relation/create", dependencies=[Depends(combined_auth)])
|
||||
async def create_relation(request: RelationCreateRequest):
|
||||
"""
|
||||
Create a new relationship between two entities in the knowledge graph
|
||||
|
||||
This endpoint establishes an undirected relationship between two existing entities.
|
||||
The provided source/target order is accepted for convenience, but the backend
|
||||
stored edge is undirected and may be returned with the entities swapped.
|
||||
Both entities must already exist in the knowledge graph. The system automatically
|
||||
generates vector embeddings for the relationship to enable semantic search and graph traversal.
|
||||
|
||||
Prerequisites:
|
||||
- Both source_entity and target_entity must exist in the knowledge graph
|
||||
- Use /graph/entity/create to create entities first if they don't exist
|
||||
|
||||
Request Body:
|
||||
source_entity (str): Name of the source entity (relationship origin)
|
||||
target_entity (str): Name of the target entity (relationship destination)
|
||||
relation_data (dict): Relationship properties including:
|
||||
- description (str): Textual description of the relationship
|
||||
- keywords (str): Comma-separated keywords describing the relationship type
|
||||
- source_id (str): Related chunk_id from which the description originates
|
||||
- weight (float): Relationship strength/importance (default: 1.0)
|
||||
- Additional custom properties as needed
|
||||
|
||||
Response Schema:
|
||||
{
|
||||
"status": "success",
|
||||
"message": "Relation created successfully between 'Elon Musk' and 'Tesla'",
|
||||
"data": {
|
||||
"src_id": "Elon Musk",
|
||||
"tgt_id": "Tesla",
|
||||
"description": "Elon Musk is the CEO of Tesla",
|
||||
"keywords": "CEO, founder",
|
||||
"source_id": "chunk-123<SEP>chunk-456"
|
||||
"weight": 1.0,
|
||||
... (other relationship properties)
|
||||
}
|
||||
}
|
||||
|
||||
HTTP Status Codes:
|
||||
200: Relationship created successfully
|
||||
400: Invalid request (e.g., missing entities, invalid data, duplicate relationship)
|
||||
500: Internal server error
|
||||
|
||||
Example Request:
|
||||
POST /graph/relation/create
|
||||
{
|
||||
"source_entity": "Elon Musk",
|
||||
"target_entity": "Tesla",
|
||||
"relation_data": {
|
||||
"description": "Elon Musk is the CEO of Tesla",
|
||||
"keywords": "CEO, founder",
|
||||
"weight": 1.0
|
||||
}
|
||||
}
|
||||
"""
|
||||
try:
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
# Use the proper acreate_relation method which handles:
|
||||
# - Graph lock for concurrency
|
||||
# - Entity existence validation
|
||||
# - Duplicate relation checks
|
||||
# - Vector embedding creation in relationships_vdb
|
||||
# - Index consistency via _edit_relation_done
|
||||
result = await rag.acreate_relation(
|
||||
source_entity=request.source_entity,
|
||||
target_entity=request.target_entity,
|
||||
relation_data=request.relation_data,
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Relation created successfully between '{request.source_entity}' and '{request.target_entity}'",
|
||||
"data": result,
|
||||
}
|
||||
except HTTPException:
|
||||
raise
|
||||
except ValueError as ve:
|
||||
logger.error(
|
||||
f"Validation error creating relation between '{request.source_entity}' and '{request.target_entity}': {str(ve)}"
|
||||
)
|
||||
raise HTTPException(status_code=400, detail=str(ve))
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error creating relation between '{request.source_entity}' and '{request.target_entity}': {str(e)}"
|
||||
)
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error creating relation: {str(e)}"
|
||||
)
|
||||
|
||||
@router.post("/graph/entities/merge", dependencies=[Depends(combined_auth)])
|
||||
async def merge_entities(request: EntityMergeRequest):
|
||||
"""
|
||||
Merge multiple entities into a single entity, preserving all relationships
|
||||
|
||||
This endpoint consolidates duplicate or misspelled entities while preserving the entire
|
||||
graph structure. It's particularly useful for cleaning up knowledge graphs after document
|
||||
processing or correcting entity name variations.
|
||||
|
||||
What the Merge Operation Does:
|
||||
1. Deletes the specified source entities from the knowledge graph
|
||||
2. Transfers all relationships from source entities to the target entity
|
||||
3. Intelligently merges duplicate relationships (if multiple sources have the same relationship)
|
||||
4. Updates vector embeddings for accurate retrieval and search
|
||||
5. Preserves the complete graph structure and connectivity
|
||||
6. Maintains relationship properties and metadata
|
||||
|
||||
Use Cases:
|
||||
- Fixing spelling errors in entity names (e.g., "Elon Msk" -> "Elon Musk")
|
||||
- Consolidating duplicate entities discovered after document processing
|
||||
- Merging name variations (e.g., "NY", "New York", "New York City")
|
||||
- Cleaning up the knowledge graph for better query performance
|
||||
- Standardizing entity names across the knowledge base
|
||||
|
||||
Request Body:
|
||||
entities_to_change (list[str]): List of entity names to be merged and deleted
|
||||
entity_to_change_into (str): Target entity that will receive all relationships
|
||||
|
||||
Response Schema:
|
||||
{
|
||||
"status": "success",
|
||||
"message": "Successfully merged 2 entities into 'Elon Musk'",
|
||||
"data": {
|
||||
"merged_entity": "Elon Musk",
|
||||
"deleted_entities": ["Elon Msk", "Ellon Musk"],
|
||||
"relationships_transferred": 15,
|
||||
... (merge operation details)
|
||||
}
|
||||
}
|
||||
|
||||
HTTP Status Codes:
|
||||
200: Entities merged successfully
|
||||
400: Invalid request (e.g., empty entity list, target entity doesn't exist)
|
||||
500: Internal server error
|
||||
|
||||
Example Request:
|
||||
POST /graph/entities/merge
|
||||
{
|
||||
"entities_to_change": ["Elon Msk", "Ellon Musk"],
|
||||
"entity_to_change_into": "Elon Musk"
|
||||
}
|
||||
|
||||
Note:
|
||||
- The target entity (entity_to_change_into) must exist in the knowledge graph
|
||||
- Source entities will be permanently deleted after the merge
|
||||
- This operation cannot be undone, so verify entity names before merging
|
||||
"""
|
||||
try:
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
result = await rag.amerge_entities(
|
||||
source_entities=request.entities_to_change,
|
||||
target_entity=request.entity_to_change_into,
|
||||
)
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Successfully merged {len(request.entities_to_change)} entities into '{request.entity_to_change_into}'",
|
||||
"data": result,
|
||||
}
|
||||
except HTTPException:
|
||||
raise
|
||||
except ValueError as ve:
|
||||
logger.error(
|
||||
f"Validation error merging entities {request.entities_to_change} into '{request.entity_to_change_into}': {str(ve)}"
|
||||
)
|
||||
raise HTTPException(status_code=400, detail=str(ve))
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error merging entities {request.entities_to_change} into '{request.entity_to_change_into}': {str(e)}"
|
||||
)
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Error merging entities: {str(e)}"
|
||||
)
|
||||
|
||||
@router.delete(
|
||||
"/graph/entity/delete",
|
||||
response_model=DeletionResult,
|
||||
dependencies=[Depends(combined_auth)],
|
||||
)
|
||||
async def delete_entity(request: DeleteEntityRequest):
|
||||
"""
|
||||
Delete an entity and all its relationships from the knowledge graph.
|
||||
|
||||
Args:
|
||||
request (DeleteEntityRequest): The request body containing the entity name.
|
||||
|
||||
Returns:
|
||||
DeletionResult: An object containing the outcome of the deletion process.
|
||||
|
||||
Raises:
|
||||
HTTPException: If the entity is not found (404) or an error occurs (500).
|
||||
"""
|
||||
try:
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
result = await rag.adelete_by_entity(entity_name=request.entity_name)
|
||||
if result.status == "not_found":
|
||||
raise HTTPException(status_code=404, detail=result.message)
|
||||
if result.status == "fail":
|
||||
raise HTTPException(status_code=500, detail=result.message)
|
||||
# Set doc_id to empty string since this is an entity operation, not document
|
||||
result.doc_id = ""
|
||||
return result
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_msg = f"Error deleting entity '{request.entity_name}': {str(e)}"
|
||||
logger.error(error_msg)
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=500, detail=error_msg)
|
||||
|
||||
@router.delete(
|
||||
"/graph/relation/delete",
|
||||
response_model=DeletionResult,
|
||||
dependencies=[Depends(combined_auth)],
|
||||
)
|
||||
async def delete_relation(request: DeleteRelationRequest):
|
||||
"""
|
||||
Delete a relationship between two entities from the knowledge graph.
|
||||
|
||||
Args:
|
||||
request (DeleteRelationRequest): The request body containing the source and target entity names.
|
||||
|
||||
Returns:
|
||||
DeletionResult: An object containing the outcome of the deletion process.
|
||||
|
||||
Raises:
|
||||
HTTPException: If the relation is not found (404) or an error occurs (500).
|
||||
"""
|
||||
try:
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
result = await rag.adelete_by_relation(
|
||||
source_entity=request.source_entity,
|
||||
target_entity=request.target_entity,
|
||||
)
|
||||
if result.status == "not_found":
|
||||
raise HTTPException(status_code=404, detail=result.message)
|
||||
if result.status == "fail":
|
||||
raise HTTPException(status_code=500, detail=result.message)
|
||||
# Set doc_id to empty string since this is a relation operation, not document
|
||||
result.doc_id = ""
|
||||
return result
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_msg = f"Error deleting relation from '{request.source_entity}' to '{request.target_entity}': {str(e)}"
|
||||
logger.error(error_msg)
|
||||
logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=500, detail=error_msg)
|
||||
|
||||
return router
|
||||
@@ -0,0 +1,747 @@
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Dict, Any, Optional, Type
|
||||
from lightrag.utils import logger
|
||||
import time
|
||||
import json
|
||||
import re
|
||||
from enum import Enum
|
||||
from fastapi.responses import StreamingResponse
|
||||
import asyncio
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.constants import DEFAULT_QUERY_PRIORITY
|
||||
from lightrag.utils import TiktokenTokenizer
|
||||
from lightrag.api.utils_api import get_combined_auth_dependency
|
||||
from fastapi import Depends
|
||||
|
||||
|
||||
# query mode according to query prefix (bypass is not LightRAG quer mode)
|
||||
class SearchMode(str, Enum):
|
||||
naive = "naive"
|
||||
local = "local"
|
||||
global_ = "global"
|
||||
hybrid = "hybrid"
|
||||
mix = "mix"
|
||||
bypass = "bypass"
|
||||
context = "context"
|
||||
|
||||
|
||||
class OllamaMessage(BaseModel):
|
||||
role: str
|
||||
content: str
|
||||
images: Optional[List[str]] = None
|
||||
|
||||
|
||||
class OllamaChatRequest(BaseModel):
|
||||
model: str
|
||||
messages: List[OllamaMessage]
|
||||
stream: bool = True
|
||||
options: Optional[Dict[str, Any]] = None
|
||||
system: Optional[str] = None
|
||||
|
||||
|
||||
class OllamaChatResponse(BaseModel):
|
||||
model: str
|
||||
created_at: str
|
||||
message: OllamaMessage
|
||||
done: bool
|
||||
|
||||
|
||||
class OllamaGenerateRequest(BaseModel):
|
||||
model: str
|
||||
prompt: str
|
||||
system: Optional[str] = None
|
||||
stream: bool = False
|
||||
options: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class OllamaGenerateResponse(BaseModel):
|
||||
model: str
|
||||
created_at: str
|
||||
response: str
|
||||
done: bool
|
||||
context: Optional[List[int]]
|
||||
total_duration: Optional[int]
|
||||
load_duration: Optional[int]
|
||||
prompt_eval_count: Optional[int]
|
||||
prompt_eval_duration: Optional[int]
|
||||
eval_count: Optional[int]
|
||||
eval_duration: Optional[int]
|
||||
|
||||
|
||||
class OllamaVersionResponse(BaseModel):
|
||||
version: str
|
||||
|
||||
|
||||
class OllamaModelDetails(BaseModel):
|
||||
parent_model: str
|
||||
format: str
|
||||
family: str
|
||||
families: List[str]
|
||||
parameter_size: str
|
||||
quantization_level: str
|
||||
|
||||
|
||||
class OllamaModel(BaseModel):
|
||||
name: str
|
||||
model: str
|
||||
size: int
|
||||
digest: str
|
||||
modified_at: str
|
||||
details: OllamaModelDetails
|
||||
|
||||
|
||||
class OllamaTagResponse(BaseModel):
|
||||
models: List[OllamaModel]
|
||||
|
||||
|
||||
class OllamaRunningModelDetails(BaseModel):
|
||||
parent_model: str
|
||||
format: str
|
||||
family: str
|
||||
families: List[str]
|
||||
parameter_size: str
|
||||
quantization_level: str
|
||||
|
||||
|
||||
class OllamaRunningModel(BaseModel):
|
||||
name: str
|
||||
model: str
|
||||
size: int
|
||||
digest: str
|
||||
details: OllamaRunningModelDetails
|
||||
expires_at: str
|
||||
size_vram: int
|
||||
|
||||
|
||||
class OllamaPsResponse(BaseModel):
|
||||
models: List[OllamaRunningModel]
|
||||
|
||||
|
||||
async def parse_request_body(
|
||||
request: Request, model_class: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Parse request body based on Content-Type header.
|
||||
Supports both application/json and application/octet-stream.
|
||||
|
||||
Args:
|
||||
request: The FastAPI Request object
|
||||
model_class: The Pydantic model class to parse the request into
|
||||
|
||||
Returns:
|
||||
An instance of the provided model_class
|
||||
"""
|
||||
content_type = request.headers.get("content-type", "").lower()
|
||||
|
||||
try:
|
||||
if content_type.startswith("application/json"):
|
||||
# FastAPI already handles JSON parsing for us
|
||||
body = await request.json()
|
||||
elif content_type.startswith("application/octet-stream"):
|
||||
# Manually parse octet-stream as JSON
|
||||
body_bytes = await request.body()
|
||||
body = json.loads(body_bytes.decode("utf-8"))
|
||||
else:
|
||||
# Try to parse as JSON for any other content type
|
||||
body_bytes = await request.body()
|
||||
body = json.loads(body_bytes.decode("utf-8"))
|
||||
|
||||
# Create an instance of the model
|
||||
return model_class(**body)
|
||||
except json.JSONDecodeError:
|
||||
raise HTTPException(status_code=400, detail="Invalid JSON in request body")
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=400, detail=f"Error parsing request body: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""Estimate the number of tokens in text using tiktoken"""
|
||||
tokens = TiktokenTokenizer().encode(text)
|
||||
return len(tokens)
|
||||
|
||||
|
||||
def parse_query_mode(query: str) -> tuple[str, SearchMode, bool, Optional[str]]:
|
||||
"""Parse query prefix to determine search mode
|
||||
Returns tuple of (cleaned_query, search_mode, only_need_context, user_prompt)
|
||||
|
||||
Examples:
|
||||
- "/local[use mermaid format for diagrams] query string" -> (cleaned_query, SearchMode.local, False, "use mermaid format for diagrams")
|
||||
- "/[use mermaid format for diagrams] query string" -> (cleaned_query, SearchMode.hybrid, False, "use mermaid format for diagrams")
|
||||
- "/local query string" -> (cleaned_query, SearchMode.local, False, None)
|
||||
"""
|
||||
# Initialize user_prompt as None
|
||||
user_prompt = None
|
||||
|
||||
# First check if there's a bracket format for user prompt
|
||||
bracket_pattern = r"^/([a-z]*)\[(.*?)\](.*)"
|
||||
bracket_match = re.match(bracket_pattern, query)
|
||||
|
||||
if bracket_match:
|
||||
mode_prefix = bracket_match.group(1)
|
||||
user_prompt = bracket_match.group(2)
|
||||
remaining_query = bracket_match.group(3).lstrip()
|
||||
|
||||
# Reconstruct query, removing the bracket part
|
||||
query = f"/{mode_prefix} {remaining_query}".strip()
|
||||
|
||||
# Unified handling of mode and only_need_context determination
|
||||
mode_map = {
|
||||
"/local ": (SearchMode.local, False),
|
||||
"/global ": (
|
||||
SearchMode.global_,
|
||||
False,
|
||||
), # global_ is used because 'global' is a Python keyword
|
||||
"/naive ": (SearchMode.naive, False),
|
||||
"/hybrid ": (SearchMode.hybrid, False),
|
||||
"/mix ": (SearchMode.mix, False),
|
||||
"/bypass ": (SearchMode.bypass, False),
|
||||
"/context": (
|
||||
SearchMode.mix,
|
||||
True,
|
||||
),
|
||||
"/localcontext": (SearchMode.local, True),
|
||||
"/globalcontext": (SearchMode.global_, True),
|
||||
"/hybridcontext": (SearchMode.hybrid, True),
|
||||
"/naivecontext": (SearchMode.naive, True),
|
||||
"/mixcontext": (SearchMode.mix, True),
|
||||
}
|
||||
|
||||
for prefix, (mode, only_need_context) in mode_map.items():
|
||||
if query.startswith(prefix):
|
||||
# After removing prefix and leading spaces
|
||||
cleaned_query = query[len(prefix) :].lstrip()
|
||||
return cleaned_query, mode, only_need_context, user_prompt
|
||||
|
||||
return query, SearchMode.mix, False, user_prompt
|
||||
|
||||
|
||||
class OllamaAPI:
|
||||
def __init__(self, rag: LightRAG, top_k: int = 60, api_key: Optional[str] = None):
|
||||
self.rag = rag
|
||||
self.ollama_server_infos = rag.ollama_server_infos
|
||||
self.top_k = top_k
|
||||
self.api_key = api_key
|
||||
self.router = APIRouter(tags=["ollama"])
|
||||
self.setup_routes()
|
||||
|
||||
def setup_routes(self):
|
||||
# Create combined auth dependency for Ollama API routes
|
||||
combined_auth = get_combined_auth_dependency(self.api_key)
|
||||
|
||||
@self.router.get("/version", dependencies=[Depends(combined_auth)])
|
||||
async def get_version():
|
||||
"""Get Ollama version information"""
|
||||
return OllamaVersionResponse(version="0.9.3")
|
||||
|
||||
@self.router.get("/tags", dependencies=[Depends(combined_auth)])
|
||||
async def get_tags():
|
||||
"""Return available models acting as an Ollama server"""
|
||||
return OllamaTagResponse(
|
||||
models=[
|
||||
{
|
||||
"name": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"modified_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"size": self.ollama_server_infos.LIGHTRAG_SIZE,
|
||||
"digest": self.ollama_server_infos.LIGHTRAG_DIGEST,
|
||||
"details": {
|
||||
"parent_model": "",
|
||||
"format": "gguf",
|
||||
"family": self.ollama_server_infos.LIGHTRAG_NAME,
|
||||
"families": [self.ollama_server_infos.LIGHTRAG_NAME],
|
||||
"parameter_size": "13B",
|
||||
"quantization_level": "Q4_0",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
@self.router.get("/ps", dependencies=[Depends(combined_auth)])
|
||||
async def get_running_models():
|
||||
"""List Running Models - returns currently running models"""
|
||||
return OllamaPsResponse(
|
||||
models=[
|
||||
{
|
||||
"name": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"size": self.ollama_server_infos.LIGHTRAG_SIZE,
|
||||
"digest": self.ollama_server_infos.LIGHTRAG_DIGEST,
|
||||
"details": {
|
||||
"parent_model": "",
|
||||
"format": "gguf",
|
||||
"family": "llama",
|
||||
"families": ["llama"],
|
||||
"parameter_size": "7.2B",
|
||||
"quantization_level": "Q4_0",
|
||||
},
|
||||
"expires_at": "2050-12-31T14:38:31.83753-07:00",
|
||||
"size_vram": self.ollama_server_infos.LIGHTRAG_SIZE,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
@self.router.post(
|
||||
"/generate", dependencies=[Depends(combined_auth)], include_in_schema=True
|
||||
)
|
||||
async def generate(raw_request: Request):
|
||||
"""Handle generate completion requests acting as an Ollama model
|
||||
For compatibility purpose, the request is not processed by LightRAG,
|
||||
and will be handled by underlying LLM model.
|
||||
Supports both application/json and application/octet-stream Content-Types.
|
||||
"""
|
||||
try:
|
||||
# Parse the request body manually
|
||||
request = await parse_request_body(raw_request, OllamaGenerateRequest)
|
||||
|
||||
query = request.prompt
|
||||
start_time = time.time_ns()
|
||||
prompt_tokens = estimate_tokens(query)
|
||||
|
||||
role_kwargs = (
|
||||
dict(self.rag.role_llm_kwargs["query"])
|
||||
if self.rag.role_llm_kwargs["query"] is not None
|
||||
else dict(self.rag.llm_model_kwargs)
|
||||
)
|
||||
if request.system:
|
||||
role_kwargs["system_prompt"] = request.system
|
||||
|
||||
if request.stream:
|
||||
response = await (self.rag.role_llm_funcs["query"])(
|
||||
query,
|
||||
stream=True,
|
||||
_priority=DEFAULT_QUERY_PRIORITY,
|
||||
**role_kwargs,
|
||||
)
|
||||
|
||||
async def stream_generator():
|
||||
first_chunk_time = None
|
||||
last_chunk_time = time.time_ns()
|
||||
total_response = ""
|
||||
|
||||
# Ensure response is an async generator
|
||||
if isinstance(response, str):
|
||||
# If it's a string, send in two parts
|
||||
first_chunk_time = start_time
|
||||
last_chunk_time = time.time_ns()
|
||||
total_response = response
|
||||
|
||||
data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"response": response,
|
||||
"done": False,
|
||||
}
|
||||
yield f"{json.dumps(data, ensure_ascii=False)}\n"
|
||||
|
||||
completion_tokens = estimate_tokens(total_response)
|
||||
total_time = last_chunk_time - start_time
|
||||
prompt_eval_time = first_chunk_time - start_time
|
||||
eval_time = last_chunk_time - first_chunk_time
|
||||
|
||||
data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"response": "",
|
||||
"done": True,
|
||||
"done_reason": "stop",
|
||||
"context": [],
|
||||
"total_duration": total_time,
|
||||
"load_duration": 0,
|
||||
"prompt_eval_count": prompt_tokens,
|
||||
"prompt_eval_duration": prompt_eval_time,
|
||||
"eval_count": completion_tokens,
|
||||
"eval_duration": eval_time,
|
||||
}
|
||||
yield f"{json.dumps(data, ensure_ascii=False)}\n"
|
||||
else:
|
||||
try:
|
||||
async for chunk in response:
|
||||
if chunk:
|
||||
if first_chunk_time is None:
|
||||
first_chunk_time = time.time_ns()
|
||||
|
||||
last_chunk_time = time.time_ns()
|
||||
|
||||
total_response += chunk
|
||||
data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"response": chunk,
|
||||
"done": False,
|
||||
}
|
||||
yield f"{json.dumps(data, ensure_ascii=False)}\n"
|
||||
except (asyncio.CancelledError, Exception) as e:
|
||||
error_msg = str(e)
|
||||
if isinstance(e, asyncio.CancelledError):
|
||||
error_msg = "Stream was cancelled by server"
|
||||
else:
|
||||
error_msg = f"Provider error: {error_msg}"
|
||||
|
||||
logger.error(f"Stream error: {error_msg}")
|
||||
|
||||
# Send error message to client
|
||||
error_data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"response": f"\n\nError: {error_msg}",
|
||||
"error": f"\n\nError: {error_msg}",
|
||||
"done": False,
|
||||
}
|
||||
yield f"{json.dumps(error_data, ensure_ascii=False)}\n"
|
||||
|
||||
# Send final message to close the stream
|
||||
final_data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"response": "",
|
||||
"done": True,
|
||||
}
|
||||
yield f"{json.dumps(final_data, ensure_ascii=False)}\n"
|
||||
return
|
||||
if first_chunk_time is None:
|
||||
first_chunk_time = start_time
|
||||
completion_tokens = estimate_tokens(total_response)
|
||||
total_time = last_chunk_time - start_time
|
||||
prompt_eval_time = first_chunk_time - start_time
|
||||
eval_time = last_chunk_time - first_chunk_time
|
||||
|
||||
data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"response": "",
|
||||
"done": True,
|
||||
"done_reason": "stop",
|
||||
"context": [],
|
||||
"total_duration": total_time,
|
||||
"load_duration": 0,
|
||||
"prompt_eval_count": prompt_tokens,
|
||||
"prompt_eval_duration": prompt_eval_time,
|
||||
"eval_count": completion_tokens,
|
||||
"eval_duration": eval_time,
|
||||
}
|
||||
yield f"{json.dumps(data, ensure_ascii=False)}\n"
|
||||
return
|
||||
|
||||
return StreamingResponse(
|
||||
stream_generator(),
|
||||
media_type="application/x-ndjson",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"Content-Type": "application/x-ndjson",
|
||||
"X-Accel-Buffering": "no", # Ensure proper handling of streaming responses in Nginx proxy
|
||||
},
|
||||
)
|
||||
else:
|
||||
first_chunk_time = time.time_ns()
|
||||
response_text = await (self.rag.role_llm_funcs["query"])(
|
||||
query,
|
||||
stream=False,
|
||||
_priority=DEFAULT_QUERY_PRIORITY,
|
||||
**role_kwargs,
|
||||
)
|
||||
last_chunk_time = time.time_ns()
|
||||
|
||||
if not response_text:
|
||||
response_text = "No response generated"
|
||||
|
||||
completion_tokens = estimate_tokens(str(response_text))
|
||||
total_time = last_chunk_time - start_time
|
||||
prompt_eval_time = first_chunk_time - start_time
|
||||
eval_time = last_chunk_time - first_chunk_time
|
||||
|
||||
return {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"response": str(response_text),
|
||||
"done": True,
|
||||
"done_reason": "stop",
|
||||
"context": [],
|
||||
"total_duration": total_time,
|
||||
"load_duration": 0,
|
||||
"prompt_eval_count": prompt_tokens,
|
||||
"prompt_eval_duration": prompt_eval_time,
|
||||
"eval_count": completion_tokens,
|
||||
"eval_duration": eval_time,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Ollama generate error: {str(e)}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@self.router.post(
|
||||
"/chat", dependencies=[Depends(combined_auth)], include_in_schema=True
|
||||
)
|
||||
async def chat(raw_request: Request):
|
||||
"""Process chat completion requests by acting as an Ollama model.
|
||||
Routes user queries through LightRAG by selecting query mode based on query prefix.
|
||||
Detects and forwards OpenWebUI session-related requests (for meta data generation task) directly to LLM.
|
||||
Supports both application/json and application/octet-stream Content-Types.
|
||||
"""
|
||||
try:
|
||||
# Parse the request body manually
|
||||
request = await parse_request_body(raw_request, OllamaChatRequest)
|
||||
|
||||
# Get all messages
|
||||
messages = request.messages
|
||||
if not messages:
|
||||
raise HTTPException(status_code=400, detail="No messages provided")
|
||||
|
||||
# Validate that the last message is from a user
|
||||
if messages[-1].role != "user":
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Last message must be from user role"
|
||||
)
|
||||
|
||||
# Get the last message as query and previous messages as history
|
||||
query = messages[-1].content
|
||||
# Convert OllamaMessage objects to dictionaries
|
||||
conversation_history = [
|
||||
{"role": msg.role, "content": msg.content} for msg in messages[:-1]
|
||||
]
|
||||
|
||||
# Check for query prefix
|
||||
cleaned_query, mode, only_need_context, user_prompt = parse_query_mode(
|
||||
query
|
||||
)
|
||||
|
||||
start_time = time.time_ns()
|
||||
prompt_tokens = estimate_tokens(cleaned_query)
|
||||
|
||||
param_dict = {
|
||||
"mode": mode.value,
|
||||
"stream": request.stream,
|
||||
"only_need_context": only_need_context,
|
||||
"conversation_history": conversation_history,
|
||||
"top_k": self.top_k,
|
||||
}
|
||||
|
||||
# Add user_prompt to param_dict
|
||||
if user_prompt is not None:
|
||||
param_dict["user_prompt"] = user_prompt
|
||||
|
||||
query_param = QueryParam(**param_dict)
|
||||
|
||||
if request.stream:
|
||||
# Determine if the request is prefix with "/bypass"
|
||||
if mode == SearchMode.bypass:
|
||||
role_kwargs = (
|
||||
dict(self.rag.role_llm_kwargs["query"])
|
||||
if self.rag.role_llm_kwargs["query"] is not None
|
||||
else dict(self.rag.llm_model_kwargs)
|
||||
)
|
||||
if request.system:
|
||||
role_kwargs["system_prompt"] = request.system
|
||||
response = await (self.rag.role_llm_funcs["query"])(
|
||||
cleaned_query,
|
||||
stream=True,
|
||||
history_messages=conversation_history,
|
||||
_priority=DEFAULT_QUERY_PRIORITY,
|
||||
**role_kwargs,
|
||||
)
|
||||
else:
|
||||
response = await self.rag.aquery(
|
||||
cleaned_query, param=query_param
|
||||
)
|
||||
|
||||
async def stream_generator():
|
||||
first_chunk_time = None
|
||||
last_chunk_time = time.time_ns()
|
||||
total_response = ""
|
||||
|
||||
# Ensure response is an async generator
|
||||
if isinstance(response, str):
|
||||
# If it's a string, send in two parts
|
||||
first_chunk_time = start_time
|
||||
last_chunk_time = time.time_ns()
|
||||
total_response = response
|
||||
|
||||
data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": response,
|
||||
"images": None,
|
||||
},
|
||||
"done": False,
|
||||
}
|
||||
yield f"{json.dumps(data, ensure_ascii=False)}\n"
|
||||
|
||||
completion_tokens = estimate_tokens(total_response)
|
||||
total_time = last_chunk_time - start_time
|
||||
prompt_eval_time = first_chunk_time - start_time
|
||||
eval_time = last_chunk_time - first_chunk_time
|
||||
|
||||
data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"images": None,
|
||||
},
|
||||
"done_reason": "stop",
|
||||
"done": True,
|
||||
"total_duration": total_time,
|
||||
"load_duration": 0,
|
||||
"prompt_eval_count": prompt_tokens,
|
||||
"prompt_eval_duration": prompt_eval_time,
|
||||
"eval_count": completion_tokens,
|
||||
"eval_duration": eval_time,
|
||||
}
|
||||
yield f"{json.dumps(data, ensure_ascii=False)}\n"
|
||||
else:
|
||||
try:
|
||||
async for chunk in response:
|
||||
if chunk:
|
||||
if first_chunk_time is None:
|
||||
first_chunk_time = time.time_ns()
|
||||
|
||||
last_chunk_time = time.time_ns()
|
||||
|
||||
total_response += chunk
|
||||
data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": chunk,
|
||||
"images": None,
|
||||
},
|
||||
"done": False,
|
||||
}
|
||||
yield f"{json.dumps(data, ensure_ascii=False)}\n"
|
||||
except (asyncio.CancelledError, Exception) as e:
|
||||
error_msg = str(e)
|
||||
if isinstance(e, asyncio.CancelledError):
|
||||
error_msg = "Stream was cancelled by server"
|
||||
else:
|
||||
error_msg = f"Provider error: {error_msg}"
|
||||
|
||||
logger.error(f"Stream error: {error_msg}")
|
||||
|
||||
# Send error message to client
|
||||
error_data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": f"\n\nError: {error_msg}",
|
||||
"images": None,
|
||||
},
|
||||
"error": f"\n\nError: {error_msg}",
|
||||
"done": False,
|
||||
}
|
||||
yield f"{json.dumps(error_data, ensure_ascii=False)}\n"
|
||||
|
||||
# Send final message to close the stream
|
||||
final_data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"images": None,
|
||||
},
|
||||
"done": True,
|
||||
}
|
||||
yield f"{json.dumps(final_data, ensure_ascii=False)}\n"
|
||||
return
|
||||
|
||||
if first_chunk_time is None:
|
||||
first_chunk_time = start_time
|
||||
completion_tokens = estimate_tokens(total_response)
|
||||
total_time = last_chunk_time - start_time
|
||||
prompt_eval_time = first_chunk_time - start_time
|
||||
eval_time = last_chunk_time - first_chunk_time
|
||||
|
||||
data = {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"images": None,
|
||||
},
|
||||
"done_reason": "stop",
|
||||
"done": True,
|
||||
"total_duration": total_time,
|
||||
"load_duration": 0,
|
||||
"prompt_eval_count": prompt_tokens,
|
||||
"prompt_eval_duration": prompt_eval_time,
|
||||
"eval_count": completion_tokens,
|
||||
"eval_duration": eval_time,
|
||||
}
|
||||
yield f"{json.dumps(data, ensure_ascii=False)}\n"
|
||||
|
||||
return StreamingResponse(
|
||||
stream_generator(),
|
||||
media_type="application/x-ndjson",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"Content-Type": "application/x-ndjson",
|
||||
"X-Accel-Buffering": "no", # Ensure proper handling of streaming responses in Nginx proxy
|
||||
},
|
||||
)
|
||||
else:
|
||||
first_chunk_time = time.time_ns()
|
||||
|
||||
# Determine if the request is prefix with "/bypass" or from Open WebUI's session title and session keyword generation task
|
||||
match_result = re.search(
|
||||
r"\n<chat_history>\nUSER:", cleaned_query, re.MULTILINE
|
||||
)
|
||||
if match_result or mode == SearchMode.bypass:
|
||||
role_kwargs = (
|
||||
dict(self.rag.role_llm_kwargs["query"])
|
||||
if self.rag.role_llm_kwargs["query"] is not None
|
||||
else dict(self.rag.llm_model_kwargs)
|
||||
)
|
||||
if request.system:
|
||||
role_kwargs["system_prompt"] = request.system
|
||||
|
||||
response_text = await (self.rag.role_llm_funcs["query"])(
|
||||
cleaned_query,
|
||||
stream=False,
|
||||
history_messages=conversation_history,
|
||||
_priority=DEFAULT_QUERY_PRIORITY,
|
||||
**role_kwargs,
|
||||
)
|
||||
else:
|
||||
response_text = await self.rag.aquery(
|
||||
cleaned_query, param=query_param
|
||||
)
|
||||
|
||||
last_chunk_time = time.time_ns()
|
||||
|
||||
if not response_text:
|
||||
response_text = "No response generated"
|
||||
|
||||
completion_tokens = estimate_tokens(str(response_text))
|
||||
total_time = last_chunk_time - start_time
|
||||
prompt_eval_time = first_chunk_time - start_time
|
||||
eval_time = last_chunk_time - first_chunk_time
|
||||
|
||||
return {
|
||||
"model": self.ollama_server_infos.LIGHTRAG_MODEL,
|
||||
"created_at": self.ollama_server_infos.LIGHTRAG_CREATED_AT,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": str(response_text),
|
||||
"images": None,
|
||||
},
|
||||
"done_reason": "stop",
|
||||
"done": True,
|
||||
"total_duration": total_time,
|
||||
"load_duration": 0,
|
||||
"prompt_eval_count": prompt_tokens,
|
||||
"prompt_eval_duration": prompt_eval_time,
|
||||
"eval_count": completion_tokens,
|
||||
"eval_duration": eval_time,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Ollama chat error: {str(e)}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,302 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
Start LightRAG server with Gunicorn
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import platform
|
||||
import pipmaster as pm
|
||||
|
||||
# Capture this before importing LightRAG modules, because those imports load .env.
|
||||
# On macOS, libobjc needs this value in the inherited process environment.
|
||||
_PROCESS_START_OBJC_FORK_SAFETY = os.environ.get("OBJC_DISABLE_INITIALIZE_FORK_SAFETY")
|
||||
|
||||
|
||||
def check_and_install_dependencies():
|
||||
"""Check and install required dependencies"""
|
||||
required_packages = [
|
||||
"gunicorn",
|
||||
"tiktoken",
|
||||
"psutil",
|
||||
# Add other required packages here
|
||||
]
|
||||
|
||||
for package in required_packages:
|
||||
if not pm.is_installed(package):
|
||||
print(f"Installing {package}...")
|
||||
pm.install(package)
|
||||
print(f"{package} installed successfully")
|
||||
|
||||
|
||||
def _build_global_concurrency_limits(args) -> dict:
|
||||
"""Derive cross-worker concurrency limits from the MAX_ASYNC settings.
|
||||
|
||||
Under gunicorn multi-worker, every MAX_ASYNC value keeps its documented
|
||||
meaning — the maximum number of concurrent provider calls — by acting
|
||||
BOTH as each worker's local limit and as the cross-worker global cap
|
||||
(without the gate the real total would be ~ MAX_ASYNC x workers).
|
||||
Group names must match the ``concurrency_group`` values passed to
|
||||
``priority_limit_async_func_call``: ``llm:{role}`` for the LLM roles,
|
||||
plus ``embedding`` and ``rerank``. Role fallbacks mirror the runtime
|
||||
resolution in ``_get_effective_role_llm_max_async``.
|
||||
"""
|
||||
from lightrag.llm_roles import ROLES
|
||||
|
||||
limits = {}
|
||||
for spec in ROLES:
|
||||
role_limit = getattr(args, f"{spec.name}_llm_max_async", None)
|
||||
if role_limit is None:
|
||||
role_limit = args.max_async
|
||||
if role_limit is not None and role_limit > 0:
|
||||
limits[f"llm:{spec.name}"] = role_limit
|
||||
embedding_limit = getattr(args, "embedding_func_max_async", None)
|
||||
if embedding_limit is not None and embedding_limit > 0:
|
||||
limits["embedding"] = embedding_limit
|
||||
rerank_limit = getattr(args, "rerank_max_async", None)
|
||||
if rerank_limit is not None and rerank_limit > 0:
|
||||
limits["rerank"] = rerank_limit
|
||||
return limits
|
||||
|
||||
|
||||
def main():
|
||||
from lightrag.api.utils_api import display_splash_screen, check_env_file
|
||||
from lightrag.api.config import global_args, initialize_config
|
||||
from lightrag.utils import get_env_value
|
||||
from lightrag.kg.shared_storage import initialize_share_data
|
||||
from lightrag.constants import (
|
||||
DEFAULT_WOKERS,
|
||||
DEFAULT_TIMEOUT,
|
||||
)
|
||||
|
||||
# Explicitly initialize configuration for Gunicorn mode
|
||||
initialize_config()
|
||||
|
||||
# Set Gunicorn mode flag for lifespan cleanup detection
|
||||
os.environ["LIGHTRAG_GUNICORN_MODE"] = "1"
|
||||
|
||||
# Check .env file
|
||||
if not check_env_file():
|
||||
sys.exit(1)
|
||||
|
||||
# Check macOS fork safety environment variable for multi-worker mode
|
||||
if (
|
||||
platform.system() == "Darwin"
|
||||
and global_args.workers > 1
|
||||
and _PROCESS_START_OBJC_FORK_SAFETY != "YES"
|
||||
):
|
||||
current_objc_fork_safety = os.environ.get("OBJC_DISABLE_INITIALIZE_FORK_SAFETY")
|
||||
print("\n" + "=" * 80)
|
||||
print("❌ ERROR: Missing required environment variable on macOS!")
|
||||
print("=" * 80)
|
||||
print("\nmacOS with Gunicorn multi-worker mode requires:")
|
||||
print(" OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES")
|
||||
print("\nReason:")
|
||||
print(" NumPy uses macOS's Accelerate framework (Objective-C based) for")
|
||||
print(" vector computations. The Objective-C runtime has fork safety checks")
|
||||
print(" that will crash worker processes when embedding functions are called.")
|
||||
print("\nCurrent configuration:")
|
||||
print(" - Operating System: macOS (Darwin)")
|
||||
print(f" - Workers: {global_args.workers}")
|
||||
print(
|
||||
" - Process Environment at Startup: "
|
||||
f"{_PROCESS_START_OBJC_FORK_SAFETY or 'NOT SET'}"
|
||||
)
|
||||
print(
|
||||
f" - Environment After .env Load: {current_objc_fork_safety or 'NOT SET'}"
|
||||
)
|
||||
if current_objc_fork_safety == "YES":
|
||||
print("\nNote:")
|
||||
print(" OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES was loaded from .env,")
|
||||
print(" but that is too late for the macOS Objective-C runtime.")
|
||||
print(" Export it before starting lightrag-gunicorn.")
|
||||
print("\nHow to fix:")
|
||||
print(" Option 1 - Set environment variable before starting (recommended):")
|
||||
print(" export OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES")
|
||||
print(" lightrag-gunicorn --workers 2")
|
||||
print("\n Option 2 - Add to your shell profile (~/.zshrc or ~/.bash_profile):")
|
||||
print(" echo 'export OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES' >> ~/.zshrc")
|
||||
print(" source ~/.zshrc")
|
||||
print("\n Option 3 - Use single worker mode (no multiprocessing):")
|
||||
print(" lightrag-server --workers 1")
|
||||
print("=" * 80 + "\n")
|
||||
sys.exit(1)
|
||||
|
||||
# Check and install dependencies
|
||||
check_and_install_dependencies()
|
||||
|
||||
# Note: Signal handlers are NOT registered here because:
|
||||
# - Master cleanup already handled by gunicorn_config.on_exit()
|
||||
|
||||
# Display startup information
|
||||
display_splash_screen(global_args)
|
||||
|
||||
print("🚀 Starting LightRAG with Gunicorn")
|
||||
print(f"🔄 Worker management: Gunicorn (workers={global_args.workers})")
|
||||
print("🔍 Preloading app: Enabled")
|
||||
print("📝 Note: Using Gunicorn's preload feature for shared data initialization")
|
||||
print("\n\n" + "=" * 80)
|
||||
print("MAIN PROCESS INITIALIZATION")
|
||||
print(f"Process ID: {os.getpid()}")
|
||||
print(f"Workers setting: {global_args.workers}")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
# Import Gunicorn's StandaloneApplication
|
||||
from gunicorn.app.base import BaseApplication
|
||||
|
||||
# Define a custom application class that loads our config
|
||||
class GunicornApp(BaseApplication):
|
||||
def __init__(self, app, options=None):
|
||||
self.options = options or {}
|
||||
self.application = app
|
||||
super().__init__()
|
||||
|
||||
def load_config(self):
|
||||
# Define valid Gunicorn configuration options
|
||||
valid_options = {
|
||||
"bind",
|
||||
"workers",
|
||||
"worker_class",
|
||||
"timeout",
|
||||
"keepalive",
|
||||
"preload_app",
|
||||
"errorlog",
|
||||
"accesslog",
|
||||
"loglevel",
|
||||
"certfile",
|
||||
"keyfile",
|
||||
"limit_request_line",
|
||||
"limit_request_fields",
|
||||
"limit_request_field_size",
|
||||
"graceful_timeout",
|
||||
"max_requests",
|
||||
"max_requests_jitter",
|
||||
}
|
||||
|
||||
# Special hooks that need to be set separately
|
||||
special_hooks = {
|
||||
"on_starting",
|
||||
"on_reload",
|
||||
"on_exit",
|
||||
"pre_fork",
|
||||
"post_fork",
|
||||
"pre_exec",
|
||||
"pre_request",
|
||||
"post_request",
|
||||
"worker_init",
|
||||
"worker_exit",
|
||||
"nworkers_changed",
|
||||
"child_exit",
|
||||
}
|
||||
|
||||
# Import and configure the gunicorn_config module
|
||||
from lightrag.api import gunicorn_config
|
||||
|
||||
# Set configuration variables in gunicorn_config, prioritizing command line arguments
|
||||
gunicorn_config.workers = (
|
||||
global_args.workers
|
||||
if global_args.workers
|
||||
else get_env_value("WORKERS", DEFAULT_WOKERS, int)
|
||||
)
|
||||
|
||||
# Bind configuration prioritizes command line arguments
|
||||
host = (
|
||||
global_args.host
|
||||
if global_args.host != "0.0.0.0"
|
||||
else os.getenv("HOST", "0.0.0.0")
|
||||
)
|
||||
port = (
|
||||
global_args.port
|
||||
if global_args.port != 9621
|
||||
else get_env_value("PORT", 9621, int)
|
||||
)
|
||||
gunicorn_config.bind = f"{host}:{port}"
|
||||
|
||||
# Log level configuration prioritizes command line arguments
|
||||
gunicorn_config.loglevel = (
|
||||
global_args.log_level.lower()
|
||||
if global_args.log_level
|
||||
else os.getenv("LOG_LEVEL", "info")
|
||||
)
|
||||
|
||||
# Timeout configuration prioritizes command line arguments
|
||||
gunicorn_config.timeout = (
|
||||
global_args.timeout + 30
|
||||
if global_args.timeout is not None
|
||||
else get_env_value(
|
||||
"TIMEOUT", DEFAULT_TIMEOUT + 30, int, special_none=True
|
||||
)
|
||||
)
|
||||
|
||||
# Keepalive configuration
|
||||
gunicorn_config.keepalive = get_env_value("KEEPALIVE", 5, int)
|
||||
|
||||
# SSL configuration prioritizes command line arguments
|
||||
if global_args.ssl or os.getenv("SSL", "").lower() in (
|
||||
"true",
|
||||
"1",
|
||||
"yes",
|
||||
"t",
|
||||
"on",
|
||||
):
|
||||
gunicorn_config.certfile = (
|
||||
global_args.ssl_certfile
|
||||
if global_args.ssl_certfile
|
||||
else os.getenv("SSL_CERTFILE")
|
||||
)
|
||||
gunicorn_config.keyfile = (
|
||||
global_args.ssl_keyfile
|
||||
if global_args.ssl_keyfile
|
||||
else os.getenv("SSL_KEYFILE")
|
||||
)
|
||||
|
||||
# Set configuration options from the module
|
||||
for key in dir(gunicorn_config):
|
||||
if key in valid_options:
|
||||
value = getattr(gunicorn_config, key)
|
||||
# Skip functions like on_starting and None values
|
||||
if not callable(value) and value is not None:
|
||||
self.cfg.set(key, value)
|
||||
# Set special hooks
|
||||
elif key in special_hooks:
|
||||
value = getattr(gunicorn_config, key)
|
||||
if callable(value):
|
||||
self.cfg.set(key, value)
|
||||
|
||||
if hasattr(gunicorn_config, "logconfig_dict"):
|
||||
self.cfg.set(
|
||||
"logconfig_dict", getattr(gunicorn_config, "logconfig_dict")
|
||||
)
|
||||
|
||||
def load(self):
|
||||
# Import the application
|
||||
from lightrag.api.lightrag_server import get_application
|
||||
|
||||
return get_application(global_args)
|
||||
|
||||
# Create the application
|
||||
app = GunicornApp("")
|
||||
|
||||
# Force workers to be an integer and greater than 1 for multi-process mode
|
||||
workers_count = global_args.workers
|
||||
if workers_count > 1:
|
||||
# Set a flag to indicate we're in the main process
|
||||
os.environ["LIGHTRAG_MAIN_PROCESS"] = "1"
|
||||
# Cross-worker global concurrency limits derived from MAX_ASYNC
|
||||
# (read-only after this point; forked workers inherit them as
|
||||
# module globals). Single-worker mode needs no cross-process gate —
|
||||
# the per-process max_async already IS the total limit there.
|
||||
initialize_share_data(
|
||||
workers_count,
|
||||
global_concurrency_limits=_build_global_concurrency_limits(global_args),
|
||||
)
|
||||
else:
|
||||
initialize_share_data(1)
|
||||
|
||||
# Run the application
|
||||
print("\nStarting Gunicorn with direct Python API...")
|
||||
app.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Helpers for validating startup runtime expectations from `.env`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import dotenv_values
|
||||
|
||||
_CONTAINER_RUNTIME_TARGETS = {"compose", "docker"}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RuntimeEnvironment:
|
||||
"""Describes whether the current process is running in a container runtime."""
|
||||
|
||||
in_container: bool
|
||||
in_docker: bool
|
||||
in_kubernetes: bool
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
if self.in_kubernetes:
|
||||
return "Kubernetes"
|
||||
if self.in_docker:
|
||||
return "Docker"
|
||||
return "host"
|
||||
|
||||
|
||||
def _read_cgroup_content() -> str:
|
||||
"""Best-effort read of cgroup metadata for container detection."""
|
||||
|
||||
for candidate in ("/proc/1/cgroup", "/proc/self/cgroup"):
|
||||
try:
|
||||
return Path(candidate).read_text(encoding="utf-8")
|
||||
except OSError:
|
||||
continue
|
||||
return ""
|
||||
|
||||
|
||||
def detect_runtime_environment(
|
||||
environ: dict[str, str] | None = None,
|
||||
) -> RuntimeEnvironment:
|
||||
"""Detect whether the current process is running on host, Docker, or Kubernetes."""
|
||||
|
||||
environ = environ or os.environ
|
||||
cgroup_content = _read_cgroup_content().lower()
|
||||
|
||||
in_kubernetes = bool(
|
||||
environ.get("KUBERNETES_SERVICE_HOST")
|
||||
or Path("/var/run/secrets/kubernetes.io/serviceaccount").exists()
|
||||
or "kubepods" in cgroup_content
|
||||
or "kubernetes" in cgroup_content
|
||||
)
|
||||
in_docker = bool(
|
||||
Path("/.dockerenv").exists()
|
||||
or Path("/run/.containerenv").exists()
|
||||
or any(
|
||||
marker in cgroup_content
|
||||
for marker in ("docker", "containerd", "libpod", "podman")
|
||||
)
|
||||
)
|
||||
|
||||
return RuntimeEnvironment(
|
||||
in_container=in_kubernetes or in_docker,
|
||||
in_docker=in_docker,
|
||||
in_kubernetes=in_kubernetes,
|
||||
)
|
||||
|
||||
|
||||
def load_runtime_target_from_env_file(env_path: str | Path = ".env") -> str | None:
|
||||
"""Return the raw LIGHTRAG_RUNTIME_TARGET value from the `.env` file, if present."""
|
||||
|
||||
env_values = dotenv_values(str(env_path))
|
||||
runtime_target = env_values.get("LIGHTRAG_RUNTIME_TARGET")
|
||||
if runtime_target is None:
|
||||
return None
|
||||
return runtime_target.strip()
|
||||
|
||||
|
||||
def validate_runtime_target(
|
||||
runtime_target: str | None,
|
||||
runtime_environment: RuntimeEnvironment | None = None,
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Validate `.env` runtime target against the current runtime environment."""
|
||||
|
||||
if runtime_target is None:
|
||||
return True, None
|
||||
|
||||
normalized_target = runtime_target.strip().lower()
|
||||
runtime_environment = runtime_environment or detect_runtime_environment()
|
||||
|
||||
if normalized_target == "host":
|
||||
if runtime_environment.in_container:
|
||||
return (
|
||||
False,
|
||||
"Configuration error in .env: LIGHTRAG_RUNTIME_TARGET=host.\n"
|
||||
"This value from .env requires the server process to run on the host, "
|
||||
f"but the current process is running inside {runtime_environment.label}.",
|
||||
)
|
||||
return True, None
|
||||
|
||||
if normalized_target in _CONTAINER_RUNTIME_TARGETS:
|
||||
if runtime_environment.in_container:
|
||||
return True, None
|
||||
return (
|
||||
False,
|
||||
f"Configuration error in .env: LIGHTRAG_RUNTIME_TARGET={runtime_target}.\n"
|
||||
"This value from .env requires the server process to run inside Docker or "
|
||||
"Kubernetes, but the current process is running on the host.",
|
||||
)
|
||||
|
||||
return (
|
||||
False,
|
||||
f"Configuration error in .env: LIGHTRAG_RUNTIME_TARGET={runtime_target!r}.\n"
|
||||
"This value from .env must be 'host' or 'compose' (alias: 'docker').",
|
||||
)
|
||||
|
||||
|
||||
def validate_runtime_target_from_env_file(
|
||||
env_path: str | Path = ".env",
|
||||
runtime_environment: RuntimeEnvironment | None = None,
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Load LIGHTRAG_RUNTIME_TARGET from `.env` and validate it if declared."""
|
||||
|
||||
runtime_target = load_runtime_target_from_env_file(env_path)
|
||||
return validate_runtime_target(runtime_target, runtime_environment)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 628 B |
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,575 @@
|
||||
"""
|
||||
Utility functions for the LightRAG API.
|
||||
"""
|
||||
|
||||
import os
|
||||
import argparse
|
||||
from typing import Optional, List, Tuple
|
||||
import sys
|
||||
import time
|
||||
import logging
|
||||
from ascii_colors import ASCIIColors
|
||||
from .._version import __api_version__ as api_version
|
||||
from .._version import __version__ as core_version
|
||||
from lightrag.constants import (
|
||||
DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE,
|
||||
)
|
||||
from lightrag.api.runtime_validation import validate_runtime_target_from_env_file
|
||||
from fastapi import HTTPException, Security, Request, Response, status
|
||||
from fastapi.security import APIKeyHeader, OAuth2PasswordBearer
|
||||
from starlette.status import HTTP_403_FORBIDDEN
|
||||
from .auth import auth_handler
|
||||
from .config import ollama_server_infos, global_args, get_env_value
|
||||
|
||||
logger = logging.getLogger("lightrag")
|
||||
|
||||
# ========== Token Renewal Rate Limiting ==========
|
||||
# Cache to track last renewal time per user (username as key)
|
||||
# Format: {username: last_renewal_timestamp}
|
||||
_token_renewal_cache: dict[str, float] = {}
|
||||
_RENEWAL_MIN_INTERVAL = 60 # Minimum 60 seconds between renewals for same user
|
||||
|
||||
# ========== Token Renewal Path Exclusions ==========
|
||||
# Paths that should NOT trigger token auto-renewal
|
||||
# - /health: Health check endpoint, no login required
|
||||
# - /documents/paginated: Client polls this frequently (5-30s), renewal not needed
|
||||
# - /documents/pipeline_status: Client polls this very frequently (2s), renewal not needed
|
||||
_TOKEN_RENEWAL_SKIP_PATHS = [
|
||||
"/health",
|
||||
"/documents/paginated",
|
||||
"/documents/pipeline_status",
|
||||
]
|
||||
|
||||
|
||||
def check_env_file():
|
||||
"""
|
||||
Check if .env file exists and handle user confirmation if needed.
|
||||
Returns True if should continue, False if should exit.
|
||||
"""
|
||||
env_path = ".env"
|
||||
|
||||
if not os.path.exists(env_path):
|
||||
warning_msg = "Warning: Startup directory must contain .env file for multi-instance support."
|
||||
ASCIIColors.yellow(warning_msg)
|
||||
|
||||
# Check if running in interactive terminal
|
||||
if sys.stdin.isatty():
|
||||
response = input("Do you want to continue? (yes/NO): ")
|
||||
if response.lower() != "yes":
|
||||
ASCIIColors.red("Server startup cancelled")
|
||||
return False
|
||||
return True
|
||||
|
||||
is_valid, error_message = validate_runtime_target_from_env_file(env_path)
|
||||
if not is_valid:
|
||||
for line in error_message.splitlines():
|
||||
ASCIIColors.red(line)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
# Get whitelist paths from global_args, only once during initialization
|
||||
whitelist_paths = global_args.whitelist_paths.split(",")
|
||||
|
||||
# Pre-compile path matching patterns
|
||||
whitelist_patterns: List[Tuple[str, bool]] = []
|
||||
for path in whitelist_paths:
|
||||
path = path.strip()
|
||||
if path:
|
||||
# If path ends with /*, match all paths with that prefix
|
||||
if path.endswith("/*"):
|
||||
prefix = path[:-2]
|
||||
whitelist_patterns.append((prefix, True)) # (prefix, is_prefix_match)
|
||||
else:
|
||||
whitelist_patterns.append((path, False)) # (exact_path, is_prefix_match)
|
||||
|
||||
# Global authentication configuration
|
||||
auth_configured = bool(auth_handler.accounts)
|
||||
|
||||
|
||||
def get_combined_auth_dependency(api_key: Optional[str] = None):
|
||||
"""
|
||||
Create a combined authentication dependency that implements authentication logic
|
||||
based on API key, OAuth2 token, and whitelist paths.
|
||||
|
||||
Args:
|
||||
api_key (Optional[str]): API key for validation
|
||||
|
||||
Returns:
|
||||
Callable: A dependency function that implements the authentication logic
|
||||
"""
|
||||
# Use global whitelist_patterns and auth_configured variables
|
||||
# whitelist_patterns and auth_configured are already initialized at module level
|
||||
|
||||
# Only calculate api_key_configured as it depends on the function parameter
|
||||
api_key_configured = bool(api_key)
|
||||
|
||||
# Create security dependencies with proper descriptions for Swagger UI
|
||||
oauth2_scheme = OAuth2PasswordBearer(
|
||||
tokenUrl="login", auto_error=False, description="OAuth2 Password Authentication"
|
||||
)
|
||||
|
||||
# If API key is configured, create an API key header security
|
||||
api_key_header = None
|
||||
if api_key_configured:
|
||||
api_key_header = APIKeyHeader(
|
||||
name="X-API-Key", auto_error=False, description="API Key Authentication"
|
||||
)
|
||||
|
||||
async def combined_dependency(
|
||||
request: Request,
|
||||
response: Response, # Added: needed to return new token via response header
|
||||
token: str = Security(oauth2_scheme),
|
||||
api_key_header_value: Optional[str] = None
|
||||
if api_key_header is None
|
||||
else Security(api_key_header),
|
||||
):
|
||||
# 1. Check if path is in whitelist
|
||||
path = request.url.path
|
||||
for pattern, is_prefix in whitelist_patterns:
|
||||
if (is_prefix and path.startswith(pattern)) or (
|
||||
not is_prefix and path == pattern
|
||||
):
|
||||
return # Whitelist path, allow access
|
||||
|
||||
# 2. Validate token first if provided in the request (Ensure 401 error if token is invalid)
|
||||
if token:
|
||||
try:
|
||||
token_info = auth_handler.validate_token(token)
|
||||
|
||||
# ========== Token Auto-Renewal Logic ==========
|
||||
from lightrag.api.config import global_args
|
||||
from datetime import datetime, timezone
|
||||
|
||||
if global_args.token_auto_renew:
|
||||
# Check if current path should skip token renewal
|
||||
skip_renewal = any(
|
||||
path == skip_path or path.startswith(skip_path + "/")
|
||||
for skip_path in _TOKEN_RENEWAL_SKIP_PATHS
|
||||
)
|
||||
|
||||
if skip_renewal:
|
||||
logger.debug(f"Token auto-renewal skipped for path: {path}")
|
||||
else:
|
||||
try:
|
||||
expire_time = token_info.get("exp")
|
||||
if expire_time:
|
||||
# Calculate remaining time ratio
|
||||
now = datetime.now(timezone.utc)
|
||||
remaining_seconds = (expire_time - now).total_seconds()
|
||||
|
||||
# Get original token expiration duration
|
||||
role = token_info.get("role", "user")
|
||||
total_hours = (
|
||||
auth_handler.guest_expire_hours
|
||||
if role == "guest"
|
||||
else auth_handler.expire_hours
|
||||
)
|
||||
total_seconds = total_hours * 3600
|
||||
|
||||
# Issue new token if remaining time < threshold
|
||||
if (
|
||||
remaining_seconds
|
||||
< total_seconds * global_args.token_renew_threshold
|
||||
):
|
||||
# ========== Rate Limiting Check ==========
|
||||
username = token_info["username"]
|
||||
current_time = time.time()
|
||||
last_renewal = _token_renewal_cache.get(username, 0)
|
||||
time_since_last_renewal = (
|
||||
current_time - last_renewal
|
||||
)
|
||||
|
||||
# Only renew if enough time has passed since last renewal
|
||||
if time_since_last_renewal >= _RENEWAL_MIN_INTERVAL:
|
||||
new_token = auth_handler.create_token(
|
||||
username=username,
|
||||
role=role,
|
||||
metadata=token_info.get("metadata", {}),
|
||||
)
|
||||
# Return new token via response header
|
||||
response.headers["X-New-Token"] = new_token
|
||||
|
||||
# Update renewal cache
|
||||
_token_renewal_cache[username] = current_time
|
||||
|
||||
# Optional: log renewal
|
||||
logger.info(
|
||||
f"Token auto-renewed for user {username} "
|
||||
f"(role: {role}, remaining: {remaining_seconds:.0f}s)"
|
||||
)
|
||||
else:
|
||||
# Log skip due to rate limit
|
||||
logger.debug(
|
||||
f"Token renewal skipped for {username} "
|
||||
f"(rate limit: last renewal {time_since_last_renewal:.0f}s ago)"
|
||||
)
|
||||
# ========== End of Rate Limiting Check ==========
|
||||
except Exception as e:
|
||||
# Renewal failure should not affect normal request, just log
|
||||
logger.warning(f"Token auto-renew failed: {e}")
|
||||
# ========== End of Token Auto-Renewal Logic ==========
|
||||
|
||||
# A token only authenticates when it matches the configured auth mode:
|
||||
# - password auth (AUTH_ACCOUNTS set): accept non-guest user tokens
|
||||
# - fully open (no AUTH_ACCOUNTS, no API key): accept guest tokens
|
||||
# In the API-key-only profile (API key set, no AUTH_ACCOUNTS) a guest
|
||||
# token must NOT authenticate: anyone can obtain one (via /auth-status,
|
||||
# /login, or by signing it with the public default secret), so honoring
|
||||
# it here would let a forged guest token bypass the X-API-Key check
|
||||
# below (GHSA-f4vv-55c2-5789 / GHSA-xr5c-v5r6-c9f9). Instead, fall
|
||||
# through so the API key stays mandatory in that mode.
|
||||
if not auth_configured and token_info.get("role") == "guest":
|
||||
if not api_key_configured:
|
||||
return
|
||||
# API-key-only mode: ignore the guest token; the X-API-Key check
|
||||
# below is the sole authority. Fall through (no return, no raise).
|
||||
elif auth_configured and token_info.get("role") != "guest":
|
||||
# Accept non-guest token if password auth is configured
|
||||
return
|
||||
else:
|
||||
# Token present but not valid for the configured auth mode.
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid token. Please login again.",
|
||||
)
|
||||
except HTTPException as e:
|
||||
# If already a 401 error, re-raise it
|
||||
if e.status_code == status.HTTP_401_UNAUTHORIZED:
|
||||
raise
|
||||
# For other exceptions, continue processing
|
||||
|
||||
# 3. Acept all request if no API protection needed
|
||||
if not auth_configured and not api_key_configured:
|
||||
return
|
||||
|
||||
# 4. Validate API key if provided and API-Key authentication is configured
|
||||
if (
|
||||
api_key_configured
|
||||
and api_key_header_value
|
||||
and api_key_header_value == api_key
|
||||
):
|
||||
return # API key validation successful
|
||||
|
||||
### Authentication failed ####
|
||||
|
||||
# if password authentication is configured but not provided, ensure 401 error if auth_configured
|
||||
if auth_configured and not token:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="No credentials provided. Please login.",
|
||||
)
|
||||
|
||||
# if api key is provided but validation failed
|
||||
if api_key_header_value:
|
||||
raise HTTPException(
|
||||
status_code=HTTP_403_FORBIDDEN,
|
||||
detail="Invalid API Key",
|
||||
)
|
||||
|
||||
# if api_key_configured but not provided
|
||||
if api_key_configured and not api_key_header_value:
|
||||
raise HTTPException(
|
||||
status_code=HTTP_403_FORBIDDEN,
|
||||
detail="API Key required",
|
||||
)
|
||||
|
||||
# Otherwise: refuse access and return 403 error
|
||||
raise HTTPException(
|
||||
status_code=HTTP_403_FORBIDDEN,
|
||||
detail="API Key required or login authentication required.",
|
||||
)
|
||||
|
||||
return combined_dependency
|
||||
|
||||
|
||||
def get_auth_status_dependency(api_key: Optional[str] = None):
|
||||
"""Create a dependency that reports whether the request carries accepted
|
||||
credentials, WITHOUT enforcing authentication (it never raises).
|
||||
|
||||
Used by endpoints such as ``/health`` that must stay reachable for
|
||||
unauthenticated liveness probes (always HTTP 200) while only revealing
|
||||
sensitive configuration to authenticated callers. The acceptance rules
|
||||
mirror ``get_combined_auth_dependency`` exactly:
|
||||
|
||||
- fully open (no AUTH_ACCOUNTS, no API key): nothing is protected
|
||||
anywhere, so the request is treated as authenticated.
|
||||
- password auth (AUTH_ACCOUNTS set): a valid non-guest token, or a
|
||||
valid API key when one is configured, authenticates.
|
||||
- API-key-only (API key set, no AUTH_ACCOUNTS): only a valid API key
|
||||
authenticates; a guest token is forgeable and must NOT count
|
||||
(GHSA-f4vv-55c2-5789 / GHSA-xr5c-v5r6-c9f9).
|
||||
"""
|
||||
api_key_configured = bool(api_key)
|
||||
oauth2_scheme = OAuth2PasswordBearer(
|
||||
tokenUrl="login", auto_error=False, description="OAuth2 Password Authentication"
|
||||
)
|
||||
api_key_header = None
|
||||
if api_key_configured:
|
||||
api_key_header = APIKeyHeader(
|
||||
name="X-API-Key", auto_error=False, description="API Key Authentication"
|
||||
)
|
||||
|
||||
async def auth_status_dependency(
|
||||
token: str = Security(oauth2_scheme),
|
||||
api_key_header_value: Optional[str] = None
|
||||
if api_key_header is None
|
||||
else Security(api_key_header),
|
||||
) -> bool:
|
||||
# Fully-open mode: nothing is protected anywhere, so reveal config too.
|
||||
if not auth_configured and not api_key_configured:
|
||||
return True
|
||||
|
||||
# A valid API key authenticates in any mode where one is configured.
|
||||
if (
|
||||
api_key_configured
|
||||
and api_key_header_value
|
||||
and api_key_header_value == api_key
|
||||
):
|
||||
return True
|
||||
|
||||
if token:
|
||||
try:
|
||||
token_info = auth_handler.validate_token(token)
|
||||
except Exception:
|
||||
token_info = None
|
||||
if token_info:
|
||||
role = token_info.get("role")
|
||||
# Password auth: accept a non-guest token. A guest token never
|
||||
# authenticates here (in API-key-only mode it is forgeable).
|
||||
if auth_configured and role != "guest":
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
return auth_status_dependency
|
||||
|
||||
|
||||
def whitelist_exposes_api_routes(whitelist_paths: str) -> bool:
|
||||
"""Return True if WHITELIST_PATHS exempts any Ollama-compatible /api route.
|
||||
|
||||
Mirrors the prefix/exact matching in get_combined_auth_dependency so that a
|
||||
catch-all entry such as ``/*`` (which strips to an empty prefix and matches
|
||||
every request path, including ``/api/chat``) is recognized as exposing the
|
||||
/api routes — not just literal ``/api...`` entries.
|
||||
"""
|
||||
for entry in whitelist_paths.split(","):
|
||||
entry = entry.strip()
|
||||
if not entry:
|
||||
continue
|
||||
if entry.endswith("/*"):
|
||||
# Prefix match: this entry exempts an /api route when some /api path
|
||||
# starts with the prefix ("/api".startswith(prefix) also covers the
|
||||
# empty catch-all prefix from "/*") or the prefix is itself under
|
||||
# /api/. The "/api/" boundary matters: "/apiary/*" only exempts
|
||||
# /apiary..., not /api/chat, so it must NOT be flagged.
|
||||
prefix = entry[:-2]
|
||||
if "/api".startswith(prefix) or prefix.startswith("/api/"):
|
||||
return True
|
||||
else:
|
||||
# Exact match: only the literal path is exempted.
|
||||
if entry == "/api" or entry.startswith("/api/"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def display_splash_screen(args: argparse.Namespace) -> None:
|
||||
"""
|
||||
Display a colorful splash screen showing LightRAG server configuration
|
||||
|
||||
Args:
|
||||
args: Parsed command line arguments
|
||||
"""
|
||||
# Banner
|
||||
# Banner
|
||||
top_border = "╔══════════════════════════════════════════════════════════════╗"
|
||||
bottom_border = "╚══════════════════════════════════════════════════════════════╝"
|
||||
width = len(top_border) - 4 # width inside the borders
|
||||
|
||||
line1_text = f"LightRAG Server v{core_version}/{api_version}"
|
||||
line2_text = "Fast, Lightweight RAG Server Implementation"
|
||||
|
||||
line1 = f"║ {line1_text.center(width)} ║"
|
||||
line2 = f"║ {line2_text.center(width)} ║"
|
||||
|
||||
banner = f"""
|
||||
{top_border}
|
||||
{line1}
|
||||
{line2}
|
||||
{bottom_border}
|
||||
"""
|
||||
ASCIIColors.cyan(banner)
|
||||
|
||||
# Server Configuration
|
||||
ASCIIColors.magenta("\n📡 Server Configuration:")
|
||||
ASCIIColors.white(" ├─ Host: ", end="")
|
||||
ASCIIColors.yellow(f"{args.host}")
|
||||
ASCIIColors.white(" ├─ Port: ", end="")
|
||||
ASCIIColors.yellow(f"{args.port}")
|
||||
ASCIIColors.white(" ├─ Workers: ", end="")
|
||||
ASCIIColors.yellow(f"{args.workers}")
|
||||
ASCIIColors.white(" ├─ Timeout: ", end="")
|
||||
ASCIIColors.yellow(f"{args.timeout}")
|
||||
ASCIIColors.white(" ├─ CORS Origins: ", end="")
|
||||
ASCIIColors.yellow(f"{args.cors_origins}")
|
||||
ASCIIColors.white(" ├─ SSL Enabled: ", end="")
|
||||
ASCIIColors.yellow(f"{args.ssl}")
|
||||
if args.ssl:
|
||||
ASCIIColors.white(" ├─ SSL Cert: ", end="")
|
||||
ASCIIColors.yellow(f"{args.ssl_certfile}")
|
||||
ASCIIColors.white(" ├─ SSL Key: ", end="")
|
||||
ASCIIColors.yellow(f"{args.ssl_keyfile}")
|
||||
ASCIIColors.white(" ├─ Ollama Emulating Model: ", end="")
|
||||
ASCIIColors.yellow(f"{ollama_server_infos.LIGHTRAG_MODEL}")
|
||||
ASCIIColors.white(" ├─ Log Level: ", end="")
|
||||
ASCIIColors.yellow(f"{args.log_level}")
|
||||
ASCIIColors.white(" ├─ Verbose Debug: ", end="")
|
||||
ASCIIColors.yellow(f"{args.verbose}")
|
||||
ASCIIColors.white(" ├─ API Key: ", end="")
|
||||
ASCIIColors.yellow("Set" if args.key else "Not Set")
|
||||
ASCIIColors.white(" └─ JWT Auth: ", end="")
|
||||
ASCIIColors.yellow("Enabled" if args.auth_accounts else "Disabled")
|
||||
|
||||
# Directory Configuration
|
||||
ASCIIColors.magenta("\n📂 Directory Configuration:")
|
||||
ASCIIColors.white(" ├─ Working Directory: ", end="")
|
||||
ASCIIColors.yellow(f"{args.working_dir}")
|
||||
ASCIIColors.white(" └─ Input Directory: ", end="")
|
||||
ASCIIColors.yellow(f"{args.input_dir}")
|
||||
# Embedding Configuration
|
||||
ASCIIColors.magenta("\n📊 Embedding Configuration:")
|
||||
ASCIIColors.white(" ├─ Binding: ", end="")
|
||||
ASCIIColors.yellow(f"{args.embedding_binding}")
|
||||
ASCIIColors.white(" ├─ Host: ", end="")
|
||||
ASCIIColors.yellow(f"{args.embedding_binding_host}")
|
||||
ASCIIColors.white(" ├─ Model: ", end="")
|
||||
ASCIIColors.yellow(f"{args.embedding_model}")
|
||||
ASCIIColors.white(" ├─ Dimensions: ", end="")
|
||||
ASCIIColors.yellow(f"{args.embedding_dim}")
|
||||
ASCIIColors.white(" └─ Asymmetric: ", end="")
|
||||
ASCIIColors.yellow(f"{args.embedding_asymmetric}")
|
||||
|
||||
# RAG Configuration
|
||||
ASCIIColors.magenta("\n⚙️ RAG Configuration:")
|
||||
ASCIIColors.white(" ├─ Summary Language: ", end="")
|
||||
ASCIIColors.yellow(f"{args.summary_language}")
|
||||
ASCIIColors.white(" ├─ Max Parallel Insert: ", end="")
|
||||
ASCIIColors.yellow(f"{args.max_parallel_insert}")
|
||||
ASCIIColors.white(" ├─ Chunk Size: ", end="")
|
||||
ASCIIColors.yellow(f"{args.chunk_size}")
|
||||
ASCIIColors.white(" ├─ Chunk Overlap Size: ", end="")
|
||||
ASCIIColors.yellow(f"{args.chunk_overlap_size}")
|
||||
ASCIIColors.white(" ├─ Cosine Threshold: ", end="")
|
||||
ASCIIColors.yellow(f"{args.cosine_threshold}")
|
||||
ASCIIColors.white(" ├─ Top-K: ", end="")
|
||||
ASCIIColors.yellow(f"{args.top_k}")
|
||||
ASCIIColors.white(" └─ Force LLM Summary on Merge: ", end="")
|
||||
ASCIIColors.yellow(
|
||||
f"{get_env_value('FORCE_LLM_SUMMARY_ON_MERGE', DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE, int)}"
|
||||
)
|
||||
|
||||
# System Configuration
|
||||
ASCIIColors.magenta("\n💾 Storage Configuration:")
|
||||
ASCIIColors.white(" ├─ KV Storage: ", end="")
|
||||
ASCIIColors.yellow(f"{args.kv_storage}")
|
||||
ASCIIColors.white(" ├─ Vector Storage: ", end="")
|
||||
ASCIIColors.yellow(f"{args.vector_storage}")
|
||||
ASCIIColors.white(" ├─ Graph Storage: ", end="")
|
||||
ASCIIColors.yellow(f"{args.graph_storage}")
|
||||
ASCIIColors.white(" ├─ Document Status Storage: ", end="")
|
||||
ASCIIColors.yellow(f"{args.doc_status_storage}")
|
||||
ASCIIColors.white(" └─ Workspace: ", end="")
|
||||
ASCIIColors.yellow(f"{args.workspace if args.workspace else '-'}")
|
||||
|
||||
# Server Status
|
||||
ASCIIColors.green("\n✨ Server starting up...\n")
|
||||
|
||||
# Server Access Information
|
||||
protocol = "https" if args.ssl else "http"
|
||||
if args.host == "0.0.0.0":
|
||||
ASCIIColors.magenta("\n🌐 Server Access Information:")
|
||||
ASCIIColors.white(" ├─ WebUI (local): ", end="")
|
||||
ASCIIColors.yellow(f"{protocol}://localhost:{args.port}")
|
||||
ASCIIColors.white(" ├─ Remote Access: ", end="")
|
||||
ASCIIColors.yellow(f"{protocol}://<your-ip-address>:{args.port}")
|
||||
ASCIIColors.white(" ├─ API Documentation (local): ", end="")
|
||||
ASCIIColors.yellow(f"{protocol}://localhost:{args.port}/docs")
|
||||
ASCIIColors.white(" └─ Alternative Documentation (local): ", end="")
|
||||
ASCIIColors.yellow(f"{protocol}://localhost:{args.port}/redoc")
|
||||
|
||||
ASCIIColors.magenta("\n📝 Note:")
|
||||
ASCIIColors.cyan(""" Since the server is running on 0.0.0.0:
|
||||
- Use 'localhost' or '127.0.0.1' for local access
|
||||
- Use your machine's IP address for remote access
|
||||
- To find your IP address:
|
||||
• Windows: Run 'ipconfig' in terminal
|
||||
• Linux/Mac: Run 'ifconfig' or 'ip addr' in terminal
|
||||
""")
|
||||
else:
|
||||
base_url = f"{protocol}://{args.host}:{args.port}"
|
||||
ASCIIColors.magenta("\n🌐 Server Access Information:")
|
||||
ASCIIColors.white(" ├─ WebUI (local): ", end="")
|
||||
ASCIIColors.yellow(f"{base_url}")
|
||||
ASCIIColors.white(" ├─ API Documentation: ", end="")
|
||||
ASCIIColors.yellow(f"{base_url}/docs")
|
||||
ASCIIColors.white(" └─ Alternative Documentation: ", end="")
|
||||
ASCIIColors.yellow(f"{base_url}/redoc")
|
||||
|
||||
# Security Notice
|
||||
if args.key:
|
||||
ASCIIColors.white("✅ Security Notice:")
|
||||
ASCIIColors.white(""" API Key authentication is enabled.
|
||||
Make sure to include the X-API-Key header in all your requests.
|
||||
""")
|
||||
if args.auth_accounts:
|
||||
ASCIIColors.white("✅ Security Notice:")
|
||||
ASCIIColors.white(""" JWT authentication is enabled.
|
||||
Make sure to login before making the request, and include the 'Authorization' in the header.
|
||||
""")
|
||||
|
||||
# Warn when the server runs without any authentication. In this mode every
|
||||
# endpoint is publicly reachable (see get_combined_auth_dependency: with
|
||||
# neither AUTH_ACCOUNTS nor LIGHTRAG_API_KEY set, all requests are allowed).
|
||||
if not args.key and not args.auth_accounts:
|
||||
loopback_hosts = {"127.0.0.1", "::1", "localhost"}
|
||||
if args.host in loopback_hosts:
|
||||
ASCIIColors.yellow("\n⚠️ Security Warning:")
|
||||
ASCIIColors.white(f""" No authentication is configured (no API Key, no login accounts).
|
||||
The server is bound to a loopback address ('{args.host}'), so it is only
|
||||
reachable from this machine. Set LIGHTRAG_API_KEY, or AUTH_ACCOUNTS together
|
||||
with TOKEN_SECRET, before binding to a non-loopback address (e.g. HOST=0.0.0.0).
|
||||
""")
|
||||
else:
|
||||
ASCIIColors.red("\n🔴 SECURITY ALERT:")
|
||||
ASCIIColors.white(f""" The server is listening on '{args.host}' WITHOUT any authentication.
|
||||
Every endpoint (document upload, query, knowledge graph, deletion) is
|
||||
publicly accessible to anyone who can reach this address.
|
||||
|
||||
Secure the server before exposing it to a network by setting at least one of:
|
||||
- LIGHTRAG_API_KEY=<a-strong-secret> (X-API-Key header authentication)
|
||||
- AUTH_ACCOUNTS=user:password together with TOKEN_SECRET=<a-strong-secret>
|
||||
(JWT login authentication; AUTH_ACCOUNTS
|
||||
without TOKEN_SECRET fails to start)
|
||||
Or restrict access by binding to loopback only: HOST=127.0.0.1
|
||||
""")
|
||||
|
||||
# When authentication IS configured but the server is exposed on a
|
||||
# non-loopback address, warn that the default whitelist still exempts the
|
||||
# Ollama-compatible /api/* routes (kept open for Ollama-client compatibility).
|
||||
# Those routes invoke the LLM and read the knowledge base, so they stay
|
||||
# public unless the operator narrows WHITELIST_PATHS (e.g. to /health).
|
||||
if args.key or args.auth_accounts:
|
||||
loopback_hosts = {"127.0.0.1", "::1", "localhost"}
|
||||
ollama_open = whitelist_exposes_api_routes(args.whitelist_paths)
|
||||
if args.host not in loopback_hosts and ollama_open:
|
||||
ASCIIColors.yellow("\n⚠️ Security Warning:")
|
||||
ASCIIColors.white(f""" WHITELIST_PATHS ('{args.whitelist_paths}') exempts the Ollama-compatible
|
||||
/api/* routes (/api/chat, /api/generate, ...) from authentication, so they
|
||||
remain publicly accessible on '{args.host}' even though auth is enabled.
|
||||
These routes invoke the LLM and read your knowledge base. If you do not need
|
||||
open Ollama access, set WHITELIST_PATHS=/health to require authentication.
|
||||
""")
|
||||
|
||||
# Ensure splash output flush to system log
|
||||
sys.stdout.flush()
|
||||
+1068
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,383 @@
|
||||
"""Chunk schema helpers shared across the chunking + extraction pipeline.
|
||||
|
||||
Three responsibilities live here so chunker implementations and the pipeline
|
||||
both consume identical normalization rules:
|
||||
|
||||
- :func:`normalize_chunk_heading` collapses the legacy flat
|
||||
``heading``/``parent_headings``/``level`` triple and the new nested form
|
||||
into the canonical ``{"level", "heading", "parent_headings"}`` dict.
|
||||
- :func:`normalize_chunk_sidecar` validates the new ``sidecar`` payload and
|
||||
ensures ``refs`` is always present as a list (single-source items may omit
|
||||
it before normalization; we materialize a single-element list for the
|
||||
storage layer).
|
||||
- :func:`strip_internal_multimodal_markup_for_extraction` rewrites
|
||||
``<cite>`` / ``<drawing>`` / ``<equation>`` markup so the entity-extraction
|
||||
LLM sees a clean text body. The original ``chunk["content"]`` is never
|
||||
mutated; the cleaned string is only used to build the extraction prompt.
|
||||
|
||||
The clean function is intentionally conservative: it only strips
|
||||
parser-emitted identifier attributes that have no business reaching the LLM
|
||||
(``id``, ``refid``, ``path``, ``src``). Visible captions and equation bodies
|
||||
are preserved so the extracted entities can still ground against them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import unicodedata
|
||||
from typing import Any
|
||||
|
||||
from lightrag.constants import DEFAULT_HEADING_LEVEL_MAX_CHARS
|
||||
|
||||
|
||||
_SIDECAR_TYPES = frozenset({"block", "drawing", "table", "equation"})
|
||||
|
||||
# Separator joining heading levels into a single breadcrumb line. Shared so the
|
||||
# extraction-side token budgeter (see ``operate._truncate_section_context``) can
|
||||
# split the rendered breadcrumb back into levels without a drifting magic string.
|
||||
HEADING_BREADCRUMB_SEP = " → "
|
||||
|
||||
|
||||
def normalize_chunk_heading(dp: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""Return the canonical nested heading dict or ``None`` when absent.
|
||||
|
||||
Accepts:
|
||||
|
||||
- ``dp["heading"]`` already a dict ``{"level", "heading", "parent_headings"}``.
|
||||
- Legacy flat fields ``heading: str`` + ``parent_headings: list[str]`` +
|
||||
``level: int``.
|
||||
|
||||
Empty / missing inputs collapse to ``None`` so callers can simply omit
|
||||
the field when writing the chunk record.
|
||||
"""
|
||||
nested = dp.get("heading")
|
||||
if isinstance(nested, dict):
|
||||
heading_text = str(nested.get("heading") or "").strip()
|
||||
parents_raw = nested.get("parent_headings") or []
|
||||
level_raw = nested.get("level", 0)
|
||||
else:
|
||||
heading_text = str(nested or "").strip()
|
||||
parents_raw = dp.get("parent_headings") or []
|
||||
level_raw = dp.get("level", 0)
|
||||
|
||||
parent_headings: list[str] = []
|
||||
if isinstance(parents_raw, list):
|
||||
for entry in parents_raw:
|
||||
text = str(entry or "").strip()
|
||||
if text:
|
||||
parent_headings.append(text)
|
||||
|
||||
try:
|
||||
level = int(level_raw or 0)
|
||||
except (TypeError, ValueError):
|
||||
level = 0
|
||||
|
||||
if not heading_text and not parent_headings and level == 0:
|
||||
return None
|
||||
|
||||
return {
|
||||
"level": level,
|
||||
"heading": heading_text,
|
||||
"parent_headings": parent_headings,
|
||||
}
|
||||
|
||||
|
||||
_HEADING_WHITESPACE_RE = re.compile(r"\s+")
|
||||
# Unicode categories stripped from a heading: control (Cc) and format (Cf)
|
||||
# chars — NULs, zero-width marks (ZWSP/ZWNJ/ZWJ/WORD JOINER/BOM are all Cf),
|
||||
# directional/format codes — that only add token noise to the LLM prompt.
|
||||
_HEADING_STRIP_CATEGORIES = frozenset({"Cc", "Cf"})
|
||||
# The only Cc chars that ``\s`` folds into a space; keep them through the strip
|
||||
# pass so they become a space below instead of gluing adjacent words. (Cannot
|
||||
# use ``str.isspace()`` for this: Python treats \x1c-\x1f as whitespace but
|
||||
# ``\s`` does not, so those would survive the strip — an explicit set matches
|
||||
# the regex.)
|
||||
_HEADING_KEEP_WS = frozenset("\t\n\r\f\v")
|
||||
|
||||
|
||||
def _clean_heading_text(text: str) -> str:
|
||||
"""Flatten a heading into one clean line for the LLM.
|
||||
|
||||
Converts the breadcrumb separator ``→`` to a space, drops every Unicode
|
||||
control (Cc) / format (Cf) char (zero-width marks, NULs, directional/format
|
||||
codes), and collapses every run of whitespace (tab, newline, NBSP,
|
||||
full-width / ideographic space, ...) into a single regular space as the
|
||||
final step. Normal CJK, Latin, digits, and punctuation are left untouched,
|
||||
and no spaces are inserted between adjacent CJK characters, so it is safe
|
||||
for Chinese headings.
|
||||
"""
|
||||
# ``→`` (U+2192, the breadcrumb separator char) must never survive inside a
|
||||
# single heading, or it would forge an extra level when the breadcrumb is
|
||||
# split back on " → " (see operate._truncate_section_context).
|
||||
text = text.replace("→", " ")
|
||||
text = "".join(
|
||||
ch
|
||||
for ch in text
|
||||
if ch in _HEADING_KEEP_WS
|
||||
or unicodedata.category(ch) not in _HEADING_STRIP_CATEGORIES
|
||||
)
|
||||
# Collapse LAST so the spaces introduced above (→, kept whitespace controls,
|
||||
# and any gap left around a removed control char) all fold into one space.
|
||||
text = _HEADING_WHITESPACE_RE.sub(" ", text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _truncate_heading_level(text: str, max_chars: int) -> str:
|
||||
"""Hard-cap a single heading level, marking elision with an ellipsis."""
|
||||
if max_chars <= 0 or len(text) <= max_chars:
|
||||
return text
|
||||
# Reserve one char for the ellipsis so the result length stays <= max_chars.
|
||||
return text[: max_chars - 1].rstrip() + "…"
|
||||
|
||||
|
||||
def _clean_and_cap_headings(headings: list[str], max_heading_len: int) -> list[str]:
|
||||
"""Clean each heading, drop empties, then hard-cap each level's length.
|
||||
|
||||
Shared by :func:`format_parent_headings` and :func:`format_heading_context`
|
||||
so both the query-stage and extraction-stage breadcrumbs apply identical
|
||||
cleaning (:func:`_clean_heading_text`) and per-level truncation, and cannot
|
||||
drift apart. Order matters: clean → drop empties → cap.
|
||||
"""
|
||||
return [
|
||||
_truncate_heading_level(c, max_heading_len)
|
||||
for c in (_clean_heading_text(h) for h in headings)
|
||||
if c
|
||||
]
|
||||
|
||||
|
||||
def format_parent_headings(
|
||||
dp: dict[str, Any],
|
||||
*,
|
||||
max_heading_len: int = DEFAULT_HEADING_LEVEL_MAX_CHARS,
|
||||
) -> str:
|
||||
"""Join a chunk's parent heading chain into ``h1 → h2 → h3``.
|
||||
|
||||
Reuses :func:`normalize_chunk_heading` so both the nested and legacy flat
|
||||
heading shapes are handled, then cleans each heading via
|
||||
:func:`_clean_heading_text` so the string sent to the LLM is a single tidy
|
||||
line. Returns an empty string when the chunk has no (non-empty) parent
|
||||
headings, so callers can simply omit the field when it is empty.
|
||||
|
||||
Each individual heading level is capped at ``max_heading_len`` characters
|
||||
(set ``<= 0`` to disable), matching :func:`format_heading_context`, so one
|
||||
runaway title cannot bloat the query context before its token truncation.
|
||||
"""
|
||||
normalized = normalize_chunk_heading(dp)
|
||||
if not normalized:
|
||||
return ""
|
||||
cleaned = _clean_and_cap_headings(normalized["parent_headings"], max_heading_len)
|
||||
return HEADING_BREADCRUMB_SEP.join(cleaned)
|
||||
|
||||
|
||||
def format_heading_context(
|
||||
dp: dict[str, Any],
|
||||
*,
|
||||
max_heading_len: int = DEFAULT_HEADING_LEVEL_MAX_CHARS,
|
||||
) -> str:
|
||||
"""Join a chunk's full heading chain (parents + current) into ``h1 → h2 → h3``.
|
||||
|
||||
Like :func:`format_parent_headings` but appends the chunk's own section
|
||||
heading after the parent chain, so the entity-extraction LLM sees the
|
||||
complete breadcrumb of the section the input text belongs to. Reuses
|
||||
:func:`normalize_chunk_heading` (handles both nested and legacy flat shapes)
|
||||
and :func:`_clean_heading_text`. Returns an empty string when the chunk
|
||||
carries no (non-empty) heading information, so callers can simply omit the
|
||||
field when it is empty.
|
||||
|
||||
Each individual heading level is capped at ``max_heading_len`` characters
|
||||
(set ``<= 0`` to disable) so one runaway title cannot bloat the prompt; the
|
||||
caller is still responsible for token-budgeting the joined breadcrumb.
|
||||
"""
|
||||
normalized = normalize_chunk_heading(dp)
|
||||
if not normalized:
|
||||
return ""
|
||||
chain = list(normalized["parent_headings"])
|
||||
if normalized["heading"]:
|
||||
chain.append(normalized["heading"])
|
||||
cleaned = _clean_and_cap_headings(chain, max_heading_len)
|
||||
return HEADING_BREADCRUMB_SEP.join(cleaned)
|
||||
|
||||
|
||||
def normalize_chunk_sidecar(dp: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""Return the canonical sidecar dict or ``None`` when absent / invalid.
|
||||
|
||||
Output shape::
|
||||
|
||||
{"type": <one of block|drawing|table|equation>,
|
||||
"id": <primary source id>,
|
||||
"refs": [{"type": ..., "id": ...}, ...]}
|
||||
|
||||
``refs`` is always materialized as a list with at least the primary id.
|
||||
Single-source chunks therefore land in storage with ``refs=[{type,id}]``
|
||||
so downstream consumers don't need to special-case the field's presence.
|
||||
"""
|
||||
sidecar = dp.get("sidecar")
|
||||
if not isinstance(sidecar, dict):
|
||||
return None
|
||||
sidecar_type = str(sidecar.get("type") or "").strip()
|
||||
sidecar_id = str(sidecar.get("id") or "").strip()
|
||||
if sidecar_type not in _SIDECAR_TYPES or not sidecar_id:
|
||||
return None
|
||||
|
||||
refs_raw = sidecar.get("refs")
|
||||
refs: list[dict[str, str]] = []
|
||||
if isinstance(refs_raw, list):
|
||||
for entry in refs_raw:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
ref_type = str(entry.get("type") or "").strip()
|
||||
ref_id = str(entry.get("id") or "").strip()
|
||||
if ref_type in _SIDECAR_TYPES and ref_id:
|
||||
refs.append({"type": ref_type, "id": ref_id})
|
||||
if not refs:
|
||||
refs = [{"type": sidecar_type, "id": sidecar_id}]
|
||||
|
||||
return {"type": sidecar_type, "id": sidecar_id, "refs": refs}
|
||||
|
||||
|
||||
# `<cite type="..." refid="...">visible text</cite>` → `visible text`.
|
||||
_CITE_RE = re.compile(
|
||||
r"<cite\b[^>]*>(.*?)</cite>",
|
||||
flags=re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
|
||||
# Inner attribute stripper used when the caller wants to *preserve* the
|
||||
# `<cite type="…">…</cite>` wrapper but drop the parser-internal `refid`.
|
||||
# Matches ` refid="…"` (leading whitespace + quoted value) so the
|
||||
# surrounding attribute layout (e.g. `type="table"`) stays intact.
|
||||
_CITE_REFID_ATTR_RE = re.compile(
|
||||
r'\s+refid\s*=\s*"[^"]*"',
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
|
||||
# Self-closing `<drawing ...>` placeholder. We keep `caption` (visible) and
|
||||
# drop `id`, `path`, `src`, `format`, etc. Tags without any caption are
|
||||
# removed entirely so they don't pollute extraction input.
|
||||
_DRAWING_RE = re.compile(
|
||||
r"<drawing\b([^>]*)/>",
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
|
||||
# Container `<equation id="..." format="...">latex</equation>`. Strip
|
||||
# identifier attributes; preserve the body and the `format` attribute so
|
||||
# extraction still sees the equation is a structured element.
|
||||
_EQUATION_RE = re.compile(
|
||||
r"<equation\b([^>]*)>(.*?)</equation>",
|
||||
flags=re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
|
||||
# Container `<table id="tb-..." format="json" caption="...">rows</table>`.
|
||||
# Native parser emits the internal ``tb-<doc>-NNNN`` identifier here, which
|
||||
# would otherwise leak into the entity-extraction prompt and become a noisy
|
||||
# entity. Strip ``id``; keep ``format`` / ``caption`` (and the body verbatim)
|
||||
# so the extractor still recognizes the element as a structured table.
|
||||
_TABLE_RE = re.compile(
|
||||
r"<table\b([^>]*)>(.*?)</table>",
|
||||
flags=re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
|
||||
# Match attribute pairs like ``caption="text with \"escapes\""``. We treat
|
||||
# only the safe identifier-style attributes; complex quoting is rare in
|
||||
# parser output.
|
||||
_ATTR_RE = re.compile(
|
||||
r'(\w+)\s*=\s*"((?:[^"\\]|\\.)*)"',
|
||||
)
|
||||
|
||||
|
||||
def _attrs_to_dict(attr_string: str) -> dict[str, str]:
|
||||
return {
|
||||
match.group(1).lower(): match.group(2)
|
||||
for match in _ATTR_RE.finditer(attr_string)
|
||||
}
|
||||
|
||||
|
||||
def _format_attrs(pairs: list[tuple[str, str]]) -> str:
|
||||
return "".join(f' {k}="{v}"' for k, v in pairs if v)
|
||||
|
||||
|
||||
def _replace_drawing(match: re.Match[str]) -> str:
|
||||
attrs = _attrs_to_dict(match.group(1))
|
||||
caption = attrs.get("caption", "")
|
||||
if not caption.strip():
|
||||
return ""
|
||||
return f"<drawing{_format_attrs([('caption', caption)])} />"
|
||||
|
||||
|
||||
def _replace_equation(match: re.Match[str]) -> str:
|
||||
attrs = _attrs_to_dict(match.group(1))
|
||||
body = match.group(2)
|
||||
keep: list[tuple[str, str]] = []
|
||||
fmt = attrs.get("format", "")
|
||||
if fmt:
|
||||
keep.append(("format", fmt))
|
||||
caption = attrs.get("caption", "")
|
||||
if caption.strip():
|
||||
keep.append(("caption", caption))
|
||||
return f"<equation{_format_attrs(keep)}>{body}</equation>"
|
||||
|
||||
|
||||
def _replace_table(match: re.Match[str]) -> str:
|
||||
attrs = _attrs_to_dict(match.group(1))
|
||||
body = match.group(2)
|
||||
keep: list[tuple[str, str]] = []
|
||||
fmt = attrs.get("format", "")
|
||||
if fmt:
|
||||
keep.append(("format", fmt))
|
||||
caption = attrs.get("caption", "")
|
||||
if caption.strip():
|
||||
keep.append(("caption", caption))
|
||||
return f"<table{_format_attrs(keep)}>{body}</table>"
|
||||
|
||||
|
||||
def strip_internal_multimodal_markup_for_extraction(
|
||||
content: str, *, keep_cite_tag: bool = False
|
||||
) -> str:
|
||||
"""Strip parser-internal identifiers from a chunk content string.
|
||||
|
||||
Only the entity-extraction prompt should receive the cleaned form;
|
||||
callers must NOT mutate the stored chunk ``content`` so query-time
|
||||
citations still resolve back to the original parser output.
|
||||
|
||||
Transformations always applied:
|
||||
|
||||
- ``<drawing id="im-…" path="…" src="…" caption="Fig 1" />``
|
||||
→ ``<drawing caption="Fig 1" />``
|
||||
(drops the entire tag when no caption is present)
|
||||
- ``<table id="tb-…" format="json" caption="…">rows</table>``
|
||||
→ ``<table format="json" caption="…">rows</table>``
|
||||
- ``<equation id="eq-…" format="latex">…</equation>``
|
||||
→ ``<equation format="latex">…</equation>``
|
||||
|
||||
Cite-tag handling depends on ``keep_cite_tag``:
|
||||
|
||||
- ``keep_cite_tag=False`` (default — entity-extraction path):
|
||||
``<cite type="…" refid="…">Table 1</cite>`` → ``Table 1``. The
|
||||
cite wrapper is dropped so the extractor does not surface it as
|
||||
a noisy structural entity.
|
||||
- ``keep_cite_tag=True`` (multimodal-analysis surrounding path):
|
||||
``<cite type="table" refid="…">Table 1</cite>`` →
|
||||
``<cite type="table">Table 1</cite>``. Only the internal
|
||||
``refid`` is removed; the wrapper survives so the VLM/LLM can
|
||||
tell visible reference labels (e.g. "Table 1") apart from inline
|
||||
prose.
|
||||
"""
|
||||
if not content:
|
||||
return content
|
||||
if keep_cite_tag:
|
||||
cleaned = _CITE_REFID_ATTR_RE.sub("", content)
|
||||
else:
|
||||
cleaned = _CITE_RE.sub(lambda m: m.group(1), content)
|
||||
cleaned = _DRAWING_RE.sub(_replace_drawing, cleaned)
|
||||
cleaned = _TABLE_RE.sub(_replace_table, cleaned)
|
||||
cleaned = _EQUATION_RE.sub(_replace_equation, cleaned)
|
||||
return cleaned
|
||||
|
||||
|
||||
__all__ = [
|
||||
"normalize_chunk_heading",
|
||||
"format_parent_headings",
|
||||
"format_heading_context",
|
||||
"HEADING_BREADCRUMB_SEP",
|
||||
"normalize_chunk_sidecar",
|
||||
"strip_internal_multimodal_markup_for_extraction",
|
||||
]
|
||||
@@ -0,0 +1,66 @@
|
||||
"""LightRAG chunking strategies.
|
||||
|
||||
Two contracts coexist intentionally:
|
||||
|
||||
- **Legacy contract** — :func:`chunking_by_token_size` keeps its
|
||||
historical 6-positional-arg signature
|
||||
|
||||
``(tokenizer, content, split_by_character,
|
||||
split_by_character_only, chunk_overlap_token_size,
|
||||
chunk_token_size)``
|
||||
|
||||
so externally-supplied :attr:`lightrag.LightRAG.chunking_func`
|
||||
implementations continue to work unchanged. The legacy contract is
|
||||
only invoked when ``process_options`` does NOT specify a chunking
|
||||
selector (i.e. ``chunking_explicit`` is False) — typically direct
|
||||
:meth:`LightRAG.ainsert` calls with raw text.
|
||||
|
||||
- **File-chunker contract** — for documents whose ``process_options``
|
||||
explicitly selects a chunking strategy, the file-based dispatcher in
|
||||
``_PipelineMixin.process_single_document`` reads
|
||||
``doc_process_opts.chunking`` and routes to a chunker following the
|
||||
standardized signature
|
||||
|
||||
``(tokenizer, content, chunk_token_size, *,
|
||||
<strategy-specific kwargs>)``
|
||||
|
||||
Currently shipped file chunkers:
|
||||
|
||||
- :func:`chunking_by_fixed_token` — the ``"F"`` strategy. Same
|
||||
algorithm as :func:`chunking_by_token_size`, surfaced under the
|
||||
new contract.
|
||||
- :func:`chunking_by_recursive_character` — the ``"R"`` strategy.
|
||||
Wraps LangChain ``RecursiveCharacterTextSplitter``; recursively
|
||||
splits on a separator cascade with token-aware sizing.
|
||||
- :func:`chunking_by_semantic_vector` — the ``"V"`` strategy.
|
||||
Wraps LangChain ``SemanticChunker``; sentence-level embedding
|
||||
similarity finds breakpoints. Async; needs an
|
||||
:class:`~lightrag.utils.EmbeddingFunc`.
|
||||
- :func:`chunking_by_paragraph_semantic` — the ``"P"`` strategy.
|
||||
Heading-aware semantic chunker; consumes the docx-native
|
||||
``.blocks.jsonl`` sidecar. Falls back to R when the sidecar is
|
||||
missing or unreadable.
|
||||
|
||||
See ``docs/ParagraphSemanticChunking-zh.md`` for the algorithm behind
|
||||
the ``"P"`` strategy and ``docs/FileProcessingConfiguration-zh.md`` for
|
||||
how ``process_options`` and the new ``chunk_options`` snapshot drive
|
||||
chunker selection per document.
|
||||
"""
|
||||
|
||||
from lightrag.chunker.paragraph_semantic import chunking_by_paragraph_semantic
|
||||
from lightrag.chunker.recursive_character import (
|
||||
chunking_by_recursive_character,
|
||||
)
|
||||
from lightrag.chunker.semantic_vector import chunking_by_semantic_vector
|
||||
from lightrag.chunker.token_size import (
|
||||
chunking_by_fixed_token,
|
||||
chunking_by_token_size,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"chunking_by_fixed_token",
|
||||
"chunking_by_paragraph_semantic",
|
||||
"chunking_by_recursive_character",
|
||||
"chunking_by_semantic_vector",
|
||||
"chunking_by_token_size",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,402 @@
|
||||
"""Recursive character chunking — the ``"R"`` strategy.
|
||||
|
||||
Wraps LangChain's :class:`RecursiveCharacterTextSplitter` and delivers
|
||||
output rows in the LightRAG file-chunker schema. The splitter walks the
|
||||
``separators`` list from longest semantic boundary (``\\n\\n`` by default)
|
||||
to weakest (the empty string), recursively re-splitting any segment that
|
||||
still exceeds the token cap.
|
||||
|
||||
Token accounting goes through the LightRAG :class:`Tokenizer` via the
|
||||
``length_function`` plug-in — without that, ``chunk_size`` would be
|
||||
measured in characters and ``chunk_token_size`` would lose its meaning.
|
||||
|
||||
Output cap is *not* enforced internally: oversized segments are produced
|
||||
when no separator can break them, and
|
||||
:func:`lightrag.utils.enforce_chunk_token_limit_before_embedding` does the
|
||||
final hard split before embedding.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any
|
||||
|
||||
from lightrag.utils import Tokenizer, logger
|
||||
|
||||
try:
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
_LANGCHAIN_TEXT_SPLITTERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
_LANGCHAIN_TEXT_SPLITTERS_AVAILABLE = False
|
||||
RecursiveCharacterTextSplitter = None # type: ignore[assignment]
|
||||
|
||||
|
||||
_SpanPiece = tuple[str, int, int]
|
||||
|
||||
|
||||
def _split_text_with_regex_spans(
|
||||
text: str,
|
||||
separator_pattern: str,
|
||||
*,
|
||||
keep_separator: bool | str,
|
||||
base_offset: int,
|
||||
) -> list[_SpanPiece]:
|
||||
"""Mirror LangChain's regex split while retaining source offsets."""
|
||||
if not separator_pattern:
|
||||
return [
|
||||
(char, base_offset + index, base_offset + index + 1)
|
||||
for index, char in enumerate(text)
|
||||
if char
|
||||
]
|
||||
|
||||
matches = list(re.finditer(separator_pattern, text))
|
||||
if not matches:
|
||||
return [(text, base_offset, base_offset + len(text))] if text else []
|
||||
|
||||
pieces: list[_SpanPiece] = []
|
||||
if keep_separator:
|
||||
if keep_separator == "end":
|
||||
cursor = 0
|
||||
for match in matches:
|
||||
if match.end() > cursor:
|
||||
pieces.append(
|
||||
(
|
||||
text[cursor : match.end()],
|
||||
base_offset + cursor,
|
||||
base_offset + match.end(),
|
||||
)
|
||||
)
|
||||
cursor = match.end()
|
||||
if cursor < len(text):
|
||||
pieces.append(
|
||||
(text[cursor:], base_offset + cursor, base_offset + len(text))
|
||||
)
|
||||
else:
|
||||
first = matches[0]
|
||||
if first.start() > 0:
|
||||
pieces.append(
|
||||
(text[: first.start()], base_offset, base_offset + first.start())
|
||||
)
|
||||
for index, match in enumerate(matches):
|
||||
end = (
|
||||
matches[index + 1].start()
|
||||
if index + 1 < len(matches)
|
||||
else len(text)
|
||||
)
|
||||
if end > match.start():
|
||||
pieces.append(
|
||||
(
|
||||
text[match.start() : end],
|
||||
base_offset + match.start(),
|
||||
base_offset + end,
|
||||
)
|
||||
)
|
||||
else:
|
||||
cursor = 0
|
||||
for match in matches:
|
||||
if match.start() > cursor:
|
||||
pieces.append(
|
||||
(
|
||||
text[cursor : match.start()],
|
||||
base_offset + cursor,
|
||||
base_offset + match.start(),
|
||||
)
|
||||
)
|
||||
cursor = match.end()
|
||||
if cursor < len(text):
|
||||
pieces.append(
|
||||
(text[cursor:], base_offset + cursor, base_offset + len(text))
|
||||
)
|
||||
|
||||
return [piece for piece in pieces if piece[0]]
|
||||
|
||||
|
||||
def _join_span_pieces(
|
||||
pieces: list[_SpanPiece],
|
||||
separator: str,
|
||||
*,
|
||||
strip_whitespace: bool,
|
||||
) -> _SpanPiece | None:
|
||||
"""Join split pieces exactly as LangChain does and compute the trimmed span."""
|
||||
if not pieces:
|
||||
return None
|
||||
|
||||
chars: list[str] = []
|
||||
char_offsets: list[int] = []
|
||||
for index, (fragment, start, end) in enumerate(pieces):
|
||||
if index > 0 and separator:
|
||||
previous_end = pieces[index - 1][2]
|
||||
for sep_index, sep_char in enumerate(separator):
|
||||
chars.append(sep_char)
|
||||
char_offsets.append(previous_end + sep_index)
|
||||
chars.extend(fragment)
|
||||
char_offsets.extend(range(start, end))
|
||||
|
||||
text = "".join(chars)
|
||||
if strip_whitespace:
|
||||
left = 0
|
||||
right = len(text)
|
||||
while left < right and text[left].isspace():
|
||||
left += 1
|
||||
while right > left and text[right - 1].isspace():
|
||||
right -= 1
|
||||
else:
|
||||
left, right = 0, len(text)
|
||||
|
||||
if left >= right:
|
||||
return None
|
||||
return text[left:right], char_offsets[left], char_offsets[right - 1] + 1
|
||||
|
||||
|
||||
def _merge_splits_with_spans(
|
||||
splits: Sequence[_SpanPiece],
|
||||
separator: str,
|
||||
*,
|
||||
chunk_size: int,
|
||||
chunk_overlap: int,
|
||||
length_function: Callable[[str], int],
|
||||
strip_whitespace: bool,
|
||||
) -> list[_SpanPiece]:
|
||||
"""Mirror ``TextSplitter._merge_splits`` while preserving source spans."""
|
||||
separator_len = length_function(separator)
|
||||
docs: list[_SpanPiece] = []
|
||||
current_doc: list[_SpanPiece] = []
|
||||
total = 0
|
||||
|
||||
for split in splits:
|
||||
split_len = length_function(split[0])
|
||||
if (
|
||||
total + split_len + (separator_len if len(current_doc) > 0 else 0)
|
||||
> chunk_size
|
||||
):
|
||||
if total > chunk_size:
|
||||
logger.warning(
|
||||
"Created a chunk of size %d, which is longer than the specified %d",
|
||||
total,
|
||||
chunk_size,
|
||||
)
|
||||
if len(current_doc) > 0:
|
||||
doc = _join_span_pieces(
|
||||
current_doc,
|
||||
separator,
|
||||
strip_whitespace=strip_whitespace,
|
||||
)
|
||||
if doc is not None:
|
||||
docs.append(doc)
|
||||
while total > chunk_overlap or (
|
||||
total + split_len + (separator_len if len(current_doc) > 0 else 0)
|
||||
> chunk_size
|
||||
and total > 0
|
||||
):
|
||||
total -= length_function(current_doc[0][0]) + (
|
||||
separator_len if len(current_doc) > 1 else 0
|
||||
)
|
||||
current_doc = current_doc[1:]
|
||||
current_doc.append(split)
|
||||
total += split_len + (separator_len if len(current_doc) > 1 else 0)
|
||||
|
||||
doc = _join_span_pieces(
|
||||
current_doc,
|
||||
separator,
|
||||
strip_whitespace=strip_whitespace,
|
||||
)
|
||||
if doc is not None:
|
||||
docs.append(doc)
|
||||
return docs
|
||||
|
||||
|
||||
def _split_text_with_spans(
|
||||
text: str,
|
||||
*,
|
||||
base_offset: int,
|
||||
separators: Sequence[str],
|
||||
chunk_size: int,
|
||||
chunk_overlap: int,
|
||||
length_function: Callable[[str], int],
|
||||
keep_separator: bool | str,
|
||||
is_separator_regex: bool,
|
||||
strip_whitespace: bool,
|
||||
) -> list[_SpanPiece]:
|
||||
"""Mirror ``RecursiveCharacterTextSplitter._split_text`` with offsets."""
|
||||
separator = separators[-1]
|
||||
new_separators: Sequence[str] = []
|
||||
for index, candidate in enumerate(separators):
|
||||
separator_pattern = candidate if is_separator_regex else re.escape(candidate)
|
||||
if not candidate:
|
||||
separator = candidate
|
||||
break
|
||||
if re.search(separator_pattern, text):
|
||||
separator = candidate
|
||||
new_separators = separators[index + 1 :]
|
||||
break
|
||||
|
||||
separator_pattern = separator if is_separator_regex else re.escape(separator)
|
||||
splits = _split_text_with_regex_spans(
|
||||
text,
|
||||
separator_pattern,
|
||||
keep_separator=keep_separator,
|
||||
base_offset=base_offset,
|
||||
)
|
||||
|
||||
final_chunks: list[_SpanPiece] = []
|
||||
good_splits: list[_SpanPiece] = []
|
||||
merge_separator = "" if keep_separator else separator
|
||||
for split in splits:
|
||||
if length_function(split[0]) < chunk_size:
|
||||
good_splits.append(split)
|
||||
else:
|
||||
if good_splits:
|
||||
final_chunks.extend(
|
||||
_merge_splits_with_spans(
|
||||
good_splits,
|
||||
merge_separator,
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
length_function=length_function,
|
||||
strip_whitespace=strip_whitespace,
|
||||
)
|
||||
)
|
||||
good_splits = []
|
||||
if not new_separators:
|
||||
final_chunks.append(split)
|
||||
else:
|
||||
final_chunks.extend(
|
||||
_split_text_with_spans(
|
||||
split[0],
|
||||
base_offset=split[1],
|
||||
separators=new_separators,
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
length_function=length_function,
|
||||
keep_separator=keep_separator,
|
||||
is_separator_regex=is_separator_regex,
|
||||
strip_whitespace=strip_whitespace,
|
||||
)
|
||||
)
|
||||
if good_splits:
|
||||
final_chunks.extend(
|
||||
_merge_splits_with_spans(
|
||||
good_splits,
|
||||
merge_separator,
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
length_function=length_function,
|
||||
strip_whitespace=strip_whitespace,
|
||||
)
|
||||
)
|
||||
return final_chunks
|
||||
|
||||
|
||||
def chunking_by_recursive_character(
|
||||
tokenizer: Tokenizer,
|
||||
content: str,
|
||||
chunk_token_size: int = 1200,
|
||||
*,
|
||||
chunk_overlap_token_size: int = 100,
|
||||
separators: list[str] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Recursive character splitter — the ``"R"`` chunking strategy.
|
||||
|
||||
Args:
|
||||
tokenizer: LightRAG tokenizer; used as the length function so
|
||||
``chunk_token_size`` and ``chunk_overlap_token_size`` are
|
||||
interpreted in tokens, not characters.
|
||||
content: Text to split.
|
||||
chunk_token_size: Hard target size for each chunk (tokens).
|
||||
chunk_overlap_token_size: Token overlap between adjacent chunks.
|
||||
separators: Cascade of split candidates. ``None`` defers to
|
||||
LangChain's defaults: ``["\\n\\n", "\\n", " ", ""]``.
|
||||
|
||||
Returns:
|
||||
Ordered list of ``{"tokens", "content", "chunk_order_index"}``
|
||||
dicts.
|
||||
"""
|
||||
if not _LANGCHAIN_TEXT_SPLITTERS_AVAILABLE:
|
||||
raise ImportError(
|
||||
"langchain-text-splitters is required for the 'R' chunking "
|
||||
"strategy; install with `pip install langchain-text-splitters>=0.3`."
|
||||
)
|
||||
|
||||
if not content or not content.strip():
|
||||
return []
|
||||
|
||||
def length_function(text: str) -> int:
|
||||
return len(tokenizer.encode(text))
|
||||
|
||||
splitter_kwargs: dict[str, Any] = {
|
||||
"chunk_size": max(int(chunk_token_size), 1),
|
||||
"chunk_overlap": max(int(chunk_overlap_token_size), 0),
|
||||
"length_function": length_function,
|
||||
"strip_whitespace": True,
|
||||
}
|
||||
if separators is not None:
|
||||
splitter_kwargs["separators"] = list(separators)
|
||||
|
||||
splitter = RecursiveCharacterTextSplitter(**splitter_kwargs)
|
||||
|
||||
# We deliberately do *not* request LangChain's ``add_start_index``. That
|
||||
# offset is computed with a character-vs-token unit mismatch when a
|
||||
# token-based ``length_function`` is in play, and text-search recovery is
|
||||
# ambiguous for repeated blocks. Instead we mirror LangChain's split/merge
|
||||
# control flow while carrying each split unit's source offsets through it.
|
||||
pieces = _split_text_with_spans(
|
||||
content,
|
||||
base_offset=0,
|
||||
separators=list(splitter._separators),
|
||||
chunk_size=int(splitter._chunk_size),
|
||||
chunk_overlap=int(splitter._chunk_overlap),
|
||||
length_function=length_function,
|
||||
keep_separator=splitter._keep_separator,
|
||||
is_separator_regex=bool(splitter._is_separator_regex),
|
||||
strip_whitespace=bool(splitter._strip_whitespace),
|
||||
)
|
||||
results: list[dict[str, Any]] = []
|
||||
for raw_body, start_index, end_index in pieces:
|
||||
left = 0
|
||||
right = len(raw_body)
|
||||
while left < right and raw_body[left].isspace():
|
||||
left += 1
|
||||
while right > left and raw_body[right - 1].isspace():
|
||||
right -= 1
|
||||
body = raw_body[left:right]
|
||||
if not body:
|
||||
continue
|
||||
start_index += left
|
||||
end_index -= len(raw_body) - right
|
||||
results.append(
|
||||
{
|
||||
"tokens": len(tokenizer.encode(body)),
|
||||
"content": body,
|
||||
"chunk_order_index": len(results),
|
||||
"_source_span": {"start": start_index, "end": end_index},
|
||||
}
|
||||
)
|
||||
|
||||
if not results:
|
||||
# Defensive: splitter returned only whitespace fragments. Fall
|
||||
# through with a single chunk of stripped content so downstream
|
||||
# callers always receive at least one row when input is non-empty.
|
||||
logger.warning(
|
||||
"[recursive_character] splitter produced no non-empty chunks "
|
||||
"for %d-char input; emitting single fallback chunk.",
|
||||
len(content),
|
||||
)
|
||||
body = content.strip()
|
||||
if body:
|
||||
start = content.find(body)
|
||||
results.append(
|
||||
{
|
||||
"tokens": len(tokenizer.encode(body)),
|
||||
"content": body,
|
||||
"chunk_order_index": 0,
|
||||
**(
|
||||
{"_source_span": {"start": start, "end": start + len(body)}}
|
||||
if start >= 0
|
||||
else {}
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,349 @@
|
||||
"""Semantic vector chunking — the ``"V"`` strategy.
|
||||
|
||||
Wraps LangChain's :class:`SemanticChunker` (from ``langchain-experimental``)
|
||||
which splits text by sentence embeddings: it first segments the input into
|
||||
sentences, embeds each sentence (in adjacent windows of ``buffer_size``),
|
||||
and finds breakpoints where the cosine distance between consecutive
|
||||
windows crosses a threshold derived from the chosen distribution
|
||||
(``percentile`` / ``standard_deviation`` / ``interquartile`` /
|
||||
``gradient``).
|
||||
|
||||
The chunker exposed here is ``async`` because LightRAG's
|
||||
:class:`EmbeddingFunc` is async. Internally we call SemanticChunker
|
||||
synchronously inside :func:`asyncio.to_thread` and bridge the embedding
|
||||
calls back to the main event loop via
|
||||
:func:`asyncio.run_coroutine_threadsafe`.
|
||||
|
||||
Caveats:
|
||||
- SemanticChunker does NOT enforce a maximum chunk size; the caller's
|
||||
``chunk_token_size`` is *advisory* here. Oversized chunks will be
|
||||
hard-split before embedding by
|
||||
:func:`lightrag.utils.enforce_chunk_token_limit_before_embedding`.
|
||||
- When ``embedding_func`` is ``None`` we log a warning and fall back to
|
||||
:func:`lightrag.chunker.chunking_by_recursive_character` — V's only
|
||||
differentiator is embeddings, and R is the closest structural-only
|
||||
alternative.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lightrag.constants import DEFAULT_SENTENCE_SPLIT_REGEX
|
||||
from lightrag.utils import EmbeddingFunc, Tokenizer, logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_experimental.text_splitter import (
|
||||
SemanticChunker as SemanticChunkerType,
|
||||
)
|
||||
else:
|
||||
SemanticChunkerType = Any
|
||||
|
||||
try:
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langchain_experimental.text_splitter import SemanticChunker
|
||||
|
||||
_LANGCHAIN_EXPERIMENTAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
_LANGCHAIN_EXPERIMENTAL_AVAILABLE = False
|
||||
Embeddings = object # type: ignore[assignment,misc]
|
||||
SemanticChunker = None # type: ignore[assignment]
|
||||
|
||||
|
||||
class _AsyncEmbeddingFuncAdapter(Embeddings):
|
||||
"""Bridge a LightRAG :class:`EmbeddingFunc` (async) to LangChain's
|
||||
sync :class:`Embeddings` interface used by ``SemanticChunker``.
|
||||
|
||||
The adapter must be constructed inside the running event loop so it
|
||||
can capture the loop reference; the blocking ``embed_documents`` /
|
||||
``embed_query`` calls are then made from a worker thread (via
|
||||
:func:`asyncio.to_thread` in the public chunker) and bounce back to
|
||||
the captured loop with :func:`asyncio.run_coroutine_threadsafe`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedding_func: EmbeddingFunc,
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
) -> None:
|
||||
self._embedding_func = embedding_func
|
||||
self._loop = loop
|
||||
|
||||
def _run(self, texts: list[str], context: str) -> list[list[float]]:
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
self._embedding_func(texts, context=context),
|
||||
self._loop,
|
||||
)
|
||||
result = future.result()
|
||||
return [list(map(float, vec)) for vec in result]
|
||||
|
||||
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._run(list(texts), context="document")
|
||||
|
||||
def embed_query(self, text: str) -> list[float]:
|
||||
return self._run([text], context="query")[0]
|
||||
|
||||
|
||||
def _sentence_spans(text: str, sentences: list[str]) -> list[tuple[int, int]]:
|
||||
spans: list[tuple[int, int]] = []
|
||||
cursor = 0
|
||||
for sentence in sentences:
|
||||
if not sentence:
|
||||
spans.append((cursor, cursor))
|
||||
continue
|
||||
start = text.find(sentence, cursor)
|
||||
if start < 0:
|
||||
start = text.find(sentence)
|
||||
if start < 0:
|
||||
start = cursor
|
||||
end = start + len(sentence)
|
||||
spans.append((start, end))
|
||||
cursor = end
|
||||
return spans
|
||||
|
||||
|
||||
def _trim_span(text: str, start: int, end: int) -> tuple[int, int]:
|
||||
start = max(0, min(start, len(text)))
|
||||
end = max(start, min(end, len(text)))
|
||||
while start < end and text[start].isspace():
|
||||
start += 1
|
||||
while end > start and text[end - 1].isspace():
|
||||
end -= 1
|
||||
return start, end
|
||||
|
||||
|
||||
def _semantic_groups_with_spans(
|
||||
splitter: SemanticChunkerType,
|
||||
text: str,
|
||||
) -> list[tuple[str, int, int]]:
|
||||
"""Mirror SemanticChunker grouping while keeping original source spans.
|
||||
|
||||
.. warning::
|
||||
This re-implements the body of ``SemanticChunker.split_text`` so each group
|
||||
carries its exact source span (``text[start:end]``) instead of the upstream
|
||||
``" ".join(sentences)`` reflow. It relies on **private** members
|
||||
(``sentence_split_regex``, ``breakpoint_threshold_type``, ``min_chunk_size``,
|
||||
``number_of_chunks``, ``_calculate_sentence_distances``,
|
||||
``_threshold_from_clusters``, ``_calculate_breakpoint_threshold``). Verified
|
||||
byte-for-byte against ``langchain-experimental`` 0.3.2–0.4.x (the range pinned
|
||||
in ``pyproject.toml``: ``langchain-experimental>=0.3.2,<1``). If that pin is
|
||||
widened, re-verify against the new upstream ``split_text`` —
|
||||
``tests/chunker/test_chunker_semantic_vector.py`` has a drift guard that
|
||||
compares this mirror's grouping to the live ``splitter.split_text`` output.
|
||||
"""
|
||||
single_sentences_list = re.split(splitter.sentence_split_regex, text)
|
||||
spans = _sentence_spans(text, single_sentences_list)
|
||||
|
||||
def _group(start_index: int, end_index: int) -> tuple[str, int, int] | None:
|
||||
start, _ = spans[start_index]
|
||||
_, end = spans[end_index]
|
||||
start, end = _trim_span(text, start, end)
|
||||
if start >= end:
|
||||
return None
|
||||
return text[start:end], start, end
|
||||
|
||||
if len(single_sentences_list) == 1:
|
||||
group = _group(0, 0)
|
||||
return [group] if group else []
|
||||
if (
|
||||
splitter.breakpoint_threshold_type == "gradient"
|
||||
and len(single_sentences_list) == 2
|
||||
):
|
||||
return [g for i in range(2) if (g := _group(i, i)) is not None]
|
||||
|
||||
distances, sentences = splitter._calculate_sentence_distances(single_sentences_list)
|
||||
if splitter.number_of_chunks is not None:
|
||||
breakpoint_distance_threshold = splitter._threshold_from_clusters(distances)
|
||||
breakpoint_array = distances
|
||||
else:
|
||||
breakpoint_distance_threshold, breakpoint_array = (
|
||||
splitter._calculate_breakpoint_threshold(distances)
|
||||
)
|
||||
|
||||
indices_above_thresh = [
|
||||
i for i, x in enumerate(breakpoint_array) if x > breakpoint_distance_threshold
|
||||
]
|
||||
|
||||
chunks: list[tuple[str, int, int]] = []
|
||||
start_index = 0
|
||||
for index in indices_above_thresh:
|
||||
end_index = index
|
||||
group_sentences = sentences[start_index : end_index + 1]
|
||||
combined_text = " ".join([d["sentence"] for d in group_sentences])
|
||||
if (
|
||||
splitter.min_chunk_size is not None
|
||||
and len(combined_text) < splitter.min_chunk_size
|
||||
):
|
||||
continue
|
||||
group = _group(start_index, end_index)
|
||||
if group is not None:
|
||||
chunks.append(group)
|
||||
start_index = index + 1
|
||||
|
||||
if start_index < len(sentences):
|
||||
group = _group(start_index, len(sentences) - 1)
|
||||
if group is not None:
|
||||
chunks.append(group)
|
||||
return chunks
|
||||
|
||||
|
||||
async def chunking_by_semantic_vector(
|
||||
tokenizer: Tokenizer,
|
||||
content: str,
|
||||
chunk_token_size: int = 1200,
|
||||
*,
|
||||
embedding_func: EmbeddingFunc | None = None,
|
||||
breakpoint_threshold_type: str = "percentile",
|
||||
breakpoint_threshold_amount: float | None = None,
|
||||
buffer_size: int = 1,
|
||||
sentence_split_regex: str = DEFAULT_SENTENCE_SPLIT_REGEX,
|
||||
number_of_chunks: int | None = None,
|
||||
min_chunk_size: int | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Semantic vector chunker — the ``"V"`` chunking strategy.
|
||||
|
||||
Args:
|
||||
tokenizer: LightRAG tokenizer (used for output token counts).
|
||||
content: Text to split.
|
||||
chunk_token_size: Hard upper bound (tokens). SemanticChunker does
|
||||
NOT enforce a maximum natively, so any piece that exceeds
|
||||
this value is re-split via
|
||||
:func:`chunking_by_recursive_character` before being emitted.
|
||||
embedding_func: LightRAG :class:`EmbeddingFunc`. When ``None``
|
||||
this chunker logs a warning and falls back to
|
||||
:func:`chunking_by_recursive_character`.
|
||||
breakpoint_threshold_type: ``percentile`` | ``standard_deviation``
|
||||
| ``interquartile`` | ``gradient`` (LangChain default:
|
||||
``percentile``).
|
||||
breakpoint_threshold_amount: Threshold magnitude. ``None`` lets
|
||||
LangChain pick the per-type default (e.g. 95 for percentile).
|
||||
buffer_size: Number of adjacent sentences combined when computing
|
||||
distances (LangChain default: 1).
|
||||
sentence_split_regex: Pattern fed to LangChain's
|
||||
:class:`SemanticChunker` for the initial sentence split.
|
||||
Default extends the upstream English-only pattern with
|
||||
Chinese sentence terminators ``。？！`` so mixed-language and
|
||||
pure-Chinese inputs split correctly.
|
||||
number_of_chunks: Optional target chunk count (LangChain SemanticChunker).
|
||||
min_chunk_size: Optional minimum character size for semantic groups.
|
||||
|
||||
Returns:
|
||||
Ordered list of ``{"tokens", "content", "chunk_order_index"}``
|
||||
dicts.
|
||||
"""
|
||||
if not content or not content.strip():
|
||||
return []
|
||||
|
||||
if embedding_func is None:
|
||||
# V's only differentiator is embeddings — without them the
|
||||
# closest neighbour is R's structural splitting. V chunks are
|
||||
# non-overlapping by design (semantic boundaries), so the
|
||||
# fallback uses ``chunk_overlap_token_size=0`` to preserve that
|
||||
# semantic and avoid LangChain's "overlap > chunk_size" guard
|
||||
# for very small ``chunk_token_size``.
|
||||
logger.warning(
|
||||
"[semantic_vector] embedding_func is None; falling back to "
|
||||
"recursive-character chunking."
|
||||
)
|
||||
from lightrag.chunker.recursive_character import (
|
||||
chunking_by_recursive_character,
|
||||
)
|
||||
|
||||
return chunking_by_recursive_character(
|
||||
tokenizer,
|
||||
content,
|
||||
chunk_token_size,
|
||||
chunk_overlap_token_size=0,
|
||||
)
|
||||
|
||||
if not _LANGCHAIN_EXPERIMENTAL_AVAILABLE:
|
||||
raise ImportError(
|
||||
"langchain-experimental is required for the 'V' chunking "
|
||||
"strategy; install with `pip install langchain-experimental>=0.3.2`."
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
adapter = _AsyncEmbeddingFuncAdapter(embedding_func, loop)
|
||||
|
||||
chunker_kwargs: dict[str, Any] = {
|
||||
"embeddings": adapter,
|
||||
"buffer_size": int(buffer_size),
|
||||
"breakpoint_threshold_type": breakpoint_threshold_type,
|
||||
"sentence_split_regex": sentence_split_regex,
|
||||
"number_of_chunks": number_of_chunks,
|
||||
"min_chunk_size": min_chunk_size,
|
||||
}
|
||||
if breakpoint_threshold_amount is not None:
|
||||
chunker_kwargs["breakpoint_threshold_amount"] = float(
|
||||
breakpoint_threshold_amount
|
||||
)
|
||||
|
||||
splitter = SemanticChunker(**chunker_kwargs)
|
||||
pieces = await asyncio.to_thread(_semantic_groups_with_spans, splitter, content)
|
||||
|
||||
# SemanticChunker has no internal size cap; oversized pieces here
|
||||
# would otherwise rely on the embedding-time hard fallback (which
|
||||
# uses ``embedding_token_limit``, not ``chunk_token_size``) to split
|
||||
# them. Enforce ``chunk_token_size`` directly via R for any piece
|
||||
# that exceeds it so the user-configured size is actually honored.
|
||||
# Lazy import dodges the recursive_character ↔ semantic_vector
|
||||
# circular dependency (same pattern as the embedding-None fallback
|
||||
# above).
|
||||
from lightrag.chunker.recursive_character import (
|
||||
chunking_by_recursive_character,
|
||||
)
|
||||
|
||||
target_max = max(int(chunk_token_size), 1)
|
||||
results: list[dict[str, Any]] = []
|
||||
for piece, source_start, source_end in pieces:
|
||||
body = piece.strip()
|
||||
if not body:
|
||||
continue
|
||||
piece_tokens = len(tokenizer.encode(body))
|
||||
if piece_tokens <= target_max:
|
||||
results.append(
|
||||
{
|
||||
"tokens": piece_tokens,
|
||||
"content": body,
|
||||
"chunk_order_index": len(results),
|
||||
"_source_span": {
|
||||
"start": source_start,
|
||||
"end": source_end,
|
||||
},
|
||||
}
|
||||
)
|
||||
continue
|
||||
# Oversized semantic piece: re-split via R while preserving the
|
||||
# surrounding chunk order. ``chunk_overlap_token_size=0`` keeps
|
||||
# V's non-overlapping semantics.
|
||||
sub_pieces = chunking_by_recursive_character(
|
||||
tokenizer,
|
||||
body,
|
||||
target_max,
|
||||
chunk_overlap_token_size=0,
|
||||
)
|
||||
for sub in sub_pieces:
|
||||
sub_body = sub.get("content", "")
|
||||
if not sub_body:
|
||||
continue
|
||||
sub_span = sub.get("_source_span")
|
||||
source_span = None
|
||||
if isinstance(sub_span, dict):
|
||||
try:
|
||||
source_span = {
|
||||
"start": source_start + int(sub_span["start"]),
|
||||
"end": source_start + int(sub_span["end"]),
|
||||
}
|
||||
except (KeyError, TypeError, ValueError):
|
||||
source_span = None
|
||||
results.append(
|
||||
{
|
||||
"tokens": sub.get("tokens", len(tokenizer.encode(sub_body))),
|
||||
"content": sub_body,
|
||||
"chunk_order_index": len(results),
|
||||
**({"_source_span": source_span} if source_span else {}),
|
||||
}
|
||||
)
|
||||
return results
|
||||
@@ -0,0 +1,263 @@
|
||||
"""Fixed-size token-window chunking — the LightRAG default strategy.
|
||||
|
||||
Chunks the input text into windows of at most ``chunk_token_size`` tokens
|
||||
with ``chunk_overlap_token_size`` of overlap between adjacent windows.
|
||||
When ``split_by_character`` is supplied, the splitter first segments on
|
||||
that delimiter and then either tokenizes each segment as-is
|
||||
(``split_by_character_only=True``) or further sub-splits any segment
|
||||
that exceeds the token cap.
|
||||
|
||||
Two entry points are exported:
|
||||
|
||||
- :func:`chunking_by_token_size` — the **legacy 6-arg signature**
|
||||
used as the default value for :attr:`lightrag.LightRAG.chunking_func`.
|
||||
Kept for backward compatibility so externally-supplied chunking
|
||||
functions can continue to drop in unchanged.
|
||||
|
||||
- :func:`chunking_by_fixed_token` — the same algorithm exposed under
|
||||
the **new file-chunker contract** (standard prefix
|
||||
``(tokenizer, content, chunk_token_size)`` plus keyword-only
|
||||
knobs). Used by the file-based chunking dispatcher in
|
||||
``process_single_document`` for ``doc_process_opts.chunking == "F"``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from lightrag.exceptions import ChunkTokenLimitExceededError
|
||||
from lightrag.utils import Tokenizer, logger
|
||||
|
||||
|
||||
def _trimmed_span(content: str, start: int, end: int) -> tuple[int, int]:
|
||||
"""Return the source span after applying the chunker's ``.strip()``."""
|
||||
start = max(0, min(start, len(content)))
|
||||
end = max(start, min(end, len(content)))
|
||||
while start < end and content[start].isspace():
|
||||
start += 1
|
||||
while end > start and content[end - 1].isspace():
|
||||
end -= 1
|
||||
return start, end
|
||||
|
||||
|
||||
def _source_span(content: str, start: int, end: int) -> dict[str, int] | None:
|
||||
start, end = _trimmed_span(content, start, end)
|
||||
if start >= end:
|
||||
return None
|
||||
return {"start": start, "end": end}
|
||||
|
||||
|
||||
def _token_window_source_span(
|
||||
tokenizer: Tokenizer,
|
||||
content: str,
|
||||
tokens: list[int],
|
||||
start_token: int,
|
||||
end_token: int,
|
||||
*,
|
||||
anchor: tuple[int, int],
|
||||
) -> tuple[dict[str, int] | None, tuple[int, int]]:
|
||||
"""Map a decoded token window back to its exact source span.
|
||||
|
||||
``anchor`` is the previous window's *verified* ``(start_token, start_char)``.
|
||||
Window starts are monotonically increasing, so instead of re-decoding the whole
|
||||
``tokens[:start_token]`` prefix (O(N) per window → O(N²) overall) we decode only
|
||||
the delta ``tokens[anchor_token:start_token]`` (≈ one chunking step) to predict
|
||||
the start char. The predicted offset is then verified against ``content`` exactly
|
||||
as a full prefix decode would be: byte-level BPE decode is non-concatenative at a
|
||||
multi-byte UTF-8 boundary, so a delta can be off by the few chars of one split
|
||||
char — the ±32 ``find`` fallback corrects that, and re-anchoring on the verified
|
||||
position each call keeps the error from accumulating. Net cost is O(N) total
|
||||
while the located span stays byte-exact.
|
||||
|
||||
Returns ``(span, new_anchor)``. On an unlocatable (U+FFFD) window the span is
|
||||
``None`` and the anchor is returned unchanged so the next window still predicts
|
||||
from the last verified position.
|
||||
"""
|
||||
anchor_token, anchor_char = anchor
|
||||
window = tokenizer.decode(tokens[start_token:end_token])
|
||||
if start_token >= anchor_token:
|
||||
start = anchor_char + len(tokenizer.decode(tokens[anchor_token:start_token]))
|
||||
else: # non-monotonic caller (not expected) — fall back to a full prefix decode
|
||||
start = len(tokenizer.decode(tokens[:start_token]))
|
||||
end = start + len(window)
|
||||
if content[start:end] != window:
|
||||
found = content.find(
|
||||
window,
|
||||
max(0, start - 32),
|
||||
min(len(content), end + 32 + len(window)),
|
||||
)
|
||||
if found < 0:
|
||||
return None, anchor
|
||||
start = found
|
||||
end = found + len(window)
|
||||
return _source_span(content, start, end), (start_token, start)
|
||||
|
||||
|
||||
def _make_chunk(
|
||||
*,
|
||||
content: str,
|
||||
tokens: int,
|
||||
order: int,
|
||||
source_span: dict[str, int] | None,
|
||||
emit_source_span: bool,
|
||||
) -> dict[str, Any]:
|
||||
item: dict[str, Any] = {
|
||||
"tokens": tokens,
|
||||
"content": content.strip(),
|
||||
"chunk_order_index": order,
|
||||
}
|
||||
if emit_source_span and source_span is not None:
|
||||
item["_source_span"] = source_span
|
||||
return item
|
||||
|
||||
|
||||
def chunking_by_token_size(
|
||||
tokenizer: Tokenizer,
|
||||
content: str,
|
||||
split_by_character: str | None = None,
|
||||
split_by_character_only: bool = False,
|
||||
chunk_overlap_token_size: int = 100,
|
||||
chunk_token_size: int = 1200,
|
||||
*,
|
||||
_emit_source_span: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Legacy 6-arg fixed-token chunker (default for ``LightRAG.chunking_func``).
|
||||
|
||||
Signature is preserved for backward compatibility with externally
|
||||
supplied ``chunking_func`` implementations. New file-based chunking
|
||||
dispatch uses :func:`chunking_by_fixed_token` instead.
|
||||
"""
|
||||
tokens = tokenizer.encode(content)
|
||||
results: list[dict[str, Any]] = []
|
||||
if split_by_character:
|
||||
raw_chunks = content.split(split_by_character)
|
||||
raw_spans: list[tuple[int, int]] = []
|
||||
cursor = 0
|
||||
for raw_chunk in raw_chunks:
|
||||
start = cursor
|
||||
end = start + len(raw_chunk)
|
||||
raw_spans.append((start, end))
|
||||
cursor = end + len(split_by_character)
|
||||
new_chunks = []
|
||||
if split_by_character_only:
|
||||
for chunk, (chunk_start, chunk_end) in zip(raw_chunks, raw_spans):
|
||||
_tokens = tokenizer.encode(chunk)
|
||||
if len(_tokens) > chunk_token_size:
|
||||
logger.warning(
|
||||
"Chunk split_by_character exceeds token limit: len=%d limit=%d",
|
||||
len(_tokens),
|
||||
chunk_token_size,
|
||||
)
|
||||
raise ChunkTokenLimitExceededError(
|
||||
chunk_tokens=len(_tokens),
|
||||
chunk_token_limit=chunk_token_size,
|
||||
chunk_preview=chunk[:120],
|
||||
)
|
||||
span = (
|
||||
_source_span(content, chunk_start, chunk_end)
|
||||
if _emit_source_span
|
||||
else None
|
||||
)
|
||||
new_chunks.append((len(_tokens), chunk, span))
|
||||
else:
|
||||
for chunk, (chunk_start, chunk_end) in zip(raw_chunks, raw_spans):
|
||||
_tokens = tokenizer.encode(chunk)
|
||||
if len(_tokens) > chunk_token_size:
|
||||
# Anchor is chunk-relative (offsets are shifted by chunk_start
|
||||
# below), so it resets per split-by-character segment.
|
||||
anchor = (0, 0)
|
||||
for start in range(
|
||||
0, len(_tokens), chunk_token_size - chunk_overlap_token_size
|
||||
):
|
||||
end_token = min(start + chunk_token_size, len(_tokens))
|
||||
chunk_content = tokenizer.decode(_tokens[start:end_token])
|
||||
span = None
|
||||
if _emit_source_span:
|
||||
span, anchor = _token_window_source_span(
|
||||
tokenizer,
|
||||
chunk,
|
||||
_tokens,
|
||||
start,
|
||||
end_token,
|
||||
anchor=anchor,
|
||||
)
|
||||
if span is not None:
|
||||
span = {
|
||||
"start": chunk_start + span["start"],
|
||||
"end": chunk_start + span["end"],
|
||||
}
|
||||
new_chunks.append(
|
||||
(
|
||||
min(chunk_token_size, len(_tokens) - start),
|
||||
chunk_content,
|
||||
span,
|
||||
)
|
||||
)
|
||||
else:
|
||||
span = (
|
||||
_source_span(content, chunk_start, chunk_end)
|
||||
if _emit_source_span
|
||||
else None
|
||||
)
|
||||
new_chunks.append((len(_tokens), chunk, span))
|
||||
for index, (_len, chunk, span) in enumerate(new_chunks):
|
||||
results.append(
|
||||
_make_chunk(
|
||||
content=chunk,
|
||||
tokens=_len,
|
||||
order=index,
|
||||
source_span=span,
|
||||
emit_source_span=_emit_source_span,
|
||||
)
|
||||
)
|
||||
else:
|
||||
anchor = (0, 0)
|
||||
for index, start in enumerate(
|
||||
range(0, len(tokens), chunk_token_size - chunk_overlap_token_size)
|
||||
):
|
||||
end = min(start + chunk_token_size, len(tokens))
|
||||
chunk_content = tokenizer.decode(tokens[start:end])
|
||||
span = None
|
||||
if _emit_source_span:
|
||||
span, anchor = _token_window_source_span(
|
||||
tokenizer, content, tokens, start, end, anchor=anchor
|
||||
)
|
||||
results.append(
|
||||
_make_chunk(
|
||||
content=chunk_content,
|
||||
tokens=min(chunk_token_size, len(tokens) - start),
|
||||
order=index,
|
||||
source_span=span,
|
||||
emit_source_span=_emit_source_span,
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def chunking_by_fixed_token(
|
||||
tokenizer: Tokenizer,
|
||||
content: str,
|
||||
chunk_token_size: int = 1200,
|
||||
*,
|
||||
chunk_overlap_token_size: int = 100,
|
||||
split_by_character: str | None = None,
|
||||
split_by_character_only: bool = False,
|
||||
_emit_source_span: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Fixed-token chunker — file-chunker contract for the ``"F"`` strategy.
|
||||
|
||||
Implements the same fixed-window algorithm as
|
||||
:func:`chunking_by_token_size`, exposed under the standard
|
||||
file-chunker signature ``(tokenizer, content, chunk_token_size, *,
|
||||
<strategy kwargs>)`` so the file-based chunking dispatcher in
|
||||
``process_single_document`` can call every strategy uniformly.
|
||||
"""
|
||||
return chunking_by_token_size(
|
||||
tokenizer,
|
||||
content,
|
||||
split_by_character=split_by_character,
|
||||
split_by_character_only=split_by_character_only,
|
||||
chunk_overlap_token_size=chunk_overlap_token_size,
|
||||
chunk_token_size=chunk_token_size,
|
||||
_emit_source_span=_emit_source_span,
|
||||
)
|
||||
@@ -0,0 +1,337 @@
|
||||
"""
|
||||
Centralized configuration constants for LightRAG.
|
||||
|
||||
This module defines default values for configuration constants used across
|
||||
different parts of the LightRAG system. Centralizing these values ensures
|
||||
consistency and makes maintenance easier.
|
||||
"""
|
||||
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
# Default values for server settings
|
||||
DEFAULT_WOKERS = 2
|
||||
DEFAULT_MAX_GRAPH_NODES = 1000
|
||||
|
||||
# Default values for extraction settings
|
||||
DEFAULT_SUMMARY_LANGUAGE = "English" # Default language for document processing
|
||||
DEFAULT_MAX_GLEANING = 1
|
||||
DEFAULT_ENTITY_NAME_MAX_LENGTH = 256
|
||||
# Max UTF-8 byte length for entity identifiers. Milvus enforces VARCHAR
|
||||
# max_length in BYTES (not characters), so a CJK name within the character
|
||||
# limit can still exceed the field limit. MUST stay <= the max_length of the
|
||||
# entity_name / src_id / tgt_id fields in lightrag/kg/milvus_impl.py.
|
||||
DEFAULT_ENTITY_NAME_MAX_BYTES = 512
|
||||
|
||||
# Per-response output limits for entity extraction prompts
|
||||
DEFAULT_MAX_EXTRACTION_RECORDS = 100
|
||||
DEFAULT_MAX_EXTRACTION_ENTITIES = 40
|
||||
|
||||
# Number of description fragments to trigger LLM summary
|
||||
DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE = 8
|
||||
# Max description token size to trigger LLM summary
|
||||
DEFAULT_SUMMARY_MAX_TOKENS = 1200
|
||||
# Recommended LLM summary output length in tokens
|
||||
DEFAULT_SUMMARY_LENGTH_RECOMMENDED = 600
|
||||
# Maximum token size sent to LLM for summary
|
||||
DEFAULT_SUMMARY_CONTEXT_SIZE = 12000
|
||||
# Maximum token size allowed for entity extraction input context
|
||||
DEFAULT_MAX_EXTRACT_INPUT_TOKENS = 20480
|
||||
# Maximum token size for the per-chunk `---Section Context---` heading
|
||||
# breadcrumb injected into the extraction prompt. Keeps section metadata from
|
||||
# pushing an otherwise-valid chunk past the provider context window; over budget
|
||||
# the breadcrumb collapses to ``first → … → leaf`` (top-level + nearest section).
|
||||
DEFAULT_MAX_SECTION_CONTEXT_TOKENS = 256
|
||||
# Per-level character cap for each heading in that breadcrumb. Must stay below
|
||||
# 1/3 of DEFAULT_MAX_SECTION_CONTEXT_TOKENS so the collapsed two-level form
|
||||
# (first + leaf, plus separator/ellipsis) always fits within the token budget.
|
||||
DEFAULT_HEADING_LEVEL_MAX_CHARS = 80
|
||||
# Separator for: description, source_id and relation-key fields(Can not be changed after data inserted)
|
||||
GRAPH_FIELD_SEP = "<SEP>"
|
||||
|
||||
# Query and retrieval configuration defaults
|
||||
DEFAULT_TOP_K = 40
|
||||
DEFAULT_CHUNK_TOP_K = 20
|
||||
DEFAULT_MAX_ENTITY_TOKENS = 6000
|
||||
DEFAULT_MAX_RELATION_TOKENS = 8000
|
||||
DEFAULT_MAX_TOTAL_TOKENS = 30000
|
||||
DEFAULT_COSINE_THRESHOLD = 0.2
|
||||
DEFAULT_RELATED_CHUNK_NUMBER = 5
|
||||
DEFAULT_KG_CHUNK_PICK_METHOD = "VECTOR"
|
||||
|
||||
# Rerank configuration defaults
|
||||
DEFAULT_MIN_RERANK_SCORE = 0.0
|
||||
DEFAULT_RERANK_BINDING = "null"
|
||||
|
||||
# Default source ids limit in meta data for entity and relation
|
||||
DEFAULT_MAX_SOURCE_IDS_PER_ENTITY = 200
|
||||
DEFAULT_MAX_SOURCE_IDS_PER_RELATION = 200
|
||||
### control chunk_ids limitation method: FIFO, FIFO
|
||||
### FIFO: First in first out
|
||||
### KEEP: Keep oldest (less merge action and faster)
|
||||
SOURCE_IDS_LIMIT_METHOD_KEEP = "KEEP"
|
||||
SOURCE_IDS_LIMIT_METHOD_FIFO = "FIFO"
|
||||
DEFAULT_SOURCE_IDS_LIMIT_METHOD = SOURCE_IDS_LIMIT_METHOD_KEEP
|
||||
VALID_SOURCE_IDS_LIMIT_METHODS = {
|
||||
SOURCE_IDS_LIMIT_METHOD_KEEP,
|
||||
SOURCE_IDS_LIMIT_METHOD_FIFO,
|
||||
}
|
||||
# Maximum number of file paths stored in entity/relation file_path field (For displayed only, does not affect query performance)
|
||||
DEFAULT_MAX_FILE_PATHS = 75
|
||||
|
||||
# Field length of file_path in Milvus Schema for entity and relation (Should not be changed)
|
||||
# file_path must store all file paths up to the DEFAULT_MAX_FILE_PATHS limit within the metadata.
|
||||
DEFAULT_MAX_FILE_PATH_LENGTH = 32768
|
||||
# Placeholder for more file paths in meta data for entity and relation (Should not be changed)
|
||||
DEFAULT_FILE_PATH_MORE_PLACEHOLDER = "truncated"
|
||||
|
||||
# Default temperature for LLM
|
||||
DEFAULT_TEMPERATURE = 1.0
|
||||
|
||||
# Async configuration defaults
|
||||
DEFAULT_MAX_ASYNC = 4 # Default maximum async operations
|
||||
DEFAULT_MAX_PARALLEL_INSERT = 3 # Default maximum parallel insert operations
|
||||
|
||||
# Chunker defaults — i18n-aware so Chinese / mixed-language documents
|
||||
# split correctly out of the box. Override per deployment via
|
||||
# CHUNK_R_SEPARATORS / CHUNK_V_SENTENCE_SPLIT_REGEX env vars.
|
||||
#
|
||||
# DEFAULT_R_SEPARATORS: cascade tried by langchain RecursiveCharacterTextSplitter.
|
||||
# Order matters — strongest boundary first: paragraph (\n\n) > line (\n) >
|
||||
# Chinese sentence-end (。！？) > Chinese semi-clause (；，) > space > char.
|
||||
# English sentence-ending punctuation (.?!) is intentionally NOT included
|
||||
# because RecursiveCharacterTextSplitter does literal-string splitting, so
|
||||
# "." would also split numerals (``0.95``) and abbreviations (``e.g.``).
|
||||
# The English path falls through space / char as before.
|
||||
DEFAULT_R_SEPARATORS: tuple[str, ...] = (
|
||||
"\n\n",
|
||||
"\n",
|
||||
"。",
|
||||
"！",
|
||||
"？",
|
||||
"；",
|
||||
"，",
|
||||
" ",
|
||||
"",
|
||||
)
|
||||
# DEFAULT_SENTENCE_SPLIT_REGEX: pattern fed to langchain SemanticChunker.
|
||||
# Two alternates so the English branch keeps its ``\s+`` requirement
|
||||
# (avoiding ``0.95`` mid-token splits) while the Chinese branch matches
|
||||
# bare ``。？！`` (CJK has no inter-sentence whitespace).
|
||||
DEFAULT_SENTENCE_SPLIT_REGEX = r"(?<=[.?!])\s+|(?<=[。？！])"
|
||||
|
||||
# DEFAULT_CHUNK_P_SIZE: paragraph-semantic chunker target size when
|
||||
# CHUNK_P_SIZE env is unset. Deliberately larger than the global
|
||||
# CHUNK_SIZE default — heading-aligned paragraph merging needs more
|
||||
# headroom to keep semantically related paragraphs together; falling
|
||||
# back to CHUNK_SIZE (1200) would force premature splits and defeat
|
||||
# the strategy's purpose.
|
||||
DEFAULT_CHUNK_P_SIZE = 2000
|
||||
|
||||
# Paragraph-semantic "drop references" detection defaults (the chunking="P"
|
||||
# drop_references option). DEFAULT_P_REFERENCES_TAIL_N: a reference block is
|
||||
# only dropped when it sits within the last N content blocks of the document
|
||||
# (a safety window so a mid-document "References" subsection is not removed).
|
||||
# DEFAULT_P_REFERENCES_HEADINGS: heading prefixes that mark a reference
|
||||
# section — English words matched case-insensitively at a word boundary,
|
||||
# the Chinese "参考文献" matched as a plain prefix. Both are tunable via env
|
||||
# (CHUNK_P_REFERENCES_TAIL_N / CHUNK_P_REFERENCES_HEADINGS, the latter
|
||||
# pipe-separated) read live by the chunker at run time.
|
||||
DEFAULT_P_REFERENCES_TAIL_N = 2
|
||||
DEFAULT_P_REFERENCES_HEADINGS = ("References", "Bibliography", "参考文献")
|
||||
|
||||
# LightRAG Document pipeline
|
||||
FULL_DOCS_FORMAT_RAW = "raw" # content in full_docs["content"]
|
||||
# Post-parse persistence marker: full_docs rows written by the parsers carry
|
||||
# this parse_format; on resume/retry they route to ReuseParser. Not a valid
|
||||
# enqueue docs_format (the 'lightrag' ingestion entrypoint was removed).
|
||||
FULL_DOCS_FORMAT_LIGHTRAG = "lightrag" # content in LightRAG Document files
|
||||
FULL_DOCS_FORMAT_PENDING_PARSE = (
|
||||
"pending_parse" # file saved but not yet parsed; parse_native will read from disk
|
||||
)
|
||||
# Marker prefix for full_docs.content when format=lightrag.
|
||||
# Per docs/FileProcessingConfiguration-zh.md, the content is "{{LRdoc}}" + a
|
||||
# leading summary of the parsed document so paginated APIs can show a real
|
||||
# preview without loading the full LightRAG Document file.
|
||||
LIGHTRAG_DOC_CONTENT_PREFIX = "{{LRdoc}}"
|
||||
# Engine identifier strings (registry keys). The set of user-selectable
|
||||
# engines and their suffix capabilities now live in
|
||||
# lightrag.parser.registry (ParserSpec table) — the single source of truth.
|
||||
PARSER_ENGINE_LEGACY = "legacy"
|
||||
PARSER_ENGINE_NATIVE = "native"
|
||||
PARSER_ENGINE_MINERU = "mineru"
|
||||
PARSER_ENGINE_DOCLING = "docling"
|
||||
PARSED_DIR_NAME = "__parsed__" # Dir for parsed files (renamed from __enqueued__)
|
||||
# Prefix marking a doc_status content_summary as GENERATED from a file
|
||||
# extraction error (enqueue-time error documents and parse-stage FAILED
|
||||
# upserts). Doubles as the match sentinel that lets a later failure replace
|
||||
# a stale generated summary while real raw-document summaries are preserved —
|
||||
# keep every producer on this constant so the match never drifts.
|
||||
FILE_EXTRACTION_SUMMARY_PREFIX = "[File Extraction]"
|
||||
|
||||
# Suffixes for parser artifact subdirectories under ``<input>/__parsed__/``.
|
||||
# Centralising them here keeps the sidecar writer, engine cache modules and
|
||||
# the delete-path whitelist in sync — new engines should add their raw-dir
|
||||
# suffix to ``PARSED_ARTIFACT_DIR_SUFFIXES`` so deletion picks them up
|
||||
# automatically.
|
||||
PARSED_DIR_SUFFIX = ".parsed" # spec sidecar layout (every engine)
|
||||
MINERU_RAW_DIR_SUFFIX = ".mineru_raw" # preserved MinerU raw bundle
|
||||
DOCLING_RAW_DIR_SUFFIX = ".docling_raw" # preserved Docling raw bundle
|
||||
NATIVE_RAW_DIR_SUFFIX = ".native_raw" # native md downloaded-image cache bundle
|
||||
PARSED_ARTIFACT_DIR_SUFFIXES: tuple[str, ...] = (
|
||||
PARSED_DIR_SUFFIX,
|
||||
MINERU_RAW_DIR_SUFFIX,
|
||||
DOCLING_RAW_DIR_SUFFIX,
|
||||
NATIVE_RAW_DIR_SUFFIX,
|
||||
)
|
||||
|
||||
# Per-file processing options carried by filename hints / LIGHTRAG_PARSER rules.
|
||||
# See docs/FileProcessingConfiguration-zh.md for the full specification.
|
||||
PROCESS_OPTION_IMAGES = "i" # Enable VLM analysis for drawings/images
|
||||
PROCESS_OPTION_TABLES = "t" # Enable VLM analysis for tables
|
||||
PROCESS_OPTION_EQUATIONS = "e" # Enable VLM analysis for equations
|
||||
PROCESS_OPTION_SKIP_KG = "!" # Skip entity/relation extraction (no KG build)
|
||||
ProcessChunkingOption: TypeAlias = Literal["F", "R", "V", "P"]
|
||||
PROCESS_OPTION_CHUNK_FIXED: ProcessChunkingOption = (
|
||||
"F" # Fixed-length / separator chunking (default)
|
||||
)
|
||||
PROCESS_OPTION_CHUNK_RECURSIVE: ProcessChunkingOption = (
|
||||
"R" # Recursive semantic chunking
|
||||
)
|
||||
PROCESS_OPTION_CHUNK_VECTOR: ProcessChunkingOption = (
|
||||
"V" # Vector-driven semantic chunking
|
||||
)
|
||||
PROCESS_OPTION_CHUNK_PARAGRAH: ProcessChunkingOption = (
|
||||
"P" # Paragrah-driven semantic chunking
|
||||
)
|
||||
|
||||
PROCESS_OPTION_CHUNK_CHARS: frozenset[ProcessChunkingOption] = frozenset(
|
||||
{
|
||||
PROCESS_OPTION_CHUNK_FIXED,
|
||||
PROCESS_OPTION_CHUNK_RECURSIVE,
|
||||
PROCESS_OPTION_CHUNK_VECTOR,
|
||||
PROCESS_OPTION_CHUNK_PARAGRAH,
|
||||
}
|
||||
)
|
||||
SUPPORTED_PROCESS_OPTIONS = frozenset(
|
||||
{
|
||||
PROCESS_OPTION_IMAGES,
|
||||
PROCESS_OPTION_TABLES,
|
||||
PROCESS_OPTION_EQUATIONS,
|
||||
PROCESS_OPTION_SKIP_KG,
|
||||
PROCESS_OPTION_CHUNK_FIXED,
|
||||
PROCESS_OPTION_CHUNK_RECURSIVE,
|
||||
PROCESS_OPTION_CHUNK_VECTOR,
|
||||
PROCESS_OPTION_CHUNK_PARAGRAH,
|
||||
}
|
||||
)
|
||||
|
||||
DEFAULT_MAX_PARALLEL_ANALYZE = 5 # Multimodal analysis (VLM) concurrency
|
||||
|
||||
# Per-engine parsing concurrency defaults. mineru / docling are
|
||||
# resource-intensive (GPU/CPU + memory), so they default to a modest amount of
|
||||
# parallelism (2); lower to 1 when resources are tight, or raise via the
|
||||
# MAX_PARALLEL_PARSE_* env vars when you have spare capacity.
|
||||
DEFAULT_MAX_PARALLEL_PARSE_NATIVE = 5
|
||||
DEFAULT_MAX_PARALLEL_PARSE_MINERU = 2
|
||||
DEFAULT_MAX_PARALLEL_PARSE_DOCLING = 2
|
||||
|
||||
# Staged pipeline queue size defaults.
|
||||
DEFAULT_QUEUE_SIZE_PARSE = 20
|
||||
DEFAULT_QUEUE_SIZE_ANALYZE = 100
|
||||
DEFAULT_QUEUE_SIZE_INSERT = 4
|
||||
|
||||
# LLM / embedding call priority levels. Lower values run first
|
||||
# (asyncio.PriorityQueue semantics); priority only orders calls *within* a
|
||||
# single role queue (extract / keyword / query / vlm). These name the values
|
||||
# passed as the ``_priority`` argument to the priority_limit_async_func_call
|
||||
# wrapper, centralizing the magic numbers that were previously inlined at each
|
||||
# call site.
|
||||
#
|
||||
# Query stage (interactive: query/keyword LLM calls and query-time embeddings)
|
||||
# gets the highest priority so user requests stay responsive.
|
||||
DEFAULT_QUERY_PRIORITY = 5
|
||||
# Entity/relation description summary generation — ahead of raw extraction but
|
||||
# behind interactive query work.
|
||||
DEFAULT_SUMMARY_PRIORITY = 8
|
||||
# Processing stage entity/relation extraction (ingestion). Also the wrapper's
|
||||
# baseline default for any call that does not pass ``_priority``.
|
||||
DEFAULT_PROCESSING_PRIORITY = 10
|
||||
# Priority used for all multimodal analysis LLM calls. Set equal to
|
||||
# DEFAULT_PROCESSING_PRIORITY so analysis and ingestion work share the EXTRACT
|
||||
# queue fairly and advance evenly — otherwise a busy ingestion queue starves
|
||||
# analysis tasks, stalling analysis nodes and dragging down overall throughput.
|
||||
DEFAULT_MM_ANALYSIS_PRIORITY = DEFAULT_PROCESSING_PRIORITY
|
||||
|
||||
# Multimodal analysis / chunk thresholds
|
||||
# Minimum token count retained when truncating a multimodal chunk's
|
||||
# description to fit within DEFAULT_MAX_EXTRACT_INPUT_TOKENS. Falling below
|
||||
# this floor leaves the description too thin to ground a useful entity
|
||||
# description, so the pipeline raises instead of producing a stub.
|
||||
DEFAULT_MM_CHUNK_DESCRIPTION_MIN_TOKENS = 100
|
||||
# Minimum image side (width or height) in pixels accepted for VLM analysis.
|
||||
# Anything smaller is treated as decorative (icons, separators, etc.) and
|
||||
# written as status="skipped".
|
||||
DEFAULT_MM_IMAGE_MIN_PIXEL = 64
|
||||
|
||||
# Embedding configuration defaults
|
||||
DEFAULT_EMBEDDING_FUNC_MAX_ASYNC = 8 # Default max async for embedding functions
|
||||
DEFAULT_EMBEDDING_BATCH_NUM = 10 # Default batch size for embedding computations
|
||||
|
||||
# Gunicorn worker timeout
|
||||
DEFAULT_TIMEOUT = 300
|
||||
|
||||
# Default llm and embedding timeout
|
||||
DEFAULT_LLM_TIMEOUT = 240
|
||||
DEFAULT_EMBEDDING_TIMEOUT = 30
|
||||
|
||||
# Rerank async / timeout defaults
|
||||
# Concurrency falls back to base MAX_ASYNC_LLM when env unset; timeout has its own
|
||||
# default since reranker calls are typically much faster than full LLM generation.
|
||||
DEFAULT_RERANK_MAX_ASYNC = DEFAULT_MAX_ASYNC
|
||||
DEFAULT_RERANK_TIMEOUT = 30
|
||||
|
||||
# Cross-worker global concurrency gate (gunicorn multi-worker) defaults.
|
||||
# A lease whose heartbeat is older than the TTL marks its owner as suspect;
|
||||
# a suspect lease is reclaimed only after the additional grace elapses while
|
||||
# the owner PID is still alive (dead PIDs are reclaimed immediately).
|
||||
DEFAULT_GLOBAL_SLOT_HEARTBEAT_TTL = 20.0 # ~4x the 5s health-check heartbeat
|
||||
DEFAULT_GLOBAL_SLOT_SUSPECT_GRACE = 20.0 # ~1x heartbeat TTL
|
||||
# Polling backoff bounds while a worker waits for a free global slot.
|
||||
# The first acquisition attempt is always immediate (backoff applies only
|
||||
# after a failure). The longest-waiting live process keeps polling at the
|
||||
# MIN interval so it usually claims the next freed slot (soft FIFO across
|
||||
# workers); other waiters back off exponentially up to the DEFERRED cap.
|
||||
# The cap stays small on purpose: when the favored waiter leaves, the
|
||||
# promoted one is asleep at most one deferred period, bounding slot idling.
|
||||
DEFAULT_GLOBAL_SLOT_POLL_MIN = 0.05
|
||||
DEFAULT_GLOBAL_SLOT_POLL_DEFERRED_MAX = 0.4
|
||||
# Waiter records not refreshed within this TTL are ignored for the
|
||||
# longest-waiter ranking and reaped: a crashed or stalled poller must not
|
||||
# keep occupying the favored seat (which would push every live waiter onto
|
||||
# the deferred backoff and waste slots). Keep > 2x the deferred poll cap.
|
||||
DEFAULT_GLOBAL_SLOT_WAITER_STALE_TTL = 1.0
|
||||
# Max consecutive zombie (cancelled) queue entries a worker drains while
|
||||
# holding a global slot before returning the slot to other processes.
|
||||
DEFAULT_GLOBAL_SLOT_DRAIN_LIMIT = 16
|
||||
# Physical queue compaction (global-limit mode only): triggered when the
|
||||
# estimated zombie count exceeds the threshold; each maintenance pass
|
||||
# processes at most the batch limit to keep the event loop responsive.
|
||||
DEFAULT_ZOMBIE_COMPACT_THRESHOLD = 64
|
||||
DEFAULT_COMPACT_BATCH_LIMIT = 512
|
||||
# Cross-worker queue stats: snapshots older than the stale TTL (and entries
|
||||
# owned by dead PIDs) are reaped during aggregation; publishes triggered by
|
||||
# counter updates are debounced to the min interval.
|
||||
DEFAULT_QUEUE_STATS_STALE_TTL = 15.0
|
||||
DEFAULT_QUEUE_STATS_MIN_PUBLISH_INTERVAL = 0.1
|
||||
|
||||
# Logging configuration defaults
|
||||
DEFAULT_LOG_MAX_BYTES = 10485760 # Default 10MB
|
||||
DEFAULT_LOG_BACKUP_COUNT = 5 # Default 5 backups
|
||||
DEFAULT_LOG_FILENAME = "lightrag.log" # Default log filename
|
||||
|
||||
# Ollama server configuration defaults
|
||||
DEFAULT_OLLAMA_MODEL_NAME = "lightrag"
|
||||
DEFAULT_OLLAMA_MODEL_TAG = "latest"
|
||||
DEFAULT_OLLAMA_MODEL_SIZE = 7365960935
|
||||
DEFAULT_OLLAMA_CREATED_AT = "2024-01-15T00:00:00Z"
|
||||
DEFAULT_OLLAMA_DIGEST = "sha256:lightrag"
|
||||
@@ -0,0 +1,496 @@
|
||||
# 📊 RAGAS-based Evaluation Framework
|
||||
|
||||
## What is RAGAS?
|
||||
|
||||
**RAGAS** (Retrieval Augmented Generation Assessment) is a framework for reference-free evaluation of RAG systems using LLMs. RAGAS uses state-of-the-art evaluation metrics:
|
||||
|
||||
### Core Metrics
|
||||
|
||||
| Metric | What It Measures | Good Score |
|
||||
| ------ | ---------------- | ---------- |
|
||||
| **Faithfulness** | Is the answer factually accurate based on retrieved context? | > 0.80 |
|
||||
| **Answer Relevance** | Is the answer relevant to the user's question? | > 0.80 |
|
||||
| **Context Recall** | Was all relevant information retrieved from documents? | > 0.80 |
|
||||
| **Context Precision** | Is retrieved context clean without irrelevant noise? | > 0.80 |
|
||||
| **RAGAS Score** | Overall quality metric (average of above) | > 0.80 |
|
||||
|
||||
### 📁 LightRAG Evalua'tion Framework Directory Structure
|
||||
|
||||
```
|
||||
lightrag/evaluation/
|
||||
├── eval_rag_quality.py # Main evaluation script
|
||||
├── sample_dataset.json # 3 test questions about LightRAG
|
||||
├── sample_documents/ # Matching markdown files for testing
|
||||
│ ├── 01_lightrag_overview.md
|
||||
│ ├── 02_rag_architecture.md
|
||||
│ ├── 03_lightrag_improvements.md
|
||||
│ ├── 04_supported_databases.md
|
||||
│ ├── 05_evaluation_and_deployment.md
|
||||
│ └── README.md
|
||||
├── __init__.py # Package init
|
||||
├── results/ # Output directory
|
||||
│ ├── results_YYYYMMDD_HHMMSS.json # Raw metrics in JSON
|
||||
│ └── results_YYYYMMDD_HHMMSS.csv # Metrics in CSV format
|
||||
└── README.md # This file
|
||||
```
|
||||
|
||||
**Quick Test:** Index files from `sample_documents/` into LightRAG, then run the evaluator to reproduce results (~89-100% RAGAS score per question).
|
||||
|
||||
|
||||
|
||||
## 🚀 Quick Start
|
||||
|
||||
### 1. Install Dependencies
|
||||
|
||||
```bash
|
||||
pip install ragas datasets langfuse
|
||||
```
|
||||
|
||||
Or use your project dependencies (already included in pyproject.toml):
|
||||
|
||||
```bash
|
||||
pip install -e ".[evaluation]"
|
||||
```
|
||||
|
||||
### 2. Run Evaluation
|
||||
|
||||
**Optional offline sample retrieval check (no API/model calls):**
|
||||
```bash
|
||||
python lightrag/evaluation/offline_retrieval_check.py --strict
|
||||
```
|
||||
|
||||
This checks whether the bundled sample questions can lexically retrieve their
|
||||
expected sample documents before running LightRAG, embeddings, LLM calls, or
|
||||
RAGAS.
|
||||
|
||||
**Basic usage (uses defaults):**
|
||||
```bash
|
||||
cd /path/to/LightRAG
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
|
||||
**Specify custom dataset:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py --dataset my_test.json
|
||||
```
|
||||
|
||||
**Specify custom RAG endpoint:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py --ragendpoint http://my-server.com:9621
|
||||
```
|
||||
|
||||
**Specify both (short form):**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py -d my_test.json -r http://localhost:9621
|
||||
```
|
||||
|
||||
**Get help:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py --help
|
||||
```
|
||||
|
||||
### 3. View Results
|
||||
|
||||
Results are saved automatically in `lightrag/evaluation/results/`:
|
||||
|
||||
```
|
||||
results/
|
||||
├── results_20241023_143022.json ← Raw metrics in JSON format
|
||||
└── results_20241023_143022.csv ← Metrics in CSV format (for spreadsheets)
|
||||
```
|
||||
|
||||
**Results include:**
|
||||
- ✅ Overall RAGAS score
|
||||
- 📊 Per-metric averages (Faithfulness, Answer Relevance, Context Recall, Context Precision)
|
||||
- 📋 Individual test case results
|
||||
- 📈 Performance breakdown by question
|
||||
|
||||
|
||||
|
||||
## 📋 Command-Line Arguments
|
||||
|
||||
The evaluation script supports command-line arguments for easy configuration:
|
||||
|
||||
| Argument | Short | Default | Description |
|
||||
| -------- | ----- | ------- | ----------- |
|
||||
| `--dataset` | `-d` | `sample_dataset.json` | Path to test dataset JSON file |
|
||||
| `--ragendpoint` | `-r` | `http://localhost:9621` or `$LIGHTRAG_API_URL` | LightRAG API endpoint URL |
|
||||
|
||||
### Usage Examples
|
||||
|
||||
**Use default dataset and endpoint:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
|
||||
**Custom dataset with default endpoint:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py --dataset path/to/my_dataset.json
|
||||
```
|
||||
|
||||
**Default dataset with custom endpoint:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py --ragendpoint http://my-server.com:9621
|
||||
```
|
||||
|
||||
**Custom dataset and endpoint:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py -d my_dataset.json -r http://localhost:9621
|
||||
```
|
||||
|
||||
**Absolute path to dataset:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py -d /path/to/custom_dataset.json
|
||||
```
|
||||
|
||||
**Show help message:**
|
||||
```bash
|
||||
python lightrag/evaluation/eval_rag_quality.py --help
|
||||
```
|
||||
|
||||
|
||||
|
||||
## ⚙️ Configuration
|
||||
|
||||
### Environment Variables
|
||||
|
||||
The evaluation framework supports customization through environment variables:
|
||||
|
||||
**⚠️ IMPORTANT: Both LLM and Embedding endpoints MUST be OpenAI-compatible**
|
||||
- The RAGAS framework requires OpenAI-compatible API interfaces
|
||||
- Custom endpoints must implement the OpenAI API format (e.g., vLLM, SGLang, LocalAI)
|
||||
- Non-compatible endpoints will cause evaluation failures
|
||||
|
||||
| Variable | Default | Description |
|
||||
| -------- | ------- | ----------- |
|
||||
| **LLM Configuration** | | |
|
||||
| `EVAL_LLM_MODEL` | `gpt-4o-mini` | LLM model used for RAGAS evaluation |
|
||||
| `EVAL_LLM_BINDING_API_KEY` | falls back to `OPENAI_API_KEY` | API key for LLM evaluation |
|
||||
| `EVAL_LLM_BINDING_HOST` | (optional) | Custom OpenAI-compatible endpoint URL for LLM |
|
||||
| **Embedding Configuration** | | |
|
||||
| `EVAL_EMBEDDING_MODEL` | `text-embedding-3-large` | Embedding model for evaluation |
|
||||
| `EVAL_EMBEDDING_BINDING_API_KEY` | falls back to `EVAL_LLM_BINDING_API_KEY` → `OPENAI_API_KEY` | API key for embeddings |
|
||||
| `EVAL_EMBEDDING_BINDING_HOST` | falls back to `EVAL_LLM_BINDING_HOST` | Custom OpenAI-compatible endpoint URL for embeddings |
|
||||
| **Performance Tuning** | | |
|
||||
| `EVAL_MAX_CONCURRENT` | 2 | Number of concurrent test case evaluations (1=serial) |
|
||||
| `EVAL_QUERY_TOP_K` | 10 | Number of documents to retrieve per query |
|
||||
| `EVAL_LLM_MAX_RETRIES` | 5 | Maximum LLM request retries |
|
||||
| `EVAL_LLM_TIMEOUT` | 180 | LLM request timeout in seconds |
|
||||
|
||||
### Usage Examples
|
||||
|
||||
**Example 1: Default Configuration (OpenAI Official API)**
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
Both LLM and embeddings use OpenAI's official API with default models.
|
||||
|
||||
**Example 2: Custom Models on OpenAI**
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export EVAL_LLM_MODEL=gpt-4o-mini
|
||||
export EVAL_EMBEDDING_MODEL=text-embedding-3-large
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
|
||||
**Example 3: Same Custom OpenAI-Compatible Endpoint for Both**
|
||||
```bash
|
||||
# Both LLM and embeddings use the same custom endpoint
|
||||
export EVAL_LLM_BINDING_API_KEY=your-custom-key
|
||||
export EVAL_LLM_BINDING_HOST=http://localhost:8000/v1
|
||||
export EVAL_LLM_MODEL=qwen-plus
|
||||
export EVAL_EMBEDDING_MODEL=BAAI/bge-m3
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
Embeddings automatically inherit LLM endpoint configuration.
|
||||
|
||||
**Example 4: Separate Endpoints (Cost Optimization)**
|
||||
```bash
|
||||
# Use OpenAI for LLM (high quality)
|
||||
export EVAL_LLM_BINDING_API_KEY=sk-openai-key
|
||||
export EVAL_LLM_MODEL=gpt-4o-mini
|
||||
# No EVAL_LLM_BINDING_HOST means use OpenAI official API
|
||||
|
||||
# Use local vLLM for embeddings (cost-effective)
|
||||
export EVAL_EMBEDDING_BINDING_API_KEY=local-key
|
||||
export EVAL_EMBEDDING_BINDING_HOST=http://localhost:8001/v1
|
||||
export EVAL_EMBEDDING_MODEL=BAAI/bge-m3
|
||||
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
LLM uses OpenAI official API, embeddings use local custom endpoint.
|
||||
|
||||
**Example 5: Different Custom Endpoints for LLM and Embeddings**
|
||||
```bash
|
||||
# LLM on one OpenAI-compatible server
|
||||
export EVAL_LLM_BINDING_API_KEY=key1
|
||||
export EVAL_LLM_BINDING_HOST=http://llm-server:8000/v1
|
||||
export EVAL_LLM_MODEL=custom-llm
|
||||
|
||||
# Embeddings on another OpenAI-compatible server
|
||||
export EVAL_EMBEDDING_BINDING_API_KEY=key2
|
||||
export EVAL_EMBEDDING_BINDING_HOST=http://embedding-server:8001/v1
|
||||
export EVAL_EMBEDDING_MODEL=custom-embedding
|
||||
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
Both use different custom OpenAI-compatible endpoints.
|
||||
|
||||
**Example 6: Using Environment Variables from .env File**
|
||||
```bash
|
||||
# Create .env file in project root
|
||||
cat > .env << EOF
|
||||
EVAL_LLM_BINDING_API_KEY=your-key
|
||||
EVAL_LLM_BINDING_HOST=http://localhost:8000/v1
|
||||
EVAL_LLM_MODEL=qwen-plus
|
||||
EVAL_EMBEDDING_MODEL=BAAI/bge-m3
|
||||
EOF
|
||||
|
||||
# Run evaluation (automatically loads .env)
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
|
||||
### Concurrency Control & Rate Limiting
|
||||
|
||||
The evaluation framework includes built-in concurrency control to prevent API rate limiting issues:
|
||||
|
||||
**Why Concurrency Control Matters:**
|
||||
- RAGAS internally makes many concurrent LLM calls for each test case
|
||||
- Context Precision metric calls LLM once per retrieved document
|
||||
- Without control, this can easily exceed API rate limits
|
||||
|
||||
**Default Configuration (Conservative):**
|
||||
```bash
|
||||
EVAL_MAX_CONCURRENT=2 # Serial evaluation (one test at a time)
|
||||
EVAL_QUERY_TOP_K=10 # OP_K query parameter of LightRAG
|
||||
EVAL_LLM_MAX_RETRIES=5 # Retry failed requests 5 times
|
||||
EVAL_LLM_TIMEOUT=180 # 3-minute timeout per request
|
||||
```
|
||||
|
||||
**Common Issues and Solutions:**
|
||||
|
||||
| Issue | Solution |
|
||||
| ----- | -------- |
|
||||
| **Warning: "LM returned 1 generations instead of 3"** | Reduce `EVAL_MAX_CONCURRENT` to 1 or decrease `EVAL_QUERY_TOP_K` |
|
||||
| **Context Precision returns NaN** | Lower `EVAL_QUERY_TOP_K` to reduce LLM calls per test case |
|
||||
| **Rate limit errors (429)** | Increase `EVAL_LLM_MAX_RETRIES` and decrease `EVAL_MAX_CONCURRENT` |
|
||||
| **Request timeouts** | Increase `EVAL_LLM_TIMEOUT` to 180 or higher |
|
||||
|
||||
|
||||
|
||||
## 📝 Test Dataset
|
||||
|
||||
`sample_dataset.json` contains 3 generic questions about LightRAG. Replace with questions matching YOUR indexed documents.
|
||||
|
||||
**Custom Test Cases:**
|
||||
|
||||
```json
|
||||
{
|
||||
"test_cases": [
|
||||
{
|
||||
"question": "Your question here",
|
||||
"ground_truth": "Expected answer from your data",
|
||||
"project": "evaluation_project_name"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📊 Interpreting Results
|
||||
|
||||
### Score Ranges
|
||||
|
||||
- **0.80-1.00**: ✅ Excellent (Production-ready)
|
||||
- **0.60-0.80**: ⚠️ Good (Room for improvement)
|
||||
- **0.40-0.60**: ❌ Poor (Needs optimization)
|
||||
- **0.00-0.40**: 🔴 Critical (Major issues)
|
||||
|
||||
### What Low Scores Mean
|
||||
|
||||
| Metric | Low Score Indicates |
|
||||
| ------ | ------------------- |
|
||||
| **Faithfulness** | Responses contain hallucinations or incorrect information |
|
||||
| **Answer Relevance** | Answers don't match what users asked |
|
||||
| **Context Recall** | Missing important information in retrieval |
|
||||
| **Context Precision** | Retrieved documents contain irrelevant noise |
|
||||
|
||||
### Optimization Tips
|
||||
|
||||
1. **Low Faithfulness**:
|
||||
- Improve entity extraction quality
|
||||
- Better document chunking
|
||||
- Tune retrieval temperature
|
||||
|
||||
2. **Low Answer Relevance**:
|
||||
- Improve prompt engineering
|
||||
- Better query understanding
|
||||
- Check semantic similarity threshold
|
||||
|
||||
3. **Low Context Recall**:
|
||||
- Increase retrieval `top_k` results
|
||||
- Improve embedding model
|
||||
- Better document preprocessing
|
||||
|
||||
4. **Low Context Precision**:
|
||||
- Smaller, focused chunks
|
||||
- Better filtering
|
||||
- Improve chunking strategy
|
||||
|
||||
---
|
||||
|
||||
## 📚 Resources
|
||||
|
||||
- [RAGAS Documentation](https://docs.ragas.io/)
|
||||
- [RAGAS GitHub](https://github.com/explodinggradients/ragas)
|
||||
|
||||
---
|
||||
|
||||
## 🐛 Troubleshooting
|
||||
|
||||
### "ModuleNotFoundError: No module named 'ragas'"
|
||||
|
||||
```bash
|
||||
pip install ragas datasets
|
||||
```
|
||||
|
||||
### "Warning: LM returned 1 generations instead of requested 3" or Context Precision NaN
|
||||
|
||||
**Cause**: This warning indicates API rate limiting or concurrent request overload:
|
||||
- RAGAS makes multiple LLM calls per test case (faithfulness, relevancy, recall, precision)
|
||||
- Context Precision calls LLM once per retrieved document (with `EVAL_QUERY_TOP_K=10`, that's 10 calls)
|
||||
- Concurrent evaluation multiplies these calls: `EVAL_MAX_CONCURRENT × LLM calls per test`
|
||||
|
||||
**Solutions** (in order of effectiveness):
|
||||
|
||||
1. **Serial Evaluation** (Default):
|
||||
```bash
|
||||
export EVAL_MAX_CONCURRENT=1
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
|
||||
2. **Reduce Retrieved Documents**:
|
||||
```bash
|
||||
export EVAL_QUERY_TOP_K=5 # Halves Context Precision LLM calls
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
|
||||
3. **Increase Retry & Timeout**:
|
||||
```bash
|
||||
export EVAL_LLM_MAX_RETRIES=10
|
||||
export EVAL_LLM_TIMEOUT=180
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
|
||||
4. **Use Higher Quota API** (if available):
|
||||
- Upgrade to OpenAI Tier 2+ for higher RPM limits
|
||||
- Use self-hosted OpenAI-compatible service with no rate limits
|
||||
|
||||
### "AttributeError: 'InstructorLLM' object has no attribute 'agenerate_prompt'" or NaN results
|
||||
|
||||
This error occurs with RAGAS 0.3.x when LLM and Embeddings are not explicitly configured. The evaluation framework now handles this automatically by:
|
||||
- Using environment variables to configure evaluation models
|
||||
- Creating proper LLM and Embeddings instances for RAGAS
|
||||
|
||||
**Solution**: Ensure you have set one of the following:
|
||||
- `OPENAI_API_KEY` environment variable (default)
|
||||
- `EVAL_LLM_BINDING_API_KEY` for custom API key
|
||||
|
||||
The framework will automatically configure the evaluation models.
|
||||
|
||||
### "No sample_dataset.json found"
|
||||
|
||||
Make sure you're running from the project root:
|
||||
|
||||
```bash
|
||||
cd /path/to/LightRAG
|
||||
python lightrag/evaluation/eval_rag_quality.py
|
||||
```
|
||||
|
||||
### "LightRAG query API errors during evaluation"
|
||||
|
||||
The evaluation uses your configured LLM (OpenAI by default). Ensure:
|
||||
- API keys are set in `.env`
|
||||
- Network connection is stable
|
||||
|
||||
### Evaluation requires running LightRAG API
|
||||
|
||||
The evaluator queries a running LightRAG API server at `http://localhost:9621`. Make sure:
|
||||
1. LightRAG API server is running (`python lightrag/api/lightrag_server.py`)
|
||||
2. Documents are indexed in your LightRAG instance
|
||||
3. API is accessible at the configured URL
|
||||
|
||||
|
||||
|
||||
## 📝 Next Steps
|
||||
|
||||
1. Start LightRAG API server
|
||||
2. Upload sample documents into LightRAG throught WebUI
|
||||
3. Run `python lightrag/evaluation/eval_rag_quality.py`
|
||||
4. Review results (JSON/CSV) in `results/` folder
|
||||
|
||||
Evaluation Result Sample:
|
||||
|
||||
```
|
||||
INFO: ======================================================================
|
||||
INFO: 🔍 RAGAS Evaluation - Using Real LightRAG API
|
||||
INFO: ======================================================================
|
||||
INFO: Evaluation Models:
|
||||
INFO: • LLM Model: gpt-4.1
|
||||
INFO: • Embedding Model: text-embedding-3-large
|
||||
INFO: • Endpoint: OpenAI Official API
|
||||
INFO: Concurrency & Rate Limiting:
|
||||
INFO: • Query Top-K: 10 Entities/Relations
|
||||
INFO: • LLM Max Retries: 5
|
||||
INFO: • LLM Timeout: 180 seconds
|
||||
INFO: Test Configuration:
|
||||
INFO: • Total Test Cases: 6
|
||||
INFO: • Test Dataset: sample_dataset.json
|
||||
INFO: • LightRAG API: http://localhost:9621
|
||||
INFO: • Results Directory: results
|
||||
INFO: ======================================================================
|
||||
INFO: 🚀 Starting RAGAS Evaluation of LightRAG System
|
||||
INFO: 🔧 RAGAS Evaluation (Stage 2): 2 concurrent
|
||||
INFO: ======================================================================
|
||||
INFO:
|
||||
INFO: ===================================================================================================================
|
||||
INFO: 📊 EVALUATION RESULTS SUMMARY
|
||||
INFO: ===================================================================================================================
|
||||
INFO: # | Question | Faith | AnswRel | CtxRec | CtxPrec | RAGAS | Status
|
||||
INFO: -------------------------------------------------------------------------------------------------------------------
|
||||
INFO: 1 | How does LightRAG solve the hallucination probl... | 1.0000 | 1.0000 | 1.0000 | 1.0000 | 1.0000 | ✓
|
||||
INFO: 2 | What are the three main components required in ... | 0.8500 | 0.5790 | 1.0000 | 1.0000 | 0.8573 | ✓
|
||||
INFO: 3 | How does LightRAG's retrieval performance compa... | 0.8056 | 1.0000 | 1.0000 | 1.0000 | 0.9514 | ✓
|
||||
INFO: 4 | What vector databases does LightRAG support and... | 0.8182 | 0.9807 | 1.0000 | 1.0000 | 0.9497 | ✓
|
||||
INFO: 5 | What are the four key metrics for evaluating RA... | 1.0000 | 0.7452 | 1.0000 | 1.0000 | 0.9363 | ✓
|
||||
INFO: 6 | What are the core benefits of LightRAG and how ... | 0.9583 | 0.8829 | 1.0000 | 1.0000 | 0.9603 | ✓
|
||||
INFO: ===================================================================================================================
|
||||
INFO:
|
||||
INFO: ======================================================================
|
||||
INFO: 📊 EVALUATION COMPLETE
|
||||
INFO: ======================================================================
|
||||
INFO: Total Tests: 6
|
||||
INFO: Successful: 6
|
||||
INFO: Failed: 0
|
||||
INFO: Success Rate: 100.00%
|
||||
INFO: Elapsed Time: 161.10 seconds
|
||||
INFO: Avg Time/Test: 26.85 seconds
|
||||
INFO:
|
||||
INFO: ======================================================================
|
||||
INFO: 📈 BENCHMARK RESULTS (Average)
|
||||
INFO: ======================================================================
|
||||
INFO: Average Faithfulness: 0.9053
|
||||
INFO: Average Answer Relevance: 0.8646
|
||||
INFO: Average Context Recall: 1.0000
|
||||
INFO: Average Context Precision: 1.0000
|
||||
INFO: Average RAGAS Score: 0.9425
|
||||
INFO: ----------------------------------------------------------------------
|
||||
INFO: Min RAGAS Score: 0.8573
|
||||
INFO: Max RAGAS Score: 1.0000
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
**Happy Evaluating! 🚀**
|
||||
@@ -0,0 +1,25 @@
|
||||
"""
|
||||
LightRAG Evaluation Module
|
||||
|
||||
RAGAS-based evaluation framework for assessing RAG system quality.
|
||||
|
||||
Usage:
|
||||
from lightrag.evaluation import RAGEvaluator
|
||||
|
||||
evaluator = RAGEvaluator()
|
||||
results = await evaluator.run()
|
||||
|
||||
Note: RAGEvaluator is imported lazily to avoid import errors
|
||||
when ragas/datasets are not installed.
|
||||
"""
|
||||
|
||||
__all__ = ["RAGEvaluator"]
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
"""Lazy import to avoid dependency errors when ragas is not installed."""
|
||||
if name == "RAGEvaluator":
|
||||
from .eval_rag_quality import RAGEvaluator
|
||||
|
||||
return RAGEvaluator
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,254 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Offline sanity check for the bundled LightRAG evaluation samples.
|
||||
|
||||
The check uses a small deterministic lexical ranker. It does not start
|
||||
LightRAG, call the API server, compute embeddings, or call LLM/RAGAS services.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
import sys
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
EVAL_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_DATASET = EVAL_DIR / "sample_dataset.json"
|
||||
DEFAULT_DOCS_DIR = EVAL_DIR / "sample_documents"
|
||||
DEFAULT_ORACLE = EVAL_DIR / "sample_retrieval_oracle.json"
|
||||
|
||||
STOPWORDS = {
|
||||
"a",
|
||||
"an",
|
||||
"and",
|
||||
"are",
|
||||
"as",
|
||||
"at",
|
||||
"be",
|
||||
"by",
|
||||
"for",
|
||||
"from",
|
||||
"how",
|
||||
"in",
|
||||
"into",
|
||||
"is",
|
||||
"it",
|
||||
"its",
|
||||
"of",
|
||||
"on",
|
||||
"or",
|
||||
"that",
|
||||
"the",
|
||||
"their",
|
||||
"to",
|
||||
"what",
|
||||
"with",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
name: str
|
||||
tokens: Counter[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryResult:
|
||||
question: str
|
||||
expected: list[str]
|
||||
ranked: list[str]
|
||||
|
||||
def recall_at(self, top_k: int) -> float:
|
||||
hits = set(self.expected) & set(self.ranked[:top_k])
|
||||
return len(hits) / len(self.expected)
|
||||
|
||||
def reciprocal_rank(self) -> float:
|
||||
for rank, document in enumerate(self.ranked, start=1):
|
||||
if document in self.expected:
|
||||
return 1 / rank
|
||||
return 0.0
|
||||
|
||||
|
||||
def tokenize(text: str) -> list[str]:
|
||||
tokens = re.findall(r"[a-z0-9]+", text.lower())
|
||||
return [token for token in tokens if token not in STOPWORDS and len(token) > 1]
|
||||
|
||||
|
||||
def load_cases(dataset_path: Path) -> list[dict[str, Any]]:
|
||||
payload = json.loads(dataset_path.read_text(encoding="utf-8"))
|
||||
cases = payload.get("test_cases")
|
||||
if not isinstance(cases, list):
|
||||
raise ValueError(f"{dataset_path} must contain a test_cases list")
|
||||
return cases
|
||||
|
||||
|
||||
def load_oracle(oracle_path: Path) -> dict[str, list[str]]:
|
||||
payload = json.loads(oracle_path.read_text(encoding="utf-8"))
|
||||
entries = payload.get("oracle")
|
||||
if not isinstance(entries, list):
|
||||
raise ValueError(f"{oracle_path} must contain an oracle list")
|
||||
|
||||
oracle: dict[str, list[str]] = {}
|
||||
for entry in entries:
|
||||
question = str(entry.get("question", "")).strip()
|
||||
expected = entry.get("expected_documents")
|
||||
if not question or not isinstance(expected, list) or not expected:
|
||||
raise ValueError("Each oracle entry needs question and expected_documents")
|
||||
oracle[question] = [str(document) for document in expected]
|
||||
return oracle
|
||||
|
||||
|
||||
def load_documents(docs_dir: Path) -> list[Document]:
|
||||
documents: list[Document] = []
|
||||
for path in sorted(docs_dir.glob("*.md")):
|
||||
if path.name.lower() == "readme.md":
|
||||
continue
|
||||
documents.append(
|
||||
Document(
|
||||
name=path.name,
|
||||
tokens=Counter(tokenize(path.read_text(encoding="utf-8"))),
|
||||
)
|
||||
)
|
||||
if not documents:
|
||||
raise ValueError(f"No markdown sample documents found in {docs_dir}")
|
||||
return documents
|
||||
|
||||
|
||||
def inverse_document_frequency(documents: list[Document]) -> dict[str, float]:
|
||||
document_frequency: Counter[str] = Counter()
|
||||
for document in documents:
|
||||
document_frequency.update(document.tokens.keys())
|
||||
doc_count = len(documents)
|
||||
return {
|
||||
token: math.log((doc_count + 1) / (frequency + 1)) + 1
|
||||
for token, frequency in document_frequency.items()
|
||||
}
|
||||
|
||||
|
||||
def score_query(
|
||||
query_tokens: list[str],
|
||||
document: Document,
|
||||
idf: dict[str, float],
|
||||
) -> float:
|
||||
score = 0.0
|
||||
for token in query_tokens:
|
||||
if token in document.tokens:
|
||||
score += (1 + math.log(document.tokens[token])) * idf.get(token, 0.0)
|
||||
return score
|
||||
|
||||
|
||||
def audit_samples(
|
||||
cases: list[dict[str, Any]],
|
||||
oracle: dict[str, list[str]],
|
||||
documents: list[Document],
|
||||
) -> list[QueryResult]:
|
||||
idf = inverse_document_frequency(documents)
|
||||
results: list[QueryResult] = []
|
||||
|
||||
for case in cases:
|
||||
question = str(case.get("question", "")).strip()
|
||||
if question not in oracle:
|
||||
raise ValueError(f"No oracle entry for question: {question}")
|
||||
|
||||
query_tokens = tokenize(question)
|
||||
scored_documents = [
|
||||
(score_query(query_tokens, document, idf), document)
|
||||
for document in documents
|
||||
]
|
||||
ranked = [
|
||||
document
|
||||
for score, document in sorted(
|
||||
scored_documents,
|
||||
key=lambda item: (-item[0], item[1].name),
|
||||
)
|
||||
if score > 0
|
||||
]
|
||||
results.append(
|
||||
QueryResult(
|
||||
question=question,
|
||||
expected=oracle[question],
|
||||
ranked=[document.name for document in ranked],
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def summarize(results: list[QueryResult], top_k: int) -> dict[str, Any]:
|
||||
if not results:
|
||||
raise ValueError("No query results to summarize")
|
||||
recalls = [result.recall_at(top_k) for result in results]
|
||||
reciprocal_ranks = [result.reciprocal_rank() for result in results]
|
||||
return {
|
||||
"queries": len(results),
|
||||
"top_k": top_k,
|
||||
"average_recall_at_k": sum(recalls) / len(recalls),
|
||||
"mean_reciprocal_rank": sum(reciprocal_ranks) / len(reciprocal_ranks),
|
||||
"full_recall_queries": sum(recall == 1.0 for recall in recalls),
|
||||
"no_hit_queries": sum(recall == 0.0 for recall in recalls),
|
||||
}
|
||||
|
||||
|
||||
def print_report(results: list[QueryResult], top_k: int) -> None:
|
||||
summary = summarize(results, top_k)
|
||||
print("LightRAG sample retrieval check")
|
||||
print(f"Queries: {summary['queries']}")
|
||||
print(f"Top-k: {summary['top_k']}")
|
||||
print(f"Average recall@k: {summary['average_recall_at_k']:.3f}")
|
||||
print(f"Mean reciprocal rank: {summary['mean_reciprocal_rank']:.3f}")
|
||||
print(f"Full-recall queries: {summary['full_recall_queries']}/{summary['queries']}")
|
||||
print(f"No-hit queries: {summary['no_hit_queries']}")
|
||||
print()
|
||||
for index, result in enumerate(results, start=1):
|
||||
top_docs = ", ".join(result.ranked[:top_k])
|
||||
expected = ", ".join(result.expected)
|
||||
print(f"{index}. recall@{top_k}={result.recall_at(top_k):.3f}")
|
||||
print(f" expected: {expected}")
|
||||
print(f" top docs: {top_docs}")
|
||||
|
||||
|
||||
def parse_args(argv: list[str]) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Run an offline retrieval check for LightRAG evaluation samples."
|
||||
)
|
||||
parser.add_argument("--dataset", default=str(DEFAULT_DATASET))
|
||||
parser.add_argument("--docs-dir", default=str(DEFAULT_DOCS_DIR))
|
||||
parser.add_argument("--oracle", default=str(DEFAULT_ORACLE))
|
||||
parser.add_argument("--top-k", type=int, default=2)
|
||||
parser.add_argument(
|
||||
"--strict",
|
||||
action="store_true",
|
||||
help="Exit non-zero unless every sample query has full recall@k.",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = parse_args(argv or sys.argv[1:])
|
||||
if args.top_k <= 0:
|
||||
print("--top-k must be positive", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
try:
|
||||
cases = load_cases(Path(args.dataset).expanduser())
|
||||
oracle = load_oracle(Path(args.oracle).expanduser())
|
||||
documents = load_documents(Path(args.docs_dir).expanduser())
|
||||
results = audit_samples(cases, oracle, documents)
|
||||
print_report(results, args.top_k)
|
||||
summary = summarize(results, args.top_k)
|
||||
except (OSError, ValueError, json.JSONDecodeError) as exc:
|
||||
print(f"Sample retrieval check failed: {exc}", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
if args.strict and summary["full_recall_queries"] != summary["queries"]:
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,34 @@
|
||||
{
|
||||
"test_cases": [
|
||||
{
|
||||
"question": "How does LightRAG solve the hallucination problem in large language models?",
|
||||
"ground_truth": "LightRAG solves the hallucination problem by combining large language models with external knowledge retrieval. The framework ensures accurate responses by grounding LLM outputs in actual documents. LightRAG provides contextual responses that reduce hallucinations significantly.",
|
||||
"project": "lightrag_evaluation_sample"
|
||||
},
|
||||
{
|
||||
"question": "What are the three main components required in a RAG system?",
|
||||
"ground_truth": "A RAG system requires three main components: a retrieval system (vector database or search engine) to find relevant documents, an embedding model to convert text into vector representations for similarity search, and a large language model (LLM) to generate responses based on retrieved context.",
|
||||
"project": "lightrag_evaluation_sample"
|
||||
},
|
||||
{
|
||||
"question": "How does LightRAG's retrieval performance compare to traditional RAG approaches?",
|
||||
"ground_truth": "LightRAG delivers faster retrieval performance than traditional RAG approaches. The framework optimizes document retrieval operations for speed. Traditional RAG systems often suffer from slow query response times. LightRAG achieves high quality results with improved performance. The framework combines speed with accuracy in retrieval operations, prioritizing ease of use without sacrificing quality.",
|
||||
"project": "lightrag_evaluation_sample"
|
||||
},
|
||||
{
|
||||
"question": "What vector databases does LightRAG support and what are their key characteristics?",
|
||||
"ground_truth": "LightRAG supports multiple vector databases including ChromaDB for simple deployment and efficient similarity search, Neo4j for graph-based knowledge representation with vector capabilities, Milvus for high-performance vector search at scale, Qdrant for fast similarity search with filtering and production-ready infrastructure, MongoDB Atlas for combined document storage and vector search, Redis for in-memory low-latency vector search, and a built-in nano-vectordb that eliminates external dependencies for small projects. This multi-database support enables developers to choose appropriate backends based on scale, performance, and infrastructure requirements.",
|
||||
"project": "lightrag_evaluation_sample"
|
||||
},
|
||||
{
|
||||
"question": "What are the four key metrics for evaluating RAG system quality and what does each metric measure?",
|
||||
"ground_truth": "RAG system quality is measured through four key metrics: Faithfulness measures whether answers are factually grounded in retrieved context and detects hallucinations. Answer Relevance measures how well answers address the user question and evaluates response appropriateness. Context Recall measures completeness of retrieval and whether all relevant information was retrieved from documents. Context Precision measures quality and relevance of retrieved documents without noise or irrelevant content.",
|
||||
"project": "lightrag_evaluation_sample"
|
||||
},
|
||||
{
|
||||
"question": "What are the core benefits of LightRAG and how does it improve upon traditional RAG systems?",
|
||||
"ground_truth": "LightRAG offers five core benefits: accuracy through document-grounded responses, up-to-date information without model retraining, domain expertise through specialized document collections, cost-effectiveness by avoiding expensive fine-tuning, and transparency by showing source documents. Compared to traditional RAG systems, LightRAG provides a simpler API with intuitive interfaces, faster retrieval performance with optimized operations, better integration with multiple vector database backends for flexible selection, and optimized prompting strategies with refined templates. LightRAG prioritizes ease of use while maintaining quality and combines speed with accuracy.",
|
||||
"project": "lightrag_evaluation_sample"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
# LightRAG Framework Overview
|
||||
|
||||
## What is LightRAG?
|
||||
|
||||
**LightRAG** is a Simple and Fast Retrieval-Augmented Generation framework. LightRAG was developed by HKUDS (Hong Kong University Data Science Lab). The framework provides developers with tools to build RAG applications efficiently.
|
||||
|
||||
## Problem Statement
|
||||
|
||||
Large language models face several limitations. LLMs have a knowledge cutoff date that prevents them from accessing recent information. Large language models generate hallucinations when providing responses without factual grounding. LLMs lack domain-specific expertise in specialized fields.
|
||||
|
||||
## How LightRAG Solves These Problems
|
||||
|
||||
LightRAG solves the hallucination problem by combining large language models with external knowledge retrieval. The framework ensures accurate responses by grounding LLM outputs in actual documents. LightRAG provides contextual responses that reduce hallucinations significantly. The system enables efficient retrieval from external knowledge bases to supplement LLM capabilities.
|
||||
|
||||
## Core Benefits
|
||||
|
||||
LightRAG offers accuracy through document-grounded responses. The framework provides up-to-date information without model retraining. LightRAG enables domain expertise through specialized document collections. The system delivers cost-effectiveness by avoiding expensive model fine-tuning. LightRAG ensures transparency by showing source documents for each response.
|
||||
@@ -0,0 +1,21 @@
|
||||
# RAG System Architecture
|
||||
|
||||
## Main Components of RAG Systems
|
||||
|
||||
A RAG system consists of three main components that work together to provide intelligent responses.
|
||||
|
||||
### Component 1: Retrieval System
|
||||
|
||||
The retrieval system is the first component of a RAG system. A retrieval system finds relevant documents from large document collections. Vector databases serve as the primary storage for the retrieval system. Search engines can also function as retrieval systems in RAG architectures.
|
||||
|
||||
### Component 2: Embedding Model
|
||||
|
||||
The embedding model is the second component of a RAG system. An embedding model converts text into vector representations for similarity search. The embedding model transforms documents and queries into numerical vectors. These vector representations enable semantic similarity matching between queries and documents.
|
||||
|
||||
### Component 3: Large Language Model
|
||||
|
||||
The large language model is the third component of a RAG system. An LLM generates responses based on retrieved context from documents. The large language model synthesizes information from multiple sources into coherent answers. LLMs provide natural language generation capabilities for the RAG system.
|
||||
|
||||
## How Components Work Together
|
||||
|
||||
The retrieval system fetches relevant documents for a user query. The embedding model enables similarity matching between query and documents. The LLM generates the final response using retrieved context. These three components collaborate to provide accurate, contextual responses.
|
||||
@@ -0,0 +1,25 @@
|
||||
# LightRAG Improvements Over Traditional RAG
|
||||
|
||||
## Key Improvements
|
||||
|
||||
LightRAG improves upon traditional RAG approaches in several significant ways.
|
||||
|
||||
### Simpler API Design
|
||||
|
||||
LightRAG offers a simpler API compared to traditional RAG frameworks. The framework provides intuitive interfaces for developers. Traditional RAG systems often require complex configuration and setup. LightRAG focuses on ease of use while maintaining functionality.
|
||||
|
||||
### Faster Retrieval Performance
|
||||
|
||||
LightRAG delivers faster retrieval performance than traditional RAG approaches. The framework optimizes document retrieval operations for speed. Traditional RAG systems often suffer from slow query response times. LightRAG achieves high quality results with improved performance.
|
||||
|
||||
### Better Vector Database Integration
|
||||
|
||||
LightRAG provides better integration with various vector databases. The framework supports multiple vector database backends seamlessly. Traditional RAG approaches typically lock developers into specific database choices. LightRAG enables flexible storage backend selection.
|
||||
|
||||
### Optimized Prompting Strategies
|
||||
|
||||
LightRAG implements optimized prompting strategies for better results. The framework uses refined prompt templates for accurate responses. Traditional RAG systems often use generic prompting approaches. LightRAG balances simplicity with high quality output.
|
||||
|
||||
## Design Philosophy
|
||||
|
||||
LightRAG prioritizes ease of use without sacrificing quality. The framework combines speed with accuracy in retrieval operations. LightRAG maintains flexibility in database and model selection.
|
||||
@@ -0,0 +1,37 @@
|
||||
# LightRAG Vector Database Support
|
||||
|
||||
## Supported Vector Databases
|
||||
|
||||
LightRAG supports multiple vector databases for flexible deployment options.
|
||||
|
||||
### ChromaDB
|
||||
|
||||
ChromaDB is a vector database supported by LightRAG. ChromaDB provides simple deployment for development environments. The database offers efficient vector similarity search capabilities.
|
||||
|
||||
### Neo4j
|
||||
|
||||
Neo4j is a graph database supported by LightRAG. Neo4j enables graph-based knowledge representation alongside vector search. The database combines relationship modeling with vector capabilities.
|
||||
|
||||
### Milvus
|
||||
|
||||
Milvus is a vector database supported by LightRAG. Milvus provides high-performance vector search at scale. The database handles large-scale vector collections efficiently.
|
||||
|
||||
### Qdrant
|
||||
|
||||
Qdrant is a vector database supported by LightRAG. Qdrant offers fast similarity search with filtering capabilities. The database provides production-ready vector search infrastructure.
|
||||
|
||||
### MongoDB Atlas Vector Search
|
||||
|
||||
MongoDB Atlas Vector Search is supported by LightRAG. MongoDB Atlas combines document storage with vector search capabilities. The database enables unified data management for RAG applications.
|
||||
|
||||
### Redis
|
||||
|
||||
Redis is supported by LightRAG for vector search operations. Redis provides in-memory vector search with low latency. The database offers fast retrieval for real-time applications.
|
||||
|
||||
### Built-in Nano-VectorDB
|
||||
|
||||
LightRAG includes a built-in nano-vectordb for simple deployments. Nano-vectordb eliminates external database dependencies for small projects. The built-in database provides basic vector search functionality without additional setup.
|
||||
|
||||
## Database Selection Benefits
|
||||
|
||||
The multiple database support enables developers to choose appropriate storage backends. LightRAG adapts to different deployment scenarios from development to production. Users can select databases based on scale, performance, and infrastructure requirements.
|
||||
@@ -0,0 +1,41 @@
|
||||
# RAG Evaluation Metrics and Deployment
|
||||
|
||||
## Key RAG Evaluation Metrics
|
||||
|
||||
RAG system quality is measured through four key metrics.
|
||||
|
||||
### Faithfulness Metric
|
||||
|
||||
Faithfulness measures whether answers are factually grounded in retrieved context. The faithfulness metric detects hallucinations in LLM responses. High faithfulness scores indicate answers based on actual document content. The metric evaluates factual accuracy of generated responses.
|
||||
|
||||
### Answer Relevance Metric
|
||||
|
||||
Answer Relevance measures how well answers address the user question. The answer relevance metric evaluates response quality and appropriateness. High answer relevance scores show responses that directly answer user queries. The metric assesses the connection between questions and generated answers.
|
||||
|
||||
### Context Recall Metric
|
||||
|
||||
Context Recall measures completeness of retrieval from documents. The context recall metric evaluates whether all relevant information was retrieved. High context recall scores indicate comprehensive document retrieval. The metric assesses retrieval system effectiveness.
|
||||
|
||||
### Context Precision Metric
|
||||
|
||||
Context Precision measures quality and relevance of retrieved documents. The context precision metric evaluates retrieval accuracy without noise. High context precision scores show clean retrieval without irrelevant content. The metric measures retrieval system selectivity.
|
||||
|
||||
## LightRAG Deployment Options
|
||||
|
||||
LightRAG can be deployed in production through multiple approaches.
|
||||
|
||||
### Docker Container Deployment
|
||||
|
||||
Docker containers enable consistent LightRAG deployment across environments. Docker provides isolated runtime environments for the framework. Container deployment simplifies dependency management and scaling.
|
||||
|
||||
### REST API Server with FastAPI
|
||||
|
||||
FastAPI serves as the REST API framework for LightRAG deployment. The FastAPI server exposes LightRAG functionality through HTTP endpoints. REST API deployment enables client-server architecture for RAG applications.
|
||||
|
||||
### Direct Python Integration
|
||||
|
||||
Direct Python integration embeds LightRAG into Python applications. Python integration provides programmatic access to RAG capabilities. Direct integration supports custom application workflows and pipelines.
|
||||
|
||||
### Deployment Features
|
||||
|
||||
LightRAG supports environment-based configuration for different deployment scenarios. The framework integrates with multiple LLM providers for flexibility. LightRAG enables horizontal scaling for production workloads.
|
||||
@@ -0,0 +1,21 @@
|
||||
# Sample Documents for Evaluation
|
||||
|
||||
These markdown files correspond to test questions in `../sample_dataset.json`.
|
||||
|
||||
## Usage
|
||||
|
||||
1. **Index documents** into LightRAG (via WebUI, API, or Python)
|
||||
2. **Run evaluation**: `python lightrag/evaluation/eval_rag_quality.py`
|
||||
3. **Expected results**: ~91-100% RAGAS score per question
|
||||
|
||||
## Files
|
||||
|
||||
- `01_lightrag_overview.md` - LightRAG framework and hallucination problem
|
||||
- `02_rag_architecture.md` - RAG system components
|
||||
- `03_lightrag_improvements.md` - LightRAG vs traditional RAG
|
||||
- `04_supported_databases.md` - Vector database support
|
||||
- `05_evaluation_and_deployment.md` - Metrics and deployment
|
||||
|
||||
## Note
|
||||
|
||||
Documents use clear entity-relationship patterns for LightRAG's default entity extraction prompts. For better results with your data, customize `lightrag/prompt.py`.
|
||||
@@ -0,0 +1,31 @@
|
||||
{
|
||||
"oracle": [
|
||||
{
|
||||
"question": "How does LightRAG solve the hallucination problem in large language models?",
|
||||
"expected_documents": ["01_lightrag_overview.md"]
|
||||
},
|
||||
{
|
||||
"question": "What are the three main components required in a RAG system?",
|
||||
"expected_documents": ["02_rag_architecture.md"]
|
||||
},
|
||||
{
|
||||
"question": "How does LightRAG's retrieval performance compare to traditional RAG approaches?",
|
||||
"expected_documents": ["03_lightrag_improvements.md"]
|
||||
},
|
||||
{
|
||||
"question": "What vector databases does LightRAG support and what are their key characteristics?",
|
||||
"expected_documents": ["04_supported_databases.md"]
|
||||
},
|
||||
{
|
||||
"question": "What are the four key metrics for evaluating RAG system quality and what does each metric measure?",
|
||||
"expected_documents": ["05_evaluation_and_deployment.md"]
|
||||
},
|
||||
{
|
||||
"question": "What are the core benefits of LightRAG and how does it improve upon traditional RAG systems?",
|
||||
"expected_documents": [
|
||||
"01_lightrag_overview.md",
|
||||
"03_lightrag_improvements.md"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
from typing import Literal
|
||||
|
||||
|
||||
class APIStatusError(Exception):
|
||||
"""Raised when an API response has a status code of 4xx or 5xx."""
|
||||
|
||||
response: httpx.Response
|
||||
status_code: int
|
||||
request_id: str | None
|
||||
|
||||
def __init__(
|
||||
self, message: str, *, response: httpx.Response, body: object | None
|
||||
) -> None:
|
||||
super().__init__(message)
|
||||
self.request = response.request
|
||||
self.body = body
|
||||
self.response = response
|
||||
self.status_code = response.status_code
|
||||
self.request_id = response.headers.get("x-request-id")
|
||||
|
||||
|
||||
class APIConnectionError(Exception):
|
||||
def __init__(
|
||||
self, *, message: str = "Connection error.", request: httpx.Request | None
|
||||
) -> None:
|
||||
super().__init__(message)
|
||||
self.request = request
|
||||
|
||||
|
||||
class BadRequestError(APIStatusError):
|
||||
status_code: Literal[400] = 400 # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
|
||||
|
||||
class AuthenticationError(APIStatusError):
|
||||
status_code: Literal[401] = 401 # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
|
||||
|
||||
class PermissionDeniedError(APIStatusError):
|
||||
status_code: Literal[403] = 403 # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
|
||||
|
||||
class NotFoundError(APIStatusError):
|
||||
status_code: Literal[404] = 404 # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
|
||||
|
||||
class ConflictError(APIStatusError):
|
||||
status_code: Literal[409] = 409 # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
|
||||
|
||||
class UnprocessableEntityError(APIStatusError):
|
||||
status_code: Literal[422] = 422 # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
|
||||
|
||||
class RateLimitError(APIStatusError):
|
||||
status_code: Literal[429] = 429 # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
|
||||
|
||||
class APITimeoutError(APIConnectionError):
|
||||
def __init__(self, request: httpx.Request | None) -> None:
|
||||
super().__init__(message="Request timed out.", request=request)
|
||||
|
||||
|
||||
class StorageNotInitializedError(RuntimeError):
|
||||
"""Raised when storage operations are attempted before initialization."""
|
||||
|
||||
def __init__(self, storage_type: str = "Storage"):
|
||||
super().__init__(
|
||||
f"{storage_type} not initialized. Please ensure proper initialization:\n"
|
||||
f"\n"
|
||||
f" rag = LightRAG(...)\n"
|
||||
f" await rag.initialize_storages() # Required - auto-initializes pipeline_status\n"
|
||||
f"\n"
|
||||
f"See: https://github.com/HKUDS/LightRAG#important-initialization-requirements"
|
||||
)
|
||||
|
||||
|
||||
class PipelineNotInitializedError(KeyError):
|
||||
"""Raised when pipeline status is accessed before initialization."""
|
||||
|
||||
def __init__(self, namespace: str = ""):
|
||||
msg = (
|
||||
f"Pipeline namespace '{namespace}' not found.\n"
|
||||
f"\n"
|
||||
f"Pipeline status should be auto-initialized by initialize_storages().\n"
|
||||
f"If you see this error, please ensure:\n"
|
||||
f"\n"
|
||||
f" 1. You called await rag.initialize_storages()\n"
|
||||
f" 2. For multi-workspace setups, each LightRAG instance was properly initialized\n"
|
||||
f"\n"
|
||||
f"Standard initialization:\n"
|
||||
f" rag = LightRAG(workspace='your_workspace')\n"
|
||||
f" await rag.initialize_storages() # Auto-initializes pipeline_status\n"
|
||||
f"\n"
|
||||
f"If you need manual control (advanced):\n"
|
||||
f" from lightrag.kg.shared_storage import initialize_pipeline_status\n"
|
||||
f" await initialize_pipeline_status(workspace='your_workspace')"
|
||||
)
|
||||
super().__init__(msg)
|
||||
|
||||
|
||||
class PipelineCancelledException(Exception):
|
||||
"""Raised when pipeline processing is cancelled by user request."""
|
||||
|
||||
def __init__(self, message: str = "User cancelled"):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
|
||||
|
||||
class IndexFlushError(Exception):
|
||||
"""Raised when a storage backend fails to flush buffered index ops.
|
||||
|
||||
Carries the storage driver name and namespace so the pipeline can abort
|
||||
the batch with an actionable reason. The underlying error is preserved as
|
||||
the exception ``__cause__`` (set via ``raise ... from cause``).
|
||||
"""
|
||||
|
||||
def __init__(self, storage_name: str, namespace: str, cause: BaseException):
|
||||
self.storage_name = storage_name
|
||||
self.namespace = namespace
|
||||
super().__init__(f"{storage_name}[{namespace}] index flush failed: {cause}")
|
||||
|
||||
|
||||
class ChunkTokenLimitExceededError(ValueError):
|
||||
"""Raised when a chunk exceeds the configured token limit."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_tokens: int,
|
||||
chunk_token_limit: int,
|
||||
chunk_preview: str | None = None,
|
||||
) -> None:
|
||||
preview = chunk_preview.strip() if chunk_preview else None
|
||||
truncated_preview = preview[:80] if preview else None
|
||||
preview_note = f" Preview: '{truncated_preview}'" if truncated_preview else ""
|
||||
message = (
|
||||
f"Chunk token length {chunk_tokens} exceeds chunk_token_size {chunk_token_limit}."
|
||||
f"{preview_note}"
|
||||
)
|
||||
super().__init__(message)
|
||||
self.chunk_tokens = chunk_tokens
|
||||
self.chunk_token_limit = chunk_token_limit
|
||||
self.chunk_preview = truncated_preview
|
||||
|
||||
|
||||
class ChunkBlockMatchError(ValueError):
|
||||
"""Raised when a chunk cannot be located in the document's blocks.jsonl.
|
||||
|
||||
Sidecar backfill (``lightrag.sidecar.backfill``) maps F/R/V chunks back to
|
||||
their source block(s) by matching chunk content against the parse-time
|
||||
``*.blocks.jsonl`` merged text. When a sidecar-less chunk cannot be located,
|
||||
this is raised so the pipeline marks the document FAILED rather than
|
||||
persisting chunks with missing/incorrect provenance.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_order_index: int,
|
||||
chunk_preview: str | None = None,
|
||||
blocks_path: str | None = None,
|
||||
) -> None:
|
||||
preview = chunk_preview.strip() if chunk_preview else None
|
||||
truncated_preview = preview[:80] if preview else None
|
||||
preview_note = f" Preview: '{truncated_preview}'" if truncated_preview else ""
|
||||
path_note = f" (blocks: {blocks_path})" if blocks_path else ""
|
||||
message = (
|
||||
f"Chunk #{chunk_order_index} could not be located in the document "
|
||||
f"blocks during sidecar backfill.{preview_note}{path_note}"
|
||||
)
|
||||
super().__init__(message)
|
||||
self.chunk_order_index = chunk_order_index
|
||||
self.chunk_preview = truncated_preview
|
||||
self.blocks_path = blocks_path
|
||||
|
||||
|
||||
class DataMigrationError(Exception):
|
||||
"""Raised when data migration from legacy collection/table fails."""
|
||||
|
||||
def __init__(self, message: str):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
|
||||
|
||||
class MultimodalAnalysisError(RuntimeError):
|
||||
"""Raised when multimodal analysis must fail the current document.
|
||||
|
||||
Hard failures (missing required field, schema mismatch, model not
|
||||
available, sidecar already carries ``status="failure"``) bubble this
|
||||
exception so the pipeline marks the document failed instead of writing
|
||||
an unusable analyze result. Callers persist a ``status="failure"``
|
||||
sidecar entry alongside the raise so a re-run sees the failure.
|
||||
"""
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Shared atomic file-write helpers.
|
||||
|
||||
Why this lives at the package root rather than under ``lightrag/kg/``:
|
||||
``lightrag.utils.write_json`` needs ``atomic_write`` to gain crash safety,
|
||||
and several ``lightrag/kg/*`` modules need both. Hosting the helpers under
|
||||
``lightrag/kg/`` would create a ``utils -> kg -> utils`` import cycle.
|
||||
Keeping this module dependency-free (stdlib only) avoids that.
|
||||
|
||||
Semantics
|
||||
---------
|
||||
``atomic_write`` writes through a per-writer ``.tmp.<pid>.<tid>.<ns>`` sibling
|
||||
and renames into place with ``os.replace`` — atomic on the same filesystem on
|
||||
both POSIX (``rename(2)``) and Windows (``MoveFileEx`` with
|
||||
``MOVEFILE_REPLACE_EXISTING``). Two failure modes are handled differently:
|
||||
|
||||
- A Python exception (``write_fn`` raised, ``os.replace`` failed, etc.):
|
||||
``finally`` runs, the in-flight tmp is removed best-effort, and the
|
||||
exception propagates. The on-disk destination is the prior snapshot.
|
||||
- A process-level kill (SIGKILL, OOM, hard reboot) between writing the tmp
|
||||
and the rename: ``finally`` does not run, the tmp survives as an orphan,
|
||||
and ``reap_orphan_tmp_files`` cleans it on the next startup once it ages
|
||||
past the threshold.
|
||||
|
||||
What is *not* preserved across the inode swap done by ``os.replace``: owner,
|
||||
group, ACLs, xattrs, hard-link relationships, and any symlink-target identity.
|
||||
The mode bits (rwx) are preserved explicitly — see ``_preserve_mode``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import glob
|
||||
import logging
|
||||
import os
|
||||
import stat
|
||||
import threading
|
||||
import time
|
||||
from typing import Callable
|
||||
|
||||
logger = logging.getLogger("lightrag")
|
||||
|
||||
# Orphan .tmp files older than this are reaped on startup. Large enough that
|
||||
# an in-flight write from another live process cannot plausibly still be
|
||||
# running (multi-million-node graphml writes finish in minutes, not hours).
|
||||
TMP_REAP_AGE_SECONDS = 3600
|
||||
|
||||
|
||||
def tmp_path_for(file_name: str) -> str:
|
||||
"""Return a per-writer tmp sibling for ``file_name``.
|
||||
|
||||
The suffix embeds PID, thread id, and a nanosecond timestamp so that
|
||||
multiple concurrent writers — separate processes sharing the same working
|
||||
directory, or multiple threads inside one process — cannot trample each
|
||||
other's in-flight tmp and leave a "no such file" rename error behind.
|
||||
"""
|
||||
return f"{file_name}.tmp.{os.getpid()}.{threading.get_ident()}.{time.time_ns()}"
|
||||
|
||||
|
||||
def _preserve_mode(tmp: str, dst: str, workspace: str) -> None:
|
||||
"""Carry ``dst``'s existing mode bits onto ``tmp`` before the rename.
|
||||
|
||||
Without this, ``os.replace`` swaps the inode and the new file inherits
|
||||
umask defaults — any intentional restriction (e.g. chmod 0600) on the
|
||||
prior snapshot would be silently widened.
|
||||
"""
|
||||
if not os.path.exists(dst):
|
||||
return
|
||||
try:
|
||||
os.chmod(tmp, stat.S_IMODE(os.stat(dst).st_mode))
|
||||
except OSError as exc:
|
||||
logger.warning(f"[{workspace}] Could not preserve mode of {dst}: {exc}")
|
||||
|
||||
|
||||
def reap_orphan_tmp_files(
|
||||
file_name: str,
|
||||
workspace: str = "_",
|
||||
age_seconds: int = TMP_REAP_AGE_SECONDS,
|
||||
extra_patterns: tuple[str, ...] = (),
|
||||
) -> None:
|
||||
"""Delete stale tmp siblings of ``file_name`` left behind by hard kills.
|
||||
|
||||
Default pattern matches ``glob.escape(file_name) + ".tmp.*"`` — the suffix
|
||||
shape produced by ``tmp_path_for``. ``extra_patterns`` accepts already-built
|
||||
glob patterns and is intended for migrating away from legacy naming
|
||||
schemes (e.g. Faiss's previous fixed ``<meta>.tmp`` suffix, which the
|
||||
default pattern's trailing ``.*`` will not match).
|
||||
|
||||
``glob.escape`` is required because ``file_name`` is composed from
|
||||
``working_dir + namespace`` and can legitimately contain glob
|
||||
metacharacters (workspace ``[v2]``, ``*``, ``?``). Concatenating naively
|
||||
would silently miss the real orphan or widen the match to tmp files of
|
||||
unrelated storage types.
|
||||
"""
|
||||
patterns = [glob.escape(file_name) + ".tmp.*", *extra_patterns]
|
||||
now = time.time()
|
||||
for pattern in patterns:
|
||||
for path in glob.glob(pattern):
|
||||
try:
|
||||
age = now - os.path.getmtime(path)
|
||||
except OSError:
|
||||
continue
|
||||
if age < age_seconds:
|
||||
continue
|
||||
try:
|
||||
os.remove(path)
|
||||
logger.info(
|
||||
f"[{workspace}] Reaped orphan tmp file: {path} (age {age:.0f}s)"
|
||||
)
|
||||
except OSError as exc:
|
||||
logger.warning(
|
||||
f"[{workspace}] Failed to reap orphan tmp file {path}: {exc}"
|
||||
)
|
||||
|
||||
|
||||
def atomic_write(
|
||||
file_name: str,
|
||||
write_fn: Callable[[str], None],
|
||||
workspace: str = "_",
|
||||
) -> None:
|
||||
"""Run ``write_fn(tmp_path)`` then atomically replace ``file_name`` with it.
|
||||
|
||||
``write_fn`` is responsible for actually producing the file contents at
|
||||
the path it receives. It must not assume the tmp path equals ``file_name``
|
||||
— Faiss/Nano callers rely on the tmp path being a real sibling.
|
||||
|
||||
On any exception from ``write_fn`` or from the rename, the tmp is removed
|
||||
best-effort and the exception propagates. The destination file is not
|
||||
touched in that case.
|
||||
"""
|
||||
tmp = tmp_path_for(file_name)
|
||||
try:
|
||||
write_fn(tmp)
|
||||
_preserve_mode(tmp, file_name, workspace)
|
||||
os.replace(tmp, file_name)
|
||||
except BaseException:
|
||||
try:
|
||||
if os.path.exists(tmp):
|
||||
os.remove(tmp)
|
||||
except OSError as exc:
|
||||
logger.warning(
|
||||
f"[{workspace}] Failed to remove tmp after failed atomic write: {exc}"
|
||||
)
|
||||
raise
|
||||
@@ -0,0 +1,161 @@
|
||||
STORAGE_IMPLEMENTATIONS = {
|
||||
"KV_STORAGE": {
|
||||
"implementations": [
|
||||
"JsonKVStorage",
|
||||
"RedisKVStorage",
|
||||
"PGKVStorage",
|
||||
"MongoKVStorage",
|
||||
"OpenSearchKVStorage",
|
||||
],
|
||||
"required_methods": ["get_by_id", "upsert"],
|
||||
},
|
||||
"GRAPH_STORAGE": {
|
||||
"implementations": [
|
||||
"NetworkXStorage",
|
||||
"Neo4JStorage",
|
||||
"PGGraphStorage",
|
||||
"MongoGraphStorage",
|
||||
"MemgraphStorage",
|
||||
"OpenSearchGraphStorage",
|
||||
],
|
||||
"required_methods": ["upsert_node", "upsert_edge"],
|
||||
},
|
||||
"VECTOR_STORAGE": {
|
||||
"implementations": [
|
||||
"NanoVectorDBStorage",
|
||||
"MilvusVectorDBStorage",
|
||||
"PGVectorStorage",
|
||||
"FaissVectorDBStorage",
|
||||
"QdrantVectorDBStorage",
|
||||
"MongoVectorDBStorage",
|
||||
"OpenSearchVectorDBStorage",
|
||||
# "ChromaVectorDBStorage",
|
||||
],
|
||||
"required_methods": ["query", "upsert"],
|
||||
},
|
||||
"DOC_STATUS_STORAGE": {
|
||||
"implementations": [
|
||||
"JsonDocStatusStorage",
|
||||
"RedisDocStatusStorage",
|
||||
"PGDocStatusStorage",
|
||||
"MongoDocStatusStorage",
|
||||
"OpenSearchDocStatusStorage",
|
||||
],
|
||||
"required_methods": ["get_docs_by_status"],
|
||||
},
|
||||
}
|
||||
|
||||
# Storage implementation environment variable without default value
|
||||
STORAGE_ENV_REQUIREMENTS: dict[str, list[str]] = {
|
||||
# KV Storage Implementations
|
||||
"JsonKVStorage": [],
|
||||
"MongoKVStorage": [
|
||||
"MONGO_URI",
|
||||
"MONGO_DATABASE",
|
||||
],
|
||||
"RedisKVStorage": ["REDIS_URI"],
|
||||
"PGKVStorage": ["POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_DATABASE"],
|
||||
# Graph Storage Implementations
|
||||
"NetworkXStorage": [],
|
||||
"Neo4JStorage": ["NEO4J_URI", "NEO4J_USERNAME", "NEO4J_PASSWORD"],
|
||||
"MongoGraphStorage": [
|
||||
"MONGO_URI",
|
||||
"MONGO_DATABASE",
|
||||
],
|
||||
"MemgraphStorage": ["MEMGRAPH_URI"],
|
||||
"AGEStorage": [
|
||||
"AGE_POSTGRES_DB",
|
||||
"AGE_POSTGRES_USER",
|
||||
"AGE_POSTGRES_PASSWORD",
|
||||
],
|
||||
"PGGraphStorage": [
|
||||
"POSTGRES_USER",
|
||||
"POSTGRES_PASSWORD",
|
||||
"POSTGRES_DATABASE",
|
||||
],
|
||||
# Vector Storage Implementations
|
||||
"NanoVectorDBStorage": [],
|
||||
"MilvusVectorDBStorage": [
|
||||
"MILVUS_URI",
|
||||
"MILVUS_DB_NAME",
|
||||
],
|
||||
# "ChromaVectorDBStorage": [],
|
||||
"PGVectorStorage": ["POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_DATABASE"],
|
||||
"FaissVectorDBStorage": [],
|
||||
"QdrantVectorDBStorage": ["QDRANT_URL"], # QDRANT_API_KEY has default value None
|
||||
"MongoVectorDBStorage": [
|
||||
"MONGO_URI",
|
||||
"MONGO_DATABASE",
|
||||
],
|
||||
# Document Status Storage Implementations
|
||||
"JsonDocStatusStorage": [],
|
||||
"RedisDocStatusStorage": ["REDIS_URI"],
|
||||
"PGDocStatusStorage": ["POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_DATABASE"],
|
||||
"MongoDocStatusStorage": [
|
||||
"MONGO_URI",
|
||||
"MONGO_DATABASE",
|
||||
],
|
||||
# OpenSearch Storage Implementations
|
||||
"OpenSearchKVStorage": [
|
||||
"OPENSEARCH_HOSTS",
|
||||
],
|
||||
"OpenSearchDocStatusStorage": [
|
||||
"OPENSEARCH_HOSTS",
|
||||
],
|
||||
"OpenSearchGraphStorage": [
|
||||
"OPENSEARCH_HOSTS",
|
||||
],
|
||||
"OpenSearchVectorDBStorage": [
|
||||
"OPENSEARCH_HOSTS",
|
||||
],
|
||||
}
|
||||
|
||||
# Storage implementation module mapping
|
||||
STORAGES = {
|
||||
"NetworkXStorage": ".kg.networkx_impl",
|
||||
"JsonKVStorage": ".kg.json_kv_impl",
|
||||
"NanoVectorDBStorage": ".kg.nano_vector_db_impl",
|
||||
"JsonDocStatusStorage": ".kg.json_doc_status_impl",
|
||||
"Neo4JStorage": ".kg.neo4j_impl",
|
||||
"MilvusVectorDBStorage": ".kg.milvus_impl",
|
||||
"MongoKVStorage": ".kg.mongo_impl",
|
||||
"MongoDocStatusStorage": ".kg.mongo_impl",
|
||||
"MongoGraphStorage": ".kg.mongo_impl",
|
||||
"MongoVectorDBStorage": ".kg.mongo_impl",
|
||||
"RedisKVStorage": ".kg.redis_impl",
|
||||
"RedisDocStatusStorage": ".kg.redis_impl",
|
||||
"ChromaVectorDBStorage": ".kg.chroma_impl",
|
||||
"PGKVStorage": ".kg.postgres_impl",
|
||||
"PGVectorStorage": ".kg.postgres_impl",
|
||||
"AGEStorage": ".kg.age_impl",
|
||||
"PGGraphStorage": ".kg.postgres_impl",
|
||||
"PGDocStatusStorage": ".kg.postgres_impl",
|
||||
"FaissVectorDBStorage": ".kg.faiss_impl",
|
||||
"QdrantVectorDBStorage": ".kg.qdrant_impl",
|
||||
"MemgraphStorage": ".kg.memgraph_impl",
|
||||
"OpenSearchKVStorage": ".kg.opensearch_impl",
|
||||
"OpenSearchDocStatusStorage": ".kg.opensearch_impl",
|
||||
"OpenSearchGraphStorage": ".kg.opensearch_impl",
|
||||
"OpenSearchVectorDBStorage": ".kg.opensearch_impl",
|
||||
}
|
||||
|
||||
|
||||
def verify_storage_implementation(storage_type: str, storage_name: str) -> None:
|
||||
"""Verify if storage implementation is compatible with specified storage type
|
||||
|
||||
Args:
|
||||
storage_type: Storage type (KV_STORAGE, GRAPH_STORAGE etc.)
|
||||
storage_name: Storage implementation name
|
||||
|
||||
Raises:
|
||||
ValueError: If storage implementation is incompatible or missing required methods
|
||||
"""
|
||||
if storage_type not in STORAGE_IMPLEMENTATIONS:
|
||||
raise ValueError(f"Unknown storage type: {storage_type}")
|
||||
|
||||
storage_info = STORAGE_IMPLEMENTATIONS[storage_type]
|
||||
if storage_name not in storage_info["implementations"]:
|
||||
raise ValueError(
|
||||
f"Storage implementation '{storage_name}' is not compatible with {storage_type}. "
|
||||
f"Compatible implementations are: {', '.join(storage_info['implementations'])}"
|
||||
)
|
||||
@@ -0,0 +1,344 @@
|
||||
import asyncio
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, final
|
||||
import numpy as np
|
||||
|
||||
from lightrag.base import BaseVectorStorage
|
||||
from lightrag.constants import DEFAULT_QUERY_PRIORITY
|
||||
from lightrag.utils import logger
|
||||
import pipmaster as pm
|
||||
|
||||
if not pm.is_installed("chromadb"):
|
||||
pm.install("chromadb")
|
||||
|
||||
from chromadb import HttpClient, PersistentClient # type: ignore
|
||||
from chromadb.config import Settings # type: ignore
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
class ChromaVectorDBStorage(BaseVectorStorage):
|
||||
"""ChromaDB vector storage implementation."""
|
||||
|
||||
def __post_init__(self):
|
||||
self._validate_embedding_func()
|
||||
try:
|
||||
config = self.global_config.get("vector_db_storage_cls_kwargs", {})
|
||||
cosine_threshold = config.get("cosine_better_than_threshold")
|
||||
if cosine_threshold is None:
|
||||
raise ValueError(
|
||||
"cosine_better_than_threshold must be specified in vector_db_storage_cls_kwargs"
|
||||
)
|
||||
self.cosine_better_than_threshold = cosine_threshold
|
||||
|
||||
user_collection_settings = config.get("collection_settings", {})
|
||||
# Default HNSW index settings for ChromaDB
|
||||
default_collection_settings = {
|
||||
# Distance metric used for similarity search (cosine similarity)
|
||||
"hnsw:space": "cosine",
|
||||
# Number of nearest neighbors to explore during index construction
|
||||
# Higher values = better recall but slower indexing
|
||||
"hnsw:construction_ef": 128,
|
||||
# Number of nearest neighbors to explore during search
|
||||
# Higher values = better recall but slower search
|
||||
"hnsw:search_ef": 128,
|
||||
# Number of connections per node in the HNSW graph
|
||||
# Higher values = better recall but more memory usage
|
||||
"hnsw:M": 16,
|
||||
# Number of vectors to process in one batch during indexing
|
||||
"hnsw:batch_size": 100,
|
||||
# Number of updates before forcing index synchronization
|
||||
# Lower values = more frequent syncs but slower indexing
|
||||
"hnsw:sync_threshold": 1000,
|
||||
}
|
||||
collection_settings = {
|
||||
**default_collection_settings,
|
||||
**user_collection_settings,
|
||||
}
|
||||
|
||||
local_path = config.get("local_path", None)
|
||||
if local_path:
|
||||
self._client = PersistentClient(
|
||||
path=local_path,
|
||||
settings=Settings(
|
||||
allow_reset=True,
|
||||
anonymized_telemetry=False,
|
||||
),
|
||||
)
|
||||
else:
|
||||
auth_provider = config.get(
|
||||
"auth_provider", "chromadb.auth.token_authn.TokenAuthClientProvider"
|
||||
)
|
||||
auth_credentials = config.get("auth_token", "secret-token")
|
||||
headers = {}
|
||||
|
||||
if "token_authn" in auth_provider:
|
||||
headers = {
|
||||
config.get(
|
||||
"auth_header_name", "X-Chroma-Token"
|
||||
): auth_credentials
|
||||
}
|
||||
elif "basic_authn" in auth_provider:
|
||||
auth_credentials = config.get("auth_credentials", "admin:admin")
|
||||
|
||||
self._client = HttpClient(
|
||||
host=config.get("host", "localhost"),
|
||||
port=config.get("port", 8000),
|
||||
headers=headers,
|
||||
settings=Settings(
|
||||
chroma_api_impl="rest",
|
||||
chroma_client_auth_provider=auth_provider,
|
||||
chroma_client_auth_credentials=auth_credentials,
|
||||
allow_reset=True,
|
||||
anonymized_telemetry=False,
|
||||
),
|
||||
)
|
||||
|
||||
self._collection = self._client.get_or_create_collection(
|
||||
name=self.namespace,
|
||||
metadata={
|
||||
**collection_settings,
|
||||
"dimension": self.embedding_func.embedding_dim,
|
||||
},
|
||||
)
|
||||
# Use batch size from collection settings if specified
|
||||
self._max_batch_size = self.global_config.get(
|
||||
"embedding_batch_num", collection_settings.get("hnsw:batch_size", 32)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"ChromaDB initialization failed: {str(e)}")
|
||||
raise
|
||||
|
||||
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
||||
logger.debug(f"Inserting {len(data)} to {self.namespace}")
|
||||
if not data:
|
||||
return
|
||||
|
||||
try:
|
||||
import time
|
||||
|
||||
current_time = int(time.time())
|
||||
|
||||
ids = list(data.keys())
|
||||
documents = [v["content"] for v in data.values()]
|
||||
metadatas = [
|
||||
{
|
||||
**{k: v for k, v in item.items() if k in self.meta_fields},
|
||||
"created_at": current_time,
|
||||
}
|
||||
or {"_default": "true", "created_at": current_time}
|
||||
for item in data.values()
|
||||
]
|
||||
|
||||
# Process in batches
|
||||
batches = [
|
||||
documents[i : i + self._max_batch_size]
|
||||
for i in range(0, len(documents), self._max_batch_size)
|
||||
]
|
||||
|
||||
embedding_tasks = [self.embedding_func(batch) for batch in batches]
|
||||
embeddings_list = []
|
||||
|
||||
# Pre-allocate embeddings_list with known size
|
||||
embeddings_list = [None] * len(embedding_tasks)
|
||||
|
||||
# Use asyncio.gather instead of as_completed if order doesn't matter
|
||||
embeddings_results = await asyncio.gather(*embedding_tasks)
|
||||
embeddings_list = list(embeddings_results)
|
||||
|
||||
embeddings = np.concatenate(embeddings_list)
|
||||
|
||||
# Upsert in batches
|
||||
for i in range(0, len(ids), self._max_batch_size):
|
||||
batch_slice = slice(i, i + self._max_batch_size)
|
||||
|
||||
self._collection.upsert(
|
||||
ids=ids[batch_slice],
|
||||
embeddings=embeddings[batch_slice].tolist(),
|
||||
documents=documents[batch_slice],
|
||||
metadatas=metadatas[batch_slice],
|
||||
)
|
||||
|
||||
return ids
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during ChromaDB upsert: {str(e)}")
|
||||
raise
|
||||
|
||||
async def query(self, query: str, top_k: int) -> list[dict[str, Any]]:
|
||||
try:
|
||||
embedding = await self.embedding_func(
|
||||
[query], _priority=DEFAULT_QUERY_PRIORITY
|
||||
) # higher priority for query
|
||||
|
||||
results = self._collection.query(
|
||||
query_embeddings=embedding.tolist()
|
||||
if not isinstance(embedding, list)
|
||||
else embedding,
|
||||
n_results=top_k * 2, # Request more results to allow for filtering
|
||||
include=["metadatas", "distances", "documents"],
|
||||
)
|
||||
|
||||
# Filter results by cosine similarity threshold and take top k
|
||||
# We request 2x results initially to have enough after filtering
|
||||
# ChromaDB returns cosine similarity (1 = identical, 0 = orthogonal)
|
||||
# We convert to distance (0 = identical, 1 = orthogonal) via (1 - similarity)
|
||||
# Only keep results with distance below threshold, then take top k
|
||||
return [
|
||||
{
|
||||
"id": results["ids"][0][i],
|
||||
"distance": 1 - results["distances"][0][i],
|
||||
"content": results["documents"][0][i],
|
||||
"created_at": results["metadatas"][0][i].get("created_at"),
|
||||
**results["metadatas"][0][i],
|
||||
}
|
||||
for i in range(len(results["ids"][0]))
|
||||
if (1 - results["distances"][0][i]) >= self.cosine_better_than_threshold
|
||||
][:top_k]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during ChromaDB query: {str(e)}")
|
||||
raise
|
||||
|
||||
async def index_done_callback(self) -> None:
|
||||
# ChromaDB handles persistence automatically
|
||||
pass
|
||||
|
||||
async def delete_entity(self, entity_name: str) -> None:
|
||||
"""Delete an entity by its ID.
|
||||
|
||||
Args:
|
||||
entity_name: The ID of the entity to delete
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Deleting entity with ID {entity_name} from {self.namespace}")
|
||||
self._collection.delete(ids=[entity_name])
|
||||
except Exception as e:
|
||||
logger.error(f"Error during entity deletion: {str(e)}")
|
||||
raise
|
||||
|
||||
async def delete_entity_relation(self, entity_name: str) -> None:
|
||||
"""Delete an entity and its relations by ID.
|
||||
In vector DB context, this is equivalent to delete_entity.
|
||||
|
||||
Args:
|
||||
entity_name: The ID of the entity to delete
|
||||
"""
|
||||
await self.delete_entity(entity_name)
|
||||
|
||||
async def delete(self, ids: list[str]) -> None:
|
||||
"""Delete vectors with specified IDs
|
||||
|
||||
Args:
|
||||
ids: List of vector IDs to be deleted
|
||||
"""
|
||||
try:
|
||||
self._collection.delete(ids=ids)
|
||||
logger.debug(
|
||||
f"Successfully deleted {len(ids)} vectors from {self.namespace}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error while deleting vectors from {self.namespace}: {e}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during prefix search in ChromaDB: {str(e)}")
|
||||
raise
|
||||
|
||||
async def get_by_id(self, id: str) -> dict[str, Any] | None:
|
||||
"""Get vector data by its ID
|
||||
|
||||
Args:
|
||||
id: The unique identifier of the vector
|
||||
|
||||
Returns:
|
||||
The vector data if found, or None if not found
|
||||
"""
|
||||
try:
|
||||
# Query the collection for a single vector by ID
|
||||
result = self._collection.get(
|
||||
ids=[id], include=["metadatas", "embeddings", "documents"]
|
||||
)
|
||||
|
||||
if not result or not result["ids"] or len(result["ids"]) == 0:
|
||||
return None
|
||||
|
||||
# Format the result to match the expected structure
|
||||
return {
|
||||
"id": result["ids"][0],
|
||||
"vector": result["embeddings"][0],
|
||||
"content": result["documents"][0],
|
||||
"created_at": result["metadatas"][0].get("created_at"),
|
||||
**result["metadatas"][0],
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving vector data for ID {id}: {e}")
|
||||
return None
|
||||
|
||||
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
||||
"""Get multiple vector data by their IDs
|
||||
|
||||
Args:
|
||||
ids: List of unique identifiers
|
||||
|
||||
Returns:
|
||||
List of vector data objects that were found
|
||||
"""
|
||||
if not ids:
|
||||
return []
|
||||
|
||||
try:
|
||||
# Query the collection for multiple vectors by IDs
|
||||
result = self._collection.get(
|
||||
ids=ids, include=["metadatas", "embeddings", "documents"]
|
||||
)
|
||||
|
||||
if not result or not result["ids"] or len(result["ids"]) == 0:
|
||||
return []
|
||||
|
||||
# Format the results to match the expected structure and preserve ordering
|
||||
formatted_map: dict[str, dict[str, Any]] = {}
|
||||
for i, result_id in enumerate(result["ids"]):
|
||||
record = {
|
||||
"id": result_id,
|
||||
"vector": result["embeddings"][i],
|
||||
"content": result["documents"][i],
|
||||
"created_at": result["metadatas"][i].get("created_at"),
|
||||
**result["metadatas"][i],
|
||||
}
|
||||
formatted_map[str(result_id)] = record
|
||||
|
||||
ordered_results: list[dict[str, Any] | None] = []
|
||||
for requested_id in ids:
|
||||
ordered_results.append(formatted_map.get(str(requested_id)))
|
||||
|
||||
return ordered_results
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving vector data for IDs {ids}: {e}")
|
||||
return []
|
||||
|
||||
async def drop(self) -> dict[str, str]:
|
||||
"""Drop all vector data from storage and clean up resources
|
||||
|
||||
This method will delete all documents from the ChromaDB collection.
|
||||
|
||||
Returns:
|
||||
dict[str, str]: Operation status and message
|
||||
- On success: {"status": "success", "message": "data dropped"}
|
||||
- On failure: {"status": "error", "message": "<error details>"}
|
||||
"""
|
||||
try:
|
||||
# Get all IDs in the collection
|
||||
result = self._collection.get(include=[])
|
||||
if result and result["ids"] and len(result["ids"]) > 0:
|
||||
# Delete all documents
|
||||
self._collection.delete(ids=result["ids"])
|
||||
|
||||
logger.info(
|
||||
f"Process {os.getpid()} drop ChromaDB collection {self.namespace}"
|
||||
)
|
||||
return {"status": "success", "message": "data dropped"}
|
||||
except Exception as e:
|
||||
logger.error(f"Error dropping ChromaDB collection {self.namespace}: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Storage backend class factory.
|
||||
|
||||
Resolves a storage backend name (e.g. ``"JsonKVStorage"``) to its concrete
|
||||
implementation class. The four default backends are imported directly so
|
||||
they always work without depending on the ``STORAGES`` registry; everything
|
||||
else is resolved lazily through the registry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from typing import Any, Callable
|
||||
|
||||
from lightrag.kg import STORAGES
|
||||
|
||||
|
||||
def get_storage_class(storage_name: str) -> Callable[..., Any]:
|
||||
"""Return the storage backend class for ``storage_name``."""
|
||||
if storage_name == "JsonKVStorage":
|
||||
from lightrag.kg.json_kv_impl import JsonKVStorage
|
||||
|
||||
return JsonKVStorage
|
||||
if storage_name == "NanoVectorDBStorage":
|
||||
from lightrag.kg.nano_vector_db_impl import NanoVectorDBStorage
|
||||
|
||||
return NanoVectorDBStorage
|
||||
if storage_name == "NetworkXStorage":
|
||||
from lightrag.kg.networkx_impl import NetworkXStorage
|
||||
|
||||
return NetworkXStorage
|
||||
if storage_name == "JsonDocStatusStorage":
|
||||
from lightrag.kg.json_doc_status_impl import JsonDocStatusStorage
|
||||
|
||||
return JsonDocStatusStorage
|
||||
|
||||
# Fallback to dynamic import for other storage implementations.
|
||||
# STORAGES values are relative paths (e.g. ".kg.postgres_impl") authored
|
||||
# against the top-level ``lightrag`` package, so anchor the import there
|
||||
# rather than letting it resolve against this module's own package.
|
||||
import_path = STORAGES[storage_name]
|
||||
module = importlib.import_module(import_path, package="lightrag")
|
||||
return getattr(module, storage_name)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,554 @@
|
||||
from dataclasses import dataclass
|
||||
import os
|
||||
from typing import Any, Union, final
|
||||
|
||||
from lightrag.base import (
|
||||
DocProcessingStatus,
|
||||
DocStatus,
|
||||
DocStatusStorage,
|
||||
)
|
||||
from lightrag.file_atomic import reap_orphan_tmp_files
|
||||
from lightrag.utils import (
|
||||
_cooperative_yield,
|
||||
load_json,
|
||||
logger,
|
||||
validate_workspace,
|
||||
write_json,
|
||||
get_pinyin_sort_key,
|
||||
)
|
||||
from lightrag.exceptions import StorageNotInitializedError
|
||||
from .shared_storage import (
|
||||
get_namespace_data,
|
||||
get_namespace_lock,
|
||||
get_data_init_lock,
|
||||
get_update_flag,
|
||||
set_all_update_flags,
|
||||
clear_all_update_flags,
|
||||
try_initialize_namespace,
|
||||
)
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
class JsonDocStatusStorage(DocStatusStorage):
|
||||
"""JSON-file-backed document-status storage, sharing memory across processes.
|
||||
|
||||
Uses the **same shared-memory + dirty-flag protocol** as
|
||||
``JsonKVStorage`` — see that class's docstring for the canonical
|
||||
description of:
|
||||
* how ``self._data`` is a cross-process
|
||||
``multiprocessing.Manager().dict()`` proxy obtained via
|
||||
``get_namespace_data``;
|
||||
* how ``try_initialize_namespace`` ensures exactly one process
|
||||
reads the JSON file on first init;
|
||||
* how ``set_all_update_flags`` marks dirty state (semantics
|
||||
*reversed* from the file-backed classes
|
||||
``NanoVectorDBStorage`` / ``FaissVectorDBStorage`` /
|
||||
``NetworkXStorage``);
|
||||
* how ``index_done_callback`` flushes and calls
|
||||
``clear_all_update_flags``;
|
||||
* why ``_storage_lock`` wraps **every** ``self._data`` access
|
||||
(not just commit / reload).
|
||||
|
||||
Differences from ``JsonKVStorage`` (in this class only):
|
||||
* ``upsert`` calls ``index_done_callback`` synchronously after
|
||||
mutating shared memory, so doc-status changes hit disk
|
||||
immediately rather than being deferred to the pipeline's
|
||||
batched ``_insert_done()``. Rationale: doc-status is the
|
||||
recovery anchor for the ingest pipeline — if the process
|
||||
crashes after an in-memory upsert but before the next batch
|
||||
commit, the doc must still be visible as PENDING/PROCESSING
|
||||
on restart. The other writes (``delete``, ``drop``) follow
|
||||
the standard deferred-commit pattern.
|
||||
* Pre-upsert preparation (``chunks_list`` default) runs
|
||||
*outside* the lock because it only mutates the caller-
|
||||
supplied dict, not the shared store.
|
||||
* Read methods are richer (``get_docs_by_status`` /
|
||||
``get_docs_by_track_id`` / ``get_docs_paginated`` /
|
||||
``get_doc_by_file_path`` / etc.), but they all follow the
|
||||
same "acquire ``_storage_lock``, scan ``self._data``, copy
|
||||
values out before returning" template.
|
||||
|
||||
Non-pipeline write paths:
|
||||
* ``drop`` — destructive, **not** serialized; the caller must
|
||||
hold the pipeline ``busy`` reservation (the
|
||||
``/documents/clear`` endpoint does this).
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
# Reject path traversal before using workspace in a file path
|
||||
validate_workspace(self.workspace)
|
||||
working_dir = self.global_config["working_dir"]
|
||||
if self.workspace:
|
||||
# Include workspace in the file path for data isolation
|
||||
workspace_dir = os.path.join(working_dir, self.workspace)
|
||||
else:
|
||||
# Default behavior when workspace is empty
|
||||
workspace_dir = working_dir
|
||||
self.workspace = ""
|
||||
|
||||
os.makedirs(workspace_dir, exist_ok=True)
|
||||
self._file_name = os.path.join(workspace_dir, f"kv_store_{self.namespace}.json")
|
||||
self._data = None
|
||||
self._storage_lock = None
|
||||
self.storage_updated = None
|
||||
|
||||
reap_orphan_tmp_files(self._file_name, self.workspace or "_")
|
||||
|
||||
async def initialize(self):
|
||||
"""Bind to the shared namespace dict and load from disk on first init.
|
||||
|
||||
Same protocol as ``JsonKVStorage.initialize``: a global init
|
||||
lock (``try_initialize_namespace``) elects one process to read
|
||||
the JSON file into the shared ``self._data``; other processes
|
||||
skip the read and see the same shared dict.
|
||||
"""
|
||||
self._storage_lock = get_namespace_lock(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
self.storage_updated = await get_update_flag(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
async with get_data_init_lock():
|
||||
# check need_init must before get_namespace_data
|
||||
need_init = await try_initialize_namespace(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
self._data = await get_namespace_data(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
if need_init:
|
||||
loaded_data = load_json(self._file_name) or {}
|
||||
async with self._storage_lock:
|
||||
self._data.update(loaded_data)
|
||||
logger.info(
|
||||
f"[{self.workspace}] Process {os.getpid()} doc status load {self.namespace} with {len(loaded_data)} records"
|
||||
)
|
||||
|
||||
async def filter_keys(self, keys: set[str]) -> set[str]:
|
||||
"""Return keys that should be processed (not in storage or not successfully processed)"""
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonDocStatusStorage")
|
||||
async with self._storage_lock:
|
||||
return set(keys) - set(self._data.keys())
|
||||
|
||||
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
||||
ordered_results: list[dict[str, Any] | None] = []
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonDocStatusStorage")
|
||||
async with self._storage_lock:
|
||||
for id in ids:
|
||||
data = self._data.get(id, None)
|
||||
if data:
|
||||
ordered_results.append(data.copy())
|
||||
else:
|
||||
ordered_results.append(None)
|
||||
return ordered_results
|
||||
|
||||
async def get_status_counts(self) -> dict[str, int]:
|
||||
"""Get counts of documents in each status"""
|
||||
counts = {status.value: 0 for status in DocStatus}
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonDocStatusStorage")
|
||||
async with self._storage_lock:
|
||||
for doc in self._data.values():
|
||||
counts[doc["status"]] += 1
|
||||
return counts
|
||||
|
||||
async def get_docs_by_status(
|
||||
self, status: DocStatus
|
||||
) -> dict[str, DocProcessingStatus]:
|
||||
"""Get all documents with a specific status"""
|
||||
return await self.get_docs_by_statuses([status])
|
||||
|
||||
async def get_docs_by_statuses(
|
||||
self, statuses: list[DocStatus]
|
||||
) -> dict[str, DocProcessingStatus]:
|
||||
"""Get all documents matching any of the given statuses in a single pass.
|
||||
|
||||
Acquires the storage lock once and scans the in-memory dict once,
|
||||
filtering against a set of status values. More efficient than N separate
|
||||
get_docs_by_status() calls, which would acquire the lock N times and scan
|
||||
the data N times.
|
||||
"""
|
||||
if not statuses:
|
||||
return {}
|
||||
status_values = {s.value for s in statuses}
|
||||
result = {}
|
||||
async with self._storage_lock:
|
||||
for k, v in self._data.items():
|
||||
if v["status"] not in status_values:
|
||||
continue
|
||||
try:
|
||||
data = v.copy()
|
||||
data.pop("content", None)
|
||||
if not data.get("file_path"):
|
||||
data["file_path"] = "no-file-path"
|
||||
if "metadata" not in data:
|
||||
data["metadata"] = {}
|
||||
if "error_msg" not in data:
|
||||
data["error_msg"] = None
|
||||
result[k] = DocProcessingStatus(**data)
|
||||
except (KeyError, TypeError) as e:
|
||||
logger.error(
|
||||
f"[{self.workspace}] Missing required field for document {k}: {e}"
|
||||
)
|
||||
continue
|
||||
return result
|
||||
|
||||
async def get_docs_by_track_id(
|
||||
self, track_id: str
|
||||
) -> dict[str, DocProcessingStatus]:
|
||||
"""Get all documents with a specific track_id"""
|
||||
result = {}
|
||||
async with self._storage_lock:
|
||||
for k, v in self._data.items():
|
||||
if v.get("track_id") == track_id:
|
||||
try:
|
||||
# Make a copy of the data to avoid modifying the original
|
||||
data = v.copy()
|
||||
# Remove deprecated content field if it exists
|
||||
data.pop("content", None)
|
||||
# Normalize missing or null file_path
|
||||
if not data.get("file_path"):
|
||||
data["file_path"] = "no-file-path"
|
||||
# Ensure new fields exist with default values
|
||||
if "metadata" not in data:
|
||||
data["metadata"] = {}
|
||||
if "error_msg" not in data:
|
||||
data["error_msg"] = None
|
||||
result[k] = DocProcessingStatus(**data)
|
||||
except KeyError as e:
|
||||
logger.error(
|
||||
f"[{self.workspace}] Missing required field for document {k}: {e}"
|
||||
)
|
||||
continue
|
||||
return result
|
||||
|
||||
async def index_done_callback(self) -> None:
|
||||
"""Flush dirty shared memory to disk and clear all dirty flags.
|
||||
|
||||
Identical commit protocol to ``JsonKVStorage.index_done_callback``
|
||||
(snapshot the shared dict → ``write_json`` → if sanitization
|
||||
happened reload the cleaned data → ``clear_all_update_flags``).
|
||||
See ``JsonKVStorage`` docstring for details.
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
if self.storage_updated.value:
|
||||
data_dict = (
|
||||
dict(self._data) if hasattr(self._data, "_getvalue") else self._data
|
||||
)
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Process {os.getpid()} doc status writting {len(data_dict)} records to {self.namespace}"
|
||||
)
|
||||
|
||||
# Write JSON and check if sanitization was applied
|
||||
needs_reload = write_json(data_dict, self._file_name)
|
||||
|
||||
# If data was sanitized, reload cleaned data to update shared memory
|
||||
if needs_reload:
|
||||
logger.info(
|
||||
f"[{self.workspace}] Reloading sanitized data into shared memory for {self.namespace}"
|
||||
)
|
||||
cleaned_data = load_json(self._file_name)
|
||||
if cleaned_data is not None:
|
||||
self._data.clear()
|
||||
self._data.update(cleaned_data)
|
||||
|
||||
await clear_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
|
||||
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
||||
"""Insert/update doc-status records and **persist immediately**.
|
||||
|
||||
Differs from ``JsonKVStorage.upsert`` in that it calls
|
||||
``index_done_callback`` synchronously at the end, so changes
|
||||
are flushed to disk before this coroutine returns. Rationale:
|
||||
doc-status is the recovery anchor for the ingest pipeline — if
|
||||
the process crashes after an in-memory upsert but before the
|
||||
next batch commit, the doc must still be visible as
|
||||
PENDING/PROCESSING on restart.
|
||||
|
||||
Steps:
|
||||
1. Pre-process the caller's dict (default ``chunks_list``)
|
||||
**outside** the lock — only mutates the caller-supplied
|
||||
value dicts, not shared state.
|
||||
2. Under ``_storage_lock``, ``self._data.update(data)`` and
|
||||
``set_all_update_flags`` to mark every process dirty.
|
||||
3. Await ``index_done_callback`` for an immediate flush.
|
||||
|
||||
See ``JsonKVStorage`` class docstring for the shared-memory +
|
||||
dirty-flag protocol that underpins step 2.
|
||||
"""
|
||||
if not data:
|
||||
return
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Inserting {len(data)} records to {self.namespace}"
|
||||
)
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonDocStatusStorage")
|
||||
# Prepare data outside the lock: this only mutates the caller-supplied
|
||||
# dict values, not shared storage state, so no lock needed here.
|
||||
for i, (doc_id, doc_data) in enumerate(data.items(), start=1):
|
||||
if "chunks_list" not in doc_data:
|
||||
doc_data["chunks_list"] = []
|
||||
await _cooperative_yield(i)
|
||||
async with self._storage_lock:
|
||||
self._data.update(data)
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
|
||||
await self.index_done_callback()
|
||||
|
||||
async def is_empty(self) -> bool:
|
||||
"""Check if the storage is empty
|
||||
|
||||
Returns:
|
||||
bool: True if storage is empty, False otherwise
|
||||
|
||||
Raises:
|
||||
StorageNotInitializedError: If storage is not initialized
|
||||
"""
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonDocStatusStorage")
|
||||
async with self._storage_lock:
|
||||
return len(self._data) == 0
|
||||
|
||||
async def get_by_id(self, id: str) -> Union[dict[str, Any], None]:
|
||||
async with self._storage_lock:
|
||||
return self._data.get(id)
|
||||
|
||||
async def get_docs_paginated(
|
||||
self,
|
||||
status_filter: DocStatus | None = None,
|
||||
status_filters: list[DocStatus] | None = None,
|
||||
page: int = 1,
|
||||
page_size: int = 50,
|
||||
sort_field: str = "updated_at",
|
||||
sort_direction: str = "desc",
|
||||
) -> tuple[list[tuple[str, DocProcessingStatus]], int]:
|
||||
"""Get documents with pagination support
|
||||
|
||||
Args:
|
||||
status_filter: Filter by document status, None for all statuses
|
||||
page: Page number (1-based)
|
||||
page_size: Number of documents per page (10-200)
|
||||
sort_field: Field to sort by ('created_at', 'updated_at', 'id')
|
||||
sort_direction: Sort direction ('asc' or 'desc')
|
||||
|
||||
Returns:
|
||||
Tuple of (list of (doc_id, DocProcessingStatus) tuples, total_count)
|
||||
"""
|
||||
status_filter_values = self.resolve_status_filter_values(
|
||||
status_filter=status_filter,
|
||||
status_filters=status_filters,
|
||||
)
|
||||
|
||||
# Validate parameters
|
||||
if page < 1:
|
||||
page = 1
|
||||
if page_size < 10:
|
||||
page_size = 10
|
||||
elif page_size > 200:
|
||||
page_size = 200
|
||||
|
||||
if sort_field not in ["created_at", "updated_at", "id", "file_path"]:
|
||||
sort_field = "updated_at"
|
||||
|
||||
if sort_direction.lower() not in ["asc", "desc"]:
|
||||
sort_direction = "desc"
|
||||
|
||||
# For JSON storage, we load all data and sort/filter in memory
|
||||
all_docs = []
|
||||
|
||||
async with self._storage_lock:
|
||||
for doc_id, doc_data in self._data.items():
|
||||
# Apply status filter
|
||||
if (
|
||||
status_filter_values is not None
|
||||
and doc_data.get("status") not in status_filter_values
|
||||
):
|
||||
continue
|
||||
|
||||
try:
|
||||
# Prepare document data
|
||||
data = doc_data.copy()
|
||||
data.pop("content", None)
|
||||
if not data.get("file_path"):
|
||||
data["file_path"] = "no-file-path"
|
||||
if "metadata" not in data:
|
||||
data["metadata"] = {}
|
||||
if "error_msg" not in data:
|
||||
data["error_msg"] = None
|
||||
|
||||
doc_status = DocProcessingStatus(**data)
|
||||
|
||||
# Add sort key for sorting
|
||||
if sort_field == "id":
|
||||
doc_status._sort_key = doc_id
|
||||
elif sort_field == "file_path":
|
||||
# Use pinyin sorting for file_path field to support Chinese characters
|
||||
file_path_value = getattr(doc_status, sort_field, "")
|
||||
doc_status._sort_key = get_pinyin_sort_key(file_path_value)
|
||||
else:
|
||||
doc_status._sort_key = getattr(doc_status, sort_field, "")
|
||||
|
||||
all_docs.append((doc_id, doc_status))
|
||||
|
||||
except KeyError as e:
|
||||
logger.error(
|
||||
f"[{self.workspace}] Error processing document {doc_id}: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
# Sort documents
|
||||
reverse_sort = sort_direction.lower() == "desc"
|
||||
all_docs.sort(
|
||||
key=lambda x: getattr(x[1], "_sort_key", ""), reverse=reverse_sort
|
||||
)
|
||||
|
||||
# Remove sort key from documents
|
||||
for doc_id, doc in all_docs:
|
||||
if hasattr(doc, "_sort_key"):
|
||||
delattr(doc, "_sort_key")
|
||||
|
||||
total_count = len(all_docs)
|
||||
|
||||
# Apply pagination
|
||||
start_idx = (page - 1) * page_size
|
||||
end_idx = start_idx + page_size
|
||||
paginated_docs = all_docs[start_idx:end_idx]
|
||||
|
||||
return paginated_docs, total_count
|
||||
|
||||
async def get_all_status_counts(self) -> dict[str, int]:
|
||||
"""Get counts of documents in each status for all documents
|
||||
|
||||
Returns:
|
||||
Dictionary mapping status names to counts, including 'all' field
|
||||
"""
|
||||
counts = await self.get_status_counts()
|
||||
|
||||
# Add 'all' field with total count
|
||||
total_count = sum(counts.values())
|
||||
counts["all"] = total_count
|
||||
|
||||
return counts
|
||||
|
||||
async def delete(self, doc_ids: list[str]) -> None:
|
||||
"""Remove doc-status records from shared memory.
|
||||
|
||||
Unlike ``upsert``, ``delete`` does **not** force an immediate
|
||||
flush — it follows the standard deferred-commit pattern.
|
||||
Persistence happens at the next ``index_done_callback``
|
||||
(driven by the pipeline's ``_insert_done()`` at end of batch).
|
||||
|
||||
Only calls ``set_all_update_flags`` if at least one key was
|
||||
actually present (avoids creating spurious dirty state for
|
||||
no-op deletes).
|
||||
|
||||
Args:
|
||||
doc_ids: List of document IDs to be deleted from storage
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
any_deleted = False
|
||||
for doc_id in doc_ids:
|
||||
result = self._data.pop(doc_id, None)
|
||||
if result is not None:
|
||||
any_deleted = True
|
||||
|
||||
if any_deleted:
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
|
||||
async def get_doc_by_file_path(self, file_path: str) -> Union[dict[str, Any], None]:
|
||||
"""Get document by file path
|
||||
|
||||
Args:
|
||||
file_path: The file path to search for
|
||||
|
||||
Returns:
|
||||
Union[dict[str, Any], None]: Document data if found, None otherwise
|
||||
Returns the same format as get_by_ids method
|
||||
"""
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonDocStatusStorage")
|
||||
|
||||
async with self._storage_lock:
|
||||
for doc_id, doc_data in self._data.items():
|
||||
if doc_data.get("file_path") == file_path:
|
||||
# Return complete document data, consistent with get_by_ids method
|
||||
return doc_data
|
||||
|
||||
return None
|
||||
|
||||
async def get_doc_by_file_basename(
|
||||
self, basename: str
|
||||
) -> Union[tuple[str, dict[str, Any]], None]:
|
||||
"""Find an existing record whose canonical basename matches.
|
||||
|
||||
The caller is responsible for passing an already-canonical basename.
|
||||
Stored ``file_path`` values are canonicalized by the business layer, so
|
||||
this lookup intentionally performs an exact match only.
|
||||
"""
|
||||
if not basename:
|
||||
return None
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonDocStatusStorage")
|
||||
|
||||
if basename == "unknown_source":
|
||||
return None
|
||||
async with self._storage_lock:
|
||||
for doc_id, doc_data in self._data.items():
|
||||
if doc_data.get("file_path") == basename:
|
||||
return doc_id, doc_data
|
||||
return None
|
||||
|
||||
async def get_doc_by_content_hash(
|
||||
self, content_hash: str
|
||||
) -> Union[tuple[str, dict[str, Any]], None]:
|
||||
"""Find an existing record whose content_hash field matches."""
|
||||
if not content_hash:
|
||||
return None
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonDocStatusStorage")
|
||||
|
||||
async with self._storage_lock:
|
||||
for doc_id, doc_data in self._data.items():
|
||||
if doc_data.get("content_hash") == content_hash:
|
||||
return doc_id, doc_data
|
||||
return None
|
||||
|
||||
async def drop(self) -> dict[str, str]:
|
||||
"""Clear shared memory and immediately persist the empty state.
|
||||
|
||||
This method will:
|
||||
1. Clear the shared ``self._data`` dict under
|
||||
``_storage_lock`` (visible to all processes immediately).
|
||||
2. ``set_all_update_flags`` so every process knows there is
|
||||
dirty state pending persistence.
|
||||
3. Call ``index_done_callback`` synchronously to flush the
|
||||
empty state to disk and clear the dirty flags.
|
||||
|
||||
Caller contract:
|
||||
``drop`` is destructive and **not** serialized by this
|
||||
storage class. The caller must hold the pipeline ``busy``
|
||||
reservation (the ``/documents/clear`` endpoint does this)
|
||||
before invoking it. See class docstring,
|
||||
*Non-pipeline write paths*.
|
||||
|
||||
Returns:
|
||||
dict[str, str]: Operation status and message
|
||||
- On success: {"status": "success", "message": "data dropped"}
|
||||
- On failure: {"status": "error", "message": "<error details>"}
|
||||
"""
|
||||
try:
|
||||
async with self._storage_lock:
|
||||
self._data.clear()
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
|
||||
await self.index_done_callback()
|
||||
logger.info(
|
||||
f"[{self.workspace}] Process {os.getpid()} drop {self.namespace}"
|
||||
)
|
||||
return {"status": "success", "message": "data dropped"}
|
||||
except Exception as e:
|
||||
logger.error(f"[{self.workspace}] Error dropping {self.namespace}: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
@@ -0,0 +1,492 @@
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, final
|
||||
|
||||
from lightrag.base import (
|
||||
BaseKVStorage,
|
||||
)
|
||||
from lightrag.file_atomic import reap_orphan_tmp_files
|
||||
from lightrag.utils import (
|
||||
_cooperative_yield,
|
||||
load_json,
|
||||
logger,
|
||||
validate_workspace,
|
||||
write_json,
|
||||
)
|
||||
from lightrag.exceptions import StorageNotInitializedError
|
||||
from .shared_storage import (
|
||||
get_namespace_data,
|
||||
get_namespace_lock,
|
||||
get_data_init_lock,
|
||||
get_update_flag,
|
||||
set_all_update_flags,
|
||||
clear_all_update_flags,
|
||||
try_initialize_namespace,
|
||||
)
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
class JsonKVStorage(BaseKVStorage):
|
||||
"""JSON-file-backed KV storage with **shared in-memory state across processes**.
|
||||
|
||||
This class uses a *fundamentally different* cross-process model from
|
||||
``NanoVectorDBStorage`` / ``FaissVectorDBStorage`` / ``NetworkXStorage``
|
||||
(which keep one in-memory copy per process and reconcile via file
|
||||
reloads). Compare carefully before changing either side.
|
||||
|
||||
Storage model:
|
||||
``self._data`` is **not** a per-process dict — it is the value
|
||||
returned by ``get_namespace_data(namespace, workspace=...)``, i.e.
|
||||
a reference into ``shared_storage._shared_dicts``. In multi-
|
||||
process mode this is a ``multiprocessing.Manager().dict()`` proxy
|
||||
that every worker sees the **same instance** of; in single-
|
||||
process mode it degrades to a plain ``dict``. Either way, a
|
||||
mutation in any process is *immediately* visible to every other
|
||||
process — there is no reload needed.
|
||||
|
||||
The on-disk file at
|
||||
``working_dir/[workspace/]kv_store_<namespace>.json`` exists for
|
||||
durability only. It is the source of truth at startup and the
|
||||
target of ``index_done_callback`` flushes, but is **not** part of
|
||||
the steady-state read/write path.
|
||||
|
||||
First-time load (``initialize``):
|
||||
``try_initialize_namespace`` is a global init lock that returns
|
||||
``True`` to exactly one process per ``(namespace, workspace)``.
|
||||
That process reads the JSON file and populates ``self._data``
|
||||
under ``_storage_lock``. Other processes skip the load — they
|
||||
will see the data through the same shared ``self._data`` proxy.
|
||||
|
||||
Cross-process sync protocol (note: reversed semantics vs file-backed
|
||||
classes):
|
||||
Anyone writing (``upsert`` / ``delete`` / ``drop``):
|
||||
1. Mutate ``self._data`` under ``_storage_lock`` (same lock,
|
||||
same dict, all processes see the change immediately).
|
||||
2. Call ``set_all_update_flags`` to mark **every** process's
|
||||
``storage_updated`` flag ``True``. Here ``True`` means
|
||||
*"there is dirty data that still needs to be flushed"*,
|
||||
not *"there is fresher data on disk that I need to
|
||||
reload"* as in the file-backed implementations.
|
||||
Commit (``index_done_callback``):
|
||||
1. Under ``_storage_lock``, if ``storage_updated.value`` is
|
||||
``True``, snapshot ``self._data`` and write it to disk
|
||||
via ``write_json`` (atomic).
|
||||
2. ``clear_all_update_flags`` — wipe every process's flag
|
||||
back to ``False``. Because the in-memory state is already
|
||||
consistent across processes, there is nothing for the
|
||||
*other* processes to do; the clear is just a
|
||||
"the dirty data has been persisted" signal.
|
||||
|
||||
Lock scope:
|
||||
``_storage_lock`` is a per-``(namespace, workspace)`` keyed lock
|
||||
spanning intra-process coroutines **and** inter-process workers.
|
||||
Unlike the file-backed classes (which only lock reload/commit
|
||||
critical sections), this class **holds the lock over every
|
||||
``self._data`` access** — read or write — because the underlying
|
||||
``Manager().dict()`` is not free-threaded across processes.
|
||||
|
||||
Two places intentionally do work outside the lock for latency
|
||||
reasons:
|
||||
* ``upsert`` performs its per-key timestamp prep loop inside
|
||||
the lock but yields to the event loop via
|
||||
``_cooperative_yield`` between keys (safe: ``NamespaceLock``
|
||||
is non-reentrant, so siblings blocked on it stay blocked).
|
||||
* ``JsonDocStatusStorage.upsert`` prepares its caller-supplied
|
||||
dict outside the lock (it only mutates the input, not the
|
||||
shared store).
|
||||
|
||||
Who can write:
|
||||
Pipeline ``busy`` still serializes the document ingest / purge
|
||||
flows, but the *file-flush trigger* is symmetric: any process
|
||||
whose ``storage_updated.value`` is ``True`` when
|
||||
``index_done_callback`` fires will perform the write. In a
|
||||
single-writer pipeline this is always the same process; if you
|
||||
ever permit multiple writers, two processes may race to flush
|
||||
the same in-memory state — that race is safe (both flush the
|
||||
same shared dict, ``write_json`` is atomic per file) but
|
||||
wasteful, and the ``clear_all_update_flags`` after each flush
|
||||
means subsequent re-flushes are no-ops.
|
||||
|
||||
Caveats vs file-backed implementations:
|
||||
* **No reload path.** If something writes to the on-disk file
|
||||
out of band, this class will not pick it up until restart.
|
||||
The file is only ever written by ``index_done_callback`` and
|
||||
read once in ``initialize``.
|
||||
* **No ``_get_*`` entry method.** Adding one would be wrong —
|
||||
there's nothing to "get fresher than" since the in-memory
|
||||
state is already the shared, authoritative view.
|
||||
* **``write_json`` may sanitize.** If sanitization happens, the
|
||||
on-disk JSON differs from what was in memory; the callback
|
||||
re-reads the cleaned file back into ``self._data`` under the
|
||||
same lock so the shared view stays consistent with disk.
|
||||
|
||||
Non-pipeline write paths:
|
||||
* ``drop`` — destructive, **not** serialized by this storage
|
||||
class. Currently gated by the API layer
|
||||
(``/documents/clear``); any new caller must hold the pipeline
|
||||
``busy`` reservation.
|
||||
* ``upsert`` / ``delete`` invoked from non-pipeline admin flows
|
||||
(cache management, etc.) — safe under the shared-lock model,
|
||||
but consumers should still respect the pipeline gate to avoid
|
||||
interleaving with batched ingest work.
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
# Reject path traversal before using workspace in a file path
|
||||
validate_workspace(self.workspace)
|
||||
working_dir = self.global_config["working_dir"]
|
||||
if self.workspace:
|
||||
# Include workspace in the file path for data isolation
|
||||
workspace_dir = os.path.join(working_dir, self.workspace)
|
||||
else:
|
||||
# Default behavior when workspace is empty
|
||||
workspace_dir = working_dir
|
||||
self.workspace = ""
|
||||
|
||||
os.makedirs(workspace_dir, exist_ok=True)
|
||||
self._file_name = os.path.join(workspace_dir, f"kv_store_{self.namespace}.json")
|
||||
self._data = None
|
||||
self._storage_lock = None
|
||||
self.storage_updated = None
|
||||
|
||||
reap_orphan_tmp_files(self._file_name, self.workspace or "_")
|
||||
|
||||
async def initialize(self):
|
||||
"""Bind to the shared namespace dict and load from disk on first init.
|
||||
|
||||
``try_initialize_namespace`` is a global init lock that returns
|
||||
``True`` for exactly one process per ``(namespace, workspace)``;
|
||||
that process reads the JSON file and populates the shared
|
||||
``self._data`` under ``_storage_lock``. Subsequent processes
|
||||
skip the file read — they will see the same shared dict via
|
||||
``get_namespace_data``.
|
||||
|
||||
For ``*_cache`` namespaces an extra
|
||||
``_migrate_legacy_cache_structure`` pass runs against the loaded
|
||||
data and may rewrite the on-disk file if a migration was applied.
|
||||
"""
|
||||
self._storage_lock = get_namespace_lock(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
self.storage_updated = await get_update_flag(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
async with get_data_init_lock():
|
||||
# check need_init must before get_namespace_data
|
||||
need_init = await try_initialize_namespace(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
self._data = await get_namespace_data(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
if need_init:
|
||||
loaded_data = load_json(self._file_name) or {}
|
||||
async with self._storage_lock:
|
||||
# Migrate legacy cache structure if needed
|
||||
if self.namespace.endswith("_cache"):
|
||||
loaded_data = await self._migrate_legacy_cache_structure(
|
||||
loaded_data
|
||||
)
|
||||
|
||||
self._data.update(loaded_data)
|
||||
data_count = len(loaded_data)
|
||||
|
||||
logger.info(
|
||||
f"[{self.workspace}] Process {os.getpid()} KV load {self.namespace} with {data_count} records"
|
||||
)
|
||||
|
||||
async def index_done_callback(self) -> None:
|
||||
"""Flush dirty in-memory state to disk and clear all dirty flags.
|
||||
|
||||
Commit point in the shared-memory protocol (see class docstring,
|
||||
*Cross-process sync protocol*). Steps:
|
||||
1. Under ``_storage_lock``, check this process's
|
||||
``storage_updated.value``. If ``False``, nothing to do —
|
||||
return.
|
||||
2. Snapshot ``self._data`` (converting from ``Manager.dict``
|
||||
proxy to a plain ``dict`` so the JSON encoder doesn't trip
|
||||
over the proxy) and write it via ``write_json``.
|
||||
3. If ``write_json`` reports sanitization was applied, the
|
||||
on-disk file no longer matches what was in memory — reload
|
||||
the cleaned data back into ``self._data`` under the same
|
||||
lock so the shared view stays consistent.
|
||||
4. ``clear_all_update_flags`` — wipe every process's
|
||||
``storage_updated`` flag back to ``False``, signaling
|
||||
that the dirty data has been persisted.
|
||||
|
||||
Note the **semantic difference** from the file-backed classes'
|
||||
commit: there is no ``set_all_update_flags`` here. The shared
|
||||
dict is already consistent across processes; the only thing
|
||||
``index_done_callback`` does globally is *clear* the dirty
|
||||
flags.
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
if self.storage_updated.value:
|
||||
data_dict = (
|
||||
dict(self._data) if hasattr(self._data, "_getvalue") else self._data
|
||||
)
|
||||
|
||||
# Calculate data count - all data is now flattened
|
||||
data_count = len(data_dict)
|
||||
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Process {os.getpid()} KV writting {data_count} records to {self.namespace}"
|
||||
)
|
||||
|
||||
# Write JSON and check if sanitization was applied
|
||||
needs_reload = write_json(data_dict, self._file_name)
|
||||
|
||||
# If data was sanitized, reload cleaned data to update shared memory
|
||||
if needs_reload:
|
||||
logger.info(
|
||||
f"[{self.workspace}] Reloading sanitized data into shared memory for {self.namespace}"
|
||||
)
|
||||
cleaned_data = load_json(self._file_name)
|
||||
if cleaned_data is not None:
|
||||
self._data.clear()
|
||||
self._data.update(cleaned_data)
|
||||
|
||||
await clear_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
|
||||
async def get_by_id(self, id: str) -> dict[str, Any] | None:
|
||||
async with self._storage_lock:
|
||||
result = self._data.get(id)
|
||||
if result:
|
||||
# Create a copy to avoid modifying the original data
|
||||
result = dict(result)
|
||||
# Ensure time fields are present, provide default values for old data
|
||||
result.setdefault("create_time", 0)
|
||||
result.setdefault("update_time", 0)
|
||||
# Ensure _id field contains the clean ID
|
||||
result["_id"] = id
|
||||
return result
|
||||
|
||||
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
||||
async with self._storage_lock:
|
||||
results = []
|
||||
for id in ids:
|
||||
data = self._data.get(id, None)
|
||||
if data:
|
||||
# Create a copy to avoid modifying the original data
|
||||
result = {k: v for k, v in data.items()}
|
||||
# Ensure time fields are present, provide default values for old data
|
||||
result.setdefault("create_time", 0)
|
||||
result.setdefault("update_time", 0)
|
||||
# Ensure _id field contains the clean ID
|
||||
result["_id"] = id
|
||||
results.append(result)
|
||||
else:
|
||||
results.append(None)
|
||||
return results
|
||||
|
||||
async def filter_keys(self, keys: set[str]) -> set[str]:
|
||||
async with self._storage_lock:
|
||||
return set(keys) - set(self._data.keys())
|
||||
|
||||
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
||||
"""Insert or update KV records in shared memory; mark all processes dirty.
|
||||
|
||||
Two side effects under ``_storage_lock``:
|
||||
1. Stamp ``create_time`` / ``update_time`` / ``_id`` on each
|
||||
value, then ``self._data.update(data)``. Because
|
||||
``self._data`` is the shared ``Manager.dict()`` proxy, the
|
||||
update is visible to all processes immediately — no
|
||||
reload needed.
|
||||
2. ``set_all_update_flags`` — flip every process's
|
||||
``storage_updated.value`` to ``True``. Here ``True``
|
||||
means *"there is dirty data that still needs to be
|
||||
flushed to disk"*, **not** *"there is fresher data on
|
||||
disk"* as in the file-backed classes (see class docstring
|
||||
for the contrast).
|
||||
|
||||
Persistence is deferred to the next ``index_done_callback`` (the
|
||||
pipeline calls this via ``_insert_done()`` after each batch).
|
||||
|
||||
Note: the per-key prep loop calls ``_cooperative_yield`` inside
|
||||
the lock. That is safe because ``NamespaceLock`` is non-
|
||||
reentrant — siblings waiting on this lock stay blocked across
|
||||
the yield; only unrelated coroutines benefit from the yield.
|
||||
"""
|
||||
if not data:
|
||||
return
|
||||
|
||||
import time
|
||||
|
||||
current_time = int(time.time()) # Get current Unix timestamp
|
||||
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Inserting {len(data)} records to {self.namespace}"
|
||||
)
|
||||
if self._storage_lock is None:
|
||||
raise StorageNotInitializedError("JsonKVStorage")
|
||||
async with self._storage_lock:
|
||||
# Add timestamps to data based on whether key exists.
|
||||
# The loop reads self._data (k in self._data) so it must stay inside
|
||||
# the lock. _cooperative_yield is safe here: NamespaceLock is
|
||||
# non-reentrant, so other coroutines waiting on this lock will block
|
||||
# until we release it; the yield only benefits unrelated coroutines.
|
||||
for i, (k, v) in enumerate(data.items(), start=1):
|
||||
# For text_chunks namespace, ensure llm_cache_list field exists
|
||||
if self.namespace.endswith("text_chunks"):
|
||||
if "llm_cache_list" not in v:
|
||||
v["llm_cache_list"] = []
|
||||
|
||||
# Add timestamps based on whether key exists
|
||||
if k in self._data: # Key exists, only update update_time
|
||||
v["update_time"] = current_time
|
||||
else: # New key, set both create_time and update_time
|
||||
v["create_time"] = current_time
|
||||
v["update_time"] = current_time
|
||||
|
||||
v["_id"] = k
|
||||
await _cooperative_yield(i)
|
||||
|
||||
self._data.update(data)
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
|
||||
async def delete(self, ids: list[str]) -> None:
|
||||
"""Remove records from shared memory; mark all processes dirty if any deleted.
|
||||
|
||||
Under ``_storage_lock``: ``self._data.pop(doc_id, None)`` for
|
||||
each id. Only calls ``set_all_update_flags`` if at least one key
|
||||
was actually present (avoids creating spurious dirty state for
|
||||
no-op deletes).
|
||||
|
||||
See class docstring for the shared-memory + dirty-flag protocol
|
||||
and the semantic contrast vs file-backed classes.
|
||||
|
||||
Args:
|
||||
ids: List of document IDs to be deleted from storage
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
any_deleted = False
|
||||
for doc_id in ids:
|
||||
result = self._data.pop(doc_id, None)
|
||||
if result is not None:
|
||||
any_deleted = True
|
||||
|
||||
if any_deleted:
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
|
||||
async def is_empty(self) -> bool:
|
||||
"""Check if the storage is empty
|
||||
|
||||
Returns:
|
||||
bool: True if storage contains no data, False otherwise
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
return len(self._data) == 0
|
||||
|
||||
async def drop(self) -> dict[str, str]:
|
||||
"""Clear shared memory and immediately persist the empty state.
|
||||
|
||||
This method will:
|
||||
1. Clear the shared ``self._data`` dict under
|
||||
``_storage_lock`` (visible to all processes immediately).
|
||||
2. ``set_all_update_flags`` so every process knows there is
|
||||
dirty state pending persistence.
|
||||
3. Call ``index_done_callback`` synchronously to flush the
|
||||
empty state to disk and clear the dirty flags.
|
||||
|
||||
Caller contract:
|
||||
``drop`` is destructive and **not** serialized by this
|
||||
storage class. The caller must hold the pipeline ``busy``
|
||||
reservation (the ``/documents/clear`` endpoint does this)
|
||||
before invoking it — running ``drop`` concurrently with an
|
||||
active document pipeline will wipe out in-flight work and
|
||||
silently lose data. See class docstring,
|
||||
*Non-pipeline write paths*.
|
||||
|
||||
Returns:
|
||||
dict[str, str]: Operation status and message
|
||||
- On success: {"status": "success", "message": "data dropped"}
|
||||
- On failure: {"status": "error", "message": "<error details>"}
|
||||
"""
|
||||
try:
|
||||
async with self._storage_lock:
|
||||
self._data.clear()
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
|
||||
await self.index_done_callback()
|
||||
logger.info(
|
||||
f"[{self.workspace}] Process {os.getpid()} drop {self.namespace}"
|
||||
)
|
||||
return {"status": "success", "message": "data dropped"}
|
||||
except Exception as e:
|
||||
logger.error(f"[{self.workspace}] Error dropping {self.namespace}: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
async def _migrate_legacy_cache_structure(self, data: dict) -> dict:
|
||||
"""Migrate legacy nested cache structure to flattened structure
|
||||
|
||||
Args:
|
||||
data: Original data dictionary that may contain legacy structure
|
||||
|
||||
Returns:
|
||||
Migrated data dictionary with flattened cache keys (sanitized if needed)
|
||||
"""
|
||||
from lightrag.utils import generate_cache_key
|
||||
|
||||
# Early return if data is empty
|
||||
if not data:
|
||||
return data
|
||||
|
||||
# Check first entry to see if it's already in new format
|
||||
first_key = next(iter(data.keys()))
|
||||
if ":" in first_key and len(first_key.split(":")) == 3:
|
||||
# Already in flattened format, return as-is
|
||||
return data
|
||||
|
||||
migrated_data = {}
|
||||
migration_count = 0
|
||||
|
||||
for key, value in data.items():
|
||||
# Check if this is a legacy nested cache structure
|
||||
if isinstance(value, dict) and all(
|
||||
isinstance(v, dict) and "return" in v for v in value.values()
|
||||
):
|
||||
# This looks like a legacy cache mode with nested structure
|
||||
mode = key
|
||||
for cache_hash, cache_entry in value.items():
|
||||
cache_type = cache_entry.get("cache_type", "extract")
|
||||
flattened_key = generate_cache_key(mode, cache_type, cache_hash)
|
||||
migrated_data[flattened_key] = cache_entry
|
||||
migration_count += 1
|
||||
else:
|
||||
# Keep non-cache data or already flattened cache data as-is
|
||||
migrated_data[key] = value
|
||||
|
||||
if migration_count > 0:
|
||||
logger.info(
|
||||
f"[{self.workspace}] Migrated {migration_count} legacy cache entries to flattened structure"
|
||||
)
|
||||
# Persist migrated data immediately and check if sanitization was applied
|
||||
needs_reload = write_json(migrated_data, self._file_name)
|
||||
|
||||
# If data was sanitized during write, reload cleaned data
|
||||
if needs_reload:
|
||||
logger.info(
|
||||
f"[{self.workspace}] Reloading sanitized migration data for {self.namespace}"
|
||||
)
|
||||
cleaned_data = load_json(self._file_name)
|
||||
if cleaned_data is not None:
|
||||
return cleaned_data # Return cleaned data to update shared memory
|
||||
|
||||
return migrated_data
|
||||
|
||||
async def finalize(self):
|
||||
"""On shutdown, flush ``*_cache`` namespaces to disk.
|
||||
|
||||
Cache namespaces are routinely written to during query/extract
|
||||
without triggering an immediate ``index_done_callback`` (caches
|
||||
churn fast and the pipeline doesn't always end at a natural
|
||||
commit point). This hook ensures whatever dirty cache state is
|
||||
in shared memory at process exit gets persisted, so the next
|
||||
run can pick it up.
|
||||
|
||||
Non-cache namespaces don't need this — their writes already
|
||||
flow through pipeline-driven ``_insert_done()`` commits.
|
||||
"""
|
||||
if self.namespace.endswith("_cache"):
|
||||
await self.index_done_callback()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,918 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
import zlib
|
||||
from typing import Any, final
|
||||
from dataclasses import dataclass
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
from lightrag.file_atomic import atomic_write, reap_orphan_tmp_files
|
||||
from lightrag.utils import (
|
||||
logger,
|
||||
compute_mdhash_id,
|
||||
validate_workspace,
|
||||
)
|
||||
|
||||
from lightrag.base import BaseVectorStorage
|
||||
from lightrag.constants import DEFAULT_QUERY_PRIORITY
|
||||
from nano_vectordb import NanoVectorDB
|
||||
from .shared_storage import (
|
||||
get_namespace_lock,
|
||||
get_update_flag,
|
||||
set_all_update_flags,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PendingNanoDoc:
|
||||
"""A buffered upsert waiting for deferred embedding and materialization.
|
||||
|
||||
``record`` holds ``__id__`` / ``__created_at__`` plus the ``meta_fields``
|
||||
(which always include ``content`` for the entity/relation/chunk vdbs), so
|
||||
the content needed for deferred embedding lives in the record itself — no
|
||||
separate copy is kept. ``vector`` starts as ``None`` and is filled either
|
||||
during the lock-held flush or by a lazy ``get_vectors_by_ids`` embedding;
|
||||
once set it is reused by the next flush instead of re-calling the model.
|
||||
The compressed ``vector`` / raw ``__vector__`` keys are added to ``record``
|
||||
only at flush time, right before ``client.upsert``.
|
||||
"""
|
||||
|
||||
record: dict[str, Any]
|
||||
vector: np.ndarray | None = None
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
class NanoVectorDBStorage(BaseVectorStorage):
|
||||
"""File-backed vector storage built on the in-memory ``NanoVectorDB``.
|
||||
|
||||
Storage model:
|
||||
A single ``NanoVectorDB`` instance lives in process memory; its full
|
||||
state is serialized to one JSON file at
|
||||
``working_dir/[workspace/]vdb_<namespace>.json``. That JSON file is
|
||||
the **only** cross-process synchronization surface — there is no
|
||||
shared memory, no message bus, and no network channel between
|
||||
processes. All cross-process visibility is therefore mediated by
|
||||
(a) an atomic file write at commit time and (b) a per-namespace
|
||||
``storage_updated`` flag distributed through
|
||||
``lightrag.kg.shared_storage``.
|
||||
|
||||
Concurrency invariants (the code in this file is correct *only* while
|
||||
all three hold):
|
||||
1. **Single writer per workspace.** The document pipeline's
|
||||
``busy`` / ``destructive_busy`` flags (see ``AGENTS.md``
|
||||
*Pipeline concurrency contract*) guarantee that at most one
|
||||
process performs ``upsert`` / ``delete`` /
|
||||
``index_done_callback`` at any time. Every other process is
|
||||
read-only with respect to this storage.
|
||||
2. **Eventual consistency is sufficient.** Read-only processes
|
||||
only need to observe the writer's data *after* the writer's
|
||||
``index_done_callback`` completes. Reads that land in the gap
|
||||
between a writer's in-memory mutation and its commit may
|
||||
legitimately return the pre-update snapshot.
|
||||
3. **NanoVectorDB operations are fully synchronous.** Under a
|
||||
single-threaded asyncio event loop, ``client.upsert`` /
|
||||
``client.query`` / ``client.delete`` cannot be preempted by
|
||||
another coroutine, which gives them implicit mutual exclusion
|
||||
over ``self._client.__storage``. This is why the methods below
|
||||
don't have to hold ``_storage_lock`` while calling into
|
||||
``client``.
|
||||
|
||||
Cross-process sync protocol:
|
||||
Writer side (``index_done_callback``):
|
||||
1. Atomically write the in-memory state to disk
|
||||
(``atomic_write`` swaps a tmp file into place).
|
||||
2. Call ``set_all_update_flags`` to flip every process's
|
||||
``storage_updated`` flag (including the writer's own).
|
||||
3. Immediately reset the writer's own flag to ``False`` so
|
||||
the next call to ``_get_client`` does not trigger a
|
||||
self-reload of the data this process just wrote.
|
||||
Reader side (any method that goes through ``_get_client``):
|
||||
1. Inside ``_storage_lock``, observe
|
||||
``storage_updated.value is True``.
|
||||
2. **Fully reload** ``self._client`` from disk — NanoVectorDB
|
||||
has no incremental sync API, so the entire JSON file is
|
||||
re-parsed and a fresh in-memory matrix is rebuilt.
|
||||
3. Reset the reader's own flag to ``False`` so concurrent
|
||||
coroutines in the same process don't double-reload.
|
||||
|
||||
Lock scope:
|
||||
``_storage_lock`` is a per-``(namespace, workspace)`` keyed lock
|
||||
spanning both intra-process coroutines and inter-process workers.
|
||||
It only wraps the *reload* and *commit* critical sections, not
|
||||
every ``client.xxx`` call. Operating on ``client`` outside the
|
||||
lock is safe today *because of invariant (3)* — if either premise
|
||||
is ever broken (e.g. ``client.xxx`` is moved to a thread pool, or
|
||||
NanoVectorDB is swapped for an async vector library), the lock
|
||||
scope must be widened to cover the mutation/read itself.
|
||||
|
||||
Non-pipeline write paths:
|
||||
The pipeline's ``busy`` gate serializes ``upsert`` / ``delete`` /
|
||||
``index_done_callback`` called from the document ingestion and
|
||||
purge flows. The following entry points are **not** serialized by
|
||||
the pipeline gate and must be guarded externally:
|
||||
* ``drop`` — currently gated by the API layer (the
|
||||
``/documents/clear`` endpoint takes the pipeline busy
|
||||
reservation before invoking it).
|
||||
* ``delete_entity`` / ``delete_entity_relation`` — currently
|
||||
not exposed in the WebUI. If you wire them up to a new
|
||||
caller, that caller must arrange single-writer
|
||||
serialization the same way the pipeline does.
|
||||
|
||||
Deferred-embedding protocol:
|
||||
``upsert`` does **not** call the embedding model. It only buffers a
|
||||
``_PendingNanoDoc`` (content-bearing record + ``vector=None``) in the
|
||||
minimal ``self._pending_upserts`` area, overwriting any prior pending
|
||||
doc for the same id (which also clears a temp vector a previous
|
||||
``get_vectors_by_ids`` may have cached). The model is called once per
|
||||
id at flush time (``_flush_pending_locked``), so repeated upserts of
|
||||
the same id — and many small upsert calls — embed only once. See
|
||||
issue #2785 and the ``OpenSearchVectorDBStorage`` equivalent.
|
||||
|
||||
Embedding runs **inside ``_storage_lock``** during the flush (not in
|
||||
``upsert``): under the single-writer invariant this keeps the content
|
||||
used for embedding consistent with the record written to disk and
|
||||
prevents a destructive op from interleaving between embed and write.
|
||||
The lock is non-reentrant, so ``_flush_pending_locked`` requires the
|
||||
caller to already hold it and operates on ``self._client`` directly
|
||||
(never through ``_get_client``).
|
||||
|
||||
Reads are read-your-writes: ``get_by_id`` / ``get_by_ids`` /
|
||||
``get_vectors_by_ids`` consult ``_pending_upserts`` first.
|
||||
``get_vectors_by_ids`` lazily embeds a pending doc on demand and
|
||||
caches the vector back for the next flush. ``query`` and
|
||||
``client_storage`` see only data already materialized into
|
||||
``self._client`` — unflushed pending data is intentionally not
|
||||
queryable. A flush failure (embedding error, count mismatch, or save
|
||||
IO error) raises through ``index_done_callback``; the pending buffer
|
||||
is preserved, and if only the save failed ``_client_dirty`` stays
|
||||
``True`` so a subsequent ``finalize`` retries the save.
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
# Reject path traversal before using workspace in a file path
|
||||
validate_workspace(self.workspace)
|
||||
self._validate_embedding_func()
|
||||
# Initialize basic attributes
|
||||
self._client = None
|
||||
self._storage_lock = None
|
||||
self.storage_updated = None
|
||||
|
||||
# Use global config value if specified, otherwise use default
|
||||
kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {})
|
||||
cosine_threshold = kwargs.get("cosine_better_than_threshold")
|
||||
if cosine_threshold is None:
|
||||
raise ValueError(
|
||||
"cosine_better_than_threshold must be specified in vector_db_storage_cls_kwargs"
|
||||
)
|
||||
self.cosine_better_than_threshold = cosine_threshold
|
||||
|
||||
working_dir = self.global_config["working_dir"]
|
||||
if self.workspace:
|
||||
# Include workspace in the file path for data isolation
|
||||
workspace_dir = os.path.join(working_dir, self.workspace)
|
||||
self.final_namespace = f"{self.workspace}_{self.namespace}"
|
||||
else:
|
||||
# Default behavior when workspace is empty
|
||||
self.final_namespace = self.namespace
|
||||
self.workspace = ""
|
||||
workspace_dir = working_dir
|
||||
|
||||
os.makedirs(workspace_dir, exist_ok=True)
|
||||
self._client_file_name = os.path.join(
|
||||
workspace_dir, f"vdb_{self.namespace}.json"
|
||||
)
|
||||
|
||||
self._max_batch_size = self.global_config["embedding_batch_num"]
|
||||
|
||||
# Sweep orphan tmp siblings left behind by hard kills mid-save before
|
||||
# NanoVectorDB opens the target file.
|
||||
reap_orphan_tmp_files(self._client_file_name, self.workspace or "_")
|
||||
|
||||
self._client = NanoVectorDB(
|
||||
self.embedding_func.embedding_dim,
|
||||
storage_file=self._client_file_name,
|
||||
)
|
||||
|
||||
# Minimal pending area for deferred embedding: id -> _PendingNanoDoc.
|
||||
# Holds only records not yet embedded+materialized into self._client;
|
||||
# it never duplicates rows already written to the client. Flushed
|
||||
# under _storage_lock by _flush_pending_locked().
|
||||
self._pending_upserts: dict[str, _PendingNanoDoc] = {}
|
||||
# True when self._client has materialized changes that have not been
|
||||
# successfully saved to disk yet. This lets finalize retry a save even
|
||||
# after a previous flush cleared the pending buffer.
|
||||
self._client_dirty = False
|
||||
|
||||
async def initialize(self):
|
||||
"""Initialize storage data"""
|
||||
# Get the update flag for cross-process update notification
|
||||
self.storage_updated = await get_update_flag(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
# Get the storage lock for use in other methods
|
||||
self._storage_lock = get_namespace_lock(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
|
||||
def _reload_client_from_disk_locked(self, *, for_write: bool = False) -> bool:
|
||||
"""Reload ``self._client`` if another process committed newer data.
|
||||
|
||||
Precondition: the caller must already hold ``_storage_lock``. This is
|
||||
used by write paths as well as reads because deferred upserts mean a
|
||||
stale writer must merge its pending buffer into the latest on-disk
|
||||
snapshot, not save over it or return without flushing.
|
||||
"""
|
||||
if not self.storage_updated.value:
|
||||
return False
|
||||
|
||||
log_message = (
|
||||
f"[{self.workspace}] Process {os.getpid()} reloading {self.namespace} "
|
||||
"due to update by another process"
|
||||
)
|
||||
if for_write:
|
||||
logger.warning(log_message)
|
||||
else:
|
||||
logger.info(log_message)
|
||||
|
||||
self._client = NanoVectorDB(
|
||||
self.embedding_func.embedding_dim,
|
||||
storage_file=self._client_file_name,
|
||||
)
|
||||
self.storage_updated.value = False
|
||||
return True
|
||||
|
||||
async def _get_client(self):
|
||||
"""Return the live ``NanoVectorDB`` instance, reloading from disk if needed.
|
||||
|
||||
This is the **single entry point** every public method funnels
|
||||
through to obtain ``self._client``. It is also the **only place
|
||||
readers transition to a fresher on-disk snapshot**: when another
|
||||
process has committed (via ``index_done_callback``) and flipped
|
||||
this process's ``storage_updated`` flag, the next call here
|
||||
rebuilds ``self._client`` by re-parsing the entire JSON file.
|
||||
NanoVectorDB has no incremental sync API — the reload is
|
||||
unconditionally a full file reload.
|
||||
|
||||
Under the *Single writer* invariant (see class docstring), the
|
||||
reload branch never fires in the writer process: the writer
|
||||
resets its own flag at the end of every ``index_done_callback``.
|
||||
The branch exists for readers.
|
||||
|
||||
``_storage_lock`` is held during the check-and-reload to (a)
|
||||
serialize concurrent reload attempts by sibling coroutines in
|
||||
the same process and (b) interlock with ``index_done_callback``
|
||||
so a reader cannot observe a partially-saved file.
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
self._reload_client_from_disk_locked()
|
||||
return self._client
|
||||
|
||||
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
||||
"""Buffer vectors for deferred embedding; persistence is deferred too.
|
||||
|
||||
Embedding is **not** performed here. Each record is buffered in
|
||||
``self._pending_upserts`` with ``vector=None`` and the embedding model
|
||||
is called once per id at flush time (``_flush_pending_locked`` during
|
||||
``index_done_callback`` / ``finalize``). This coalesces repeated
|
||||
upserts of the same id and many small upsert calls into a single
|
||||
embedding pass (see class docstring, *Deferred-embedding protocol*,
|
||||
and issue #2785).
|
||||
|
||||
Persistence:
|
||||
Changes live only in this process's memory until the next
|
||||
``index_done_callback``. Cross-process readers will not see
|
||||
them until that commit fires (see class docstring,
|
||||
*Cross-process sync protocol*). Until the flush, an upserted id
|
||||
is observable only through the read-your-writes read paths, not
|
||||
through ``query``.
|
||||
"""
|
||||
# logger.debug(f"[{self.workspace}] Buffering {len(data)} to {self.namespace}")
|
||||
if not data:
|
||||
return
|
||||
|
||||
current_time = int(time.time())
|
||||
pending = [
|
||||
(
|
||||
k,
|
||||
{
|
||||
"__id__": k,
|
||||
"__created_at__": current_time,
|
||||
**{k1: v1 for k1, v1 in v.items() if k1 in self.meta_fields},
|
||||
},
|
||||
)
|
||||
for k, v in data.items()
|
||||
]
|
||||
|
||||
# Buffer under the lock to interlock with the lock-held flush. A new
|
||||
# _PendingNanoDoc(vector=None) overwrites any prior pending doc for the
|
||||
# same id, discarding a temp vector a previous get_vectors_by_ids may
|
||||
# have cached (content-version change -> must re-embed new content).
|
||||
async with self._storage_lock:
|
||||
for doc_id, record in pending:
|
||||
self._pending_upserts[doc_id] = _PendingNanoDoc(record=record)
|
||||
|
||||
async def _flush_pending_locked(self) -> None:
|
||||
"""Embed pending docs and materialize them into ``self._client``.
|
||||
|
||||
Precondition: the caller **must already hold** ``_storage_lock``. The
|
||||
lock is non-reentrant, so this helper never calls ``_get_client`` and
|
||||
operates on ``self._client`` directly. Embedding runs inside the lock
|
||||
on purpose (see class docstring, *Deferred-embedding protocol*).
|
||||
|
||||
Failure handling: if embedding raises or the returned count does not
|
||||
match, the exception propagates and ``_pending_upserts`` is left intact
|
||||
so the next flush retries; nothing is written to ``self._client``.
|
||||
"""
|
||||
if not self._pending_upserts:
|
||||
return
|
||||
|
||||
# Snapshot for stable ordering between the embed list and the write.
|
||||
pending_items = list(self._pending_upserts.items())
|
||||
to_embed = [
|
||||
(doc_id, pdoc) for doc_id, pdoc in pending_items if pdoc.vector is None
|
||||
]
|
||||
|
||||
if to_embed:
|
||||
contents = [pdoc.record["content"] for _, pdoc in to_embed]
|
||||
batches = [
|
||||
contents[i : i + self._max_batch_size]
|
||||
for i in range(0, len(contents), self._max_batch_size)
|
||||
]
|
||||
logger.info(
|
||||
f"[{self.workspace}] {self.namespace} flush: embedding "
|
||||
f"{len(to_embed)} vectors in {len(batches)} batch(es) "
|
||||
f"(batch_num={self._max_batch_size})"
|
||||
)
|
||||
try:
|
||||
embeddings_list = await asyncio.gather(
|
||||
*[
|
||||
self.embedding_func(batch, context="document")
|
||||
for batch in batches
|
||||
]
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[{self.workspace}] Error embedding pending vector ops "
|
||||
f"(upserts={len(to_embed)}): {e}"
|
||||
)
|
||||
raise
|
||||
embeddings = np.concatenate(embeddings_list)
|
||||
if len(embeddings) != len(to_embed):
|
||||
# Explicit raise (not a log): a mismatch would mis-pair vectors
|
||||
# with records. Keep pending intact so the next flush retries.
|
||||
raise RuntimeError(
|
||||
f"[{self.workspace}] embedding is not 1-1 with pending data, "
|
||||
f"{len(embeddings)} != {len(to_embed)}"
|
||||
)
|
||||
for (_, pdoc), embedding in zip(to_embed, embeddings):
|
||||
pdoc.vector = embedding
|
||||
|
||||
list_data = []
|
||||
for _, pdoc in pending_items:
|
||||
vector = pdoc.vector
|
||||
# Compress vector using Float16 + zlib + Base64 for storage optimization
|
||||
vector_f16 = vector.astype(np.float16)
|
||||
compressed_vector = zlib.compress(vector_f16.tobytes())
|
||||
encoded_vector = base64.b64encode(compressed_vector).decode("utf-8")
|
||||
record = pdoc.record
|
||||
record["vector"] = encoded_vector
|
||||
record["__vector__"] = vector
|
||||
list_data.append(record)
|
||||
|
||||
self._client.upsert(datas=list_data)
|
||||
self._client_dirty = True
|
||||
|
||||
# Clear only the entries we just flushed (an upsert that arrived after
|
||||
# the snapshot would have re-set vector=None and must not be dropped).
|
||||
for doc_id, pdoc in pending_items:
|
||||
if self._pending_upserts.get(doc_id) is pdoc:
|
||||
del self._pending_upserts[doc_id]
|
||||
|
||||
def _save_to_disk_locked(self) -> None:
|
||||
"""Atomically persist ``self._client`` and notify other processes.
|
||||
|
||||
Precondition: the caller must already hold ``_storage_lock``. Factored
|
||||
out of ``index_done_callback`` so ``finalize`` reuses the exact same
|
||||
save+notify sequence. ``NanoVectorDB.save()`` always writes to whatever
|
||||
path is on the instance, so we temporarily redirect ``storage_file`` to
|
||||
the per-writer tmp and let ``atomic_write`` own the rename; the original
|
||||
path is restored on every path (success and exception).
|
||||
"""
|
||||
|
||||
def _save_atomic(tmp: str) -> None:
|
||||
original = self._client.storage_file
|
||||
self._client.storage_file = tmp
|
||||
try:
|
||||
self._client.save()
|
||||
finally:
|
||||
self._client.storage_file = original
|
||||
|
||||
atomic_write(self._client_file_name, _save_atomic, self.workspace or "_")
|
||||
|
||||
async def query(
|
||||
self, query: str, top_k: int, query_embedding: list[float] = None
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Similarity search over data already materialized into ``self._client``.
|
||||
|
||||
Buffered (unflushed) upserts are **not** searchable — only rows that a
|
||||
prior ``index_done_callback`` / ``finalize`` flushed are considered.
|
||||
Use the read-your-writes paths (``get_by_id`` / ``get_by_ids`` /
|
||||
``get_vectors_by_ids``) to observe pending data before a flush.
|
||||
"""
|
||||
# Use provided embedding or compute it
|
||||
if query_embedding is not None:
|
||||
embedding = query_embedding
|
||||
else:
|
||||
# Execute embedding outside of lock to avoid improve cocurrent
|
||||
embedding = await self.embedding_func(
|
||||
[query], context="query", _priority=DEFAULT_QUERY_PRIORITY
|
||||
) # higher priority for query
|
||||
embedding = embedding[0]
|
||||
|
||||
client = await self._get_client()
|
||||
results = client.query(
|
||||
query=embedding,
|
||||
top_k=top_k,
|
||||
better_than_threshold=self.cosine_better_than_threshold,
|
||||
)
|
||||
results = [
|
||||
{
|
||||
**{k: v for k, v in dp.items() if k != "vector"},
|
||||
"id": dp["__id__"],
|
||||
"distance": dp["__metrics__"],
|
||||
"created_at": dp.get("__created_at__"),
|
||||
}
|
||||
for dp in results
|
||||
]
|
||||
return results
|
||||
|
||||
@property
|
||||
async def client_storage(self):
|
||||
"""Return a **live reference** to ``NanoVectorDB.__storage``.
|
||||
|
||||
The returned dict is the same object NanoVectorDB mutates in
|
||||
place during ``upsert`` / ``delete``. Reading it outside
|
||||
``_storage_lock`` is safe today only because NanoVectorDB
|
||||
mutations are fully synchronous (see class docstring,
|
||||
*Lock scope*). Callers must not retain this reference across an
|
||||
``await`` that might cross into ``_get_client`` again: a reload
|
||||
will swap ``self._client`` for a fresh instance and leave the
|
||||
held reference pointing at the old (now-stale) storage.
|
||||
"""
|
||||
client = await self._get_client()
|
||||
return getattr(client, "_NanoVectorDB__storage")
|
||||
|
||||
async def delete(self, ids: list[str]):
|
||||
"""Delete vectors with specified IDs.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires a
|
||||
subsequent ``index_done_callback``. In ``lightrag.py`` this is
|
||||
handled by ``_insert_done()`` at the end of the document batch.
|
||||
Callers outside the pipeline must persist explicitly.
|
||||
|
||||
Args:
|
||||
ids: List of vector IDs to be deleted
|
||||
"""
|
||||
try:
|
||||
# Hold the lock so the pending-cancel and the client delete are a
|
||||
# single critical section against a concurrent flush. Operate on
|
||||
# self._client directly (the lock is non-reentrant; no _get_client).
|
||||
async with self._storage_lock:
|
||||
self._reload_client_from_disk_locked(for_write=True)
|
||||
|
||||
for doc_id in ids:
|
||||
self._pending_upserts.pop(doc_id, None)
|
||||
|
||||
# Record count before deletion
|
||||
before_count = len(self._client)
|
||||
|
||||
self._client.delete(ids)
|
||||
|
||||
# Calculate actual deleted count
|
||||
after_count = len(self._client)
|
||||
deleted_count = before_count - after_count
|
||||
if deleted_count:
|
||||
self._client_dirty = True
|
||||
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Successfully deleted {deleted_count} vectors from {self.namespace}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[{self.workspace}] Error while deleting vectors from {self.namespace}: {e}"
|
||||
)
|
||||
|
||||
async def delete_entity(self, entity_name: str) -> None:
|
||||
"""Delete the vector associated with a single entity name.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. Callers outside the
|
||||
pipeline must persist explicitly.
|
||||
|
||||
Buffer semantics — post-prune with caller short-circuit contract:
|
||||
The materialized client delete runs first; the matching
|
||||
pending upsert (if any) is popped **only after** it
|
||||
succeeds. If the materialized delete raises, the pending
|
||||
buffer stays intact and the exception is re-raised so the
|
||||
caller can short-circuit before ``index_done_callback``
|
||||
flushes a half-cleaned buffer.
|
||||
|
||||
**Not pipeline-gated** — see class docstring
|
||||
*Non-pipeline write paths*. The caller is responsible for
|
||||
ensuring single-writer serialization.
|
||||
"""
|
||||
try:
|
||||
entity_id = compute_mdhash_id(entity_name, prefix="ent-")
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Attempting to delete entity {entity_name} with ID {entity_id}"
|
||||
)
|
||||
|
||||
async with self._storage_lock:
|
||||
self._reload_client_from_disk_locked(for_write=True)
|
||||
|
||||
# Materialized side first so a failure leaves the
|
||||
# pending buffer intact for the caller's retry path.
|
||||
if self._client.get([entity_id]):
|
||||
self._client.delete([entity_id])
|
||||
self._client_dirty = True
|
||||
deleted = True
|
||||
else:
|
||||
deleted = False
|
||||
|
||||
# Materialized delete succeeded — safe to cancel any
|
||||
# buffered upsert for this entity.
|
||||
pending_cancelled = (
|
||||
self._pending_upserts.pop(entity_id, None) is not None
|
||||
)
|
||||
|
||||
if deleted or pending_cancelled:
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Successfully deleted entity {entity_name}"
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Entity {entity_name} not found in storage"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"[{self.workspace}] Error deleting entity {entity_name}: {e}")
|
||||
raise
|
||||
|
||||
async def delete_entity_relation(self, entity_name: str) -> None:
|
||||
"""Delete every relation vector incident to ``entity_name``.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. Callers outside the
|
||||
pipeline must persist explicitly.
|
||||
|
||||
Buffer semantics — post-prune with caller short-circuit contract:
|
||||
The materialized client delete runs first; matching pending
|
||||
upserts are pruned **only after** it succeeds. If the
|
||||
materialized delete raises, the pending buffer stays intact
|
||||
and the exception is re-raised so the caller (e.g.
|
||||
``adelete_by_entity``) can short-circuit before
|
||||
``_persist_graph_updates`` triggers ``index_done_callback``
|
||||
on a half-cleaned buffer.
|
||||
|
||||
Previously the buffer was pre-pruned and the outer
|
||||
``except`` swallowed exceptions into ``logger.error`` — that
|
||||
combination silently dropped both buffered relation vectors
|
||||
and the failure signal.
|
||||
|
||||
**Not pipeline-gated** — see class docstring
|
||||
*Non-pipeline write paths*. The caller is responsible for
|
||||
ensuring single-writer serialization.
|
||||
"""
|
||||
try:
|
||||
async with self._storage_lock:
|
||||
self._reload_client_from_disk_locked(for_write=True)
|
||||
|
||||
# Materialized side first so a failure leaves the
|
||||
# pending buffer intact for the caller's retry path.
|
||||
# Use .get() for src_id / tgt_id so rows from foreign
|
||||
# namespaces without those keys silently don't match.
|
||||
storage = getattr(self._client, "_NanoVectorDB__storage")
|
||||
ids_to_delete = [
|
||||
dp["__id__"]
|
||||
for dp in storage["data"]
|
||||
if dp.get("src_id") == entity_name
|
||||
or dp.get("tgt_id") == entity_name
|
||||
]
|
||||
if ids_to_delete:
|
||||
self._client.delete(ids_to_delete)
|
||||
self._client_dirty = True
|
||||
|
||||
# Materialized delete succeeded — safe to prune matching
|
||||
# buffered upserts so a subsequent flush won't re-upsert
|
||||
# the just-deleted relations.
|
||||
pending_ids = [
|
||||
doc_id
|
||||
for doc_id, pdoc in self._pending_upserts.items()
|
||||
if pdoc.record.get("src_id") == entity_name
|
||||
or pdoc.record.get("tgt_id") == entity_name
|
||||
]
|
||||
for doc_id in pending_ids:
|
||||
del self._pending_upserts[doc_id]
|
||||
|
||||
total = len(pending_ids) + len(ids_to_delete)
|
||||
if total:
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Deleted {total} relations for {entity_name}"
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
f"[{self.workspace}] No relations found for entity {entity_name}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[{self.workspace}] Error deleting relations for {entity_name}: {e}"
|
||||
)
|
||||
raise
|
||||
|
||||
async def drop_pending_index_ops(self) -> None:
|
||||
"""Discard buffered upserts on an aborting batch.
|
||||
|
||||
Only the pending buffer is dropped; records already materialized into
|
||||
``self._client`` by a prior ``_flush_pending_locked`` whose save step
|
||||
then failed (``_client_dirty=True``) are intentionally NOT rolled back.
|
||||
|
||||
The pipeline treats each file as an atomic unit: an abort marks the
|
||||
affected documents FAILED and the whole file is reprocessed on the
|
||||
next run. Because upserts are keyed by deterministic ids (entity-name
|
||||
/ relation / chunk hashes), reprocessing overwrites those vectors
|
||||
idempotently, so the final state is identical whether or not we roll
|
||||
back here. This matches the server-backed backends (Milvus / OpenSearch
|
||||
/ Postgres / Mongo / Qdrant), which likewise keep a sibling flush's
|
||||
already-committed partial data on abort rather than rolling it back;
|
||||
and if the process crashes before the next save, these in-memory
|
||||
writes are dropped anyway. Rolling back only FAISS/Nano would add an
|
||||
inconsistent, non-load-bearing "FAILED == clean" guarantee, so it is
|
||||
deliberately omitted.
|
||||
"""
|
||||
if self._storage_lock is None:
|
||||
self._pending_upserts.clear()
|
||||
return
|
||||
async with self._storage_lock:
|
||||
self._pending_upserts.clear()
|
||||
|
||||
async def index_done_callback(self) -> bool:
|
||||
"""Flush deferred embeddings, commit to disk, and notify other processes.
|
||||
|
||||
This is the writer's **commit point** in the cross-process sync
|
||||
protocol (see class docstring). Effects, in order:
|
||||
1. If another process committed first, reload the latest on-disk
|
||||
snapshot while preserving this process's pending buffer.
|
||||
2. ``_flush_pending_locked`` embeds every buffered upsert (once
|
||||
per id) and materializes it into ``self._client``. A failure
|
||||
here **raises** — pending is kept, nothing is written.
|
||||
3. ``_save_to_disk_locked`` (``atomic_write``) lays a tmp file
|
||||
beside the target and renames it into place — readers either
|
||||
see the previous file in full or the new file in full, never a
|
||||
torn write. A failure here **also raises**; ``_client_dirty``
|
||||
stays ``True`` so a later ``finalize`` retries the save.
|
||||
4. ``set_all_update_flags`` flips every registered process's
|
||||
``storage_updated`` flag, then we immediately reset our own
|
||||
flag to ``False`` so the writer does not self-reload on the
|
||||
next call to ``_get_client``.
|
||||
|
||||
Either failure surfaces loudly through ``_insert_done`` so the caller
|
||||
can abort the document batch instead of silently losing vectors. The
|
||||
bool return is kept for legacy callers but is effectively always
|
||||
``True`` on the success path.
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
self._reload_client_from_disk_locked(for_write=True)
|
||||
|
||||
# Flush + save both raise on failure (embedding mismatch / save IO
|
||||
# error). The exception propagates out of the lock so _insert_done
|
||||
# aborts the batch; pending stays intact and _client_dirty stays
|
||||
# True (if only the save failed) for a later retry.
|
||||
await self._flush_pending_locked()
|
||||
self._save_to_disk_locked()
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
self.storage_updated.value = False
|
||||
self._client_dirty = False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _format_record(dp: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Shape a stored/pending record into the public read result."""
|
||||
return {
|
||||
**{k: v for k, v in dp.items() if k not in ("vector", "__vector__")},
|
||||
"id": dp.get("__id__"),
|
||||
"created_at": dp.get("__created_at__"),
|
||||
}
|
||||
|
||||
async def get_by_id(self, id: str) -> dict[str, Any] | None:
|
||||
"""Get vector data by its ID (read-your-writes against the pending buffer).
|
||||
|
||||
Args:
|
||||
id: The unique identifier of the vector
|
||||
|
||||
Returns:
|
||||
The vector data if found, or None if not found
|
||||
"""
|
||||
# Read-your-writes: a buffered upsert is visible before its flush.
|
||||
async with self._storage_lock:
|
||||
pending = self._pending_upserts.get(id)
|
||||
if pending is not None:
|
||||
return self._format_record(pending.record)
|
||||
|
||||
client = await self._get_client()
|
||||
result = client.get([id])
|
||||
if result:
|
||||
return self._format_record(result[0])
|
||||
return None
|
||||
|
||||
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
||||
"""Get multiple vector data by their IDs (read-your-writes), preserving order.
|
||||
|
||||
Args:
|
||||
ids: List of unique identifiers
|
||||
|
||||
Returns:
|
||||
List of vector data objects that were found
|
||||
"""
|
||||
if not ids:
|
||||
return []
|
||||
|
||||
# Read-your-writes: serve buffered upserts from the pending area and
|
||||
# only query the materialized client for the remaining ids.
|
||||
result_map: dict[str, dict[str, Any]] = {}
|
||||
remaining: list[str] = []
|
||||
async with self._storage_lock:
|
||||
for requested_id in ids:
|
||||
pending = self._pending_upserts.get(requested_id)
|
||||
if pending is not None:
|
||||
result_map[str(requested_id)] = self._format_record(pending.record)
|
||||
else:
|
||||
remaining.append(requested_id)
|
||||
|
||||
if remaining:
|
||||
client = await self._get_client()
|
||||
for dp in client.get(remaining):
|
||||
if not dp:
|
||||
continue
|
||||
record = self._format_record(dp)
|
||||
key = record.get("id")
|
||||
if key is not None:
|
||||
result_map[str(key)] = record
|
||||
|
||||
return [result_map.get(str(requested_id)) for requested_id in ids]
|
||||
|
||||
async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]:
|
||||
"""Get vectors by their IDs (read-your-writes), returning only ID and vector.
|
||||
|
||||
For buffered upserts the vector is computed lazily (and cached back onto
|
||||
the pending doc so the next flush reuses it instead of re-embedding);
|
||||
for materialized rows the stored compressed vector is decoded.
|
||||
|
||||
Args:
|
||||
ids: List of unique identifiers
|
||||
|
||||
Returns:
|
||||
Dictionary mapping IDs to their vector embeddings
|
||||
Format: {id: [vector_values], ...}
|
||||
"""
|
||||
if not ids:
|
||||
return {}
|
||||
|
||||
vectors_dict: dict[str, list[float]] = {}
|
||||
remaining: list[str] = []
|
||||
async with self._storage_lock:
|
||||
to_embed: list[tuple[str, _PendingNanoDoc]] = []
|
||||
for requested_id in ids:
|
||||
pending = self._pending_upserts.get(requested_id)
|
||||
if pending is None:
|
||||
remaining.append(requested_id)
|
||||
elif pending.vector is not None:
|
||||
vectors_dict[requested_id] = pending.vector.astype(
|
||||
np.float32
|
||||
).tolist()
|
||||
else:
|
||||
to_embed.append((requested_id, pending))
|
||||
|
||||
if to_embed:
|
||||
contents = [pdoc.record["content"] for _, pdoc in to_embed]
|
||||
batches = [
|
||||
contents[i : i + self._max_batch_size]
|
||||
for i in range(0, len(contents), self._max_batch_size)
|
||||
]
|
||||
embeddings_list = await asyncio.gather(
|
||||
*[
|
||||
self.embedding_func(batch, context="document")
|
||||
for batch in batches
|
||||
]
|
||||
)
|
||||
embeddings = np.concatenate(embeddings_list)
|
||||
if len(embeddings) != len(to_embed):
|
||||
raise RuntimeError(
|
||||
f"[{self.workspace}] embedding is not 1-1 with pending data, "
|
||||
f"{len(embeddings)} != {len(to_embed)}"
|
||||
)
|
||||
for (requested_id, pdoc), embedding in zip(to_embed, embeddings):
|
||||
# Cache the vector back so the next flush reuses it.
|
||||
pdoc.vector = embedding
|
||||
vectors_dict[requested_id] = embedding.astype(np.float32).tolist()
|
||||
|
||||
if remaining:
|
||||
client = await self._get_client()
|
||||
for result in client.get(remaining):
|
||||
if result and "vector" in result and "__id__" in result:
|
||||
# Decompress vector data (Base64 + zlib + Float16 compressed)
|
||||
decoded = base64.b64decode(result["vector"])
|
||||
decompressed = zlib.decompress(decoded)
|
||||
vector_f16 = np.frombuffer(decompressed, dtype=np.float16)
|
||||
vector_f32 = vector_f16.astype(np.float32).tolist()
|
||||
vectors_dict[result["__id__"]] = vector_f32
|
||||
|
||||
return vectors_dict
|
||||
|
||||
async def drop(self) -> dict[str, str]:
|
||||
"""Drop all vector data from storage and reinitialize the client.
|
||||
|
||||
This method will:
|
||||
1. Remove the vector database storage file if it exists
|
||||
2. Reinitialize the vector database client
|
||||
3. Update flags to notify other processes
|
||||
4. Changes are persisted to disk immediately
|
||||
|
||||
Caller contract:
|
||||
``drop`` is destructive and **not** serialized by this storage
|
||||
class. The caller must hold the pipeline ``busy`` reservation
|
||||
(the ``/documents/clear`` endpoint does this) before invoking
|
||||
it — running ``drop`` concurrently with an active document
|
||||
pipeline will tear down storage out from under the writer and
|
||||
silently lose data. See class docstring,
|
||||
*Non-pipeline write paths*.
|
||||
|
||||
Returns:
|
||||
dict[str, str]: Operation status and message
|
||||
- On success: {"status": "success", "message": "data dropped"}
|
||||
- On failure: {"status": "error", "message": "<error details>"}
|
||||
"""
|
||||
try:
|
||||
async with self._storage_lock:
|
||||
# Discard buffered (unflushed) upserts along with the data.
|
||||
self._pending_upserts.clear()
|
||||
|
||||
# delete _client_file_name
|
||||
if os.path.exists(self._client_file_name):
|
||||
os.remove(self._client_file_name)
|
||||
|
||||
self._client = NanoVectorDB(
|
||||
self.embedding_func.embedding_dim,
|
||||
storage_file=self._client_file_name,
|
||||
)
|
||||
self._client_dirty = False
|
||||
|
||||
# Notify other processes that data has been updated
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
# Reset own update flag to avoid self-reloading
|
||||
self.storage_updated.value = False
|
||||
|
||||
logger.info(
|
||||
f"[{self.workspace}] Process {os.getpid()} drop {self.namespace}(file:{self._client_file_name})"
|
||||
)
|
||||
return {"status": "success", "message": "data dropped"}
|
||||
except Exception as e:
|
||||
logger.error(f"[{self.workspace}] Error dropping {self.namespace}: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
async def finalize(self):
|
||||
"""Flush any buffered upserts and persist before shutdown (safety net).
|
||||
|
||||
Normally ``index_done_callback`` has already drained the pending buffer
|
||||
and synced to disk, but two paths land here with work to do:
|
||||
|
||||
- **Pending upserts only** (no prior ``index_done_callback``): flush
|
||||
and save. We reload first so a stale process picks up other writers'
|
||||
commits before merging its pending buffer in.
|
||||
- **Unsaved materialized changes** (``_client_dirty=True``): an earlier
|
||||
``index_done_callback`` flushed pending into ``self._client`` but
|
||||
its save raised. Skip the reload — reloading would drop those
|
||||
materialized-but-unsaved rows — and just retry the save.
|
||||
|
||||
Flush / save failures propagate (same contract as
|
||||
``index_done_callback``); a partially flushed buffer is preserved for
|
||||
a future retry.
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
if not self._pending_upserts and not self._client_dirty:
|
||||
return
|
||||
if self._pending_upserts:
|
||||
# Only reload when we have nothing un-persisted in self._client.
|
||||
# A dirty client carries successfully-flushed-but-unsaved rows
|
||||
# from a prior index_done_callback; reloading would silently
|
||||
# drop them.
|
||||
if not self._client_dirty:
|
||||
self._reload_client_from_disk_locked(for_write=True)
|
||||
await self._flush_pending_locked()
|
||||
self._save_to_disk_locked()
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
self.storage_updated.value = False
|
||||
self._client_dirty = False
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,796 @@
|
||||
import os
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import final
|
||||
|
||||
from lightrag.file_atomic import atomic_write, reap_orphan_tmp_files
|
||||
from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
|
||||
from lightrag.utils import logger, validate_workspace
|
||||
from lightrag.base import BaseGraphStorage
|
||||
import networkx as nx
|
||||
from .shared_storage import (
|
||||
get_namespace_lock,
|
||||
get_update_flag,
|
||||
set_all_update_flags,
|
||||
)
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# 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)
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
class NetworkXStorage(BaseGraphStorage):
|
||||
"""File-backed knowledge-graph storage built on ``networkx.Graph``.
|
||||
|
||||
Storage model:
|
||||
A single ``networkx.Graph`` instance lives in process memory; its
|
||||
full state is serialized to one GraphML file at
|
||||
``working_dir/[workspace/]graph_<namespace>.graphml``. That GraphML
|
||||
file is the **only** cross-process synchronization surface — there
|
||||
is no shared memory, no message bus, and no network channel
|
||||
between processes. Cross-process visibility is mediated by (a) an
|
||||
atomic file write at commit time and (b) a per-namespace
|
||||
``storage_updated`` flag distributed through
|
||||
``lightrag.kg.shared_storage``.
|
||||
|
||||
Concurrency invariants (the code in this file is correct *only* while
|
||||
all three hold):
|
||||
1. **Single writer per workspace.** The document pipeline's
|
||||
``busy`` / ``destructive_busy`` flags (see ``AGENTS.md``
|
||||
*Pipeline concurrency contract*) guarantee at most one process
|
||||
performs ``upsert_*`` / ``delete_*`` / ``remove_*`` /
|
||||
``index_done_callback`` at any time. Every other process is
|
||||
read-only.
|
||||
2. **Eventual consistency is sufficient.** Read-only processes
|
||||
only need to observe the writer's data *after* the writer's
|
||||
``index_done_callback`` completes. Reads landing in the gap
|
||||
between a writer's in-memory mutation and its commit may
|
||||
legitimately return the pre-update snapshot.
|
||||
3. **networkx operations are fully synchronous.** Under a
|
||||
single-threaded asyncio event loop, ``graph.add_node`` /
|
||||
``graph.remove_node`` / ``graph.degree`` / etc. cannot be
|
||||
preempted by another coroutine, which gives them implicit
|
||||
mutual exclusion over ``self._graph``. This is why the methods
|
||||
below don't have to hold ``_storage_lock`` while calling into
|
||||
``graph``.
|
||||
|
||||
Cross-process sync protocol (identical in shape to
|
||||
``NanoVectorDBStorage`` — see that class's docstring for the canonical
|
||||
description):
|
||||
Writer side (``index_done_callback``):
|
||||
1. ``write_nx_graph`` atomically writes the GraphML file
|
||||
(``atomic_write`` lays a tmp file beside the target and
|
||||
renames it into place — readers either see the previous
|
||||
file in full or the new file in full, never a torn write).
|
||||
2. ``set_all_update_flags`` flips every process's
|
||||
``storage_updated`` flag (including the writer's own).
|
||||
3. Immediately reset the writer's own flag to ``False`` so
|
||||
the next call to ``_get_graph`` does not trigger a
|
||||
self-reload of the data this process just wrote.
|
||||
Reader side (any method that goes through ``_get_graph``):
|
||||
1. Inside ``_storage_lock``, observe
|
||||
``storage_updated.value is True``.
|
||||
2. **Fully reload** ``self._graph`` from disk via
|
||||
``load_nx_graph``. networkx GraphML has no incremental
|
||||
sync API, so the entire file is re-parsed.
|
||||
3. Reset the reader's own flag.
|
||||
|
||||
Lock scope:
|
||||
``_storage_lock`` is a per-``(namespace, workspace)`` keyed lock
|
||||
spanning both intra-process coroutines and inter-process workers.
|
||||
It wraps only the *reload* and *commit* critical sections, not
|
||||
every ``graph.xxx`` call. Operating on ``graph`` outside the lock
|
||||
is safe today *because of invariant (3)* — if either premise is
|
||||
ever broken (e.g. ``graph.xxx`` is moved to a thread pool, or
|
||||
networkx is swapped for an async graph library), the lock scope
|
||||
must be widened to cover the mutation/read itself.
|
||||
|
||||
Implementation differences from ``NanoVectorDBStorage`` (same design,
|
||||
different surface):
|
||||
* No ``client_storage`` property — there is no equivalent live
|
||||
reference being exposed to callers, so NanoVectorDB's
|
||||
"do-not-retain-across-await" caveat does not apply here.
|
||||
* ``write_nx_graph`` passes the tmp path directly to
|
||||
``nx.write_graphml``, so the writer needs no equivalent of
|
||||
NanoVectorDB's "temporarily reassign ``storage_file``" trick.
|
||||
* Mutation surface is finer-grained (``upsert_node`` /
|
||||
``upsert_edge`` / ``upsert_nodes_batch`` /
|
||||
``upsert_edges_batch`` / ``delete_node`` / ``remove_nodes`` /
|
||||
``remove_edges``); each goes through ``_get_graph`` once and
|
||||
then operates synchronously on ``self._graph``.
|
||||
|
||||
Non-pipeline write paths:
|
||||
The pipeline's ``busy`` gate serializes mutation calls reached
|
||||
through the document ingestion and purge flows. The following
|
||||
entry points are **not** serialized by the pipeline gate and
|
||||
must be guarded externally:
|
||||
* ``drop`` — currently gated by the API layer (the
|
||||
``/documents/clear`` endpoint takes the pipeline busy
|
||||
reservation before invoking it).
|
||||
* ``delete_node`` / ``remove_nodes`` / ``remove_edges`` /
|
||||
``upsert_node`` / ``upsert_edge`` when invoked from
|
||||
``utils_graph.py`` admin flows (``adelete_by_entity`` /
|
||||
``adelete_by_relation`` / entity-edit flows). These flows
|
||||
are currently not exposed in the WebUI; any future caller
|
||||
must arrange single-writer serialization the same way the
|
||||
pipeline does.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def load_nx_graph(file_name) -> nx.Graph:
|
||||
if os.path.exists(file_name):
|
||||
return nx.read_graphml(file_name)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def write_nx_graph(graph: nx.Graph, file_name, workspace="_"):
|
||||
logger.info(
|
||||
f"[{workspace}] Writing graph with {graph.number_of_nodes()} nodes, {graph.number_of_edges()} edges"
|
||||
)
|
||||
atomic_write(
|
||||
file_name,
|
||||
lambda tmp: nx.write_graphml(graph, tmp),
|
||||
workspace,
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
# Reject path traversal before using workspace in a file path
|
||||
validate_workspace(self.workspace)
|
||||
working_dir = self.global_config["working_dir"]
|
||||
if self.workspace:
|
||||
# Include workspace in the file path for data isolation
|
||||
workspace_dir = os.path.join(working_dir, self.workspace)
|
||||
else:
|
||||
# Default behavior when workspace is empty
|
||||
workspace_dir = working_dir
|
||||
self.workspace = ""
|
||||
|
||||
os.makedirs(workspace_dir, exist_ok=True)
|
||||
self._graphml_xml_file = os.path.join(
|
||||
workspace_dir, f"graph_{self.namespace}.graphml"
|
||||
)
|
||||
self._storage_lock = None
|
||||
self.storage_updated = None
|
||||
self._graph = None
|
||||
|
||||
reap_orphan_tmp_files(self._graphml_xml_file, workspace=self.workspace or "_")
|
||||
|
||||
# Load initial graph
|
||||
preloaded_graph = NetworkXStorage.load_nx_graph(self._graphml_xml_file)
|
||||
if preloaded_graph is not None:
|
||||
logger.info(
|
||||
f"[{self.workspace}] Loaded graph from {self._graphml_xml_file} with {preloaded_graph.number_of_nodes()} nodes, {preloaded_graph.number_of_edges()} edges"
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"[{self.workspace}] Created new empty graph file: {self._graphml_xml_file}"
|
||||
)
|
||||
self._graph = preloaded_graph or nx.Graph()
|
||||
|
||||
async def initialize(self):
|
||||
"""Initialize storage data"""
|
||||
# Get the update flag for cross-process update notification
|
||||
self.storage_updated = await get_update_flag(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
# Get the storage lock for use in other methods
|
||||
self._storage_lock = get_namespace_lock(
|
||||
self.namespace, workspace=self.workspace
|
||||
)
|
||||
|
||||
async def _get_graph(self):
|
||||
"""Return the live ``networkx.Graph``, reloading from disk if needed.
|
||||
|
||||
This is the **single entry point** every public method funnels
|
||||
through to obtain ``self._graph``. It is also the **only place
|
||||
readers transition to a fresher on-disk snapshot**: when another
|
||||
process has committed (via ``index_done_callback``) and flipped
|
||||
this process's ``storage_updated`` flag, the next call here
|
||||
rebuilds ``self._graph`` by re-parsing the entire GraphML file.
|
||||
networkx has no incremental sync API — the reload is
|
||||
unconditionally a full file reload.
|
||||
|
||||
Under the *Single writer* invariant (see class docstring), the
|
||||
reload branch never fires in the writer process: the writer
|
||||
resets its own flag at the end of every ``index_done_callback``.
|
||||
The branch exists for readers.
|
||||
|
||||
``_storage_lock`` is held during the check-and-reload to (a)
|
||||
serialize concurrent reload attempts by sibling coroutines in
|
||||
the same process and (b) interlock with ``index_done_callback``
|
||||
so a reader cannot observe a partially-saved file.
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
# Check if data needs to be reloaded
|
||||
if self.storage_updated.value:
|
||||
logger.info(
|
||||
f"[{self.workspace}] Process {os.getpid()} reloading graph {self._graphml_xml_file} due to modifications by another process"
|
||||
)
|
||||
# Reload data
|
||||
self._graph = (
|
||||
NetworkXStorage.load_nx_graph(self._graphml_xml_file) or nx.Graph()
|
||||
)
|
||||
# Reset update flag
|
||||
self.storage_updated.value = False
|
||||
|
||||
return self._graph
|
||||
|
||||
async def has_node(self, node_id: str) -> bool:
|
||||
graph = await self._get_graph()
|
||||
return graph.has_node(node_id)
|
||||
|
||||
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
|
||||
graph = await self._get_graph()
|
||||
return graph.has_edge(source_node_id, target_node_id)
|
||||
|
||||
async def get_node(self, node_id: str) -> dict[str, str] | None:
|
||||
graph = await self._get_graph()
|
||||
return graph.nodes.get(node_id)
|
||||
|
||||
async def node_degree(self, node_id: str) -> int:
|
||||
graph = await self._get_graph()
|
||||
if graph.has_node(node_id):
|
||||
return graph.degree(node_id)
|
||||
return 0
|
||||
|
||||
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
|
||||
graph = await self._get_graph()
|
||||
src_degree = graph.degree(src_id) if graph.has_node(src_id) else 0
|
||||
tgt_degree = graph.degree(tgt_id) if graph.has_node(tgt_id) else 0
|
||||
return src_degree + tgt_degree
|
||||
|
||||
async def get_edge(
|
||||
self, source_node_id: str, target_node_id: str
|
||||
) -> dict[str, str] | None:
|
||||
graph = await self._get_graph()
|
||||
return graph.edges.get((source_node_id, target_node_id))
|
||||
|
||||
async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
|
||||
graph = await self._get_graph()
|
||||
if graph.has_node(source_node_id):
|
||||
return list(graph.edges(source_node_id))
|
||||
return None
|
||||
|
||||
async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
|
||||
"""Insert or update a single node; persistence is deferred.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. In ``lightrag.py`` this
|
||||
is handled by ``_insert_done()`` at the end of the document
|
||||
batch. Callers outside the pipeline must persist explicitly.
|
||||
|
||||
Correctness relies on the class docstring *Lock scope* invariant
|
||||
(synchronous networkx ops + single-writer pipeline gate).
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
graph.add_node(node_id, **node_data)
|
||||
|
||||
async def upsert_edge(
|
||||
self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
|
||||
) -> None:
|
||||
"""Insert or update a single edge; persistence is deferred.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. Callers outside the
|
||||
pipeline must persist explicitly.
|
||||
|
||||
Correctness relies on the class docstring *Lock scope* invariant.
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
graph.add_edge(source_node_id, target_node_id, **edge_data)
|
||||
|
||||
async def upsert_nodes_batch(self, nodes: list[tuple[str, dict[str, str]]]) -> None:
|
||||
"""Batch insert/update multiple nodes in a single call.
|
||||
|
||||
Much faster than calling upsert_node() in a loop for large imports
|
||||
because it avoids per-call async event loop overhead.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. Callers outside the
|
||||
pipeline must persist explicitly.
|
||||
|
||||
Args:
|
||||
nodes: List of (node_id, node_data) tuples.
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
for node_id, node_data in nodes:
|
||||
graph.add_node(node_id, **node_data)
|
||||
|
||||
async def has_nodes_batch(self, node_ids: list[str]) -> set[str]:
|
||||
"""Check existence of multiple nodes in a single call.
|
||||
|
||||
Returns:
|
||||
Set of node_ids that exist in the graph.
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
return {nid for nid in node_ids if graph.has_node(nid)}
|
||||
|
||||
async def upsert_edges_batch(
|
||||
self, edges: list[tuple[str, str, dict[str, str]]]
|
||||
) -> None:
|
||||
"""Batch insert/update multiple edges in a single call.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. Callers outside the
|
||||
pipeline must persist explicitly.
|
||||
|
||||
Args:
|
||||
edges: List of (source_id, target_id, edge_data) tuples.
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
for src, tgt, edge_data in edges:
|
||||
graph.add_edge(src, tgt, **edge_data)
|
||||
|
||||
async def delete_node(self, node_id: str) -> None:
|
||||
"""Remove a single node from the graph; persistence is deferred.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. Callers outside the
|
||||
pipeline must persist explicitly.
|
||||
|
||||
Pipeline-gating depends on the caller: invocations from the
|
||||
document purge flow are serialized by ``pipeline busy``;
|
||||
invocations from ``utils_graph.py`` admin flows are **not** —
|
||||
see class docstring *Non-pipeline write paths*.
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
if graph.has_node(node_id):
|
||||
graph.remove_node(node_id)
|
||||
logger.debug(f"[{self.workspace}] Node {node_id} deleted from the graph")
|
||||
else:
|
||||
logger.warning(
|
||||
f"[{self.workspace}] Node {node_id} not found in the graph for deletion"
|
||||
)
|
||||
|
||||
async def remove_nodes(self, nodes: list[str]):
|
||||
"""Delete multiple nodes from the graph.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. Callers outside the
|
||||
pipeline must persist explicitly.
|
||||
|
||||
Pipeline-gating depends on the caller — see ``delete_node`` and
|
||||
class docstring *Non-pipeline write paths*.
|
||||
|
||||
Args:
|
||||
nodes: List of node IDs to be deleted
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
for node in nodes:
|
||||
if graph.has_node(node):
|
||||
graph.remove_node(node)
|
||||
|
||||
async def remove_edges(self, edges: list[tuple[str, str]]):
|
||||
"""Delete multiple edges from the graph.
|
||||
|
||||
Persistence:
|
||||
Changes are in-memory only; cross-process visibility requires
|
||||
a subsequent ``index_done_callback``. Callers outside the
|
||||
pipeline must persist explicitly.
|
||||
|
||||
Pipeline-gating depends on the caller — see ``delete_node`` and
|
||||
class docstring *Non-pipeline write paths*.
|
||||
|
||||
Args:
|
||||
edges: List of edges to be deleted, each edge is a (source, target) tuple
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
for source, target in edges:
|
||||
if graph.has_edge(source, target):
|
||||
graph.remove_edge(source, target)
|
||||
|
||||
async def get_all_labels(self) -> list[str]:
|
||||
"""
|
||||
Get all node labels(entity names) in the graph
|
||||
Returns:
|
||||
[label1, label2, ...] # Alphabetically sorted label list
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
labels = set()
|
||||
for node in graph.nodes():
|
||||
labels.add(str(node)) # Add node id as a label
|
||||
|
||||
# Return sorted list
|
||||
return sorted(list(labels))
|
||||
|
||||
async def get_popular_labels(self, limit: int = 300) -> list[str]:
|
||||
"""
|
||||
Get popular labels(entity names) by node degree (most connected entities)
|
||||
|
||||
Args:
|
||||
limit: Maximum number of labels to return
|
||||
|
||||
Returns:
|
||||
List of labels sorted by degree (highest first)
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
|
||||
# Get degrees of all nodes and sort by degree descending
|
||||
degrees = dict(graph.degree())
|
||||
sorted_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Return top labels limited by the specified limit
|
||||
popular_labels = [str(node) for node, _ in sorted_nodes[:limit]]
|
||||
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Retrieved {len(popular_labels)} popular labels (limit: {limit})"
|
||||
)
|
||||
|
||||
return popular_labels
|
||||
|
||||
async def search_labels(self, query: str, limit: int = 50) -> list[str]:
|
||||
"""
|
||||
Search labels(entity names) with fuzzy matching
|
||||
|
||||
Args:
|
||||
query: Search query string
|
||||
limit: Maximum number of results to return
|
||||
|
||||
Returns:
|
||||
List of matching labels sorted by relevance
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
query_lower = query.lower().strip()
|
||||
|
||||
if not query_lower:
|
||||
return []
|
||||
|
||||
# Collect matching nodes with relevance scores
|
||||
matches = []
|
||||
for node in graph.nodes():
|
||||
node_str = str(node)
|
||||
node_lower = node_str.lower()
|
||||
|
||||
# Skip if no match
|
||||
if query_lower not in node_lower:
|
||||
continue
|
||||
|
||||
# Calculate relevance score
|
||||
# Exact match gets highest score
|
||||
if node_lower == query_lower:
|
||||
score = 1000
|
||||
# Prefix match gets high score
|
||||
elif node_lower.startswith(query_lower):
|
||||
score = 500
|
||||
# Contains match gets base score, with bonus for shorter strings
|
||||
else:
|
||||
# Shorter strings with matches are more relevant
|
||||
score = 100 - len(node_str)
|
||||
# Bonus for word boundary matches
|
||||
if f" {query_lower}" in node_lower or f"_{query_lower}" in node_lower:
|
||||
score += 50
|
||||
|
||||
matches.append((node_str, score))
|
||||
|
||||
# Sort by relevance score (desc) then alphabetically
|
||||
matches.sort(key=lambda x: (-x[1], x[0]))
|
||||
|
||||
# Return top matches limited by the specified limit
|
||||
search_results = [match[0] for match in matches[:limit]]
|
||||
|
||||
logger.debug(
|
||||
f"[{self.workspace}] Search query '{query}' returned {len(search_results)} results (limit: {limit})"
|
||||
)
|
||||
|
||||
return search_results
|
||||
|
||||
async def get_knowledge_graph(
|
||||
self,
|
||||
node_label: str,
|
||||
max_depth: int = 3,
|
||||
max_nodes: int = None,
|
||||
) -> KnowledgeGraph:
|
||||
"""
|
||||
Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
|
||||
|
||||
Args:
|
||||
node_label: Label of the starting node，* means all nodes
|
||||
max_depth: Maximum depth of the subgraph, Defaults to 3
|
||||
max_nodes: Maxiumu nodes to return by BFS, Defaults to 1000
|
||||
|
||||
Returns:
|
||||
KnowledgeGraph object containing nodes and edges, with an is_truncated flag
|
||||
indicating whether the graph was truncated due to max_nodes limit
|
||||
"""
|
||||
# Get max_nodes from global_config if not provided
|
||||
if max_nodes is None:
|
||||
max_nodes = self.global_config.get("max_graph_nodes", 1000)
|
||||
else:
|
||||
# Limit max_nodes to not exceed global_config max_graph_nodes
|
||||
max_nodes = min(max_nodes, self.global_config.get("max_graph_nodes", 1000))
|
||||
|
||||
graph = await self._get_graph()
|
||||
|
||||
result = KnowledgeGraph()
|
||||
|
||||
# Handle special case for "*" label
|
||||
if node_label == "*":
|
||||
# Get degrees of all nodes
|
||||
degrees = dict(graph.degree())
|
||||
# Sort nodes by degree in descending order and take top max_nodes
|
||||
sorted_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Check if graph is truncated
|
||||
if len(sorted_nodes) > max_nodes:
|
||||
result.is_truncated = True
|
||||
logger.info(
|
||||
f"[{self.workspace}] Graph truncated: {len(sorted_nodes)} nodes found, limited to {max_nodes}"
|
||||
)
|
||||
|
||||
limited_nodes = [node for node, _ in sorted_nodes[:max_nodes]]
|
||||
# Create subgraph with the highest degree nodes
|
||||
subgraph = graph.subgraph(limited_nodes)
|
||||
else:
|
||||
# Check if node exists
|
||||
if node_label not in graph:
|
||||
logger.warning(
|
||||
f"[{self.workspace}] Node {node_label} not found in the graph"
|
||||
)
|
||||
return KnowledgeGraph() # Return empty graph
|
||||
|
||||
# Use modified BFS to get nodes, prioritizing high-degree nodes at the same depth
|
||||
bfs_nodes = []
|
||||
visited = set()
|
||||
# Store (node, depth, degree) in the queue
|
||||
queue = deque([(node_label, 0, graph.degree(node_label))])
|
||||
|
||||
# Flag to track if there are unexplored neighbors due to depth limit
|
||||
has_unexplored_neighbors = False
|
||||
|
||||
# Modified breadth-first search with degree-based prioritization
|
||||
while queue and len(bfs_nodes) < max_nodes:
|
||||
# Get the current depth from the first node in queue
|
||||
current_depth = queue[0][1]
|
||||
|
||||
# Collect all nodes at the current depth
|
||||
current_level_nodes = []
|
||||
while queue and queue[0][1] == current_depth:
|
||||
current_level_nodes.append(queue.popleft())
|
||||
|
||||
# Sort nodes at current depth by degree (highest first)
|
||||
current_level_nodes.sort(key=lambda x: x[2], reverse=True)
|
||||
|
||||
# Process all nodes at current depth in order of degree
|
||||
for current_node, depth, degree in current_level_nodes:
|
||||
if current_node not in visited:
|
||||
visited.add(current_node)
|
||||
bfs_nodes.append(current_node)
|
||||
|
||||
# Only explore neighbors if we haven't reached max_depth
|
||||
if depth < max_depth:
|
||||
# Add neighbor nodes to queue with incremented depth
|
||||
neighbors = list(graph.neighbors(current_node))
|
||||
# Filter out already visited neighbors
|
||||
unvisited_neighbors = [
|
||||
n for n in neighbors if n not in visited
|
||||
]
|
||||
# Add neighbors to the queue with their degrees
|
||||
for neighbor in unvisited_neighbors:
|
||||
neighbor_degree = graph.degree(neighbor)
|
||||
queue.append((neighbor, depth + 1, neighbor_degree))
|
||||
else:
|
||||
# Check if there are unexplored neighbors (skipped due to depth limit)
|
||||
neighbors = list(graph.neighbors(current_node))
|
||||
unvisited_neighbors = [
|
||||
n for n in neighbors if n not in visited
|
||||
]
|
||||
if unvisited_neighbors:
|
||||
has_unexplored_neighbors = True
|
||||
|
||||
# Check if we've reached max_nodes
|
||||
if len(bfs_nodes) >= max_nodes:
|
||||
break
|
||||
|
||||
# Check if graph is truncated - either due to max_nodes limit or depth limit
|
||||
if (queue and len(bfs_nodes) >= max_nodes) or has_unexplored_neighbors:
|
||||
if len(bfs_nodes) >= max_nodes:
|
||||
result.is_truncated = True
|
||||
logger.info(
|
||||
f"[{self.workspace}] Graph truncated: max_nodes limit {max_nodes} reached"
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"[{self.workspace}] Graph truncated: found {len(bfs_nodes)} nodes within max_depth {max_depth}"
|
||||
)
|
||||
|
||||
# Create subgraph with BFS discovered nodes
|
||||
subgraph = graph.subgraph(bfs_nodes)
|
||||
|
||||
# Add nodes to result
|
||||
seen_nodes = set()
|
||||
seen_edges = set()
|
||||
for node in subgraph.nodes():
|
||||
if str(node) in seen_nodes:
|
||||
continue
|
||||
|
||||
node_data = dict(subgraph.nodes[node])
|
||||
# Get entity_type as labels
|
||||
labels = []
|
||||
if "entity_type" in node_data:
|
||||
if isinstance(node_data["entity_type"], list):
|
||||
labels.extend(node_data["entity_type"])
|
||||
else:
|
||||
labels.append(node_data["entity_type"])
|
||||
|
||||
# Create node with properties
|
||||
node_properties = {k: v for k, v in node_data.items()}
|
||||
|
||||
result.nodes.append(
|
||||
KnowledgeGraphNode(
|
||||
id=str(node), labels=[str(node)], properties=node_properties
|
||||
)
|
||||
)
|
||||
seen_nodes.add(str(node))
|
||||
|
||||
# Add edges to result
|
||||
for edge in subgraph.edges():
|
||||
source, target = edge
|
||||
# Esure unique edge_id for undirect graph
|
||||
if str(source) > str(target):
|
||||
source, target = target, source
|
||||
edge_id = f"{source}-{target}"
|
||||
if edge_id in seen_edges:
|
||||
continue
|
||||
|
||||
edge_data = dict(subgraph.edges[edge])
|
||||
|
||||
# Create edge with complete information
|
||||
result.edges.append(
|
||||
KnowledgeGraphEdge(
|
||||
id=edge_id,
|
||||
type="DIRECTED",
|
||||
source=str(source),
|
||||
target=str(target),
|
||||
properties=edge_data,
|
||||
)
|
||||
)
|
||||
seen_edges.add(edge_id)
|
||||
|
||||
logger.info(
|
||||
f"[{self.workspace}] Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
|
||||
)
|
||||
return result
|
||||
|
||||
async def get_all_nodes(self) -> list[dict]:
|
||||
"""Get all nodes in the graph.
|
||||
|
||||
Returns:
|
||||
A list of all nodes, where each node is a dictionary of its properties
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
all_nodes = []
|
||||
for node_id, node_data in graph.nodes(data=True):
|
||||
node_data_with_id = node_data.copy()
|
||||
node_data_with_id["id"] = node_id
|
||||
all_nodes.append(node_data_with_id)
|
||||
return all_nodes
|
||||
|
||||
async def get_all_edges(self) -> list[dict]:
|
||||
"""Get all edges in the graph.
|
||||
|
||||
Returns:
|
||||
A list of all edges, where each edge is a dictionary of its properties
|
||||
"""
|
||||
graph = await self._get_graph()
|
||||
all_edges = []
|
||||
for u, v, edge_data in graph.edges(data=True):
|
||||
edge_data_with_nodes = edge_data.copy()
|
||||
edge_data_with_nodes["source"] = u
|
||||
edge_data_with_nodes["target"] = v
|
||||
all_edges.append(edge_data_with_nodes)
|
||||
return all_edges
|
||||
|
||||
async def index_done_callback(self) -> bool:
|
||||
"""Commit in-memory graph to disk and notify other processes.
|
||||
|
||||
This is the writer's **commit point** in the cross-process sync
|
||||
protocol (see class docstring). Two effects, in order:
|
||||
1. ``write_nx_graph`` atomically writes the GraphML file
|
||||
(``atomic_write`` swaps a tmp file into place).
|
||||
2. ``set_all_update_flags`` flips every registered process's
|
||||
``storage_updated`` flag, then we immediately reset our
|
||||
own flag to ``False`` so the writer does not self-reload
|
||||
on the next call to ``_get_graph``.
|
||||
|
||||
Two-block structure (intentional, do not collapse):
|
||||
* **First ``async with``** — early-return path for a
|
||||
hypothetical second writer. Under the current single-writer
|
||||
pipeline contract (class docstring, invariant 1) the
|
||||
``storage_updated.value`` check is permanently ``False`` in
|
||||
the writer, so this branch is **dead code in production**.
|
||||
It is kept as defensive scaffolding for any future
|
||||
relaxation of the single-writer invariant; removing it
|
||||
would silently re-enable lost-write bugs the moment a
|
||||
second writer is introduced.
|
||||
* **Second ``async with``** — the actual save + notify.
|
||||
"""
|
||||
async with self._storage_lock:
|
||||
# Check if storage was updated by another process
|
||||
if self.storage_updated.value:
|
||||
# Storage was updated by another process, reload data instead of saving
|
||||
logger.info(
|
||||
f"[{self.workspace}] Graph was updated by another process, reloading..."
|
||||
)
|
||||
self._graph = (
|
||||
NetworkXStorage.load_nx_graph(self._graphml_xml_file) or nx.Graph()
|
||||
)
|
||||
# Reset update flag
|
||||
self.storage_updated.value = False
|
||||
return False # Return error
|
||||
|
||||
# Acquire lock and perform persistence
|
||||
async with self._storage_lock:
|
||||
try:
|
||||
# Save data to disk
|
||||
NetworkXStorage.write_nx_graph(
|
||||
self._graph, self._graphml_xml_file, self.workspace
|
||||
)
|
||||
# Notify other processes that data has been updated
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
# Reset own update flag to avoid self-reloading
|
||||
self.storage_updated.value = False
|
||||
return True # Return success
|
||||
except Exception as e:
|
||||
# Raise (do NOT swallow + return False): _insert_done's
|
||||
# _flush_one only detects failures via exceptions, so a
|
||||
# swallowed graph-save error would let the document be marked
|
||||
# PROCESSED with the graph changes unpersisted. Surfacing it
|
||||
# aligns this backend with the others (faiss/nano raise too).
|
||||
logger.error(f"[{self.workspace}] Error saving graph: {e}")
|
||||
raise
|
||||
|
||||
return True
|
||||
|
||||
async def drop(self) -> dict[str, str]:
|
||||
"""Drop all graph data from storage and reinitialize the graph.
|
||||
|
||||
This method will:
|
||||
1. Remove the graph storage file if it exists
|
||||
2. Reset the graph to an empty ``nx.Graph()``
|
||||
3. Update flags to notify other processes
|
||||
4. Changes are persisted to disk immediately
|
||||
|
||||
Caller contract:
|
||||
``drop`` is destructive and **not** serialized by this storage
|
||||
class. The caller must hold the pipeline ``busy`` reservation
|
||||
(the ``/documents/clear`` endpoint does this) before invoking
|
||||
it — running ``drop`` concurrently with an active document
|
||||
pipeline will tear down storage out from under the writer and
|
||||
silently lose data. See class docstring,
|
||||
*Non-pipeline write paths*.
|
||||
|
||||
Returns:
|
||||
dict[str, str]: Operation status and message
|
||||
- On success: {"status": "success", "message": "data dropped"}
|
||||
- On failure: {"status": "error", "message": "<error details>"}
|
||||
"""
|
||||
try:
|
||||
async with self._storage_lock:
|
||||
# delete _client_file_name
|
||||
if os.path.exists(self._graphml_xml_file):
|
||||
os.remove(self._graphml_xml_file)
|
||||
self._graph = nx.Graph()
|
||||
# Notify other processes that data has been updated
|
||||
await set_all_update_flags(self.namespace, workspace=self.workspace)
|
||||
# Reset own update flag to avoid self-reloading
|
||||
self.storage_updated.value = False
|
||||
logger.info(
|
||||
f"[{self.workspace}] Process {os.getpid()} drop graph file:{self._graphml_xml_file}"
|
||||
)
|
||||
return {"status": "success", "message": "data dropped"}
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[{self.workspace}] Error dropping graph file:{self._graphml_xml_file}: {e}"
|
||||
)
|
||||
return {"status": "error", "message": str(e)}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,301 @@
|
||||
"""Shared image-input normalization for LLM bindings.
|
||||
|
||||
All LLM bindings accept a unified ``image_inputs`` keyword parameter. Each
|
||||
element may be:
|
||||
|
||||
- a raw base64 string (the MIME type is inferred via ``imghdr`` / magic bytes,
|
||||
defaulting to ``image/png``);
|
||||
- a data URL of the form ``data:<mime>;base64,<payload>``;
|
||||
- a dict with keys ``base64`` (required) and optional ``mime_type``,
|
||||
``source_id``, ``source_file``, ``modality``, ``doc_id``.
|
||||
|
||||
The provider-specific binding code converts the normalized result to its own
|
||||
content-block format. The VLM pipeline uses :func:`image_cache_metadata` for
|
||||
cache-key inputs (deliberately excluding ``source_id`` / ``source_file`` so the
|
||||
same image at different filenames still hits the same entry) and
|
||||
:func:`image_audit_metadata` for the human-readable ``original_prompt`` audit
|
||||
block.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import re
|
||||
import struct
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
DATA_URL_RE = re.compile(
|
||||
r"^data:(?P<mime>[\w./+-]+);base64,(?P<data>[A-Za-z0-9+/=\s]+)$"
|
||||
)
|
||||
|
||||
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
|
||||
_JPEG_SIGNATURE = b"\xff\xd8\xff"
|
||||
_GIF_SIGNATURES = (b"GIF87a", b"GIF89a")
|
||||
_WEBP_RIFF = b"RIFF"
|
||||
_WEBP_TAG = b"WEBP"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NormalizedImage:
|
||||
index: int
|
||||
raw_bytes: bytes
|
||||
mime_type: str
|
||||
sha256: str
|
||||
base64_str: str
|
||||
source_id: str | None
|
||||
source_file: str | None
|
||||
modality: str | None
|
||||
doc_id: str | None
|
||||
# Pixel dimensions parsed from the raster header (None when the format
|
||||
# is recognized but dimensions could not be extracted).
|
||||
width: int | None = None
|
||||
height: int | None = None
|
||||
|
||||
|
||||
def _detect_mime(raw: bytes) -> str:
|
||||
if raw.startswith(_PNG_SIGNATURE):
|
||||
return "image/png"
|
||||
if raw.startswith(_JPEG_SIGNATURE):
|
||||
return "image/jpeg"
|
||||
if any(raw.startswith(sig) for sig in _GIF_SIGNATURES):
|
||||
return "image/gif"
|
||||
if len(raw) >= 12 and raw[0:4] == _WEBP_RIFF and raw[8:12] == _WEBP_TAG:
|
||||
return "image/webp"
|
||||
return "image/png"
|
||||
|
||||
|
||||
def _decode_base64(data: str) -> bytes:
|
||||
cleaned = re.sub(r"\s+", "", data)
|
||||
try:
|
||||
return base64.b64decode(cleaned, validate=True)
|
||||
except (base64.binascii.Error, ValueError) as exc:
|
||||
raise ValueError(f"invalid base64 image data: {exc}") from exc
|
||||
|
||||
|
||||
def _coerce_item(item: Any) -> dict[str, Any]:
|
||||
if isinstance(item, str):
|
||||
match = DATA_URL_RE.match(item.strip())
|
||||
if match:
|
||||
return {"base64": match.group("data"), "mime_type": match.group("mime")}
|
||||
return {"base64": item}
|
||||
if isinstance(item, dict):
|
||||
if "base64" not in item:
|
||||
raise ValueError("image_inputs dict element must contain a 'base64' key")
|
||||
return item
|
||||
raise TypeError(
|
||||
f"image_inputs element must be str or dict, got {type(item).__name__}"
|
||||
)
|
||||
|
||||
|
||||
def normalize_image_inputs(
|
||||
image_inputs: list[Any] | None,
|
||||
) -> list[NormalizedImage]:
|
||||
"""Normalize the unified ``image_inputs`` parameter.
|
||||
|
||||
Returns an empty list when ``image_inputs`` is falsy, so callers can do a
|
||||
plain ``if normalized:`` check.
|
||||
"""
|
||||
if not image_inputs:
|
||||
return []
|
||||
|
||||
result: list[NormalizedImage] = []
|
||||
for idx, raw_item in enumerate(image_inputs):
|
||||
item = _coerce_item(raw_item)
|
||||
raw_bytes = _decode_base64(item["base64"])
|
||||
if not raw_bytes:
|
||||
raise ValueError(f"image_inputs[{idx}] decoded to empty bytes")
|
||||
mime_type = item.get("mime_type") or _detect_mime(raw_bytes)
|
||||
sha = hashlib.sha256(raw_bytes).hexdigest()
|
||||
clean_b64 = base64.b64encode(raw_bytes).decode("ascii")
|
||||
dims = _dimensions_from_bytes(raw_bytes)
|
||||
width, height = (dims[0], dims[1]) if dims else (None, None)
|
||||
result.append(
|
||||
NormalizedImage(
|
||||
index=idx,
|
||||
raw_bytes=raw_bytes,
|
||||
mime_type=mime_type,
|
||||
sha256=sha,
|
||||
base64_str=clean_b64,
|
||||
source_id=item.get("source_id"),
|
||||
source_file=item.get("source_file"),
|
||||
modality=item.get("modality"),
|
||||
doc_id=item.get("doc_id"),
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def image_cache_metadata(images: list[NormalizedImage]) -> list[dict[str, Any]]:
|
||||
"""Return cache-key-safe image metadata (no source identifiers).
|
||||
|
||||
Includes ``width`` / ``height`` so the cache key reflects the full
|
||||
image digest the design contract specifies (mime, sha256, bytes,
|
||||
width, height). The sha256 alone is sufficient for identity, but
|
||||
surfacing dimensions matches the documented audit shape and gives
|
||||
diagnostics a one-line "what was sent" without re-decoding.
|
||||
"""
|
||||
return [
|
||||
{
|
||||
"index": img.index,
|
||||
"mime_type": img.mime_type,
|
||||
"sha256": img.sha256,
|
||||
"bytes": len(img.raw_bytes),
|
||||
"width": img.width,
|
||||
"height": img.height,
|
||||
}
|
||||
for img in images
|
||||
]
|
||||
|
||||
|
||||
def image_audit_metadata(images: list[NormalizedImage]) -> list[dict[str, Any]]:
|
||||
"""Return audit metadata suitable for the ``original_prompt`` block.
|
||||
|
||||
Never includes the raw base64 payload — only digests and source pointers.
|
||||
"""
|
||||
return [
|
||||
{
|
||||
"index": img.index,
|
||||
"mime_type": img.mime_type,
|
||||
"sha256": img.sha256,
|
||||
"bytes": len(img.raw_bytes),
|
||||
"width": img.width,
|
||||
"height": img.height,
|
||||
"source_id": img.source_id,
|
||||
"source_file": img.source_file,
|
||||
"modality": img.modality,
|
||||
"doc_id": img.doc_id,
|
||||
}
|
||||
for img in images
|
||||
]
|
||||
|
||||
|
||||
def _read_png_dimensions(data: bytes) -> tuple[int, int] | None:
|
||||
# IHDR is the first chunk; width/height are big-endian uint32 at offsets
|
||||
# 16/20 (8-byte signature + 4 length + 4 "IHDR" + 4 width + 4 height).
|
||||
if len(data) < 24 or not data.startswith(_PNG_SIGNATURE):
|
||||
return None
|
||||
width, height = struct.unpack(">II", data[16:24])
|
||||
return width, height
|
||||
|
||||
|
||||
def _read_gif_dimensions(data: bytes) -> tuple[int, int] | None:
|
||||
# Logical screen descriptor: width/height are little-endian uint16 at
|
||||
# offsets 6/8.
|
||||
if len(data) < 10 or not any(data.startswith(sig) for sig in _GIF_SIGNATURES):
|
||||
return None
|
||||
width, height = struct.unpack("<HH", data[6:10])
|
||||
return width, height
|
||||
|
||||
|
||||
def _read_jpeg_dimensions(data: bytes) -> tuple[int, int] | None:
|
||||
# Scan for a Start-Of-Frame marker (SOF0 / SOF2 / etc.). Skip segments by
|
||||
# their length field. We deliberately accept any SOF variant the codec
|
||||
# might emit rather than enumerating each one.
|
||||
if len(data) < 4 or not data.startswith(_JPEG_SIGNATURE):
|
||||
return None
|
||||
i = 2
|
||||
n = len(data)
|
||||
while i < n:
|
||||
if data[i] != 0xFF:
|
||||
return None
|
||||
# Skip fill bytes.
|
||||
while i < n and data[i] == 0xFF:
|
||||
i += 1
|
||||
if i >= n:
|
||||
return None
|
||||
marker = data[i]
|
||||
i += 1
|
||||
# Standalone markers without a length field.
|
||||
if marker in (0xD8, 0xD9) or 0xD0 <= marker <= 0xD7:
|
||||
continue
|
||||
if i + 2 > n:
|
||||
return None
|
||||
segment_len = struct.unpack(">H", data[i : i + 2])[0]
|
||||
if segment_len < 2 or i + segment_len > n:
|
||||
return None
|
||||
# SOF0..SOF15 except 0xC4 (DHT), 0xC8 (JPG reserved), 0xCC (DAC).
|
||||
if 0xC0 <= marker <= 0xCF and marker not in (0xC4, 0xC8, 0xCC):
|
||||
# SOF payload: precision(1) + height(2) + width(2) + …
|
||||
if i + 7 > n:
|
||||
return None
|
||||
height, width = struct.unpack(">HH", data[i + 3 : i + 7])
|
||||
return width, height
|
||||
i += segment_len
|
||||
return None
|
||||
|
||||
|
||||
def _read_webp_dimensions(data: bytes) -> tuple[int, int] | None:
|
||||
if len(data) < 30 or data[0:4] != _WEBP_RIFF or data[8:12] != _WEBP_TAG:
|
||||
return None
|
||||
chunk_type = data[12:16]
|
||||
if chunk_type == b"VP8 ":
|
||||
# Lossy: 3-byte tag + 3-byte sync code at offset 23, then 4 bytes
|
||||
# holding 14-bit width / 14-bit height in little-endian halves.
|
||||
if len(data) < 30:
|
||||
return None
|
||||
width = struct.unpack("<H", data[26:28])[0] & 0x3FFF
|
||||
height = struct.unpack("<H", data[28:30])[0] & 0x3FFF
|
||||
return width, height
|
||||
if chunk_type == b"VP8L":
|
||||
# Lossless: signature(0x2F) + 4 bytes encoding 14-bit width-1 / 14-bit
|
||||
# height-1 starting at offset 21.
|
||||
if len(data) < 25 or data[20] != 0x2F:
|
||||
return None
|
||||
b0, b1, b2, b3 = data[21], data[22], data[23], data[24]
|
||||
width = ((b1 & 0x3F) << 8 | b0) + 1
|
||||
height = ((b3 & 0x0F) << 10 | b2 << 2 | (b1 & 0xC0) >> 6) + 1
|
||||
return width, height
|
||||
if chunk_type == b"VP8X":
|
||||
# Extended: 3 bytes width-1 / 3 bytes height-1, little-endian, at
|
||||
# offsets 24/27.
|
||||
if len(data) < 30:
|
||||
return None
|
||||
width = (data[24] | data[25] << 8 | data[26] << 16) + 1
|
||||
height = (data[27] | data[28] << 8 | data[29] << 16) + 1
|
||||
return width, height
|
||||
return None
|
||||
|
||||
|
||||
def read_image_dimensions(path: Path) -> tuple[int, int] | None:
|
||||
"""Return ``(width, height)`` for a raster image, or ``None`` if unknown.
|
||||
|
||||
Reads only the file header — no Pillow dependency. Supports PNG, JPEG,
|
||||
GIF and WebP (VP8 / VP8L / VP8X). Returns ``None`` for unsupported
|
||||
formats and on any I/O or parse error so callers can fall back to a
|
||||
skipped/failure decision without raising.
|
||||
"""
|
||||
try:
|
||||
with open(path, "rb") as fh:
|
||||
header = fh.read(64 * 1024)
|
||||
except OSError:
|
||||
return None
|
||||
return _dimensions_from_bytes(header)
|
||||
|
||||
|
||||
def _dimensions_from_bytes(data: bytes) -> tuple[int, int] | None:
|
||||
"""Run the four header readers against a byte buffer.
|
||||
|
||||
Shared between the file-path entry point (:func:`read_image_dimensions`)
|
||||
and :func:`normalize_image_inputs`, which receives raster payloads
|
||||
decoded from the unified ``image_inputs`` parameter.
|
||||
"""
|
||||
if not data:
|
||||
return None
|
||||
for reader in (
|
||||
_read_png_dimensions,
|
||||
_read_gif_dimensions,
|
||||
_read_jpeg_dimensions,
|
||||
_read_webp_dimensions,
|
||||
):
|
||||
try:
|
||||
dims = reader(data)
|
||||
except (struct.error, IndexError, ValueError):
|
||||
continue
|
||||
if dims:
|
||||
return dims
|
||||
return None
|
||||
@@ -0,0 +1,376 @@
|
||||
from ..utils import verbose_debug, VERBOSE_DEBUG
|
||||
import sys
|
||||
import os
|
||||
import logging
|
||||
import warnings
|
||||
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")
|
||||
|
||||
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,
|
||||
image_inputs: list[Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Union[str, AsyncIterator[str]]:
|
||||
"""Call Anthropic Messages API with LightRAG-compatible shims.
|
||||
|
||||
Structured output note:
|
||||
- This adapter does not support OpenAI-style ``response_format`` JSON mode.
|
||||
- If callers pass ``response_format``, it is stripped before the request.
|
||||
- Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are
|
||||
accepted only as compatibility shims; they emit warnings and are ignored.
|
||||
"""
|
||||
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)
|
||||
# Anthropic Messages API has no JSON mode; drop legacy flags and
|
||||
# response_format. Emit DeprecationWarning when the booleans were set.
|
||||
if kwargs.pop("keyword_extraction", False):
|
||||
warnings.warn(
|
||||
"anthropic_complete_if_cache(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if kwargs.pop("entity_extraction", False):
|
||||
warnings.warn(
|
||||
"anthropic_complete_if_cache(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs.pop("response_format", None)
|
||||
timeout = kwargs.pop("timeout", None)
|
||||
|
||||
# Require max_tokens; the Anthropic SDK errors if it's missing
|
||||
kwargs.setdefault("max_tokens", 8192)
|
||||
# Pop stream from kwargs so it doesn't leak into create_params;
|
||||
# default to False (non-streaming) for consistency with other providers
|
||||
stream = kwargs.pop("stream", False)
|
||||
|
||||
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]] = []
|
||||
messages.extend(history_messages)
|
||||
if image_inputs:
|
||||
from lightrag.llm._vision_utils import normalize_image_inputs
|
||||
|
||||
normalized_images = normalize_image_inputs(image_inputs)
|
||||
user_content: list[dict[str, Any]] = []
|
||||
for img in normalized_images:
|
||||
user_content.append(
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": img.mime_type,
|
||||
"data": img.base64_str,
|
||||
},
|
||||
}
|
||||
)
|
||||
user_content.append({"type": "text", "text": prompt})
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
else:
|
||||
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:
|
||||
create_params = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"stream": stream,
|
||||
**kwargs,
|
||||
}
|
||||
if system_prompt:
|
||||
create_params["system"] = system_prompt
|
||||
response = await anthropic_async_client.messages.create(**create_params)
|
||||
except APITimeoutError as e:
|
||||
# APITimeoutError subclasses APIConnectionError, so it must come first
|
||||
# or it would be swallowed by the broader handler below.
|
||||
logger.error(f"Anthropic API Timeout Error: {e}")
|
||||
try:
|
||||
await anthropic_async_client.close()
|
||||
except Exception as close_error:
|
||||
logger.warning(f"Failed to close Anthropic client: {close_error}")
|
||||
raise
|
||||
except APIConnectionError as e:
|
||||
logger.error(f"Anthropic API Connection Error: {e}")
|
||||
try:
|
||||
await anthropic_async_client.close()
|
||||
except Exception as close_error:
|
||||
logger.warning(f"Failed to close Anthropic client: {close_error}")
|
||||
raise
|
||||
except RateLimitError as e:
|
||||
logger.error(f"Anthropic API Rate Limit Error: {e}")
|
||||
try:
|
||||
await anthropic_async_client.close()
|
||||
except Exception as close_error:
|
||||
logger.warning(f"Failed to close Anthropic client: {close_error}")
|
||||
raise
|
||||
except Exception as e:
|
||||
body = getattr(e, "body", None)
|
||||
request_id = getattr(e, "request_id", None)
|
||||
req = getattr(e, "request", None)
|
||||
extra_parts = []
|
||||
if body:
|
||||
extra_parts.append(f"Response body: {body}")
|
||||
if request_id:
|
||||
extra_parts.append(f"Request ID: {request_id}")
|
||||
if req is not None:
|
||||
extra_parts.append(f"Request URL: {req.url}")
|
||||
extra = ("\n" + "\n".join(extra_parts)) if extra_parts else ""
|
||||
logger.error(
|
||||
f"Anthropic API Call Failed,\nModel: {model},\nParams: {kwargs}, Got: {e}{extra}"
|
||||
)
|
||||
try:
|
||||
await anthropic_async_client.close()
|
||||
except Exception as close_error:
|
||||
logger.warning(f"Failed to close Anthropic client: {close_error}")
|
||||
raise
|
||||
except BaseException:
|
||||
# CancelledError (and other BaseExceptions) aren't caught above, yet a
|
||||
# cancellation while awaiting create() — before we hold the response or
|
||||
# stream — would otherwise leak the client. Close it, then re-raise the
|
||||
# cancellation untouched (cooperative cancellation).
|
||||
try:
|
||||
await anthropic_async_client.close()
|
||||
except Exception as close_error:
|
||||
logger.warning(f"Failed to close Anthropic client: {close_error}")
|
||||
raise
|
||||
|
||||
if not stream:
|
||||
try:
|
||||
return response.content[0].text
|
||||
finally:
|
||||
try:
|
||||
await anthropic_async_client.close()
|
||||
except Exception as close_error:
|
||||
logger.warning(f"Failed to close Anthropic client: {close_error}")
|
||||
|
||||
async def stream_response():
|
||||
try:
|
||||
async for event in response:
|
||||
content = (
|
||||
event.delta.text
|
||||
if hasattr(event, "delta")
|
||||
and hasattr(event.delta, "text")
|
||||
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
|
||||
finally:
|
||||
# Release the streaming response and the client. The generator owns
|
||||
# the client lifetime, so the caller never closes it; finally also
|
||||
# runs on GeneratorExit (early break) which the except clause above
|
||||
# cannot catch. AsyncStream exposes close() (not aclose()).
|
||||
#
|
||||
# The client close lives in an outer finally so it still runs if
|
||||
# response.close() is interrupted by CancelledError (a BaseException,
|
||||
# not caught by `except Exception`) — otherwise task cancellation /
|
||||
# client disconnect during stream teardown would leak the client.
|
||||
try:
|
||||
try:
|
||||
await response.close()
|
||||
except Exception as close_error:
|
||||
logger.warning(f"Failed to close Anthropic stream: {close_error}")
|
||||
finally:
|
||||
try:
|
||||
await anthropic_async_client.close()
|
||||
except Exception as close_error:
|
||||
logger.warning(f"Failed to close Anthropic client: {close_error}")
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
# Backward-compatibility shim: the previous embedding implementation lived in
|
||||
# this module under the (misleading) name ``anthropic_embed`` even though it
|
||||
# called Voyage AI under the hood. The real implementation now lives in
|
||||
# ``lightrag.llm.voyageai.voyageai_embed``. Keep the old name importable for one
|
||||
# release cycle so downstream users get a clear deprecation warning instead of
|
||||
# an ImportError. Remove in a future major version.
|
||||
def anthropic_embed(*args, **kwargs):
|
||||
"""Deprecated alias for :func:`lightrag.llm.voyageai.voyageai_embed`.
|
||||
|
||||
This shim accepts the same arguments as the original ``anthropic_embed``
|
||||
function (which was always backed by VoyageAI) and forwards them to
|
||||
:func:`voyageai_embed`. It will be removed in a future release.
|
||||
"""
|
||||
|
||||
warnings.warn(
|
||||
"lightrag.llm.anthropic.anthropic_embed is deprecated and will be "
|
||||
"removed in a future release. Import "
|
||||
"lightrag.llm.voyageai.voyageai_embed instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
from lightrag.llm.voyageai import voyageai_embed
|
||||
|
||||
return voyageai_embed.func(*args, **kwargs)
|
||||
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
Azure OpenAI compatibility layer.
|
||||
|
||||
This module provides backward compatibility by re-exporting Azure OpenAI functions
|
||||
from the main openai module where the actual implementation resides.
|
||||
|
||||
All core logic for both OpenAI and Azure OpenAI now lives in lightrag.llm.openai,
|
||||
with this module serving as a thin compatibility wrapper for existing code that
|
||||
imports from lightrag.llm.azure_openai.
|
||||
"""
|
||||
|
||||
from lightrag.llm.openai import (
|
||||
azure_openai_complete_if_cache,
|
||||
azure_openai_complete,
|
||||
azure_openai_embed,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"azure_openai_complete_if_cache",
|
||||
"azure_openai_complete",
|
||||
"azure_openai_embed",
|
||||
]
|
||||
@@ -0,0 +1,609 @@
|
||||
import copy
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import warnings
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, Union
|
||||
|
||||
from lightrag.utils import wrap_embedding_func_with_attrs
|
||||
|
||||
# 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 _normalize_bedrock_endpoint_url(endpoint_url: str | None) -> str | None:
|
||||
"""Return a usable Bedrock endpoint override or None for SDK defaults."""
|
||||
if endpoint_url is None:
|
||||
return None
|
||||
|
||||
normalized = endpoint_url.strip()
|
||||
if not normalized or normalized == "DEFAULT_BEDROCK_ENDPOINT":
|
||||
return None
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def _bedrock_client_kwargs(
|
||||
region: str | None,
|
||||
endpoint_url: str | None,
|
||||
aws_access_key_id: str | None = None,
|
||||
aws_secret_access_key: str | None = None,
|
||||
aws_session_token: str | None = None,
|
||||
) -> dict:
|
||||
"""Build kwargs for aioboto3 ``session.client("bedrock-runtime", ...)``."""
|
||||
client_kwargs: dict = {"region_name": region}
|
||||
if endpoint_url is not None:
|
||||
client_kwargs["endpoint_url"] = endpoint_url
|
||||
if aws_access_key_id:
|
||||
client_kwargs["aws_access_key_id"] = aws_access_key_id
|
||||
if aws_secret_access_key:
|
||||
client_kwargs["aws_secret_access_key"] = aws_secret_access_key
|
||||
if aws_session_token:
|
||||
client_kwargs["aws_session_token"] = aws_session_token
|
||||
return client_kwargs
|
||||
|
||||
|
||||
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,
|
||||
aws_region: str | None = None,
|
||||
api_key: str | None = None,
|
||||
endpoint_url: str | None = None,
|
||||
image_inputs: list[Any] | None = None,
|
||||
**kwargs,
|
||||
) -> Union[str, AsyncIterator[str]]:
|
||||
"""Call Amazon Bedrock Converse API with LightRAG-compatible shims.
|
||||
|
||||
Structured output note:
|
||||
- This adapter does not support OpenAI-style ``response_format`` JSON mode.
|
||||
- If callers pass ``response_format``, it is stripped before the request.
|
||||
- Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are
|
||||
accepted only as compatibility shims; they emit warnings and are ignored.
|
||||
|
||||
Authentication note:
|
||||
- Bedrock does not use LightRAG's generic ``api_key`` fields.
|
||||
- ``LLM_BINDING_API_KEY`` and ``EMBEDDING_BINDING_API_KEY`` are ignored for
|
||||
Bedrock.
|
||||
- To use Bedrock API key / bearer-token auth, set
|
||||
``AWS_BEARER_TOKEN_BEDROCK`` before starting the process; this is a
|
||||
process-level AWS SDK setting.
|
||||
- For role-specific Bedrock LLMs, use explicit SigV4 parameters
|
||||
(``aws_access_key_id``, ``aws_secret_access_key``, ``aws_session_token``,
|
||||
``aws_region``). Per-role bearer-token overrides are not supported.
|
||||
|
||||
Endpoint note:
|
||||
- ``endpoint_url`` overrides the default regional Bedrock endpoint. Pass
|
||||
``None``, an empty string, or the sentinel ``DEFAULT_BEDROCK_ENDPOINT``
|
||||
to let the AWS SDK select its default endpoint.
|
||||
"""
|
||||
if enable_cot:
|
||||
logging.debug(
|
||||
"enable_cot=True is not supported for Bedrock and will be ignored."
|
||||
)
|
||||
|
||||
# Bedrock Converse API has no JSON mode; drop legacy extraction flags and
|
||||
# response_format below and rely on the prompt template plus downstream
|
||||
# tolerant JSON parsing.
|
||||
keyword_extraction = kwargs.pop("keyword_extraction", False)
|
||||
entity_extraction = kwargs.pop("entity_extraction", False)
|
||||
if keyword_extraction:
|
||||
warnings.warn(
|
||||
"bedrock_complete_if_cache(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if entity_extraction:
|
||||
warnings.warn(
|
||||
"bedrock_complete_if_cache(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if api_key:
|
||||
warnings.warn(
|
||||
"bedrock_complete_if_cache(api_key=...) is ignored; use SigV4 "
|
||||
"parameters or set AWS_BEARER_TOKEN_BEDROCK before process start.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
region = aws_region or kwargs.pop("aws_region", None)
|
||||
endpoint_url = _normalize_bedrock_endpoint_url(endpoint_url)
|
||||
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",
|
||||
]:
|
||||
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
|
||||
if image_inputs:
|
||||
from lightrag.llm._vision_utils import normalize_image_inputs
|
||||
|
||||
normalized_images = normalize_image_inputs(image_inputs)
|
||||
user_content: list[dict[str, Any]] = [{"text": prompt}]
|
||||
for img in normalized_images:
|
||||
fmt = img.mime_type.split("/", 1)[1] if "/" in img.mime_type else "png"
|
||||
user_content.append(
|
||||
{"image": {"format": fmt, "source": {"bytes": img.raw_bytes}}}
|
||||
)
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
if stream:
|
||||
logging.getLogger(__name__).debug(
|
||||
"[bedrock] image_inputs provided; forcing non-stream Converse "
|
||||
"(stream + image combination has SDK limitations)"
|
||||
)
|
||||
stream = False
|
||||
else:
|
||||
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",
|
||||
}
|
||||
inference_config: dict[str, Any] = {}
|
||||
for param in ("max_tokens", "temperature", "top_p", "stop_sequences"):
|
||||
if param not in kwargs:
|
||||
continue
|
||||
value = kwargs.pop(param)
|
||||
# Bedrock rejects None; a None default means "inherit provider default"
|
||||
if value is None:
|
||||
continue
|
||||
inference_config[inference_params_map.get(param, param)] = value
|
||||
if inference_config:
|
||||
args["inferenceConfig"] = inference_config
|
||||
|
||||
# Pass-through for model-specific parameters (e.g. Anthropic reasoning_config,
|
||||
# Nova inferenceConfig extensions). Mirrors OpenAI's `extra_body`.
|
||||
extra_fields = kwargs.pop("extra_fields", None)
|
||||
if extra_fields:
|
||||
args["additionalModelRequestFields"] = extra_fields
|
||||
|
||||
# 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_kwargs = _bedrock_client_kwargs(
|
||||
region,
|
||||
endpoint_url,
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
)
|
||||
|
||||
# Define the generator function that will manage the client lifecycle
|
||||
async def stream_generator():
|
||||
# async with ensures the aioboto3 client is closed even under
|
||||
# task cancellation, avoiding aiohttp "Unclosed connection" warnings.
|
||||
async with session.client("bedrock-runtime", **client_kwargs) as client:
|
||||
event_stream = None
|
||||
try:
|
||||
# Make the API call
|
||||
response = await client.converse_stream(**args, **kwargs)
|
||||
event_stream = response.get("stream")
|
||||
|
||||
# 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:
|
||||
# Convert to appropriate exception type
|
||||
_handle_bedrock_exception(e, "Bedrock streaming")
|
||||
|
||||
finally:
|
||||
# Close the event stream once; client cleanup is handled by async with.
|
||||
# aiobotocore's EventStream exposes sync `close()`, while generic
|
||||
# async iterators expose async `aclose()` — handle both and dispatch
|
||||
# awaitable results accordingly.
|
||||
if event_stream is not None:
|
||||
close_fn = getattr(event_stream, "close", None) or getattr(
|
||||
event_stream, "aclose", None
|
||||
)
|
||||
if callable(close_fn):
|
||||
try:
|
||||
result = close_fn()
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
except Exception as close_error:
|
||||
logging.warning(
|
||||
f"Failed to close Bedrock event stream: {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",
|
||||
**_bedrock_client_kwargs(
|
||||
region,
|
||||
endpoint_url,
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
),
|
||||
) 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")
|
||||
|
||||
# When thinking/reasoning is enabled, the first content block is a
|
||||
# `reasoningContent` block and the visible text follows in a later
|
||||
# block. Pick the first block that carries a text payload.
|
||||
content = next(
|
||||
(
|
||||
block["text"]
|
||||
for block in response["output"]["message"]["content"]
|
||||
if isinstance(block, dict) and block.get("text")
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
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,
|
||||
entity_extraction=False,
|
||||
**kwargs,
|
||||
) -> Union[str, AsyncIterator[str]]:
|
||||
# Bedrock Converse API has no JSON mode; the shim booleans are absorbed
|
||||
# and forwarded so bedrock_complete_if_cache can emit DeprecationWarnings
|
||||
# with accurate stack frames.
|
||||
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,
|
||||
keyword_extraction=keyword_extraction,
|
||||
entity_extraction=entity_extraction,
|
||||
**kwargs,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@wrap_embedding_func_with_attrs(
|
||||
embedding_dim=1024, max_token_size=8192, model_name="amazon.titan-embed-text-v2:0"
|
||||
)
|
||||
@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,
|
||||
aws_region: str | None = None,
|
||||
api_key: str | None = None,
|
||||
endpoint_url: str | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Generate embeddings with Amazon Bedrock Runtime.
|
||||
|
||||
Authentication note:
|
||||
- Bedrock does not use LightRAG's generic ``api_key`` fields.
|
||||
- ``LLM_BINDING_API_KEY`` and ``EMBEDDING_BINDING_API_KEY`` are ignored for
|
||||
Bedrock.
|
||||
- To use Bedrock API key / bearer-token auth, set
|
||||
``AWS_BEARER_TOKEN_BEDROCK`` before starting the process; this is a
|
||||
process-level AWS SDK setting.
|
||||
- For role-specific Bedrock configuration, use explicit SigV4 parameters
|
||||
(``aws_access_key_id``, ``aws_secret_access_key``, ``aws_session_token``,
|
||||
``aws_region``). Per-role bearer-token overrides are not supported.
|
||||
"""
|
||||
if api_key:
|
||||
warnings.warn(
|
||||
"bedrock_embed(api_key=...) is ignored; use SigV4 parameters or "
|
||||
"set AWS_BEARER_TOKEN_BEDROCK before process start.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
region = aws_region
|
||||
endpoint_url = _normalize_bedrock_endpoint_url(endpoint_url)
|
||||
|
||||
session = aioboto3.Session()
|
||||
async with session.client(
|
||||
"bedrock-runtime",
|
||||
**_bedrock_client_kwargs(
|
||||
region,
|
||||
endpoint_url,
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
),
|
||||
) 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(
|
||||
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 "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")
|
||||
@@ -0,0 +1,833 @@
|
||||
"""
|
||||
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
|
||||
|
||||
@classmethod
|
||||
def options_dict_for_role(
|
||||
cls, args: Namespace, role: str, is_cross_provider: bool = False
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Extract role-specific provider options with proper inheritance.
|
||||
|
||||
Same provider:
|
||||
- inherit the base binding options from parsed args
|
||||
- overlay any role-specific environment variable overrides
|
||||
|
||||
Cross provider:
|
||||
- start from empty provider options
|
||||
- overlay any role-specific environment variable overrides
|
||||
|
||||
Role env vars follow the pattern:
|
||||
`{ROLE}_{BINDING_PREFIX}_{FIELD}`
|
||||
e.g. `EXTRACT_OPENAI_LLM_TEMPERATURE`
|
||||
"""
|
||||
import os
|
||||
|
||||
if is_cross_provider:
|
||||
base: dict[str, Any] = {}
|
||||
else:
|
||||
base = cls.options_dict(args)
|
||||
|
||||
role_upper = role.upper()
|
||||
env_prefix = cls._binding_name.upper() + "_"
|
||||
|
||||
for arg_item in cls.args_env_name_type_value():
|
||||
original_env = arg_item["env_name"]
|
||||
role_env = f"{role_upper}_{original_env}"
|
||||
field_name = original_env[len(env_prefix) :].lower()
|
||||
|
||||
env_raw = os.getenv(role_env)
|
||||
if env_raw is None:
|
||||
continue
|
||||
|
||||
field_type = _resolve_optional_type(arg_item["type"])
|
||||
try:
|
||||
if field_type is bool:
|
||||
base[field_name] = env_raw.lower() in (
|
||||
"true",
|
||||
"1",
|
||||
"yes",
|
||||
"t",
|
||||
"on",
|
||||
)
|
||||
elif field_type in (list, List[str]):
|
||||
base[field_name] = json.loads(env_raw)
|
||||
elif field_type is dict:
|
||||
base[field_name] = json.loads(env_raw)
|
||||
elif field_type is int:
|
||||
base[field_name] = int(env_raw)
|
||||
elif field_type is float:
|
||||
base[field_name] = float(env_raw)
|
||||
else:
|
||||
base[field_name] = env_raw
|
||||
except (ValueError, json.JSONDecodeError):
|
||||
base[field_name] = env_raw
|
||||
|
||||
return base
|
||||
|
||||
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 | None = None
|
||||
|
||||
_help: ClassVar[dict[str, str]] = {
|
||||
"task_type": "Task type for embedding optimization. If not specified, automatically determined from context (RETRIEVAL_QUERY for queries, RETRIEVAL_DOCUMENT for documents). Supported types: 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}}\')',
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Binding Options for AWS Bedrock
|
||||
# =============================================================================
|
||||
#
|
||||
# Bedrock binding options map to the subset of the Bedrock Converse API
|
||||
# inferenceConfig that LightRAG's bedrock driver actually forwards. See
|
||||
# ``lightrag/llm/bedrock.py`` for the whitelist — any field added here that is
|
||||
# not in that whitelist will be silently dropped by the driver.
|
||||
# =============================================================================
|
||||
@dataclass
|
||||
class BedrockLLMOptions(BindingOptions):
|
||||
"""Options for AWS Bedrock LLM (Converse API inferenceConfig)."""
|
||||
|
||||
_binding_name: ClassVar[str] = "bedrock_llm"
|
||||
|
||||
temperature: float = DEFAULT_TEMPERATURE
|
||||
max_tokens: int | None = None
|
||||
top_p: float = 1.0
|
||||
stop_sequences: List[str] = field(default_factory=list)
|
||||
extra_fields: dict = None # Converse API additionalModelRequestFields
|
||||
|
||||
_help: ClassVar[dict[str, str]] = {
|
||||
"temperature": "Controls randomness (0.0-1.0 for most Bedrock models)",
|
||||
"max_tokens": "Maximum tokens generated in the response (leave empty for model default)",
|
||||
"top_p": "Nucleus sampling parameter (0.0-1.0)",
|
||||
"stop_sequences": "Stop sequences (JSON array of strings, e.g., '[\"</s>\"]')",
|
||||
"extra_fields": 'Model-specific request fields forwarded as Converse API additionalModelRequestFields (JSON dict, e.g., \'{"reasoning_config": {"type": "enabled"}}\')',
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 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())
|
||||
@@ -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)
|
||||
@@ -0,0 +1,751 @@
|
||||
"""
|
||||
Gemini LLM binding for LightRAG.
|
||||
|
||||
This module provides asynchronous helpers that adapt Google's Gemini models
|
||||
to the same interface used by the rest of the LightRAG LLM bindings. The
|
||||
implementation mirrors the OpenAI helpers while relying on the official
|
||||
``google-genai`` client under the hood.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from collections.abc import AsyncIterator
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from tenacity import (
|
||||
retry,
|
||||
stop_after_attempt,
|
||||
wait_exponential,
|
||||
retry_if_exception_type,
|
||||
)
|
||||
|
||||
from lightrag.utils import (
|
||||
logger,
|
||||
remove_think_tags,
|
||||
safe_unicode_decode,
|
||||
wrap_embedding_func_with_attrs,
|
||||
)
|
||||
|
||||
import pipmaster as pm
|
||||
|
||||
# Install the Google Gemini client and its dependencies on demand
|
||||
if not pm.is_installed("google-genai"):
|
||||
pm.install("google-genai")
|
||||
if not pm.is_installed("google-api-core"):
|
||||
pm.install("google-api-core")
|
||||
|
||||
from google import genai # type: ignore
|
||||
from google.genai import types # type: ignore
|
||||
from google.api_core import exceptions as google_api_exceptions # type: ignore
|
||||
|
||||
|
||||
class InvalidResponseError(Exception):
|
||||
"""Custom exception class for triggering retry mechanism when Gemini returns empty responses"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
_DEFAULT_GEMINI_BASE_URLS = {
|
||||
"https://generativelanguage.googleapis.com",
|
||||
"https://generativelanguage.googleapis.com/",
|
||||
"https://generativelanguage.googleapis.com/v1beta",
|
||||
"https://generativelanguage.googleapis.com/v1beta/",
|
||||
"https://generativelanguage.googleapis.com/v1",
|
||||
"https://generativelanguage.googleapis.com/v1/",
|
||||
}
|
||||
|
||||
|
||||
def _normalize_gemini_base_url(base_url: str | None) -> str | None:
|
||||
"""Treat Google's default Gemini API service roots as SDK defaults."""
|
||||
if not base_url:
|
||||
return None
|
||||
|
||||
normalized = base_url.strip()
|
||||
if not normalized or normalized == "DEFAULT_GEMINI_ENDPOINT":
|
||||
return None
|
||||
|
||||
if normalized.rstrip("/") in {
|
||||
service_root.rstrip("/") for service_root in _DEFAULT_GEMINI_BASE_URLS
|
||||
}:
|
||||
return None
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _get_gemini_client(
|
||||
api_key: str, base_url: str | None, timeout: int | None = None
|
||||
) -> genai.Client:
|
||||
"""
|
||||
Create (or fetch cached) Gemini client.
|
||||
|
||||
Args:
|
||||
api_key: Google Gemini API key (not used in Vertex AI mode).
|
||||
base_url: Optional custom API endpoint.
|
||||
timeout: Optional request timeout in milliseconds.
|
||||
|
||||
Returns:
|
||||
genai.Client: Configured Gemini client instance.
|
||||
"""
|
||||
client_kwargs: dict[str, Any] = {}
|
||||
normalized_base_url = _normalize_gemini_base_url(base_url)
|
||||
|
||||
# Add Vertex AI support
|
||||
use_vertexai = os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
|
||||
if use_vertexai:
|
||||
# Vertex AI mode: use project/location, NOT api_key
|
||||
client_kwargs["vertexai"] = True
|
||||
project = os.getenv("GOOGLE_CLOUD_PROJECT")
|
||||
if project:
|
||||
location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
|
||||
client_kwargs["project"] = project
|
||||
if location:
|
||||
client_kwargs["location"] = location
|
||||
else:
|
||||
raise ValueError(
|
||||
"GOOGLE_CLOUD_PROJECT must be set when using Vertex AI mode"
|
||||
)
|
||||
else:
|
||||
# Standard Gemini API mode: use api_key
|
||||
client_kwargs["api_key"] = api_key
|
||||
|
||||
if normalized_base_url is not None or timeout is not None:
|
||||
try:
|
||||
http_options_kwargs = {}
|
||||
if normalized_base_url is not None:
|
||||
http_options_kwargs["base_url"] = normalized_base_url
|
||||
if timeout is not None:
|
||||
http_options_kwargs["timeout"] = timeout
|
||||
|
||||
client_kwargs["http_options"] = types.HttpOptions(**http_options_kwargs)
|
||||
except Exception as e:
|
||||
logger.error("Failed to apply custom Gemini http_options: %s", e)
|
||||
raise e
|
||||
|
||||
return genai.Client(**client_kwargs)
|
||||
|
||||
|
||||
def _ensure_api_key(api_key: str | None) -> str:
|
||||
# In Vertex AI mode, API key is not required
|
||||
use_vertexai = os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
|
||||
if use_vertexai:
|
||||
# Return empty string for Vertex AI mode (not used)
|
||||
return ""
|
||||
|
||||
key = api_key or os.getenv("LLM_BINDING_API_KEY") or os.getenv("GEMINI_API_KEY")
|
||||
if not key:
|
||||
raise ValueError(
|
||||
"Gemini API key not provided. "
|
||||
"Set LLM_BINDING_API_KEY or GEMINI_API_KEY in the environment."
|
||||
)
|
||||
return key
|
||||
|
||||
|
||||
def _build_generation_config(
|
||||
base_config: dict[str, Any] | None,
|
||||
system_prompt: str | None,
|
||||
response_format: Any | None,
|
||||
) -> types.GenerateContentConfig | None:
|
||||
config_data = dict(base_config or {})
|
||||
|
||||
if system_prompt:
|
||||
if config_data.get("system_instruction"):
|
||||
config_data["system_instruction"] = (
|
||||
f"{config_data['system_instruction']}\n{system_prompt}"
|
||||
)
|
||||
else:
|
||||
config_data["system_instruction"] = system_prompt
|
||||
|
||||
# Translate response_format to Gemini's native generation config fields.
|
||||
if response_format is not None:
|
||||
config_data.setdefault("response_mime_type", "application/json")
|
||||
schema = _normalize_gemini_response_schema(response_format)
|
||||
if schema is not None and "response_json_schema" not in config_data:
|
||||
config_data["response_json_schema"] = schema
|
||||
|
||||
# Remove entries that are explicitly set to None to avoid type errors
|
||||
sanitized = {
|
||||
key: value
|
||||
for key, value in config_data.items()
|
||||
if value is not None and value != ""
|
||||
}
|
||||
|
||||
if not sanitized:
|
||||
return None
|
||||
|
||||
return types.GenerateContentConfig(**sanitized)
|
||||
|
||||
|
||||
def _normalize_gemini_response_schema(response_format: Any) -> Any | None:
|
||||
"""Extract a Gemini-compatible JSON schema from LightRAG/OpenAI inputs."""
|
||||
if response_format is None:
|
||||
return None
|
||||
|
||||
if isinstance(response_format, dict):
|
||||
if response_format.get("type") == "json_object":
|
||||
return None
|
||||
|
||||
if response_format.get("type") == "json_schema":
|
||||
json_schema = response_format.get("json_schema")
|
||||
if isinstance(json_schema, dict):
|
||||
schema = json_schema.get("schema")
|
||||
if isinstance(schema, dict):
|
||||
return schema
|
||||
return json_schema
|
||||
|
||||
return response_format
|
||||
|
||||
return response_format
|
||||
|
||||
|
||||
def _validate_gemini_response_format(response_format: Any | None) -> None:
|
||||
"""Reject typed structured-output helpers; only dict payloads are supported."""
|
||||
if response_format is None or isinstance(response_format, dict):
|
||||
return
|
||||
|
||||
raise TypeError(
|
||||
"gemini_complete_if_cache only supports dict response_format payloads; "
|
||||
"typed/Pydantic response_format values are not supported."
|
||||
)
|
||||
|
||||
|
||||
def _format_history_messages(history_messages: list[dict[str, Any]] | None) -> str:
|
||||
if not history_messages:
|
||||
return ""
|
||||
|
||||
history_lines: list[str] = []
|
||||
for message in history_messages:
|
||||
role = message.get("role", "user")
|
||||
content = message.get("content", "")
|
||||
history_lines.append(f"[{role}] {content}")
|
||||
|
||||
return "\n".join(history_lines)
|
||||
|
||||
|
||||
def _extract_response_text(
|
||||
response: Any, extract_thoughts: bool = False
|
||||
) -> tuple[str, str]:
|
||||
"""
|
||||
Extract text content from Gemini response, separating regular content from thoughts.
|
||||
|
||||
Args:
|
||||
response: Gemini API response object
|
||||
extract_thoughts: Whether to extract thought content separately
|
||||
|
||||
Returns:
|
||||
Tuple of (regular_text, thought_text)
|
||||
"""
|
||||
candidates = getattr(response, "candidates", None)
|
||||
if not candidates:
|
||||
return ("", "")
|
||||
|
||||
regular_parts: list[str] = []
|
||||
thought_parts: list[str] = []
|
||||
|
||||
for candidate in candidates:
|
||||
if not getattr(candidate, "content", None):
|
||||
continue
|
||||
# Use 'or []' to handle None values from parts attribute
|
||||
for part in getattr(candidate.content, "parts", None) or []:
|
||||
text = getattr(part, "text", None)
|
||||
if not text:
|
||||
continue
|
||||
|
||||
# Check if this part is thought content using the 'thought' attribute
|
||||
is_thought = getattr(part, "thought", False)
|
||||
|
||||
if is_thought and extract_thoughts:
|
||||
thought_parts.append(text)
|
||||
elif not is_thought:
|
||||
regular_parts.append(text)
|
||||
|
||||
return ("\n".join(regular_parts), "\n".join(thought_parts))
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=60),
|
||||
retry=(
|
||||
retry_if_exception_type(google_api_exceptions.InternalServerError)
|
||||
| retry_if_exception_type(google_api_exceptions.ServiceUnavailable)
|
||||
| retry_if_exception_type(google_api_exceptions.ResourceExhausted)
|
||||
| retry_if_exception_type(google_api_exceptions.GatewayTimeout)
|
||||
| retry_if_exception_type(google_api_exceptions.BadGateway)
|
||||
| retry_if_exception_type(google_api_exceptions.DeadlineExceeded)
|
||||
| retry_if_exception_type(google_api_exceptions.Aborted)
|
||||
| retry_if_exception_type(google_api_exceptions.Unknown)
|
||||
| retry_if_exception_type(InvalidResponseError)
|
||||
),
|
||||
)
|
||||
async def gemini_complete_if_cache(
|
||||
model: str,
|
||||
prompt: str,
|
||||
system_prompt: str | None = None,
|
||||
history_messages: list[dict[str, Any]] | None = None,
|
||||
enable_cot: bool = False,
|
||||
base_url: str | None = None,
|
||||
api_key: str | None = None,
|
||||
token_tracker: Any | None = None,
|
||||
stream: bool | None = None,
|
||||
response_format: Any | None = None,
|
||||
keyword_extraction: bool = False,
|
||||
entity_extraction: bool = False,
|
||||
generation_config: dict[str, Any] | None = None,
|
||||
timeout: int | None = None,
|
||||
image_inputs: list[Any] | None = None,
|
||||
**_: Any,
|
||||
) -> str | AsyncIterator[str]:
|
||||
"""
|
||||
Complete a prompt using Gemini's API with Chain of Thought (COT) support.
|
||||
|
||||
This function supports automatic integration of reasoning content from Gemini models
|
||||
that provide Chain of Thought capabilities via the thinking_config API feature.
|
||||
|
||||
Structured output note:
|
||||
- This adapter accepts OpenAI-style ``response_format`` and translates it
|
||||
to Gemini's native generation config fields.
|
||||
- ``response_format={"type": "json_object"}`` maps to
|
||||
``response_mime_type="application/json"``.
|
||||
- Dict-form ``json_schema`` payloads map to
|
||||
``response_mime_type="application/json"`` plus
|
||||
``response_json_schema=<schema>``.
|
||||
- Typed/Pydantic ``response_format`` helpers are rejected explicitly.
|
||||
- Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are
|
||||
compatibility shims; when no explicit ``response_format`` is supplied,
|
||||
they are mapped to ``{"type": "json_object"}``.
|
||||
|
||||
COT Integration:
|
||||
- When enable_cot=True: Thought content is wrapped in <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.
|
||||
response_format: OpenAI-style structured output control translated to
|
||||
Gemini generation config. ``{"type": "json_object"}`` maps to
|
||||
``response_mime_type="application/json"``; dict-form
|
||||
``json_schema`` payloads map to ``response_json_schema``.
|
||||
Typed/Pydantic response_format values are rejected.
|
||||
token_tracker: Optional token usage tracker for monitoring API usage.
|
||||
stream: Whether to stream the response.
|
||||
hashing_kv: Storage interface (for interface parity with other bindings).
|
||||
enable_cot: Whether to include Chain of Thought content in the response.
|
||||
timeout: Request timeout in seconds (will be converted to milliseconds for Gemini API).
|
||||
**_: Additional keyword arguments (ignored).
|
||||
|
||||
Returns:
|
||||
The completed text (with COT content if enable_cot=True) or an async iterator
|
||||
of text chunks if streaming. COT content is wrapped in <think>...</think> tags.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the response from Gemini is empty.
|
||||
ValueError: If API key is not provided or configured.
|
||||
"""
|
||||
key = _ensure_api_key(api_key)
|
||||
# Convert timeout from seconds to milliseconds for Gemini API
|
||||
timeout_ms = timeout * 1000 if timeout else None
|
||||
client = _get_gemini_client(key, base_url, timeout_ms)
|
||||
|
||||
# Deprecation shims: map legacy boolean flags to response_format only when
|
||||
# an explicit response_format was not supplied.
|
||||
if response_format is None:
|
||||
if entity_extraction:
|
||||
warnings.warn(
|
||||
"gemini_complete_if_cache(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
response_format = {"type": "json_object"}
|
||||
elif keyword_extraction:
|
||||
warnings.warn(
|
||||
"gemini_complete_if_cache(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
response_format = {"type": "json_object"}
|
||||
_validate_gemini_response_format(response_format)
|
||||
if response_format is not None:
|
||||
enable_cot = False
|
||||
|
||||
history_block = _format_history_messages(history_messages)
|
||||
prompt_sections = []
|
||||
if history_block:
|
||||
prompt_sections.append(history_block)
|
||||
prompt_sections.append(f"[user] {prompt}")
|
||||
combined_prompt = "\n".join(prompt_sections)
|
||||
|
||||
config_obj = _build_generation_config(
|
||||
generation_config,
|
||||
system_prompt=system_prompt,
|
||||
response_format=response_format,
|
||||
)
|
||||
|
||||
if image_inputs:
|
||||
from lightrag.llm._vision_utils import normalize_image_inputs
|
||||
|
||||
normalized_images = normalize_image_inputs(image_inputs)
|
||||
parts: list[Any] = [combined_prompt]
|
||||
parts.extend(
|
||||
types.Part.from_bytes(data=img.raw_bytes, mime_type=img.mime_type)
|
||||
for img in normalized_images
|
||||
)
|
||||
contents: list[Any] = [parts]
|
||||
else:
|
||||
contents = [combined_prompt]
|
||||
|
||||
request_kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"contents": contents,
|
||||
}
|
||||
if config_obj is not None:
|
||||
request_kwargs["config"] = config_obj
|
||||
|
||||
if stream:
|
||||
|
||||
async def _async_stream() -> AsyncIterator[str]:
|
||||
# COT state tracking for streaming
|
||||
cot_active = False
|
||||
cot_started = False
|
||||
initial_content_seen = False
|
||||
usage_metadata = None
|
||||
|
||||
try:
|
||||
# Use native async streaming from genai SDK
|
||||
# Note: generate_content_stream returns Awaitable[AsyncIterator], need to await first
|
||||
stream_iter = await client.aio.models.generate_content_stream(
|
||||
**request_kwargs
|
||||
)
|
||||
async for chunk in stream_iter:
|
||||
usage = getattr(chunk, "usage_metadata", None)
|
||||
if usage is not None:
|
||||
usage_metadata = usage
|
||||
|
||||
# Extract both regular and thought content
|
||||
regular_text, thought_text = _extract_response_text(
|
||||
chunk, extract_thoughts=True
|
||||
)
|
||||
|
||||
if enable_cot:
|
||||
# Process regular content
|
||||
if regular_text:
|
||||
if not initial_content_seen:
|
||||
initial_content_seen = True
|
||||
|
||||
# Close COT section if it was active
|
||||
if cot_active:
|
||||
yield "</think>"
|
||||
cot_active = False
|
||||
|
||||
# Process and yield regular content
|
||||
if "\\u" in regular_text:
|
||||
regular_text = safe_unicode_decode(
|
||||
regular_text.encode("utf-8")
|
||||
)
|
||||
yield regular_text
|
||||
|
||||
# Process thought content
|
||||
if thought_text:
|
||||
if not initial_content_seen and not cot_started:
|
||||
# Start COT section
|
||||
yield "<think>"
|
||||
cot_active = True
|
||||
cot_started = True
|
||||
|
||||
# Yield thought content if COT is active
|
||||
if cot_active:
|
||||
if "\\u" in thought_text:
|
||||
thought_text = safe_unicode_decode(
|
||||
thought_text.encode("utf-8")
|
||||
)
|
||||
yield thought_text
|
||||
else:
|
||||
# COT disabled - only yield regular content
|
||||
if regular_text:
|
||||
if "\\u" in regular_text:
|
||||
regular_text = safe_unicode_decode(
|
||||
regular_text.encode("utf-8")
|
||||
)
|
||||
yield regular_text
|
||||
|
||||
# Ensure COT is properly closed if still active
|
||||
if cot_active:
|
||||
yield "</think>"
|
||||
cot_active = False
|
||||
|
||||
except Exception:
|
||||
# Try to close COT tag before re-raising
|
||||
if cot_active:
|
||||
try:
|
||||
yield "</think>"
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
finally:
|
||||
# Track token usage after streaming completes
|
||||
if token_tracker and usage_metadata:
|
||||
token_tracker.add_usage(
|
||||
{
|
||||
"prompt_tokens": getattr(
|
||||
usage_metadata, "prompt_token_count", 0
|
||||
),
|
||||
"completion_tokens": getattr(
|
||||
usage_metadata, "candidates_token_count", 0
|
||||
),
|
||||
"total_tokens": getattr(
|
||||
usage_metadata, "total_token_count", 0
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
return _async_stream()
|
||||
|
||||
# Non-streaming: use native async client
|
||||
response = await client.aio.models.generate_content(**request_kwargs)
|
||||
|
||||
# Extract both regular text and thought text
|
||||
regular_text, thought_text = _extract_response_text(response, extract_thoughts=True)
|
||||
|
||||
# Apply COT filtering logic based on enable_cot parameter
|
||||
if enable_cot:
|
||||
# Include thought content wrapped in <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,
|
||||
response_format: Any | None = None,
|
||||
keyword_extraction: bool = False,
|
||||
entity_extraction: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> str | AsyncIterator[str]:
|
||||
# Accept legacy keyword if passed via kwargs to preserve backwards compat.
|
||||
entity_extraction = kwargs.pop("entity_extraction", entity_extraction)
|
||||
hashing_kv = kwargs.get("hashing_kv")
|
||||
model_name = None
|
||||
if hashing_kv is not None:
|
||||
model_name = hashing_kv.global_config.get("llm_model_name")
|
||||
if model_name is None:
|
||||
model_name = kwargs.pop("model_name", None)
|
||||
if model_name is None:
|
||||
raise ValueError("Gemini model name not provided in configuration.")
|
||||
|
||||
return await gemini_complete_if_cache(
|
||||
model_name,
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
response_format=response_format,
|
||||
keyword_extraction=keyword_extraction,
|
||||
entity_extraction=entity_extraction,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@wrap_embedding_func_with_attrs(
|
||||
embedding_dim=1536,
|
||||
max_token_size=2048,
|
||||
model_name="gemini-embedding-001",
|
||||
supports_asymmetric=True,
|
||||
)
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=60),
|
||||
retry=(
|
||||
retry_if_exception_type(google_api_exceptions.InternalServerError)
|
||||
| retry_if_exception_type(google_api_exceptions.ServiceUnavailable)
|
||||
| retry_if_exception_type(google_api_exceptions.ResourceExhausted)
|
||||
| retry_if_exception_type(google_api_exceptions.GatewayTimeout)
|
||||
| retry_if_exception_type(google_api_exceptions.BadGateway)
|
||||
| retry_if_exception_type(google_api_exceptions.DeadlineExceeded)
|
||||
| retry_if_exception_type(google_api_exceptions.Aborted)
|
||||
| retry_if_exception_type(google_api_exceptions.Unknown)
|
||||
),
|
||||
)
|
||||
async def gemini_embed(
|
||||
texts: list[str],
|
||||
model: str = "gemini-embedding-001",
|
||||
base_url: str | None = None,
|
||||
api_key: str | None = None,
|
||||
embedding_dim: int | None = None,
|
||||
max_token_size: int | None = None,
|
||||
task_type: str | None = None,
|
||||
timeout: int | None = None,
|
||||
token_tracker: Any | None = None,
|
||||
context: str = "document",
|
||||
) -> np.ndarray:
|
||||
"""Generate embeddings for a list of texts using Gemini's API.
|
||||
|
||||
This function uses Google's Gemini embedding model to generate text embeddings.
|
||||
It supports dynamic dimension control and automatic normalization for dimensions
|
||||
less than 3072.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed.
|
||||
model: The Gemini embedding model to use. Default is "gemini-embedding-001".
|
||||
base_url: Optional custom API endpoint.
|
||||
api_key: Optional Gemini API key. If None, uses environment variables.
|
||||
embedding_dim: Optional embedding dimension for dynamic dimension reduction.
|
||||
**IMPORTANT**: This parameter is automatically injected by the EmbeddingFunc wrapper.
|
||||
Do NOT manually pass this parameter when calling the function directly.
|
||||
The dimension is controlled by the @wrap_embedding_func_with_attrs decorator
|
||||
or the EMBEDDING_DIM environment variable.
|
||||
Supported range: 128-3072. Recommended values: 768, 1536, 3072.
|
||||
max_token_size: Maximum tokens per text. This parameter is automatically
|
||||
injected by the EmbeddingFunc wrapper when the underlying function
|
||||
signature supports it (via inspect.signature check). Gemini API will
|
||||
automatically truncate texts exceeding this limit (autoTruncate=True
|
||||
by default), so no client-side truncation is needed.
|
||||
task_type: Task type for embedding optimization. Default is "RETRIEVAL_DOCUMENT".
|
||||
Supported types: SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING,
|
||||
RETRIEVAL_DOCUMENT, RETRIEVAL_QUERY, CODE_RETRIEVAL_QUERY,
|
||||
QUESTION_ANSWERING, FACT_VERIFICATION.
|
||||
timeout: Request timeout in seconds (will be converted to milliseconds for Gemini API).
|
||||
token_tracker: Optional token usage tracker for monitoring API usage.
|
||||
context: The embedding context - "query" for search queries, "document" for indexed content.
|
||||
**IMPORTANT**: This parameter is automatically injected by the EmbeddingFunc wrapper
|
||||
when supports_asymmetric=True. Default is "document".
|
||||
|
||||
Returns:
|
||||
A numpy array of embeddings, one per input text. For dimensions < 3072,
|
||||
the embeddings are L2-normalized to ensure optimal semantic similarity performance.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided or configured.
|
||||
RuntimeError: If the response from Gemini is invalid or empty.
|
||||
|
||||
Note:
|
||||
- For dimension 3072: Embeddings are already normalized by the API
|
||||
- For dimensions < 3072: Embeddings are L2-normalized after retrieval
|
||||
- Normalization ensures accurate semantic similarity via cosine distance
|
||||
- Gemini API automatically truncates texts exceeding max_token_size (autoTruncate=True)
|
||||
"""
|
||||
# Note: max_token_size is received but not used for client-side truncation.
|
||||
# Gemini API handles truncation automatically with autoTruncate=True (default).
|
||||
_ = max_token_size # Acknowledge parameter to avoid unused variable warning
|
||||
|
||||
key = _ensure_api_key(api_key)
|
||||
# Convert timeout from seconds to milliseconds for Gemini API
|
||||
timeout_ms = timeout * 1000 if timeout else None
|
||||
client = _get_gemini_client(key, base_url, timeout_ms)
|
||||
|
||||
# Prepare embedding configuration
|
||||
config_kwargs: dict[str, Any] = {}
|
||||
|
||||
# Add task_type to config
|
||||
if task_type is None:
|
||||
if context == "query":
|
||||
task_type = "RETRIEVAL_QUERY"
|
||||
elif context == "document":
|
||||
task_type = "RETRIEVAL_DOCUMENT"
|
||||
else:
|
||||
task_type = "RETRIEVAL_DOCUMENT" # Default for backward compatibility
|
||||
config_kwargs["task_type"] = task_type
|
||||
|
||||
# Add output_dimensionality if embedding_dim is provided
|
||||
if embedding_dim is not None:
|
||||
config_kwargs["output_dimensionality"] = embedding_dim
|
||||
|
||||
# Create config object if we have parameters
|
||||
config_obj = types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
request_kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"contents": texts,
|
||||
}
|
||||
if config_obj is not None:
|
||||
request_kwargs["config"] = config_obj
|
||||
|
||||
# Use native async client for embedding
|
||||
response = await client.aio.models.embed_content(**request_kwargs)
|
||||
|
||||
# Extract embeddings from response
|
||||
if not hasattr(response, "embeddings") or not response.embeddings:
|
||||
raise RuntimeError("Gemini response did not contain embeddings.")
|
||||
|
||||
# Convert embeddings to numpy array
|
||||
embeddings = np.array(
|
||||
[np.array(e.values, dtype=np.float32) for e in response.embeddings]
|
||||
)
|
||||
|
||||
# Apply L2 normalization for dimensions < 3072
|
||||
# The 3072 dimension embedding is already normalized by Gemini API
|
||||
if embedding_dim and embedding_dim < 3072:
|
||||
# Normalize each embedding vector to unit length
|
||||
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||
# Avoid division by zero
|
||||
norms = np.where(norms == 0, 1, norms)
|
||||
embeddings = embeddings / norms
|
||||
logger.debug(
|
||||
f"Applied L2 normalization to {len(embeddings)} embeddings of dimension {embedding_dim}"
|
||||
)
|
||||
|
||||
# Track token usage if tracker is provided
|
||||
# Note: Gemini embedding API may not provide usage metadata
|
||||
if token_tracker and hasattr(response, "usage_metadata"):
|
||||
usage = response.usage_metadata
|
||||
token_counts = {
|
||||
"prompt_tokens": getattr(usage, "prompt_token_count", 0),
|
||||
"total_tokens": getattr(usage, "total_token_count", 0),
|
||||
}
|
||||
token_tracker.add_usage(token_counts)
|
||||
|
||||
logger.debug(
|
||||
f"Generated {len(embeddings)} Gemini embeddings with dimension {embeddings.shape[1]}"
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
__all__ = [
|
||||
"gemini_complete_if_cache",
|
||||
"gemini_model_complete",
|
||||
"gemini_embed",
|
||||
]
|
||||
@@ -0,0 +1,232 @@
|
||||
import copy
|
||||
import os
|
||||
import warnings
|
||||
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,
|
||||
entity_extraction=False,
|
||||
enable_cot: bool = False,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""Run local Hugging Face inference with LightRAG-compatible shims.
|
||||
|
||||
Structured output note:
|
||||
- This adapter does not support OpenAI-style ``response_format`` JSON mode.
|
||||
- If callers pass ``response_format``, it is stripped before generation.
|
||||
- Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are
|
||||
accepted only as compatibility shims; they emit warnings and are ignored.
|
||||
"""
|
||||
# HuggingFace local inference has no JSON mode; drop response_format and
|
||||
# warn when legacy shim flags are set.
|
||||
if kwargs.pop("keyword_extraction", False) or keyword_extraction:
|
||||
warnings.warn(
|
||||
"hf_model_complete(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if kwargs.pop("entity_extraction", False) or entity_extraction:
|
||||
warnings.warn(
|
||||
"hf_model_complete(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs.pop("response_format", 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,
|
||||
model_name="hf_embedding_model",
|
||||
supports_asymmetric=True,
|
||||
)
|
||||
async def hf_embed(
|
||||
texts: list[str],
|
||||
tokenizer,
|
||||
embed_model,
|
||||
context: str = "document",
|
||||
query_prefix: str | None = None,
|
||||
document_prefix: str | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Generate embeddings for a list of texts using a Hugging Face model.
|
||||
|
||||
Args:
|
||||
texts (list[str]): List of input texts to embed.
|
||||
tokenizer: Hugging Face tokenizer.
|
||||
embed_model: Hugging Face model for generating embeddings.
|
||||
context (str): Context indicating whether the texts are "query" or "document".
|
||||
query_prefix (str | None): Optional prefix to add to query texts.
|
||||
document_prefix (str | None): Optional prefix to add to document texts.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Array of embeddings.
|
||||
"""
|
||||
# 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)
|
||||
|
||||
# Apply context-based prefixes if provided
|
||||
if context == "query" and query_prefix:
|
||||
texts = [query_prefix + text for text in texts]
|
||||
elif context == "document" and document_prefix:
|
||||
texts = [document_prefix + text for text in texts]
|
||||
|
||||
# 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()
|
||||
@@ -0,0 +1,183 @@
|
||||
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,
|
||||
model_name="jina-embeddings-v4",
|
||||
supports_asymmetric=True,
|
||||
)
|
||||
@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],
|
||||
model: str = "jina-embeddings-v4",
|
||||
embedding_dim: int = 2048,
|
||||
late_chunking: bool = False,
|
||||
base_url: str = None,
|
||||
api_key: str = None,
|
||||
context: str | None = None,
|
||||
task: str | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Generate embeddings for a list of texts using Jina AI's API.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed.
|
||||
model: The Jina embedding model to use (default: jina-embeddings-v4).
|
||||
Supported models: jina-embeddings-v3, jina-embeddings-v4, etc.
|
||||
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.
|
||||
context: The embedding context - "query" for search queries, "document" for indexed content.
|
||||
**IMPORTANT**: This parameter is automatically injected by the EmbeddingFunc wrapper
|
||||
when supports_asymmetric=True. When ``task`` is left at its default of None,
|
||||
``context`` drives the task selection.
|
||||
task: Embedding task mode. Default is None so that ``context`` (when present)
|
||||
picks the right Jina task:
|
||||
- "retrieval.query" for context="query"
|
||||
- "retrieval.passage" for context="document"
|
||||
- "text-matching" otherwise (true backward-compatible default)
|
||||
Any explicit non-None task value overrides context-based selection.
|
||||
|
||||
|
||||
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']}",
|
||||
}
|
||||
|
||||
# Determine task based on context if not explicitly provided
|
||||
if task is None:
|
||||
if context == "query":
|
||||
task = "retrieval.query"
|
||||
elif context == "document":
|
||||
task = "retrieval.passage"
|
||||
else:
|
||||
task = "text-matching" # Default for backward compatibility
|
||||
|
||||
data = {
|
||||
"model": model,
|
||||
"task": task,
|
||||
"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
|
||||
@@ -0,0 +1,235 @@
|
||||
import warnings
|
||||
|
||||
import pipmaster as pm
|
||||
from llama_index.core.llms import (
|
||||
ChatMessage,
|
||||
MessageRole,
|
||||
ChatResponse,
|
||||
)
|
||||
from typing import Any, 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: Any = 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,
|
||||
entity_extraction=False,
|
||||
settings: Any = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""
|
||||
Main completion function for LlamaIndex.
|
||||
|
||||
Args:
|
||||
prompt: Input prompt
|
||||
system_prompt: Optional system prompt
|
||||
history_messages: Optional chat history
|
||||
keyword_extraction: Deprecated compatibility shim. Emits a warning and
|
||||
is ignored.
|
||||
entity_extraction: Deprecated compatibility shim. Emits a warning and
|
||||
is ignored.
|
||||
settings: Optional LlamaIndex settings
|
||||
**kwargs: Additional arguments. ``response_format`` is not supported by
|
||||
this adapter and is stripped before calling LlamaIndex.
|
||||
|
||||
Structured output note:
|
||||
- This adapter does not support OpenAI-style ``response_format`` JSON mode.
|
||||
- If callers pass ``response_format``, it is stripped before generation.
|
||||
"""
|
||||
if history_messages is None:
|
||||
history_messages = []
|
||||
|
||||
# LlamaIndex adapters have no JSON mode; drop response_format and warn
|
||||
# when legacy boolean shim flags are set.
|
||||
if kwargs.pop("keyword_extraction", False) or keyword_extraction:
|
||||
warnings.warn(
|
||||
"llama_index_complete(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if kwargs.pop("entity_extraction", False) or entity_extraction:
|
||||
warnings.warn(
|
||||
"llama_index_complete(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs.pop("response_format", 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: Any = 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)
|
||||
@@ -0,0 +1,179 @@
|
||||
import warnings
|
||||
|
||||
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:
|
||||
"""Run lmdeploy generation with LightRAG-compatible shims.
|
||||
|
||||
Structured output note:
|
||||
- This adapter does not support OpenAI-style ``response_format`` JSON mode.
|
||||
- If callers pass ``response_format``, it is stripped before generation.
|
||||
- Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are
|
||||
accepted only as compatibility shims; they emit warnings and are ignored.
|
||||
|
||||
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)
|
||||
# lmdeploy has no JSON mode; drop response_format and warn when legacy
|
||||
# boolean shim flags are set.
|
||||
if kwargs.pop("keyword_extraction", False):
|
||||
warnings.warn(
|
||||
"lmdeploy_model_if_cache(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if kwargs.pop("entity_extraction", False):
|
||||
warnings.warn(
|
||||
"lmdeploy_model_if_cache(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
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
|
||||
@@ -0,0 +1,212 @@
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
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 Any, List, Union
|
||||
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",
|
||||
image_inputs: list[Any] | None = None,
|
||||
**kwargs,
|
||||
) -> Union[str, AsyncIterator[str]]:
|
||||
"""Client implementation for lollms generation.
|
||||
|
||||
Structured output note:
|
||||
- This adapter does not support OpenAI-style ``response_format`` JSON mode.
|
||||
- If callers pass ``response_format``, it is stripped before the request.
|
||||
- Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are
|
||||
accepted only as compatibility shims; they emit warnings and are ignored.
|
||||
|
||||
Vision note:
|
||||
- lollms does not support image inputs. Passing a non-empty
|
||||
``image_inputs`` raises :class:`NotImplementedError`.
|
||||
"""
|
||||
if image_inputs:
|
||||
raise NotImplementedError(
|
||||
"lollms binding does not support image_inputs; configure a "
|
||||
"vision-capable VLM provider (openai/azure_openai/gemini/bedrock/"
|
||||
"ollama/anthropic) for VLM_LLM_BINDING."
|
||||
)
|
||||
|
||||
if enable_cot:
|
||||
from lightrag.utils import logger
|
||||
|
||||
logger.debug("enable_cot=True is not supported for lollms and will be ignored.")
|
||||
|
||||
# lollms has no JSON mode; drop response_format and warn when legacy
|
||||
# boolean shim flags are set.
|
||||
if kwargs.pop("keyword_extraction", False):
|
||||
warnings.warn(
|
||||
"lollms_model_if_cache(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if kwargs.pop("entity_extraction", False):
|
||||
warnings.warn(
|
||||
"lollms_model_if_cache(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs.pop("response_format", None)
|
||||
|
||||
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,
|
||||
entity_extraction=False,
|
||||
**kwargs,
|
||||
) -> Union[str, AsyncIterator[str]]:
|
||||
"""Complete function for lollms model generation."""
|
||||
|
||||
# Forward legacy extraction flags as kwargs so lollms_model_if_cache can
|
||||
# emit a single DeprecationWarning with the correct stack frame.
|
||||
if keyword_extraction:
|
||||
kwargs.setdefault("keyword_extraction", True)
|
||||
if entity_extraction:
|
||||
kwargs.setdefault("entity_extraction", True)
|
||||
model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
|
||||
|
||||
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, model_name="lollms_embedding_model"
|
||||
)
|
||||
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)
|
||||
@@ -0,0 +1,72 @@
|
||||
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, model_name="nvidia_embedding_model"
|
||||
)
|
||||
@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)
|
||||
)
|
||||
# Hold the client in an async-with so its httpx connection pool is
|
||||
# released on every exit path (success, error, and each @retry attempt),
|
||||
# instead of leaking one pool per call until GC. Mirrors ``openai_embed``.
|
||||
async with openai_async_client:
|
||||
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])
|
||||
@@ -0,0 +1,334 @@
|
||||
from collections.abc import AsyncIterator
|
||||
import os
|
||||
import re
|
||||
import warnings
|
||||
|
||||
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 Any, 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
|
||||
|
||||
|
||||
def _normalize_ollama_response_format(kwargs: dict) -> None:
|
||||
"""Translate OpenAI-style response_format into Ollama's native format field.
|
||||
|
||||
Precedence: an explicit ``format`` value (Ollama's native field) wins over
|
||||
``response_format`` — if ``format`` is already set, ``response_format`` is
|
||||
dropped silently. Otherwise, ``{"type": "json_object"}`` maps to
|
||||
``format="json"`` and any other payload is passed through unchanged so
|
||||
callers can supply JSON schemas directly.
|
||||
"""
|
||||
|
||||
response_format = kwargs.pop("response_format", None)
|
||||
if kwargs.get("format") is not None or response_format is None:
|
||||
return
|
||||
|
||||
if isinstance(response_format, dict):
|
||||
if response_format.get("type") == "json_object":
|
||||
kwargs["format"] = "json"
|
||||
return
|
||||
if response_format.get("type") == "json_schema":
|
||||
json_schema = response_format.get("json_schema")
|
||||
if isinstance(json_schema, dict):
|
||||
kwargs["format"] = json_schema.get("schema", json_schema)
|
||||
return
|
||||
|
||||
# Fall back to passing through schema-like payloads for native Ollama support.
|
||||
kwargs["format"] = response_format
|
||||
|
||||
|
||||
@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,
|
||||
image_inputs: list[Any] | None = None,
|
||||
**kwargs,
|
||||
) -> Union[str, AsyncIterator[str]]:
|
||||
"""Call Ollama chat API with OpenAI-style structured-output compatibility.
|
||||
|
||||
Structured output note:
|
||||
- This adapter accepts OpenAI-style ``response_format`` and translates it
|
||||
to Ollama's native ``format`` field.
|
||||
- ``response_format={"type": "json_object"}`` maps to ``format="json"``.
|
||||
- Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are
|
||||
compatibility shims; when no explicit ``response_format`` is supplied,
|
||||
they are mapped to ``{"type": "json_object"}``.
|
||||
"""
|
||||
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)
|
||||
# Deprecation shims: map legacy boolean flags to response_format only when
|
||||
# an explicit response_format was not supplied by the caller.
|
||||
if kwargs.get("response_format") is None:
|
||||
if kwargs.pop("entity_extraction", False):
|
||||
warnings.warn(
|
||||
"_ollama_model_if_cache(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
elif kwargs.pop("keyword_extraction", False):
|
||||
warnings.warn(
|
||||
"_ollama_model_if_cache(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
else:
|
||||
# response_format was supplied explicitly; drop legacy flags silently.
|
||||
kwargs.pop("entity_extraction", None)
|
||||
kwargs.pop("keyword_extraction", None)
|
||||
|
||||
_normalize_ollama_response_format(kwargs)
|
||||
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)
|
||||
user_message: dict[str, Any] = {"role": "user", "content": prompt}
|
||||
if image_inputs:
|
||||
from lightrag.llm._vision_utils import normalize_image_inputs
|
||||
|
||||
normalized_images = normalize_image_inputs(image_inputs)
|
||||
user_message["images"] = [img.base64_str for img in normalized_images]
|
||||
messages.append(user_message)
|
||||
|
||||
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,
|
||||
entity_extraction=False,
|
||||
**kwargs,
|
||||
) -> Union[str, AsyncIterator[str]]:
|
||||
# Forward legacy extraction flags as kwargs so _ollama_model_if_cache can
|
||||
# emit a single DeprecationWarning with the correct stack frame.
|
||||
if keyword_extraction:
|
||||
kwargs.setdefault("keyword_extraction", True)
|
||||
if entity_extraction:
|
||||
kwargs.setdefault("entity_extraction", True)
|
||||
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,
|
||||
model_name="bge-m3:latest",
|
||||
supports_asymmetric=True,
|
||||
)
|
||||
async def ollama_embed(
|
||||
texts: list[str],
|
||||
embed_model: str = "bge-m3:latest",
|
||||
max_token_size: int | None = None,
|
||||
context: str = "document",
|
||||
query_prefix: str | None = None,
|
||||
document_prefix: str | None = None,
|
||||
**kwargs,
|
||||
) -> np.ndarray:
|
||||
"""Generate embeddings using Ollama's API.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed.
|
||||
embed_model: The Ollama embedding model to use. Default is "bge-m3:latest".
|
||||
max_token_size: Maximum tokens per text. This parameter is automatically
|
||||
injected by the EmbeddingFunc wrapper when the underlying function
|
||||
signature supports it (via inspect.signature check). Ollama will
|
||||
automatically truncate texts exceeding the model's context length
|
||||
(num_ctx), so no client-side truncation is needed.
|
||||
context: The embedding context - "query" for search queries, "document" for indexed content.
|
||||
**IMPORTANT**: This parameter is automatically injected by the EmbeddingFunc wrapper
|
||||
when supports_asymmetric=True. Default is "document".
|
||||
query_prefix: Optional prefix to prepend to texts when context="query" (e.g., "search_query: ").
|
||||
document_prefix: Optional prefix to prepend to texts when context="document" (e.g., "search_document: ").
|
||||
**kwargs: Additional arguments passed to the Ollama client.
|
||||
|
||||
Returns:
|
||||
A numpy array of embeddings, one per input text.
|
||||
|
||||
Note:
|
||||
- Ollama API automatically truncates texts exceeding the model's context length
|
||||
- The max_token_size parameter is received but not used for client-side truncation
|
||||
"""
|
||||
# Apply context-based prefixes if provided
|
||||
if context == "query" and query_prefix:
|
||||
texts = [query_prefix + text for text in texts]
|
||||
elif context == "document" and document_prefix:
|
||||
texts = [document_prefix + text for text in texts]
|
||||
|
||||
# Note: max_token_size is received but not used for client-side truncation.
|
||||
# Ollama API handles truncation automatically based on the model's num_ctx setting.
|
||||
_ = max_token_size # Acknowledge parameter to avoid unused variable warning
|
||||
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}")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,197 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import pipmaster as pm # Pipmaster for dynamic library install
|
||||
|
||||
# Add Voyage AI import
|
||||
if not pm.is_installed("voyageai"):
|
||||
pm.install("voyageai")
|
||||
|
||||
import voyageai
|
||||
from voyageai.error import (
|
||||
APIConnectionError,
|
||||
RateLimitError,
|
||||
ServerError,
|
||||
ServiceUnavailableError,
|
||||
Timeout,
|
||||
TryAgain,
|
||||
)
|
||||
|
||||
from tenacity import (
|
||||
retry,
|
||||
stop_after_attempt,
|
||||
wait_exponential,
|
||||
retry_if_exception_type,
|
||||
)
|
||||
from lightrag.utils import wrap_embedding_func_with_attrs, logger
|
||||
|
||||
|
||||
# Custom exceptions for VoyageAI errors
|
||||
class VoyageAIError(Exception):
|
||||
"""Generic VoyageAI API error"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@wrap_embedding_func_with_attrs(
|
||||
embedding_dim=1024, max_token_size=32000, supports_asymmetric=True
|
||||
)
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=60),
|
||||
retry=retry_if_exception_type(
|
||||
(
|
||||
APIConnectionError,
|
||||
RateLimitError,
|
||||
ServerError,
|
||||
ServiceUnavailableError,
|
||||
Timeout,
|
||||
TryAgain,
|
||||
)
|
||||
),
|
||||
)
|
||||
async def voyageai_embed(
|
||||
texts: list[str],
|
||||
model: str = "voyage-3",
|
||||
api_key: str | None = None,
|
||||
embedding_dim: int | None = None,
|
||||
input_type: str | None = None,
|
||||
truncation: bool | None = True,
|
||||
context: str | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Generate embeddings for a list of texts using VoyageAI's API.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed.
|
||||
model: The VoyageAI embedding model to use. Options include:
|
||||
- "voyage-3": General purpose (1024 dims, 32K context)
|
||||
- "voyage-3-lite": Lightweight (512 dims, 32K context)
|
||||
- "voyage-3-large": Highest accuracy (1024 dims, 32K context)
|
||||
- "voyage-code-3": Code optimized (1024 dims, 32K context)
|
||||
- "voyage-law-2": Legal documents (1024 dims, 16K context)
|
||||
- "voyage-finance-2": Finance (1024 dims, 32K context)
|
||||
api_key: Optional VoyageAI API key. If None, falls back to the
|
||||
``VOYAGE_API_KEY`` environment variable (the name VoyageAI's own
|
||||
SDK uses), then to ``VOYAGEAI_API_KEY`` for backward compatibility.
|
||||
embedding_dim: Optional Matryoshka output dimension. Only honored by
|
||||
models that support dimension reduction (e.g. voyage-3-large);
|
||||
ignored otherwise. The decorator default is 1024 to match
|
||||
``voyage-3``; if you select ``voyage-3-lite`` (512 dims) override
|
||||
``EMBEDDING_DIM`` accordingly so the vector store size matches.
|
||||
input_type: Optional input type hint for the model. Options:
|
||||
- "query": For search queries
|
||||
- "document": For documents to be indexed
|
||||
- None: Let the model decide (default)
|
||||
truncation: Whether the API should truncate texts that exceed the model's
|
||||
token limit. Defaults to True (matches the VoyageAI SDK default).
|
||||
context: Optional LightRAG embedding context. When ``input_type`` is not
|
||||
set, "query" maps to ``input_type="query"`` and "document" maps to
|
||||
``input_type="document"``.
|
||||
|
||||
Returns:
|
||||
A numpy array of embeddings, one per input text.
|
||||
|
||||
Raises:
|
||||
VoyageAIError: If the API call fails or returns invalid data.
|
||||
|
||||
"""
|
||||
if not api_key:
|
||||
api_key = os.environ.get("VOYAGE_API_KEY") or os.environ.get("VOYAGEAI_API_KEY")
|
||||
if not api_key:
|
||||
logger.error(
|
||||
"VoyageAI API key not provided and neither VOYAGE_API_KEY nor "
|
||||
"VOYAGEAI_API_KEY environment variable is set"
|
||||
)
|
||||
raise ValueError(
|
||||
"VoyageAI API key is required: pass api_key, or set the "
|
||||
"VOYAGE_API_KEY (preferred) or VOYAGEAI_API_KEY environment variable"
|
||||
)
|
||||
|
||||
if input_type is None and context in {"query", "document"}:
|
||||
input_type = context
|
||||
|
||||
try:
|
||||
client = voyageai.AsyncClient(api_key=api_key)
|
||||
|
||||
total_chars = sum(len(t) for t in texts)
|
||||
avg_chars = total_chars / len(texts) if texts else 0
|
||||
logger.debug(
|
||||
f"VoyageAI embedding request: {len(texts)} texts, "
|
||||
f"total_chars={total_chars}, avg_chars={avg_chars:.0f}, model={model}, "
|
||||
f"input_type={input_type}"
|
||||
)
|
||||
|
||||
# Prepare API call parameters
|
||||
embed_params = dict(
|
||||
texts=texts,
|
||||
model=model,
|
||||
# Optional parameters -- if None, voyageai client uses defaults
|
||||
output_dimension=embedding_dim,
|
||||
truncation=truncation,
|
||||
input_type=input_type,
|
||||
)
|
||||
# Make API call with timing
|
||||
result = await client.embed(**embed_params)
|
||||
|
||||
if not result.embeddings:
|
||||
err_msg = "VoyageAI API returned empty embeddings"
|
||||
logger.error(err_msg)
|
||||
raise VoyageAIError(err_msg)
|
||||
|
||||
if len(result.embeddings) != len(texts):
|
||||
err_msg = f"VoyageAI API returned {len(result.embeddings)} embeddings for {len(texts)} texts"
|
||||
logger.error(err_msg)
|
||||
raise VoyageAIError(err_msg)
|
||||
|
||||
# Convert to numpy array with timing
|
||||
embeddings = np.array(result.embeddings, dtype=np.float32)
|
||||
logger.debug(f"VoyageAI embeddings generated: shape {embeddings.shape}")
|
||||
|
||||
return embeddings
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"VoyageAI embedding error: {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",
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
import sys
|
||||
import warnings
|
||||
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,
|
||||
)
|
||||
|
||||
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, # LightRAG output switch: include reasoning_content as <think>...</think>
|
||||
thinking: Optional[
|
||||
Dict[str, object]
|
||||
] = None, # Zhipu request param: use {"type": "enabled"} to enable thinking
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""Call Zhipu chat completions with optional official thinking support.
|
||||
|
||||
Parameter roles:
|
||||
- `thinking`: forwarded to the Zhipu API as-is. To enable thinking output,
|
||||
pass a config such as `{"type": "enabled"}`.
|
||||
- `enable_cot`: LightRAG-only formatting switch. When True and the API
|
||||
returns `reasoning_content`, it is preserved in the final string as
|
||||
`<think>...</think>`.
|
||||
- `response_format`: forwarded as Zhipu's OpenAI-compatible structured
|
||||
output parameter when supplied by callers.
|
||||
- Deprecated `keyword_extraction` and `entity_extraction` booleans are
|
||||
compatibility shims; when no explicit `response_format` is supplied,
|
||||
they are mapped to `{"type": "json_object"}`.
|
||||
"""
|
||||
# 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}")
|
||||
|
||||
# Deprecation shims: map legacy extraction booleans to response_format only
|
||||
# when an explicit response_format was not supplied by the caller. The
|
||||
# legacy path also forces enable_cot=False so reasoning_content cannot
|
||||
# corrupt the JSON payload expected by callers relying on it.
|
||||
keyword_extraction = kwargs.pop("keyword_extraction", False)
|
||||
entity_extraction = kwargs.pop("entity_extraction", False)
|
||||
if kwargs.get("response_format") is None:
|
||||
if entity_extraction:
|
||||
warnings.warn(
|
||||
"zhipu_complete_if_cache(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
enable_cot = False
|
||||
elif keyword_extraction:
|
||||
warnings.warn(
|
||||
"zhipu_complete_if_cache(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
enable_cot = False
|
||||
|
||||
# Structured output and COT are mutually exclusive here because
|
||||
# reasoning_content would corrupt the JSON payload expected by callers.
|
||||
if kwargs.get("response_format") is not None:
|
||||
enable_cot = False
|
||||
|
||||
# Remove unsupported kwargs
|
||||
kwargs = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if k not in ["hashing_kv", "keyword_extraction", "entity_extraction"]
|
||||
}
|
||||
# `thinking` is an official Zhipu request field. Example:
|
||||
# {"type": "enabled"} enables reasoning output on supported models.
|
||||
if thinking is not None:
|
||||
kwargs["thinking"] = thinking
|
||||
|
||||
response = client.chat.completions.create(model=model, messages=messages, **kwargs)
|
||||
if not response.choices or response.choices[0].message is None:
|
||||
return ""
|
||||
message = response.choices[0].message
|
||||
content = message.content or ""
|
||||
reasoning_content = getattr(message, "reasoning_content", "") or ""
|
||||
|
||||
if enable_cot and reasoning_content.strip():
|
||||
if content:
|
||||
return f"<think>{reasoning_content}</think>{content}"
|
||||
return f"<think>{reasoning_content}</think>"
|
||||
|
||||
return content
|
||||
|
||||
|
||||
async def zhipu_complete(
|
||||
prompt,
|
||||
system_prompt=None,
|
||||
history_messages=[],
|
||||
keyword_extraction=False,
|
||||
entity_extraction=False,
|
||||
enable_cot: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
"""Zhipu completion wrapper with LightRAG structured-output shims.
|
||||
|
||||
Structured output note:
|
||||
- This adapter accepts OpenAI-style ``response_format`` and forwards it to
|
||||
Zhipu's compatible chat-completions API.
|
||||
- Deprecated ``keyword_extraction`` and ``entity_extraction`` booleans are
|
||||
compatibility shims; when no explicit ``response_format`` is supplied,
|
||||
they are mapped to ``{"type": "json_object"}``.
|
||||
"""
|
||||
# Pop legacy extraction flags from kwargs to avoid passing them downstream.
|
||||
keyword_extraction = kwargs.pop("keyword_extraction", keyword_extraction)
|
||||
entity_extraction = kwargs.pop("entity_extraction", entity_extraction)
|
||||
|
||||
# Deprecation shims: map legacy boolean flags to response_format only when
|
||||
# an explicit response_format was not supplied by the caller. The legacy
|
||||
# path also forces enable_cot=False so that reasoning_content cannot
|
||||
# corrupt the JSON payload expected by callers that were relying on it.
|
||||
if kwargs.get("response_format") is None:
|
||||
if entity_extraction:
|
||||
warnings.warn(
|
||||
"zhipu_complete(entity_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
enable_cot = False
|
||||
elif keyword_extraction:
|
||||
warnings.warn(
|
||||
"zhipu_complete(keyword_extraction=True) is deprecated; "
|
||||
"pass response_format={'type': 'json_object'} instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
enable_cot = False
|
||||
|
||||
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, max_token_size=8192, model_name="embedding-3"
|
||||
)
|
||||
@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,
|
||||
embedding_dim: int | None = 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:
|
||||
request_kwargs = dict(kwargs)
|
||||
if embedding_dim is not None:
|
||||
request_kwargs["dimensions"] = embedding_dim
|
||||
response = client.embeddings.create(
|
||||
model=model, input=[text], **request_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)
|
||||
@@ -0,0 +1,579 @@
|
||||
"""LLM role registry, configuration types, and runtime mixin.
|
||||
|
||||
LightRAG can route different stages of work (entity extraction, keyword
|
||||
extraction, query, vlm) to distinct LLM bindings. This module owns the
|
||||
static role registry (:data:`ROLES`), the per-role configuration
|
||||
(:class:`RoleLLMConfig`), and the :class:`_RoleLLMMixin` that drives the
|
||||
runtime: builder registration, wrapper rebuilding, hot config updates,
|
||||
queue cleanup, and queue-status reporting.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
from typing import Any, Callable, Mapping
|
||||
|
||||
from lightrag.utils import (
|
||||
get_env_value,
|
||||
logger,
|
||||
priority_limit_async_func_call,
|
||||
)
|
||||
|
||||
|
||||
def _optional_env_int(env_key: str) -> int | None:
|
||||
return get_env_value(env_key, None, int, special_none=True)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RoleSpec:
|
||||
"""Static descriptor for a known LLM role.
|
||||
|
||||
Adding a new role anywhere in LightRAG is a single-line edit: append a
|
||||
``RoleSpec`` to :data:`ROLES`. Every other component (env var loop in
|
||||
``api/config.py``, queue observability, role config update flow) iterates
|
||||
this registry rather than hard-coding role names.
|
||||
"""
|
||||
|
||||
name: str
|
||||
"""Canonical lowercase role key (used in ``role_llm_configs`` dict and CLI/log output)."""
|
||||
|
||||
env_prefix: str
|
||||
"""Uppercase prefix used by the API env-var layer, e.g. ``"EXTRACT"`` for
|
||||
``EXTRACT_LLM_BINDING`` / ``EXTRACT_MAX_ASYNC_LLM`` / ``EXTRACT_LLM_TIMEOUT``."""
|
||||
|
||||
queue_name: str
|
||||
"""Display name passed to ``priority_limit_async_func_call`` for log lines."""
|
||||
|
||||
|
||||
ROLES: tuple[RoleSpec, ...] = (
|
||||
RoleSpec("extract", "EXTRACT", "extract LLM func"),
|
||||
RoleSpec("keyword", "KEYWORD", "keyword LLM func"),
|
||||
RoleSpec("query", "QUERY", "query LLM func"),
|
||||
RoleSpec("vlm", "VLM", "vlm LLM func"),
|
||||
)
|
||||
ROLE_NAMES: frozenset[str] = frozenset(spec.name for spec in ROLES)
|
||||
ROLES_BY_NAME: dict[str, RoleSpec] = {spec.name: spec for spec in ROLES}
|
||||
|
||||
|
||||
@dataclass
|
||||
class RoleLLMConfig:
|
||||
"""Per-role LLM override accepted at :class:`LightRAG` init time.
|
||||
|
||||
Any field left as ``None`` falls back to the corresponding base LLM
|
||||
setting (``llm_model_func`` / ``llm_model_kwargs`` / ``llm_model_max_async``
|
||||
/ ``default_llm_timeout``). When ``max_async`` is None at init and the
|
||||
user did not pass a ``role_llm_configs`` entry for the role, the value is
|
||||
additionally seeded from ``{ROLE_PREFIX}_MAX_ASYNC_LLM``. ``metadata`` seeds
|
||||
runtime observability and role-builder context.
|
||||
"""
|
||||
|
||||
func: Callable[..., object] | None = None
|
||||
kwargs: dict[str, Any] | None = None
|
||||
max_async: int | None = None
|
||||
timeout: int | None = None
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RoleLLMState:
|
||||
"""Runtime state for one role. Internal — not part of the public API."""
|
||||
|
||||
raw_func: Callable[..., object]
|
||||
kwargs: dict[str, Any] | None
|
||||
max_async: int | None
|
||||
timeout: int | None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
wrapped: Callable[..., object] | None = None
|
||||
|
||||
|
||||
class _RoleLLMMixin:
|
||||
"""Mixin that owns the role LLM runtime on :class:`LightRAG`.
|
||||
|
||||
Mixed into LightRAG only. Relies on attributes that the main class
|
||||
initializes in ``__post_init__`` (``_role_llm_states``, ``_role_llm_builders``,
|
||||
``llm_model_func``, ``llm_model_kwargs``, ``llm_model_max_async``,
|
||||
``default_llm_timeout``, ``embedding_func``, ``rerank_model_func``).
|
||||
"""
|
||||
|
||||
_SECRET_MARKERS = (
|
||||
"api_key",
|
||||
"api-key",
|
||||
"apikey",
|
||||
"access_key",
|
||||
"access-key",
|
||||
"secret",
|
||||
"token",
|
||||
"credential",
|
||||
"password",
|
||||
"passphrase",
|
||||
"pwd",
|
||||
"auth",
|
||||
"session",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_llm_role(role: str) -> str:
|
||||
normalized = role.strip().lower()
|
||||
if normalized not in ROLE_NAMES:
|
||||
raise ValueError(f"Invalid LLM role: {role}")
|
||||
return normalized
|
||||
|
||||
def register_role_llm_builder(
|
||||
self,
|
||||
builder: Callable[
|
||||
[str, dict[str, Any]], tuple[Callable[..., object], dict[str, Any] | None]
|
||||
],
|
||||
) -> None:
|
||||
"""Register a runtime builder used by update_llm_role_config for binding/model updates."""
|
||||
self._llm_role_builder = builder
|
||||
|
||||
def set_role_llm_metadata(self, role: str, **metadata: Any) -> None:
|
||||
"""Store role metadata used when rebuilding a role-specific LLM function."""
|
||||
role = self._normalize_llm_role(role)
|
||||
state = self._role_llm_states[role]
|
||||
for key, value in metadata.items():
|
||||
if value is None:
|
||||
continue
|
||||
state.metadata[key] = value
|
||||
|
||||
@property
|
||||
def role_llm_funcs(self) -> Mapping[str, Callable[..., object]]:
|
||||
"""Read-only mapping of role name → wrapped (queue-managed) LLM func."""
|
||||
return {
|
||||
name: state.wrapped
|
||||
for name, state in self._role_llm_states.items()
|
||||
if state.wrapped is not None
|
||||
}
|
||||
|
||||
@property
|
||||
def role_llm_kwargs(self) -> Mapping[str, dict[str, Any] | None]:
|
||||
"""Read-only mapping of role name → effective LLM kwargs (None means inherit base)."""
|
||||
return {name: state.kwargs for name, state in self._role_llm_states.items()}
|
||||
|
||||
def _get_effective_role_llm_kwargs(self, role: str) -> dict[str, Any]:
|
||||
state = self._role_llm_states[self._normalize_llm_role(role)]
|
||||
if state.kwargs is not None:
|
||||
return state.kwargs
|
||||
if state.metadata.get("is_cross_provider"):
|
||||
return {}
|
||||
return self.llm_model_kwargs
|
||||
|
||||
def _get_effective_role_llm_timeout(self, role: str) -> int:
|
||||
state = self._role_llm_states[self._normalize_llm_role(role)]
|
||||
return state.timeout if state.timeout is not None else self.default_llm_timeout
|
||||
|
||||
def _get_effective_role_llm_max_async(self, role: str) -> int:
|
||||
state = self._role_llm_states[self._normalize_llm_role(role)]
|
||||
return (
|
||||
state.max_async if state.max_async is not None else self.llm_model_max_async
|
||||
)
|
||||
|
||||
def _wrap_llm_role_func(
|
||||
self,
|
||||
role_name: str,
|
||||
raw_func: Callable[..., object],
|
||||
max_async: int,
|
||||
timeout: int,
|
||||
model_kwargs: dict[str, Any],
|
||||
) -> Callable[..., object]:
|
||||
spec = ROLES_BY_NAME[role_name]
|
||||
return priority_limit_async_func_call(
|
||||
max_async,
|
||||
llm_timeout=timeout,
|
||||
queue_name=spec.queue_name,
|
||||
concurrency_group=f"llm:{role_name}",
|
||||
)(
|
||||
partial(
|
||||
raw_func,
|
||||
hashing_kv=self.llm_response_cache,
|
||||
**model_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
def _rebuild_role_llm_funcs(self) -> None:
|
||||
"""Wrap each role's raw_func with its own priority queue.
|
||||
|
||||
Base ``llm_model_func`` is intentionally NOT wrapped — concurrency
|
||||
for the base function is enforced at the role layer (every code path
|
||||
that calls an LLM goes through a role wrapper).
|
||||
"""
|
||||
for spec in ROLES:
|
||||
self._rebuild_single_role_llm_func(spec.name)
|
||||
|
||||
def _rebuild_single_role_llm_func(self, role: str) -> None:
|
||||
role = self._normalize_llm_role(role)
|
||||
state = self._role_llm_states[role]
|
||||
state.wrapped = self._wrap_llm_role_func(
|
||||
role,
|
||||
state.raw_func,
|
||||
self._get_effective_role_llm_max_async(role),
|
||||
self._get_effective_role_llm_timeout(role),
|
||||
self._get_effective_role_llm_kwargs(role),
|
||||
)
|
||||
|
||||
async def _shutdown_llm_wrapper(self, wrapped_func: Callable[..., object]) -> None:
|
||||
shutdown = getattr(wrapped_func, "shutdown", None)
|
||||
if callable(shutdown):
|
||||
await shutdown(graceful=True)
|
||||
|
||||
def _schedule_retired_llm_queue_cleanup(
|
||||
self, wrapped_func: Callable[..., object] | None
|
||||
) -> None:
|
||||
if wrapped_func is None or not callable(
|
||||
getattr(wrapped_func, "shutdown", None)
|
||||
):
|
||||
return
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
# The retired wrapper's queue and worker tasks are tied to the
|
||||
# event loop that first used them. Spinning up a fresh loop via
|
||||
# asyncio.run would either hang on queue.join() or touch
|
||||
# primitives bound to a closed loop. Skip cleanup with a warning
|
||||
# — call aupdate_llm_role_config() from an async context for
|
||||
# deterministic shutdown.
|
||||
logger.warning(
|
||||
"update_llm_role_config: skipping retired LLM queue cleanup "
|
||||
"because no event loop is running; call aupdate_llm_role_config() "
|
||||
"from an async context for deterministic shutdown"
|
||||
)
|
||||
return
|
||||
|
||||
task = loop.create_task(self._shutdown_llm_wrapper(wrapped_func))
|
||||
self._retired_llm_queue_cleanup_tasks.add(task)
|
||||
task.add_done_callback(self._finalize_retired_llm_queue_cleanup)
|
||||
|
||||
def _finalize_retired_llm_queue_cleanup(self, task: asyncio.Task) -> None:
|
||||
self._retired_llm_queue_cleanup_tasks.discard(task)
|
||||
try:
|
||||
task.result()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning(f"Retired LLM queue cleanup failed: {e}")
|
||||
|
||||
async def wait_for_retired_llm_queues(self) -> None:
|
||||
"""Wait until all retired role LLM queues have drained and shut down.
|
||||
|
||||
Cleanup failures are logged by ``_finalize_retired_llm_queue_cleanup``
|
||||
and intentionally swallowed here so callers can rely on this method
|
||||
always returning once every retired wrapper has finished.
|
||||
"""
|
||||
while self._retired_llm_queue_cleanup_tasks:
|
||||
tasks = list(self._retired_llm_queue_cleanup_tasks)
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
def _apply_llm_role_config_update(
|
||||
self,
|
||||
role: str,
|
||||
*,
|
||||
model_func: Callable[..., object] | None = None,
|
||||
model_kwargs: dict[str, Any] | None = None,
|
||||
max_async: int | None = None,
|
||||
timeout: int | None = None,
|
||||
binding: str | None = None,
|
||||
model: str | None = None,
|
||||
host: str | None = None,
|
||||
api_key: str | None = None,
|
||||
provider_options: dict[str, Any] | None = None,
|
||||
) -> Callable[..., object] | None:
|
||||
role = self._normalize_llm_role(role)
|
||||
state = self._role_llm_states[role]
|
||||
old_wrapped = state.wrapped
|
||||
|
||||
snapshot = _RoleLLMState(
|
||||
raw_func=state.raw_func,
|
||||
kwargs=deepcopy(state.kwargs),
|
||||
max_async=state.max_async,
|
||||
timeout=state.timeout,
|
||||
metadata=deepcopy(state.metadata),
|
||||
wrapped=state.wrapped,
|
||||
)
|
||||
|
||||
try:
|
||||
if model_func is not None and not callable(model_func):
|
||||
raise TypeError("model_func must be callable")
|
||||
|
||||
if model_kwargs is not None:
|
||||
state.kwargs = model_kwargs
|
||||
if max_async is not None:
|
||||
state.max_async = max_async
|
||||
if timeout is not None:
|
||||
state.timeout = timeout
|
||||
if model_func is not None:
|
||||
state.raw_func = model_func
|
||||
|
||||
metadata_updated = any(
|
||||
value is not None
|
||||
for value in (binding, model, host, api_key, provider_options)
|
||||
)
|
||||
if binding is not None:
|
||||
state.metadata["binding"] = binding
|
||||
if model is not None:
|
||||
state.metadata["model"] = model
|
||||
if host is not None:
|
||||
state.metadata["host"] = host
|
||||
if api_key is not None:
|
||||
state.metadata["api_key"] = api_key
|
||||
if provider_options is not None:
|
||||
state.metadata["provider_options"] = provider_options
|
||||
if "base_binding" in state.metadata and "binding" in state.metadata:
|
||||
state.metadata["is_cross_provider"] = (
|
||||
state.metadata["binding"] != state.metadata["base_binding"]
|
||||
)
|
||||
|
||||
if metadata_updated:
|
||||
builder = getattr(self, "_llm_role_builder", None)
|
||||
if builder is None and model_func is None:
|
||||
raise ValueError(
|
||||
"Runtime role builder is not configured; provide model_func or register_role_llm_builder() first"
|
||||
)
|
||||
if builder is not None:
|
||||
built_func, built_kwargs = builder(role, state.metadata)
|
||||
state.raw_func = built_func
|
||||
if model_kwargs is None and built_kwargs is not None:
|
||||
state.kwargs = built_kwargs
|
||||
|
||||
self._rebuild_single_role_llm_func(role)
|
||||
except Exception:
|
||||
state.raw_func = snapshot.raw_func
|
||||
state.kwargs = snapshot.kwargs
|
||||
state.max_async = snapshot.max_async
|
||||
state.timeout = snapshot.timeout
|
||||
state.metadata = snapshot.metadata
|
||||
state.wrapped = snapshot.wrapped
|
||||
raise
|
||||
|
||||
self._log_llm_role_config("updated", role=role)
|
||||
return old_wrapped
|
||||
|
||||
def update_llm_role_config(
|
||||
self,
|
||||
role: str,
|
||||
*,
|
||||
model_func: Callable[..., object] | None = None,
|
||||
model_kwargs: dict[str, Any] | None = None,
|
||||
max_async: int | None = None,
|
||||
timeout: int | None = None,
|
||||
binding: str | None = None,
|
||||
model: str | None = None,
|
||||
host: str | None = None,
|
||||
api_key: str | None = None,
|
||||
provider_options: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Update a role-specific LLM configuration at runtime.
|
||||
|
||||
Supports lightweight updates (kwargs/max_async/timeout/model_func) directly.
|
||||
For binding/model/host/api_key/provider_options updates, a role builder must
|
||||
be registered via register_role_llm_builder().
|
||||
"""
|
||||
old_wrapped = self._apply_llm_role_config_update(
|
||||
role,
|
||||
model_func=model_func,
|
||||
model_kwargs=model_kwargs,
|
||||
max_async=max_async,
|
||||
timeout=timeout,
|
||||
binding=binding,
|
||||
model=model,
|
||||
host=host,
|
||||
api_key=api_key,
|
||||
provider_options=provider_options,
|
||||
)
|
||||
self._schedule_retired_llm_queue_cleanup(old_wrapped)
|
||||
|
||||
async def aupdate_llm_role_config(
|
||||
self,
|
||||
role: str,
|
||||
*,
|
||||
model_func: Callable[..., object] | None = None,
|
||||
model_kwargs: dict[str, Any] | None = None,
|
||||
max_async: int | None = None,
|
||||
timeout: int | None = None,
|
||||
binding: str | None = None,
|
||||
model: str | None = None,
|
||||
host: str | None = None,
|
||||
api_key: str | None = None,
|
||||
provider_options: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Async variant of update_llm_role_config that waits for queue cleanup.
|
||||
|
||||
Blocking behavior:
|
||||
This coroutine awaits a graceful shutdown of the retired role
|
||||
wrapper's priority queue. The shutdown blocks on
|
||||
``queue.join()`` until every already-queued LLM call has been
|
||||
executed (workers always call ``task_done()`` in ``finally``,
|
||||
so in-flight requests are not cut off).
|
||||
|
||||
The wait is bounded by ``max_task_duration`` of the retired
|
||||
queue, which is computed as ``llm_timeout * 2 + 15`` seconds
|
||||
(default ``180 * 2 + 15 = 375`` seconds, ~6 min 15 s). When
|
||||
this bound is reached, the drain times out and the shutdown
|
||||
falls through to forced cancellation: pending futures are
|
||||
cancelled, the queue is cleared, workers are stopped. So this
|
||||
method **never blocks indefinitely**, but with a deep backlog
|
||||
of slow LLM calls it can take up to that bound to return, and
|
||||
in-flight calls past the bound will be cancelled.
|
||||
|
||||
If you need a non-blocking switch, use the sync
|
||||
``update_llm_role_config()`` (which schedules cleanup as a
|
||||
background task) and await ``wait_for_retired_llm_queues()``
|
||||
separately when you want to confirm the old queue is gone.
|
||||
"""
|
||||
old_wrapped = self._apply_llm_role_config_update(
|
||||
role,
|
||||
model_func=model_func,
|
||||
model_kwargs=model_kwargs,
|
||||
max_async=max_async,
|
||||
timeout=timeout,
|
||||
binding=binding,
|
||||
model=model,
|
||||
host=host,
|
||||
api_key=api_key,
|
||||
provider_options=provider_options,
|
||||
)
|
||||
if old_wrapped is not None:
|
||||
await self._shutdown_llm_wrapper(old_wrapped)
|
||||
|
||||
@classmethod
|
||||
def _is_secret_key(cls, key: str) -> bool:
|
||||
lowered = key.lower()
|
||||
return any(marker in lowered for marker in cls._SECRET_MARKERS)
|
||||
|
||||
def _scrubbed_llm_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return a deep copy of ``metadata`` with auth-bearing fields removed.
|
||||
|
||||
Auth-bearing fields are stripped entirely — not masked — because a
|
||||
masked ``"***"`` carries no information for an external consumer
|
||||
(operators already see ``binding`` / ``host`` to confirm a role is
|
||||
configured). Stripping makes the invariant simple: anything that
|
||||
appears in this output is safe to log, cache, ship over the wire.
|
||||
|
||||
Components that legitimately need the raw secret (the role builder,
|
||||
provider clients) read it directly off the private
|
||||
``_role_llm_states[role].metadata`` dict.
|
||||
"""
|
||||
|
||||
def scrub_value(value: Any) -> Any:
|
||||
if isinstance(value, Mapping):
|
||||
return {
|
||||
key: scrub_value(inner_value)
|
||||
for key, inner_value in value.items()
|
||||
if not self._is_secret_key(str(key))
|
||||
}
|
||||
if isinstance(value, list):
|
||||
return [scrub_value(item) for item in value]
|
||||
if isinstance(value, tuple):
|
||||
return tuple(scrub_value(item) for item in value)
|
||||
return deepcopy(value)
|
||||
|
||||
return scrub_value(metadata)
|
||||
|
||||
def get_llm_role_config(self, role: str | None = None) -> dict[str, Any]:
|
||||
"""Return effective role LLM runtime configuration (observability snapshot).
|
||||
|
||||
Each role entry exposes ``binding`` / ``model`` / ``host`` at the top
|
||||
level for convenience and again inside ``metadata`` as part of the
|
||||
full runtime snapshot (which may contain extra builder-specific
|
||||
keys). Auth-bearing fields (``api_key``, ``aws_secret_access_key``,
|
||||
``password``, …) are **stripped entirely** from ``metadata`` — this
|
||||
method is intended for ``/health`` / WebUI / audit output and must
|
||||
never leak credentials. There is no escape hatch; runtime components
|
||||
that legitimately need the raw value read it from
|
||||
``_role_llm_states[role].metadata`` directly.
|
||||
"""
|
||||
|
||||
def role_config(role_name: str) -> dict[str, Any]:
|
||||
state = self._role_llm_states[role_name]
|
||||
metadata = self._scrubbed_llm_metadata(state.metadata)
|
||||
return {
|
||||
"binding": metadata.get("binding"),
|
||||
"model": metadata.get("model"),
|
||||
"host": metadata.get("host"),
|
||||
"is_cross_provider": metadata.get("is_cross_provider", False),
|
||||
"max_async": self._get_effective_role_llm_max_async(role_name),
|
||||
"timeout": self._get_effective_role_llm_timeout(role_name),
|
||||
"has_model_kwargs": state.kwargs is not None,
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
if role is not None:
|
||||
return role_config(self._normalize_llm_role(role))
|
||||
|
||||
return {spec.name: role_config(spec.name) for spec in ROLES}
|
||||
|
||||
def _log_llm_role_config(self, reason: str, role: str | None = None) -> None:
|
||||
"""Log the sanitized role LLM runtime configuration."""
|
||||
if role is None:
|
||||
configs = self.get_llm_role_config()
|
||||
role_names = [spec.name for spec in ROLES]
|
||||
logger.info(f"Role LLM Configuration ({reason}):")
|
||||
else:
|
||||
normalized_role = self._normalize_llm_role(role)
|
||||
configs = {normalized_role: self.get_llm_role_config(normalized_role)}
|
||||
role_names = [normalized_role]
|
||||
logger.info(f"Role LLM Configuration ({reason}: {normalized_role}):")
|
||||
|
||||
for role_name in role_names:
|
||||
cfg = configs[role_name]
|
||||
logger.info(
|
||||
" - %s: %s/%s, host=%s, max_async=%s, timeout=%s",
|
||||
role_name,
|
||||
cfg["binding"],
|
||||
cfg["model"],
|
||||
cfg["host"],
|
||||
cfg["max_async"],
|
||||
cfg["timeout"],
|
||||
)
|
||||
|
||||
async def _queue_status_for_func(
|
||||
self, func: Callable[..., object] | None
|
||||
) -> dict[str, Any]:
|
||||
if func is None:
|
||||
return {"available": False}
|
||||
# Prefer the cross-worker aggregated view (sums every gunicorn
|
||||
# worker's published snapshot; falls back to the local snapshot
|
||||
# internally on any shared-storage failure, so "available" keeps
|
||||
# meaning "this wrapper exists", never "aggregation succeeded").
|
||||
get_stats = getattr(func, "get_aggregated_queue_stats", None)
|
||||
if not callable(get_stats):
|
||||
get_stats = getattr(func, "get_queue_stats", None)
|
||||
if not callable(get_stats):
|
||||
return {"available": False}
|
||||
stats = get_stats()
|
||||
if inspect.isawaitable(stats):
|
||||
stats = await stats
|
||||
stats["available"] = True
|
||||
return stats
|
||||
|
||||
async def get_llm_queue_status(self, include_base: bool = True) -> dict[str, Any]:
|
||||
"""Return queue status for each role's wrapped LLM func.
|
||||
|
||||
The base ``llm_model_func`` is no longer queue-wrapped, so it is not
|
||||
reported here. ``include_base`` is kept for signature compatibility
|
||||
but has no effect.
|
||||
"""
|
||||
del include_base # base is unwrapped — see docstring
|
||||
|
||||
result: dict[str, Any] = {}
|
||||
for spec in ROLES:
|
||||
state = self._role_llm_states.get(spec.name)
|
||||
result[spec.name] = await self._queue_status_for_func(
|
||||
state.wrapped if state else None
|
||||
)
|
||||
return result
|
||||
|
||||
async def get_embedding_queue_status(self) -> dict[str, Any]:
|
||||
"""Return queue status for the wrapped embedding function."""
|
||||
return await self._queue_status_for_func(
|
||||
self.embedding_func.func if self.embedding_func is not None else None
|
||||
)
|
||||
|
||||
async def get_rerank_queue_status(self) -> dict[str, Any]:
|
||||
"""Return queue status for the wrapped rerank function."""
|
||||
return await self._queue_status_for_func(self.rerank_model_func)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,28 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable
|
||||
|
||||
|
||||
# All namespace should not be changed
|
||||
class NameSpace:
|
||||
KV_STORE_FULL_DOCS = "full_docs"
|
||||
KV_STORE_TEXT_CHUNKS = "text_chunks"
|
||||
KV_STORE_LLM_RESPONSE_CACHE = "llm_response_cache"
|
||||
KV_STORE_FULL_ENTITIES = "full_entities"
|
||||
KV_STORE_FULL_RELATIONS = "full_relations"
|
||||
KV_STORE_ENTITY_CHUNKS = "entity_chunks"
|
||||
KV_STORE_RELATION_CHUNKS = "relation_chunks"
|
||||
|
||||
VECTOR_STORE_ENTITIES = "entities"
|
||||
VECTOR_STORE_RELATIONSHIPS = "relationships"
|
||||
VECTOR_STORE_CHUNKS = "chunks"
|
||||
|
||||
GRAPH_STORE_CHUNK_ENTITY_RELATION = "chunk_entity_relation"
|
||||
|
||||
DOC_STATUS = "doc_status"
|
||||
|
||||
|
||||
def is_namespace(namespace: str, base_namespace: str | Iterable[str]):
|
||||
if isinstance(base_namespace, str):
|
||||
return namespace.endswith(base_namespace)
|
||||
return any(is_namespace(namespace, ns) for ns in base_namespace)
|
||||
+6003
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,13 @@
|
||||
"""Document parsing layer.
|
||||
|
||||
Sub-packages and modules:
|
||||
|
||||
- ``routing``: parser engine selection and per-file parser directives.
|
||||
- ``debug``: minimal ``LightRAG`` stand-in for offline parser debugging.
|
||||
- ``cli``: ``python -m lightrag.parser.cli`` entry point for single-file
|
||||
parser debugging across all engines.
|
||||
- ``docx``: native ``.docx`` parser. Additional native format parsers
|
||||
should live as sibling sub-packages here (e.g. ``parser/pdf/``).
|
||||
- ``external``: adapters for external parsing services (``mineru``,
|
||||
``docling``) that post to a remote API and cache the raw bundle.
|
||||
"""
|
||||
@@ -0,0 +1,217 @@
|
||||
"""Shared HTML-table parsing helpers for parser engines.
|
||||
|
||||
Pure functions that recover structural facts from an HTML ``<table>`` string
|
||||
without a heavy dependency: row/column counts (colspan-aware), the verbatim
|
||||
``<thead>`` substring, table-payload detection, and ``<html>/<body>`` wrapper
|
||||
stripping. Originally private to the mineru IR builder; lifted into a leaf
|
||||
module so the native markdown IR builder can reuse the exact same logic
|
||||
(merged-cell semantics must survive identically across engines).
|
||||
|
||||
Leaf module with no parser-layer imports — safe for any engine to import.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from html.parser import HTMLParser
|
||||
|
||||
from lightrag.utils import logger
|
||||
|
||||
|
||||
@dataclass
|
||||
class HTMLTableInfo:
|
||||
num_rows: int = 0
|
||||
num_cols: int = 0
|
||||
|
||||
|
||||
class _HTMLTableInfoParser(HTMLParser):
|
||||
"""Count ``<tr>`` rows and their (colspan-aware) column widths.
|
||||
|
||||
Used only to recover ``num_rows`` / ``num_cols`` when the engine did not
|
||||
supply them; the ``<thead>`` header itself is preserved verbatim by
|
||||
:func:`extract_thead_html`, not reconstructed here.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(convert_charrefs=True)
|
||||
# ``col_count`` (sum of colspans) for each completed top-level ``<tr>``.
|
||||
self.row_col_counts: list[int] = []
|
||||
self._tr_depth = 0
|
||||
self._cell_depth = 0
|
||||
self._row_col_count = 0
|
||||
self._cell_colspan = 1
|
||||
|
||||
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
|
||||
tag = tag.lower()
|
||||
if tag == "tr":
|
||||
if self._tr_depth == 0:
|
||||
self._row_col_count = 0
|
||||
self._tr_depth += 1
|
||||
return
|
||||
if tag in {"td", "th"} and self._tr_depth > 0:
|
||||
if self._cell_depth == 0:
|
||||
self._cell_colspan = _cell_span(attrs, "colspan")
|
||||
self._cell_depth += 1
|
||||
|
||||
def handle_endtag(self, tag: str) -> None:
|
||||
tag = tag.lower()
|
||||
if tag in {"td", "th"} and self._cell_depth > 0:
|
||||
self._cell_depth -= 1
|
||||
if self._cell_depth == 0:
|
||||
self._row_col_count += self._cell_colspan
|
||||
self._cell_colspan = 1
|
||||
return
|
||||
if tag == "tr" and self._tr_depth > 0:
|
||||
self._tr_depth -= 1
|
||||
# ``col_count > 0`` ⇔ the row held at least one cell (colspan ≥ 1),
|
||||
# so an empty ``<tr></tr>`` is skipped exactly as before.
|
||||
if self._tr_depth == 0 and self._row_col_count > 0:
|
||||
self.row_col_counts.append(self._row_col_count)
|
||||
self._row_col_count = 0
|
||||
|
||||
|
||||
def _cell_span(attrs: list[tuple[str, str | None]], name: str) -> int:
|
||||
"""Read a ``colspan``/``rowspan`` attribute as an int ``>= 1`` (default 1)."""
|
||||
for key, value in attrs:
|
||||
if key.lower() != name:
|
||||
continue
|
||||
try:
|
||||
return max(int(value or "1"), 1)
|
||||
except ValueError:
|
||||
return 1
|
||||
return 1
|
||||
|
||||
|
||||
def extract_html_table_info(html: str) -> HTMLTableInfo:
|
||||
parser = _HTMLTableInfoParser()
|
||||
try:
|
||||
parser.feed(html or "")
|
||||
parser.close()
|
||||
except Exception as exc: # pragma: no cover - HTMLParser is forgiving.
|
||||
logger.debug("[html_table] failed to parse table HTML: %s", exc)
|
||||
return HTMLTableInfo()
|
||||
return HTMLTableInfo(
|
||||
num_rows=len(parser.row_col_counts),
|
||||
num_cols=max(parser.row_col_counts, default=0),
|
||||
)
|
||||
|
||||
|
||||
def extract_thead_html(html: str) -> str | None:
|
||||
"""Return the first top-level ``<thead …>…</thead>`` substring verbatim.
|
||||
|
||||
The raw markup is kept so merged-cell semantics (``rowspan`` / ``colspan``)
|
||||
survive into ``tables.json`` and, later, into every repeated header chunk
|
||||
of a split table. Returns ``None`` when the table has no ``<thead>`` or the
|
||||
``<thead>`` carries no visible text (a blank spacer row, which would
|
||||
otherwise emit empty ``<th>`` headers).
|
||||
"""
|
||||
stripped = (html or "").strip()
|
||||
lower = stripped.lower()
|
||||
start = find_html_tag(lower, "thead")
|
||||
if start < 0:
|
||||
return None
|
||||
close = lower.find("</thead>", start)
|
||||
if close < 0:
|
||||
return None
|
||||
thead = stripped[start : close + len("</thead>")]
|
||||
# Blank check: drop a header whose cells hold no non-whitespace text.
|
||||
if not re.sub(r"<[^>]+>", "", thead).strip():
|
||||
return None
|
||||
return thead
|
||||
|
||||
|
||||
def looks_like_html_table_payload(body: str) -> bool:
|
||||
lower = (body or "").lstrip().lower()
|
||||
return any(
|
||||
starts_with_html_tag(lower, tag)
|
||||
for tag in ("table", "thead", "tbody", "tfoot", "tr", "html", "body")
|
||||
)
|
||||
|
||||
|
||||
def unwrap_html_table(payload: str) -> str:
|
||||
"""Strip a ``<html>/<body>`` document wrapper that a table model
|
||||
sometimes emits, returning the outermost ``<table…>…</table>`` span. Keeps
|
||||
a single clean ``<table>`` so the writer does not nest tables and the
|
||||
non-greedy ``TABLE_TAG_RE`` is not truncated at an inner ``</table>``.
|
||||
Falls back to the stripped payload when no ``<table>`` element exists."""
|
||||
stripped = (payload or "").strip()
|
||||
lower = stripped.lower()
|
||||
start = _find_table_open(lower)
|
||||
if start < 0:
|
||||
return stripped
|
||||
close = lower.rfind("</table>")
|
||||
if close < start:
|
||||
return stripped
|
||||
return stripped[start : close + len("</table>")]
|
||||
|
||||
|
||||
def _find_table_open(lower: str) -> int:
|
||||
"""First index of a real ``<table`` start tag (not e.g. ``<tablefoo``).
|
||||
Returns -1 when none is present."""
|
||||
return find_html_tag(lower, "table")
|
||||
|
||||
|
||||
def find_html_tag(lower: str, tag: str) -> int:
|
||||
"""First index of a real ``<tag`` start tag (not e.g. ``<tablefoo`` for
|
||||
``tag="table"``). ``lower`` must already be lower-cased. Returns -1 when
|
||||
none is present."""
|
||||
needle = f"<{tag}"
|
||||
idx = 0
|
||||
while True:
|
||||
idx = lower.find(needle, idx)
|
||||
if idx < 0:
|
||||
return -1
|
||||
nxt = idx + len(needle)
|
||||
if nxt >= len(lower) or lower[nxt] in {" ", "\t", "\r", "\n", ">", "/"}:
|
||||
return idx
|
||||
idx = nxt
|
||||
|
||||
|
||||
def starts_with_html_tag(lower: str, tag: str) -> bool:
|
||||
prefix = f"<{tag}"
|
||||
if not lower.startswith(prefix):
|
||||
return False
|
||||
if len(lower) == len(prefix):
|
||||
return True
|
||||
return lower[len(prefix)] in {" ", "\t", "\r", "\n", ">", "/"}
|
||||
|
||||
|
||||
def html_table_inner_body(html: str) -> str:
|
||||
stripped = (html or "").strip()
|
||||
lower = stripped.lower()
|
||||
if not starts_with_html_tag(lower, "table"):
|
||||
return stripped
|
||||
open_end = _open_tag_end(stripped)
|
||||
close_start = lower.rfind("</table>")
|
||||
if open_end < 0 or close_start <= open_end:
|
||||
return stripped
|
||||
return stripped[open_end + 1 : close_start].strip()
|
||||
|
||||
|
||||
def _open_tag_end(html: str) -> int:
|
||||
"""Index of the ``>`` closing the leading tag, skipping quoted attribute
|
||||
values so a ``>`` inside an attribute (e.g. ``<table data-x="a>b">``) does
|
||||
not terminate the tag early. Returns -1 when no closing ``>`` is found."""
|
||||
quote: str | None = None
|
||||
for idx, ch in enumerate(html):
|
||||
if quote is not None:
|
||||
if ch == quote:
|
||||
quote = None
|
||||
elif ch in {'"', "'"}:
|
||||
quote = ch
|
||||
elif ch == ">":
|
||||
return idx
|
||||
return -1
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HTMLTableInfo",
|
||||
"extract_html_table_info",
|
||||
"extract_thead_html",
|
||||
"looks_like_html_table_payload",
|
||||
"unwrap_html_table",
|
||||
"find_html_tag",
|
||||
"starts_with_html_tag",
|
||||
"html_table_inner_body",
|
||||
]
|
||||
@@ -0,0 +1,72 @@
|
||||
"""Shared markdown rendering helpers for parser engines.
|
||||
|
||||
Used by the native docx parser and the external (mineru / docling) IR
|
||||
builders so heading-line rendering stays identical across engines. This is
|
||||
a leaf module with no heavy imports — ``lightrag/parser/__init__.py`` only
|
||||
carries a docstring — so all three engines can import it without risking a
|
||||
circular dependency.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
# Markdown caps heading levels at 6 (``######``); deeper outline levels are
|
||||
# clamped to 6 rather than emitting an illegal 7+ run of ``#``.
|
||||
MAX_HEADING_LEVEL = 6
|
||||
|
||||
# A heading line that is ALREADY markdown: 1-6 ``#`` followed by one or more
|
||||
# spaces. Used to avoid double-prefixing text that an upstream engine emitted
|
||||
# with its own markdown heading marker (e.g. mineru/docling extracting
|
||||
# ``# Foo``).
|
||||
#
|
||||
# Ambiguity note: this is a heuristic, not a parse. A heading whose text
|
||||
# genuinely starts with ``#`` plus a space (an author literally writing
|
||||
# ``"# Note"`` as heading content) is indistinguishable from an
|
||||
# already-rendered markdown heading and will be treated as the latter. This
|
||||
# is accepted as a rare edge case in exchange for engine-agnostic dedup.
|
||||
_MD_HEADING_RE = re.compile(r"^#{1,6} +")
|
||||
|
||||
|
||||
def strip_heading_markdown_prefix(text: str) -> str:
|
||||
"""Return heading metadata without an existing markdown heading prefix.
|
||||
|
||||
The content renderer may keep a source line such as ``"# Foo"`` verbatim
|
||||
to avoid double-prefixing, but structured metadata (``heading``,
|
||||
``parent_headings``, doc title) must stay clean. The whole ``#`` run and
|
||||
all following spaces are removed, so ``"# Extra"`` yields ``"Extra"``
|
||||
(no leading space leaks into the metadata).
|
||||
|
||||
See the module-level ambiguity note: text that genuinely begins with a
|
||||
``#`` + space run is stripped here too.
|
||||
"""
|
||||
return _MD_HEADING_RE.sub("", text, count=1)
|
||||
|
||||
|
||||
def render_heading_line(level: int, text: str) -> str:
|
||||
"""Render a heading as a markdown-prefixed content line.
|
||||
|
||||
Args:
|
||||
level: 1-based heading level (1 = H1). Values < 1 are treated as 1;
|
||||
values > :data:`MAX_HEADING_LEVEL` are clamped so a level >= 7
|
||||
heading still gets ``######``.
|
||||
text: The heading text.
|
||||
|
||||
Returns:
|
||||
``text`` unchanged when it already starts with a markdown heading
|
||||
prefix (``^#{1,6} +`` — 1-6 ``#`` then one or more spaces); otherwise
|
||||
``"#" * clamped_level + " " + text``. See the module-level ambiguity
|
||||
note: a heading whose content genuinely begins with such a run is
|
||||
kept verbatim rather than re-prefixed.
|
||||
"""
|
||||
if _MD_HEADING_RE.match(text):
|
||||
return text
|
||||
hashes = "#" * min(max(level, 1), MAX_HEADING_LEVEL)
|
||||
return f"{hashes} {text}"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_HEADING_LEVEL",
|
||||
"render_heading_line",
|
||||
"strip_heading_markdown_prefix",
|
||||
]
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Unified parser contract for native + external parser engines.
|
||||
|
||||
Every engine (native docx, mineru, docling, legacy, plus the internal
|
||||
``reuse``/``passthrough`` format handlers) implements :class:`BaseParser`.
|
||||
The pipeline dispatches through the registry
|
||||
(:mod:`lightrag.parser.registry`) instead of a growing ``if engine == …``
|
||||
chain.
|
||||
|
||||
Design notes:
|
||||
|
||||
- ``BaseParser`` carries *behaviour only* (the ``parse`` coroutine and any
|
||||
engine-private hooks). Capability metadata (supported suffixes, queue
|
||||
group, endpoint requirements) lives in the registry's lightweight
|
||||
``ParserSpec`` table so capability queries never import a parser
|
||||
implementation.
|
||||
- ``ParseResult.to_dict()`` emits only semantically-present fields so the
|
||||
returned dict stays byte-for-byte compatible with the pre-refactor
|
||||
``parse_native``/``parse_mineru``/``parse_docling`` return shapes (the
|
||||
worker reads these by ``.get(...)``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc # noqa: F401
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResolvedSource:
|
||||
"""The common parse preamble shared by every engine."""
|
||||
|
||||
source_path: Path
|
||||
document_name: str
|
||||
parsed_dir: Path
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParseContext:
|
||||
"""Inputs handed to :meth:`BaseParser.parse`.
|
||||
|
||||
Wraps the LightRAG instance plus the per-document handles so a parser
|
||||
does not receive a bare ``self``. Convenience helpers lazily import
|
||||
pipeline-layer functions at call time, keeping this module import-cheap
|
||||
and free of import cycles.
|
||||
"""
|
||||
|
||||
rag: Any
|
||||
doc_id: str
|
||||
file_path: str
|
||||
content_data: dict[str, Any]
|
||||
|
||||
def source_path(self, parser_engine: str) -> Path:
|
||||
"""Resolve the on-disk source file for this document."""
|
||||
from lightrag.pipeline import (
|
||||
_call_source_file_resolver,
|
||||
_read_source_file,
|
||||
)
|
||||
|
||||
return Path(
|
||||
_call_source_file_resolver(
|
||||
self.rag,
|
||||
self.file_path,
|
||||
source_file=_read_source_file(self.content_data),
|
||||
parser_engine=parser_engine,
|
||||
)
|
||||
)
|
||||
|
||||
def resolve(self, parser_engine: str) -> ResolvedSource:
|
||||
"""Resolve ``(source_path, document_name, parsed_dir)``.
|
||||
|
||||
Mirrors the preamble shared by ``parse_mineru``/``parse_docling``/
|
||||
``parse_native``: canonicalize the document name defensively (so
|
||||
direct callers may pass absolute or hint-bearing paths) and derive
|
||||
the ``__parsed__/<base>.parsed/`` output directory.
|
||||
"""
|
||||
from lightrag.utils_pipeline import (
|
||||
normalize_document_file_path,
|
||||
parsed_artifact_dir_for,
|
||||
)
|
||||
|
||||
source_path = self.source_path(parser_engine)
|
||||
document_name = normalize_document_file_path(self.file_path)
|
||||
if document_name == "unknown_source":
|
||||
document_name = source_path.name or f"{self.doc_id}.bin"
|
||||
parsed_dir = parsed_artifact_dir_for(
|
||||
document_name, parent_hint=source_path.parent
|
||||
)
|
||||
return ResolvedSource(source_path, document_name, parsed_dir)
|
||||
|
||||
async def archive_source(self, source_path: str) -> str | None:
|
||||
"""Archive the source after a successful parse + full_docs sync.
|
||||
|
||||
Resolved through the pipeline module's namespace (where the function
|
||||
is imported) so existing tests that patch
|
||||
``lightrag.pipeline.archive_docx_source_after_full_docs_sync`` keep
|
||||
intercepting it now that the call site lives in the parser layer.
|
||||
"""
|
||||
import lightrag.pipeline as _pipeline
|
||||
|
||||
return await _pipeline.archive_docx_source_after_full_docs_sync(source_path)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParseResult:
|
||||
"""Structured parser output.
|
||||
|
||||
``to_dict`` only emits fields that carry meaning so the dict matches the
|
||||
pre-refactor return shapes exactly (no spurious ``None``/``False`` keys).
|
||||
"""
|
||||
|
||||
doc_id: str
|
||||
file_path: str
|
||||
parse_format: str
|
||||
content: str
|
||||
blocks_path: str = ""
|
||||
parse_engine: str | None = None
|
||||
parse_stage_skipped: bool = False
|
||||
parse_warnings: dict[str, Any] | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
out: dict[str, Any] = {
|
||||
"doc_id": self.doc_id,
|
||||
"file_path": self.file_path,
|
||||
"parse_format": self.parse_format,
|
||||
"content": self.content,
|
||||
"blocks_path": self.blocks_path,
|
||||
}
|
||||
if self.parse_engine is not None:
|
||||
out["parse_engine"] = self.parse_engine
|
||||
if self.parse_stage_skipped:
|
||||
out["parse_stage_skipped"] = True
|
||||
if self.parse_warnings:
|
||||
out["parse_warnings"] = self.parse_warnings
|
||||
return out
|
||||
|
||||
|
||||
class BaseParser(ABC):
|
||||
"""Abstract base for every parser engine.
|
||||
|
||||
Subclasses set ``engine_name`` and implement :meth:`parse`. Capability
|
||||
metadata (suffixes/queue group/endpoint) is declared in the registry
|
||||
``ParserSpec``, not here.
|
||||
"""
|
||||
|
||||
engine_name: str
|
||||
|
||||
@abstractmethod
|
||||
async def parse(self, ctx: ParseContext) -> ParseResult:
|
||||
"""Parse one document and return its :class:`ParseResult`."""
|
||||
...
|
||||
@@ -0,0 +1,314 @@
|
||||
"""Unified sidecar debug CLI for any registered parser engine.
|
||||
|
||||
Dispatches one source file through the parser registry
|
||||
(``get_parser(engine).parse(...)``) and writes the resulting sidecar (and
|
||||
raw bundle, for external engines) into a flat layout — no ``__parsed__/``
|
||||
middle layer, source file never archived — so the artifacts can be
|
||||
inspected next to the input file. Because dispatch goes through the
|
||||
registry, a third-party engine registered via ``register_parser`` is a
|
||||
valid ``--engine`` choice with no CLI changes.
|
||||
|
||||
Invocation::
|
||||
|
||||
python -m lightrag.parser.cli path/to/sample.docx --engine native
|
||||
python -m lightrag.parser.cli path/to/sample.pdf --engine mineru
|
||||
python -m lightrag.parser.cli path/to/sample.pdf --engine docling --force-reparse
|
||||
|
||||
See ``docs/ParserDebugCLI-zh.md`` for the full reference.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from contextlib import ExitStack
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
|
||||
def _normalize_direct_script_sys_path() -> None:
|
||||
if __package__:
|
||||
return
|
||||
|
||||
parser_dir = Path(__file__).resolve().parent
|
||||
repo_root = parser_dir.parent.parent
|
||||
|
||||
# Direct execution adds lightrag/parser to sys.path, which makes the
|
||||
# native parser's third-party ``docx`` import resolve to parser/docx.
|
||||
sys.path[:] = [
|
||||
entry for entry in sys.path if Path(entry or ".").resolve() != parser_dir
|
||||
]
|
||||
repo_root_str = str(repo_root)
|
||||
if repo_root_str in sys.path:
|
||||
sys.path.remove(repo_root_str)
|
||||
sys.path.insert(0, repo_root_str)
|
||||
|
||||
|
||||
_normalize_direct_script_sys_path()
|
||||
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
# Registry import is cheap by design (no parser impl is pulled in), so
|
||||
# deriving --engine choices here keeps third-party engines selectable
|
||||
# without making --help pay for a heavy import.
|
||||
from lightrag.parser.registry import supported_parser_engines
|
||||
|
||||
engines = sorted(supported_parser_engines())
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="parse_sidecar",
|
||||
description=(
|
||||
"Run a registered parser engine on a single file and emit sidecar "
|
||||
"artifacts (plus a raw bundle for external engines) into a flat "
|
||||
"layout alongside the source. No __parsed__/ middle layer; the "
|
||||
"source file is never moved."
|
||||
),
|
||||
)
|
||||
parser.add_argument("input_file", type=Path, help="Source file to parse.")
|
||||
parser.add_argument(
|
||||
"--engine",
|
||||
required=True,
|
||||
choices=engines,
|
||||
help="Parser engine to drive.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--sidecar-parent-dir",
|
||||
type=Path,
|
||||
default=None,
|
||||
help=(
|
||||
"Parent directory for <name>.parsed/ and <name>.<engine>_raw/. "
|
||||
"Default: the source file's parent directory."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--doc-id",
|
||||
default=None,
|
||||
help="Override the doc id. Default: doc-<md5(absolute input path)>.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force-reparse",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Only affects mineru/docling. By default a non-empty raw_dir is "
|
||||
"treated as a valid cache and reused without manifest checks; "
|
||||
"this flag clears raw_dir and forces a fresh download/parse."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preview",
|
||||
type=int,
|
||||
default=5,
|
||||
metavar="N",
|
||||
help="Number of block rows to preview after parsing (0 disables).",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def _print_summary(blocks_path: Path, raw_dir: Path | None, preview: int) -> None:
|
||||
with blocks_path.open("r", encoding="utf-8") as fh:
|
||||
meta_line = fh.readline().strip()
|
||||
if not meta_line:
|
||||
raise SystemExit(f"empty blocks file at {blocks_path}")
|
||||
meta = json.loads(meta_line)
|
||||
rows = [json.loads(line) for line in fh if line.strip()]
|
||||
parsed_dir = blocks_path.parent
|
||||
print(f"parsed dir : {parsed_dir} (exists={parsed_dir.exists()})")
|
||||
if raw_dir is not None:
|
||||
print(f"raw dir : {raw_dir} (exists={raw_dir.exists()})")
|
||||
print(f"document : {meta.get('document_name')}")
|
||||
print(f"doc_id : {meta.get('doc_id')}")
|
||||
print(f"engine : {meta.get('parse_engine')}")
|
||||
print(f"blocks : {meta.get('blocks')}")
|
||||
print(
|
||||
f"sidecars : tables={meta.get('table_file')} "
|
||||
f"drawings={meta.get('drawing_file')} "
|
||||
f"equations={meta.get('equation_file')} "
|
||||
f"asset_dir={meta.get('asset_dir')}"
|
||||
)
|
||||
if preview > 0 and rows:
|
||||
shown = min(preview, len(rows))
|
||||
print(f"--- preview (first {shown} of {len(rows)} blocks) ---")
|
||||
for row in rows[:preview]:
|
||||
heading = row.get("heading") or ""
|
||||
content = (row.get("content") or "").replace("\n", " ")
|
||||
snippet = content if len(content) <= 80 else content[:77] + "..."
|
||||
print(f" [{row.get('blockid', '')[:8]}] heading={heading!r} :: {snippet}")
|
||||
|
||||
|
||||
def _print_raw_summary(result: dict, preview: int) -> None:
|
||||
"""Summary for engines that emit plain content with no sidecar (legacy /
|
||||
any non-sidecar third-party engine: ``blocks_path`` is empty)."""
|
||||
content = result.get("content") or ""
|
||||
print(f"engine : {result.get('parse_engine')}")
|
||||
print(f"format : {result.get('parse_format')}")
|
||||
print(f"content : {len(content)} chars")
|
||||
if preview > 0 and content:
|
||||
snippet = content[:400]
|
||||
print(f"--- preview (first {len(snippet)} of {len(content)} chars) ---")
|
||||
print(snippet + ("..." if len(content) > len(snippet) else ""))
|
||||
|
||||
|
||||
async def _run(args: argparse.Namespace) -> int:
|
||||
# Pipeline + heavy parser imports are deferred so ``--help`` and the
|
||||
# input-file existence check don't pay for them.
|
||||
from lightrag.constants import FULL_DOCS_FORMAT_PENDING_PARSE
|
||||
from lightrag.parser.base import ParseContext
|
||||
from lightrag.parser.external._base import ExternalParserBase
|
||||
from lightrag.parser.registry import get_parser, suffix_capabilities
|
||||
from lightrag.parser.debug import build_debug_rag
|
||||
from lightrag.parser.docx.parse_document import DocxContentError
|
||||
from lightrag.utils import compute_mdhash_id
|
||||
import lightrag.pipeline as pipeline_mod
|
||||
import lightrag.utils_pipeline as utils_pipeline_mod
|
||||
|
||||
source = args.input_file.resolve()
|
||||
if not source.is_file():
|
||||
print(f"error: input file does not exist: {source}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
# Reject suffix/engine mismatches up-front: the pipeline would otherwise
|
||||
# fail deep inside the IR builder with a less helpful message.
|
||||
suffix = source.suffix.lstrip(".").lower()
|
||||
supported = suffix_capabilities(args.engine)
|
||||
if suffix not in supported:
|
||||
print(
|
||||
f"error: engine '{args.engine}' does not support .{suffix or '<no suffix>'} "
|
||||
f"files (supported: {', '.join(sorted(supported))})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
parser = get_parser(args.engine)
|
||||
if parser is None:
|
||||
print(f"error: engine '{args.engine}' is not registered", file=sys.stderr)
|
||||
return 1
|
||||
is_external = isinstance(parser, ExternalParserBase)
|
||||
|
||||
sidecar_parent = (args.sidecar_parent_dir or source.parent).resolve()
|
||||
sidecar_parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
parsed_dir = sidecar_parent / f"{source.name}.parsed"
|
||||
# External engines preserve a raw bundle next to the sidecar; its name is
|
||||
# derived from the engine's own raw_dir_suffix (mirrors the flattened
|
||||
# raw_dir_for_parsed_dir layout), so any external engine works here.
|
||||
raw_dir = (
|
||||
sidecar_parent / f"{source.name}{parser.raw_dir_suffix}"
|
||||
if is_external
|
||||
else None
|
||||
)
|
||||
|
||||
doc_id = args.doc_id or compute_mdhash_id(str(source), prefix="doc-")
|
||||
|
||||
def _patched_artifact_dir(
|
||||
file_path: str | None = None,
|
||||
*,
|
||||
parent_hint: Any | None = None,
|
||||
) -> Path:
|
||||
# Flatten the production "<INPUT_DIR>/__parsed__/<base>.parsed/"
|
||||
# layout to "<sidecar_parent>/<source.name>.parsed/" so the sidecar
|
||||
# and the source file sit side by side.
|
||||
return parsed_dir
|
||||
|
||||
def _lenient_bundle(
|
||||
raw_dir_arg: Path, _source_file: Path | None = None, *_args: Any, **_kwargs: Any
|
||||
) -> bool:
|
||||
# Tolerate the ExternalParserBase hook's keyword-only ``engine_params``
|
||||
# (and any future args) — this stub replaces ``is_bundle_valid``.
|
||||
return raw_dir_arg.exists() and any(raw_dir_arg.iterdir())
|
||||
|
||||
def _force_miss(*_args: Any, **_kwargs: Any) -> bool:
|
||||
return False
|
||||
|
||||
bundle_check = _force_miss if args.force_reparse else _lenient_bundle
|
||||
|
||||
async def _noop_archive(*_args: Any, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
rag = build_debug_rag()
|
||||
|
||||
with ExitStack() as stack:
|
||||
# Patch 1: redirect sidecar output to the flat layout.
|
||||
# parsed_artifact_dir_for is from-imported into pipeline at
|
||||
# module load, so patch both namespaces.
|
||||
stack.enter_context(
|
||||
mock.patch.object(
|
||||
utils_pipeline_mod,
|
||||
"parsed_artifact_dir_for",
|
||||
_patched_artifact_dir,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
mock.patch.object(
|
||||
pipeline_mod,
|
||||
"parsed_artifact_dir_for",
|
||||
_patched_artifact_dir,
|
||||
)
|
||||
)
|
||||
|
||||
# Patch 2: raw cache strategy. ExternalParserBase.parse calls
|
||||
# ``self.is_bundle_valid(...)``, so patch the resolved instance method
|
||||
# directly — works for any external engine without knowing its module.
|
||||
if is_external:
|
||||
stack.enter_context(
|
||||
mock.patch.object(parser, "is_bundle_valid", bundle_check)
|
||||
)
|
||||
|
||||
# Patch 3: keep the source file in place. Every engine archives via
|
||||
# ctx.archive_source -> lightrag.pipeline.archive_docx_source_after_...
|
||||
stack.enter_context(
|
||||
mock.patch.object(
|
||||
pipeline_mod,
|
||||
"archive_docx_source_after_full_docs_sync",
|
||||
_noop_archive,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
result = (
|
||||
await parser.parse(
|
||||
ParseContext(
|
||||
rag,
|
||||
doc_id,
|
||||
str(source),
|
||||
{
|
||||
"parse_format": FULL_DOCS_FORMAT_PENDING_PARSE,
|
||||
"content": "",
|
||||
},
|
||||
)
|
||||
)
|
||||
).to_dict()
|
||||
except DocxContentError as exc:
|
||||
# The native DOCX parser now raises (instead of sys.exit) on a
|
||||
# content-limit violation so the server pipeline can fail just the
|
||||
# offending document. Preserve the friendly CLI UX: print the
|
||||
# formatted message and return a non-zero exit code, no traceback.
|
||||
print(str(exc), file=sys.stderr)
|
||||
return 1
|
||||
|
||||
blocks_path = result.get("blocks_path") or ""
|
||||
if blocks_path:
|
||||
_print_summary(Path(blocks_path), raw_dir, args.preview)
|
||||
else:
|
||||
# No sidecar (legacy / non-sidecar third-party engine): the result
|
||||
# carries plain content, so summarize that instead of a blocks file.
|
||||
_print_raw_summary(result, args.preview)
|
||||
return 0
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
# Discover third-party parser engines before --engine choices are built,
|
||||
# so a `lightrag.parsers` entry-point engine is selectable here exactly
|
||||
# like in the server (which loads plugins in create_app).
|
||||
from lightrag.parser.plugins import load_third_party_parsers
|
||||
|
||||
load_third_party_parsers()
|
||||
args = _build_parser().parse_args(argv)
|
||||
return asyncio.run(_run(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Shared debug LightRAG stand-in for registry-dispatched parsing.
|
||||
|
||||
A minimal ``LightRAG`` stand-in plus a deterministic ``datetime`` shim,
|
||||
shared by the unified parser debug CLI (``lightrag/parser/cli.py``),
|
||||
the golden-fixture regen script (``scripts/regen_native_docx_golden.py``),
|
||||
and the byte-equivalence golden tests
|
||||
(``tests/parser/docx/test_native_docx_golden.py``).
|
||||
|
||||
Every engine is driven the same way — ``get_parser(engine).parse(
|
||||
ParseContext(rag, ...))`` — and ``ParseContext`` reads the same ``rag``
|
||||
surface (``_persist_parsed_full_docs``, ``_resolve_source_file_for_parser``,
|
||||
``self.full_docs``, ``self.doc_status``), so a single stand-in covers every
|
||||
engine — when a parser grows a new dependency, extend this module rather
|
||||
than copy-pasting parallel stubs into each call site.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
|
||||
class DebugFullDocs:
|
||||
"""In-memory ``full_docs`` shim — captures the persisted record."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.data: dict[str, Any] = {}
|
||||
|
||||
async def upsert(self, payload: dict[str, Any]) -> None:
|
||||
self.data.update(payload)
|
||||
|
||||
async def get_by_id(self, doc_id: str) -> Any:
|
||||
return self.data.get(doc_id)
|
||||
|
||||
async def index_done_callback(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class DebugDocStatus:
|
||||
"""No-op ``doc_status`` shim — the parse_* methods never read/write it."""
|
||||
|
||||
async def get_by_id(self, doc_id: str) -> Any:
|
||||
return None
|
||||
|
||||
async def upsert(self, data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def build_debug_rag():
|
||||
"""Build a minimal LightRAG stand-in that exposes what a parser reads.
|
||||
|
||||
The import of ``LightRAG`` is intentionally function-local: deferring
|
||||
it avoids a circular import when this helper is loaded during package
|
||||
init (the parser CLI resolves ``lightrag.parser.debug`` before
|
||||
``lightrag`` itself is fully bound).
|
||||
|
||||
A parser is driven via ``get_parser(engine).parse(ParseContext(rag, ...))``
|
||||
(CLI / golden tests / regen script). ``ParseContext`` reads these off the
|
||||
``rag`` it is handed — every entry MUST be provided by this stand-in:
|
||||
|
||||
- **methods** (rebound from :class:`LightRAG`):
|
||||
- ``_persist_parsed_full_docs(doc_id, payload)`` — async; touches
|
||||
``self.full_docs``.
|
||||
- ``_resolve_source_file_for_parser(file_path)`` — returns the
|
||||
on-disk source path. Stubbed to identity here since the CLI / tests
|
||||
feed an already-resolved path.
|
||||
- **storages**:
|
||||
- ``self.full_docs.upsert(...)`` / ``.get_by_id(...)`` /
|
||||
``.index_done_callback()`` — :class:`DebugFullDocs` covers all three.
|
||||
- ``self.doc_status.get_by_id(...)`` / ``.upsert(...)`` —
|
||||
:class:`DebugDocStatus` covers both.
|
||||
|
||||
When a parser grows a new dependency on the ``rag`` handed to
|
||||
``ParseContext``, extend this stand-in (and update the list above) rather
|
||||
than copy-pasting a parallel stub into the call sites.
|
||||
"""
|
||||
from lightrag import LightRAG
|
||||
|
||||
class _DebugRag:
|
||||
_persist_parsed_full_docs = LightRAG._persist_parsed_full_docs
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.full_docs = DebugFullDocs()
|
||||
self.doc_status = DebugDocStatus()
|
||||
|
||||
def _resolve_source_file_for_parser(self, file_path: str) -> str:
|
||||
return file_path
|
||||
|
||||
return _DebugRag()
|
||||
|
||||
|
||||
_FROZEN_NOW = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
class FrozenDateTime(datetime):
|
||||
"""Pin ``datetime.now`` so ``write_sidecar`` stamps a deterministic time."""
|
||||
|
||||
@classmethod
|
||||
def now(cls, tz=None): # noqa: D401
|
||||
return _FROZEN_NOW if tz is None else _FROZEN_NOW.astimezone(tz)
|
||||
@@ -0,0 +1,14 @@
|
||||
"""LightRAG native DOCX parser package.
|
||||
|
||||
The :mod:`parse_document` / :mod:`numbering_resolver` / :mod:`table_extractor` /
|
||||
:mod:`drawing_image_extractor` / :mod:`utils` / :mod:`omml` modules ship the
|
||||
upstream DOCX extraction logic verbatim (with imports localized for the new
|
||||
package path).
|
||||
|
||||
The pipeline-side orchestration (extract → IR → sidecar) now lives in
|
||||
:meth:`lightrag.pipeline._PipelineMixin.parse_native` so the native and
|
||||
MinerU engines share one shape; see :mod:`lightrag.parser.docx.ir_builder`
|
||||
for the engine IR builder.
|
||||
"""
|
||||
|
||||
__all__: list[str] = []
|
||||
@@ -0,0 +1,445 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ABOUTME: Shared drawing/image extraction utilities for DOCX parsing and editing
|
||||
ABOUTME: Resolves w:drawing -> a:blip relationships, exports embedded images, builds placeholders
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import posixpath
|
||||
import re
|
||||
import shutil
|
||||
import zipfile
|
||||
from dataclasses import dataclass, field
|
||||
from html import escape, unescape
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Dict, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
try:
|
||||
from defusedxml import ElementTree as ET
|
||||
except ImportError: # pragma: no cover
|
||||
from xml.etree import ElementTree as ET
|
||||
|
||||
|
||||
NS = {
|
||||
"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
|
||||
"wp": "http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing",
|
||||
"a": "http://schemas.openxmlformats.org/drawingml/2006/main",
|
||||
"r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
|
||||
"v": "urn:schemas-microsoft-com:vml",
|
||||
}
|
||||
|
||||
REL_NS = "http://schemas.openxmlformats.org/package/2006/relationships"
|
||||
CONTENT_TYPE_NS = "http://schemas.openxmlformats.org/package/2006/content-types"
|
||||
IMAGE_REL_TYPE = (
|
||||
"http://schemas.openxmlformats.org/officeDocument/2006/relationships/image"
|
||||
)
|
||||
SOURCE_DOCUMENT_PART = "/word/document.xml"
|
||||
|
||||
# Match old and new drawing placeholders (requires id/name, allows extra attributes)
|
||||
DRAWING_PATTERN = re.compile(
|
||||
r'<drawing\b(?=[^>]*\bid="[^"]*")(?=[^>]*\bname="[^"]*")[^>]*/>'
|
||||
)
|
||||
DRAWING_TAG_PATTERN = re.compile(r"<drawing\b[^>]*/>")
|
||||
DRAWING_ATTR_PATTERN = re.compile(r'([a-zA-Z_][\w:.-]*)="([^"]*)"')
|
||||
|
||||
|
||||
@dataclass
|
||||
class DrawingRelationship:
|
||||
"""Relationship metadata for a single relationship ID."""
|
||||
|
||||
rel_id: str
|
||||
target: str
|
||||
target_mode: str
|
||||
rel_type: str
|
||||
part_name: Optional[str] = None
|
||||
content_type: Optional[str] = None
|
||||
image_format: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DrawingExtractionContext:
|
||||
"""Context used to resolve and export drawing images for one DOCX file."""
|
||||
|
||||
docx_path: Path
|
||||
blocks_output_path: Optional[Path] = None
|
||||
export_dir_name: Optional[str] = None
|
||||
export_dir_path: Optional[Path] = None
|
||||
relationships: Dict[str, DrawingRelationship] = field(default_factory=dict)
|
||||
_exported_part_to_relpath: Dict[str, str] = field(default_factory=dict)
|
||||
_used_filenames: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
def resolve_relationship(self, rel_id: str) -> Optional[DrawingRelationship]:
|
||||
return self.relationships.get(rel_id)
|
||||
|
||||
def export_embedded_image(self, rel: DrawingRelationship) -> Optional[str]:
|
||||
"""
|
||||
Export an embedded image relationship target to export_dir.
|
||||
|
||||
Returns:
|
||||
Relative path like "<blocks_stem>.image/image1.png" if exported,
|
||||
or None when export is not applicable.
|
||||
"""
|
||||
if not self.export_dir_path or not self.export_dir_name:
|
||||
return None
|
||||
if rel.target_mode.lower() == "external":
|
||||
return None
|
||||
if not rel.part_name:
|
||||
return None
|
||||
if rel.part_name in self._exported_part_to_relpath:
|
||||
return self._exported_part_to_relpath[rel.part_name]
|
||||
|
||||
zip_member = rel.part_name.lstrip("/")
|
||||
try:
|
||||
with zipfile.ZipFile(self.docx_path, "r") as zf:
|
||||
blob = zf.read(zip_member)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
filename = self._dedupe_filename(PurePosixPath(rel.part_name).name or "image")
|
||||
output_file = self.export_dir_path / filename
|
||||
output_file.write_bytes(blob)
|
||||
|
||||
rel_path = str(PurePosixPath(self.export_dir_name) / filename)
|
||||
self._exported_part_to_relpath[rel.part_name] = rel_path
|
||||
return rel_path
|
||||
|
||||
def _dedupe_filename(self, base_name: str) -> str:
|
||||
if base_name not in self._used_filenames:
|
||||
self._used_filenames[base_name] = base_name
|
||||
return base_name
|
||||
|
||||
stem = Path(base_name).stem
|
||||
suffix = Path(base_name).suffix
|
||||
index = 2
|
||||
while True:
|
||||
candidate = f"{stem}_{index}{suffix}"
|
||||
if candidate not in self._used_filenames:
|
||||
self._used_filenames[candidate] = candidate
|
||||
return candidate
|
||||
index += 1
|
||||
|
||||
|
||||
def _normalize_image_format(ext_or_type: str) -> Optional[str]:
|
||||
if not ext_or_type:
|
||||
return None
|
||||
value = ext_or_type.strip().lower()
|
||||
|
||||
# Content-Type
|
||||
if value.startswith("image/"):
|
||||
value = value.split("/", 1)[1]
|
||||
if "+" in value:
|
||||
value = value.split("+", 1)[0]
|
||||
if value.startswith("x-"):
|
||||
value = value[2:]
|
||||
|
||||
# Extension (with or without leading dot)
|
||||
value = value.lstrip(".")
|
||||
if value == "jpg":
|
||||
return "jpeg"
|
||||
if value in {"jpeg", "png", "gif", "bmp", "tiff", "webp", "svg", "emf", "wmf"}:
|
||||
return value
|
||||
return value or None
|
||||
|
||||
|
||||
def _infer_format_from_target(target: str) -> Optional[str]:
|
||||
if not target:
|
||||
return None
|
||||
parsed = urlparse(target)
|
||||
path = parsed.path if parsed.scheme else target
|
||||
suffix = PurePosixPath(path).suffix
|
||||
return _normalize_image_format(suffix)
|
||||
|
||||
|
||||
def _resolve_part_name(source_part_name: str, target: str) -> str:
|
||||
if target.startswith("/"):
|
||||
return posixpath.normpath(target)
|
||||
source_dir = posixpath.dirname(source_part_name)
|
||||
joined = posixpath.join(source_dir, target)
|
||||
normalized = posixpath.normpath(joined)
|
||||
if not normalized.startswith("/"):
|
||||
normalized = "/" + normalized
|
||||
return normalized
|
||||
|
||||
|
||||
def create_drawing_context(
|
||||
docx_path: str,
|
||||
blocks_output_path: Optional[str] = None,
|
||||
) -> DrawingExtractionContext:
|
||||
"""
|
||||
Create extraction context for a DOCX file.
|
||||
|
||||
If blocks_output_path is provided, this also prepares `<blocks_stem>.image/`
|
||||
beside the blocks file and clears any previous content.
|
||||
"""
|
||||
docx_file = Path(docx_path)
|
||||
ctx = DrawingExtractionContext(docx_path=docx_file)
|
||||
|
||||
if blocks_output_path:
|
||||
output_path = Path(blocks_output_path)
|
||||
export_dir_name = f"{output_path.stem}.image"
|
||||
export_dir_path = output_path.parent / export_dir_name
|
||||
if export_dir_path.exists():
|
||||
shutil.rmtree(export_dir_path)
|
||||
export_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
ctx.blocks_output_path = output_path
|
||||
ctx.export_dir_name = export_dir_name
|
||||
ctx.export_dir_path = export_dir_path
|
||||
|
||||
load_relationships(ctx)
|
||||
return ctx
|
||||
|
||||
|
||||
def load_relationships(ctx: DrawingExtractionContext) -> None:
|
||||
rels_xml = "word/_rels/document.xml.rels"
|
||||
content_types_xml = "[Content_Types].xml"
|
||||
|
||||
overrides: Dict[str, str] = {}
|
||||
defaults: Dict[str, str] = {}
|
||||
|
||||
try:
|
||||
with zipfile.ZipFile(ctx.docx_path, "r") as zf:
|
||||
if content_types_xml in zf.namelist():
|
||||
ct_root = ET.parse(zf.open(content_types_xml)).getroot()
|
||||
for node in ct_root.findall(f".//{{{CONTENT_TYPE_NS}}}Override"):
|
||||
part_name = node.get("PartName")
|
||||
content_type = node.get("ContentType")
|
||||
if part_name and content_type:
|
||||
overrides[part_name] = content_type
|
||||
for node in ct_root.findall(f".//{{{CONTENT_TYPE_NS}}}Default"):
|
||||
ext = node.get("Extension")
|
||||
content_type = node.get("ContentType")
|
||||
if ext and content_type:
|
||||
defaults[ext.lower()] = content_type
|
||||
|
||||
if rels_xml not in zf.namelist():
|
||||
return
|
||||
rels_root = ET.parse(zf.open(rels_xml)).getroot()
|
||||
except Exception:
|
||||
return
|
||||
|
||||
for rel in rels_root.findall(f".//{{{REL_NS}}}Relationship"):
|
||||
rel_id = rel.get("Id")
|
||||
target = rel.get("Target", "")
|
||||
target_mode = rel.get("TargetMode", "")
|
||||
rel_type = rel.get("Type", "")
|
||||
if not rel_id:
|
||||
continue
|
||||
|
||||
part_name = None
|
||||
content_type = None
|
||||
image_format = None
|
||||
|
||||
if target_mode.lower() != "external":
|
||||
part_name = _resolve_part_name(SOURCE_DOCUMENT_PART, target)
|
||||
if part_name:
|
||||
content_type = overrides.get(part_name)
|
||||
if not content_type:
|
||||
ext = PurePosixPath(part_name).suffix.lower().lstrip(".")
|
||||
content_type = defaults.get(ext)
|
||||
image_format = _normalize_image_format(
|
||||
content_type or _infer_format_from_target(part_name)
|
||||
)
|
||||
else:
|
||||
image_format = _normalize_image_format(_infer_format_from_target(target))
|
||||
|
||||
ctx.relationships[rel_id] = DrawingRelationship(
|
||||
rel_id=rel_id,
|
||||
target=target,
|
||||
target_mode=target_mode,
|
||||
rel_type=rel_type,
|
||||
part_name=part_name,
|
||||
content_type=content_type,
|
||||
image_format=image_format,
|
||||
)
|
||||
|
||||
|
||||
def _extract_blip_relationship(drawing_elem) -> Optional[Tuple[str, str]]:
|
||||
for blip in drawing_elem.findall(".//a:blip", NS):
|
||||
# Prefer explicit external links when both link/embed are present on one blip.
|
||||
# Word may keep an embedded cache for linked pictures.
|
||||
rel_link = blip.get(f"{{{NS['r']}}}link")
|
||||
if rel_link:
|
||||
return "link", rel_link
|
||||
rel_embed = blip.get(f"{{{NS['r']}}}embed")
|
||||
if rel_embed:
|
||||
return "embed", rel_embed
|
||||
return None
|
||||
|
||||
|
||||
def _extract_imagedata_relationship(container_elem) -> Optional[str]:
|
||||
"""Find an image relationship id from a w:pict / w:object via v:imagedata.
|
||||
|
||||
These legacy VML containers are how Word references EMF/WMF metafiles
|
||||
(and the rendered preview of any embedded OLE object). v:imagedata uses
|
||||
``r:id`` to point at the image part for both embedded and externally
|
||||
linked images — the relationship's ``TargetMode`` is what disambiguates
|
||||
the two cases, so the caller must inspect the resolved relationship.
|
||||
"""
|
||||
r_id_attr = f"{{{NS['r']}}}id"
|
||||
for imgdata in container_elem.findall(".//v:imagedata", NS):
|
||||
rel_id = imgdata.get(r_id_attr)
|
||||
if rel_id:
|
||||
return rel_id
|
||||
return None
|
||||
|
||||
|
||||
def _build_placeholder(attrs: Dict[str, str]) -> str:
|
||||
ordered_keys = ["id", "name", "path", "format"]
|
||||
pieces = []
|
||||
for key in ordered_keys:
|
||||
if key in attrs and attrs[key] is not None:
|
||||
pieces.append(f'{key}="{escape(str(attrs[key]), quote=True)}"')
|
||||
|
||||
# Preserve extra attributes deterministically (sorted by name)
|
||||
for key in sorted(k for k in attrs.keys() if k not in ordered_keys):
|
||||
value = attrs[key]
|
||||
if value is not None:
|
||||
pieces.append(f'{key}="{escape(str(value), quote=True)}"')
|
||||
|
||||
return f"<drawing {' '.join(pieces)} />"
|
||||
|
||||
|
||||
def extract_drawing_placeholder_from_element(
|
||||
drawing_elem,
|
||||
context: Optional[DrawingExtractionContext] = None,
|
||||
include_extended_attrs: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
Build a <drawing ... /> placeholder from a w:drawing element.
|
||||
|
||||
Behavior:
|
||||
- Always emits id/name from wp:docPr when present.
|
||||
- For embedded images (a:blip@r:embed): exports image and sets path/format.
|
||||
- For linked images (a:blip@r:link): does not download; path is original link target.
|
||||
- When no image reference exists (e.g. chart drawing): keeps id/name only.
|
||||
"""
|
||||
doc_pr = drawing_elem.find(".//wp:docPr", NS)
|
||||
attrs = {
|
||||
"id": doc_pr.get("id", "") if doc_pr is not None else "",
|
||||
"name": doc_pr.get("name", "") if doc_pr is not None else "",
|
||||
}
|
||||
|
||||
if include_extended_attrs:
|
||||
rel_ref = _extract_blip_relationship(drawing_elem)
|
||||
if rel_ref is not None and context is not None:
|
||||
rel_kind, rel_id = rel_ref
|
||||
rel = context.resolve_relationship(rel_id)
|
||||
if rel is not None:
|
||||
if rel_kind == "embed" and rel.rel_type == IMAGE_REL_TYPE:
|
||||
rel_path = context.export_embedded_image(rel)
|
||||
if rel_path:
|
||||
attrs["path"] = rel_path
|
||||
if rel.image_format:
|
||||
attrs["format"] = rel.image_format
|
||||
elif rel_kind == "link":
|
||||
if rel.target:
|
||||
attrs["path"] = rel.target
|
||||
if rel.image_format:
|
||||
attrs["format"] = rel.image_format
|
||||
|
||||
return _build_placeholder(attrs)
|
||||
|
||||
|
||||
def extract_vml_image_placeholder_from_element(
|
||||
container_elem,
|
||||
context: Optional[DrawingExtractionContext] = None,
|
||||
include_extended_attrs: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
Build a <drawing ... /> placeholder from a w:pict or w:object element.
|
||||
|
||||
Legacy Word documents and OLE-embedded objects (Visio diagrams, equation
|
||||
editor previews, etc.) expose their rendered image via VML rather than
|
||||
DrawingML. The image is referenced through ``<v:imagedata r:id="..."/>``
|
||||
inside ``<v:shape>``, and the underlying bytes are commonly EMF/WMF
|
||||
metafiles. This function exports those bytes through the same context as
|
||||
DrawingML images so EMF/WMF assets land in the blocks.assets directory
|
||||
alongside PNG/JPEG ones.
|
||||
|
||||
The output placeholder format matches
|
||||
``extract_drawing_placeholder_from_element`` so downstream consumers
|
||||
treat both paths uniformly.
|
||||
"""
|
||||
shape = container_elem.find(".//v:shape", NS)
|
||||
attrs = {
|
||||
"id": shape.get("id", "") if shape is not None else "",
|
||||
"name": shape.get("alt", "") if shape is not None else "",
|
||||
}
|
||||
|
||||
if include_extended_attrs:
|
||||
rel_id = _extract_imagedata_relationship(container_elem)
|
||||
if rel_id and context is not None:
|
||||
rel = context.resolve_relationship(rel_id)
|
||||
if rel is not None and rel.rel_type == IMAGE_REL_TYPE:
|
||||
# VML reuses r:id for both embedded image parts and externally
|
||||
# linked images; only the resolved TargetMode tells us which.
|
||||
# Treating an external relationship as embedded would call
|
||||
# export_embedded_image() (which short-circuits on external)
|
||||
# and silently drop the linked path.
|
||||
if rel.target_mode.lower() == "external":
|
||||
if rel.target:
|
||||
attrs["path"] = rel.target
|
||||
if rel.image_format:
|
||||
attrs["format"] = rel.image_format
|
||||
else:
|
||||
rel_path = context.export_embedded_image(rel)
|
||||
if rel_path:
|
||||
attrs["path"] = rel_path
|
||||
if rel.image_format:
|
||||
attrs["format"] = rel.image_format
|
||||
|
||||
return _build_placeholder(attrs)
|
||||
|
||||
|
||||
def parse_drawing_attributes(placeholder: str) -> Dict[str, str]:
|
||||
"""Parse attributes from a <drawing ... /> placeholder."""
|
||||
return {
|
||||
name: unescape(value)
|
||||
for name, value in DRAWING_ATTR_PATTERN.findall(placeholder)
|
||||
}
|
||||
|
||||
|
||||
def normalize_drawing_placeholder(
|
||||
placeholder: str,
|
||||
include_extended_attrs: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
Normalize one drawing placeholder into canonical attribute order.
|
||||
|
||||
Args:
|
||||
placeholder: Input placeholder string
|
||||
include_extended_attrs: If False, keeps only id/name.
|
||||
"""
|
||||
attrs = parse_drawing_attributes(placeholder)
|
||||
normalized = {
|
||||
"id": attrs.get("id", ""),
|
||||
"name": attrs.get("name", ""),
|
||||
}
|
||||
if include_extended_attrs:
|
||||
if "path" in attrs:
|
||||
normalized["path"] = attrs["path"]
|
||||
if "format" in attrs:
|
||||
normalized["format"] = attrs["format"]
|
||||
for key, value in attrs.items():
|
||||
if key not in {"id", "name", "path", "format"}:
|
||||
normalized[key] = value
|
||||
return _build_placeholder(normalized)
|
||||
|
||||
|
||||
def normalize_drawing_placeholders_in_text(
|
||||
text: str,
|
||||
include_extended_attrs: bool = False,
|
||||
) -> str:
|
||||
"""Normalize all drawing placeholders inside a text blob."""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
def _replace(match: re.Match) -> str:
|
||||
return normalize_drawing_placeholder(
|
||||
match.group(0),
|
||||
include_extended_attrs=include_extended_attrs,
|
||||
)
|
||||
|
||||
return DRAWING_TAG_PATTERN.sub(_replace, text)
|
||||
@@ -0,0 +1,339 @@
|
||||
"""Native DOCX IR builder: ``extract_docx_blocks`` output → :class:`IRDoc`.
|
||||
|
||||
Input contract: a list of block dicts as produced by
|
||||
``lightrag.parser.docx.parse_document.extract_docx_blocks``. Each
|
||||
block carries ``content`` text in which ``<table>``, ``<equation>`` and
|
||||
``<drawing …/>`` placeholders are already embedded by the upstream parser.
|
||||
The builder rewrites those placeholders into IR placeholder tokens
|
||||
(``{{TBL:k}} / {{EQ:k}} / {{EQI:k}} / {{IMG:k}}``) and builds the matching
|
||||
``IRTable`` / ``IREquation`` / ``IRDrawing`` items.
|
||||
|
||||
Asset bytes are extracted to disk by the upstream parser *before* this
|
||||
builder runs (via ``DrawingExtractionContext`` passed to
|
||||
``extract_docx_blocks``). The builder therefore declares assets with
|
||||
``AssetSpec.source=None`` — the writer records each entry's size without
|
||||
copying.
|
||||
|
||||
Block-vs-inline equation distinction follows the legacy native rule: an
|
||||
``<equation>…</equation>`` tag is *block* iff each side is either the
|
||||
content boundary or a ``\\n`` character. Anything else stays inline,
|
||||
keeps its tag in block text without an id, and never enters
|
||||
``equations.json``.
|
||||
|
||||
Positions are always emitted as ``IRPosition(type="paraid", range=[start,
|
||||
end])`` where each side may be ``None`` (legacy / non-Word docx authors
|
||||
sometimes omit ``w14:paraId``). The writer's ``to_jsonable`` faithfully
|
||||
preserves the per-side null so consumers can distinguish "start missing"
|
||||
vs "both missing".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Any
|
||||
|
||||
from lightrag.parser.docx.drawing_image_extractor import (
|
||||
DRAWING_TAG_PATTERN,
|
||||
parse_drawing_attributes,
|
||||
)
|
||||
from lightrag.sidecar.ir import (
|
||||
AssetSpec,
|
||||
IRBlock,
|
||||
IRDoc,
|
||||
IRDrawing,
|
||||
IREquation,
|
||||
IRPosition,
|
||||
IRTable,
|
||||
)
|
||||
|
||||
|
||||
_TABLE_TAG_RE = re.compile(r"<table>(.*?)</table>", re.DOTALL)
|
||||
_EQUATION_TAG_RE = re.compile(r"<equation>(.*?)</equation>", re.DOTALL)
|
||||
|
||||
|
||||
def _normalize_dimension(rows_value: Any) -> tuple[int, int]:
|
||||
if not isinstance(rows_value, list):
|
||||
return 0, 0
|
||||
num_rows = len(rows_value)
|
||||
num_cols = max((len(r) for r in rows_value if isinstance(r, list)), default=0)
|
||||
return num_rows, num_cols
|
||||
|
||||
|
||||
def _placeholder_keyspace() -> Callable[[str], str]:
|
||||
"""Return a fresh counter producing ``{prefix}{N}`` keys (1-indexed)."""
|
||||
counter = itertools.count(1)
|
||||
return lambda prefix: f"{prefix}{next(counter)}"
|
||||
|
||||
|
||||
def _safe_asset_ref_from_path(path_val: str, asset_prefix: str) -> str | None:
|
||||
"""Return the path inside ``asset_prefix`` only when it is safe.
|
||||
|
||||
Native DOCX images are pre-extracted into ``<base>.blocks.assets/``.
|
||||
Treat a drawing path as local only when the suffix is a clean POSIX
|
||||
relative path. Unsafe local-looking paths are dropped instead of being
|
||||
registered as assets or preserved as linked references.
|
||||
"""
|
||||
if not asset_prefix or not path_val.startswith(asset_prefix):
|
||||
return None
|
||||
|
||||
rel_raw = path_val[len(asset_prefix) :]
|
||||
if not rel_raw or "\\" in rel_raw:
|
||||
return None
|
||||
|
||||
rel_path = PurePosixPath(rel_raw)
|
||||
if rel_path.is_absolute():
|
||||
return None
|
||||
if any(part == ".." for part in rel_path.parts):
|
||||
return None
|
||||
|
||||
rel = rel_path.as_posix()
|
||||
if rel in {"", "."}:
|
||||
return None
|
||||
return rel
|
||||
|
||||
|
||||
@dataclass
|
||||
class _BlockBuilder:
|
||||
"""Per-block scratch state for the three ``re.sub`` rewrite passes.
|
||||
|
||||
Keeping the replacer routines as bound methods (rather than closures
|
||||
redefined inside the per-block loop) means they're compiled once at
|
||||
class-load and the state they mutate — ``tables`` / ``drawings`` /
|
||||
``equations`` / ``table_position`` — is held explicitly rather than
|
||||
captured implicitly from the enclosing frame.
|
||||
"""
|
||||
|
||||
next_key: Callable[[str], str]
|
||||
assets: list[AssetSpec]
|
||||
seen_asset_refs: set[str]
|
||||
asset_prefix: str
|
||||
block_table_headers: list[Any]
|
||||
tables: list[IRTable] = field(default_factory=list)
|
||||
drawings: list[IRDrawing] = field(default_factory=list)
|
||||
equations: list[IREquation] = field(default_factory=list)
|
||||
# Position of the *next* ``<table>`` placeholder within this block,
|
||||
# used to look up the matching entry in ``block_table_headers``.
|
||||
table_position: int = 0
|
||||
|
||||
def replace_table(self, match: "re.Match[str]") -> str:
|
||||
table_body_raw = match.group(1)
|
||||
try:
|
||||
rows = json.loads(table_body_raw)
|
||||
if not isinstance(rows, list):
|
||||
rows = None
|
||||
except json.JSONDecodeError:
|
||||
rows = None
|
||||
|
||||
if rows is not None:
|
||||
parsed_rows: list[list[str]] | None = [
|
||||
[str(c) for c in r] if isinstance(r, list) else [str(r)] for r in rows
|
||||
]
|
||||
html: str | None = None
|
||||
else:
|
||||
parsed_rows = None
|
||||
html = table_body_raw
|
||||
|
||||
num_rows, num_cols = _normalize_dimension(parsed_rows)
|
||||
|
||||
header_pos = self.table_position
|
||||
self.table_position += 1
|
||||
header_rows = (
|
||||
self.block_table_headers[header_pos]
|
||||
if header_pos < len(self.block_table_headers)
|
||||
else None
|
||||
)
|
||||
# Treat empty list / explicit None identically: no header
|
||||
# entry on the sidecar item.
|
||||
table_header = header_rows if header_rows else None
|
||||
|
||||
placeholder = self.next_key("tb")
|
||||
self.tables.append(
|
||||
IRTable(
|
||||
placeholder_key=placeholder,
|
||||
rows=parsed_rows,
|
||||
html=html,
|
||||
num_rows=num_rows,
|
||||
num_cols=num_cols,
|
||||
caption="",
|
||||
footnotes=[],
|
||||
table_header=table_header,
|
||||
body_override=table_body_raw,
|
||||
)
|
||||
)
|
||||
return f"{{{{TBL:{placeholder}}}}}"
|
||||
|
||||
def replace_equation(self, match: "re.Match[str]") -> str:
|
||||
latex = match.group(1)
|
||||
source = match.string
|
||||
start, end = match.start(), match.end()
|
||||
is_block = (start == 0 or source[start - 1] == "\n") and (
|
||||
end == len(source) or source[end] == "\n"
|
||||
)
|
||||
placeholder = self.next_key("eq")
|
||||
self.equations.append(
|
||||
IREquation(
|
||||
placeholder_key=placeholder,
|
||||
latex=latex,
|
||||
is_block=is_block,
|
||||
caption="",
|
||||
footnotes=[],
|
||||
)
|
||||
)
|
||||
token = "EQ" if is_block else "EQI"
|
||||
return f"{{{{{token}:{placeholder}}}}}"
|
||||
|
||||
def replace_drawing(self, match: "re.Match[str]") -> str:
|
||||
attrs = parse_drawing_attributes(match.group(0))
|
||||
path_val = attrs.get("path", "") or ""
|
||||
src_val = attrs.get("src", "") or ""
|
||||
fmt = attrs.get("format", "") or ""
|
||||
if not fmt and path_val:
|
||||
fmt = Path(path_val).suffix.lower().lstrip(".")
|
||||
|
||||
# Two flavours of <drawing path="…">:
|
||||
# 1. Local asset under <base>.blocks.assets/ — already
|
||||
# extracted to disk by DrawingExtractionContext;
|
||||
# register as AssetSpec(source=None) and let the
|
||||
# writer resolve the path via asset_paths.
|
||||
# 2. External/linked path (URL, or any path that does
|
||||
# not live under asset_prefix) — pass through
|
||||
# verbatim via IRDrawing.path_override; do NOT emit
|
||||
# an AssetSpec (no on-disk bytes to materialize).
|
||||
rel_inside_assets = _safe_asset_ref_from_path(path_val, self.asset_prefix)
|
||||
if rel_inside_assets is not None:
|
||||
asset_ref = rel_inside_assets
|
||||
suggested_name = Path(rel_inside_assets).name or rel_inside_assets
|
||||
if asset_ref and asset_ref not in self.seen_asset_refs:
|
||||
self.assets.append(
|
||||
AssetSpec(
|
||||
ref=asset_ref,
|
||||
suggested_name=suggested_name,
|
||||
source=None, # already extracted to disk
|
||||
)
|
||||
)
|
||||
self.seen_asset_refs.add(asset_ref)
|
||||
path_override: str | None = None
|
||||
else:
|
||||
asset_ref = ""
|
||||
# Only mark as an external/linked reference when the
|
||||
# upstream parser actually emitted a path. An empty
|
||||
# ``path=""`` should fall back to the regular asset-
|
||||
# resolution path (which will also produce ``path=""``
|
||||
# downstream) rather than masquerading as an explicit
|
||||
# builder override.
|
||||
path_override = (
|
||||
None
|
||||
if self.asset_prefix and path_val.startswith(self.asset_prefix)
|
||||
else path_val or None
|
||||
)
|
||||
|
||||
placeholder = self.next_key("im")
|
||||
self.drawings.append(
|
||||
IRDrawing(
|
||||
placeholder_key=placeholder,
|
||||
asset_ref=asset_ref,
|
||||
fmt=fmt,
|
||||
caption="",
|
||||
footnotes=[],
|
||||
src=src_val,
|
||||
path_override=path_override,
|
||||
)
|
||||
)
|
||||
return f"{{{{IMG:{placeholder}}}}}"
|
||||
|
||||
|
||||
class NativeDocxIRBuilder:
|
||||
"""Translate ``extract_docx_blocks`` output into an :class:`IRDoc`.
|
||||
|
||||
The builder is stateless — instantiate per call. ``asset_dir_name`` is
|
||||
the relative name (without trailing slash) of ``<base>.blocks.assets/``
|
||||
that the upstream parser used when emitting ``<drawing path>``
|
||||
attributes; the builder strips that prefix when building
|
||||
:attr:`AssetSpec.ref` so the writer's ref↔filename mapping has
|
||||
predictable keys.
|
||||
"""
|
||||
|
||||
def normalize(
|
||||
self,
|
||||
blocks: list[dict[str, Any]],
|
||||
*,
|
||||
document_name: str,
|
||||
asset_dir_name: str,
|
||||
parse_metadata: dict[str, Any] | None = None,
|
||||
) -> IRDoc:
|
||||
next_key = _placeholder_keyspace()
|
||||
ir_blocks: list[IRBlock] = []
|
||||
assets: list[AssetSpec] = []
|
||||
seen_asset_refs: set[str] = set()
|
||||
|
||||
asset_prefix = f"{asset_dir_name}/" if asset_dir_name else ""
|
||||
|
||||
for block in blocks:
|
||||
raw_content = block.get("content") or ""
|
||||
heading = block.get("heading") or ""
|
||||
level = int(block.get("level", 0) or 0)
|
||||
parent_headings = list(block.get("parent_headings") or [])
|
||||
# Preserve per-side nulls in [start, end].
|
||||
uuid_start = block.get("uuid") or None
|
||||
uuid_end = block.get("uuid_end") or None
|
||||
|
||||
builder = _BlockBuilder(
|
||||
next_key=next_key,
|
||||
assets=assets,
|
||||
seen_asset_refs=seen_asset_refs,
|
||||
asset_prefix=asset_prefix,
|
||||
block_table_headers=list(block.get("table_headers") or []),
|
||||
)
|
||||
|
||||
# Rewrite order matches the legacy native flow: tables, then
|
||||
# equations, then drawings — each ``re.sub`` operates on the
|
||||
# output of the previous pass.
|
||||
content_template = _TABLE_TAG_RE.sub(builder.replace_table, raw_content)
|
||||
content_template = _EQUATION_TAG_RE.sub(
|
||||
builder.replace_equation, content_template
|
||||
)
|
||||
content_template = DRAWING_TAG_PATTERN.sub(
|
||||
builder.replace_drawing, content_template
|
||||
)
|
||||
|
||||
positions = [
|
||||
IRPosition(type="paraid", range=[uuid_start, uuid_end]),
|
||||
]
|
||||
|
||||
ir_blocks.append(
|
||||
IRBlock(
|
||||
content_template=content_template,
|
||||
heading=heading,
|
||||
level=level,
|
||||
parent_headings=parent_headings,
|
||||
positions=positions,
|
||||
tables=builder.tables,
|
||||
drawings=builder.drawings,
|
||||
equations=builder.equations,
|
||||
)
|
||||
)
|
||||
|
||||
# doc_title: parse_metadata["first_heading"] when present, else file
|
||||
# stem fallback (resolved here so the writer doesn't have to know).
|
||||
first_heading = ""
|
||||
if isinstance(parse_metadata, dict):
|
||||
first_heading = str(parse_metadata.get("first_heading") or "")
|
||||
doc_title = first_heading or (Path(document_name).stem or document_name)
|
||||
|
||||
return IRDoc(
|
||||
document_name=document_name,
|
||||
document_format=Path(document_name).suffix.lower().lstrip("."),
|
||||
doc_title=doc_title,
|
||||
split_option={},
|
||||
blocks=ir_blocks,
|
||||
assets=assets,
|
||||
bbox_attributes=None,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["NativeDocxIRBuilder"]
|
||||
@@ -0,0 +1,423 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ABOUTME: Resolves automatic numbering labels from DOCX documents
|
||||
ABOUTME: Parses numbering.xml and computes rendered number strings
|
||||
"""
|
||||
|
||||
import zipfile
|
||||
from defusedxml import ElementTree as ET
|
||||
from typing import Dict
|
||||
|
||||
NSMAP = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
|
||||
|
||||
|
||||
class NumberingResolver:
|
||||
"""
|
||||
Resolves paragraph numbering to rendered label strings.
|
||||
|
||||
DOCX stores numbering definitions in numbering.xml:
|
||||
- abstractNum: Defines format templates (lvlText like "%1.%2.")
|
||||
- num: Links numId to abstractNumId
|
||||
|
||||
Each paragraph references: numId (which definition) + ilvl (which level)
|
||||
"""
|
||||
|
||||
# Number format converters
|
||||
FORMAT_CONVERTERS = {
|
||||
"decimal": lambda n: str(n),
|
||||
"lowerLetter": lambda n: chr(ord("a") + (n - 1) % 26),
|
||||
"upperLetter": lambda n: chr(ord("A") + (n - 1) % 26),
|
||||
"lowerRoman": lambda n: NumberingResolver._to_roman(n).lower(),
|
||||
"upperRoman": lambda n: NumberingResolver._to_roman(n),
|
||||
"chineseCountingThousand": lambda n: NumberingResolver._to_chinese(n),
|
||||
"ideographTraditional": lambda n: "甲乙丙丁戊己庚辛壬癸"[(n - 1) % 10],
|
||||
"bullet": lambda n: "•",
|
||||
"none": lambda n: "",
|
||||
}
|
||||
|
||||
def __init__(self, docx_path: str):
|
||||
self.abstract_nums: Dict[str, dict] = {} # abstractNumId -> level definitions
|
||||
self.num_to_abstract: Dict[str, str] = {} # numId -> abstractNumId
|
||||
self.counters: Dict[
|
||||
str, Dict[int, int]
|
||||
] = {} # numId -> {ilvl -> current_count}
|
||||
self.start_overrides: Dict[
|
||||
str, Dict[int, int]
|
||||
] = {} # numId -> {ilvl -> start_value}
|
||||
self.style_numpr: Dict[
|
||||
str, dict
|
||||
] = {} # styleId -> {numId, ilvl} from styles.xml
|
||||
self.style_based_on: Dict[str, str] = {} # styleId -> basedOn styleId
|
||||
# Smart numbering merge state (Word's rendering behavior)
|
||||
self.last_numId: str = None # Previous paragraph's numId
|
||||
self.last_abstract_id: str = None # Previous paragraph's abstractNumId
|
||||
self.last_style_id: str = None # Previous paragraph's style ID
|
||||
self._parse_numbering_xml(docx_path)
|
||||
self._parse_styles_xml(docx_path)
|
||||
|
||||
def _parse_numbering_xml(self, docx_path: str):
|
||||
"""Parse numbering.xml from DOCX archive"""
|
||||
try:
|
||||
with zipfile.ZipFile(docx_path, "r") as zf:
|
||||
if "word/numbering.xml" not in zf.namelist():
|
||||
return
|
||||
|
||||
tree = ET.parse(zf.open("word/numbering.xml"))
|
||||
root = tree.getroot()
|
||||
|
||||
# Parse abstractNum definitions
|
||||
for abstract in root.findall(".//w:abstractNum", NSMAP):
|
||||
abstract_id = abstract.get(f"{{{NSMAP['w']}}}abstractNumId")
|
||||
levels = {}
|
||||
|
||||
for lvl in abstract.findall("w:lvl", NSMAP):
|
||||
ilvl = int(lvl.get(f"{{{NSMAP['w']}}}ilvl"))
|
||||
|
||||
start_elem = lvl.find("w:start", NSMAP)
|
||||
start = (
|
||||
int(start_elem.get(f"{{{NSMAP['w']}}}val"))
|
||||
if start_elem is not None
|
||||
else 1
|
||||
)
|
||||
|
||||
num_fmt_elem = lvl.find("w:numFmt", NSMAP)
|
||||
num_fmt = (
|
||||
num_fmt_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
if num_fmt_elem is not None
|
||||
else "decimal"
|
||||
)
|
||||
|
||||
lvl_text_elem = lvl.find("w:lvlText", NSMAP)
|
||||
lvl_text = (
|
||||
lvl_text_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
if lvl_text_elem is not None
|
||||
else "%1."
|
||||
)
|
||||
|
||||
is_lgl_elem = lvl.find("w:isLgl", NSMAP)
|
||||
is_lgl = False
|
||||
if is_lgl_elem is not None:
|
||||
val = is_lgl_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
is_lgl = val is None or val not in ("0", "false")
|
||||
|
||||
levels[ilvl] = {
|
||||
"start": start,
|
||||
"numFmt": num_fmt,
|
||||
"lvlText": lvl_text,
|
||||
"isLgl": is_lgl,
|
||||
}
|
||||
|
||||
self.abstract_nums[abstract_id] = levels
|
||||
|
||||
# Parse num -> abstractNum mapping and startOverride
|
||||
for num in root.findall(".//w:num", NSMAP):
|
||||
num_id = num.get(f"{{{NSMAP['w']}}}numId")
|
||||
abstract_ref = num.find("w:abstractNumId", NSMAP)
|
||||
if abstract_ref is not None:
|
||||
self.num_to_abstract[num_id] = abstract_ref.get(
|
||||
f"{{{NSMAP['w']}}}val"
|
||||
)
|
||||
|
||||
# Parse lvlOverride/startOverride for this num
|
||||
for lvl_override in num.findall("w:lvlOverride", NSMAP):
|
||||
ilvl = int(lvl_override.get(f"{{{NSMAP['w']}}}ilvl"))
|
||||
start_override = lvl_override.find("w:startOverride", NSMAP)
|
||||
if start_override is not None:
|
||||
start_val = int(start_override.get(f"{{{NSMAP['w']}}}val"))
|
||||
if num_id not in self.start_overrides:
|
||||
self.start_overrides[num_id] = {}
|
||||
self.start_overrides[num_id][ilvl] = start_val
|
||||
except Exception:
|
||||
# Silently ignore parsing errors - document may not have numbering
|
||||
pass
|
||||
|
||||
def _parse_styles_xml(self, docx_path: str):
|
||||
"""Parse styles.xml to get style-inherited numbering definitions"""
|
||||
try:
|
||||
with zipfile.ZipFile(docx_path, "r") as zf:
|
||||
if "word/styles.xml" not in zf.namelist():
|
||||
return
|
||||
|
||||
tree = ET.parse(zf.open("word/styles.xml"))
|
||||
root = tree.getroot()
|
||||
|
||||
# Parse style definitions
|
||||
for style in root.findall(".//w:style", NSMAP):
|
||||
style_id = style.get(f"{{{NSMAP['w']}}}styleId")
|
||||
if not style_id:
|
||||
continue
|
||||
|
||||
# Check for basedOn (style inheritance)
|
||||
based_on = style.find("w:basedOn", NSMAP)
|
||||
if based_on is not None:
|
||||
parent_id = based_on.get(f"{{{NSMAP['w']}}}val")
|
||||
if parent_id:
|
||||
self.style_based_on[style_id] = parent_id
|
||||
|
||||
# Check for numPr in style's pPr
|
||||
pPr = style.find("w:pPr", NSMAP)
|
||||
if pPr is not None:
|
||||
numPr = pPr.find("w:numPr", NSMAP)
|
||||
if numPr is not None:
|
||||
num_id_elem = numPr.find("w:numId", NSMAP)
|
||||
ilvl_elem = numPr.find("w:ilvl", NSMAP)
|
||||
|
||||
if num_id_elem is not None:
|
||||
num_id = num_id_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
ilvl = (
|
||||
int(ilvl_elem.get(f"{{{NSMAP['w']}}}val"))
|
||||
if ilvl_elem is not None
|
||||
else 0
|
||||
)
|
||||
self.style_numpr[style_id] = {
|
||||
"numId": num_id,
|
||||
"ilvl": ilvl,
|
||||
}
|
||||
except Exception:
|
||||
# Silently ignore parsing errors
|
||||
pass
|
||||
|
||||
def _get_numbering_from_style(self, style_id: str, visited=None) -> dict:
|
||||
"""
|
||||
Get numbering definition from style, following inheritance chain.
|
||||
|
||||
Args:
|
||||
style_id: Style ID to look up
|
||||
visited: Set of visited style IDs (to prevent circular references)
|
||||
|
||||
Returns:
|
||||
dict with 'numId' and 'ilvl', or None
|
||||
"""
|
||||
if visited is None:
|
||||
visited = set()
|
||||
|
||||
# Prevent circular references
|
||||
if style_id in visited:
|
||||
return None
|
||||
visited.add(style_id)
|
||||
|
||||
# Check if this style has numPr
|
||||
if style_id in self.style_numpr:
|
||||
return self.style_numpr[style_id]
|
||||
|
||||
# Check parent style
|
||||
if style_id in self.style_based_on:
|
||||
parent_id = self.style_based_on[style_id]
|
||||
return self._get_numbering_from_style(parent_id, visited)
|
||||
|
||||
return None
|
||||
|
||||
def reset_tracking_state(self):
|
||||
"""
|
||||
Reset numbering tracking state.
|
||||
|
||||
Call this when encountering structural breaks that should
|
||||
interrupt numbering continuity:
|
||||
- Section breaks (sectPr)
|
||||
- Table boundaries (before and after tables)
|
||||
|
||||
This prevents incorrect numbering continuation across
|
||||
document structure boundaries.
|
||||
"""
|
||||
self.last_numId = None
|
||||
self.last_abstract_id = None
|
||||
self.last_style_id = None
|
||||
|
||||
def get_label(self, para_element) -> str:
|
||||
"""
|
||||
Get rendered numbering label for a paragraph.
|
||||
|
||||
Checks both direct numPr and style-inherited numbering. Direct numPr
|
||||
is a paragraph-local override and applies only to the current
|
||||
paragraph; subsequent paragraphs that carry only pStyle fall back to
|
||||
the style's numPr declared in styles.xml.
|
||||
|
||||
Args:
|
||||
para_element: lxml Element for <w:p>
|
||||
|
||||
Returns:
|
||||
Rendered label string (e.g., "1.1", "a)", "第一章") or empty string
|
||||
"""
|
||||
try:
|
||||
pPr = para_element.find(f"{{{NSMAP['w']}}}pPr")
|
||||
if pPr is None:
|
||||
return ""
|
||||
|
||||
num_id = None
|
||||
ilvl = 0
|
||||
style_id = None
|
||||
|
||||
# Get pStyle (if present)
|
||||
pStyle = pPr.find(f"{{{NSMAP['w']}}}pStyle")
|
||||
if pStyle is not None:
|
||||
style_id = pStyle.get(f"{{{NSMAP['w']}}}val")
|
||||
|
||||
# Check for direct numPr in paragraph
|
||||
numPr = pPr.find(f"{{{NSMAP['w']}}}numPr")
|
||||
if numPr is not None:
|
||||
num_id_elem = numPr.find(f"{{{NSMAP['w']}}}numId")
|
||||
ilvl_elem = numPr.find(f"{{{NSMAP['w']}}}ilvl")
|
||||
|
||||
if num_id_elem is not None:
|
||||
num_id = num_id_elem.get(f"{{{NSMAP['w']}}}val")
|
||||
ilvl = (
|
||||
int(ilvl_elem.get(f"{{{NSMAP['w']}}}val"))
|
||||
if ilvl_elem is not None
|
||||
else 0
|
||||
)
|
||||
|
||||
# If no direct numPr, fall back to style-inherited numbering.
|
||||
# Direct numPr is a paragraph-local override in Word; it must not
|
||||
# persist as a runtime default for the style, otherwise subsequent
|
||||
# paragraphs that only carry pStyle will keep following the local
|
||||
# override instead of the style's declared numPr.
|
||||
if num_id is None and style_id:
|
||||
style_num = self._get_numbering_from_style(style_id)
|
||||
if style_num:
|
||||
num_id = style_num["numId"]
|
||||
ilvl = style_num["ilvl"]
|
||||
|
||||
# If still no numbering found, clear state and return empty
|
||||
if num_id is None:
|
||||
# We should use list structure breaking logic to reset last_numId, last_abstract_id and last_style_id
|
||||
return ""
|
||||
|
||||
# Get abstract definition
|
||||
abstract_id = self.num_to_abstract.get(num_id)
|
||||
if abstract_id is None or abstract_id not in self.abstract_nums:
|
||||
# Clear state for invalid numbering
|
||||
self.last_numId = None
|
||||
self.last_abstract_id = None
|
||||
return ""
|
||||
|
||||
levels = self.abstract_nums[abstract_id]
|
||||
if ilvl not in levels:
|
||||
# Clear state for invalid level
|
||||
self.last_numId = None
|
||||
self.last_abstract_id = None
|
||||
return ""
|
||||
|
||||
# Smart numbering merge: (Word's rendering behavior)
|
||||
# When consecutive paragraphs have different numId but same abstractNumId,
|
||||
# Word continues the numbering sequence rather than restarting.
|
||||
# This happens regardless of whether the numId is new or style matches.
|
||||
|
||||
if (
|
||||
self.last_numId is not None
|
||||
and self.last_numId != num_id
|
||||
and self.last_abstract_id == abstract_id
|
||||
and self.last_numId in self.counters
|
||||
):
|
||||
# Merge: copy previous numId's counter to current numId
|
||||
self.counters[num_id] = self.counters[self.last_numId].copy()
|
||||
|
||||
# Initialize/update counter
|
||||
if num_id not in self.counters:
|
||||
self.counters[num_id] = {}
|
||||
|
||||
# Initialize all parent levels if not present (for deep nested numbering)
|
||||
for i in range(ilvl):
|
||||
if i not in self.counters[num_id] and i in levels:
|
||||
# Use startOverride if exists, otherwise use abstractNum's start value
|
||||
if (
|
||||
num_id in self.start_overrides
|
||||
and i in self.start_overrides[num_id]
|
||||
):
|
||||
self.counters[num_id][i] = self.start_overrides[num_id][i]
|
||||
else:
|
||||
self.counters[num_id][i] = levels[i]["start"]
|
||||
|
||||
# Reset lower levels when higher level increments
|
||||
for i in range(ilvl + 1, 10):
|
||||
if i in self.counters[num_id]:
|
||||
del self.counters[num_id][i]
|
||||
|
||||
# Initialize current level if needed
|
||||
if ilvl not in self.counters[num_id]:
|
||||
# Use startOverride if exists, otherwise use abstractNum's start value
|
||||
if (
|
||||
num_id in self.start_overrides
|
||||
and ilvl in self.start_overrides[num_id]
|
||||
):
|
||||
self.counters[num_id][ilvl] = self.start_overrides[num_id][ilvl]
|
||||
else:
|
||||
self.counters[num_id][ilvl] = levels[ilvl]["start"]
|
||||
else:
|
||||
self.counters[num_id][ilvl] += 1
|
||||
|
||||
# Format the label using lvlText template
|
||||
label = self._format_label(num_id, ilvl, levels)
|
||||
|
||||
# Update tracking state for next paragraph
|
||||
self.last_numId = num_id
|
||||
self.last_abstract_id = abstract_id
|
||||
self.last_style_id = style_id
|
||||
|
||||
return label
|
||||
except Exception:
|
||||
# Return empty on any error to avoid breaking document parsing
|
||||
return ""
|
||||
|
||||
def _format_label(self, num_id: str, ilvl: int, levels: dict) -> str:
|
||||
"""Format label string by replacing %1, %2, etc."""
|
||||
try:
|
||||
lvl_text = levels[ilvl]["lvlText"]
|
||||
result = lvl_text
|
||||
current_is_lgl = levels[ilvl].get("isLgl", False)
|
||||
|
||||
for i in range(ilvl + 1):
|
||||
if i in levels and i in self.counters.get(num_id, {}):
|
||||
num_fmt = levels[i]["numFmt"]
|
||||
if current_is_lgl and i < ilvl:
|
||||
num_fmt = "decimal"
|
||||
count = self.counters[num_id][i]
|
||||
converter = self.FORMAT_CONVERTERS.get(num_fmt, lambda n: str(n))
|
||||
formatted = converter(count)
|
||||
result = result.replace(f"%{i + 1}", formatted)
|
||||
|
||||
return result
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def _to_roman(n: int) -> str:
|
||||
"""Convert integer to Roman numeral"""
|
||||
if n <= 0 or n >= 4000:
|
||||
return str(n)
|
||||
values = [
|
||||
(1000, "M"),
|
||||
(900, "CM"),
|
||||
(500, "D"),
|
||||
(400, "CD"),
|
||||
(100, "C"),
|
||||
(90, "XC"),
|
||||
(50, "L"),
|
||||
(40, "XL"),
|
||||
(10, "X"),
|
||||
(9, "IX"),
|
||||
(5, "V"),
|
||||
(4, "IV"),
|
||||
(1, "I"),
|
||||
]
|
||||
result = ""
|
||||
for value, numeral in values:
|
||||
while n >= value:
|
||||
result += numeral
|
||||
n -= value
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _to_chinese(n: int) -> str:
|
||||
"""Convert integer to Chinese numeral"""
|
||||
digits = "零一二三四五六七八九"
|
||||
if n <= 0 or n > 99:
|
||||
return str(n)
|
||||
if n < 10:
|
||||
return digits[n]
|
||||
if n < 20:
|
||||
return "十" + (digits[n % 10] if n % 10 else "")
|
||||
if n < 100:
|
||||
tens = n // 10
|
||||
ones = n % 10
|
||||
return digits[tens] + "十" + (digits[ones] if ones else "")
|
||||
return str(n)
|
||||
@@ -0,0 +1,10 @@
|
||||
"""
|
||||
ABOUTME: OMML (Office Math Markup Language) to LaTeX conversion
|
||||
"""
|
||||
|
||||
from .ommlparser import OMMLParser
|
||||
|
||||
|
||||
def convert_omml_to_latex(omml_element) -> str:
|
||||
"""Convert an m:oMath XML element to a LaTeX string."""
|
||||
return OMMLParser().parse(omml_element)
|
||||
@@ -0,0 +1,38 @@
|
||||
"""
|
||||
Postprocessing functions for cleaning up latex equations in linear format which don't give valid LaTeX.
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
clean_exps = {
|
||||
r"\\degf": "°F",
|
||||
r"\\degc": "°C",
|
||||
r"(\\cbrt)(\w+)": r"\\sqrt[3]{\2}",
|
||||
r"(\\qdrt)(\w+)": r"\\sqrt[4]{\2}",
|
||||
r"\\sfrac": r"\\frac",
|
||||
r"(\\o[i]+nt)(\w+)": r"\1{\2}",
|
||||
r"\\bullet(\w+)": r"\\bullet \1",
|
||||
r"\\sum([a-zA-Z0-9]+)": r"\\sum{\1}",
|
||||
r"\\prod([a-zA-Z0-9]+)": r"\\prod{\1}",
|
||||
r"\\amalg([a-zA-Z0-9]+)": r"\\amalg{\1}",
|
||||
r"\\bigcup([a-zA-Z0-9]+)": r"\\bigcup{\1}",
|
||||
r"\\bigcap([a-zA-Z0-9]+)": r"\\bigcap{\1}",
|
||||
r"\\bigvee([a-zA-Z0-9]+)": r"\\bigvee{\1}",
|
||||
r"\\bigwedge([a-zA-Z0-9]+)": r"\\bigwedge{\1}",
|
||||
r"\\lfloor([a-zA-Z0-9]+)": r"\\lfloor{\1}",
|
||||
r"\\lceil([a-zA-Z0-9]+)": r"\\lceil{\1}",
|
||||
r"\\lim\\below\{(.+)\}\{(.+)\}": r"\\lim_{\1}{\2}",
|
||||
r"\\min\\below\{(.+)\}\{(.+)\}": r"\\min_{\1}{\2}",
|
||||
r"\\max\\below\{(.+)\}\{(.+)\}": r"\\max_{\1}{\2}",
|
||||
}
|
||||
|
||||
|
||||
def clean_exp(exp):
|
||||
"""
|
||||
Takes in a linear expression and converts known invalid LaTeX equations to valid LaTeX
|
||||
:param exp:str - An equation in invalid syntax
|
||||
:return :str - A valid equation
|
||||
"""
|
||||
for e in clean_exps:
|
||||
exp = re.sub(e, clean_exps[e], exp)
|
||||
return exp
|
||||
@@ -0,0 +1,511 @@
|
||||
from xml.etree.cElementTree import Element
|
||||
|
||||
from .utils import qn
|
||||
|
||||
|
||||
class OMMLParser:
|
||||
"""
|
||||
Parser class for reading OMML and converting it into LaTeX.
|
||||
"""
|
||||
|
||||
FUNCTION_MAP = {
|
||||
"sin": "\\sin",
|
||||
"cos": "\\cos",
|
||||
"tan": "\\tan",
|
||||
"cot": "\\cot",
|
||||
"sec": "\\sec",
|
||||
"csc": "\\csc",
|
||||
"sinh": "\\sinh",
|
||||
"cosh": "\\cosh",
|
||||
"tanh": "\\tanh",
|
||||
"coth": "\\coth",
|
||||
"sech": "\\operatorname{sech}",
|
||||
"csch": "\\operatorname{csch}",
|
||||
"log": "\\log",
|
||||
"ln": "\\ln",
|
||||
"min": "\\min",
|
||||
"max": "\\max",
|
||||
"lim": "\\lim",
|
||||
}
|
||||
|
||||
def _normalize_func_name(self, content: str) -> str:
|
||||
if not content:
|
||||
return content
|
||||
if content.startswith("\\"):
|
||||
return content
|
||||
key = content.strip()
|
||||
mapped = self.FUNCTION_MAP.get(key)
|
||||
return mapped if mapped else content
|
||||
|
||||
def parse(self, root: Element) -> str:
|
||||
"""
|
||||
Parses an m:oMath OMML tag into LaTeX.
|
||||
:param root: An m:oMath OMML tag
|
||||
:return: The LaTeX representation of the OMML input
|
||||
"""
|
||||
text = ""
|
||||
try:
|
||||
if root.tag == qn("m:t"):
|
||||
return self.parse_t(root)
|
||||
for child in root:
|
||||
if child.tag in self.parsers:
|
||||
text += self.parsers[child.tag](self, child)
|
||||
except AttributeError:
|
||||
# In case of missing attributes on OMML tags,
|
||||
# we return an empty string (ref:issue_14)
|
||||
return ""
|
||||
return text
|
||||
|
||||
def parse_e(self, root: Element) -> str:
|
||||
text = ""
|
||||
for child in root:
|
||||
text += self.parse(child)
|
||||
return text
|
||||
|
||||
def parse_r(self, root: Element) -> str:
|
||||
# TODO: Add support for m:rPr and m:scr to support different character styles
|
||||
# For now, we just parse the text content of m:r
|
||||
text = ""
|
||||
for child in root:
|
||||
text += self.parse(child)
|
||||
return text
|
||||
|
||||
def parse_t(self, root: Element):
|
||||
symbol_map = {
|
||||
"≜": "\\triangleq",
|
||||
"≝": "\\stackrel{\\tiny def}{=}",
|
||||
"≞": "\\stackrel{\\tiny m}{=}",
|
||||
}
|
||||
replacements = {
|
||||
"<": "\\lt ",
|
||||
">": "\\gt ",
|
||||
"≤": "\\leq ",
|
||||
"≥": "\\geq ",
|
||||
"∞": "\\infty ",
|
||||
"<": "\\lt ",
|
||||
">": "\\gt ",
|
||||
"≤": "\\leq ",
|
||||
"≥": "\\geq ",
|
||||
}
|
||||
text = root.text.split()
|
||||
if not text:
|
||||
return " "
|
||||
for i, t in enumerate(text):
|
||||
if t in symbol_map:
|
||||
text[i] = symbol_map[t]
|
||||
for key, value in replacements.items():
|
||||
for i, t in enumerate(text):
|
||||
text[i] = t.replace(key, value)
|
||||
return " ".join(text)
|
||||
|
||||
def parse_acc(self, root: Element) -> str:
|
||||
character_map = {
|
||||
768: "\\grave",
|
||||
769: "\\acute",
|
||||
770: "\\hat",
|
||||
771: "\\tilde",
|
||||
773: "\\bar",
|
||||
774: "\\breve",
|
||||
775: "\\dot",
|
||||
776: "\\ddot",
|
||||
780: "\\check",
|
||||
831: "\\overline{\\overline",
|
||||
8400: "\\overset\\leftharpoonup",
|
||||
8401: "\\overset\\rightharpoonup",
|
||||
8406: "\\overleftarrow",
|
||||
8407: "\\overrightarrow",
|
||||
8411: "\\dddot",
|
||||
8417: "\\overset\\leftrightarrow",
|
||||
}
|
||||
text = ""
|
||||
accent = 770
|
||||
for child in root:
|
||||
if child.tag == qn("m:accPr"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:chr"):
|
||||
val = child2.attrib.get(qn("m:val"))
|
||||
if val:
|
||||
try:
|
||||
accent = ord(val)
|
||||
except TypeError:
|
||||
pass
|
||||
|
||||
accent_cmd = character_map.get(accent)
|
||||
if accent_cmd is None:
|
||||
accent_cmd = character_map.get(770, "\\hat")
|
||||
text += accent_cmd + "{"
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child)
|
||||
text += "}"
|
||||
if accent == 831:
|
||||
text += "}"
|
||||
return text
|
||||
|
||||
def parse_bar(self, root: Element) -> str:
|
||||
text = "\\overline{"
|
||||
for child in root:
|
||||
if child.tag == qn("m:barPr"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:pos"):
|
||||
if child2.attrib.get(qn("m:val")) == "bot":
|
||||
text = "\\underline{"
|
||||
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child)
|
||||
text += "}"
|
||||
return text
|
||||
|
||||
def parse_border_box(self, root: Element) -> str:
|
||||
text = "\\boxed{"
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child)
|
||||
text += "}"
|
||||
return text
|
||||
|
||||
def parse_box(self, root: Element) -> str:
|
||||
text = ""
|
||||
for child in root:
|
||||
text += self.parse(child)
|
||||
return text
|
||||
|
||||
def parse_group_chr(self, root: Element) -> str:
|
||||
character_map = {
|
||||
"←": "\\leftarrow",
|
||||
"→": "\\rightarrow",
|
||||
"↔": "\\leftrightarrow",
|
||||
"⇐": "\\Leftarrow",
|
||||
"⇒": "\\Rightarrow",
|
||||
"⇔": "\\Leftrightarrow",
|
||||
}
|
||||
text = "\\underbrace{"
|
||||
bottom = False
|
||||
for child in root:
|
||||
if child.tag == qn("m:groupChrPr"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:chr"):
|
||||
char = child2.attrib.get(qn("m:val"))
|
||||
if char in character_map:
|
||||
text = character_map[char]
|
||||
for child2 in child:
|
||||
if (
|
||||
child2.tag == qn("m:pos")
|
||||
and child2.attrib.get(qn("m:val")) == "top"
|
||||
):
|
||||
# If m:pos is set to "top", the symbol is supposed to
|
||||
# be on top and the text is actually supposed to be under
|
||||
bottom = True
|
||||
|
||||
content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if text == "\\underbrace{":
|
||||
if bottom:
|
||||
text = "\\overbrace{" + content + "}"
|
||||
else:
|
||||
text += content + "}"
|
||||
else:
|
||||
if not bottom:
|
||||
text = "\\overset{" + content + "}" + "{" + text + "}"
|
||||
else:
|
||||
text = "\\underset{" + content + "}" + "{" + text + "}"
|
||||
return text
|
||||
|
||||
def parse_d(self, root: Element) -> str:
|
||||
bracket_map = {
|
||||
"(": "\\left(",
|
||||
")": "\\right)",
|
||||
"[": "\\left[",
|
||||
"]": "\\right]",
|
||||
"{": "\\left{",
|
||||
"}": "\\right}",
|
||||
"〈": "\\left\\langle",
|
||||
"〉": "\\right\\rangle",
|
||||
"⟨": "\\left\\langle",
|
||||
"⟩": "\\right\\rangle",
|
||||
"⌊": "\\left\\lfloor",
|
||||
"⌋": "\\right\\rfloor",
|
||||
"⌈": "\\left\\lceil",
|
||||
"⌉": "\\right\\rceil",
|
||||
"|": "\\left|",
|
||||
"‖": "\\left\\|",
|
||||
"⟦": "[\\![",
|
||||
"⟧": "]\\!]",
|
||||
}
|
||||
text = ""
|
||||
start_bracket = "("
|
||||
end_bracket = ")"
|
||||
seperator = "|"
|
||||
is_matrix = False
|
||||
for child in root:
|
||||
for child2 in child:
|
||||
if child.tag == qn("m:dPr"):
|
||||
if child2.tag == qn("m:begChr"):
|
||||
start_bracket = child2.attrib.get(qn("m:val"))
|
||||
if child2.tag == qn("m:endChr"):
|
||||
end_bracket = child2.attrib.get(qn("m:val"))
|
||||
if child2.tag == qn("m:sepChr"):
|
||||
seperator = child2.attrib.get(qn("m:val"))
|
||||
if child2.tag == qn("m:m"):
|
||||
is_matrix = True
|
||||
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
if text:
|
||||
text += seperator
|
||||
text += self.parse(child)
|
||||
end_bracket_replacements = {
|
||||
"|": "\\right|",
|
||||
"‖": "\\right\\|",
|
||||
"[": "\\right[",
|
||||
}
|
||||
start_bracket_replacements = {
|
||||
"]": "\\left]",
|
||||
}
|
||||
start = ""
|
||||
end = ""
|
||||
if start_bracket:
|
||||
if start_bracket in start_bracket_replacements:
|
||||
start = start_bracket_replacements[start_bracket] + " "
|
||||
elif start_bracket in bracket_map:
|
||||
start = bracket_map[start_bracket] + " "
|
||||
else:
|
||||
start = "\\left(" + " "
|
||||
if end_bracket:
|
||||
if end_bracket in end_bracket_replacements:
|
||||
end = " " + end_bracket_replacements[end_bracket]
|
||||
elif end_bracket in bracket_map:
|
||||
end = " " + bracket_map[end_bracket]
|
||||
else:
|
||||
end = " " + "\\right)"
|
||||
# If there is no end bracket and this tag contains an m:eqArr tag as a
|
||||
# child, we assume that the eqArr should be translated to a cases environment
|
||||
# instead of an eqnarray* environment.
|
||||
else:
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:eqArr"):
|
||||
text = text.replace("\\begin{eqnarray*}", "")
|
||||
text = text.replace("\\end{eqnarray*}", "")
|
||||
return "\\begin{cases} " + text + " \\end{cases}"
|
||||
if is_matrix:
|
||||
if start_bracket == "(" and end_bracket == ")":
|
||||
return text.replace("{matrix}", "{pmatrix}")
|
||||
elif start_bracket == "|" and end_bracket == "|":
|
||||
return text.replace("{matrix}", "{vmatrix}")
|
||||
elif start_bracket == "‖" and end_bracket == "‖":
|
||||
return text.replace("{matrix}", "{Vmatrix}")
|
||||
else:
|
||||
return text.replace("{matrix}", "{bmatrix}")
|
||||
return start + text + end
|
||||
|
||||
def parse_eq_arr(self, root: Element) -> str:
|
||||
text = "\\begin{eqnarray*}"
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child) + " \\\\"
|
||||
text += "\\end{eqnarray*}"
|
||||
return text
|
||||
|
||||
def parse_f(self, root: Element) -> str:
|
||||
text = "\\frac{"
|
||||
num = ""
|
||||
den = ""
|
||||
is_binom = False
|
||||
for child in root:
|
||||
if child.tag == qn("m:fPr"):
|
||||
for child2 in child:
|
||||
if (
|
||||
child2.tag == qn("m:type")
|
||||
and child2.attrib.get(qn("m:val")) == "noBar"
|
||||
):
|
||||
is_binom = True
|
||||
if child.tag == qn("m:num"):
|
||||
num = self.parse(child)
|
||||
if child.tag == qn("m:den"):
|
||||
den = self.parse(child)
|
||||
if is_binom:
|
||||
text = "\\genfrac{}{}{0pt}{}{"
|
||||
text += num + "}{" + den + "}"
|
||||
return text
|
||||
|
||||
def parse_m(self, root: Element) -> str:
|
||||
text = "\\begin{matrix} "
|
||||
text += self.parse(root)[:-3] # Remove the last ' \\'
|
||||
text += "\\end{matrix}"
|
||||
return text
|
||||
|
||||
def parse_mr(self, root: Element) -> str:
|
||||
text = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child) + " & "
|
||||
return text[:-2] + "\\\\ " # Remove the last ' & '
|
||||
|
||||
def parse_func(self, root: Element) -> str:
|
||||
subscript = ""
|
||||
superscript = ""
|
||||
text = ""
|
||||
func_name = "sin"
|
||||
for child in root:
|
||||
if child.tag == qn("m:fName"):
|
||||
for child2 in child:
|
||||
if child2.tag in [qn("m:sSup"), qn("m:sSub"), qn("m:r")]:
|
||||
for child3 in child2:
|
||||
if child3.tag == qn("m:sub"):
|
||||
subscript = self.parse(child3)
|
||||
if child3.tag == qn("m:sup"):
|
||||
superscript = self.parse(child3)
|
||||
if child3.tag == qn("m:t") or child3.tag == qn("m:e"):
|
||||
func_name = self.parse(child3)
|
||||
elif child2.tag == qn("m:limLow"):
|
||||
for child3 in child2:
|
||||
if child3.tag == qn("m:lim"):
|
||||
for child4 in child3:
|
||||
subscript += self.parse(child4)
|
||||
if child3.tag == qn("m:e"):
|
||||
func_name = self.parse(child3)
|
||||
|
||||
if child.tag == qn("m:e"):
|
||||
text += self.parse(child)
|
||||
if func_name in ["lim", "max", "min"]:
|
||||
return f"\\{func_name}\\limits_{{{subscript}}}^{{{superscript}}}{{{text}}}"
|
||||
if func_name not in self.FUNCTION_MAP:
|
||||
return f"{{{func_name}}}^{{{superscript}}}_{{{subscript}}}{{{text}}}"
|
||||
return (
|
||||
self.FUNCTION_MAP[func_name]
|
||||
+ f"_{{{subscript}}}^{{{superscript}}}{{{text}}}"
|
||||
)
|
||||
|
||||
def parse_s_sup(self, root: Element) -> str:
|
||||
content = ""
|
||||
exp_content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if child.tag == qn("m:sup"):
|
||||
exp_content = self.parse(child)
|
||||
content = self._normalize_func_name(content)
|
||||
return f"{{{content}}}^{{{exp_content}}}"
|
||||
|
||||
def parse_s_sub(self, root: Element) -> str:
|
||||
content = ""
|
||||
sub_content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if child.tag == qn("m:sub"):
|
||||
sub_content = self.parse(child)
|
||||
content = self._normalize_func_name(content)
|
||||
return f"{{{content}}}_{{{sub_content}}}"
|
||||
|
||||
def parse_s_sub_sup(self, root: Element) -> str:
|
||||
content = ""
|
||||
sub_content = ""
|
||||
exp_content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if child.tag == qn("m:sub"):
|
||||
sub_content = self.parse(child)
|
||||
if child.tag == qn("m:sup"):
|
||||
exp_content = self.parse(child)
|
||||
content = self._normalize_func_name(content)
|
||||
return f"{{{content}}}_{{{sub_content}}}^{{{exp_content}}}"
|
||||
|
||||
def parse_s_pre(self, root: Element) -> str:
|
||||
content = ""
|
||||
sub_content = ""
|
||||
exp_content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if child.tag == qn("m:sub"):
|
||||
sub_content = self.parse(child)
|
||||
if child.tag == qn("m:sup"):
|
||||
exp_content = self.parse(child)
|
||||
return "{}^{" + exp_content + "}_{" + sub_content + "}{" + content + "}"
|
||||
|
||||
def parse_rad(self, root: Element) -> str:
|
||||
content = ""
|
||||
order = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:deg"):
|
||||
order = self.parse(child)
|
||||
if child.tag == qn("m:e"):
|
||||
content += self.parse(child)
|
||||
if order:
|
||||
return f"\\sqrt[{order}]{{{content}}}"
|
||||
return f"\\sqrt{{{content}}}"
|
||||
|
||||
def parse_nary(self, root: Element) -> str:
|
||||
character_map = {
|
||||
8719: "\\prod",
|
||||
8720: "\\coprod",
|
||||
8721: "\\sum",
|
||||
8747: "\\int",
|
||||
8748: "\\iint",
|
||||
8749: "\\iiint",
|
||||
8750: "\\oint",
|
||||
8751: "\\oiint",
|
||||
8752: "\\oiiint",
|
||||
8896: "\\bigwedge",
|
||||
8897: "\\bigvee",
|
||||
8898: "\\bigcap",
|
||||
8899: "\\bigcup",
|
||||
}
|
||||
char = 8747
|
||||
for child in root:
|
||||
if child.tag == qn("m:naryPr"):
|
||||
for child2 in child:
|
||||
if child2.tag == qn("m:chr"):
|
||||
val = child2.attrib.get(qn("m:val"))
|
||||
if val:
|
||||
try:
|
||||
char = ord(val)
|
||||
except TypeError:
|
||||
pass
|
||||
text = character_map.get(char, character_map[8721])
|
||||
sub = ""
|
||||
sup = ""
|
||||
content = ""
|
||||
for child in root:
|
||||
if child.tag == qn("m:sub"):
|
||||
sub = self.parse(child)
|
||||
if child.tag == qn("m:sup"):
|
||||
sup = self.parse(child)
|
||||
if child.tag == qn("m:e"):
|
||||
content = self.parse(child)
|
||||
if sub:
|
||||
text += f"_{{{sub}}}"
|
||||
if sup:
|
||||
text += f"^{{{sup}}}"
|
||||
text += "{" + content + "}"
|
||||
return text
|
||||
|
||||
parsers = {
|
||||
qn("m:r"): parse_r,
|
||||
qn("m:acc"): parse_acc,
|
||||
qn("m:borderBox"): parse_border_box,
|
||||
qn("m:bar"): parse_bar,
|
||||
qn("m:box"): parse_box,
|
||||
qn("m:d"): parse_d,
|
||||
qn("m:e"): parse_e,
|
||||
qn("m:groupChr"): parse_group_chr,
|
||||
qn("m:f"): parse_f,
|
||||
qn("m:sSup"): parse_s_sup,
|
||||
qn("m:sSub"): parse_s_sub,
|
||||
qn("m:sSubSup"): parse_s_sub_sup,
|
||||
qn("m:sPre"): parse_s_pre,
|
||||
qn("m:t"): parse_t,
|
||||
qn("m:rad"): parse_rad,
|
||||
qn("m:nary"): parse_nary,
|
||||
qn("m:eqArr"): parse_eq_arr,
|
||||
qn("m:func"): parse_func,
|
||||
qn("m:m"): parse_m,
|
||||
qn("m:mr"): parse_mr,
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Utility functions to extract text from the supported mathematical equations from xml tags and
|
||||
convert them into LaTeX
|
||||
"""
|
||||
|
||||
from .cleaners import clean_exp
|
||||
|
||||
ns_map = {
|
||||
"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
|
||||
"m": "http://schemas.openxmlformats.org/officeDocument/2006/math",
|
||||
}
|
||||
|
||||
|
||||
def linear_expression(tag):
|
||||
"""
|
||||
Just returns the text contained in the given tag while setting docxlatex_skip_iteration flags
|
||||
for all its children.
|
||||
:param tag:defusedxml.Element - An xml element which contains a math equation in linear form
|
||||
:return text:str - The equation in valid LaTeX syntax
|
||||
"""
|
||||
text = ""
|
||||
for child in tag.iter():
|
||||
child.set("docxlatex_skip_iteration", True)
|
||||
text += child.text if child.text is not None else ""
|
||||
text = clean_exp(text)
|
||||
return text
|
||||
|
||||
|
||||
def qn(tag):
|
||||
"""
|
||||
A utility function to turn a namespace
|
||||
prefixed tag name into a Clark-notation qualified tag name for lxml. For
|
||||
example, qn('m:oMath') returns '{http://schemas.openxmlformats.org/officeDocument/2006/math}oMath'
|
||||
|
||||
:param tag:str - A namespace-prefixed tag name
|
||||
:return qn:str - A Clark-notation qualified name tag for lxml.
|
||||
"""
|
||||
prefix, tag_root = tag.split(":")
|
||||
uri = ns_map[prefix]
|
||||
return "{{{}}}{}".format(uri, tag_root)
|
||||
Executable
+932
@@ -0,0 +1,932 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ABOUTME: Parses DOCX documents into text blocks using python-docx
|
||||
ABOUTME: Extracts automatic numbering, splits by headings, converts tables to JSON
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
|
||||
try:
|
||||
from docx import Document
|
||||
from docx.opc.exceptions import PackageNotFoundError
|
||||
except ImportError as exc:
|
||||
# Raise instead of sys.exit: this module is imported in-process by the
|
||||
# gunicorn/uvicorn worker, where a SystemExit would tear down the whole
|
||||
# worker rather than surfacing a normal, catchable error.
|
||||
raise ImportError(
|
||||
"python-docx not installed. Run: pip install python-docx"
|
||||
) from exc
|
||||
|
||||
from lightrag.parser._markdown import (
|
||||
render_heading_line,
|
||||
strip_heading_markdown_prefix,
|
||||
)
|
||||
from .numbering_resolver import NumberingResolver
|
||||
from .table_extractor import TableExtractor
|
||||
from .drawing_image_extractor import (
|
||||
DrawingExtractionContext,
|
||||
extract_drawing_placeholder_from_element,
|
||||
extract_vml_image_placeholder_from_element,
|
||||
)
|
||||
|
||||
|
||||
# Constants for content validation (character-based for UI/display)
|
||||
MAX_HEADING_LENGTH = 200 # Maximum heading length in characters (UI constraint)
|
||||
|
||||
# OOXML tracked-change/comment tags whose subtree must be dropped so we only
|
||||
# surface the *final* revised text. w:ins / w:moveTo are kept via default
|
||||
# recursion so inserted/moved-in content survives.
|
||||
_SKIP_REVISION_TAGS = frozenset({"del", "moveFrom"})
|
||||
_SKIP_COMMENT_TAGS = frozenset(
|
||||
{"commentRangeStart", "commentRangeEnd", "commentReference", "annotationRef"}
|
||||
)
|
||||
_SKIP_PARAGRAPH_TAGS = _SKIP_REVISION_TAGS | _SKIP_COMMENT_TAGS
|
||||
|
||||
|
||||
class DocxContentError(ValueError):
|
||||
"""DOCX content violates a parsing constraint (heading/table/anchor limits).
|
||||
|
||||
Raised instead of calling ``sys.exit`` so the pipeline's per-document
|
||||
``except Exception`` handler marks just that document FAILED while the
|
||||
gunicorn/uvicorn worker process keeps running. Subclasses ``ValueError``
|
||||
(i.e. an ``Exception``, not ``BaseException``) so the existing pipeline
|
||||
handlers catch it.
|
||||
"""
|
||||
|
||||
|
||||
def format_error(title: str, details: str, solution: str) -> str:
|
||||
"""
|
||||
Build a friendly, formatted error message (title / details / SOLUTION).
|
||||
|
||||
Args:
|
||||
title: Error title
|
||||
details: Detailed error information
|
||||
solution: Suggested solution steps
|
||||
|
||||
Returns:
|
||||
str: The formatted multi-line message.
|
||||
"""
|
||||
return (
|
||||
"\n"
|
||||
+ "=" * 68
|
||||
+ f"\nERROR: {title}\n"
|
||||
+ "=" * 68
|
||||
+ f"\n\n{details}"
|
||||
+ "\n\nSOLUTION:\n"
|
||||
+ solution
|
||||
+ "\n\n"
|
||||
+ "=" * 68
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
|
||||
def print_error(title: str, details: str, solution: str):
|
||||
"""Print a friendly, formatted error message to stderr."""
|
||||
print(format_error(title, details, solution), file=sys.stderr)
|
||||
|
||||
|
||||
def _diagnose_invalid_docx(file_path: str) -> tuple[str, str]:
|
||||
"""Diagnose why a ``.docx`` file is not a valid OOXML/ZIP package.
|
||||
|
||||
python-docx raises ``PackageNotFoundError("Package not found at '...'")``
|
||||
both when the path is missing AND when the file exists but is not a valid
|
||||
zip. By the time this runs the file has already been confirmed to exist
|
||||
(the native worker validates ``p.exists()`` first), so the real cause is a
|
||||
corrupt file or a non-DOCX payload wearing a ``.docx`` extension. Sniff the
|
||||
magic bytes to name the actual format so the error message reflects the
|
||||
real problem instead of an empty "not found".
|
||||
|
||||
Returns a ``(details, solution)`` tuple for :func:`format_error`. Reads only
|
||||
the file header and never raises — any IO failure degrades to a generic
|
||||
"cannot read" diagnosis.
|
||||
"""
|
||||
import zipfile
|
||||
|
||||
convert_solution = (
|
||||
" 1. Open the file in Microsoft Word or WPS\n"
|
||||
' 2. Use "Save As" and choose "Word Document (*.docx)"\n'
|
||||
" 3. Re-upload the converted .docx to LightRAG"
|
||||
)
|
||||
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
head = f.read(8)
|
||||
except OSError as exc:
|
||||
return (
|
||||
f"The file at '{file_path}' could not be read: {exc}",
|
||||
" 1. Verify the file exists and is readable\n"
|
||||
" 2. Re-upload it to LightRAG",
|
||||
)
|
||||
|
||||
if not head:
|
||||
return (
|
||||
f"The file at '{file_path}' is empty (0 bytes). The upload was "
|
||||
"likely truncated or the source file is corrupt.",
|
||||
" 1. Check the original document opens correctly\n"
|
||||
" 2. Re-upload a complete copy to LightRAG",
|
||||
)
|
||||
|
||||
if head.startswith(b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"):
|
||||
# OLE2 Compound File — the legacy binary Word 97-2003 .doc format.
|
||||
return (
|
||||
f"The file at '{file_path}' is a legacy Word 97-2003 (.doc) "
|
||||
"document saved with a .docx extension. The .doc binary format is "
|
||||
"not a ZIP/OOXML package and cannot be parsed by the native engine.",
|
||||
convert_solution,
|
||||
)
|
||||
|
||||
if head.startswith(b"{\\rtf"):
|
||||
return (
|
||||
f"The file at '{file_path}' is an RTF document saved with a .docx "
|
||||
"extension. RTF is not a ZIP/OOXML package.",
|
||||
convert_solution,
|
||||
)
|
||||
|
||||
if head.startswith(b"%PDF"):
|
||||
return (
|
||||
f"The file at '{file_path}' is a PDF saved with a .docx extension. "
|
||||
"It is not a ZIP/OOXML package.",
|
||||
" 1. Convert the PDF to .docx, or upload it through a PDF-capable "
|
||||
"parser engine (e.g. mineru/docling)\n"
|
||||
" 2. Re-upload to LightRAG",
|
||||
)
|
||||
|
||||
stripped = head.lstrip()
|
||||
if stripped.startswith(b"<"):
|
||||
# <?xml ...>, <html ...>, or Word 2003 "<w:wordDocument>" flat XML.
|
||||
return (
|
||||
f"The file at '{file_path}' is an HTML or XML document saved with a "
|
||||
".docx extension, not a ZIP/OOXML package.",
|
||||
convert_solution,
|
||||
)
|
||||
|
||||
if head.startswith(b"PK\x03\x04") and not zipfile.is_zipfile(file_path):
|
||||
# Has the ZIP local-file-header magic but the archive is unreadable.
|
||||
return (
|
||||
f"The file at '{file_path}' starts like a ZIP archive but is "
|
||||
"truncated or corrupt, so it cannot be opened as a DOCX package.",
|
||||
" 1. Check the original document opens correctly\n"
|
||||
" 2. Re-upload a complete, uncorrupted copy to LightRAG",
|
||||
)
|
||||
|
||||
return (
|
||||
f"The file at '{file_path}' is not a valid DOCX (ZIP/OOXML) package. "
|
||||
"It is either corrupt or a different file format saved with a .docx "
|
||||
"extension.",
|
||||
convert_solution,
|
||||
)
|
||||
|
||||
|
||||
def truncate_heading(heading_text: str, para_id: str = None) -> str:
|
||||
"""
|
||||
Truncate heading if it exceeds MAX_HEADING_LENGTH.
|
||||
|
||||
Args:
|
||||
heading_text: The heading text to check
|
||||
para_id: Optional paragraph ID for warning message
|
||||
|
||||
Returns:
|
||||
str: Original heading if within limit, truncated heading with "..." if too long
|
||||
"""
|
||||
if len(heading_text) > MAX_HEADING_LENGTH:
|
||||
truncated = heading_text[: MAX_HEADING_LENGTH - 3] + "..."
|
||||
location = f" (para_id: {para_id})" if para_id else ""
|
||||
print(
|
||||
f"Warning: Heading truncated (length {len(heading_text)} > max {MAX_HEADING_LENGTH}){location}: "
|
||||
f'"{truncated}"',
|
||||
file=sys.stderr,
|
||||
)
|
||||
return truncated
|
||||
return heading_text
|
||||
|
||||
|
||||
def validate_heading_length(heading_text: str, para_id: str):
|
||||
"""
|
||||
Validate that heading length does not exceed MAX_HEADING_LENGTH.
|
||||
|
||||
Args:
|
||||
heading_text: The heading text to validate
|
||||
para_id: The paragraph ID for error reporting
|
||||
|
||||
Raises:
|
||||
DocxContentError: if heading exceeds maximum length
|
||||
"""
|
||||
if len(heading_text) > MAX_HEADING_LENGTH:
|
||||
preview = (
|
||||
heading_text[:100] + "..." if len(heading_text) > 100 else heading_text
|
||||
)
|
||||
raise DocxContentError(
|
||||
format_error(
|
||||
f"Heading too long ({len(heading_text)} characters, max {MAX_HEADING_LENGTH})",
|
||||
f"The following heading exceeds the maximum allowed length:\n\n{preview}\n\n"
|
||||
f"Location(para_id): {para_id}\n"
|
||||
f"Actual length: {len(heading_text)} characters",
|
||||
" 1. Open the document in Microsoft Word\n"
|
||||
f" 2. Shorten this heading to {MAX_HEADING_LENGTH} characters or less\n"
|
||||
" 3. Re-upload it to LightRAG",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def find_first_valid_para_id(para_ids: list) -> str | None:
|
||||
"""
|
||||
Find the first valid paraId in a 2D array of paraIds.
|
||||
|
||||
Args:
|
||||
para_ids: 2D list of paraIds from table cells
|
||||
|
||||
Returns:
|
||||
First non-None paraId found, or None when every cell lacks a paraId.
|
||||
Callers must tolerate ``None`` and treat it as a tracking gap rather
|
||||
than a fatal error (legacy / non-Word docx authors omit ``w14:paraId``
|
||||
attributes and we want to keep parsing).
|
||||
"""
|
||||
for row in para_ids:
|
||||
for para_id in row:
|
||||
if para_id:
|
||||
return para_id
|
||||
return None
|
||||
|
||||
|
||||
def find_last_valid_para_id(para_ids: list) -> str | None:
|
||||
"""
|
||||
Find the last valid paraId in a 2D array of paraIds.
|
||||
|
||||
Returns the last non-None paraId, falling back to the first valid one
|
||||
when reverse-iteration does not yield anything (single-paraId tables),
|
||||
and finally ``None`` when every cell lacks a paraId.
|
||||
"""
|
||||
for row in reversed(para_ids):
|
||||
for para_id in reversed(row):
|
||||
if para_id:
|
||||
return para_id
|
||||
|
||||
return find_first_valid_para_id(para_ids)
|
||||
|
||||
|
||||
def _table_has_any_paraid(para_ids: list) -> bool:
|
||||
"""True when at least one cell in the 2D paraId grid carries an id."""
|
||||
return find_first_valid_para_id(para_ids) is not None
|
||||
|
||||
|
||||
def extract_para_id(para_element) -> str:
|
||||
"""
|
||||
Extract w14:paraId attribute from paragraph element.
|
||||
|
||||
Args:
|
||||
para_element: lxml paragraph element
|
||||
|
||||
Returns:
|
||||
8-character hex paraId, or ``None`` when the paragraph carries no
|
||||
``w14:paraId`` attribute (legacy / non-Word docx authors). Callers
|
||||
propagate the ``None`` upward — the LightRAG adapter counts these
|
||||
and surfaces a single warning per document.
|
||||
"""
|
||||
return para_element.get(
|
||||
"{http://schemas.microsoft.com/office/word/2010/wordml}paraId"
|
||||
)
|
||||
|
||||
|
||||
def parse_styles_outline_levels(docx_path: str) -> dict:
|
||||
"""
|
||||
Parse styles.xml to extract outlineLvl definitions for each style,
|
||||
following style inheritance chain (basedOn).
|
||||
|
||||
Args:
|
||||
docx_path: Path to DOCX file
|
||||
|
||||
Returns:
|
||||
dict: styleId -> outlineLvl (0-8 for headings, 9 for body text)
|
||||
"""
|
||||
import zipfile
|
||||
|
||||
try:
|
||||
from defusedxml import ElementTree as ET
|
||||
except ImportError:
|
||||
from xml.etree import ElementTree as ET
|
||||
|
||||
styles_outline = {} # styleId -> outlineLvl (directly defined)
|
||||
style_based_on = {} # styleId -> parent styleId
|
||||
|
||||
try:
|
||||
with zipfile.ZipFile(docx_path, "r") as zf:
|
||||
if "word/styles.xml" not in zf.namelist():
|
||||
return styles_outline
|
||||
|
||||
tree = ET.parse(zf.open("word/styles.xml"))
|
||||
root = tree.getroot()
|
||||
|
||||
ns = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
|
||||
|
||||
# First pass: collect outlineLvl and basedOn for all styles
|
||||
for style in root.findall(f".//{{{ns}}}style"):
|
||||
style_id = style.get(f"{{{ns}}}styleId")
|
||||
if not style_id:
|
||||
continue
|
||||
|
||||
# Check for basedOn (style inheritance)
|
||||
based_on = style.find(f"{{{ns}}}basedOn")
|
||||
if based_on is not None:
|
||||
parent_id = based_on.get(f"{{{ns}}}val")
|
||||
if parent_id:
|
||||
style_based_on[style_id] = parent_id
|
||||
|
||||
# Check for outlineLvl in style's pPr
|
||||
pPr = style.find(f"{{{ns}}}pPr")
|
||||
if pPr is not None:
|
||||
outline_lvl_elem = pPr.find(f"{{{ns}}}outlineLvl")
|
||||
if outline_lvl_elem is not None:
|
||||
level = int(outline_lvl_elem.get(f"{{{ns}}}val"))
|
||||
styles_outline[style_id] = level
|
||||
|
||||
# Second pass: resolve inheritance chain for styles without direct outlineLvl
|
||||
def get_outline_level(style_id: str, visited: set = None) -> int:
|
||||
if visited is None:
|
||||
visited = set()
|
||||
if style_id in visited:
|
||||
return None # Prevent circular references
|
||||
visited.add(style_id)
|
||||
|
||||
# If this style directly defines outlineLvl, return it
|
||||
if style_id in styles_outline:
|
||||
return styles_outline[style_id]
|
||||
|
||||
# Otherwise check parent style
|
||||
if style_id in style_based_on:
|
||||
parent_id = style_based_on[style_id]
|
||||
return get_outline_level(parent_id, visited)
|
||||
|
||||
return None
|
||||
|
||||
# Fill in missing outlineLvl from inheritance chain
|
||||
all_style_ids = set(styles_outline.keys()) | set(style_based_on.keys())
|
||||
for style_id in all_style_ids:
|
||||
if style_id not in styles_outline:
|
||||
level = get_outline_level(style_id)
|
||||
if level is not None:
|
||||
styles_outline[style_id] = level
|
||||
except Exception:
|
||||
# Silently ignore parsing errors
|
||||
pass
|
||||
|
||||
return styles_outline
|
||||
|
||||
|
||||
def get_heading_level(para_element, styles_outline_map: dict) -> int:
|
||||
"""
|
||||
Get heading level from paragraph, checking both direct format and style.
|
||||
|
||||
Priority: paragraph outlineLvl > style outlineLvl
|
||||
|
||||
Args:
|
||||
para_element: lxml paragraph element
|
||||
styles_outline_map: dict of styleId -> outlineLvl from styles.xml
|
||||
|
||||
Returns:
|
||||
int: 0-8 for heading levels (0=level 1, 1=level 2, etc.), None for non-heading
|
||||
"""
|
||||
# 1. Check paragraph direct format
|
||||
pPr = para_element.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}pPr"
|
||||
)
|
||||
if pPr is not None:
|
||||
outline_elem = pPr.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}outlineLvl"
|
||||
)
|
||||
if outline_elem is not None:
|
||||
level = int(
|
||||
outline_elem.get(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val"
|
||||
)
|
||||
)
|
||||
# Only 0-8 are true heading levels (9 is body text)
|
||||
if level < 9:
|
||||
return level
|
||||
else:
|
||||
return None # Level 9 is body text
|
||||
|
||||
# 2. Check style definition's outlineLvl
|
||||
if pPr is not None:
|
||||
pStyle_elem = pPr.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}pStyle"
|
||||
)
|
||||
if pStyle_elem is not None:
|
||||
style_id = pStyle_elem.get(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val"
|
||||
)
|
||||
if style_id and style_id in styles_outline_map:
|
||||
level = styles_outline_map[style_id]
|
||||
if level < 9:
|
||||
return level
|
||||
else:
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_text_from_run(
|
||||
run,
|
||||
ns: dict,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
) -> str:
|
||||
"""
|
||||
Extract text from a run element, preserving superscript/subscript with markup.
|
||||
|
||||
Converts Word formatting to HTML-like tags:
|
||||
- Superscript: <sup>text</sup>
|
||||
- Subscript: <sub>text</sub>
|
||||
- Normal text: unchanged
|
||||
|
||||
Args:
|
||||
run: lxml run element (w:r)
|
||||
ns: XML namespace dictionary
|
||||
|
||||
Returns:
|
||||
Text string with <sup>/<sub> markup for formatted portions
|
||||
"""
|
||||
text = ""
|
||||
|
||||
# Check for vertAlign in rPr (superscript/subscript)
|
||||
vert_align = None
|
||||
rPr = run.find("w:rPr", ns)
|
||||
if rPr is not None:
|
||||
vert_elem = rPr.find("w:vertAlign", ns)
|
||||
if vert_elem is not None:
|
||||
vert_align = vert_elem.get(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val"
|
||||
)
|
||||
|
||||
# Extract text content from run children
|
||||
for child in run:
|
||||
tag = child.tag.split("}")[-1] # Remove namespace
|
||||
if tag == "t" and child.text:
|
||||
text += child.text
|
||||
elif tag == "tab":
|
||||
text += "\t"
|
||||
elif tag == "br":
|
||||
# Handle line breaks - textWrapping or no type = soft line break
|
||||
br_type = child.get(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}type"
|
||||
)
|
||||
if br_type in (None, "textWrapping"):
|
||||
text += "\n"
|
||||
# Skip page and column breaks (layout elements)
|
||||
elif tag == "drawing":
|
||||
text += extract_drawing_placeholder_from_element(
|
||||
child,
|
||||
context=drawing_context,
|
||||
include_extended_attrs=True,
|
||||
)
|
||||
elif tag in ("pict", "object"):
|
||||
text += extract_vml_image_placeholder_from_element(
|
||||
child,
|
||||
context=drawing_context,
|
||||
include_extended_attrs=True,
|
||||
)
|
||||
|
||||
# Apply superscript/subscript markup if needed
|
||||
if text and vert_align == "superscript":
|
||||
return f"<sup>{text}</sup>"
|
||||
elif text and vert_align == "subscript":
|
||||
return f"<sub>{text}</sub>"
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def extract_paragraph_content(
|
||||
element,
|
||||
ns,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
) -> str:
|
||||
"""
|
||||
Extract text and equations from a paragraph element in document order.
|
||||
|
||||
Handles w:r (text runs), m:oMath (inline equations), and m:oMathPara
|
||||
(block equations). Recurses into container elements (e.g., w:hyperlink,
|
||||
w:ins, w:sdt, w:fldSimple, w:smartTag) to avoid dropping content.
|
||||
|
||||
Args:
|
||||
element: lxml paragraph element (w:p)
|
||||
ns: XML namespace dictionary
|
||||
|
||||
Returns:
|
||||
Text string with equations wrapped in <equation> tags
|
||||
"""
|
||||
parts = []
|
||||
|
||||
def append_from(node) -> None:
|
||||
tag = node.tag.split("}")[-1]
|
||||
# Drop tracked-change deletions (w:del/w:moveFrom) and comment markers
|
||||
# (w:commentRangeStart/End, w:commentReference, w:annotationRef) so the
|
||||
# output only contains the final revised text without annotation glyphs.
|
||||
if tag in _SKIP_PARAGRAPH_TAGS:
|
||||
return
|
||||
if tag == "r":
|
||||
parts.append(
|
||||
extract_text_from_run(node, ns, drawing_context=drawing_context)
|
||||
)
|
||||
return
|
||||
if tag == "oMath":
|
||||
from .omml import convert_omml_to_latex
|
||||
|
||||
latex = convert_omml_to_latex(node)
|
||||
if latex:
|
||||
parts.append(f"<equation>{latex}</equation>")
|
||||
return
|
||||
if tag == "oMathPara":
|
||||
from .omml import convert_omml_to_latex
|
||||
|
||||
for omath in node:
|
||||
if omath.tag.split("}")[-1] == "oMath":
|
||||
latex = convert_omml_to_latex(omath)
|
||||
if latex:
|
||||
parts.append(f"<equation>{latex}</equation>")
|
||||
return
|
||||
for child in node:
|
||||
append_from(child)
|
||||
|
||||
for child in element:
|
||||
append_from(child)
|
||||
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def _is_table_empty(rows: list) -> bool:
|
||||
"""Return True iff every cell in ``rows`` is whitespace-only."""
|
||||
return all(not (cell or "").strip() for row in rows for cell in row)
|
||||
|
||||
|
||||
def _collect_table_headers(paragraphs: list) -> list:
|
||||
"""Collect per-table cross-page header rows from ``is_table`` paragraphs.
|
||||
|
||||
The returned list is aligned 1:1 with the order of ``<table>`` placeholder
|
||||
tags emitted into the block's content; entries are either the list of
|
||||
header rows captured from ``w:tblHeader`` or ``None`` when the table has
|
||||
no cross-page repeating header.
|
||||
"""
|
||||
return [p.get("_table_header") for p in paragraphs if p.get("is_table")]
|
||||
|
||||
|
||||
def _build_unsplit_block(
|
||||
heading: str, paragraphs: list, parent_headings: list, level: int
|
||||
) -> dict:
|
||||
"""Build a single block from paragraphs without size-based splitting."""
|
||||
last_para = paragraphs[-1]
|
||||
block = {
|
||||
"uuid": paragraphs[0]["para_id"],
|
||||
"uuid_end": last_para.get("para_id_end") or last_para.get("para_id"),
|
||||
"heading": heading,
|
||||
"content": "\n".join(p["text"] for p in paragraphs),
|
||||
"type": "text",
|
||||
"parent_headings": parent_headings,
|
||||
"level": level,
|
||||
}
|
||||
table_headers = _collect_table_headers(paragraphs)
|
||||
if table_headers:
|
||||
block["table_headers"] = table_headers
|
||||
return block
|
||||
|
||||
|
||||
def _flush_current_block(
|
||||
blocks: list,
|
||||
heading: str,
|
||||
paragraphs: list,
|
||||
parent_headings: list,
|
||||
level: int,
|
||||
) -> None:
|
||||
"""Flush accumulated paragraphs into a single heading-scoped block.
|
||||
|
||||
The native parser performs only heading-driven structural splitting; block
|
||||
sizing (long-block anchor splitting, table row splitting, small-block
|
||||
merging) is the downstream paragraph-semantic chunker's responsibility.
|
||||
"""
|
||||
if not paragraphs:
|
||||
return
|
||||
|
||||
blocks.append(_build_unsplit_block(heading, paragraphs, parent_headings, level))
|
||||
|
||||
|
||||
def extract_docx_blocks(
|
||||
file_path: str,
|
||||
drawing_context: DrawingExtractionContext = None,
|
||||
parse_warnings: dict | None = None,
|
||||
parse_metadata: dict | None = None,
|
||||
) -> list:
|
||||
"""
|
||||
Extract heading-scoped text blocks from a DOCX file.
|
||||
|
||||
Uses python-docx with a custom numbering resolver to:
|
||||
1. Capture automatic numbering (list labels)
|
||||
2. Split the document into one block per heading (structural splitting)
|
||||
3. Convert tables to JSON (2D array) and emit them as <table> placeholders
|
||||
4. Preserve superscript/subscript formatting with <sup>/<sub> markup
|
||||
|
||||
Block sizing — long-block anchor splitting, table row splitting, and
|
||||
small-block merging — is intentionally NOT done here; it is the downstream
|
||||
paragraph-semantic chunker's responsibility. Blocks emitted here may
|
||||
therefore be arbitrarily large.
|
||||
|
||||
Args:
|
||||
file_path: Path to the DOCX file
|
||||
parse_warnings: Optional out-dict that this function mutates with
|
||||
non-fatal warnings observed during parsing. Currently used for
|
||||
``missing_paraid_count`` — incremented once per body-level
|
||||
paragraph (heading or text) that lacks a ``w14:paraId`` and once
|
||||
per table whose every cell lacks one. Callers (the LightRAG
|
||||
adapter / debug CLI) read this to surface a one-line warning per
|
||||
document instead of crashing.
|
||||
parse_metadata: Optional out-dict that this function mutates with
|
||||
document-level metadata derived during parsing. Currently used
|
||||
for ``first_heading`` — the text of the first heading encountered
|
||||
in document order (regardless of level). Used by the LightRAG
|
||||
adapter to populate ``meta.doc_title`` in ``.blocks.jsonl``.
|
||||
|
||||
Returns:
|
||||
List of block dictionaries with heading, content, type, and metadata
|
||||
"""
|
||||
try:
|
||||
doc = Document(file_path)
|
||||
except PackageNotFoundError as exc:
|
||||
# python-docx surfaces a misleading "Package not found at '...'" for any
|
||||
# file it cannot open as a ZIP/OOXML package — including files that
|
||||
# exist but are corrupt or a different format wearing a .docx extension.
|
||||
# Diagnose the real cause from the magic bytes and raise a DocxContentError
|
||||
# (a ValueError) so the pipeline's per-document handler marks just this
|
||||
# document FAILED with an accurate, actionable message.
|
||||
details, solution = _diagnose_invalid_docx(file_path)
|
||||
raise DocxContentError(
|
||||
format_error("File is not a valid DOCX document", details, solution)
|
||||
) from exc
|
||||
resolver = NumberingResolver(file_path)
|
||||
styles_outline = parse_styles_outline_levels(file_path)
|
||||
|
||||
blocks = []
|
||||
current_heading = "Preface/Uncategorized"
|
||||
current_heading_level = 1 # Default level for "Preface/Uncategorized"
|
||||
current_heading_stack = {} # {level: heading_text} - Use dict to correctly track heading hierarchy
|
||||
current_parent_headings = [] # Parent headings for current block
|
||||
current_paragraphs = [] # Track paragraphs with metadata for splitting
|
||||
first_heading_recorded = (
|
||||
False # Track whether the document's first heading has been captured
|
||||
)
|
||||
|
||||
# Iterate through document body elements (paragraphs and tables)
|
||||
body = doc._element.body
|
||||
|
||||
for element in body:
|
||||
tag = element.tag.split("}")[-1] # Remove namespace
|
||||
|
||||
if tag == "sectPr": # Document-level section break
|
||||
resolver.reset_tracking_state()
|
||||
continue
|
||||
|
||||
if tag == "p": # Paragraph
|
||||
# Get paragraph text with superscript/subscript markup and equations
|
||||
para_text = ""
|
||||
ns = {
|
||||
"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
|
||||
"wp": "http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing",
|
||||
"m": "http://schemas.openxmlformats.org/officeDocument/2006/math",
|
||||
}
|
||||
para_text = extract_paragraph_content(
|
||||
element,
|
||||
ns,
|
||||
drawing_context=drawing_context,
|
||||
)
|
||||
|
||||
para_text = para_text.strip()
|
||||
if not para_text:
|
||||
continue
|
||||
|
||||
# Get numbering label using our resolver
|
||||
label = resolver.get_label(element)
|
||||
full_text = f"{label} {para_text}".strip() if label else para_text
|
||||
|
||||
# Check if this is a heading using the new function
|
||||
outline_level = get_heading_level(element, styles_outline)
|
||||
|
||||
# A "heading" longer than MAX_HEADING_LENGTH is not a real heading.
|
||||
# The common cause (WPS/Word): the author set an outline level on a
|
||||
# paragraph but typed the body with soft line breaks (Shift+Enter →
|
||||
# <w:br/> → '\n') instead of starting a new paragraph, so heading
|
||||
# text + body live in one <w:p>. Split at the first soft break: the
|
||||
# first line stays the heading, the remainder becomes body text. If
|
||||
# there is no usable soft break (a genuine single-line over-long
|
||||
# heading), demote the whole paragraph to body text. Either way we
|
||||
# avoid crashing via validate_heading_length() and never drop content.
|
||||
demoted_body_text = None
|
||||
if outline_level is not None and len(full_text) > MAX_HEADING_LENGTH:
|
||||
head, sep, rest = full_text.partition("\n")
|
||||
if sep and len(head) <= MAX_HEADING_LENGTH:
|
||||
full_text = head
|
||||
demoted_body_text = rest.strip() or None
|
||||
if parse_warnings is not None:
|
||||
parse_warnings["heading_softbreak_split_count"] = (
|
||||
parse_warnings.get("heading_softbreak_split_count", 0) + 1
|
||||
)
|
||||
print(
|
||||
f"Warning: heading paragraph exceeded {MAX_HEADING_LENGTH} "
|
||||
"chars; split at soft line break — kept first line as "
|
||||
"heading, rest as body.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
else:
|
||||
outline_level = None
|
||||
if parse_warnings is not None:
|
||||
parse_warnings["demoted_oversize_heading_count"] = (
|
||||
parse_warnings.get("demoted_oversize_heading_count", 0) + 1
|
||||
)
|
||||
print(
|
||||
f"Warning: paragraph has outline level but is "
|
||||
f"{len(full_text)} chars (> {MAX_HEADING_LENGTH}); treating "
|
||||
"as body text, not a heading.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
if outline_level is not None:
|
||||
# This is a heading (outline level 0-8)
|
||||
# Convert 0-based to 1-based level
|
||||
level = outline_level + 1
|
||||
|
||||
# Extract paraId for this heading
|
||||
heading_para_id = extract_para_id(element)
|
||||
if parse_warnings is not None and not heading_para_id:
|
||||
parse_warnings["missing_paraid_count"] = (
|
||||
parse_warnings.get("missing_paraid_count", 0) + 1
|
||||
)
|
||||
|
||||
# Validate heading length
|
||||
validate_heading_length(full_text, heading_para_id)
|
||||
|
||||
# Truncate heading if needed before storing
|
||||
truncated_text = truncate_heading(full_text, heading_para_id)
|
||||
clean_heading_text = strip_heading_markdown_prefix(truncated_text)
|
||||
|
||||
# Record the document's first heading (any level) for meta.doc_title.
|
||||
if not first_heading_recorded:
|
||||
if parse_metadata is not None:
|
||||
parse_metadata["first_heading"] = clean_heading_text
|
||||
first_heading_recorded = True
|
||||
|
||||
# Every recognized heading starts its own block. Always flush the
|
||||
# accumulated paragraphs so a heading with no body becomes a
|
||||
# standalone block whose content is just the heading text,
|
||||
# instead of being folded into the next heading's block.
|
||||
if current_paragraphs:
|
||||
_flush_current_block(
|
||||
blocks,
|
||||
current_heading,
|
||||
current_paragraphs,
|
||||
current_parent_headings,
|
||||
current_heading_level,
|
||||
)
|
||||
|
||||
# Reset for new block
|
||||
current_paragraphs = []
|
||||
|
||||
# Add heading to current_paragraphs. The content line gets
|
||||
# a markdown ``#`` prefix (capped at 6) via
|
||||
# render_heading_line; ``clean_heading_text`` is kept
|
||||
# for the heading field / stack / parent_headings below.
|
||||
current_paragraphs.append(
|
||||
{
|
||||
"text": render_heading_line(level, truncated_text),
|
||||
"para_id": heading_para_id,
|
||||
"is_table": False,
|
||||
}
|
||||
)
|
||||
|
||||
# Update current_heading and parent_headings for the FIRST heading in a block
|
||||
# (when current_paragraphs just had this heading added as its first element)
|
||||
if len(current_paragraphs) == 1:
|
||||
current_heading = clean_heading_text
|
||||
current_heading_level = level # Only set level when setting heading
|
||||
# Parent headings = all headings from levels strictly less than current level
|
||||
# Sort by level to maintain hierarchy order
|
||||
current_parent_headings = [
|
||||
current_heading_stack[lvl]
|
||||
for lvl in sorted(current_heading_stack.keys())
|
||||
if lvl < level
|
||||
]
|
||||
|
||||
# Update heading stack: remove current level and all lower levels, then add current
|
||||
current_heading_stack = {
|
||||
k: v for k, v in current_heading_stack.items() if k < level
|
||||
}
|
||||
current_heading_stack[level] = clean_heading_text
|
||||
|
||||
# Carry the body text that followed a soft break in an over-long
|
||||
# heading paragraph as a regular body paragraph in the same block.
|
||||
if demoted_body_text:
|
||||
current_paragraphs.append(
|
||||
{
|
||||
"text": demoted_body_text,
|
||||
"para_id": heading_para_id,
|
||||
"is_table": False,
|
||||
}
|
||||
)
|
||||
else:
|
||||
# Regular paragraph content
|
||||
para_id = extract_para_id(element)
|
||||
if parse_warnings is not None and not para_id:
|
||||
parse_warnings["missing_paraid_count"] = (
|
||||
parse_warnings.get("missing_paraid_count", 0) + 1
|
||||
)
|
||||
|
||||
# Store paragraph with metadata for potential splitting
|
||||
current_paragraphs.append(
|
||||
{"text": full_text, "para_id": para_id, "is_table": False}
|
||||
)
|
||||
|
||||
# Check for paragraph-level section break (after processing paragraph)
|
||||
# sectPr in pPr means this paragraph ends a section
|
||||
pPr = element.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}pPr"
|
||||
)
|
||||
if pPr is not None:
|
||||
sectPr = pPr.find(
|
||||
"{http://schemas.openxmlformats.org/wordprocessingml/2006/main}sectPr"
|
||||
)
|
||||
if sectPr is not None:
|
||||
# Section break after this paragraph - reset tracking
|
||||
resolver.reset_tracking_state()
|
||||
|
||||
elif tag == "tbl": # Table
|
||||
# Reset numbering tracking before table (table start boundary)
|
||||
resolver.reset_tracking_state()
|
||||
|
||||
# Directly create Table object from XML element to avoid index mismatch
|
||||
# (doc.tables may have different order due to nested tables)
|
||||
from docx.table import Table
|
||||
|
||||
table = Table(element, doc)
|
||||
table_metadata = TableExtractor.extract_with_metadata(
|
||||
table,
|
||||
numbering_resolver=resolver,
|
||||
drawing_context=drawing_context,
|
||||
)
|
||||
|
||||
table_rows = table_metadata["rows"]
|
||||
para_ids = table_metadata["para_ids"]
|
||||
para_ids_end = table_metadata["para_ids_end"] # Last paraId in each cell
|
||||
header_indices = table_metadata["header_indices"]
|
||||
|
||||
# Skip tables whose every cell is whitespace-only — otherwise an
|
||||
# empty `<table>[[""]]</table>` placeholder would leak into block
|
||||
# content and a useless IRTable would appear in tables.json.
|
||||
if _is_table_empty(table_rows):
|
||||
resolver.reset_tracking_state()
|
||||
continue
|
||||
|
||||
# Count tables whose cells carry no w14:paraId. Legacy / non-Word
|
||||
# docx authors omit these attributes; we no longer fail-fast, but
|
||||
# the adapter surfaces a single warning so the user knows the edit
|
||||
# range hints will be missing for these tables.
|
||||
if parse_warnings is not None and not _table_has_any_paraid(para_ids):
|
||||
parse_warnings["missing_paraid_count"] = (
|
||||
parse_warnings.get("missing_paraid_count", 0) + 1
|
||||
)
|
||||
|
||||
# Convert table to JSON
|
||||
table_json = json.dumps(table_rows, ensure_ascii=False)
|
||||
|
||||
# Extract cross-page repeating header rows (w:tblHeader) once per
|
||||
# table so we can surface them to the sidecar via the block-level
|
||||
# ``table_headers`` list.
|
||||
header_rows = []
|
||||
if header_indices:
|
||||
header_rows = [
|
||||
table_rows[idx] for idx in header_indices if idx < len(table_rows)
|
||||
]
|
||||
header_rows_or_none = header_rows if header_rows else None
|
||||
|
||||
# Emit the whole table as a single <table> placeholder. Token-based
|
||||
# table row splitting is the downstream chunker's responsibility.
|
||||
# Use first valid paraId from table, and last valid paraId (from
|
||||
# para_ids_end) for uuid_end.
|
||||
table_para_id = find_first_valid_para_id(para_ids)
|
||||
table_para_id_end = find_last_valid_para_id(para_ids_end)
|
||||
current_paragraphs.append(
|
||||
{
|
||||
"text": f"<table>{table_json}</table>",
|
||||
"para_id": table_para_id,
|
||||
"para_id_end": table_para_id_end, # Store end paraId for uuid_end calculation
|
||||
"is_table": True,
|
||||
"_table_header": header_rows_or_none,
|
||||
}
|
||||
)
|
||||
|
||||
# Reset numbering tracking after table (table end boundary)
|
||||
resolver.reset_tracking_state()
|
||||
|
||||
# Save final block
|
||||
_flush_current_block(
|
||||
blocks,
|
||||
current_heading,
|
||||
current_paragraphs,
|
||||
current_parent_headings,
|
||||
current_heading_level,
|
||||
)
|
||||
|
||||
return blocks
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Native DOCX engine adapter (implements NativeParserBase hooks)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lightrag.constants import PARSER_ENGINE_NATIVE
|
||||
from lightrag.parser.native_base import NativeParserBase
|
||||
from lightrag.utils import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lightrag.sidecar.ir import IRDoc
|
||||
|
||||
|
||||
class NativeDocxParser(NativeParserBase):
|
||||
"""Native DOCX parser for LightRAG's production parsing path.
|
||||
|
||||
``extract_docx_blocks`` performs only heading-driven structural splitting
|
||||
(one block per DOCX heading). Block sizing is intentionally left to the
|
||||
downstream paragraph-semantic chunker, so this parser emits the
|
||||
one-heading-one-block sidecar contract that chunking consumes.
|
||||
"""
|
||||
|
||||
engine_name = PARSER_ENGINE_NATIVE
|
||||
sidecar_path_style = "basename_only" # legacy native docx convention
|
||||
empty_content_label = "DOCX"
|
||||
|
||||
def validate_source(self, source: Path, file_path: str) -> None:
|
||||
if not (
|
||||
source.exists() and source.is_file() and source.suffix.lower() == ".docx"
|
||||
):
|
||||
raise ValueError(
|
||||
f"Native parser does not support pending file: {file_path}"
|
||||
)
|
||||
|
||||
def extract(
|
||||
self, source: Path, *, parsed_dir: Path, asset_dir: Path, base_name: str
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any], dict[str, Any]]:
|
||||
"""Extract heading-scoped DOCX blocks (sizing left to the chunker)."""
|
||||
|
||||
from lightrag.parser.docx.drawing_image_extractor import (
|
||||
DrawingExtractionContext,
|
||||
load_relationships,
|
||||
)
|
||||
from lightrag.parser.docx.parse_document import extract_docx_blocks
|
||||
|
||||
ctx = DrawingExtractionContext(
|
||||
docx_path=source,
|
||||
blocks_output_path=parsed_dir / f"{base_name}.blocks.jsonl",
|
||||
export_dir_name=asset_dir.name,
|
||||
export_dir_path=asset_dir,
|
||||
)
|
||||
load_relationships(ctx)
|
||||
warnings: dict[str, Any] = {}
|
||||
metadata: dict[str, Any] = {}
|
||||
blocks = extract_docx_blocks(
|
||||
str(source),
|
||||
drawing_context=ctx,
|
||||
parse_warnings=warnings,
|
||||
parse_metadata=metadata,
|
||||
)
|
||||
return blocks, warnings, metadata
|
||||
|
||||
def build_ir(
|
||||
self,
|
||||
blocks: list[dict[str, Any]],
|
||||
*,
|
||||
document_name: str,
|
||||
asset_dir_name: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> "IRDoc":
|
||||
from lightrag.parser.docx.ir_builder import NativeDocxIRBuilder
|
||||
|
||||
return NativeDocxIRBuilder().normalize(
|
||||
blocks,
|
||||
document_name=document_name,
|
||||
asset_dir_name=asset_dir_name,
|
||||
parse_metadata=metadata,
|
||||
)
|
||||
|
||||
def surface_warnings(
|
||||
self, warnings: dict[str, Any], source: Path
|
||||
) -> dict[str, Any] | None:
|
||||
missing = int(warnings.get("missing_paraid_count", 0) or 0)
|
||||
if missing > 0:
|
||||
# Surface once per document; affected blocks emit
|
||||
# ``positions: [{"type": "paraid", "range": null}]``.
|
||||
logger.warning(
|
||||
"[parse_native] %s: %d paragraphs lack paraId; "
|
||||
"Re-saving file in Word 2013+ to regenerate ids.",
|
||||
source.name,
|
||||
missing,
|
||||
)
|
||||
return {"missing_paraid_count": missing}
|
||||
return None
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user