gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'feat-node-expand' into webui-node-expansion

Browse Source
This commit is contained in:
yangdx
2025-03-17 05:04:53 +08:00
parent 550add8a1c 3e3e963c80
commit 0695f0a130
12 changed files with 157 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@@ -1061,7 +1061,7 @@ Valid modes are:
| **llm\_model\_func** | `callable` | Function for LLM generation | `gpt_4o_mini_complete` | | **llm\_model\_func** | `callable` | Function for LLM generation | `gpt_4o_mini_complete` |
| **llm\_model\_name** | `str` | LLM model name for generation | `meta-llama/Llama-3.2-1B-Instruct` | | **llm\_model\_name** | `str` | LLM model name for generation | `meta-llama/Llama-3.2-1B-Instruct` |
| **llm\_model\_max\_token\_size** | `int` | Maximum token size for LLM generation (affects entity relation summaries) | `32768`（default value changed by env var MAX_TOKENS) | | **llm\_model\_max\_token\_size** | `int` | Maximum token size for LLM generation (affects entity relation summaries) | `32768`（default value changed by env var MAX_TOKENS) |
| **llm\_model\_max\_async** | `int` | Maximum number of concurrent asynchronous LLM processes | `16`（default value changed by env var MAX_ASYNC) | | **llm\_model\_max\_async** | `int` | Maximum number of concurrent asynchronous LLM processes | `4`（default value changed by env var MAX_ASYNC) |
| **llm\_model\_kwargs** | `dict` | Additional parameters for LLM generation | | | **llm\_model\_kwargs** | `dict` | Additional parameters for LLM generation | |
| **vector\_db\_storage\_cls\_kwargs** | `dict` | Additional parameters for vector database, like setting the threshold for nodes and relations retrieval. | cosine_better_than_threshold: 0.2（default value changed by env var COSINE_THRESHOLD) | | **vector\_db\_storage\_cls\_kwargs** | `dict` | Additional parameters for vector database, like setting the threshold for nodes and relations retrieval. | cosine_better_than_threshold: 0.2（default value changed by env var COSINE_THRESHOLD) |
| **enable\_llm\_cache** | `bool` | If `TRUE`, stores LLM results in cache; repeated prompts return cached responses | `TRUE` | | **enable\_llm\_cache** | `bool` | If `TRUE`, stores LLM results in cache; repeated prompts return cached responses | `TRUE` |

3
env.example
View File

@@ -50,7 +50,8 @@
# MAX_TOKEN_SUMMARY=500 # Max tokens for entity or relations summary # MAX_TOKEN_SUMMARY=500 # Max tokens for entity or relations summary
# SUMMARY_LANGUAGE=English # SUMMARY_LANGUAGE=English
# MAX_EMBED_TOKENS=8192 # MAX_EMBED_TOKENS=8192
# ENABLE_LLM_CACHE_FOR_EXTRACT=false # Enable LLM cache for entity extraction, defaults to false # ENABLE_LLM_CACHE_FOR_EXTRACT=true # Enable LLM cache for entity extraction
# MAX_PARALLEL_INSERT=2 # Maximum number of parallel processing documents in pipeline
### LLM Configuration (Use valid host. For local services installed with docker, you can use host.docker.internal) ### LLM Configuration (Use valid host. For local services installed with docker, you can use host.docker.internal)
LLM_BINDING=ollama LLM_BINDING=ollama

2
lightrag/api/README.md
View File

@@ -224,7 +224,7 @@ LightRAG supports binding to various LLM/Embedding backends:
Use environment variables `LLM_BINDING` or CLI argument `--llm-binding` to select LLM backend type. Use environment variables `EMBEDDING_BINDING` or CLI argument `--embedding-binding` to select LLM backend type. Use environment variables `LLM_BINDING` or CLI argument `--llm-binding` to select LLM backend type. Use environment variables `EMBEDDING_BINDING` or CLI argument `--embedding-binding` to select LLM backend type.
### Entity Extraction Configuration ### Entity Extraction Configuration
* ENABLE_LLM_CACHE_FOR_EXTRACT: Enable LLM cache for entity extraction (default: false) * ENABLE_LLM_CACHE_FOR_EXTRACT: Enable LLM cache for entity extraction (default: true)
It's very common to set `ENABLE_LLM_CACHE_FOR_EXTRACT` to true for test environment to reduce the cost of LLM calls. It's very common to set `ENABLE_LLM_CACHE_FOR_EXTRACT` to true for test environment to reduce the cost of LLM calls.

