# RAG功能集成迁移指南

## 概述

本文档记录了为 Open-LLM-VTuber 项目集成 RAG（检索增强生成）功能的所有修改，包括新增文件、修改的文件和配置变更。此指南将帮助您将这些修改迁移到服务器上的相同项目。

## 功能说明

RAG 功能允许数字人基于外部知识库回答问题，提供更准确和相关的信息。系统支持：
- 文本块的智能分割和向量化存储
- 基于语义相似度的知识检索
- 多种嵌入模型提供商（支持千帆 OpenAI 兼容接口）
- REST API 管理知识库
- 自动将检索到的上下文注入到 LLM 提示中

## 文件修改清单

### 1. 新增文件

#### 1.1 RAG 核心模块
```
src/open_llm_vtuber/rag/
├── __init__.py                 # RAG 模块初始化
└── rag_engine.py              # RAG 引擎核心实现
```

#### 1.2 配置模块
```
src/open_llm_vtuber/config_manager/rag.py  # RAG 配置模型
```

#### 1.3 测试样例（可选）
```
doc/rag_samples/
├── faq_基础使用.txt            # 基础使用FAQ
├── 常见问题汇总.md             # 常见问题汇总
└── product_brief.json         # 产品简介示例
```

### 2. 修改的文件

#### 2.1 配置文件
- **`conf.yaml`** - 主配置文件
- **`src/open_llm_vtuber/config_manager/__init__.py`** - 配置管理器初始化
- **`src/open_llm_vtuber/config_manager/character.py`** - 角色配置模型

#### 2.2 核心功能文件  
- **`src/open_llm_vtuber/service_context.py`** - 服务上下文管理
- **`src/open_llm_vtuber/conversations/single_conversation.py`** - 单次对话处理
- **`src/open_llm_vtuber/routes.py`** - API 路由定义

## 详细修改内容

### 1. 新增文件内容

#### 1.1 `src/open_llm_vtuber/rag/__init__.py`
```python
from .rag_engine import RAGEngine, EmbeddingProviderConfig

__all__ = ["RAGEngine", "EmbeddingProviderConfig"]
```

