javis_bot/src/jarvis/llm.py

"""Direct LLM interaction utilities without extra features like temporal context."""

from __future__ import annotations
from typing import Optional, Any, Dict, List, Generator, Callable
import os
import sys
import requests
import json


def _caller_name() -> str:
    """Best-effort name of the function that invoked the LLM wrapper, used to
    label per-call timing (router / enrichment / digest / main)."""
    try:
        return sys._getframe(2).f_code.co_name
    except Exception:
        return "?"

from .debug import debug_log


# Single context-window size shared by EVERY Ollama call (chat, router,
# enrichment, digests, streaming). Ollama keeps a SEPARATE loaded model
# instance per (model, num_ctx): mixing 4096 and 8192 in one voice turn made
# it evict and cold-reload the model (~3.4s each) on every context switch —
# the dominant per-turn latency. Keeping one value collapses this to a single
# resident instance, so only the very first call of a cold model pays a load.
# 8192 is the floor the main agentic chat needs (system prompt + tool schema +
# memory) without silent truncation. Tunable via env for tight-VRAM hosts.
OLLAMA_NUM_CTX = int(os.environ.get("OLLAMA_NUM_CTX", "8192"))


class ToolsNotSupportedError(Exception):
    """Raised when the model returns HTTP 400 because native tool calling is not supported."""
    pass


def _log_ollama_timing(data: Dict[str, Any], num_ctx: int, caller: str) -> None:
    """Emit a one-line per-call latency breakdown so a slow voice turn can be
    attributed to a specific internal LLM call (router / enrichment / digest /
    main) instead of just a total. ``load_duration`` > 0 means the model was
    cold-reloaded for this call — the single most expensive thing to avoid.
    """
    if not isinstance(data, dict):
        return
    try:
        ns = 1e9
        total = data.get("total_duration", 0) / ns
        load = data.get("load_duration", 0) / ns
        peval = data.get("prompt_eval_duration", 0) / ns
        pcount = data.get("prompt_eval_count")
        gen = data.get("eval_duration", 0) / ns
        gcount = data.get("eval_count")
        reload_flag = " RELOAD" if load > 0.5 else ""
        print(
            f"  ⏱️ llm[{caller}] ctx={num_ctx} total={total:.2f}s "
            f"load={load:.2f}s{reload_flag} prompt={peval:.2f}s({pcount}t) "
            f"gen={gen:.2f}s({gcount}t)",
            flush=True,
        )
    except Exception:  # pragma: no cover - logging must never break a reply
        pass


def call_llm_direct(base_url: str, chat_model: str, system_prompt: str, user_content: str, timeout_sec: float = 10.0, thinking: bool = False, num_ctx: int = OLLAMA_NUM_CTX, temperature: Optional[float] = None) -> Optional[str]:
    """Direct LLM call without temporal context, location, or other ask_coach features.

    ``num_ctx`` controls Ollama's context window for this call. It defaults to
    the shared ``OLLAMA_NUM_CTX`` so small classification-shaped passes load the
    SAME Ollama instance as the main chat loop (no cold reload on context
    switch). Callers may still override it, but diverging from the shared value
    reintroduces a per-turn model reload.

    ``temperature`` is forwarded to Ollama when set. Pass ``0.0`` for
    classification / extraction calls where determinism beats creativity —
    Ollama defaults to ~0.8 otherwise, which can flake small models on
    rule-following tasks (e.g. the knowledge extractor's banned-form list).
    """
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content}
    ]

    options: Dict[str, Any] = {"num_ctx": num_ctx}
    if temperature is not None:
        options["temperature"] = temperature

    payload: Dict[str, Any] = {
        "model": chat_model,
        "messages": messages,
        "stream": False,
        "options": options,
        "think": thinking,
        # Keep the chat model resident between calls. Without an explicit
        # keep_alive Ollama evicts it after its default idle window and the
        # next turn pays a cold reload. We pin the chat model only (embeddings
        # pass keep_alive=0 so they unload after use) instead of a global
        # OLLAMA_KEEP_ALIVE=-1, which would keep every model resident forever.
        "keep_alive": "30m",
    }

    caller = _caller_name()
    try:
        with requests.post(f"{base_url.rstrip('/')}/api/chat", json=payload, timeout=timeout_sec) as resp:
            resp.raise_for_status()
            data = resp.json()
        _log_ollama_timing(data, num_ctx, caller)

        if isinstance(data, dict):
            content = extract_text_from_response(data)
            if isinstance(content, str) and content.strip():
                return content
            debug_log(f"call_llm_direct: empty content from response keys={list(data.keys())}", "llm")
    except requests.exceptions.Timeout:
        debug_log(f"call_llm_direct: timeout after {timeout_sec}s", "llm")
        return None
    except Exception as e:
        debug_log(f"call_llm_direct: request failed — {e}", "llm")
        return None

    return None


