siemens_ragas/rag_eval/metrics/factory.py

"""Factories for OpenAI-backed RAGAS models and metric pipelines."""

from __future__ import annotations

import logging
from typing import Any

import httpx
from openai import AsyncOpenAI

from rag_eval.compat import ensure_ragas_import_compat
from rag_eval.settings import EvaluationSettings
from rag_eval.shared.models import Scenario

from .token_tracker import get_current_tracker

ensure_ragas_import_compat()

from ragas.embeddings.base import embedding_factory
from ragas.llms import llm_factory
from ragas.metrics.collections import (
    AnswerRelevancy,
    ContextPrecision,
    ContextRecall,
    FactualCorrectness,
    Faithfulness,
    NoiseSensitivity,
    SemanticSimilarity,
)

from .judge_prompts import localize_pipeline_prompts
from .pipeline import MetricPipeline


logger = logging.getLogger("rag_eval.metrics.factory")


async def _usage_response_hook(response: httpx.Response) -> None:
    """Record token usage from an OpenAI-compatible HTTP response, if a tracker is active.

    Applies to both chat-completions and embeddings responses since both
    return top-level `model` and `usage` fields in OpenAI-compatible APIs.
    Never raises — a broken/incompatible gateway response must not affect scoring.
    """
    tracker = get_current_tracker()
    if tracker is None:
        return
    try:
        await response.aread()
        data = response.json()
        usage = data.get("usage")
        if not usage:
            # Gateway did not report usage at all — skip rather than record a
            # misleading 0/0 call.
            return
        model = data.get("model") or "unknown"
        tracker.record(
            model,
            int(usage.get("prompt_tokens", 0) or 0),
            int(usage.get("completion_tokens", 0) or 0),
        )
    except Exception:  # noqa: BLE001
        logger.debug("[factory] usage hook failed to parse response", exc_info=True)


def attach_usage_hook(client: AsyncOpenAI) -> None:
    """Attach the token-usage response hook to an AsyncOpenAI client (idempotent).

    Safe to call multiple times on the same client (e.g. when judge and
    embedding models share one client) — the hook is only appended once.
    """
    httpx_client = getattr(client, "_client", None)
    if httpx_client is None or not hasattr(httpx_client, "event_hooks"):
        return
    hooks = httpx_client.event_hooks.setdefault("response", [])
    if _usage_response_hook not in hooks:
        hooks.append(_usage_response_hook)


def _resolve_openai_client_kwargs(
    model: str,
    settings: EvaluationSettings,
) -> dict[str, Any]:
    """Return AsyncOpenAI kwargs, preferring a matching LLM Profile over .env settings.

    Lookup order:
      1. LLM Profile whose model name equals `model` (exact match)
      2. Fall back to EvaluationSettings (.env)
    """
    try:
        # Lazy import to avoid circular dependency (webapp -> rag_eval is one-way).
        from webapp.services.profile_manager import profile_manager
        profiles = profile_manager.list_all()
        for profile in profiles:
            if profile.model == model:
                kwargs: dict[str, Any] = {
                    "api_key": profile.api_key or "sk-placeholder",
                    "timeout": float(profile.timeout_seconds or 30),
                }
                if profile.base_url and profile.base_url.strip():
                    kwargs["base_url"] = profile.base_url.strip()
                logger.debug(
                    "[factory] model=%s  source=profile  base_url=%s",
                    model, kwargs.get("base_url", "(not set, using default)")
                )
                return kwargs
    except Exception as exc:  # noqa: BLE001
        # If profile lookup fails for any reason, fall through to .env settings.
        logger.warning("[factory] profile lookup failed for model=%s: %s", model, exc)

    fallback = settings.openai_client_kwargs
    logger.debug(
        "[factory] model=%s  source=env  base_url=%s",
        model, fallback.get("base_url", "(not set)")
    )
    return fallback