#### 1.2 `src/open_llm_vtuber/rag/rag_engine.py`
```python
import os
import json
import numpy as np
import aiohttp
from dataclasses import dataclass
from typing import List, Dict, Optional, Tuple
from openai import AsyncOpenAI
from loguru import logger

@dataclass
class EmbeddingProviderConfig:
    base_url: str
    api_key: str
    model: str

def _cosine_similarity_matrix(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """计算余弦相似度矩阵"""
    a_norm = np.linalg.norm(a, axis=1, keepdims=True)
    b_norm = np.linalg.norm(b, axis=1, keepdims=True)
    return np.dot(a / a_norm, (b / b_norm).T)

class RAGEngine:
    """RAG 引擎：负责文档存储、检索和上下文生成"""
    
    def __init__(
        self,
        index_dir: str = "rag_index",
        provider: EmbeddingProviderConfig = None,
        top_k: int = 4,
        max_context_chars: int = 2000,
        chunk_chars: int = 800,
        chunk_overlap: int = 100,
    ):
        self.index_dir = index_dir
        self.provider = provider
        self.top_k = top_k
        self.max_context_chars = max_context_chars
        self.chunk_chars = chunk_chars
        self.chunk_overlap = chunk_overlap

        os.makedirs(self.index_dir, exist_ok=True)
        self.meta_path = os.path.join(self.index_dir, "meta.json")
        self.vec_path = os.path.join(self.index_dir, "embeddings.npy")
        self._client = AsyncOpenAI(base_url=self.provider.base_url, api_key=self.provider.api_key)

        self._texts: List[str] = []
        self._metas: List[Dict] = []
        self._embeds: Optional[np.ndarray] = None
        self._load_index()

    def _load_index(self) -> None:
        """加载索引文件"""
        try:
            if os.path.exists(self.meta_path):
                with open(self.meta_path, "r", encoding="utf-8") as f:
                    stored = json.load(f)
                self._texts = stored.get("texts", [])
                self._metas = stored.get("metas", [])
            if os.path.exists(self.vec_path):
                self._embeds = np.load(self.vec_path)
            logger.info(f"RAG index loaded: {len(self._texts)} chunks.")
        except Exception as e:
            logger.error(f"Failed to load RAG index: {e}")
            self._texts, self._metas, self._embeds = [], [], None

    def _save_index(self) -> None:
        """保存索引文件"""
        try:
            with open(self.meta_path, "w", encoding="utf-8") as f:
                json.dump({"texts": self._texts, "metas": self._metas}, f, ensure_ascii=False)
            if self._embeds is not None:
                np.save(self.vec_path, self._embeds)
        except Exception as e:
            logger.error(f"Failed to save RAG index: {e}")

    def _chunk_text(self, text: str) -> List[str]:
        """将文本分割为块"""
        if len(text) <= self.chunk_chars:
            return [text]
        chunks = []
        for start in range(0, len(text), self.chunk_chars - self.chunk_overlap):
            end = start + self.chunk_chars
            chunk = text[start:end]
            chunks.append(chunk)
            if end >= len(text):
                break
        return chunks

    async def _embed_texts(self, texts: List[str]) -> np.ndarray:
        """生成文本向量"""
        if not texts:
            return np.zeros((0, 0), dtype=np.float32)
        
        # 如果是千帆原生 API（bce-v3 鉴权）
        if self.provider.api_key.startswith("bce-v3/"):
            return await self._embed_texts_qianfan(texts)
        else:
            # OpenAI 兼容 API
            resp = await self._client.embeddings.create(model=self.provider.model, input=texts)
            vecs = [d.embedding for d in resp.data]
            return np.array(vecs, dtype=np.float32)
    
    async def _embed_texts_qianfan(self, texts: List[str]) -> np.ndarray:
        """使用千帆原生 API 进行 embedding"""
        url = "https://qianfan.baidubce.com/v2/embeddings"
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.provider.api_key}'
        }
        payload = {
            "model": self.provider.model,
            "input": texts
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=headers, json=payload) as resp:
                if resp.status != 200:
                    error_text = await resp.text()
                    raise Exception(f"千帆 API 错误: {resp.status} - {error_text}")
                result = await resp.json()
                vecs = [d["embedding"] for d in result["data"]]
                return np.array(vecs, dtype=np.float32)

    async def add_texts(self, texts: List[str], source: str = "api") -> int:
        """添加文本到索引"""
        chunks: List[str] = []
        metas: List[Dict] = []
        for t in texts:
            for ch in self._chunk_text(t):
                chunks.append(ch)
                metas.append({"source": source})
        if not chunks:
            return 0
        embeds = await self._embed_texts(chunks)
        if self._embeds is None or self._embeds.size == 0:
            self._embeds = embeds
        else:
            self._embeds = np.vstack([self._embeds, embeds])
        self._texts.extend(chunks)
        self._metas.extend(metas)
        self._save_index()
        return len(chunks)

    async def add_file(self, filepath: str, source: str = "file") -> int:
        """添加文件到索引"""
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                content = f.read().strip()
            if content:
                return await self.add_texts([content], source)
            return 0
        except Exception as e:
            logger.error(f"Failed to read file {filepath}: {e}")
            return 0

    async def clear(self) -> None:
        """清空索引"""
        self._texts, self._metas = [], []
        self._embeds = None
        try:
            if os.path.exists(self.meta_path):
                os.remove(self.meta_path)
            if os.path.exists(self.vec_path):
                os.remove(self.vec_path)
        except Exception as e:
            logger.error(f"Failed to clear RAG index files: {e}")

    def stats(self) -> Dict:
        """获取索引统计信息"""
        dim = int(self._embeds.shape[1]) if (self._embeds is not None and self._embeds.size) else 0
        return {"chunks": len(self._texts), "dim": dim}

    async def search(self, query: str, top_k: Optional[int] = None) -> Tuple[str, List[Dict]]:
        """搜索相关文档"""
        if not query or self._embeds is None or not self._texts:
            return "", []
        qv = await self._embed_texts([query])
        sims = _cosine_similarity_matrix(qv, self._embeds)[0]
        k = min(top_k or self.top_k, len(self._texts))
        idx = np.argsort(-sims)[:k]
        ordered = [(int(i), float(sims[i])) for i in idx]

        ctx_parts: List[str] = []
        ctx_metas: List[Dict] = []
        acc_len = 0
        for i, score in ordered:
            text = self._texts[i]
            meta = self._metas[i] if i < len(self._metas) else {}
            part = f"[Doc#{i} sim={score:.3f}]\n{text}"
            if acc_len + len(part) > self.max_context_chars and ctx_parts:
                break
            ctx_parts.append(part)
            ctx_metas.append({"index": i, "score": score, **meta})
            acc_len += len(part)
        return "\n\n".join(ctx_parts), ctx_metas
```

