64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
|
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()
|