0.1.0

wolai-backend/app/services/lightrag_service.py

@@ -1,19 +1,277 @@
-"""LightRAG 集成占位。阶段 1 会在这里封装真正的查询与索引。"""
+"""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
 
-    async def queue_index(self, document_id: str, raw_text: str) -> None:
-        """预留方法：后续调用 LightRAG.update_index。"""
-        return 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
 
-    async def query(self, query_text: str, user_id: str) -> str:
-        """预留方法：后续调用 LightRAG.query，当前返回占位回答。"""
-        return f"[mock] {query_text}"
+        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()

wolai-backend/app/services/mineru_service.py

@@ -1,7 +1,60 @@
-"""MinerU OCR 占位服务。"""
+"""MinerU OCR 服务封装，支持调用本地 MinerU HTTP 接口。"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional, Tuple
+
+import httpx
+
+
+class MinerUServiceError(RuntimeError):
+    """MinerU 调用失败时抛出的异常。"""
 
 
 class MinerUService:
+    def __init__(self) -> None:
+        self.endpoint = os.getenv("MINERU_ENDPOINT", "http://127.0.0.1:18888").rstrip("/")
+        timeout = float(os.getenv("MINERU_TIMEOUT_SECONDS", "300"))
+        # MinerU 首次加载模型可能耗时较长，这里提高超时时间避免大文件 OCR 直接失败
+        self._client = httpx.Client(timeout=timeout)
+
+    def _request_markdown(self, file_path: str) -> Tuple[Optional[str], dict]:
+        """
+        调用 MinerU /file_parse，返回 (markdown, raw_response)。
+        若接口未返回内容则返回 (None, response_dict)。
+        """
+        if not self.endpoint:
+            raise MinerUServiceError("未配置 MINERU_ENDPOINT，无法调用 MinerU")
+
+        url = f"{self.endpoint}/file_parse"
+        file_name = Path(file_path).name
+        backend = os.getenv("MINERU_DEFAULT_BACKEND", "vlm-transformers")
+        model_path = os.getenv("MINERU_MODEL_PATH")
+        with open(file_path, "rb") as fp:
+            response = self._client.post(
+                url,
+                files={"files": (file_name, fp, "application/octet-stream")},
+                data={
+                    "return_md": "true",
+                    "return_content_list": "false",
+                    "return_middle_json": "false",
+                    "response_format_zip": "false",
+                    "backend": backend,
+                    # fast_api main 会从环境注入 model_path；这里双保险随请求传递
+                    **({"model_path": model_path} if model_path else {}),
+                },
+            )
+        response.raise_for_status()
+        payload = response.json()
+        results = payload.get("results") or {}
+        if not isinstance(results, dict) or not results:
+            return None, payload
+        first_key = next(iter(results.keys()))
+        md_content = results.get(first_key, {}).get("md_content")
+        return md_content, payload
+
     async def extract_markdown(self, file_url: str) -> str:
         """
         阶段 0：直接返回固定内容，保证前端流程贯通。
@@ -9,5 +62,19 @@ class MinerUService:
         """
         return f"# OCR Placeholder\n\n源文件：{file_url}"
 
+    def extract_markdown_sync(self, file_path: str) -> str:
+        """
+        Celery 任务使用的同步封装。
+        - 若 MinerU 服务可用：调用 HTTP 接口返回 markdown
+        - 若失败：抛出 MinerUServiceError 让上层标记失败
+        """
+        try:
+            markdown, raw = self._request_markdown(file_path)
+            if markdown:
+                return markdown
+            raise MinerUServiceError(f"MinerU 未返回 md_content，响应片段：{str(raw)[:300]}")
+        except Exception as exc:  # pragma: no cover - IO/网络异常
+            raise MinerUServiceError(f"MinerU 调用失败：{exc}") from exc
+
 
 mineru_service = MinerUService()

wolai-backend/app/services/storage_fetcher.py

@@ -0,0 +1,37 @@
+"""简单的文件下载器，负责将 Supabase Storage 签名 URL 暂存到临时目录。"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from pathlib import Path
+from typing import Optional
+from urllib.parse import urlparse
+
+import httpx
+
+
+class StorageFetcher:
+    def __init__(self) -> None:
+        self._client = httpx.Client(timeout=30.0)
+
+    def download(self, file_url: str) -> str:
+        response = self._client.get(file_url)
+        response.raise_for_status()
+        suffix = Path(urlparse(file_url).path).suffix or ".bin"
+        fd, path = tempfile.mkstemp(suffix=suffix)
+        try:
+            os.write(fd, response.content)
+        finally:
+            os.close(fd)
+        return path
+
+    def cleanup(self, path: Optional[str]) -> None:
+        if path and os.path.exists(path):
+            try:
+                os.remove(path)
+            except OSError:
+                pass
+
+
+storage_fetcher = StorageFetcher()

wolai-backend/app/services/supabase_rest.py

@@ -12,10 +12,11 @@ class SupabaseRestClient:
 
     def __init__(self) -> None:
         base_url = settings.supabase_url.rstrip("/")
+        apikey = settings.supabase_anon_key or settings.supabase_service_role_key
         self.client = httpx.Client(
             base_url=f"{base_url}/rest/v1",
             headers={
-                "apikey": settings.supabase_service_role_key,
+                "apikey": apikey,
                 "Authorization": f"Bearer {settings.supabase_service_role_key}",
             },
             timeout=10.0,
@@ -43,6 +44,32 @@ class SupabaseRestClient:
             return data[0]
         return None
 
+    def select(
+        self,
+        table: str,
+        filters: Optional[Dict[str, Any]] = None,
+        *,
+        columns: str = "*",
+        order: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> list[Dict[str, Any]]:
+        params: Dict[str, Any] = {"select": columns}
+        if filters:
+            for key, value in filters.items():
+                params[key] = f"eq.{value}"
+        if order:
+            params["order"] = order
+        if limit is not None:
+            params["limit"] = limit
+        response = self.client.get(f"/{table}", params=params)
+        response.raise_for_status()
+        data = response.json()
+        if isinstance(data, list):
+            return data
+        if data:
+            return [data]
+        return []
+
     def update(self, table: str, filters: Dict[str, Any], payload: Dict[str, Any]) -> None:
         params = {key: f"eq.{value}" for key, value in filters.items()}
         response = self.client.patch(