from __future__ import annotations from abc import ABC, abstractmethod from enum import Enum import os from dotenv import load_dotenv from dataclasses import dataclass, field from typing import ( Any, Literal, TypedDict, TypeVar, Callable, Optional, Dict, List, AsyncIterator, ) from .utils import EmbeddingFunc from .types import KnowledgeGraph from .constants import ( DEFAULT_TOP_K, DEFAULT_CHUNK_TOP_K, DEFAULT_MAX_ENTITY_TOKENS, DEFAULT_MAX_RELATION_TOKENS, DEFAULT_MAX_TOTAL_TOKENS, DEFAULT_HISTORY_TURNS, DEFAULT_OLLAMA_MODEL_NAME, DEFAULT_OLLAMA_MODEL_TAG, DEFAULT_OLLAMA_MODEL_SIZE, DEFAULT_OLLAMA_CREATED_AT, DEFAULT_OLLAMA_DIGEST, ) # 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 OllamaServerInfos: def __init__(self, name=None, tag=None): self._lightrag_name = name or os.getenv( "OLLAMA_EMULATING_MODEL_NAME", DEFAULT_OLLAMA_MODEL_NAME ) self._lightrag_tag = tag or os.getenv( "OLLAMA_EMULATING_MODEL_TAG", DEFAULT_OLLAMA_MODEL_TAG ) self.LIGHTRAG_SIZE = DEFAULT_OLLAMA_MODEL_SIZE self.LIGHTRAG_CREATED_AT = DEFAULT_OLLAMA_CREATED_AT self.LIGHTRAG_DIGEST = DEFAULT_OLLAMA_DIGEST @property def LIGHTRAG_NAME(self): return self._lightrag_name @LIGHTRAG_NAME.setter def LIGHTRAG_NAME(self, value): self._lightrag_name = value @property def LIGHTRAG_TAG(self): return self._lightrag_tag @LIGHTRAG_TAG.setter def LIGHTRAG_TAG(self, value): self._lightrag_tag = value @property def LIGHTRAG_MODEL(self): return f"{self._lightrag_name}:{self._lightrag_tag}" class TextChunkSchema(TypedDict): tokens: int content: str full_doc_id: str chunk_order_index: int T = TypeVar("T") @dataclass class QueryParam: """Configuration parameters for query execution in LightRAG.""" mode: Literal["local", "global", "hybrid", "naive", "mix", "bypass"] = "mix" """Specifies the retrieval mode: - "local": Focuses on context-dependent information. - "global": Utilizes global knowledge. - "hybrid": Combines local and global retrieval methods. - "naive": Performs a basic search without advanced techniques. - "mix": Integrates knowledge graph and vector retrieval. """ only_need_context: bool = False """If True, only returns the retrieved context without generating a response.""" only_need_prompt: bool = False """If True, only returns the generated prompt without producing a response.""" response_type: str = "Multiple Paragraphs" """Defines the response format. Examples: 'Multiple Paragraphs', 'Single Paragraph', 'Bullet Points'.""" stream: bool = False """If True, enables streaming output for real-time responses.""" top_k: int = int(os.getenv("TOP_K", str(DEFAULT_TOP_K))) """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode.""" chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", str(DEFAULT_CHUNK_TOP_K))) """Number of text chunks to retrieve initially from vector search and keep after reranking. If None, defaults to top_k value. """ max_entity_tokens: int = int( os.getenv("MAX_ENTITY_TOKENS", str(DEFAULT_MAX_ENTITY_TOKENS)) ) """Maximum number of tokens allocated for entity context in unified token control system.""" max_relation_tokens: int = int( os.getenv("MAX_RELATION_TOKENS", str(DEFAULT_MAX_RELATION_TOKENS)) ) """Maximum number of tokens allocated for relationship context in unified token control system.""" max_total_tokens: int = int( os.getenv("MAX_TOTAL_TOKENS", str(DEFAULT_MAX_TOTAL_TOKENS)) ) """Maximum total tokens budget for the entire query context (entities + relations + chunks + system prompt).""" hl_keywords: list[str] = field(default_factory=list) """List of high-level keywords to prioritize in retrieval.""" ll_keywords: list[str] = field(default_factory=list) """List of low-level keywords to refine retrieval focus.""" # History mesages is only send to LLM for context, not used for retrieval conversation_history: list[dict[str, str]] = field(default_factory=list) """Stores past conversation history to maintain context. Format: [{"role": "user/assistant", "content": "message"}]. """ # TODO: deprecated. No longer used in the codebase, all conversation_history messages is send to LLM history_turns: int = int(os.getenv("HISTORY_TURNS", str(DEFAULT_HISTORY_TURNS))) """Number of complete conversation turns (user-assistant pairs) to consider in the response context.""" model_func: Callable[..., object] | None = None """Optional override for the LLM model function to use for this specific query. If provided, this will be used instead of the global model function. This allows using different models for different query modes. """ user_prompt: str | None = None """User-provided prompt for the query. Addition instructions for LLM. If provided, this will be inject into the prompt template. It's purpose is the let user customize the way LLM generate the response. """ enable_rerank: bool = os.getenv("RERANK_BY_DEFAULT", "true").lower() == "true" """Enable reranking for retrieved text chunks. If True but no rerank model is configured, a warning will be issued. Default is True to enable reranking when rerank model is available. """ include_references: bool = False """If True, includes reference list in the response for supported endpoints. This parameter controls whether the API response includes a references field containing citation information for the retrieved content. """ @dataclass class StorageNameSpace(ABC): namespace: str workspace: str global_config: dict[str, Any] async def initialize(self): """Initialize the storage""" pass async def finalize(self): """Finalize the storage""" pass @abstractmethod async def index_done_callback(self) -> None: """Commit the storage operations after indexing""" @abstractmethod async def drop(self) -> dict[str, str]: """Drop all data from storage and clean up resources This abstract method defines the contract for dropping all data from a storage implementation. Each storage type must implement this method to: 1. Clear all data from memory and/or external storage 2. Remove any associated storage files if applicable 3. Reset the storage to its initial state 4. Handle cleanup of any resources 5. Notify other processes if necessary 6. This action should persistent the data to disk immediately. Returns: dict[str, str]: Operation status and message with the following format: { "status": str, # "success" or "error" "message": str # "data dropped" on success, error details on failure } Implementation specific: - On success: return {"status": "success", "message": "data dropped"} - On failure: return {"status": "error", "message": ""} - If not supported: return {"status": "error", "message": "unsupported"} """ @dataclass class BaseVectorStorage(StorageNameSpace, ABC): embedding_func: EmbeddingFunc cosine_better_than_threshold: float = field(default=0.2) meta_fields: set[str] = field(default_factory=set) def _validate_embedding_func(self): """Validate that embedding_func is provided. This method should be called at the beginning of __post_init__ in all vector storage implementations. Raises: ValueError: If embedding_func is None """ if self.embedding_func is None: raise ValueError( "embedding_func is required for vector storage. " "Please provide a valid EmbeddingFunc instance." ) def _generate_collection_suffix(self) -> str | None: """Generates collection/table suffix from embedding_func. Return suffix if model_name exists in embedding_func, otherwise return None. Note: embedding_func is guaranteed to exist (validated in __post_init__). Returns: str | None: Suffix string e.g. "text_embedding_3_large_3072d", or None if model_name not available """ import re # Check if model_name exists (model_name is optional in EmbeddingFunc) model_name = getattr(self.embedding_func, "model_name", None) if not model_name: return None # embedding_dim is required in EmbeddingFunc embedding_dim = self.embedding_func.embedding_dim # Generate suffix: clean model name and append dimension safe_model_name = re.sub(r"[^a-zA-Z0-9_]", "_", model_name.lower()) return f"{safe_model_name}_{embedding_dim}d" @abstractmethod async def query( self, query: str, top_k: int, query_embedding: list[float] = None ) -> list[dict[str, Any]]: """Query the vector storage and retrieve top_k results. Args: query: The query string to search for top_k: Number of top results to return query_embedding: Optional pre-computed embedding for the query. If provided, skips embedding computation for better performance. """ @abstractmethod async def upsert(self, data: dict[str, dict[str, Any]]) -> None: """Insert or update vectors in the storage. Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption """ @abstractmethod async def delete_entity(self, entity_name: str) -> None: """Delete a single entity by its name. Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption """ @abstractmethod async def delete_entity_relation(self, entity_name: str) -> None: """Delete relations for a given entity. Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption """ @abstractmethod 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 """ pass @abstractmethod 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 """ pass @abstractmethod async def delete(self, ids: list[str]): """Delete vectors with specified IDs Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption Args: ids: List of vector IDs to be deleted """ @abstractmethod async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]: """Get vectors by their IDs, returning only ID and vector data for efficiency Args: ids: List of unique identifiers Returns: Dictionary mapping IDs to their vector embeddings Format: {id: [vector_values], ...} """ pass @dataclass class BaseKVStorage(StorageNameSpace, ABC): embedding_func: EmbeddingFunc @abstractmethod async def get_by_id(self, id: str) -> dict[str, Any] | None: """Get value by id""" @abstractmethod async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]: """Get values by ids""" @abstractmethod async def filter_keys(self, keys: set[str]) -> set[str]: """Return un-exist keys""" @abstractmethod async def upsert(self, data: dict[str, dict[str, Any]]) -> None: """Upsert data Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. update flags to notify other processes that data persistence is needed """ @abstractmethod async def delete(self, ids: list[str]) -> None: """Delete specific records from storage by their IDs Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. update flags to notify other processes that data persistence is needed Args: ids (list[str]): List of document IDs to be deleted from storage Returns: None """ @abstractmethod async def is_empty(self) -> bool: """Check if the storage is empty Returns: bool: True if storage contains no data, False otherwise """ @dataclass class BaseGraphStorage(StorageNameSpace, ABC): """All operations related to edges in graph should be undirected.""" embedding_func: EmbeddingFunc @abstractmethod async def has_node(self, node_id: str) -> bool: """Check if a node exists in the graph. Args: node_id: The ID of the node to check Returns: True if the node exists, False otherwise """ @abstractmethod async def has_edge(self, source_node_id: str, target_node_id: str) -> bool: """Check if an edge exists between two nodes. Args: source_node_id: The ID of the source node target_node_id: The ID of the target node Returns: True if the edge exists, False otherwise """ @abstractmethod async def node_degree(self, node_id: str) -> int: """Get the degree (number of connected edges) of a node. Args: node_id: The ID of the node Returns: The number of edges connected to the node """ @abstractmethod async def edge_degree(self, src_id: str, tgt_id: str) -> int: """Get the total degree of an edge (sum of degrees of its source and target nodes). Args: src_id: The ID of the source node tgt_id: The ID of the target node Returns: The sum of the degrees of the source and target nodes """ @abstractmethod async def get_node(self, node_id: str) -> dict[str, str] | None: """Get node by its ID, returning only node properties. Args: node_id: The ID of the node to retrieve Returns: A dictionary of node properties if found, None otherwise """ @abstractmethod async def get_edge( self, source_node_id: str, target_node_id: str ) -> dict[str, str] | None: """Get edge properties between two nodes. Args: source_node_id: The ID of the source node target_node_id: The ID of the target node Returns: A dictionary of edge properties if found, None otherwise """ @abstractmethod async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None: """Get all edges connected to a node. Args: source_node_id: The ID of the node to get edges for Returns: A list of (source_id, target_id) tuples representing edges, or None if the node doesn't exist """ async def get_nodes_batch(self, node_ids: list[str]) -> dict[str, dict]: """Get nodes as a batch using UNWIND Default implementation fetches nodes one by one. Override this method for better performance in storage backends that support batch operations. """ result = {} for node_id in node_ids: node = await self.get_node(node_id) if node is not None: result[node_id] = node return result async def node_degrees_batch(self, node_ids: list[str]) -> dict[str, int]: """Node degrees as a batch using UNWIND Default implementation fetches node degrees one by one. Override this method for better performance in storage backends that support batch operations. """ result = {} for node_id in node_ids: degree = await self.node_degree(node_id) result[node_id] = degree return result async def edge_degrees_batch( self, edge_pairs: list[tuple[str, str]] ) -> dict[tuple[str, str], int]: """Edge degrees as a batch using UNWIND also uses node_degrees_batch Default implementation calculates edge degrees one by one. Override this method for better performance in storage backends that support batch operations. """ result = {} for src_id, tgt_id in edge_pairs: degree = await self.edge_degree(src_id, tgt_id) result[(src_id, tgt_id)] = degree return result async def get_edges_batch( self, pairs: list[dict[str, str]] ) -> dict[tuple[str, str], dict]: """Get edges as a batch using UNWIND Default implementation fetches edges one by one. Override this method for better performance in storage backends that support batch operations. """ result = {} for pair in pairs: src_id = pair["src"] tgt_id = pair["tgt"] edge = await self.get_edge(src_id, tgt_id) if edge is not None: result[(src_id, tgt_id)] = edge return result async def get_nodes_edges_batch( self, node_ids: list[str] ) -> dict[str, list[tuple[str, str]]]: """Get nodes edges as a batch using UNWIND Default implementation fetches node edges one by one. Override this method for better performance in storage backends that support batch operations. """ result = {} for node_id in node_ids: edges = await self.get_node_edges(node_id) result[node_id] = edges if edges is not None else [] return result @abstractmethod async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None: """Insert a new node or update an existing node in the graph. Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption Args: node_id: The ID of the node to insert or update node_data: A dictionary of node properties """ @abstractmethod async def upsert_edge( self, source_node_id: str, target_node_id: str, edge_data: dict[str, str] ) -> None: """Insert a new edge or update an existing edge in the graph. Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption Args: source_node_id: The ID of the source node target_node_id: The ID of the target node edge_data: A dictionary of edge properties """ @abstractmethod async def delete_node(self, node_id: str) -> None: """Delete a node from the graph. Importance notes for in-memory storage: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption Args: node_id: The ID of the node to delete """ @abstractmethod async def remove_nodes(self, nodes: list[str]): """Delete multiple nodes Importance notes: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption Args: nodes: List of node IDs to be deleted """ @abstractmethod async def remove_edges(self, edges: list[tuple[str, str]]): """Delete multiple edges Importance notes: 1. Changes will be persisted to disk during the next index_done_callback 2. Only one process should updating the storage at a time before index_done_callback, KG-storage-log should be used to avoid data corruption Args: edges: List of edges to be deleted, each edge is a (source, target) tuple """ # TODO: deprecated @abstractmethod async def get_all_labels(self) -> list[str]: """Get all labels in the graph. Returns: A list of all node labels in the graph, sorted alphabetically """ @abstractmethod async def get_knowledge_graph( self, node_label: str, max_depth: int = 3, max_nodes: int = 1000 ) -> 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, Defaults to 1000（BFS if possible) Returns: KnowledgeGraph object containing nodes and edges, with an is_truncated flag indicating whether the graph was truncated due to max_nodes limit """ @abstractmethod 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 (Edge is bidirectional for some storage implementation; deduplication must be handled by the caller) """ @abstractmethod 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 """ @abstractmethod async def get_popular_labels(self, limit: int = 300) -> list[str]: """Get popular labels by node degree (most connected entities) Args: limit: Maximum number of labels to return Returns: List of labels sorted by degree (highest first) """ @abstractmethod async def search_labels(self, query: str, limit: int = 50) -> list[str]: """Search labels with fuzzy matching Args: query: Search query string limit: Maximum number of results to return Returns: List of matching labels sorted by relevance """ class DocStatus(str, Enum): """Document processing status""" PENDING = "pending" PROCESSING = "processing" PREPROCESSED = "preprocessed" PROCESSED = "processed" FAILED = "failed" @dataclass class DocProcessingStatus: """Document processing status data structure""" content_summary: str """First 100 chars of document content, used for preview""" content_length: int """Total length of document""" file_path: str """File path of the document""" status: DocStatus """Current processing status""" created_at: str """ISO format timestamp when document was created""" updated_at: str """ISO format timestamp when document was last updated""" track_id: str | None = None """Tracking ID for monitoring progress""" chunks_count: int | None = None """Number of chunks after splitting, used for processing""" chunks_list: list[str] | None = field(default_factory=list) """List of chunk IDs associated with this document, used for deletion""" error_msg: str | None = None """Error message if failed""" metadata: dict[str, Any] = field(default_factory=dict) """Additional metadata""" multimodal_processed: bool | None = field(default=None, repr=False) """Internal field: indicates if multimodal processing is complete. Not shown in repr() but accessible for debugging.""" def __post_init__(self): """ Handle status conversion based on multimodal_processed field. Business rules: - If multimodal_processed is False and status is PROCESSED, then change status to PREPROCESSED - The multimodal_processed field is kept (with repr=False) for internal use and debugging """ # Apply status conversion logic if self.multimodal_processed is not None: if ( self.multimodal_processed is False and self.status == DocStatus.PROCESSED ): self.status = DocStatus.PREPROCESSED @dataclass class DocStatusStorage(BaseKVStorage, ABC): """Base class for document status storage""" @abstractmethod async def get_status_counts(self) -> dict[str, int]: """Get counts of documents in each status""" @abstractmethod async def get_docs_by_status( self, status: DocStatus ) -> dict[str, DocProcessingStatus]: """Get all documents with a specific status""" @abstractmethod async def get_docs_by_track_id( self, track_id: str ) -> dict[str, DocProcessingStatus]: """Get all documents with a specific track_id""" @abstractmethod async def get_docs_paginated( self, status_filter: 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) """ @abstractmethod 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 """ @abstractmethod async def get_doc_by_file_path(self, file_path: str) -> dict[str, Any] | None: """Get document by file path Args: file_path: The file path to search for Returns: dict[str, Any] | None: Document data if found, None otherwise Returns the same format as get_by_ids method """ class StoragesStatus(str, Enum): """Storages status""" NOT_CREATED = "not_created" CREATED = "created" INITIALIZED = "initialized" FINALIZED = "finalized" @dataclass class DeletionResult: """Represents the result of a deletion operation.""" status: Literal["success", "not_found", "fail"] doc_id: str message: str status_code: int = 200 file_path: str | None = None # Unified Query Result Data Structures for Reference List Support @dataclass class QueryResult: """ Unified query result data structure for all query modes. Attributes: content: Text content for non-streaming responses response_iterator: Streaming response iterator for streaming responses raw_data: Complete structured data including references and metadata is_streaming: Whether this is a streaming result """ content: Optional[str] = None response_iterator: Optional[AsyncIterator[str]] = None raw_data: Optional[Dict[str, Any]] = None is_streaming: bool = False @property def reference_list(self) -> List[Dict[str, str]]: """ Convenient property to extract reference list from raw_data. Returns: List[Dict[str, str]]: Reference list in format: [{"reference_id": "1", "file_path": "/path/to/file.pdf"}, ...] """ if self.raw_data: return self.raw_data.get("data", {}).get("references", []) return [] @property def metadata(self) -> Dict[str, Any]: """ Convenient property to extract metadata from raw_data. Returns: Dict[str, Any]: Query metadata including query_mode, keywords, etc. """ if self.raw_data: return self.raw_data.get("metadata", {}) return {} @dataclass class QueryContextResult: """ Unified query context result data structure. Attributes: context: LLM context string raw_data: Complete structured data including reference_list """ context: str raw_data: Dict[str, Any] @property def reference_list(self) -> List[Dict[str, str]]: """Convenient property to extract reference list from raw_data.""" return self.raw_data.get("data", {}).get("references", [])