专业的编程技术博客社区

网站首页 > 博客文章 正文

LightRAG: 简单快速的检索增强生成工具

baijin 2025-06-23 14:46:21 博客文章 2 ℃ 0 评论

这里是Aideas,每日分享AI相关资讯。本文由Aideas Agent整理并推荐。项目地址:/HKUDS/LightRAG, 程序语言:Python, 收藏: 14,287, 分支: 1,996, 今日收藏: 168 stars today。

LightRAG 服务器旨在提供 Web UI 和 API 支持。Web UI 有助于文档索引、知识图谱探索和简单的 RAG 查询界面。LightRAG 服务器还提供了一个兼容 Ollama 的接口,旨在模拟 LightRAG 作为 Ollama 聊天模型。这允许 AI 聊天机器人(例如 Open WebUI)轻松访问 LightRAG。


安装 LightRAG 核心

  • 从源代码安装
cd LightRAG
pip install -e .
  • 从 PyPI 安装
pip install lightrag-hku

安装 LightRAG 服务器

LightRAG 服务器旨在提供 Web UI 和 API 支持,方便文档索引、知识图谱探索和简单的 RAG 查询接口。LightRAG 服务器还提供 Ollama 兼容接口,旨在将 LightRAG 模拟为 Ollama 聊天模型。

  • 从 PyPI 安装
pip install "lightrag-hku[api]"
  • 从源代码安装
# 如果需要,创建 Python 虚拟环境
# 以可编辑模式安装并支持 API
pip install -e ".[api]"

快速开始

查询

初始化 LightRAG 并执行查询:

import os
import asyncio
from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import gpt_4o_mini_complete, gpt_4o_complete, openai_embed
from lightrag.kg.shared_storage import initialize_pipeline_status
from lightrag.utils import setup_logger

setup_logger("lightrag", level="INFO")

async def initialize_rag():
    rag = LightRAG(
        working_dir="your/path",
        embedding_func=openai_embed,
        llm_model_func=gpt_4o_mini_complete
    )

    await rag.initialize_storages()
    await initialize_pipeline_status()

    return rag

def main():
    # Initialize RAG instance
    rag = asyncio.run(initialize_rag())
    # Insert text
    rag.insert("Your text")

    # Perform naive search
    mode="naive"
    # Perform local search
    mode="local"
    # Perform global search
    mode="global"
    # Perform hybrid search
    mode="hybrid"
    # Mix mode Integrates knowledge graph and vector retrieval.
    mode="mix"

    rag.query(
        "What are the top themes in this story?",
        param=QueryParam(mode=mode)
    )

if __name__ == "__main__":
    main()

查询参数

class QueryParam:
    mode: Literal["local", "global", "hybrid", "naive", "mix"] = "global"
    """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. Mix mode combines knowledge graph and vector search:
        - Uses both structured (KG) and unstructured (vector) information
        - Provides comprehensive answers by analyzing relationships and context
        - Supports image content through HTML img tags
        - Allows control over retrieval depth via top_k parameter
    """
    only_need_context: bool = False
    """If True, only returns the retrieved context without generating a response."""
    response_type: str = "Multiple Paragraphs"
    """Defines the response format. Examples: 'Multiple Paragraphs', 'Single Paragraph', 'Bullet Points'."""
    top_k: int = 60
    """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
    max_token_for_text_unit: int = 4000
    """Maximum number of tokens allowed for each retrieved text chunk."""
    max_token_for_global_context: int = 4000
    """Maximum number of tokens allocated for relationship descriptions in global retrieval."""
    max_token_for_local_context: int = 4000
    """Maximum number of tokens allocated for entity descriptions in local retrieval."""
    ids: list[str] | None = None # ONLY SUPPORTED FOR PG VECTOR DBs
    """List of ids to filter the RAG."""
    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.
    """
    ...

默认值 Top_k 可以通过环境变量TOP_K来更改。

LLM 和嵌入注入

LightRAG 需要使用 LLM 和 Embedding 模型来完成文档索引和查询任务。在初始化阶段,需要将相关模型的调用方法注入到 LightRAG 中:

使用类似 Open AI 的 API

  • LightRAG 还支持类似 Open AI 的聊天/嵌入 API:
async def llm_model_func(
    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
    return await openai_complete_if_cache(
        "solar-mini",
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        api_key=os.getenv("UPSTAGE_API_KEY"),
        base_url="",
        **kwargs
    )

async def embedding_func(texts: list[str]) -> np.ndarray:
    return await openai_embed(
        texts,
        model="solar-embedding-1-large-query",
        api_key=os.getenv("UPSTAGE_API_KEY"),
        base_url=""
    )

async def initialize_rag():
    rag = LightRAG(
        working_dir=WORKING_DIR,
        llm_model_func=llm_model_func,
        embedding_func=EmbeddingFunc(
            embedding_dim=4096,
            max_token_size=8192,
            func=embedding_func
        )
    )

    await rag.initialize_storages()
    await initialize_pipeline_status()

    return rag

使用 Hugging Face 模型

# Initialize LightRAG with Hugging Face model
rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=hf_model_complete,  # Use Hugging Face model for text generation
    llm_model_name='meta-llama/Llama-3.1-8B-Instruct',  # Model name from Hugging Face
    # Use Hugging Face embedding function
    embedding_func=EmbeddingFunc(
        embedding_dim=384,
        max_token_size=5000,
        func=lambda texts: hf_embed(
            texts,
            tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"),
            embed_model=AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        )
    ),
)

使用 Ollama 模型

# Initialize LightRAG with Ollama model
rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=ollama_model_complete,  # Use Ollama model for text generation
    llm_model_name='your_model_name', # Your model name
    # Use Ollama embedding function
    embedding_func=EmbeddingFunc(
        embedding_dim=768,
        max_token_size=8192,
        func=lambda texts: ollama_embed(
            texts,
            embed_model="nomic-embed-text"
        )
    ),
)
  • 增加上下文大小

为了使 LightRAG 正常工作,上下文应至少为 32k 令牌。

  • 在 Modelfile 中增加 num_ctx 参数
  1. 拉取模型:
ollama pull qwen2
  1. 显示模型文件:
ollama show --modelfile qwen2 > Modelfile
  1. 通过添加以下行来编辑 Modelfile:
PARAMETER num_ctx 32768
  1. 创建修改后的模型:
ollama create -f Modelfile qwen2m
  • 通过 Ollama API 设置 num_ctx

可以使用 param 来配置 ollama:

rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=ollama_model_complete,  # Use Ollama model for text generation
    llm_model_name='your_model_name', # Your model name
    llm_model_kwargs={"options": {"num_ctx": 32768}},
    # Use Ollama embedding function
    embedding_func=EmbeddingFunc(
        embedding_dim=768,
        max_token_size=8192,
        func=lambda texts: ollama_embedding(
            texts,
            embed_model="nomic-embed-text"
        )
    ),
)

令牌使用跟踪

概述和用法

LightRAG 提供了一个 TokenTracker 工具,用于监控和管理大型语言模型的令牌消耗。此功能对于控制 API 成本和优化性能特别有用。

用法

from lightrag.utils import TokenTracker

# Create TokenTracker instance
token_tracker = TokenTracker()

# Method 1: Using context manager (Recommended)
# Suitable for scenarios requiring automatic token usage tracking
with token_tracker:
    result1 = await llm_model_func("your question 1")
    result2 = await llm_model_func("your question 2")

# Method 2: Manually adding token usage records
# Suitable for scenarios requiring more granular control over token statistics
token_tracker.reset()

rag.insert()

rag.query("your question 1", param=QueryParam(mode="naive"))
rag.query("your question 2", param=QueryParam(mode="mix"))

# Display total token usage (including insert and query operations)
print("Token usage:", token_tracker.get_usage())

使用技巧

  • 使用上下文管理器进行长时间会话或批处理作,以自动跟踪所有令牌消耗
  • 对于需要分段统计信息的场景,请使用手动模式并在适当时调用 reset()
  • 定期检查 Token 使用情况有助于及早发现异常消耗
  • 在开发和测试期间积极使用此功能以优化生产成本

对话历史记录支持

LightRAG 通过对话历史记录功能支持多回合对话。

# Create conversation history
conversation_history = [
    {"role": "user", "content": "What is the main character's attitude towards Christmas?"},
    {"role": "assistant", "content": "At the beginning of the story, Ebenezer Scrooge has a very negative attitude towards Christmas..."},
    {"role": "user", "content": "How does his attitude change?"}
]

