278 lines
10 KiB
Python
278 lines
10 KiB
Python
"""LightRAG 集成。负责管理实例、增量索引与问答。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
from pathlib import Path
|
||
import sys
|
||
from typing import Any, AsyncIterator, Dict, List, Optional
|
||
from urllib.parse import urlparse
|
||
|
||
|
||
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__)
|
||
|
||
try:
|
||
from lightrag import LightRAG, QueryParam
|
||
from lightrag.kg.shared_storage import initialize_pipeline_status
|
||
from lightrag.llm.openai import gpt_4o_mini_complete, openai_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")
|
||
|
||
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
|
||
|
||
|
||
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
|
||
|
||
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("/"))
|
||
# OpenAI key 也交给环境变量，缺失时依旧由外部控制
|
||
if settings.openai_api_key:
|
||
os.environ.setdefault("OPENAI_API_KEY", settings.openai_api_key)
|
||
elif self._availability_error is None:
|
||
# 允许无 key 但标记提示，便于健康检查返回可诊断信息
|
||
self._availability_error = "缺少 OPENAI_API_KEY，LightRAG 将使用占位实现"
|
||
|
||
def _namespace(self, workspace_id: Optional[str], user_id: str) -> str:
|
||
"""生成 LightRAG workspace 名称，优先 workspace，其次 user。"""
|
||
if workspace_id:
|
||
return f"workspace_{workspace_id}"
|
||
return f"user_{user_id}"
|
||
|
||
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 = 6.0) -> Any:
|
||
"""为外部调用包一层超时，避免卡住 worker / healthcheck。"""
|
||
return await asyncio.wait_for(coro, timeout=timeout)
|
||
|
||
async def _get_instance(self, workspace: str) -> 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]
|
||
|
||
rag = LightRAG(
|
||
working_dir=str(self._working_dir),
|
||
workspace=workspace,
|
||
kv_storage="PGKVStorage",
|
||
vector_storage="PGVectorStorage",
|
||
graph_storage="PGGraphStorage",
|
||
doc_status_storage="PGDocStatusStorage",
|
||
llm_model_func=gpt_4o_mini_complete,
|
||
embedding_func=openai_embed,
|
||
)
|
||
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 = "mix",
|
||
) -> Dict[str, Any]:
|
||
workspace = self._namespace(workspace_id, user_id)
|
||
rag = await self._get_instance(workspace)
|
||
param = QueryParam(mode=mode)
|
||
param.stream = stream
|
||
result = await self._wait_with_timeout(rag.aquery_llm(query_text, param))
|
||
return result
|
||
|
||
async def stream_answer(
|
||
self,
|
||
*,
|
||
query_text: str,
|
||
user_id: str,
|
||
workspace_id: Optional[str],
|
||
) -> 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},
|
||
}
|
||
result = await self.query_async(
|
||
query_text=query_text,
|
||
user_id=user_id,
|
||
workspace_id=workspace_id,
|
||
stream=True,
|
||
)
|
||
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", {}),
|
||
}
|
||
|
||
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()
|