#### 1.3 `src/open_llm_vtuber/config_manager/rag.py`
```python
from pydantic import BaseModel, Field
from typing import ClassVar, Optional, Literal
from ..i18n import I18nMixin, Description

class RAGConfig(I18nMixin, BaseModel):
    """Configuration for Retrieval-Augmented Generation (RAG)."""

    enabled: bool = Field(False, alias="enabled")
    index_dir: str = Field("rag_index", alias="index_dir")
    embed_provider_key: str = Field(..., alias="embed_provider_key")
    embed_model: str = Field(..., alias="embed_model")
    top_k: int = Field(4, alias="top_k")
    max_context_chars: int = Field(2000, alias="max_context_chars")
    chunk_chars: int = Field(800, alias="chunk_chars")
    chunk_overlap: int = Field(100, alias="chunk_overlap")

    DESCRIPTIONS: ClassVar[dict[str, Description]] = {
        "enabled": Description(en="Enable RAG for this character", zh="为该角色启用RAG"),
        "index_dir": Description(en="Directory to store RAG index files", zh="RAG索引文件存储目录"),
        "embed_provider_key": Description(
            en="Key of the LLM provider in llm_configs to use for embeddings",
            zh="llm_configs中用于嵌入的LLM提供者键名",
        ),
        "embed_model": Description(en="Embedding model name", zh="嵌入模型名称"),
        "top_k": Description(en="Number of top relevant chunks to retrieve", zh="检索最相关的Top K块"),
        "max_context_chars": Description(
            en="Maximum characters for the retrieved context to inject into prompt",
            zh="注入到提示词中的最大检索上下文字符数",
        ),
        "chunk_chars": Description(en="Character length for each text chunk", zh="每个文本块的字符长度"),
        "chunk_overlap": Description(en="Overlap characters between chunks", zh="文本块之间的重叠字符数"),
    }
```

### 2. 文件修改内容

#### 2.1 `conf.yaml` 修改
在 `character_config` 部分添加：
```yaml
character_config:
  # ... 现有配置 ...
  
  # RAG 配置 (新增)
  rag_config:
    enabled: true
    index_dir: rag_index
    embed_provider_key: openai_compatible_llm
    embed_model: embedding-v1
    top_k: 4
    max_context_chars: 2000
    chunk_chars: 800
    chunk_overlap: 100
```

同时确保 `llm_configs` 中的千帆配置正确：
```yaml
llm_configs:
  openai_compatible_llm:
    base_url: 'https://qianfan.baidubce.com/v2'
    llm_api_key: 'YOUR_QIANFAN_API_KEY'
    organization_id: null
    project_id: null
    model: 'qwen2.5:latest'
    temperature: 1.0
```

#### 2.2 `src/open_llm_vtuber/config_manager/__init__.py` 修改
添加 RAG 配置导入：
```python
# 在现有导入中添加
from .rag import RAGConfig

# 在 __all__ 列表中添加
__all__ = [
    # ... 现有项 ...
    "RAGConfig",
]
```

#### 2.3 `src/open_llm_vtuber/config_manager/character.py` 修改
在 `CharacterConfig` 类中添加：
```python
from .rag import RAGConfig

class CharacterConfig(I18nMixin, BaseModel):
    # ... 现有字段 ...
    
    # 新增 RAG 配置字段
    rag_config: RAGConfig | None = Field(None, alias="rag_config")
```

#### 2.4 `src/open_llm_vtuber/service_context.py` 修改

**导入部分添加：**
```python
from .rag.rag_engine import RAGEngine, EmbeddingProviderConfig
```

**在 `__init__` 方法中添加：**
```python
def __init__(self):
    # ... 现有代码 ...
    
    # RAG 引擎 (新增)
    self.rag_engine: RAGEngine | None = None
```

**在 `load_cache` 方法末尾添加：**
```python
async def load_cache(self, ...):
    # ... 现有代码 ...
    
    # Initialize RAG engine for this session (新增)
    await self.init_rag(self.character_config.rag_config)
    
    logger.debug(f"Loaded service context with cache: {character_config}")
```

**在 `load_from_config` 方法中添加：**
```python
async def load_from_config(self, config: Config) -> None:
    # ... 现有代码 ...
    
    # init rag from character config (optional) (新增)
    await self.init_rag(config.character_config.rag_config)
```

