Initial commit

2026-05-12 09:41:56 +08:00
commit 572283e101
936 changed files with 133949 additions and 0 deletions
@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+from config import DEFAULT_CONFIG
+
+
+def _env(name: str, default: str | None = None) -> str | None:
+    value = os.getenv(name)
+    if value is None or value.strip() == "":
+        return default
+    return value.strip()
+
+
+def _require(value: str | None, hint: str) -> str:
+    if value is None or value.strip() == "":
+        raise SystemExit(hint)
+    return value.strip()
+
+
+def _read_bytes(path: Path) -> bytes:
+    try:
+        return path.read_bytes()
+    except FileNotFoundError:
+        raise SystemExit(f"File not found: {path}")
+
+
+def _safe_get(obj: Any, attr: str, default: Any = None) -> Any:
+    try:
+        return getattr(obj, attr)
+    except Exception:
+        return default
+
+
+def main(argv: list[str]) -> int:
+    parser = argparse.ArgumentParser(
+        description="PowerRAG (RAGFlow) SDK demo: upload Markdown, parse, retrieve top-k chunks.",
+    )
+    parser.add_argument("--file", type=Path, required=True, help="Markdown file path, e.g. ./data/sample.md")
+    parser.add_argument("--question", type=str, required=True, help="User question for retrieval")
+    parser.add_argument("--top-k", type=int, default=DEFAULT_CONFIG.top_k, help="How many chunks to return (mapped to page_size)")
+    parser.add_argument(
+        "--embedding-model",
+        type=str,
+        default=DEFAULT_CONFIG.embedding_model or _env("RAGFLOW_EMBEDDING_MODEL"),
+        help=(
+            "Embedding model string in '<model>@<factory>' format. "
+            "If omitted, server tenant default is used."
+        ),
+    )
+    parser.add_argument("--candidate-k", type=int, default=DEFAULT_CONFIG.candidate_k, help="RAGFlow.retrieve(top_k=...) candidate pool size")
+    parser.add_argument("--similarity-threshold", type=float, default=DEFAULT_CONFIG.similarity_threshold, help="Filter chunks below this similarity")
+    parser.add_argument("--vector-similarity-weight", type=float, default=DEFAULT_CONFIG.vector_similarity_weight, help="Weight of vector similarity in hybrid score")
+    parser.add_argument("--keyword", action="store_true", default=DEFAULT_CONFIG.keyword, help="Enable keyword matching (hybrid retrieval)")
+    parser.add_argument("--dataset-name", type=str, default=DEFAULT_CONFIG.dataset_name, help="Dataset name to create")
+    parser.add_argument(
+        "--base-url",
+        type=str,
+        default=DEFAULT_CONFIG.base_url or _env("RAGFLOW_BASE_URL") or _env("POWERRAG_BASE_URL") or _env("BASE_URL"),
+        help="RAGFlow/PowerRAG base_url (or env RAGFLOW_BASE_URL / POWERRAG_BASE_URL / BASE_URL)",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=DEFAULT_CONFIG.api_key or _env("RAGFLOW_API_KEY") or _env("POWERRAG_API_KEY") or _env("API_KEY"),
+        help="RAGFlow/PowerRAG api_key (or env RAGFLOW_API_KEY / POWERRAG_API_KEY / API_KEY)",
+    )
+    parser.add_argument("--cleanup", action="store_true", help="Delete created dataset after finishing")
+
+    args = parser.parse_args(argv)
+
+    base_url = _require(args.base_url, "Missing base_url. Use --base-url or set env RAGFLOW_BASE_URL.")
+    api_key = _require(args.api_key, "Missing api_key. Use --api-key or set env RAGFLOW_API_KEY.")
+
+    if args.top_k <= 0:
+        raise SystemExit("--top-k must be > 0")
+    if args.candidate_k <= 0:
+        raise SystemExit("--candidate-k must be > 0")
+
+    blob = _read_bytes(args.file)
+    display_name = args.file.name
+    if not display_name.lower().endswith(".md"):
+        display_name = f"{display_name}.md"
+
+    try:
+        from ragflow_sdk import RAGFlow  # type: ignore
+    except Exception as e:
+        raise SystemExit(
+            "Failed to import ragflow_sdk. Install dependencies first:\n"
+            "  pip install -r requirements.txt\n"
+            f"Original error: {e}"
+        )
+
+    rag = RAGFlow(api_key=api_key, base_url=base_url)
+
+    dataset_kwargs: dict[str, Any] = {"name": args.dataset_name}
+    if args.embedding_model:
+        dataset_kwargs["embedding_model"] = args.embedding_model
+    dataset = rag.create_dataset(**dataset_kwargs)
+    try:
+        docs = dataset.upload_documents([{"display_name": display_name, "blob": blob}])
+        if not docs:
+            raise SystemExit("Upload succeeded but no document returned by SDK.")
+        doc = docs[0]
+
+        parse_results = dataset.parse_documents([doc.id])
+        # parse_results: list[tuple[doc_id, status, success_count, failure_count]] (per API ref)
+        print("Parse results:")
+        print(parse_results)
+        if parse_results and isinstance(parse_results, list):
+            statuses = {r[1] for r in parse_results if isinstance(r, (list, tuple)) and len(r) >= 2}
+            if statuses and statuses != {"DONE"}:
+                raise SystemExit(
+                    "Document parsing failed (status not DONE). "
+                    "Most common cause is missing/unauthorized embedding model.\n"
+                    "Try:\n"
+                    "  - set tenant default embedding model in UI or via /v1/user/set_tenant_info, OR\n"
+                    "  - rerun with --embedding-model '<model>@<factory>' (must be supported & configured for the tenant)\n"
+                    "If it still fails, check PowerRAG logs inside the container (task executor) for the detailed error.\n"
+                )
+
+        chunks = rag.retrieve(
+            question=args.question,
+            dataset_ids=[dataset.id],
+            document_ids=[doc.id],
+            page=1,
+            page_size=args.top_k,
+            similarity_threshold=args.similarity_threshold,
+            vector_similarity_weight=args.vector_similarity_weight,
+            top_k=args.candidate_k,
+            keyword=args.keyword,
+        )
+
+        print("\nRetrieved chunks:")
+        if not chunks:
+            print("(empty)")
+            return 0
+
+        for i, c in enumerate(chunks, start=1):
+            similarity = _safe_get(c, "similarity")
+            vector_similarity = _safe_get(c, "vector_similarity")
+            term_similarity = _safe_get(c, "term_similarity")
+            content = _safe_get(c, "content", "")
+            content_preview = (content or "").strip().replace("\n", " ")
+            if len(content_preview) > 260:
+                content_preview = content_preview[:260] + "…"
+            print(f"{i:02d}. similarity={similarity} vector={vector_similarity} term={term_similarity}")
+            print(f"    {content_preview}")
+
+        return 0
+    finally:
+        if args.cleanup:
+            try:
+                rag.delete_datasets(ids=[dataset.id])
+            except Exception as e:
+                print(f"Warning: failed to cleanup dataset {dataset.id}: {e}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))