src/clients.py

"""Backend clients for the comparison playground.

Exposes a single backend-agnostic streaming generator, `stream_backend`, that
the UI uses for every model. Supports two backend types:
  - custom : POST to BASE_URL with x-api-key; endpoint_id + model in the body
             (any OpenAI-compatible chat-completions endpoint).
  - azure  : Azure OpenAI via the openai SDK (base_url = endpoint/openai/v1/).

Also keeps `build_messages`, `get_api_key`, `chat_completion`, `stream_chat`, and
`BackendError` as the lower-level single-backend helpers.
"""

import json
import os
import time

import requests
from dotenv import load_dotenv

import config

# Load a local .env if present. On Hugging Face Spaces secrets are real env
# vars, so this is a harmless no-op there.
load_dotenv()


class BackendError(RuntimeError):
    """Raised for configuration or API errors so the UI can show a clear message."""


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def get_api_key() -> str:
    """Return the custom backend's API key from the environment, or raise."""
    key = os.environ.get("BACKEND_API_KEY", "").strip()
    if not key:
        raise BackendError(
            "BACKEND_API_KEY is not set. Add it to a local .env file "
            "(see .env.example) or as a Hugging Face Space Secret."
        )
    return key


def build_messages(system_prompt: str, intro: str, history: list) -> list:
    """Assemble the OpenAI-style messages list.

    - `system_prompt`: optional system message prepended first.
    - `intro`: optional assistant message seeded as the first assistant turn.
    - `history`: list of {"role", "content"} dicts, already excluding the seeded
      intro bubble.
    """
    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    if intro and intro.strip():
        messages.append({"role": "assistant", "content": intro.strip()})
    for turn in history:
        role = turn.get("role")
        content = turn.get("content")
        if role in ("user", "assistant") and content:
            messages.append({"role": role, "content": content})
    return messages


def _empty_metrics(**overrides) -> dict:
    base = {
        "__metrics__": True,
        "prompt_tokens": None,
        "completion_tokens": None,
        "total_tokens": None,
        "cached_tokens": None,
        "latency_s": None,
        # Name of the tool the model invoked this turn (e.g. "end_call"), or
        # None when the reply was ordinary text. Lets the UI/logs flag tool use.
        "tool_called": None,
        # The exact assistant message the model produced this turn, as an
        # OpenAI-style {"content", "tool_calls"} dict (raw tool arguments kept
        # verbatim). Powers the debug panel; None on error.
        "raw_response": None,
        "error": None,
    }
    base.update(overrides)
    return base


# ---------------------------------------------------------------------------
# Speech-to-Text — voice input (optional)
# ---------------------------------------------------------------------------


def _stt_api_key() -> str:
    """Key for the STT provider; falls back to the backend key if unset."""
    key = os.environ.get("STT_API_KEY", "").strip()
    if key:
        return key
    return get_api_key()


def transcribe_audio(filepath: str, language: str = "hi") -> str:
    """Transcribe an audio file and return the text.

    Uses the `ringglabs` STT SDK. Accepts standard WAV (any sample rate /
    channels — the server reads the header), which is what Gradio's microphone
    produces. The key comes from STT_API_KEY (or BACKEND_API_KEY as a fallback).
    """
    if not filepath:
        return ""
    try:
        from ringglabs.stt import Client
    except ImportError as e:
        raise BackendError(
            "The 'ringglabs' package is required for voice input. "
            "Add it to requirements.txt / pip install ringglabs."
        ) from e

    with Client(api_key=_stt_api_key()) as client:
        result = client.transcribe(filepath, language=language)
    return (getattr(result, "transcription", "") or "").strip()


# ---------------------------------------------------------------------------
# Custom OpenAI-compatible backend
# ---------------------------------------------------------------------------


