llm.py

"""Provider abstraction — one public function: complete()."""
import json
import os
import re
from typing import Optional

import config

# Lazy-initialised clients
_anthropic_client = None
_hf_client = None
_modal_session = None


def _get_anthropic():
    global _anthropic_client
    if _anthropic_client is None:
        import anthropic
        _anthropic_client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
    return _anthropic_client


def _get_hf():
    global _hf_client
    if _hf_client is None:
        from huggingface_hub import InferenceClient
        _hf_client = InferenceClient(
            token=os.getenv("HF_TOKEN"),
            provider=config.HF_INFERENCE_PROVIDER,
        )
    return _hf_client


def _get_modal_session():
    """Return a requests.Session pointed at the Modal vLLM endpoint."""
    global _modal_session
    if _modal_session is None:
        import requests

        if not config.MODAL_ENDPOINT_URL:
            raise ValueError(
                "MODAL_ENDPOINT_URL is not set. "
                "Deploy modal_serve.py first and add the URL to your .env."
            )
        session = requests.Session()
        session.headers.update({"Content-Type": "application/json"})
        _modal_session = session
    return _modal_session


def _strip_json_fences(text: str) -> str:
    """Remove ```json ... ``` or ``` ... ``` wrappers."""
    text = text.strip()
    text = re.sub(r"^```(?:json)?\s*", "", text)
    text = re.sub(r"\s*```$", "", text)
    return text.strip()


def _call_anthropic(system: str, user: str, max_tokens: int) -> str:
    client = _get_anthropic()
    response = client.messages.create(
        model=config.ANTHROPIC_MODEL,
        max_tokens=max_tokens,
        system=system,
        messages=[{"role": "user", "content": user}],
    )
    return response.content[0].text


def _call_anthropic_with_retry(system: str, messages: list, max_tokens: int) -> str:
    client = _get_anthropic()
    response = client.messages.create(
        model=config.ANTHROPIC_MODEL,
        max_tokens=max_tokens,
        system=system,
        messages=messages,
    )
    return response.content[0].text


def _call_hf(system: str, user: str, max_tokens: int) -> str:
    client = _get_hf()
    result = client.chat_completion(
        model=config.HF_MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        max_tokens=max_tokens,
    )
    return result.choices[0].message.content


def _call_hf_with_retry(system: str, messages: list, max_tokens: int) -> str:
    client = _get_hf()
    hf_messages = [{"role": "system", "content": system}] + messages
    result = client.chat_completion(
        model=config.HF_MODEL,
        messages=hf_messages,
        max_tokens=max_tokens,
    )
    return result.choices[0].message.content


def _call_modal(system: str, user: str, max_tokens: int) -> str:
    session = _get_modal_session()
    payload = {
        "model": config.MODAL_MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "max_tokens": max_tokens,
    }
    resp = session.post(
        f"{config.MODAL_ENDPOINT_URL.rstrip('/')}/v1/chat/completions",
        json=payload,
        timeout=120,
    )
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"]


def _call_modal_with_retry(system: str, messages: list, max_tokens: int) -> str:
    session = _get_modal_session()
    payload = {
        "model": config.MODAL_MODEL,
        "messages": [{"role": "system", "content": system}] + messages,
        "max_tokens": max_tokens,
    }
    resp = session.post(
        f"{config.MODAL_ENDPOINT_URL.rstrip('/')}/v1/chat/completions",
        json=payload,
        timeout=120,
    )
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"]


def complete(
    system: str,
    user: str,
    json_mode: bool = False,
    max_tokens: int = 800,
) -> str:
    """
    Call the configured LLM provider and return the response text.

    If json_mode=True:
      - Appends a JSON instruction to the system prompt.
      - Strips markdown fences from the response.
      - Retries once with a correction message if json.loads fails.
      - Raises ValueError if the retry also fails.
    """
    if json_mode:
        system = (
            system
            + "\n\nIMPORTANT: Respond ONLY with valid JSON. "
            "No markdown fences, no preamble, no trailing commentary."
        )

    provider = config.LLM_PROVIDER

    if provider == "anthropic":
        raw = _call_anthropic(system, user, max_tokens)
    elif provider == "hf":
        raw = _call_hf(system, user, max_tokens)
    elif provider == "modal":
        raw = _call_modal(system, user, max_tokens)
    else:
        raise ValueError(f"Unknown LLM_PROVIDER: {provider!r}")

    if not json_mode:
        return raw

    # --- JSON mode: attempt parse then retry once ---
    cleaned = _strip_json_fences(raw)
    try:
        json.loads(cleaned)
        return cleaned
    except json.JSONDecodeError:
        pass

    # Retry with correction
    messages = [
        {"role": "user", "content": user},
        {"role": "assistant", "content": raw},
        {
            "role": "user",
            "content": (
                "Your last output was invalid JSON. "
                "Fix it and return ONLY valid JSON with no markdown fences, no preamble."
            ),
        },
    ]
    if provider == "anthropic":
        raw2 = _call_anthropic_with_retry(system, messages, max_tokens)
    elif provider == "modal":
        raw2 = _call_modal_with_retry(system, messages, max_tokens)
    else:
        raw2 = _call_hf_with_retry(system, messages, max_tokens)

    cleaned2 = _strip_json_fences(raw2)
    try:
        json.loads(cleaned2)
        return cleaned2
    except json.JSONDecodeError as exc:
        raise ValueError(
            f"LLM returned invalid JSON after retry.\nRaw output:\n{raw2}"
        ) from exc