# Create query parameters with conversation history
query_param = QueryParam(
    mode="mix",  # or any other mode: "local", "global", "hybrid"
    conversation_history=conversation_history,  # Add the conversation history
    history_turns=3  # Number of recent conversation turns to consider
)

# Make a query that takes into account the conversation history
response = rag.query(
    "What causes this change in his character?",
    param=query_param
)

自定义提示支持

LightRAG 现在支持自定义提示,以便对系统行为进行微调控制。以下是如何使用它:

使用示例

# Create query parameters
query_param = QueryParam(
    mode="hybrid",  # or other mode: "local", "global", "hybrid", "mix" and "naive"
)

# Example 1: Using the default system prompt
response_default = rag.query(
    "What are the primary benefits of renewable energy?",
    param=query_param
)
print(response_default)

# Example 2: Using a custom prompt
custom_prompt = """
You are an expert assistant in environmental science. Provide detailed and structured answers with examples.
---Conversation History---
{history}

---Knowledge Base---
{context_data}

---Response Rules---

- Target format and length: {response_type}
"""
response_custom = rag.query(
    "What are the primary benefits of renewable energy?",
    param=query_param,
    system_prompt=custom_prompt  # Pass the custom prompt
)
print(response_custom)

插入

基本插件

# Basic Insert
rag.insert("Text")

批量插入

# Basic Batch Insert: Insert multiple texts at once
rag.insert(["TEXT1", "TEXT2",...])

# Batch Insert with custom batch size configuration
rag = LightRAG(
    working_dir=WORKING_DIR,
    addon_params={
        "insert_batch_size": 4  # Process 4 documents per batch
    }
)

rag.insert(["TEXT1", "TEXT2", "TEXT3", ...])  # Documents will be processed in batches of 4

使用 ID 插入

# Insert single text, and provide ID for it
rag.insert("TEXT1", ids=["ID_FOR_TEXT1"])

# Insert multiple texts, and provide IDs for them
rag.insert(["TEXT1", "TEXT2",...], ids=["ID_FOR_TEXT1", "ID_FOR_TEXT2"])

使用 Pipeline 插入

rag = LightRAG(..)

await rag.apipeline_enqueue_documents(input)
# Your routine in loop
await rag.apipeline_process_enqueue_documents(input)

插入多文件类型支持

import textract

file_path = 'TEXT.pdf'
text_content = textract.process(file_path)

rag.insert(text_content.decode('utf-8'))

插入自定义 KG

custom_kg = {
    "chunks": [
        {
            "content": "Alice and Bob are collaborating on quantum computing research.",
            "source_id": "doc-1"
        }
    ],
    "entities": [
        {
            "entity_name": "Alice",
            "entity_type": "person",
            "description": "Alice is a researcher specializing in quantum physics.",
            "source_id": "doc-1"
        },
        {
            "entity_name": "Bob",
            "entity_type": "person",
            "description": "Bob is a mathematician.",
            "source_id": "doc-1"
        },
        {
            "entity_name": "Quantum Computing",
            "entity_type": "technology",
            "description": "Quantum computing utilizes quantum mechanical phenomena for computation.",
            "source_id": "doc-1"
        }
    ],
    "relationships": [
        {
            "src_id": "Alice",
            "tgt_id": "Bob",
            "description": "Alice and Bob are research partners.",
            "keywords": "collaboration research",
            "weight": 1.0,
            "source_id": "doc-1"
        },
        {
            "src_id": "Alice",
            "tgt_id": "Quantum Computing",
            "description": "Alice conducts research on quantum computing.",
            "keywords": "research expertise",
            "weight": 1.0,
            "source_id": "doc-1"
        },
        {
            "src_id": "Bob",
            "tgt_id": "Quantum Computing",
            "description": "Bob researches quantum computing.",
            "keywords": "research application",
            "weight": 1.0,
            "source_id": "doc-1"
        }
    ]
}

rag.insert_custom_kg(custom_kg)

引文功能

# Define documents and their file paths
documents = ["Document content 1", "Document content 2"]
file_paths = ["path/to/doc1.txt", "path/to/doc2.txt"]

# Insert documents with file paths
rag.insert(documents, file_paths=file_paths)

存储

使用 Neo4J 进行存储

export NEO4J_URI="neo4j://localhost:7687"
export NEO4J_USERNAME="neo4j"
export NEO4J_PASSWORD="password"

