在nano-graphrag项目中集成本地Embedding服务的实践指南

背景介绍

nano-graphrag是一个基于图结构的检索增强生成(RAG)框架，它允许开发者自定义各种组件来满足特定需求。其中，文本嵌入(Embedding)功能是RAG系统的核心组件之一，负责将文本转换为向量表示。本文将详细介绍如何在nano-graphrag项目中集成本地部署的Embedding服务。

本地Embedding服务集成方案

当需要在nano-graphrag中使用本地Embedding服务时，我们需要创建一个自定义的Embedding函数。以下是实现这一目标的详细步骤：

1. 创建自定义Embedding函数

首先，我们需要定义一个异步函数来处理文本嵌入请求。这个函数需要满足nano-graphrag的接口规范：

from typing import List
import numpy as np
import aiohttp
from nano_graphrag.utils import wrap_embedding_func_with_attrs

@wrap_embedding_func_with_attrs(embedding_dim=512, max_token_size=1024)
async def custom_embedding(texts: List[str]) -> np.ndarray:
    url = "http://127.0.0.1:9997/v1/embeddings"
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer YOUR_API_KEY"
    }
    payload = {
        "model": "BAAI-bge-small-zh-v1.5",
        "input": texts
    }
    
    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=payload, headers=headers) as response:
            if response.status != 200:
                raise Exception(f"Failed to get embeddings: {response.status}")
            data = await response.json()
            embeddings = np.array([dp["embedding"] for dp in data["data"]])
            return embeddings

2. 在GraphRAG初始化时使用自定义Embedding

创建GraphRAG实例时，将自定义的Embedding函数传入：

rag = GraphRAG(
    working_dir=WORKING_DIR,
    embedding_func=custom_embedding,
    best_model_func=deepseepk_model_if_cache,
    cheap_model_func=deepseepk_model_if_cache,
)

Jupyter Notebook环境下的特殊处理

在Jupyter Notebook环境中运行异步代码时，可能会遇到事件循环冲突的问题。这是因为Jupyter本身已经运行了一个事件循环。解决方法是使用nest_asyncio库：

import nest_asyncio
nest_asyncio.apply()

这段代码允许在已经运行的事件循环中再次运行异步代码，解决了"RuntimeError: This event loop is already running"错误。

性能优化建议

1. 批处理优化：确保本地Embedding服务支持批量处理，这样可以减少网络请求次数
2. 连接池管理：重用aiohttp的ClientSession实例可以提高性能
3. 错误处理增强：添加重试机制和更详细的错误日志
4. 缓存机制：考虑实现本地缓存，避免重复计算相同文本的嵌入

常见问题排查

1. 连接问题：确保本地Embedding服务已启动且端口正确
2. 认证问题：检查API密钥是否正确配置
3. 输入格式：确认输入文本格式符合服务端要求
4. 维度匹配：确保自定义Embedding的维度与wrap_embedding_func_with_attrs中声明的一致

总结

通过自定义Embedding函数，我们可以轻松地将本地部署的Embedding服务集成到nano-graphrag项目中。这种方法不仅提供了更大的灵活性，还能根据具体需求选择最适合的Embedding模型。在实现过程中，需要注意异步编程的特殊性，特别是在Jupyter Notebook等交互式环境中使用时。