llm_client.py

"""All communication with the Ollama server.

This module is the only place that builds Ollama requests. Every call injects
the current OS context and disables the model's thinking mode, and no failure
here is allowed to crash the app.
"""

import requests

from config import CONFIG
import os_detector

NO_THINK = "/no_think"


def is_available():
    """Return True if the Ollama server responds to a request."""
    return _get_tags() is not None


def list_models():
    """Return the names of models installed on the Ollama server."""
    response = _get_tags()
    if response is None:
        return []
    try:
        return [model["name"] for model in response.get("models", [])]
    except (KeyError, ValueError, AttributeError):
        return []


def _get_tags():
    """Return the parsed /api/tags response, or None if unreachable."""
    try:
        response = requests.get(
            _endpoint("/api/tags"),
            timeout=CONFIG["ollama_timeout"],
        )
        response.raise_for_status()
        return response.json()
    except (requests.RequestException, ValueError):
        return None


def generate(prompt, system_prompt="", context=None):
    """Make one chat call to Ollama; return the text or None on failure."""
    messages = [
        {"role": "system", "content": _build_system(system_prompt, context)},
        {"role": "user", "content": prompt},
    ]
    payload = {"model": CONFIG["model"], "messages": messages, "stream": False}

    try:
        response = requests.post(
            _endpoint("/api/chat"),
            json=payload,
            timeout=CONFIG["ollama_timeout"],
        )
        response.raise_for_status()
        return response.json()["message"]["content"].strip()
    except (requests.RequestException, KeyError, ValueError) as error:
        _log_failure(error)
        return None


def ask_yes_no(question, context=None):
    """Ask a yes/no question; return a bool (unclear answers are False)."""
    system_prompt = "Answer with a single word: 'yes' or 'no'. Nothing else."
    answer = generate(question, system_prompt, context)
    if answer is None:
        return False
    return answer.strip().lower().startswith("yes")


def extract_command(text):
    """Extract the shell command from text and return it, or None on failure."""
    system_prompt = (
        "Extract the single shell command from the text. "
        "Reply with only the command, no markdown, backticks, or explanation."
    )
    response = generate(text, system_prompt)
    if response is None:
        return None
    command = _strip_code_block(response)
    return command or None


def summarize_for_user(task, outcome, command, output):
    """Return a friendly, jargon-free message describing how the task went."""
    tone = "positive and encouraging" if outcome else "calm and helpful"
    prompt = (
        f"Task: {task}\n"
        f"Command run: {command}\n"
        f"Succeeded: {bool(outcome)}\n"
        f"Output:\n{output}\n\n"
        f"Write one short, {tone} sentence for the user. "
        "Avoid technical jargon."
    )
    message = generate(prompt, "You explain command results in plain language.")
    if message is None:
        return _fallback_summary(outcome)
    return message


def _build_system(system_prompt, context):
    """Prepend the no-think directive and append OS context."""
    os_context = context
    if os_context is None:
        os_context = os_detector.get_os_context()
    lines = [NO_THINK]
    if system_prompt:
        lines.append(system_prompt)
    lines.append(_format_context(os_context))
    return "\n".join(lines)


def _format_context(os_context):
    """Render the OS context dict as readable lines for the system prompt."""
    return (
        "Environment:\n"
        f"- OS: {os_context['os']}\n"
        f"- Shell: {os_context['shell']}\n"
        f"- Package manager: {os_context['package_manager']}\n"
        f"- Home directory: {os_context['home_dir']}"
    )


def _strip_code_block(text):
    """Remove surrounding markdown code fences and backticks from a command."""
    cleaned = text.strip()
    if cleaned.startswith("```"):
        lines = cleaned.splitlines()
        inner = [line for line in lines if not line.startswith("```")]
        cleaned = "\n".join(inner)
    return cleaned.strip().strip("`").strip()


def _fallback_summary(outcome):
    """Return a plain summary used when the model cannot generate one."""
    if outcome:
        return "Done - the command finished successfully."
    return "That didn't work. You may want to try a different approach."


def _log_failure(error):
    """Print an LLM failure to the console when verbose mode is on."""
    if CONFIG["verbose"]:
        print(f"[llm_client] Ollama call failed: {error}")


def _endpoint(path):
    """Join the configured Ollama URL with an API path."""
    return CONFIG["ollama_url"].rstrip("/") + path