2
lightrag/api/docs/LightRagWithPostGRESQL.md
View File

@@ -141,7 +141,7 @@ Start the LightRAG server using specified options:
lightrag-server --port 9621 --key sk-somepassword --kv-storage PGKVStorage --graph-storage PGGraphStorage --vector-storage PGVectorStorage --doc-status-storage PGDocStatusStorage lightrag-server --port 9621 --key sk-somepassword --kv-storage PGKVStorage --graph-storage PGGraphStorage --vector-storage PGVectorStorage --doc-status-storage PGDocStatusStorage
``` ```
Replace `the-port-number` with your desired port number (default is 9621) and `your-secret-key` with a secure key. Replace the `port` number with your desired port number (default is 9621) and `your-secret-key` with a secure key.
## Conclusion ## Conclusion

2
lightrag/api/utils_api.py
View File

@@ -364,7 +364,7 @@ def parse_args(is_uvicorn_mode: bool = False) -> argparse.Namespace:
# Inject LLM cache configuration # Inject LLM cache configuration
args.enable_llm_cache_for_extract = get_env_value( args.enable_llm_cache_for_extract = get_env_value(
"ENABLE_LLM_CACHE_FOR_EXTRACT", False, bool "ENABLE_LLM_CACHE_FOR_EXTRACT", True, bool
) )
# Select Document loading tool (DOCLING, DEFAULT) # Select Document loading tool (DOCLING, DEFAULT)

4
lightrag/api/webui/assets/index-B_8hp7Xk.js → lightrag/api/webui/assets/index-nzv8EoUv.js
View File

File diff suppressed because one or more lines are too long

2
lightrag/api/webui/index.html
View File

@@ -8,7 +8,7 @@
<link rel="icon" type="image/svg+xml" href="./logo.png" /> <link rel="icon" type="image/svg+xml" href="./logo.png" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" /> <meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>Lightrag</title> <title>Lightrag</title>
<script type="module" crossorigin src="./assets/index-B_8hp7Xk.js"></script> <script type="module" crossorigin src="./assets/index-nzv8EoUv.js"></script>
<link rel="stylesheet" crossorigin href="./assets/index-TPDyec81.css"> <link rel="stylesheet" crossorigin href="./assets/index-TPDyec81.css">
</head> </head>
<body> <body>

6
lightrag/kg/postgres_impl.py
View File

@@ -755,7 +755,7 @@ class PGDocStatusStorage(DocStatusStorage):
result = await self.db.query(sql, params, True) result = await self.db.query(sql, params, True)
docs_by_status = { docs_by_status = {
element["id"]: DocProcessingStatus( element["id"]: DocProcessingStatus(
content=result[0]["content"], content=element["content"],
content_summary=element["content_summary"], content_summary=element["content_summary"],
content_length=element["content_length"], content_length=element["content_length"],
status=element["status"], status=element["status"],
@@ -1556,7 +1556,7 @@ TABLES = {
content_vector VECTOR, content_vector VECTOR,
create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
update_time TIMESTAMP, update_time TIMESTAMP,
chunk_id VARCHAR(255) NULL, chunk_id TEXT NULL,
CONSTRAINT LIGHTRAG_VDB_ENTITY_PK PRIMARY KEY (workspace, id) CONSTRAINT LIGHTRAG_VDB_ENTITY_PK PRIMARY KEY (workspace, id)
)""" )"""
}, },
@@ -1570,7 +1570,7 @@ TABLES = {
content_vector VECTOR, content_vector VECTOR,
create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
update_time TIMESTAMP, update_time TIMESTAMP,
chunk_id VARCHAR(255) NULL, chunk_id TEXT NULL,
CONSTRAINT LIGHTRAG_VDB_RELATION_PK PRIMARY KEY (workspace, id) CONSTRAINT LIGHTRAG_VDB_RELATION_PK PRIMARY KEY (workspace, id)
)""" )"""
}, },

