git clone https://github.com/HKUDS/LightRAG.git
cd LightRAG
pip install -e ".[api]"

import os
import asyncio
from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
from lightrag.kg.shared_storage import initialize_pipeline_status
from lightrag.utils import setup_logger
# 设置日志
setup_logger("lightrag", level="INFO")
# 配置工作目录
WORKING_DIR = "./tech_docs_rag"
if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)
async def initialize_rag():
    """初始化RAG实例"""
    rag = LightRAG(
        working_dir=WORKING_DIR,
        # 使用OpenAI的嵌入模型
        embedding_func=openai_embed,
        # 使用GPT-4o-mini作为主要语言模型
        llm_model_func=gpt_4o_mini_complete,
        # 设置文档分块大小
        chunk_token_size=1200,
        chunk_overlap_token_size=100,
        # 启用LLM缓存以提高效率
        enable_llm_cache=True,
    )
    # 关键步骤：必须初始化存储和管道状态
    await rag.initialize_storages()
    await initialize_pipeline_status()
    return rag
async def main():
    try:
        # 初始化RAG实例
        rag = await initialize_rag()
        print("✅ LightRAG 初始化完成")
        # 插入文档内容
        # 这里可以是你的技术文档、API说明等
        document_content = """
        Docker是一个开源的容器化平台，它使用Linux内核的资源隔离功能，
        如cgroups和namespaces，来在一个操作系统内核上运行多个独立的容器。
        容器化技术的核心优势包括：
        1. 环境一致性：开发、测试、生产环境完全一致
        2. 资源高效：比虚拟机占用更少的系统资源
        3. 快速部署：秒级启动和停止
        4. 微服务架构支持：每个服务独立打包和部署
        Kubernetes是容器编排平台，它管理Docker容器集群，
        提供服务发现、负载均衡、自动扩缩容等功能。
        """
        await rag.ainsert(document_content)
        print("✅ 文档插入完成，知识图谱构建中...")
        # 等待图谱构建完成
        await asyncio.sleep(2)
        # 演示不同的查询模式
        queries = [
            "Docker的主要优势是什么？",
            "容器化技术如何支持微服务架构？",
            "Kubernetes和Docker之间是什么关系？"
        ]
        for query in queries:
            print(f"\n🔍 查询: {query}")
            # 混合模式查询（推荐）
            response = await rag.aquery(
                query,
                param=QueryParam(mode="hybrid")
            )
            print(f"📋 回答: {response}")
    except Exception as e:
        print(f"❌ 发生错误: {e}")
    finally:
        if 'rag' in locals():
            await rag.finalize_storages()
if __name__ == "__main__":
    # 记得设置你的OpenAI API密钥
    os.environ["OPENAI_API_KEY"] = "你的API密钥"
    asyncio.run(main())

from raganything import RAGAnything
async def setup_multimodal_rag():
    """配置多模态RAG系统"""
    # 基础LightRAG实例
    lightrag_instance = LightRAG(
        working_dir="./multimodal_storage",
        llm_model_func=gpt_4o_mini_complete,
        embedding_func=openai_embed,
    )
    await lightrag_instance.initialize_storages()
    await initialize_pipeline_status()
    # RAG-Anything实例，支持多模态处理
    rag = RAGAnything(
        lightrag=lightrag_instance,
        vision_model_func=gpt_4o_complete,  # 用于处理图像
    )
    # 处理PDF文档
    await rag.process_document_complete(
        file_path="./technical_manual.pdf",
        output_dir="./processed_docs"
    )
    return rag

async def setup_production_rag():
    """生产环境配置"""
    rag = LightRAG(
        working_dir=WORKING_DIR,
        # 使用PostgreSQL作为存储后端
        kv_storage="PGKVStorage",
        vector_storage="PGVectorStorage",
        graph_storage="PGGraphStorage",
        doc_status_storage="PGDocStatusStorage",
        # 优化性能参数
        chunk_token_size=1500,
        embedding_batch_size=64,
        llm_model_max_async=8,
        # 配置工作空间以实现多租户隔离
        workspace="production_env",
        llm_model_func=gpt_4o_mini_complete,
        embedding_func=openai_embed,
    )
    await rag.initialize_storages()
    await initialize_pipeline_status()
    return rag

# 针对法律文档优化的配置
legal_rag = LightRAG(
    working_dir="./legal_analysis",
    # 增加上下文长度以处理复杂法律文本
    chunk_token_size=2000,
    # 提取更多实体关系以捕获法律概念间的联系
    entity_extract_max_gleaning=3,
    # 专门的法律领域提示
    addon_params={
        "language": "Simplified Chinese",
        "entity_types": ["法条", "案例", "当事人", "法院", "判决"],
        "domain_specific": "legal"
    }
)

# 学术论文分析配置
academic_rag = LightRAG(
    working_dir="./research_analysis",
    # 设置更大的token限制以处理长文本
    llm_model_max_token_size=64000,
    # 优化实体提取以识别学术概念
    addon_params={
        "entity_types": ["研究方法", "理论", "学者", "机构", "实验", "结论"],
        "example_number": 3
    }
)
# 批量处理学术文献
papers = ["paper1.pdf", "paper2.pdf", "paper3.pdf"]
for paper in papers:
    with open(paper, 'r') as f:
        await academic_rag.ainsert(f.read(), ids=[f"paper_{papers.index(paper)}"])

# 事实性问题：使用local模式
fact_response = await rag.aquery(
    "张三的职位是什么？",
    param=QueryParam(mode="local")
)
# 概念性问题：使用global模式
concept_response = await rag.aquery(
    "这个行业的发展趋势如何？",
    param=QueryParam(mode="global")
)
# 复杂分析：使用hybrid模式
analysis_response = await rag.aquery(
    "基于现有数据，分析市场机会在哪里？",
    param=QueryParam(mode="hybrid", top_k=80)
)