def resolve_openai_client_kwargs(
    judge_model: str,
    settings: EvaluationSettings,
) -> dict[str, Any]:
    """Public accessor for profile-aware AsyncOpenAI kwargs (matched by model name).

    Exposed so other components (e.g. the optimization advisor's direct LLM call)
    can build a client that honors the same saved-profile/.env resolution used by
    the scoring pipeline, instead of duplicating the lookup logic.
    """
    return _resolve_openai_client_kwargs(judge_model, settings)


def build_metric_registry(llm: Any, embeddings: Any) -> dict[str, Any]:
    """Instantiate the full set of supported RAGAS metrics keyed by canonical name.

    Shared by the scenario pipeline, the inline scorer, and the prompt-cache
    bootstrap so the metric set is defined in exactly one place.
    """
    return {
        "faithfulness": Faithfulness(llm=llm),
        "answer_relevancy": AnswerRelevancy(llm=llm, embeddings=embeddings),
        "context_recall": ContextRecall(llm=llm),
        "context_precision": ContextPrecision(llm=llm),
        # NoiseSensitivity mode='relevant': sensitivity to noise from relevant contexts.
        "noise_sensitivity": NoiseSensitivity(llm=llm),
        # FactualCorrectness mode='f1': balances claim precision and recall vs. ground truth.
        "factual_correctness": FactualCorrectness(llm=llm),
        # SemanticSimilarity: embedding cosine between answer and ground truth (no LLM call).
        "semantic_similarity": SemanticSimilarity(embeddings=embeddings),
    }


def build_models(
    judge_model: str,
    embedding_model: str,
    settings: EvaluationSettings,
) -> tuple[Any, Any]:
    """Create the LLM and embedding clients required by the selected RAGAS metrics.

    Resolves connection settings independently for the judge LLM and the embedding
    model by looking up each in the stored LLM Profiles (matched by model name).
    This allows judge_model and embedding_model to use different gateways / API keys.
    Falls back to .env settings when no matching profile is found.
    """
    llm_kwargs = _resolve_openai_client_kwargs(judge_model, settings)
    emb_kwargs = _resolve_openai_client_kwargs(embedding_model, settings)

    logger.info(
        "[factory] build_models  judge=%s→%s  embedding=%s→%s",
        judge_model, llm_kwargs.get("base_url", "(env default)"),
        embedding_model, emb_kwargs.get("base_url", "(env default)"),
    )

    llm_client = AsyncOpenAI(**llm_kwargs)
    attach_usage_hook(llm_client)
    # Only allocate a second client when the embedding model needs different settings.
    emb_client = AsyncOpenAI(**emb_kwargs) if emb_kwargs != llm_kwargs else llm_client
    attach_usage_hook(emb_client)

    # RAGAS structured-output judge calls can be truncated by the upstream default
    # 1024 completion budget, especially for faithfulness and GPT-5 family models.
    llm = llm_factory(
        judge_model,
        client=llm_client,
        max_tokens=max(1, int(settings.ragas_llm_max_tokens)),
    )
    embeddings = embedding_factory(provider="openai", model=embedding_model, client=emb_client)
    return llm, embeddings


def build_metric_pipeline(
    scenario: Scenario,
    settings: EvaluationSettings,
    llm: Any | None = None,
    embeddings: Any | None = None,
) -> MetricPipeline:
    """Build a metric pipeline containing only the metrics requested by the scenario.

    If llm and embeddings are provided (pre-built by the caller), they are reused.
    Otherwise, new instances are created from scenario + settings.
    """
    if llm is None or embeddings is None:
        llm, embeddings = build_models(
            scenario.judge_model,
            scenario.embedding_model,
            settings,
        )

    # Build the full registry once using the shared factory, then slice by requested metrics.
    registry = build_metric_registry(llm, embeddings)
    selected = {name: registry[name] for name in scenario.metrics}
    # Apply language-specific judge prompt localization (no-op for 'en').
    language = scenario.judge_language or settings.ragas_judge_language or "en"
    localize_pipeline_prompts(selected, language)
    return MetricPipeline(
        metrics=selected,
        metric_timeout_seconds=settings.ragas_metric_timeout_seconds,
    )