backend/ai/client.py

from __future__ import annotations

import json
from typing import Any

import httpx

from backend.config import get_settings


class GroqAIClient:
    def __init__(self) -> None:
        self.settings = get_settings()

    def generate_json(self, prompt: str) -> dict[str, Any]:
        text = self.generate_text(prompt)
        return _extract_json_object(text)

    def generate_text(self, prompt: str) -> str:
        settings = get_settings()
        response = httpx.post(
            f"{settings.groq_base_url.rstrip('/')}/chat/completions",
            headers={
                "Authorization": f"Bearer {settings.groq_api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": settings.groq_model,
                "messages": [
                    {
                        "role": "system",
                        "content": (
                            "You are a precise backend assistant. "
                            "Follow the prompt exactly and return valid JSON when requested."
                        ),
                    },
                    {"role": "user", "content": prompt},
                ],
                "temperature": 0.2,
            },
            timeout=settings.groq_timeout_seconds,
        )
        response.raise_for_status()
        payload = response.json()
        try:
            return payload["choices"][0]["message"]["content"].strip()
        except (KeyError, IndexError, AttributeError) as exc:
            raise RuntimeError("Groq returned an unexpected response shape") from exc

    def is_ready(self) -> bool:
        settings = get_settings()
        return settings.ai_provider == "groq" and bool(settings.groq_api_key.strip())


def _extract_json_object(text: str) -> dict[str, Any]:
    cleaned = text.strip()
    if cleaned.startswith("```"):
        parts = cleaned.split("```")
        cleaned = next((part for part in parts if "{" in part and "}" in part), cleaned)
        cleaned = cleaned.replace("json", "", 1).strip()
    start = cleaned.find("{")
    end = cleaned.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise ValueError("No JSON object found in model response")
    return json.loads(cleaned[start : end + 1])


llm_client = GroqAIClient()


def render_prompt(template: str, **kwargs: Any) -> str:
    rendered = template
    for key, value in kwargs.items():
        rendered = rendered.replace(f"{{{key}}}", str(value))
    return rendered