ask_ai.py

"""Lightweight AskAI wrapper with optional transformer support.

Provides an AskAI class that uses a local Python fallback responder
or, if available, a Hugging Face transformers pipeline for generation.
"""
from typing import Optional
try:
    # Prefer explicit model + tokenizer for more controllable generation
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
    import torch
    _HAS_TRANSFORMERS = True
except Exception:
    _HAS_TRANSFORMERS = False

try:
    from . import responseai
except Exception:
    import responseai

try:
    from .cpp_client import client as cpp_client
except Exception:
    cpp_client = None


class AskAI:
    def __init__(self, model_name: str = "google/flan-t5-small", use_cpp_backend: bool = True):
        self.model_name = model_name
        self.use_cpp_backend = use_cpp_backend and (cpp_client is not None and cpp_client.available())
        self.tokenizer = None
        self.model = None
        if _HAS_TRANSFORMERS:
            try:
                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
                self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
                if torch.cuda.is_available():
                    self.model = self.model.to("cuda")
            except Exception:
                self.model = None

    def ask(self, prompt: str, max_length: int = 256) -> str:
        """Return an answer for the given prompt.

        Priority:
        1. Use local Transformers model (Flan-T5) if available for more rational answers.
        2. If configured and available, ask the C++ backend via IPC.
        3. Fallback to `responseai.generate_response`.
        """
        # 1: Transformers
        if self.model is not None and self.tokenizer is not None:
            try:
                inputs = self.tokenizer(prompt, return_tensors="pt")
                if next(self.model.parameters()).is_cuda:
                    inputs = {k: v.to("cuda") for k, v in inputs.items()}
                out = self.model.generate(**inputs, max_length=max_length, do_sample=False)
                text = self.tokenizer.decode(out[0], skip_special_tokens=True)
                if text:
                    return text
            except Exception:
                pass

        # 2: C++ backend IPC
        if self.use_cpp_backend and cpp_client is not None:
            try:
                return cpp_client.ask(prompt)
            except Exception:
                pass

        # 3: fallback rule-based responder
        return responseai.generate_response(prompt)


__all__ = ["AskAI"]