rag

RAG Skill

检索增强生成（Retrieval-Augmented Generation）能力，结合文档检索和LLM生成。

功能

• 文档摄入（PDF、DOCX、TXT、Markdown）
• 文本分块（Chunking）
• 向量嵌入生成
• 向量数据库管理
• 语义搜索
• 混合搜索（向量 + 关键词）
• 知识库管理

快速开始

1. 启动 Docker 服务

cd rag/docker && docker-compose up -d

启动的服务包括：

• Weaviate（向量数据库）
• Elasticsearch（关键词检索）
• Redis（缓存存储）
• nginx（代理）

2. 配置环境变量

确保 .env 文件包含以下配置：

# Elasticsearch（必须与 docker/.env 中配置一致）
ELASTICSEARCH_USER=elastic
ELASTICSEARCH_PASSWORD="difyai123456"

# Weaviate
WEAVIATE_API_KEY=jane-secret-key

# Redis
REDIS_PASSWD="difyai123456"

# Embedding 服务（用于向量检索和 RAG Pipeline）
EMBEDDING_MODEL=bge-large-zh-v1.5
EMBEDDING_API_URL=http://localhost:9997/v1
EMBEDDING_API_KEY=sk-your-api-key

# LLM 服务（用于 RAG Pipeline）
LLM_MODEL=glm-5
LLM_API_URL=https://api.example.com/v1
LLM_API_KEY=sk-your-api-key

3. 运行示例代码

import asyncio
from rag import (
    KnowledgeBaseConfig,
    kb_manager,
)

# 必须先导入以注册内置分块策略
from rag.extensions.custom_line_split import *  # noqa: F401

async def main():
    # 创建知识库（ELAS_ONLY 模式，仅使用 Elasticsearch）
    config = KnowledgeBaseConfig(
        kb_id="my_kb",
        kb_name="我的知识库",
        KNOWLEDGE_PATH="rag/test_document/",
        ELAS_ONLY=True,
    )
    kb = await kb_manager.create_knowledge_base(**config.to_dict())
    print("✅ 知识库创建成功")

    # 加载并添加文档
    await kb_manager.load_and_add_documents("my_kb", "rag/test_document/")
    print("✅ 文档加载成功")

    # 检索
    results = await kb_manager.retrieve("my_kb", "查询内容", k=5)
    print(f"✅ 检索返回 {len(results)} 条结果")

    # 清理
    await kb_manager.delete_knowledge_base("my_kb")
    print("✅ 知识库已删除")

asyncio.run(main())

验证功能

运行完整测试脚本：

uv run python tests/test_knowledge_base_complete.py

测试覆盖：

• 多知识库管理
• 文档处理和分块
• 向量存储连接
• 检索功能

检索模式

三种模式互斥，只能启用一种：

模式	说明	配置
ELAS_ONLY	仅关键词检索（Elasticsearch）	ELAS_ONLY=True
VECTOR_ONLY	仅向量检索（Weaviate）	VECTOR_ONLY=True
MIXED	混合检索	MIXED=True

ELAS_ONLY 模式（推荐入门）