src/utils/llm.py

"""
llm.py
-------
Shared local-LLM wrapper used by the Planner, Reasoning, and Final-Answer
agents for text generation.

Primary path: google/flan-t5-base, a real open-source instruction-tuned
Hugging Face model (~250M params) that runs on free Hugging Face Space CPU
hardware in a few seconds per call. No paid API key required.

Fallback path: if the model cannot be downloaded (offline sandbox, first
run before the HF cache is warm, etc.) agents fall back to their own
deterministic, template/extractive logic instead of calling this class.
Every agent module documents its specific fallback. This wrapper simply
reports whether a real generative model is available so callers can decide.
"""
from __future__ import annotations

import logging
import os
from typing import List, Optional

logger = logging.getLogger(__name__)

GEN_MODEL_NAME = os.environ.get("RESEARCHPILOT_GEN_MODEL", "google/flan-t5-base")


class LocalLLM:
    def __init__(self, model_name: str = GEN_MODEL_NAME):
        self.model_name = model_name
        self._pipe = None
        self.available = False
        try:
            from transformers import pipeline

            self._pipe = pipeline("text2text-generation", model=model_name, device=-1)
            self.available = True
            logger.info("Loaded real generation model: %s", model_name)
        except Exception as exc:  # noqa: BLE001
            logger.warning(
                "Could not load generation model %s (%s). "
                "Agents will use their deterministic fallback logic instead.",
                model_name,
                exc,
            )

    def generate(self, prompt: str, max_new_tokens: int = 200) -> Optional[str]:
        if not self.available:
            return None
        try:
            out = self._pipe(prompt, max_new_tokens=max_new_tokens, do_sample=False)
            return out[0]["generated_text"].strip()
        except Exception as exc:  # noqa: BLE001
            logger.warning("Generation call failed (%s); falling back.", exc)
            return None


_SINGLETON: Optional[LocalLLM] = None


def get_llm() -> LocalLLM:
    global _SINGLETON
    if _SINGLETON is None:
        _SINGLETON = LocalLLM()
    return _SINGLETON