llm.py

"""
LLM wrapper with token accounting and CO2 emission estimation.

Wraps an Azure OpenAI-compatible chat/completions call and returns:
- content: the generated text
- tokens: prompt / completion / cached / total
- energy_kwh, co2_grams: environmental impact for the call

CO2 emissions are estimated with `ecologits`, which uses a model registry
(parameter counts, hardware assumptions) plus output token count and request
latency to compute global warming potential (kgCO2eq) and energy (kWh).
We chose `ecologits` over `codecarbon` because the LLM runs on a remote
endpoint — `codecarbon` measures local process energy and would only count
the client overhead, not the inference itself.
"""

import logging
import time
from typing import Optional

import requests
from fastapi import HTTPException

logger = logging.getLogger(__name__)

try:
    from ecologits.tracers.utils import llm_impacts
    _ECOLOGITS_AVAILABLE = True
except ImportError:
    _ECOLOGITS_AVAILABLE = False
    logger.warning("ecologits not installed — CO2 emission will be reported as None.")


def _to_scalar(value) -> Optional[float]:
    """
    Normalize an ecologits impact value to a single float.

    ecologits returns either a plain float or a RangeValue(min, max) depending
    on the model registry entry. For RangeValue we return the midpoint so a
    single representative number flows through the API/UI.
    """
    if value is None:
        return None
    if hasattr(value, "min") and hasattr(value, "max"):
        return (float(value.min) + float(value.max)) / 2.0
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def call_llm(
    prompt: str,
    *,
    endpoint_url: str,
    api_key: str,
    model: str,
    max_completion_tokens: int = 512,
    temperature: float = 0.7,
    top_p: float = 0.95,
    provider: str = "openai",
    timeout: Optional[float] = None,
) -> dict:
    """
    Call an Azure OpenAI-compatible chat/completions endpoint and return the
    response together with token counts and CO2 emission estimate.

    Returns a dict:
        {
            "content": str,
            "tokens": {
                "prompt": int,
                "completion": int,
                "cached": int,
                "total": int,
            },
            "energy_kwh": float | None,
            "co2_grams": float | None,
            "latency_s": float,
        }
    """
    headers = {
        "api-key": api_key,
        "Content-Type": "application/json",
    }
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_completion_tokens": max_completion_tokens,
        "temperature": temperature,
        "top_p": top_p,
    }

    start = time.perf_counter()
    try:
        resp = requests.post(endpoint_url, headers=headers, json=payload, timeout=timeout)
        resp.raise_for_status()
        data = resp.json()
    except requests.exceptions.HTTPError as e:
        logger.error(f"LLM API call failed: {e} — {resp.text}")
        raise HTTPException(status_code=503, detail=f"LLM service unavailable: {str(e)}")
    except (requests.exceptions.JSONDecodeError, ValueError):
        logger.error(
            f"LLM API returned non-JSON response (status {resp.status_code}): {repr(resp.text)}"
        )
        raise HTTPException(status_code=502, detail="LLM service returned an invalid response")
    latency_s = time.perf_counter() - start

    try:
        content = data["choices"][0]["message"]["content"].strip()
    except (KeyError, IndexError) as e:
        logger.error(f"Unexpected LLM response format: {e} — body: {data!r}")
        raise HTTPException(status_code=502, detail="Unexpected response from LLM service")

    usage = data.get("usage") or {}
    prompt_tokens = usage.get("prompt_tokens", 0)
    completion_tokens = usage.get("completion_tokens", 0)
    total_tokens = usage.get("total_tokens", prompt_tokens + completion_tokens)

    prompt_details = usage.get("prompt_tokens_details") or {}
    cached_tokens = prompt_details.get("cached_tokens", usage.get("cached_tokens", 0))

    energy_kwh: Optional[float] = None
    co2_grams: Optional[float] = None
    if _ECOLOGITS_AVAILABLE:
        try:
            impacts = llm_impacts(
                provider=provider,
                model_name=model,
                output_token_count=completion_tokens,
                request_latency=latency_s,
            )
            if impacts is not None:
                energy_kwh = _to_scalar(impacts.energy.value)
                # ecologits returns gwp in kgCO2eq; convert to grams
                gwp_kg = _to_scalar(impacts.gwp.value)
                co2_grams = gwp_kg * 1000.0 if gwp_kg is not None else None
        except Exception as e:
            logger.warning(f"ecologits impact calc failed for model={model}: {e}")

    return {
        "content": content,
        "tokens": {
            "prompt": prompt_tokens,
            "completion": completion_tokens,
            "cached": cached_tokens,
            "total": total_tokens,
        },
        "energy_kwh": energy_kwh,
        "co2_grams": co2_grams,
        "latency_s": latency_s,
    }