howl/howl/llm/client.py

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any, Optional

import httpx
import structlog
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential

from howl.config import Settings
from howl.db.queries import SenderMatch
from howl.graph.client import Message
from howl.llm.prompts import PROFILE_RECOMMEND_SYSTEM_PROMPT, SYSTEM_PROMPT, build_profile_recommend_prompt, build_prompt
from howl.llm.schemas import CLASSIFY_EMAIL_TOOL, RECOMMEND_PROFILE_TOOL, EmailClassification, ProfileRecommendation

if TYPE_CHECKING:
    from datetime import datetime

    from howl.db.models import SenderProfile

log = structlog.get_logger(__name__)


# ---------------------------------------------------------------------------
# Anthropic-format tool → OpenAI-format tool
# ---------------------------------------------------------------------------

def _to_openai_tool(tool: dict) -> dict:
    return {
        "type": "function",
        "function": {
            "name": tool["name"],
            "description": tool.get("description", ""),
            "parameters": tool["input_schema"],
        },
    }


# ---------------------------------------------------------------------------
# Provider backends
# ---------------------------------------------------------------------------

class _AnthropicBackend:
    def __init__(self, settings: Settings) -> None:
        from anthropic import Anthropic

        self._client = Anthropic(api_key=settings.anthropic_api_key)
        self._model = settings.anthropic_model
        self._max_tokens = settings.anthropic_max_tokens

    def call(
        self,
        system: str,
        user_prompt: str,
        tool: dict,
        tool_name: str,
    ) -> tuple[dict, dict]:
        """Returns (tool_input_dict, raw_response_dict)."""
        response = self._client.messages.create(
            model=self._model,
            max_tokens=self._max_tokens,
            system=system,
            tools=[tool],
            tool_choice={"type": "tool", "name": tool_name},
            messages=[{"role": "user", "content": user_prompt}],
        )

        tool_use_block = next(
            (block for block in response.content if block.type == "tool_use"),
            None,
        )
        if tool_use_block is None:
            raise ValueError(f"LLM did not call the {tool_name} tool")

        raw_response = {
            "model": response.model,
            "stop_reason": response.stop_reason,
            "usage": {
                "input_tokens": response.usage.input_tokens,
                "output_tokens": response.usage.output_tokens,
            },
            "tool_input": tool_use_block.input,
        }
        return tool_use_block.input, raw_response


class _OllamaBackend:
    """Uses Ollama's native /api/chat with JSON mode instead of tool use,
    which is more reliable across quantised models."""

    def __init__(self, settings: Settings) -> None:
        self._base_url = settings.ollama_base_url.rstrip("/")
        self._model = settings.ollama_model
        self._http = httpx.Client(timeout=180)

    @staticmethod
    def _schema_instruction(tool: dict) -> str:
        """Convert a tool definition into a JSON-mode prompt suffix."""
        schema = tool["input_schema"]
        props = schema.get("properties", {})
        required = schema.get("required", [])

        lines = [
            "",
            "Respond with a JSON object containing exactly these fields:",
        ]
        for name, prop in props.items():
            req = " (required)" if name in required else " (optional)"
            desc = prop.get("description", "")
            ptype = prop.get("type", "")
            enum = prop.get("enum")
            items_enum = prop.get("items", {}).get("enum") if prop.get("type") == "array" else None

            constraint = ""
            if enum:
                constraint = f", one of: {', '.join(enum)}"
            elif items_enum:
                constraint = f", array of values from: {', '.join(items_enum)}"

            lines.append(f"- {name} ({ptype}{constraint}){req}: {desc}")

        return "\n".join(lines)

    def call(
        self,
        system: str,
        user_prompt: str,
        tool: dict,
        tool_name: str,
    ) -> tuple[dict, dict]:
        """Returns (tool_input_dict, raw_response_dict)."""
        schema_suffix = self._schema_instruction(tool)

        payload: dict[str, Any] = {
            "model": self._model,
            "messages": [
                {"role": "system", "content": system + schema_suffix},
                {"role": "user", "content": user_prompt},
            ],
            "format": "json",
            "stream": False,
        }

        resp = self._http.post(f"{self._base_url}/api/chat", json=payload)
        resp.raise_for_status()
        data = resp.json()

        content = data.get("message", {}).get("content", "")
        try:
            tool_input = json.loads(content)
        except (json.JSONDecodeError, TypeError):
            raise ValueError(
                f"Ollama returned invalid JSON for {tool_name}. "
                f"Response: {content[:500]}"
            )

        raw_response = {
            "model": data.get("model", self._model),
            "stop_reason": data.get("done_reason", ""),
            "usage": {
                "input_tokens": data.get("prompt_eval_count", 0),
                "output_tokens": data.get("eval_count", 0),
            },
            "tool_input": tool_input,
        }
        return tool_input, raw_response


# ---------------------------------------------------------------------------
# Public LLMClient
# ---------------------------------------------------------------------------

class LLMClient:
    def __init__(self, settings: Settings) -> None:
        self._provider = settings.llm_provider
        if self._provider == "ollama":
            self._backend = _OllamaBackend(settings)
        else:
            self._backend = _AnthropicBackend(settings)
        self._max_body_chars = settings.graph_max_body_chars

    @retry(
        retry=retry_if_exception_type(Exception),
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=15),
        reraise=True,
    )
    def classify(
        self,
        message: Message,
        body: str,
        sender_match: SenderMatch,
        sender_profile: Optional["SenderProfile"] = None,
    ) -> tuple[EmailClassification, dict]:
        """
        Send the email for classification via tool use.

        Returns:
            (EmailClassification, raw_response_dict)
        """
        user_prompt = build_prompt(message, body, sender_match, self._max_body_chars, sender_profile)

        tool_input, raw_response = self._backend.call(
            system=SYSTEM_PROMPT,
            user_prompt=user_prompt,
            tool=CLASSIFY_EMAIL_TOOL,
            tool_name="classify_email",
        )

        classification = EmailClassification.model_validate(tool_input)

        log.debug(
            "email_classified",
            classification=classification.classification,
            action=classification.action,
            confidence=classification.confidence,
            model=raw_response.get("model"),
        )

        return classification, raw_response

    @retry(
        retry=retry_if_exception_type(Exception),
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=15),
        reraise=True,
    )
    def recommend_profile(
        self,
        sender_address: str,
        sender_name: str,
        sample_emails: list[tuple[str, str, "Optional[datetime]", bool]],
    ) -> tuple[ProfileRecommendation, dict]:
        """
        Analyse sample emails from a sender and recommend profile settings.

        Returns:
            (ProfileRecommendation, raw_response_dict)
        """
        user_prompt = build_profile_recommend_prompt(
            sender_address, sender_name, sample_emails, self._max_body_chars,
        )

        tool_input, raw_response = self._backend.call(
            system=PROFILE_RECOMMEND_SYSTEM_PROMPT,
            user_prompt=user_prompt,
            tool=RECOMMEND_PROFILE_TOOL,
            tool_name="recommend_sender_profile",
        )

        recommendation = ProfileRecommendation.model_validate(tool_input)

        log.debug(
            "profile_recommended",
            sender=sender_address,
            sender_type=recommendation.sender_type,
            email_types=recommendation.email_types,
            model=raw_response.get("model"),
        )

        return recommendation, raw_response