src/llm_client.py

from __future__ import annotations

import json
import os
import re
from typing import Any

from huggingface_hub import InferenceClient

from src.prompts import SYSTEM_PROMPT, ranking_prompt
from src.schemas import Opportunity, SourcePage, UserProfile


class LLMError(RuntimeError):
    pass


class LLMClient:
    """Ranks funding opportunities via the Hugging Face Inference API."""

    def __init__(self, api_key: str | None = None, model: str | None = None, timeout: int = 60) -> None:
        self.api_key = os.getenv("HF_TOKEN", "") if api_key is None else api_key
        self.model = model or os.getenv("LLM_MODEL", "deepseek-ai/DeepSeek-V4-Flash:fireworks-ai")
        self.timeout = timeout

    @property
    def enabled(self) -> bool:
        return bool(self.api_key)

    def rank_opportunities(self, profile: UserProfile, pages: list[SourcePage]) -> list[Opportunity]:
        if not self.enabled:
            return []
        client = InferenceClient(api_key=self.api_key, timeout=self.timeout)
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": ranking_prompt(profile, pages)},
        ]
        try:
            completion = self._complete(client, messages, json_mode=True)
        except Exception as first_exc:
            try:
                completion = self._complete(client, messages + [
                    {"role": "system", "content": "Respond with ONLY valid JSON, no surrounding prose."}
                ], json_mode=False)
            except Exception as retry_exc:
                raise LLMError(f"Inference request failed: {retry_exc}") from first_exc
        try:
            content = completion.choices[0].message.content
        except (AttributeError, IndexError, TypeError) as exc:
            raise LLMError("Inference returned an unexpected response shape.") from exc
        return parse_opportunities(content)

    def _complete(self, client: InferenceClient, messages: list[dict[str, str]], json_mode: bool) -> Any:
        kwargs: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "temperature": 0.2,
        }
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}
        return client.chat.completions.create(**kwargs)


def parse_opportunities(content: str) -> list[Opportunity]:
    data = _load_json(content)
    raw_items = data.get("opportunities") if isinstance(data, dict) else data
    if not isinstance(raw_items, list):
        raise LLMError("Model JSON did not contain an opportunities list.")
    opportunities: list[Opportunity] = []
    for item in raw_items:
        if not isinstance(item, dict):
            continue
        try:
            opportunities.append(Opportunity(**item))
        except ValueError:
            continue
    if not opportunities:
        raise LLMError("No valid opportunities were parsed from model output.")
    return sorted(opportunities, key=lambda opp: opp.match_score, reverse=True)


def _load_json(content: str) -> Any:
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        match = re.search(r"\{.*\}", content, flags=re.DOTALL)
        if not match:
            raise LLMError("Model response was not valid JSON.")
        return json.loads(match.group(0))