**添加新方法：**
```python
async def init_rag(self, rag_config: RAGConfig | None) -> None:
    """Initialize or update the RAG engine if configured."""
    if not rag_config or not rag_config.enabled:
        self.rag_engine = None
        logger.info("RAG is disabled or not configured.")
        return
    try:
        provider_key = rag_config.embed_provider_key
        llm_cfg_pool = self.character_config.agent_config.llm_configs
        embed_provider_config = getattr(llm_cfg_pool, provider_key, None)

        if not embed_provider_config:
            raise ValueError(f"Embedding provider '{provider_key}' not found in llm_configs.")

        embed_provider = EmbeddingProviderConfig(
            base_url=embed_provider_config.base_url,
            api_key=embed_provider_config.llm_api_key,
            model=rag_config.embed_model,
        )
        self.rag_engine = RAGEngine(
            index_dir=rag_config.index_dir,
            provider=embed_provider,
            top_k=rag_config.top_k,
            max_context_chars=rag_config.max_context_chars,
            chunk_chars=rag_config.chunk_chars,
            chunk_overlap=rag_config.chunk_overlap,
        )
        logger.info(
            f"RAG initialized: dir={rag_config.index_dir}, model={rag_config.embed_model}, top_k={rag_config.top_k}"
        )
    except Exception as e:
        logger.error(f"Failed to initialize RAG: {e}")
        self.rag_engine = None
```

#### 2.5 `src/open_llm_vtuber/conversations/single_conversation.py` 修改
在 `process_single_conversation` 函数中，在 LLM 调用前添加 RAG 检索：
```python
async def process_single_conversation(
    context: ServiceContext, input_text: str, history_uid: str
) -> Tuple[str, Optional[str]]:
    # ... 现有代码 ...
    
    # Optionally augment input with RAG context (新增)
    rag_engine = getattr(context, "rag_engine", None)
    if rag_engine:
        try:
            ctx_text, ctx_meta = await rag_engine.search(input_text)
            if ctx_text:
                logger.info(
                    f"RAG retrieved {len(ctx_meta)} chunks; augmenting user input."
                )
                input_text = (
                    f"[知识库参考]\n{ctx_text}\n\n[用户问题]\n{input_text}"
                )
        except Exception as e:
            logger.error(f"RAG retrieval failed: {e}")
    
    # ... 现有的 LLM 调用代码 ...
```

#### 2.6 `src/open_llm_vtuber/routes.py` 修改
在文件末尾的路由初始化部分之前添加 RAG 路由函数：

```python
def init_rag_routes(default_context_cache: ServiceContext) -> APIRouter:
    """Initialize RAG management routes"""
    router = APIRouter()

    @router.post("/rag/add-text")
    async def rag_add_text(payload: dict):
        """Add raw text into RAG index. payload: {"text": str}"""
        try:
            if not default_context_cache.rag_engine:
                return JSONResponse({"error": "RAG not enabled"}, status_code=400)
            text = (payload or {}).get("text", "")
            if not isinstance(text, str) or not text.strip():
                return JSONResponse({"error": "Invalid text"}, status_code=400)
            added = await default_context_cache.rag_engine.add_texts([text], source="api")
            return {"added_chunks": added}
        except Exception as e:
            logger.error(f"/rag/add-text failed: {e}")
            return JSONResponse({"error": str(e)}, status_code=500)

    @router.post("/rag/add-file")
    async def rag_add_file(file: UploadFile = File(...)):
        """Add a small plaintext file (.txt/.md/.json) into RAG index."""
        try:
            if not default_context_cache.rag_engine:
                return JSONResponse({"error": "RAG not enabled"}, status_code=400)
            filename = file.filename.lower()
            if not any(filename.endswith(ext) for ext in [".txt", ".md", ".json"]):
                return JSONResponse({"error": "Only .txt .md .json are supported"}, status_code=400)
            content_bytes = await file.read()
            try:
                text = content_bytes.decode("utf-8", errors="ignore")
            except Exception:
                text = content_bytes.decode("gbk", errors="ignore")
            if not text.strip():
                return JSONResponse({"error": "File is empty or unreadable"}, status_code=400)
            added = await default_context_cache.rag_engine.add_texts([text], source=f"file:{file.filename}")
            return {"added_chunks": added, "filename": file.filename}
        except Exception as e:
            logger.error(f"/rag/add-file failed: {e}")
            return JSONResponse({"error": str(e)}, status_code=500)

    @router.post("/rag/clear")
    async def rag_clear():
        """Clear the entire RAG index."""
        try:
            if not default_context_cache.rag_engine:
                return JSONResponse({"error": "RAG not enabled"}, status_code=400)
            await default_context_cache.rag_engine.clear()
            return {"status": "cleared"}
        except Exception as e:
            logger.error(f"/rag/clear failed: {e}")
            return JSONResponse({"error": str(e)}, status_code=500)

    @router.get("/rag/stats")
    async def rag_stats():
        """Get RAG index statistics."""
        try:
            if not default_context_cache.rag_engine:
                return JSONResponse({"error": "RAG not enabled"}, status_code=400)
            return default_context_cache.rag_engine.stats()
        except Exception as e:
            logger.error(f"/rag/stats failed: {e}")
            return JSONResponse({"error": str(e)}, status_code=500)

    return router
```