def chat_completion(
    messages: list,
    model: str,
    endpoint_id: str,
    temperature: float,
    max_tokens: int,
    tools: list | None = None,
) -> tuple:
    """Return (content, usage) from a non-streaming chat completion.

    When `tools` is provided it is sent in the body so the model may call a
    tool. If the model responds with an `end_call` tool call instead of plain
    text, its `final_message` argument is returned as the content.
    """
    payload = {
        "messages": messages,
        "model": model,
        "endpoint_id": endpoint_id,
        "temperature": float(temperature),
        "max_tokens": int(max_tokens),
        "stream": False,
    }
    if tools:
        payload["tools"] = tools
    headers = {"Content-Type": "application/json", "x-api-key": get_api_key()}
    resp = requests.post(config.BASE_URL, headers=headers, json=payload, timeout=120)
    if resp.status_code != 200:
        raise BackendError(f"API returned HTTP {resp.status_code}: {resp.text[:500]}")
    data = resp.json()
    choices = data.get("choices") or []
    content = ""
    if choices:
        message = choices[0].get("message") or {}
        content = message.get("content") or ""
        if not content:
            content = _extract_end_call_message(message.get("tool_calls"))
    usage = data.get("usage") or {}
    return content, usage


def _extract_end_call_message(tool_calls) -> str:
    """Return the `final_message` from an `end_call` tool call, or "".

    Accepts the non-streaming `tool_calls` list from a chat message. Other tool
    calls (or malformed arguments) yield an empty string.
    """
    for call in tool_calls or []:
        fn = call.get("function") or {}
        if fn.get("name") != "end_call":
            continue
        try:
            args = json.loads(fn.get("arguments") or "{}")
        except json.JSONDecodeError:
            return ""
        return (args.get("final_message") or "").strip()
    return ""


def stream_chat(
    messages: list,
    model: str,
    endpoint_id: str,
    temperature: float,
    max_tokens: int,
    tools: list | None = None,
):
    """Yield response text deltas, then a final metrics sentinel.

    When `tools` is provided it is sent in the body so the model may call a
    tool. A streamed `end_call` tool call has no text content, so its
    `final_message` argument (assembled from the streamed argument fragments)
    is yielded as the reply once the stream ends.

    Token counts come from a lightweight max_tokens=1 probe fired AFTER the
    stream (so it never competes with streaming for GPU); completion tokens are
    the streamed-piece count.
    """
    payload = {
        "messages": messages,
        "model": model,
        "endpoint_id": endpoint_id,
        "temperature": float(temperature),
        "max_tokens": int(max_tokens),
        "stream": True,
    }
    if tools:
        payload["tools"] = tools
    headers = {"Content-Type": "application/json", "x-api-key": get_api_key()}

    # Surface the exact request body (no secret header) for the debug panel.
    yield {"__request__": True, "payload": payload}

    t_start = time.monotonic()
    piece_count = 0
    ttfb = 0
    tool_args = ""  # concatenated `end_call` argument fragments
    tool_name = None  # name of the tool the model invoked, if any
    content_buf = ""  # raw text the model streamed (before any tool handling)
    saw_end_call = False
    with requests.post(
        config.BASE_URL, headers=headers, json=payload, stream=True, timeout=120
    ) as resp:
        if resp.status_code != 200:
            raise BackendError(
                f"API returned HTTP {resp.status_code}: {resp.text[:500]}"
            )
        for raw_line in resp.iter_lines(decode_unicode=True):
            if not raw_line:
                continue
            if raw_line.startswith("data: "):
                raw_line = raw_line[len("data: ") :]
            if raw_line.strip() == "[DONE]":
                break
            try:
                chunk = json.loads(raw_line)
            except json.JSONDecodeError:
                continue
            choices = chunk.get("choices") or []
            if not choices:
                continue
            else:
                ttfb = time.monotonic() - t_start
            delta = choices[0].get("delta") or {}
            piece = delta.get("content")
            if piece:
                piece_count += 1
                content_buf += piece
                yield piece
            # An `end_call` tool call streams as `tool_calls` deltas whose
            # `arguments` strings must be concatenated, then parsed at the end.
            for call in delta.get("tool_calls") or []:
                fn = call.get("function") or {}
                if fn.get("name"):
                    tool_name = fn["name"]
                    if tool_name == "end_call":
                        saw_end_call = True
                frag = fn.get("arguments")
                if frag:
                    saw_end_call = True
                    tool_args += frag

    # If the model ended the call via the tool, surface the goodbye line.
    if saw_end_call:
        try:
            final_message = (json.loads(tool_args or "{}").get("final_message") or "").strip()
        except json.JSONDecodeError:
            final_message = ""
        if final_message:
            piece_count += 1
            yield final_message

    # Probe for token counts now that streaming is done.
    try:
        _, probe = chat_completion(
            messages, model, endpoint_id, temperature, max_tokens=1, tools=tools
        )
    except Exception:
        probe = {}
    prompt_tokens = probe.get("prompt_tokens")
    details = probe.get("prompt_tokens_details") or {}
    cached_tokens = details.get("cached_tokens")
    total = (prompt_tokens + piece_count) if prompt_tokens is not None else None

    raw_response = {"content": content_buf or None}
    if tool_name:
        raw_response["tool_calls"] = [
            {"function": {"name": tool_name, "arguments": tool_args}}
        ]

    yield _empty_metrics(
        prompt_tokens=prompt_tokens,
        completion_tokens=piece_count,
        total_tokens=total,
        cached_tokens=cached_tokens,
        latency_s=ttfb,
        tool_called=tool_name,
        raw_response=raw_response,
    )


