File size: 3,622 Bytes
d44b33d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
"""Hugging Face Spaces default entry (Streamlit SDK expects `app.py`).

Local development can still use `streamlit run streamlit_app.py`; Docker Compose uses `app.py`
so the same entry path works on the Hub and in containers.

On Hugging Face Streamlit Spaces only `streamlit run app.py` is started — no separate uvicorn
process — so we spawn the FastAPI app on 127.0.0.1:8000 when `SPACE_ID` is present (see Hub
built-in env vars). Set `DOC_AUDI_EMBED_API=0` to disable. Use `DOC_AUDI_EMBED_API=1` to force
embedding elsewhere (e.g. demos).
"""

from __future__ import annotations

import atexit
import os
import socket
import subprocess
import sys
import time

_uvicorn_proc: subprocess.Popen[bytes] | None = None
_cleanup_registered = False

def _port_accepting_connections(host: str, port: int) -> bool:
    try:
        with socket.create_connection((host, port), timeout=0.3):
            return True
    except OSError:
        return False


def _want_embedded_api() -> bool:
    if os.environ.get("DOC_AUDI_EMBED_API", "").lower() in ("0", "false", "no"):
        return False
    if os.environ.get("DOC_AUDI_EMBED_API", "").lower() in ("1", "true", "yes"):
        return True
    return bool(os.environ.get("SPACE_ID"))


def _propagate_streamlit_secrets_to_environ() -> None:
    """Copy Hub tokens from Streamlit secrets into os.environ for the embedded uvicorn child.

    On Hugging Face Streamlit Spaces, repository secrets are often available as ``st.secrets``
    but are not always present in ``os.environ``. ``subprocess.Popen`` only forwards the
    process environment, so the API would miss ``HF_TOKEN`` / ``HUGGINGFACE_API_KEY`` otherwise.
    """
    try:
        import streamlit as st
    except ImportError:
        return
    secrets = getattr(st, "secrets", None)
    if secrets is None:
        return
    for key in ("HF_TOKEN", "HUGGINGFACE_API_KEY", "HUGGING_FACE_HUB_TOKEN"):
        if (os.environ.get(key) or "").strip():
            continue
        try:
            raw = secrets[key]
        except Exception:
            continue
        if raw is not None and str(raw).strip():
            os.environ[key] = str(raw).strip()


def _maybe_start_embedded_uvicorn() -> None:
    """Start uvicorn in-process when running on HF Spaces (or when DOC_AUDI_EMBED_API=1)."""
    global _uvicorn_proc, _cleanup_registered
    if not _want_embedded_api():
        return
    _propagate_streamlit_secrets_to_environ()
    if _port_accepting_connections("127.0.0.1", 8000):
        return
    if _uvicorn_proc is not None and _uvicorn_proc.poll() is None:
        for _ in range(120):
            if _port_accepting_connections("127.0.0.1", 8000):
                return
            time.sleep(0.05)
        return

    cmd = [
        sys.executable,
        "-m",
        "uvicorn",
        "api.main:app",
        "--host",
        "127.0.0.1",
        "--port",
        "8000",
    ]
    _uvicorn_proc = subprocess.Popen(cmd)
    proc = _uvicorn_proc

    if not _cleanup_registered:

        def _cleanup(p: subprocess.Popen[bytes] = proc) -> None:
            if p.poll() is None:
                p.terminate()
                try:
                    p.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    p.kill()

        atexit.register(_cleanup)
        _cleanup_registered = True

    for _ in range(120):
        if _port_accepting_connections("127.0.0.1", 8000):
            return
        time.sleep(0.05)


_maybe_start_embedded_uvicorn()

from streamlit_app import main  # noqa: E402 — start API before importing Streamlit stack

main()