517 lines
21 KiB
Python
517 lines
21 KiB
Python
"""LightRAG 集成。负责管理实例、增量索引与问答。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
from pathlib import Path
|
||
import sys
|
||
from contextvars import ContextVar
|
||
from typing import Any, AsyncIterator, Dict, List, Optional
|
||
from urllib.parse import urlparse
|
||
import httpx
|
||
import numpy as np
|
||
from openai import AsyncOpenAI
|
||
|
||
|
||
def _ensure_lightrag_available() -> None:
|
||
"""确保本地 LightRAG 代码可被 Python 找到."""
|
||
repo_root = Path(__file__).resolve().parents[3]
|
||
local_pkg = repo_root / "LightRAG"
|
||
if local_pkg.exists():
|
||
path_str = str(local_pkg)
|
||
if path_str not in sys.path:
|
||
sys.path.insert(0, path_str)
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 优先将本地 LightRAG 代码加入 sys.path，避免导入失败走占位实现
|
||
_ensure_lightrag_available()
|
||
|
||
_EMBEDDING_MODEL_OVERRIDE: ContextVar[Optional[str]] = ContextVar(
|
||
"lightrag_embedding_model_override",
|
||
default=None,
|
||
)
|
||
|
||
try:
|
||
from lightrag import LightRAG, QueryParam
|
||
from lightrag.kg.shared_storage import initialize_pipeline_status
|
||
from lightrag.llm.ollama import ollama_model_complete, ollama_embed
|
||
from lightrag.utils import logger as lightrag_logger
|
||
_LIGHTRAG_AVAILABLE = True
|
||
_IMPORT_ERROR: Optional[Exception] = None
|
||
except Exception as exc: # pragma: no cover - 本地缺失或版本不兼容时使用占位实现
|
||
_LIGHTRAG_AVAILABLE = False
|
||
_IMPORT_ERROR = exc
|
||
lightrag_logger = logging.getLogger("lightrag_stub")
|
||
# 占位的 ollama 方法，避免引用错误
|
||
async def ollama_model_complete(*_: Any, **__: Any) -> str: # type: ignore[override]
|
||
return ""
|
||
|
||
async def ollama_embed(*_: Any, **__: Any) -> list[list[float]]: # type: ignore[override]
|
||
return []
|
||
|
||
class QueryParam: # type: ignore[override]
|
||
def __init__(self, mode: str = "mix") -> None:
|
||
self.mode = mode
|
||
self.stream = True
|
||
|
||
class LightRAG: # type: ignore[override]
|
||
async def initialize_storages(self) -> None:
|
||
return None
|
||
|
||
async def ainsert(self, *_: Any, **__: Any) -> str:
|
||
return "lightrag-skipped"
|
||
|
||
async def aquery_llm(self, *_: Any, **__: Any) -> Dict[str, Any]:
|
||
return {
|
||
"llm_response": {
|
||
"response_iterator": iter(()),
|
||
"content": "LightRAG 已跳过",
|
||
"is_streaming": False,
|
||
},
|
||
"data": {"references": []},
|
||
"metadata": {"skipped": True},
|
||
}
|
||
|
||
async def initialize_pipeline_status() -> None: # type: ignore[override]
|
||
return None
|
||
|
||
def gpt_4o_mini_complete(*_: Any, **__: Any) -> str: # type: ignore[override]
|
||
return ""
|
||
|
||
def openai_embed(*_: Any, **__: Any) -> List[float]: # type: ignore[override]
|
||
return []
|
||
|
||
_ensure_lightrag_available()
|
||
|
||
from app.config import settings
|
||
from app.services.supabase_rest import supabase_rest
|
||
|
||
|
||
class LightRAGService:
|
||
"""
|
||
管理 LightRAG 单例、工作空间隔离与增量索引。
|
||
- 每个 workspace 映射为一个 LightRAG 实例（共享 Postgres）
|
||
- 支持在同步/异步上下文中调用
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self.collection = settings.lightrag_collection
|
||
self._instances: Dict[str, LightRAG] = {}
|
||
self._locks: Dict[str, asyncio.Lock] = {}
|
||
self._pipeline_ready = False
|
||
self._configure_pg_env()
|
||
if not _LIGHTRAG_AVAILABLE and _IMPORT_ERROR:
|
||
logger.warning("LightRAG 不可用，使用占位实现：%s", _IMPORT_ERROR)
|
||
self._working_dir = (
|
||
Path(__file__).resolve().parent.parent / "runtime" / "lightrag_cache"
|
||
)
|
||
self._working_dir.mkdir(parents=True, exist_ok=True)
|
||
self._availability_error: Optional[str] = None
|
||
self._ollama_host = settings.ollama_base_url
|
||
self._use_deepseek = bool(getattr(settings, "deepseek_api_key", ""))
|
||
self._deepseek_client: Optional[AsyncOpenAI] = None
|
||
if self._use_deepseek:
|
||
self._deepseek_client = AsyncOpenAI(
|
||
api_key=settings.deepseek_api_key,
|
||
base_url=getattr(settings, "deepseek_base_url", None),
|
||
)
|
||
self._llm_model_name = (
|
||
getattr(settings, "deepseek_model", "deepseek-chat")
|
||
if self._use_deepseek
|
||
else settings.lightrag_llm_model
|
||
)
|
||
|
||
def _configure_pg_env(self) -> None:
|
||
"""根据配置将 pgvector 连接信息注入 LightRAG 需要的环境变量。"""
|
||
parsed = urlparse(settings.lightrag_db_url)
|
||
if parsed.scheme not in {"postgresql", "postgres"}:
|
||
self._availability_error = "LIGHTRAG_DB_URL 必须是 Postgres 连接串"
|
||
return
|
||
|
||
if parsed.username:
|
||
os.environ.setdefault("POSTGRES_USER", parsed.username)
|
||
if parsed.password:
|
||
os.environ.setdefault("POSTGRES_PASSWORD", parsed.password)
|
||
if parsed.hostname:
|
||
os.environ.setdefault("POSTGRES_HOST", parsed.hostname)
|
||
if parsed.port:
|
||
os.environ.setdefault("POSTGRES_PORT", str(parsed.port))
|
||
if parsed.path and len(parsed.path) > 1:
|
||
os.environ.setdefault("POSTGRES_DATABASE", parsed.path.lstrip("/"))
|
||
# Ollama 走本地 HTTP，无需 OpenAI key
|
||
# 允许通过环境变量关闭 KG 抽取，避免大模型深拷贝异常
|
||
os.environ.setdefault("LIGHTRAG_DISABLE_ENTITY_RELATION", "true")
|
||
# 为大维度向量配置 IVFFlat，避免 HNSW 2000 维限制
|
||
os.environ.setdefault("POSTGRES_VECTOR_INDEX_TYPE", "IVFFlat")
|
||
os.environ.setdefault("EMBEDDING_DIM", str(settings.lightrag_embedding_dim if hasattr(settings, "lightrag_embedding_dim") else 4096))
|
||
|
||
def _namespace(self, workspace_id: Optional[str], user_id: str, *, prefer_deepseek: bool = False) -> str:
|
||
"""生成 LightRAG workspace 名称，优先 workspace，其次 user；根据模型标记区分实例。"""
|
||
base = f"workspace_{workspace_id}" if workspace_id else f"user_{user_id}"
|
||
if prefer_deepseek:
|
||
return f"{base}_deepseek"
|
||
return base
|
||
|
||
def _is_available(self) -> tuple[bool, Optional[str]]:
|
||
"""
|
||
判断 LightRAG 是否具备运行条件。
|
||
- import 失败或配置错误时返回 False 和原因
|
||
"""
|
||
if not _LIGHTRAG_AVAILABLE:
|
||
return False, f"LightRAG 导入失败: {self._availability_error or _IMPORT_ERROR}"
|
||
if self._availability_error:
|
||
# 缺少关键配置时也视为不可用
|
||
return False, self._availability_error
|
||
return True, None
|
||
|
||
async def _wait_with_timeout(self, coro: Any, *, timeout: float = 300.0) -> Any:
|
||
"""为外部调用包一层超时，避免卡住 worker / healthcheck。"""
|
||
return await asyncio.wait_for(coro, timeout=timeout)
|
||
|
||
async def _ollama_llm(self, *args: Any, **kwargs: Any) -> Any:
|
||
"""
|
||
固定使用配置好的 Ollama LLM。
|
||
LightRAG 会传入 prompt/system_prompt/history_messages。
|
||
"""
|
||
prompt = kwargs.pop("prompt", None)
|
||
if prompt is None and args:
|
||
prompt = args[0]
|
||
if prompt is None:
|
||
raise ValueError("缺少 prompt，无法调用 Ollama LLM")
|
||
stream_flag = bool(kwargs.pop("stream", False))
|
||
timeout = kwargs.pop("timeout", None)
|
||
return await ollama_model_complete(
|
||
prompt=prompt,
|
||
host=self._ollama_host,
|
||
timeout=timeout,
|
||
stream=stream_flag,
|
||
**kwargs,
|
||
)
|
||
|
||
async def _deepseek_llm(self, *args: Any, **kwargs: Any) -> Any:
|
||
"""
|
||
使用 DeepSeek 在线模型，符合 LightRAG 的 llm_model_func 接口。
|
||
"""
|
||
if not self._deepseek_client:
|
||
raise RuntimeError("DeepSeek 客户端未初始化")
|
||
prompt = kwargs.pop("prompt", None)
|
||
if prompt is None and args:
|
||
prompt = args[0]
|
||
if prompt is None:
|
||
raise ValueError("缺少 prompt，无法调用 DeepSeek LLM")
|
||
system_prompt = kwargs.pop("system_prompt", None)
|
||
history = kwargs.pop("history_messages", []) or []
|
||
stream_flag = bool(kwargs.pop("stream", False))
|
||
temperature = kwargs.pop("temperature", 0.2)
|
||
max_tokens = kwargs.pop("max_tokens", 512)
|
||
# 构建消息
|
||
messages = []
|
||
if system_prompt:
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
messages.extend(history)
|
||
messages.append({"role": "user", "content": prompt})
|
||
client = self._deepseek_client
|
||
if stream_flag:
|
||
response = await client.chat.completions.create(
|
||
model=self._llm_model_name,
|
||
messages=messages,
|
||
stream=True,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
)
|
||
|
||
async def _aiter():
|
||
async for chunk in response:
|
||
delta = chunk.choices[0].delta.content or ""
|
||
if delta:
|
||
yield delta
|
||
|
||
return _aiter()
|
||
else:
|
||
response = await client.chat.completions.create(
|
||
model=self._llm_model_name,
|
||
messages=messages,
|
||
stream=False,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
)
|
||
return response.choices[0].message.content
|
||
|
||
async def _ollama_embed(self, texts: list[str], **kwargs: Any) -> Any:
|
||
"""固定使用配置好的 Ollama embedding 模型。"""
|
||
timeout = kwargs.pop("timeout", None)
|
||
embed_model_override = kwargs.pop("embed_model", None) or _EMBEDDING_MODEL_OVERRIDE.get()
|
||
embed_model = embed_model_override or settings.lightrag_embedding_model
|
||
return await ollama_embed(
|
||
texts,
|
||
embed_model=embed_model,
|
||
host=self._ollama_host,
|
||
timeout=timeout,
|
||
**kwargs,
|
||
)
|
||
|
||
async def _embedding_rerank(self, query: str, documents: List[str], **_: Any) -> List[Dict[str, Any]]:
|
||
"""
|
||
简易基于向量的 rerank：使用指定 embedding 模型计算 query/doc 向量并按余弦相似度排序。
|
||
若调用失败则返回原序。
|
||
"""
|
||
if not documents:
|
||
return []
|
||
try:
|
||
query_vec = await self._ollama_embed([query])
|
||
doc_vecs = await self._ollama_embed(documents)
|
||
q = np.array(query_vec[0], dtype=float)
|
||
d = np.array(doc_vecs, dtype=float)
|
||
q_norm = np.linalg.norm(q) + 1e-8
|
||
d_norm = np.linalg.norm(d, axis=1) + 1e-8
|
||
scores = (d @ q) / (d_norm * q_norm)
|
||
order = np.argsort(-scores)
|
||
return [
|
||
{"index": int(idx), "relevance_score": float(scores[idx])}
|
||
for idx in order
|
||
]
|
||
except Exception as exc:
|
||
logger.warning("rerank 失败，回退原序：%s", exc)
|
||
return [{"index": i, "relevance_score": 0.0} for i in range(len(documents))]
|
||
|
||
async def _get_instance(self, workspace: str, *, prefer_deepseek: bool = False) -> LightRAG:
|
||
if workspace in self._instances:
|
||
return self._instances[workspace]
|
||
|
||
lock = self._locks.setdefault(workspace, asyncio.Lock())
|
||
async with lock:
|
||
if workspace in self._instances:
|
||
return self._instances[workspace]
|
||
|
||
llm_func = self._deepseek_llm if (self._use_deepseek and prefer_deepseek) else self._ollama_llm
|
||
rag = LightRAG(
|
||
working_dir=str(self._working_dir),
|
||
workspace=workspace,
|
||
kv_storage="PGKVStorage",
|
||
vector_storage="PGVectorStorage",
|
||
graph_storage="NetworkXStorage",
|
||
doc_status_storage="PGDocStatusStorage",
|
||
llm_model_func=llm_func,
|
||
llm_model_name=self._llm_model_name if (self._use_deepseek and prefer_deepseek) else settings.lightrag_llm_model,
|
||
llm_model_kwargs={
|
||
"options": {
|
||
# 限制生成长度，避免本地大模型回答过慢
|
||
"num_predict": 512,
|
||
"temperature": 0.2,
|
||
}
|
||
},
|
||
embedding_func=self._ollama_embed,
|
||
rerank_model_func=self._embedding_rerank,
|
||
)
|
||
await self._wait_with_timeout(rag.initialize_storages())
|
||
if not self._pipeline_ready:
|
||
await self._wait_with_timeout(initialize_pipeline_status())
|
||
self._pipeline_ready = True
|
||
self._instances[workspace] = rag
|
||
lightrag_logger.info("LightRAG workspace %s ready", workspace)
|
||
return rag
|
||
|
||
async def index_document_async(
|
||
self,
|
||
*,
|
||
document_id: str,
|
||
user_id: str,
|
||
workspace_id: Optional[str],
|
||
text: str,
|
||
title: Optional[str] = None,
|
||
) -> str:
|
||
workspace = self._namespace(workspace_id, user_id)
|
||
rag = await self._get_instance(workspace)
|
||
clean_text = text.strip()
|
||
if not clean_text:
|
||
raise ValueError("空文本无法建立 LightRAG 索引")
|
||
file_path = f"doc://{document_id}"
|
||
if title:
|
||
file_path = f"{file_path}?title={title}"
|
||
track_id = await rag.ainsert(
|
||
clean_text,
|
||
ids=[document_id],
|
||
file_paths=[file_path],
|
||
)
|
||
return track_id
|
||
|
||
def index_document(
|
||
self,
|
||
*,
|
||
document_id: str,
|
||
user_id: str,
|
||
workspace_id: Optional[str],
|
||
text: str,
|
||
title: Optional[str] = None,
|
||
) -> str:
|
||
"""同步环境（如 Celery）调用的封装。"""
|
||
return asyncio.run(
|
||
self.index_document_async(
|
||
document_id=document_id,
|
||
user_id=user_id,
|
||
workspace_id=workspace_id,
|
||
text=text,
|
||
title=title,
|
||
)
|
||
)
|
||
|
||
async def query_async(
|
||
self,
|
||
*,
|
||
query_text: str,
|
||
user_id: str,
|
||
workspace_id: Optional[str],
|
||
stream: bool = True,
|
||
mode: str = "naive",
|
||
model_choice: Optional[str] = None,
|
||
rag_settings: Optional[Dict[str, Any]] = None,
|
||
) -> Dict[str, Any]:
|
||
prefer_deepseek = model_choice == "deepseek"
|
||
workspace = self._namespace(workspace_id, user_id, prefer_deepseek=prefer_deepseek)
|
||
rag = await self._get_instance(workspace, prefer_deepseek=prefer_deepseek)
|
||
param = QueryParam(mode=mode, top_k=6, chunk_top_k=6, max_total_tokens=4096)
|
||
if rag_settings:
|
||
try:
|
||
if rag_settings.get("mode"):
|
||
param.mode = str(rag_settings["mode"])
|
||
if isinstance(rag_settings.get("top_k"), (int, float)):
|
||
param.top_k = int(rag_settings["top_k"])
|
||
if isinstance(rag_settings.get("chunk_top_k"), (int, float)):
|
||
param.chunk_top_k = int(rag_settings["chunk_top_k"])
|
||
if isinstance(rag_settings.get("max_entity_tokens"), (int, float)):
|
||
param.max_entity_tokens = int(rag_settings["max_entity_tokens"])
|
||
if isinstance(rag_settings.get("max_relation_tokens"), (int, float)):
|
||
param.max_relation_tokens = int(rag_settings["max_relation_tokens"])
|
||
if isinstance(rag_settings.get("max_total_tokens"), (int, float)):
|
||
param.max_total_tokens = int(rag_settings["max_total_tokens"])
|
||
if "enable_rerank" in rag_settings:
|
||
param.enable_rerank = bool(rag_settings["enable_rerank"])
|
||
if rag_settings.get("user_prompt"):
|
||
param.user_prompt = str(rag_settings["user_prompt"])
|
||
except Exception as exc:
|
||
logger.warning("解析 RAG 参数失败，继续使用默认值: %s", exc)
|
||
param.stream = stream
|
||
embedding_token = None
|
||
if rag_settings and isinstance(rag_settings.get("embedding_model"), str):
|
||
embedding_token = _EMBEDDING_MODEL_OVERRIDE.set(str(rag_settings["embedding_model"]))
|
||
try:
|
||
result = await self._wait_with_timeout(rag.aquery_llm(query_text, param))
|
||
finally:
|
||
if embedding_token is not None:
|
||
_EMBEDDING_MODEL_OVERRIDE.reset(embedding_token)
|
||
return result
|
||
|
||
async def stream_answer(
|
||
self,
|
||
*,
|
||
query_text: str,
|
||
user_id: str,
|
||
workspace_id: Optional[str],
|
||
model_choice: Optional[str] = None,
|
||
document_id: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
"""统一返回结构，包含 SSE 需要的 iterator 与引用。"""
|
||
ok, reason = self._is_available()
|
||
if not ok:
|
||
return {
|
||
"references": [],
|
||
"iterator": iter(()),
|
||
"content": f"LightRAG 未就绪：{reason}",
|
||
"is_streaming": False,
|
||
"metadata": {"skipped": True, "reason": reason},
|
||
}
|
||
rag_settings = self._load_rag_settings(document_id)
|
||
rag_mode = (rag_settings or {}).get("mode")
|
||
result = await self.query_async(
|
||
query_text=query_text,
|
||
user_id=user_id,
|
||
workspace_id=workspace_id,
|
||
stream=True,
|
||
model_choice=model_choice,
|
||
mode=str(rag_mode) if isinstance(rag_mode, str) else "naive",
|
||
rag_settings=rag_settings,
|
||
)
|
||
llm_resp = result.get("llm_response", {})
|
||
references: List[Dict[str, Any]] = (
|
||
result.get("data", {}).get("references", []) or []
|
||
)
|
||
return {
|
||
"references": references,
|
||
"iterator": llm_resp.get("response_iterator"),
|
||
"content": llm_resp.get("content"),
|
||
"is_streaming": llm_resp.get("is_streaming", False),
|
||
"metadata": result.get("metadata", {}),
|
||
}
|
||
|
||
def _load_rag_settings(self, document_id: Optional[str]) -> Optional[Dict[str, Any]]:
|
||
if not document_id:
|
||
return None
|
||
try:
|
||
doc = supabase_rest.select_one("documents", {"id": document_id})
|
||
except Exception as exc:
|
||
logger.warning("读取文档 %s 的 RAG 配置失败：%s", document_id, exc)
|
||
return None
|
||
value = doc.get("rag_settings") if doc else None
|
||
if isinstance(value, dict):
|
||
return value
|
||
return None
|
||
|
||
async def answer_with_context(
|
||
self,
|
||
*,
|
||
query_text: str,
|
||
context: str,
|
||
prefer_deepseek: bool = False,
|
||
) -> Dict[str, Any]:
|
||
"""基于外部上下文的简单问答，优先走 DeepSeek，无则回退 Ollama。"""
|
||
prompt = (
|
||
"你是检索结果总结助手。根据以下搜索摘要与来源回答用户问题，"
|
||
"答案需简洁且引用要点。保持中文输出，并在内容后附上引用编号。\n\n"
|
||
f"【搜索摘要】\n{context}\n\n【用户问题】{query_text}"
|
||
)
|
||
prefer_deepseek = prefer_deepseek and self._use_deepseek
|
||
stream_flag = True
|
||
iterator: Optional[AsyncIterator[str]] = None
|
||
content: Optional[str] = None
|
||
reason: Optional[str] = None
|
||
try:
|
||
if prefer_deepseek and self._deepseek_client:
|
||
resp = await self._deepseek_llm(prompt, stream=stream_flag)
|
||
iterator = resp if hasattr(resp, "__aiter__") else None
|
||
content = None if iterator else str(resp)
|
||
else:
|
||
resp = await self._ollama_llm(prompt, stream=stream_flag)
|
||
iterator = resp if hasattr(resp, "__aiter__") else None
|
||
content = None if iterator else str(resp)
|
||
except Exception as exc: # pragma: no cover - 运行时保护
|
||
iterator = None
|
||
content = f"生成失败：{exc}"
|
||
reason = str(exc)
|
||
|
||
return {
|
||
"references": [],
|
||
"iterator": iterator,
|
||
"content": content,
|
||
"is_streaming": iterator is not None,
|
||
"metadata": {"reason": reason} if reason else {},
|
||
}
|
||
|
||
async def run_healthcheck(self) -> Dict[str, Any]:
|
||
workspace = "__healthcheck__"
|
||
ok, reason = self._is_available()
|
||
if not ok:
|
||
return {"ok": False, "workspace": workspace, "error": reason}
|
||
try:
|
||
rag = await self._get_instance(workspace)
|
||
await self._wait_with_timeout(rag.doc_status.initialize(), timeout=5.0)
|
||
return {"ok": True, "workspace": workspace}
|
||
except Exception as exc: # pragma: no cover - 调试辅助
|
||
return {"ok": False, "workspace": workspace, "error": str(exc)}
|
||
|
||
|
||
lightrag_service = LightRAGService()
|