0.1.0

wolai-backend/scripts/init_lightrag.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+from typing import Iterable
+
+from app.services.lightrag_service import lightrag_service
+from app.services.supabase_rest import supabase_rest
+
+
+def fetch_documents(
+    workspace_id: str | None,
+    limit: int | None,
+) -> Iterable[dict]:
+    filters = {}
+    if workspace_id:
+        filters["workspace_id"] = workspace_id
+    docs = supabase_rest.select(
+        "documents",
+        filters,
+        columns="id,raw_text,user_id,workspace_id,title",
+        order="updated_at.desc",
+        limit=limit,
+    )
+    return docs
+
+
+async def rebuild_indexes(workspace_id: str | None, limit: int | None) -> None:
+    documents = fetch_documents(workspace_id, limit)
+    for doc in documents:
+        text = str(doc.get("raw_text") or "").strip()
+        if not text:
+            continue
+        document_id = str(doc.get("id"))
+        user_id = str(doc.get("user_id"))
+        workspace = str(doc.get("workspace_id")) if doc.get("workspace_id") else None
+        title = doc.get("title") if isinstance(doc.get("title"), str) else None
+        print(f"[LightRAG:init] indexing {document_id} ({workspace or 'user'})")  # noqa: T201
+        supabase_rest.update(
+            "documents", {"id": document_id}, {"index_status": "pending"}
+        )
+        await lightrag_service.index_document_async(
+            document_id=document_id,
+            user_id=user_id,
+            workspace_id=workspace,
+            text=text,
+            title=title,
+        )
+        supabase_rest.update(
+            "documents", {"id": document_id}, {"index_status": "completed"}
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="初始化 LightRAG 索引")
+    parser.add_argument("--workspace", help="仅处理指定 workspace", default=None)
+    parser.add_argument("--limit", type=int, help="限制处理文档数量", default=None)
+    args = parser.parse_args()
+    asyncio.run(rebuild_indexes(args.workspace, args.limit))
+
+
+if __name__ == "__main__":
+    main()