266
lightrag/lightrag.py
View File

@@ -214,7 +214,7 @@ class LightRAG:
llm_model_max_token_size: int = field(default=int(os.getenv("MAX_TOKENS", 32768))) llm_model_max_token_size: int = field(default=int(os.getenv("MAX_TOKENS", 32768)))
"""Maximum number of tokens allowed per LLM response.""" """Maximum number of tokens allowed per LLM response."""
llm_model_max_async: int = field(default=int(os.getenv("MAX_ASYNC", 16))) llm_model_max_async: int = field(default=int(os.getenv("MAX_ASYNC", 4)))
"""Maximum number of concurrent LLM calls.""" """Maximum number of concurrent LLM calls."""
llm_model_kwargs: dict[str, Any] = field(default_factory=dict) llm_model_kwargs: dict[str, Any] = field(default_factory=dict)
@@ -238,7 +238,7 @@ class LightRAG:
# Extensions # Extensions
# --- # ---
max_parallel_insert: int = field(default=int(os.getenv("MAX_PARALLEL_INSERT", 20))) max_parallel_insert: int = field(default=int(os.getenv("MAX_PARALLEL_INSERT", 2)))
"""Maximum number of parallel insert operations.""" """Maximum number of parallel insert operations."""
addon_params: dict[str, Any] = field( addon_params: dict[str, Any] = field(
@@ -553,6 +553,7 @@ class LightRAG:
Args: Args:
input: Single document string or list of document strings input: Single document string or list of document strings
split_by_character: if split_by_character is not None, split the string by character, if chunk longer than split_by_character: if split_by_character is not None, split the string by character, if chunk longer than
chunk_token_size, it will be split again by token size.
split_by_character_only: if split_by_character_only is True, split the string by character only, when split_by_character_only: if split_by_character_only is True, split the string by character only, when
split_by_character is None, this parameter is ignored. split_by_character is None, this parameter is ignored.
ids: single string of the document ID or list of unique document IDs, if not provided, MD5 hash IDs will be generated ids: single string of the document ID or list of unique document IDs, if not provided, MD5 hash IDs will be generated
@@ -574,6 +575,7 @@ class LightRAG:
Args: Args:
input: Single document string or list of document strings input: Single document string or list of document strings
split_by_character: if split_by_character is not None, split the string by character, if chunk longer than split_by_character: if split_by_character is not None, split the string by character, if chunk longer than
chunk_token_size, it will be split again by token size.
split_by_character_only: if split_by_character_only is True, split the string by character only, when split_by_character_only: if split_by_character_only is True, split the string by character only, when
split_by_character is None, this parameter is ignored. split_by_character is None, this parameter is ignored.
ids: list of unique document IDs, if not provided, MD5 hash IDs will be generated ids: list of unique document IDs, if not provided, MD5 hash IDs will be generated
@@ -767,7 +769,6 @@ class LightRAG:
async with pipeline_status_lock: async with pipeline_status_lock:
# Ensure only one worker is processing documents # Ensure only one worker is processing documents
if not pipeline_status.get("busy", False): if not pipeline_status.get("busy", False):
# 先检查是否有需要处理的文档
processing_docs, failed_docs, pending_docs = await asyncio.gather( processing_docs, failed_docs, pending_docs = await asyncio.gather(
self.doc_status.get_docs_by_status(DocStatus.PROCESSING), self.doc_status.get_docs_by_status(DocStatus.PROCESSING),
self.doc_status.get_docs_by_status(DocStatus.FAILED), self.doc_status.get_docs_by_status(DocStatus.FAILED),
@@ -779,12 +780,10 @@ class LightRAG:
to_process_docs.update(failed_docs) to_process_docs.update(failed_docs)
to_process_docs.update(pending_docs) to_process_docs.update(pending_docs)
# 如果没有需要处理的文档，直接返回，保留 pipeline_status 中的内容不变
if not to_process_docs: if not to_process_docs:
logger.info("No documents to process") logger.info("No documents to process")
return return
# 有文档需要处理，更新 pipeline_status
pipeline_status.update( pipeline_status.update(
{ {
"busy": True, "busy": True,
@@ -823,7 +822,7 @@ class LightRAG:
for i in range(0, len(to_process_docs), self.max_parallel_insert) for i in range(0, len(to_process_docs), self.max_parallel_insert)
] ]
log_message = f"Number of batches to process: {len(docs_batches)}." log_message = f"Processing {len(to_process_docs)} document(s) in {len(docs_batches)} batches"
logger.info(log_message) logger.info(log_message)
# Update pipeline status with current batch information # Update pipeline status with current batch information
@@ -832,140 +831,149 @@ class LightRAG:
pipeline_status["latest_message"] = log_message pipeline_status["latest_message"] = log_message
pipeline_status["history_messages"].append(log_message) pipeline_status["history_messages"].append(log_message)
batches: list[Any] = [] async def process_document(
# 3. iterate over batches doc_id: str,
for batch_idx, docs_batch in enumerate(docs_batches): status_doc: DocProcessingStatus,
# Update current batch in pipeline status (directly, as it's atomic) split_by_character: str | None,
pipeline_status["cur_batch"] += 1 split_by_character_only: bool,
pipeline_status: dict,
async def batch( pipeline_status_lock: asyncio.Lock,
batch_idx: int, ) -> None:
docs_batch: list[tuple[str, DocProcessingStatus]], """Process single document"""
size_batch: int, try:
) -> None: # Generate chunks from document
log_message = ( chunks: dict[str, Any] = {
f"Start processing batch {batch_idx + 1} of {size_batch}." compute_mdhash_id(dp["content"], prefix="chunk-"): {
) **dp,
logger.info(log_message) "full_doc_id": doc_id,
pipeline_status["latest_message"] = log_message
pipeline_status["history_messages"].append(log_message)
# 4. iterate over batch
for doc_id_processing_status in docs_batch:
doc_id, status_doc = doc_id_processing_status
# Generate chunks from document
chunks: dict[str, Any] = {
compute_mdhash_id(dp["content"], prefix="chunk-"): {
**dp,
"full_doc_id": doc_id,
}
for dp in self.chunking_func(
status_doc.content,
split_by_character,
split_by_character_only,
self.chunk_overlap_token_size,
self.chunk_token_size,
self.tiktoken_model_name,
)
} }
# Process document (text chunks and full docs) in parallel for dp in self.chunking_func(
# Create tasks with references for potential cancellation status_doc.content,
doc_status_task = asyncio.create_task( split_by_character,
self.doc_status.upsert( split_by_character_only,
{ self.chunk_overlap_token_size,
doc_id: { self.chunk_token_size,
"status": DocStatus.PROCESSING, self.tiktoken_model_name,
"updated_at": datetime.now().isoformat(), )
"content": status_doc.content, }
"content_summary": status_doc.content_summary, # Process document (text chunks and full docs) in parallel
"content_length": status_doc.content_length, # Create tasks with references for potential cancellation
"created_at": status_doc.created_at, doc_status_task = asyncio.create_task(
} self.doc_status.upsert(
{
doc_id: {
"status": DocStatus.PROCESSING,
"updated_at": datetime.now().isoformat(),
"content": status_doc.content,
"content_summary": status_doc.content_summary,
"content_length": status_doc.content_length,
"created_at": status_doc.created_at,
} }
) }
) )
chunks_vdb_task = asyncio.create_task( )
self.chunks_vdb.upsert(chunks) chunks_vdb_task = asyncio.create_task(
self.chunks_vdb.upsert(chunks)
)
entity_relation_task = asyncio.create_task(
self._process_entity_relation_graph(
chunks, pipeline_status, pipeline_status_lock
) )
entity_relation_task = asyncio.create_task( )
self._process_entity_relation_graph( full_docs_task = asyncio.create_task(
chunks, pipeline_status, pipeline_status_lock self.full_docs.upsert(
) {doc_id: {"content": status_doc.content}}
) )
full_docs_task = asyncio.create_task( )
self.full_docs.upsert( text_chunks_task = asyncio.create_task(
{doc_id: {"content": status_doc.content}} self.text_chunks.upsert(chunks)
) )
) tasks = [
text_chunks_task = asyncio.create_task( doc_status_task,
self.text_chunks.upsert(chunks) chunks_vdb_task,
) entity_relation_task,
tasks = [ full_docs_task,
doc_status_task, text_chunks_task,
]
await asyncio.gather(*tasks)
await self.doc_status.upsert(
{
doc_id: {
"status": DocStatus.PROCESSED,
"chunks_count": len(chunks),
"content": status_doc.content,
"content_summary": status_doc.content_summary,
"content_length": status_doc.content_length,
"created_at": status_doc.created_at,
"updated_at": datetime.now().isoformat(),
}
}
)
except Exception as e:
# Log error and update pipeline status
error_msg = f"Failed to process document {doc_id}: {str(e)}"
logger.error(error_msg)
async with pipeline_status_lock:
pipeline_status["latest_message"] = error_msg
pipeline_status["history_messages"].append(error_msg)
# Cancel other tasks as they are no longer meaningful
for task in [
chunks_vdb_task, chunks_vdb_task,
entity_relation_task, entity_relation_task,
full_docs_task, full_docs_task,
text_chunks_task, text_chunks_task,
] ]:
try: if not task.done():
await asyncio.gather(*tasks) task.cancel()
await self.doc_status.upsert( # Update document status to failed
{ await self.doc_status.upsert(
doc_id: { {
"status": DocStatus.PROCESSED, doc_id: {
"chunks_count": len(chunks), "status": DocStatus.FAILED,
"content": status_doc.content, "error": str(e),
"content_summary": status_doc.content_summary, "content": status_doc.content,
"content_length": status_doc.content_length, "content_summary": status_doc.content_summary,
"created_at": status_doc.created_at, "content_length": status_doc.content_length,
"updated_at": datetime.now().isoformat(), "created_at": status_doc.created_at,
} "updated_at": datetime.now().isoformat(),
} }
) }
except Exception as e:
# Log error and update pipeline status
error_msg = (
f"Failed to process document {doc_id}: {str(e)}"
)
logger.error(error_msg)
pipeline_status["latest_message"] = error_msg
pipeline_status["history_messages"].append(error_msg)
# Cancel other tasks as they are no longer meaningful
for task in [
chunks_vdb_task,
entity_relation_task,
full_docs_task,
text_chunks_task,
]:
if not task.done():
task.cancel()
# Update document status to failed
await self.doc_status.upsert(
{
doc_id: {
"status": DocStatus.FAILED,
"error": str(e),
"content": status_doc.content,
"content_summary": status_doc.content_summary,
"content_length": status_doc.content_length,
"created_at": status_doc.created_at,
"updated_at": datetime.now().isoformat(),
}
}
)
continue
log_message = (
f"Completed batch {batch_idx + 1} of {len(docs_batches)}."
) )
logger.info(log_message)
pipeline_status["latest_message"] = log_message
pipeline_status["history_messages"].append(log_message)
batches.append(batch(batch_idx, docs_batch, len(docs_batches))) # 3. iterate over batches
total_batches = len(docs_batches)
for batch_idx, docs_batch in enumerate(docs_batches):
current_batch = batch_idx + 1
log_message = (
f"Start processing batch {current_batch} of {total_batches}."
)
logger.info(log_message)
pipeline_status["cur_batch"] = current_batch
pipeline_status["latest_message"] = log_message
pipeline_status["history_messages"].append(log_message)
await asyncio.gather(*batches) doc_tasks = []
await self._insert_done() for doc_id, status_doc in docs_batch:
doc_tasks.append(
process_document(
doc_id,
status_doc,
split_by_character,
split_by_character_only,
pipeline_status,
pipeline_status_lock,
)
)
# Process documents in one batch parallelly
await asyncio.gather(*doc_tasks)
await self._insert_done()
log_message = f"Completed batch {current_batch} of {total_batches}."
logger.info(log_message)
pipeline_status["latest_message"] = log_message
pipeline_status["history_messages"].append(log_message)
# Check if there's a pending request to process more documents (with lock) # Check if there's a pending request to process more documents (with lock)
has_pending_request = False has_pending_request = False
@@ -1040,7 +1048,7 @@ class LightRAG:
] ]
await asyncio.gather(*tasks) await asyncio.gather(*tasks)
log_message = "All Insert done" log_message = "In memory DB persist to disk"
logger.info(log_message) logger.info(log_message)
if pipeline_status is not None and pipeline_status_lock is not None: if pipeline_status is not None and pipeline_status_lock is not None:

12
lightrag_webui/src/App.tsx
View File

@@ -56,24 +56,24 @@ function App() {
return ( return (
<ThemeProvider> <ThemeProvider>
<TabVisibilityProvider> <TabVisibilityProvider>
<main className="flex h-screen w-screen overflow-x-hidden"> <main className="flex h-screen w-screen overflow-hidden">
<Tabs <Tabs
defaultValue={currentTab} defaultValue={currentTab}
className="!m-0 flex grow flex-col !p-0" className="!m-0 flex grow flex-col !p-0 overflow-hidden"
onValueChange={handleTabChange} onValueChange={handleTabChange}
> >
<SiteHeader /> <SiteHeader />
<div className="relative grow"> <div className="relative grow">
<TabsContent value="documents" className="absolute top-0 right-0 bottom-0 left-0"> <TabsContent value="documents" className="absolute top-0 right-0 bottom-0 left-0 overflow-auto">
<DocumentManager /> <DocumentManager />
</TabsContent> </TabsContent>
<TabsContent value="knowledge-graph" className="absolute top-0 right-0 bottom-0 left-0"> <TabsContent value="knowledge-graph" className="absolute top-0 right-0 bottom-0 left-0 overflow-hidden">
<GraphViewer /> <GraphViewer />
</TabsContent> </TabsContent>
<TabsContent value="retrieval" className="absolute top-0 right-0 bottom-0 left-0"> <TabsContent value="retrieval" className="absolute top-0 right-0 bottom-0 left-0 overflow-hidden">
<RetrievalTesting /> <RetrievalTesting />
</TabsContent> </TabsContent>
<TabsContent value="api" className="absolute top-0 right-0 bottom-0 left-0"> <TabsContent value="api" className="absolute top-0 right-0 bottom-0 left-0 overflow-hidden">
<ApiSite /> <ApiSite />
</TabsContent> </TabsContent>
</div> </div>

2
lightrag_webui/src/features/GraphViewer.tsx
View File

@@ -182,7 +182,7 @@ const GraphViewer = () => {
// Always render SigmaContainer but control its visibility with CSS // Always render SigmaContainer but control its visibility with CSS
return ( return (
<div className="relative h-full w-full"> <div className="relative h-full w-full overflow-hidden">
<SigmaContainer <SigmaContainer
settings={sigmaSettings} settings={sigmaSettings}
className="!bg-background !size-full overflow-hidden" className="!bg-background !size-full overflow-hidden"

2
lightrag_webui/src/features/RetrievalTesting.tsx
View File

@@ -112,7 +112,7 @@ export default function RetrievalTesting() {
}, [setMessages]) }, [setMessages])
return ( return (
<div className="flex size-full gap-2 px-2 pb-12"> <div className="flex size-full gap-2 px-2 pb-12 overflow-hidden">
<div className="flex grow flex-col gap-4"> <div className="flex grow flex-col gap-4">
<div className="relative grow"> <div className="relative grow">
<div className="bg-primary-foreground/60 absolute inset-0 flex flex-col overflow-auto rounded-lg border p-2"> <div className="bg-primary-foreground/60 absolute inset-0 flex flex-col overflow-auto rounded-lg border p-2">