backend/gemma_client.py

"""
Gemma 4 client — dual-model architecture.
Agent: gemma-4-26B-A4B-it (MoE, generation)
Judge: gemma-4-31B-it (dense, evaluation)
"""
import os
from openai import OpenAI

_client = None


def get_client() -> OpenAI:
    global _client
    if _client is None:
        base_url = os.environ.get("GEMMA_BASE_URL", "https://api.featherless.ai/v1")
        api_key = os.environ.get("FEATHERLESS_API_KEY", "")
        _client = OpenAI(base_url=base_url, api_key=api_key)
    return _client


def get_agent_model() -> str:
    return os.environ.get("GEMMA_MODEL", "google/gemma-4-26B-A4B-it")


def get_judge_model() -> str:
    return os.environ.get("GEMMA_JUDGE_MODEL", "google/gemma-4-31B-it")


def chat_agent(messages: list[dict], temperature: float = 0.3, max_tokens: int = 1024) -> str:
    """Agent call — the model being tested."""
    resp = get_client().chat.completions.create(
        model=get_agent_model(),
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return resp.choices[0].message.content


def chat_judge(messages: list[dict], temperature: float = 0.1, max_tokens: int = 1024) -> str:
    """Judge call — the model evaluating the agent."""
    resp = get_client().chat.completions.create(
        model=get_judge_model(),
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return resp.choices[0].message.content


# Backward compat — defaults to agent
def chat(messages: list[dict], temperature: float = 0.3, max_tokens: int = 1024) -> str:
    return chat_agent(messages, temperature, max_tokens)