main.py

import datetime as dt
import json
import os
from pathlib import Path
import urllib.error
import urllib.request

import gradio as gr
from huggingface_hub import InferenceClient


def _build_label() -> str:
    version_file = Path("VERSION")
    version_from_file = ""
    if version_file.exists():
        version_from_file = version_file.read_text(encoding="utf-8").strip()

    commit = (
        os.getenv("GITHUB_SHA")
        or os.getenv("COMMIT_SHA")
        or os.getenv("SPACE_COMMIT_SHA")
        or version_from_file
        or "local"
    )
    short_commit = commit[:7] if commit != "local" else commit
    version = os.getenv("APP_VERSION") or short_commit
    deployed_at = dt.datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
    return f"Version: {version} | Commit: {short_commit} | Loaded: {deployed_at}"


def _env(name: str, default: str = "") -> str:
    return (os.getenv(name) or default).strip()


HF_TOKEN = _env("HF_TOKEN")
HF_MODEL = _env("HF_MODEL", "zai-org/GLM-5.1")

AI_BACKEND = _env("AI_BACKEND", "hf").lower()
AI_MAX_TOKENS = int(_env("AI_MAX_TOKENS", "512"))
AI_FALLBACK_ORDER = [
    p.strip().lower()
    for p in _env("AI_FALLBACK_ORDER", "hf,github,openrouter,fireworks").split(",")
    if p.strip()
]

GITHUB_TOKEN = _env("GITHUB_TOKEN")
GITHUB_MODEL = _env("GITHUB_MODEL")

OPENROUTER_API_KEY = _env("OPENROUTER_API_KEY")
OPENROUTER_MODEL = _env("OPENROUTER_MODEL")

FIREWORKS_API_KEY = _env("FIREWORKS_API_KEY")
FIREWORKS_MODEL = _env("FIREWORKS_MODEL")


# Explicit token passing helps avoid auth ambiguity across local and Space runtimes.
hf_client = InferenceClient(token=HF_TOKEN) if HF_TOKEN else InferenceClient()


def _runtime_label() -> str:
    active_model = {
        "hf": HF_MODEL,
        "github": GITHUB_MODEL,
        "openrouter": OPENROUTER_MODEL,
        "fireworks": FIREWORKS_MODEL,
    }.get(AI_BACKEND, "")
    backend_name = AI_BACKEND.upper()
    model_text = active_model or "not-set"
    return f"Backend: {backend_name} | Model: {model_text}"


def _history_to_messages(history: list, user_message: str) -> list:
    messages = []

    for item in history or []:
        if isinstance(item, dict):
            role = item.get("role")
            content = item.get("content")
            if role in {"user", "assistant", "system"} and content:
                messages.append({"role": role, "content": str(content)})
            continue

        if isinstance(item, (list, tuple)) and len(item) == 2:
            user_msg, assistant_msg = item
            if user_msg:
                messages.append({"role": "user", "content": str(user_msg)})
            if assistant_msg:
                messages.append({"role": "assistant", "content": str(assistant_msg)})

    messages.append({"role": "user", "content": user_message})
    return messages


def _extract_content(choice_message: dict) -> str:
    content = choice_message.get("content", "")
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        chunks = []
        for part in content:
            if isinstance(part, dict) and part.get("type") == "text":
                chunks.append(str(part.get("text", "")))
        return "".join(chunks).strip()
    return str(content)


def _chat_openai_compatible(
    endpoint: str,
    api_key: str,
    model: str,
    messages: list,
    extra_headers=None,
) -> str:
    if not api_key:
        raise ValueError("API key is missing.")
    if not model:
        raise ValueError("Model is not configured.")

    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": AI_MAX_TOKENS,
    }
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    if extra_headers:
        headers.update(extra_headers)

    request = urllib.request.Request(
        endpoint,
        data=json.dumps(payload).encode("utf-8"),
        headers=headers,
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=90) as response:
            body = json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        details = exc.read().decode("utf-8", errors="ignore")
        raise RuntimeError(f"HTTP {exc.code}: {details[:300]}") from exc

    choices = body.get("choices") or []
    if not choices:
        raise RuntimeError("No choices returned from provider.")
    message = choices[0].get("message") or {}
    return _extract_content(message) or "(empty response)"


def _chat_hf(messages: list) -> str:
    response = hf_client.chat_completion(
        model=HF_MODEL,
        messages=messages,
        max_tokens=AI_MAX_TOKENS,
    )
    return response.choices[0].message.content or "(empty response)"


def _chat_github(messages: list) -> str:
    return _chat_openai_compatible(
        endpoint="https://models.github.ai/inference/chat/completions",
        api_key=GITHUB_TOKEN,
        model=GITHUB_MODEL,
        messages=messages,
    )


def _chat_openrouter(messages: list) -> str:
    return _chat_openai_compatible(
        endpoint="https://openrouter.ai/api/v1/chat/completions",
        api_key=OPENROUTER_API_KEY,
        model=OPENROUTER_MODEL,
        messages=messages,
        extra_headers={
            "HTTP-Referer": _env("OPENROUTER_REFERER", "https://huggingface.co"),
            "X-Title": _env("OPENROUTER_APP_NAME", "hf-multi-provider-chat"),
        },
    )


def _chat_fireworks(messages: list) -> str:
    return _chat_openai_compatible(
        endpoint="https://api.fireworks.ai/inference/v1/chat/completions",
        api_key=FIREWORKS_API_KEY,
        model=FIREWORKS_MODEL,
        messages=messages,
    )


def _chat_once(backend: str, messages: list) -> str:
    if backend == "hf":
        return _chat_hf(messages)
    if backend == "github":
        return _chat_github(messages)
    if backend == "openrouter":
        return _chat_openrouter(messages)
    if backend == "fireworks":
        return _chat_fireworks(messages)
    raise ValueError(
        f"Unsupported AI_BACKEND='{backend}'. Use one of: hf, github, openrouter, fireworks, auto"
    )


def chat_response(message: str, history: list) -> str:
    """Send a user message using the configured backend and return assistant text."""
    if not message or not message.strip():
        return "Please enter a message."

    messages = _history_to_messages(history, message.strip())

    try:
        if AI_BACKEND == "auto":
            errors = []
            for backend in AI_FALLBACK_ORDER:
                try:
                    return _chat_once(backend, messages)
                except Exception as exc:  # noqa: BLE001
                    errors.append(f"{backend}: {exc}")
            return "All providers failed. " + " | ".join(errors)

        return _chat_once(AI_BACKEND, messages)
    except Exception as e:
        return f"Error: {str(e)}"


with gr.Blocks(title="GitHub + HuggingFace + AI Chat Demo") as demo:
    gr.Markdown("# GitHub → HuggingFace → AI Chat")
    gr.Markdown(f"**{_build_label()}**")
    gr.Markdown(
        "Multi-provider chat app for learning and testing across HF, GitHub Models, OpenRouter, and Fireworks."
    )
    gr.Markdown(f"**{_runtime_label()}**")

    gr.ChatInterface(
        chat_response,
        examples=[
            "What is the capital of France?",
            "Explain quantum computing in simple terms.",
            "Give me a low-cost model selection strategy for dev vs prod.",
        ],
        title=None,
        description="Ask me anything!",
    )


if __name__ == "__main__":
    # server_name="0.0.0.0" is required inside HF Space containers.
    # root_path ensures Gradio resolves JS/CSS assets correctly when running
    # behind a reverse proxy or custom domain.
    _root_path = os.getenv("GRADIO_ROOT_PATH", "").rstrip("/")
    demo.launch(server_name="0.0.0.0", root_path=_root_path)