rag = LightRAG(
    working_dir=WORKING_DIR,
    # 启用LLM缓存
    enable_llm_cache=True,
    # 启用实体提取缓存，方便调试
    enable_llm_cache_for_entity_extract=True,
    # 配置问答缓存
    embedding_cache_config={
        "enabled": True,
        "similarity_threshold": 0.85,
        "use_llm_check": True
    }
)

conversation_history = [
    {"role": "user", "content": "什么是容器化技术？"},
    {"role": "assistant", "content": "容器化技术是一种轻量级虚拟化技术..."},
]
# 带上下文的查询
response = await rag.aquery(
    "它相比传统虚拟机有什么优势？",
    param=QueryParam(
        mode="hybrid",
        conversation_history=conversation_history,
        history_turns=3  # 保留最近3轮对话
    )
)

class CustomerServiceRAG:
    def __init__(self):
        self.rag = None

    async def initialize(self):
        """初始化客服 RAG 系统"""
        self.rag = LightRAG(
            working_dir="./customer_service",
            # 针对客服场景优化
            chunk_token_size=800,  # 较小的块，提高响应速度
            embedding_batch_size=32,
            # 客服专用提示配置
            addon_params={
                "language": "Simplified Chinese",
                "entity_types": ["产品", "功能", "问题", "解决方案", "客户"],
                "domain_specific": "customer_service"
            }
        )
        await self.rag.initialize_storages()
        await initialize_pipeline_status()

    async def process_knowledge_base(self):
        """处理知识库文档"""
        documents = {
            "产品手册": "./docs/product_manual.txt",
            "故障排除指南": "./docs/troubleshooting.txt",
            "API文档": "./docs/api_reference.txt"
        }
        for doc_type, file_path in documents.items():
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
                await self.rag.ainsert(
                    content,
                    ids=[f"{doc_type}_{hash(content)}"]
                )
        print("✅ 知识库构建完成")

    async def handle_customer_query(self, question, conversation_context=None):
        """处理客户咨询"""
        query_param = QueryParam(
            mode="mix",  # 使用混合模式，平衡精确度和覆盖度
            response_type="Multiple Paragraphs",
            top_k=50,
            conversation_history=conversation_context or []
        )
        response = await self.rag.aquery(question, param=query_param)
        return response


# 使用示例
async def run_customer_service():
    cs_rag = CustomerServiceRAG()
    await cs_rag.initialize()
    await cs_rag.process_knowledge_base()

    # 模拟客户咨询
    response = await cs_rag.handle_customer_query(
        "我的产品无法连接数据库，可能是什么原因？"
    )
    print(f"客服回答：{response}")


if __name__ == "__main__":
    asyncio.run(run_customer_service())

# 创建新的实体
entity = await rag.acreate_entity("微服务架构", {
    "description": "微服务架构是一种软件架构模式，将应用程序构建为独立服务的集合",
    "entity_type": "architecture_pattern"
})
# 建立实体关系
relation = await rag.acreate_relation("微服务架构", "Docker", {
    "description": "Docker是实现微服务架构的重要容器化工具",
    "keywords": "容器化 部署 隔离",
    "weight": 2.5
})
# 动态更新实体信息
updated_entity = await rag.aedit_entity("Docker", {
    "description": "Docker是领先的容器化平台，广泛用于微服务部署和DevOps实践",
    "entity_type": "containerization_platform"
})

# 合并重复实体
await rag.amerge_entities(
    source_entities=["容器技术", "容器化", "Container 技术"],
    target_entity="容器化技术",
    merge_strategy={
        "description": "concatenate",  # 合并描述
        "entity_type": "keep_first",  # 保留第一个实体类型
        "source_id": "join_unique"  # 合并来源 ID
    }
)

# 本地Ollama部署示例
from lightrag.llm.ollama import ollama_model_complete, ollama_embed
rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=ollama_model_complete,
    llm_model_name='qwen2:32b',  # 使用通义千问模型
    llm_model_kwargs={"options": {"num_ctx": 32768}},  # 设置上下文长度
    embedding_func=EmbeddingFunc(
        embedding_dim=768,
        max_token_size=8192,
        func=lambda texts: ollama_embed(texts, embed_model="nomic-embed-text")
    )
)

# 性能监控示例
async def monitor_rag_performance():
    """RAG系统性能监控"""
    # 查询统计
    start_time = time.time()
    response = await rag.aquery("测试查询", param=QueryParam(mode="hybrid"))
    query_time = time.time() - start_time
    print(f"查询耗时: {query_time:.2f}秒")
    # 缓存命中率（可以通过日志分析）
    cache_stats = rag.get_cache_statistics()  # 假设的API
    print(f"缓存命中率: {cache_stats.hit_rate:.2%}")

# 清理数据目录，但保留 LLM 缓存
rm -rf ./rag_storage/vector_db/
# 保留 kv_store_llm_response_cache.json

# 优化内存使用
rag = LightRAG(
    working_dir=WORKING_DIR,
    # 减小批处理大小
    embedding_batch_size=16,
    max_parallel_insert=2,
    # 控制token使用
    chunk_token_size=800,
    llm_model_max_token_size=16000
)

# 使用重排序模型优化
from lightrag.llm.reranker import BGERerankerFunc
rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=gpt_4o_mini_complete,
    embedding_func=openai_embed,
    # 启用重排序
    enable_rerank=True,
    rerank_func=BGERerankerFunc()  # 使用BGE重排序模型
)

# 启动 Web 服务
cp env.example .env
# 在.env 中配置你的 API 密钥和模型设置
lightrag-server