app/services/llm.py

"""LLM wrapper — Vertex AI Gemini (primary) with Groq/Llama fallback.

Auto-selects backend:
  - GCP_PROJECT set  →  Vertex AI (google-genai SDK, uses Application Default Credentials)
  - GCP_PROJECT unset →  Groq free tier (GROQ_API_KEY required)

Public interface (build_system_blocks / complete / LLMResponse) unchanged —
all four agents work with either backend.

Vertex AI model: gemini-2.0-flash-001
  Pricing: ~$0.075/1M input tokens — $5 credit ≈ 66M tokens ≈ hundreds of analyses.

Rate-limit guard: asyncio.Semaphore(settings.llm_concurrency).
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from functools import lru_cache
from typing import Any

from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_attempt, wait_exponential

from app.config import get_settings

logger = logging.getLogger(__name__)

_semaphore: asyncio.Semaphore | None = None


def _get_semaphore() -> asyncio.Semaphore:
    global _semaphore
    if _semaphore is None:
        _semaphore = asyncio.Semaphore(get_settings().llm_concurrency)
    return _semaphore


@lru_cache(maxsize=1)
def _vertex_client():
    from google import genai  # type: ignore[import]

    s = get_settings()
    kwargs: dict = {"vertexai": True, "project": s.gcp_project, "location": s.gcp_location}

    if s.google_application_credentials:
        from google.oauth2 import service_account  # type: ignore[import]
        import os, pathlib
        key_path = pathlib.Path(s.google_application_credentials)
        if not key_path.is_absolute():
            # Resolve relative to backend/ dir (where .env lives)
            key_path = pathlib.Path(__file__).parents[2] / key_path
        creds = service_account.Credentials.from_service_account_file(
            str(key_path),
            scopes=["https://www.googleapis.com/auth/cloud-platform"],
        )
        kwargs["credentials"] = creds

    return genai.Client(**kwargs)


@lru_cache(maxsize=1)
def _groq_client():
    from groq import Groq  # type: ignore[import]

    s = get_settings()
    if not s.groq_api_key:
        raise RuntimeError("Neither GCP_PROJECT (Vertex AI) nor GROQ_API_KEY is set.")
    return Groq(api_key=s.groq_api_key)


# ── Shared types ──────────────────────────────────────────────────────────────

@dataclass
class LLMUsage:
    input_tokens: int = 0
    output_tokens: int = 0
    cache_creation_input_tokens: int = 0
    cache_read_input_tokens: int = 0

    def __add__(self, other: "LLMUsage") -> "LLMUsage":
        return LLMUsage(
            input_tokens=self.input_tokens + other.input_tokens,
            output_tokens=self.output_tokens + other.output_tokens,
        )


@dataclass
class LLMResponse:
    text: str
    usage: LLMUsage
    elapsed_ms: int
    raw: Any = field(repr=False)

    def extract_json(self) -> Any:
        text = self.text.strip()
        if text.startswith("```"):
            lines = text.split("\n")
            inner = []
            for line in lines[1:]:
                if line.strip() == "```":
                    break
                inner.append(line)
            text = "\n".join(inner).strip()

        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass

        for opener, closer in (("{", "}"), ("[", "]")):
            start = text.find(opener)
            end = text.rfind(closer)
            if start != -1 and end > start:
                try:
                    return json.loads(text[start : end + 1])
                except json.JSONDecodeError:
                    continue

        raise ValueError(f"could not parse JSON: {self.text[:300]}…")


# ── Public helpers ────────────────────────────────────────────────────────────

def build_system_blocks(
    *,
    shared_preamble: str,
    document: str | None,
    role_instructions: str,
) -> list[dict[str, Any]]:
    blocks = [{"type": "preamble", "text": shared_preamble}]
    if document:
        blocks.append({"type": "document", "text": document})
    blocks.append({"type": "role", "text": role_instructions})
    return blocks


async def complete(
    *,
    system_blocks: list[dict[str, Any]],
    user_prompt: str,
    max_tokens: int = 2048,
    temperature: float = 0.2,
    model: str | None = None,
) -> LLMResponse:
    settings = get_settings()

    preamble = next((b["text"] for b in system_blocks if b["type"] == "preamble"), "")
    doc = next((b["text"] for b in system_blocks if b["type"] == "document"), "")
    role = next((b["text"] for b in system_blocks if b["type"] == "role"), "")

    system_content = preamble
    if doc:
        system_content += f"\n\n<contract>\n{doc}\n</contract>"

    started = time.perf_counter()

    if settings.use_vertex:
        result = await _call_vertex(
            system_content=system_content,
            role=role,
            user_prompt=user_prompt,
            model=model or settings.vertex_model,
            max_tokens=max_tokens,
            temperature=temperature,
        )
    else:
        result = await _call_groq(
            system_content=system_content,
            role=role,
            user_prompt=user_prompt,
            model=model or settings.groq_model,
            max_tokens=max_tokens,
            temperature=temperature,
        )

    elapsed_ms = int((time.perf_counter() - started) * 1000)
    logger.info(
        "llm.complete backend=%s model=%s elapsed_ms=%d in=%d out=%d",
        "vertex" if settings.use_vertex else "groq",
        model or settings.active_model,
        elapsed_ms,
        result.usage.input_tokens,
        result.usage.output_tokens,
    )
    result.elapsed_ms = elapsed_ms
    return result


async def _call_vertex(
    *,
    system_content: str,
    role: str,
    user_prompt: str,
    model: str,
    max_tokens: int,
    temperature: float,
) -> LLMResponse:
    from google import genai  # type: ignore[import]
    from google.genai import types  # type: ignore[import]

    client = _vertex_client()
    full_user = f"{role}\n\n{user_prompt}" if role else user_prompt
    config = types.GenerateContentConfig(
        system_instruction=system_content,
        temperature=temperature,
        max_output_tokens=max_tokens,
    )

    async with _get_semaphore():
        async for attempt in AsyncRetrying(
            stop=stop_after_attempt(3),
            wait=wait_exponential(multiplier=2.0, min=2.0, max=20.0),
            retry=retry_if_exception_type(Exception),
            reraise=True,
        ):
            with attempt:
                loop = asyncio.get_event_loop()
                response = await loop.run_in_executor(
                    None,
                    lambda: client.models.generate_content(
                        model=model, contents=full_user, config=config
                    ),
                )

    text = response.text.strip() if response.text else ""
    meta = getattr(response, "usage_metadata", None)
    usage = LLMUsage(
        input_tokens=getattr(meta, "prompt_token_count", 0) or 0,
        output_tokens=getattr(meta, "candidates_token_count", 0) or 0,
    )
    return LLMResponse(text=text, usage=usage, elapsed_ms=0, raw=response)


async def _call_groq(
    *,
    system_content: str,
    role: str,
    user_prompt: str,
    model: str,
    max_tokens: int,
    temperature: float,
) -> LLMResponse:
    client = _groq_client()
    system_full = f"{system_content}\n\n{role}" if role else system_content
    messages = [
        {"role": "system", "content": system_full},
        {"role": "user", "content": user_prompt},
    ]

    async with _get_semaphore():
        async for attempt in AsyncRetrying(
            stop=stop_after_attempt(5),
            wait=wait_exponential(multiplier=2.0, min=8.0, max=60.0),
            retry=retry_if_exception_type(Exception),
            reraise=True,
        ):
            with attempt:
                loop = asyncio.get_event_loop()
                response = await loop.run_in_executor(
                    None,
                    lambda: client.chat.completions.create(
                        model=model, messages=messages,
                        temperature=temperature, max_tokens=max_tokens,
                    ),
                )

    text = response.choices[0].message.content or ""
    u = response.usage
    usage = LLMUsage(
        input_tokens=u.prompt_tokens if u else 0,
        output_tokens=u.completion_tokens if u else 0,
    )
    return LLMResponse(text=text, usage=usage, elapsed_ms=0, raw=response)