# Setup logger for LightRAG
setup_logger("lightrag", level="INFO")

# When you launch the project be sure to override the default KG: NetworkX
# by specifying kg="Neo4JStorage".

# Note: Default settings use NetworkX
# Initialize LightRAG with Neo4J implementation.
async def initialize_rag():
    rag = LightRAG(
        working_dir=WORKING_DIR,
        llm_model_func=gpt_4o_mini_complete,  # Use gpt_4o_mini_complete LLM model
        graph_storage="Neo4JStorage", #<-----------override KG default
    )

    # Initialize database connections
    await rag.initialize_storages()
    # Initialize pipeline status for document processing
    await initialize_pipeline_status()

    return rag

使用 Faiss 进行存储

  • 安装所需的依赖项:
pip install faiss-cpu
async def embedding_func(texts: list[str]) -> np.ndarray:
    model = SentenceTransformer('all-MiniLM-L6-v2')
    embeddings = model.encode(texts, convert_to_numpy=True)
    return embeddings

# Initialize LightRAG with the LLM model function and embedding function
rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=llm_model_func,
    embedding_func=EmbeddingFunc(
        embedding_dim=384,
        max_token_size=8192,
        func=embedding_func,
    ),
    vector_storage="FaissVectorDBStorage",
    vector_db_storage_cls_kwargs={
        "cosine_better_than_threshold": 0.3  # Your desired threshold
    }
)

删除

#  Delete Entity: Deleting entities by their names
rag.delete_by_entity("Project Gutenberg")

#  Delete Document: Deleting entities and relationships associated with the document by doc id
rag.delete_by_doc_id("doc_id")

编辑实体和关系

LightRAG 支持全面的知识图谱管理功能,允许创建、编辑和删除知识图谱中的实体和关系。

创建实体和关系

# Create new entity
entity = rag.create_entity("Google", {
    "description": "Google is a multinational technology company specializing in internet-related services and products.",
    "entity_type": "company"
})

# Create another entity
product = rag.create_entity("Gmail", {
    "description": "Gmail is an email service developed by Google.",
    "entity_type": "product"
})

# Create relation between entities
relation = rag.create_relation("Google", "Gmail", {
    "description": "Google develops and operates Gmail.",
    "keywords": "develops operates service",
    "weight": 2.0
})

编辑实体和关系

# Edit an existing entity
updated_entity = rag.edit_entity("Google", {
    "description": "Google is a subsidiary of Alphabet Inc., founded in 1998.",
    "entity_type": "tech_company"
})

# Rename an entity (with all its relationships properly migrated)
renamed_entity = rag.edit_entity("Gmail", {
    "entity_name": "Google Mail",
    "description": "Google Mail (formerly Gmail) is an email service."
})

# Edit a relation between entities
updated_relation = rag.edit_relation("Google", "Google Mail", {
    "description": "Google created and maintains Google Mail service.",
    "keywords": "creates maintains email service",
    "weight": 3.0
})

实体

  • create_entity:创建具有指定属性的新实体
  • edit_entity:更新现有实体的属性或重命名它

关系

  • create_relation:在现有实体之间创建新关系
  • edit_relation:更新现有关系的属性

这些可保持图形数据库和矢量数据库组件之间的数据一致性,从而确保知识图谱保持连贯性。

数据导出函数

LightRAG 以各种格式导出知识图谱数据,以便进行分析、共享和备份。系统支持导出实体、关系和关系数据。

导出函数

基本用法

# Basic CSV export (default format)
rag.export_data("knowledge_graph.csv")

# Specify any format
rag.export_data("output.xlsx", file_format="excel")

支持不同的文件格式

#Export data in CSV format
rag.export_data("graph_data.csv", file_format="csv")

# Export data in Excel sheet
rag.export_data("graph_data.xlsx", file_format="excel")

# Export data in markdown format
rag.export_data("graph_data.md", file_format="md")

# Export data in Text
rag.export_data("graph_data.txt", file_format="txt")

其他选项

在导出中包含向量嵌入:

rag.export_data("complete_data.csv", include_vector_data=True)

导出中包含的数据

所有导出包括:

  • 实体信息(名称、ID、元数据)
  • 关系数据(实体之间的连接)
  • 来自向量数据库的关系信息

实体合并

合并实体及其关系

LightRAG 支持将多个实体合并为一个实体,自动处理所有关系:

# Basic entity merging
rag.merge_entities(
    source_entities=["Artificial Intelligence", "AI", "Machine Intelligence"],
    target_entity="AI Technology"
)

使用自定义合并策略:

# Define custom merge strategy for different fields
rag.merge_entities(
    source_entities=["John Smith", "Dr. Smith", "J. Smith"],
    target_entity="John Smith",
    merge_strategy={
        "description": "concatenate",  # Combine all descriptions
        "entity_type": "keep_first",   # Keep the entity type from the first entity
        "source_id": "join_unique"     # Combine all unique source IDs
    }
)

使用自定义目标实体数据:

# Specify exact values for the merged entity
rag.merge_entities(
    source_entities=["New York", "NYC", "Big Apple"],
    target_entity="New York City",
    target_entity_data={
        "entity_type": "LOCATION",
        "description": "New York City is the most populous city in the United States.",
    }
)

结合这两种方法的高级用法:

# Merge company entities with both strategy and custom data
rag.merge_entities(
    source_entities=["Microsoft Corp", "Microsoft Corporation", "MSFT"],
    target_entity="Microsoft",
    merge_strategy={
        "description": "concatenate",  # Combine all descriptions
        "source_id": "join_unique"     # Combine source IDs
    },
    target_entity_data={
        "entity_type": "ORGANIZATION",
    }
)

合并实体时:

  • 源实体中的所有关系都将重定向到目标实体
  • 重复的关系被智能地合并
  • 阻止自关系 (循环)
  • 合并后删除源实体
  • 保留关系权重和属性

LightRAG 应用程序接口

LightRAG 服务器旨在提供 Web UI 和 API 支持。

图形可视化

LightRAG 服务器提供了全面的知识图谱可视化功能。它支持各种重力布局、节点查询、子图过滤等。


LightRAG 服务器和 WebUI

LightRAG 服务器旨在提供 Web UI 和 API 支持。Web UI 有助于文档索引、知识图谱探索和简单的 RAG 查询界面。LightRAG 服务器还提供了一个兼容 Ollama 的接口,旨在模拟 LightRAG 作为 Ollama 聊天模型。这允许 AI 聊天机器人(例如 Open WebUI)轻松访问 LightRAG。

开始使用



安装

  • 从 PyPI 安装
pip install "lightrag-hku[api]"



启动 LightRAG 服务器之前

LightRAG 需要集成 LLM(大型语言模型)和嵌入模型,以有效地执行文档索引和查询作。在初始部署 LightRAG 服务器之前,必须配置 LLM 和嵌入模型的设置。LightRAG 支持绑定到各种 LLM/Embedding 后端:

  • ollama
  • lollms
  • openai or openai compatible
  • azure_openai

以下是 LLM 和 Embedding 模型的设置:

  • OpenAI LLM + Ollama 嵌入
LLM_BINDING=openai
LLM_MODEL=gpt-4o
LLM_BINDING_HOST=
LLM_BINDING_API_KEY=your_api_key
### Max tokens send to LLM (less than model context size)
MAX_TOKENS=32768

EMBEDDING_BINDING=ollama
EMBEDDING_BINDING_HOST=http://localhost:11434
EMBEDDING_MODEL=bge-m3:latest
EMBEDDING_DIM=1024
# EMBEDDING_BINDING_API_KEY=your_api_key
  • Ollama LLM + Ollama 嵌入
LLM_BINDING=ollama
LLM_MODEL=mistral-nemo:latest
LLM_BINDING_HOST=http://localhost:11434
# LLM_BINDING_API_KEY=your_api_key
### Max tokens send to LLM (base on your Ollama Server capacity)
MAX_TOKENS=8192

EMBEDDING_BINDING=ollama
EMBEDDING_BINDING_HOST=http://localhost:11434
EMBEDDING_MODEL=bge-m3:latest
EMBEDDING_DIM=1024
# EMBEDDING_BINDING_API_KEY=your_api_key



启动 LightRAG 服务器

LightRAG 服务器支持两种作模式:

  • 简单高效的 Uvicorn 模式
lightrag-server
  • 多进程 Gunicorn + Uvicorn 模式
lightrag-gunicorn --workers 4

#AI软件技巧#

Tags:

本文暂时没有评论,来添加一个吧(●'◡'●)

欢迎 发表评论:

最近发表
标签列表