def _stream_custom(backend, messages, temperature, max_tokens):
    tools = backend.get("tools", config.TOOLS)
    yield from stream_chat(
        messages,
        backend["model"],
        backend["endpoint_id"],
        temperature,
        max_tokens,
        tools=tools,
    )


# ---------------------------------------------------------------------------
# Azure OpenAI backend
# ---------------------------------------------------------------------------


def _azure_client(backend):
    from openai import OpenAI

    key = os.environ.get(backend.get("key_env", "AZURE_API_KEY"), "").strip()
    if not key:
        raise BackendError(
            f"{backend.get('key_env', 'AZURE_API_KEY')} is not set. "
            "Add it to .env or a Space Secret to use the Azure backend."
        )
    endpoint = (backend.get("endpoint") or "").rstrip("/")
    if not endpoint:
        raise BackendError("Azure endpoint is not configured (AZURE_ENDPOINT).")
    return OpenAI(api_key=key, base_url=endpoint + "/openai/v1/")


def _stream_azure(backend, messages, temperature, max_tokens):
    client = _azure_client(backend)
    t_start = time.monotonic()
    usage = None

    content_buf = ""
    request_payload = {
        "model": backend["deployment"],
        "messages": messages,
        "temperature": float(temperature),
        "max_completion_tokens": int(max_tokens),
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    yield {"__request__": True, "payload": request_payload}
    stream = client.chat.completions.create(**request_payload)
    for chunk in stream:
        if getattr(chunk, "usage", None):
            usage = chunk.usage
        choices = chunk.choices or []
        if not choices:
            continue
        delta = choices[0].delta
        piece = getattr(delta, "content", None) if delta else None
        if piece:
            content_buf += piece
            yield piece

    prompt_tokens = completion_tokens = total_tokens = cached_tokens = None
    if usage is not None:
        prompt_tokens = getattr(usage, "prompt_tokens", None)
        completion_tokens = getattr(usage, "completion_tokens", None)
        total_tokens = getattr(usage, "total_tokens", None)
        details = getattr(usage, "prompt_tokens_details", None)
        if details is not None:
            cached_tokens = getattr(details, "cached_tokens", None)

    yield _empty_metrics(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total_tokens,
        cached_tokens=cached_tokens,
        latency_s=round(time.monotonic() - t_start, 3),
        raw_response={"content": content_buf or None},
    )


# ---------------------------------------------------------------------------
# Unified dispatch
# ---------------------------------------------------------------------------


def stream_backend(backend: dict, messages: list, temperature: float, max_tokens: int):
    """Backend-agnostic streaming generator.

    Yields: str deltas, then a final {"__metrics__", ...} sentinel. Any backend
    failure is caught and surfaced as a metrics sentinel with `error` set (the
    worker never raises).
    """
    btype = backend.get("type")
    try:
        if btype == "custom":
            yield from _stream_custom(backend, messages, temperature, max_tokens)
        elif btype == "azure":
            yield from _stream_azure(backend, messages, temperature, max_tokens)
        else:
            yield _empty_metrics(error=f"Unknown backend type: {btype}")
    except Exception as e:  # noqa: BLE001 - surface per-backend errors in the UI
        yield _empty_metrics(error=str(e))