research/llm.py

"""Thin wrapper around HuggingFace Inference Providers (chat completions).

Reads ``HF_TOKEN`` from the environment and exposes a single ``chat()`` helper used by
the research graph. The model id is configurable via ``LLM_MODEL``.
"""

from __future__ import annotations

import os
from functools import lru_cache
from typing import List, Dict

from huggingface_hub import InferenceClient

# Default kept under 32B params (project constraint). Override via the LLM_MODEL env var.
# Qwen2.5-14B-Instruct was dropped by HF Inference Providers ("not supported by any
# provider you have enabled"), so we use the 7B sibling — same family/prompt behavior,
# still served. For higher-quality scripts set LLM_MODEL=google/gemma-3-27b-it (also
# <32B and currently available). Avoid Qwen3 "thinking" variants: they emit <think>
# traces that break the strict `Speaker: text` script format.
DEFAULT_MODEL = os.environ.get("LLM_MODEL", "Qwen/Qwen2.5-7B-Instruct")


@lru_cache(maxsize=1)
def _client() -> InferenceClient:
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not token:
        raise RuntimeError(
            "HF_TOKEN is not set. Add it as a Space secret (or export it locally) so the "
            "research agents can call HuggingFace Inference Providers."
        )
    return InferenceClient(token=token)


def chat(
    messages: List[Dict[str, str]],
    *,
    model: str | None = None,
    temperature: float = 0.7,
    max_tokens: int = 2048,
) -> str:
    """Run a chat completion and return the assistant text."""
    resp = _client().chat_completion(
        messages=messages,
        model=model or DEFAULT_MODEL,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return resp.choices[0].message.content or ""


def complete(system: str, user: str, **kwargs) -> str:
    """Convenience for a single system+user turn."""
    return chat(
        [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        **kwargs,
    )