agents/cbt_agent.py

"""
CBT Agent implementation using Hugging Face Inference API directly.

This avoids smolagents version drift and works cleanly on Spaces.
Adds optional provider, billing, and timeout support plus true
token streaming for lower-latency responses in Spaces.
"""

import os
from collections.abc import Iterable

from huggingface_hub import InferenceClient

try:  # Newer hub versions expose a specific timeout error
    from huggingface_hub import InferenceTimeoutError
except Exception:  # pragma: no cover

    class InferenceTimeoutError(Exception):  # type: ignore
        pass


class CBTAgent:
    """Agent that composes warm CBT responses using HF Inference API.

    Args:
        model_name: HF model id or endpoint URL to use.
        provider: Optional provider routing (e.g., "together", "groq", "hf-inference").
        bill_to: Optional org to bill (Enterprise Hub).
        timeout: Request timeout in seconds.
        use_local: Ignored; kept for backward compatibility.
    """

    def __init__(
        self,
        model_name: str | None = None,
        *,
        provider: str | None = None,
        bill_to: str | None = None,
        timeout: float | None = None,
        use_local: bool = False,  # use_local ignored
    ):
        self.model_name = model_name or os.getenv("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct")

        # Optional routing / billing / timeout via env overrides
        self.provider = provider or os.getenv("HF_PROVIDER")
        self.bill_to = bill_to or os.getenv("HF_BILL_TO")
        # Parse env timeout if not provided
        if timeout is None:
            try:
                timeout_env = os.getenv("HF_TIMEOUT_SECONDS")
                timeout = float(timeout_env) if timeout_env else None
            except Exception:
                timeout = None
        self.timeout = timeout

        token = os.getenv("HUGGINGFACEHUB_API_TOKEN") or os.getenv("HF_TOKEN")
        if not token:
            raise RuntimeError("HF token missing: set HF_TOKEN or HUGGINGFACEHUB_API_TOKEN")

        # Initialize client with explicit named params to satisfy typing
        self.client = InferenceClient(
            model=self.model_name,
            provider=(self.provider if self.provider is not None else None),  # type: ignore[arg-type]
            token=token,
            timeout=self.timeout,
            bill_to=self.bill_to,
        )

    def _get_system_prompt(self) -> str:
        return (
            "You are a compassionate CBT assistant. Identify likely cognitive distortions, "
            "validate feelings, and suggest balanced, realistic alternatives. Keep responses "
            "concise, supportive, and practical. This is self-help guidance, not therapy."
        )

    def _generate(self, prompt: str, messages: list[dict] | None = None) -> str:
        # Prefer chat-completions if available
        try:
            msgs = messages or [
                {"role": "system", "content": self._get_system_prompt()},
                {"role": "user", "content": prompt},
            ]
            if hasattr(self.client, "chat") and hasattr(self.client.chat, "completions"):
                resp = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=msgs,
                    max_tokens=400,
                    temperature=0.6,
                )
                msg = getattr(resp.choices[0], "message", None)
                return getattr(msg, "content", "") or ""
            if hasattr(self.client, "chat_completion"):
                resp = self.client.chat_completion(
                    messages=msgs,
                    max_tokens=400,
                    temperature=0.6,
                )
                choices = resp.get("choices") or []
                if choices:
                    return choices[0].get("message", {}).get("content", "")
        except Exception:
            pass

        # Fallback to plain text generation
        return self.client.text_generation(
            prompt if prompt else (messages[-1]["content"] if messages else ""),
            max_new_tokens=400,
            temperature=0.6,
            return_full_text=False,
        )

    def _build_messages(self, message: str, context: list[dict] | None = None) -> list[dict]:
        messages: list[dict] = [
            {"role": "system", "content": self._get_system_prompt()},
        ]
        if context:
            turns = context[-3:]
            for turn in turns:
                u = turn.get("user", "")
                a = turn.get("assistant", "")
                if u:
                    messages.append({"role": "user", "content": u})
                if a:
                    messages.append({"role": "assistant", "content": a})
        messages.append({"role": "user", "content": message})
        return messages

    def analyze_thought(self, thought: str) -> dict:
        from cbt_knowledge import detect_distortions, find_similar_situations

        distortions = detect_distortions(thought)
        codes = [c for c, _ in distortions]
        prompt = (
            f"Thought: {thought}\n"
            f"Detected codes: {', '.join(codes) if codes else 'none'}\n"
            "Write a gentle, balanced reframe (2-4 sentences). Validate feelings first."
        )
        reframe = self._generate(prompt)

        primary = codes[0] if codes else None
        similar = find_similar_situations(primary) if primary else []
        return {
            "thought": thought,
            "distortions": distortions,
            "reframe": reframe,
            "similar_situations": similar,
        }

    def generate_response(self, message: str, context: list[dict] | None = None) -> str:
        """Return a full assistant response (non-streaming)."""
        messages = self._build_messages(message, context)
        return self._generate(prompt="", messages=messages)

    def stream_generate_response(
        self, message: str, context: list[dict] | None = None
    ) -> Iterable[str]:
        """Yield assistant tokens via provider streaming if supported.

        Falls back to a single full message chunk when streaming
        is not supported by the provider/client.
        """
        messages = self._build_messages(message, context)

        # Try OpenAI-compatible chat streaming first
        try:
            if hasattr(self.client, "chat") and hasattr(self.client.chat, "completions"):
                stream = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=messages,
                    temperature=0.6,
                    max_tokens=400,
                    stream=True,
                )
                for chunk in stream:  # ChatCompletionStreamOutput-like
                    # Handle both object-like and dict-like chunks
                    delta = (
                        getattr(chunk.choices[0], "delta", None)
                        if hasattr(chunk, "choices")
                        else None
                    )
                    content = None
                    if delta is not None:
                        content = getattr(delta, "content", None)
                    elif isinstance(chunk, dict):
                        content = chunk.get("choices", [{}])[0].get("delta", {}).get("content")
                    if content:
                        yield content
                return
            if hasattr(self.client, "chat_completion"):
                stream = self.client.chat_completion(
                    messages=messages,
                    temperature=0.6,
                    max_tokens=400,
                    stream=True,
                )
                for chunk in stream:
                    # older dict-shaped stream
                    content = (chunk.get("choices") or [{}])[0].get("delta", {}).get("content")
                    if content:
                        yield content
                return
        except InferenceTimeoutError as e:  # pragma: no cover
            # Surface a short message; UI will present it
            yield f"[timeout] {e}"
            return
        except Exception:
            # Fall through to plain generation
            pass

        # Fallback: text generation streaming if available
        try:
            gen = self.client.text_generation(
                messages[-1]["content"],
                max_new_tokens=400,
                temperature=0.6,
                stream=True,
                return_full_text=False,
            )
            for token in gen:
                if token:
                    yield str(token)
            return
        except Exception:
            pass

        # Final fallback: non-streaming
        yield self.generate_response(message, context)

    def get_crisis_response(self) -> str:
        return (
            "I'm concerned about what you're sharing and want you to know that support is "
            "available.\n\n"
            "Immediate help: see international resources at "
            "https://blog.opencounseling.com/suicide-hotlines/\n\n"
            "You don't have to go through this alone. Please reach out to a professional who "
            "can provide the support you deserve."
        )


def create_simple_cbt_assistant():
    """Backwards-compatible alias."""
    return CBTAgent()