然后在 `init_routes` 函数中添加 RAG 路由：
```python
def init_routes(default_context_cache: ServiceContext, server_url: str) -> FastAPI:
    # ... 现有代码 ...
    
    # Include RAG routes (新增)
    rag_router = init_rag_routes(default_context_cache)
    app.include_router(rag_router)
    
    return app
```

## 迁移步骤

### 1. 准备工作
1. 确保目标服务器的 Python 环境已安装 `aiohttp` 库：
   ```bash
   pip install aiohttp
   ```

2. 备份目标服务器上的项目文件

### 2. 文件迁移
1. **复制新增文件**：
   - 创建 `src/open_llm_vtuber/rag/` 目录
   - 复制 `rag/__init__.py` 和 `rag/rag_engine.py`
   - 复制 `config_manager/rag.py`

2. **替换修改的文件**：
   - 备份并替换修改的 6 个文件
   - 或者手动应用上述修改内容

### 3. 配置更新
1. **更新 `conf.yaml`**：
   - 添加 `rag_config` 配置段
   - 更新千帆 API 配置（`base_url` 和 `llm_api_key`）

### 4. 测试验证

#### 4.1 启动测试
```bash
python run_server.py
```
观察启动日志，确认看到：
```
INFO | RAG initialized: dir=rag_index, model=embedding-v1, top_k=4
```

#### 4.2 API 测试
```bash
# 检查 RAG 状态
curl http://localhost:PORT/rag/stats

# 添加测试数据（注意中文编码）
# PowerShell:
$testInfo = @{ text = "测试数据" } | ConvertTo-Json -Depth 10
$body = [System.Text.Encoding]::UTF8.GetBytes($testInfo)
Invoke-RestMethod -Uri "http://localhost:PORT/rag/add-text" -Method Post -Headers @{'Content-Type'='application/json; charset=utf-8'} -Body $body
```

#### 4.3 对话测试
在聊天界面测试 RAG 功能，确认：
- RAG 引擎状态为 True
- 检索到相关上下文
- 数字人回答基于知识库内容

## 重要注意事项

### 1. 编码问题
- **PowerShell 用户**：使用 UTF-8 编码发送中文数据
- **文件上传**：确保文件以 UTF-8 编码保存

### 2. API 密钥配置
- 替换 `conf.yaml` 中的 `llm_api_key` 为您的千帆 API 密钥
- 确保 `base_url` 设置为 `https://qianfan.baidubce.com/v2`

### 3. 存储目录
- RAG 索引文件将存储在项目根目录的 `rag_index/` 文件夹中
- 确保应用有读写权限

### 4. 性能考虑
- 大量数据时考虑调整 `chunk_chars` 和 `max_context_chars`
- 监控 embedding API 调用频率和成本

## 故障排除

### 常见问题
1. **RAG 引擎状态为 False**：检查配置文件和千帆 API 密钥
2. **中文显示乱码**：确认使用 UTF-8 编码发送数据
3. **API 调用失败**：验证千帆 API 密钥和网络连接

### 调试方法
- 查看服务器启动日志中的 RAG 初始化信息
- 使用 `/rag/stats` API 检查索引状态
- 在 `single_conversation.py` 中添加调试日志

## 总结

RAG 功能集成包括：
- **新增文件**: 4 个核心文件
- **修改文件**: 6 个现有文件  
- **配置更新**: 主要是 `conf.yaml`
- **API 接口**: 4 个 RAG 管理接口

按照本指南操作，您可以将 RAG 功能完整迁移到服务器环境中，为数字人提供基于知识库的智能问答能力。