def call_llm_streaming(
    base_url: str,
    chat_model: str,
    system_prompt: str,
    user_content: str,
    on_token: Optional[Callable[[str], None]] = None,
    timeout_sec: float = 30.0,
    thinking: bool = False,
) -> Optional[str]:
    """
    Streaming LLM call that invokes on_token callback for each token received.

    Args:
        base_url: Ollama base URL
        chat_model: Model name
        system_prompt: System prompt
        user_content: User message
        on_token: Callback invoked with each token as it arrives
        timeout_sec: Request timeout
        thinking: Enable thinking/reasoning mode

    Returns:
        Complete response text, or None on error
    """
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content}
    ]

    payload: Dict[str, Any] = {
        "model": chat_model,
        "messages": messages,
        "stream": True,
        "options": {"num_ctx": OLLAMA_NUM_CTX},
        "think": thinking,
        # Keep the chat model resident between calls (see call_llm_direct).
        "keep_alive": "30m",
    }

    # Use ``with`` so the streaming response (and the underlying TCP
    # connection) is released even if iter_lines exits early via an
    # exception or the caller stops consuming. Without this an aborted
    # stream pinned the connection until GC, which could happen many
    # turns later under sustained reply load.
    try:
        with requests.post(
            f"{base_url.rstrip('/')}/api/chat",
            json=payload,
            timeout=timeout_sec,
            stream=True,
        ) as resp:
            resp.raise_for_status()

            full_response = []
            for line in resp.iter_lines():
                if line:
                    try:
                        data = json.loads(line)
                        if "message" in data and isinstance(data["message"], dict):
                            content = data["message"].get("content", "")
                            if content:
                                full_response.append(content)
                                if on_token:
                                    on_token(content)
                    except json.JSONDecodeError:
                        continue

            result = "".join(full_response)
            return result if result.strip() else None

    except requests.exceptions.Timeout:
        return None
    except Exception:
        return None


def extract_text_from_response(data: Dict[str, Any]) -> Optional[str]:
    """Extract text from LLM response - supports multiple response formats."""
    # Preferred: Ollama chat non-stream format
    if "message" in data and isinstance(data["message"], dict):
        content = data["message"].get("content")
        if isinstance(content, str):
            return content
    
    # Fallback: OpenAI-style format
    if "choices" in data and isinstance(data["choices"], list) and len(data["choices"]) > 0:
        choice = data["choices"][0]
        if isinstance(choice, dict):
            if "message" in choice and isinstance(choice["message"], dict):
                content = choice["message"].get("content")
                if isinstance(content, str):
                    return content
            elif "text" in choice:
                content = choice["text"]
                if isinstance(content, str):
                    return content
    
    # Another fallback: direct "content" field
    if "content" in data:
        content = data["content"]
        if isinstance(content, str):
            return content
    
    return None


def chat_with_messages(
    base_url: str,
    chat_model: str,
    messages: List[Dict[str, str]],
    timeout_sec: float = 30.0,
    extra_options: Optional[Dict[str, Any]] = None,
    tools: Optional[List[Dict[str, Any]]] = None,
    thinking: bool = False,
) -> Optional[Dict[str, Any]]:
    """
    Send an arbitrary messages array to the LLM and return the raw response JSON.
    Caller is responsible for interpreting assistant content (including JSON/tool calls).

    Args:
        base_url: Ollama base URL
        chat_model: Model name
        messages: Conversation messages
        timeout_sec: Request timeout
        extra_options: Additional model options
        tools: Optional list of tools in OpenAI-compatible JSON schema format for native tool calling
        thinking: Enable thinking/reasoning mode

    Returns the parsed JSON response dict on success, or None on error/timeout.
    """
    # Main agentic chat uses 8192 so the system prompt (tool list + protocol
    # guidance + memory context) doesn't overflow and force ollama to truncate
    # — which previously dropped the tool schema on smaller models like
    # gemma4:e2b, tipping them into their pre-trained tool_code scaffolding.
    payload: Dict[str, Any] = {
        "model": chat_model,
        "messages": messages,
        "stream": False,
        "options": {"num_ctx": OLLAMA_NUM_CTX},
        "think": thinking,
        # Keep the chat model resident between turns (see call_llm_direct).
        "keep_alive": "30m",
    }
    if extra_options and isinstance(extra_options, dict):
        # Merge shallowly into options
        payload["options"].update(extra_options)

    # Add tools for native tool calling support (Ollama 0.4+)
    if tools and isinstance(tools, list) and len(tools) > 0:
        payload["tools"] = tools

    caller = _caller_name()
    try:
        with requests.post(f"{base_url.rstrip('/')}/api/chat", json=payload, timeout=timeout_sec) as resp:
            resp.raise_for_status()
            data = resp.json()
        _log_ollama_timing(data, OLLAMA_NUM_CTX, caller)
        if isinstance(data, dict):
            return data
    except requests.exceptions.Timeout:
        print("  ⏱️ LLM request timed out", flush=True)
        return None
    except requests.exceptions.ConnectionError as e:
        print(f"  ❌ LLM connection error: {e}", flush=True)
        return None
    except requests.exceptions.HTTPError as e:
        # Raise a specific error when the model rejects the tools parameter (HTTP 400).
        # This lets the caller fall back to text-based tool calling automatically.
        if e.response is not None and e.response.status_code == 400 and tools:
            raise ToolsNotSupportedError(
                f"Model {chat_model!r} returned HTTP 400 — native tools API not supported"
            )
        print(f"  ❌ LLM HTTP error: {e}", flush=True)
        return None
    except Exception as e:
        print(f"  ❌ LLM error: {e}", flush=True)
        return None

    return None