File size: 2,763 Bytes
9513cca
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
de13ab3
 
 
9513cca
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
"""
Launch the Streamlit demo from an installed package.

Why:
- After `pip install visual-rag-toolkit`, the repo layout isn't present.
- We package the `demo/` module and expose `visual_rag.demo()` + `visual-rag-demo`.
"""

from __future__ import annotations

import argparse
import importlib
import os
import subprocess
import sys
from pathlib import Path
from typing import Optional


def demo(
    *,
    host: str = "0.0.0.0",
    port: int = 7860,
    headless: bool = True,
    open_browser: bool = False,
    extra_args: Optional[list[str]] = None,
) -> int:
    """
    Launch the Streamlit demo UI.

    Requirements:
    - `visual-rag-toolkit[ui,qdrant,embedding,pdf]` (or `visual-rag-toolkit[all]`)

    Returns:
        Streamlit process exit code.
    """
    try:
        import streamlit  # noqa: F401
    except Exception as e:  # pragma: no cover
        raise RuntimeError(
            "Streamlit is not installed. Install with:\n"
            '  pip install "visual-rag-toolkit[ui,qdrant,embedding,pdf]"'
        ) from e

    # Resolve the installed demo entrypoint path.
    mod = importlib.import_module("demo.app")
    app_path = Path(getattr(mod, "__file__", "")).resolve()
    if not app_path.exists():  # pragma: no cover
        raise RuntimeError("Could not locate installed demo app (demo.app).")

    # Build a stable Streamlit invocation.
    cmd = [sys.executable, "-m", "streamlit", "run", str(app_path)]
    cmd += ["--server.address", str(host)]
    cmd += ["--server.port", str(int(port))]
    # headless=true prevents browser from auto-opening; open_browser overrides
    should_be_headless = headless and not open_browser
    cmd += ["--server.headless", "true" if should_be_headless else "false"]
    cmd += ["--browser.gatherUsageStats", "false"]
    cmd += ["--server.runOnSave", "false"]

    if extra_args:
        cmd += list(extra_args)

    env = os.environ.copy()
    # Make sure the demo doesn't spam internal Streamlit warnings in logs.
    env.setdefault("STREAMLIT_BROWSER_GATHER_USAGE_STATS", "false")

    return subprocess.call(cmd, env=env)


def main() -> None:
    p = argparse.ArgumentParser(description="Launch the Visual RAG Toolkit Streamlit demo.")
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=7860)
    p.add_argument(
        "--no-headless", action="store_true", help="Run with a browser window (not headless)."
    )
    p.add_argument("--open", action="store_true", help="Open browser automatically.")
    args, unknown = p.parse_known_args()

    rc = demo(
        host=args.host,
        port=args.port,
        headless=(not args.no_headless),
        open_browser=bool(args.open),
        extra_args=unknown,
    )
    raise SystemExit(rc)