"""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))