nova.py

#!/usr/bin/env python3
"""Nova — a local AI terminal assistant powered by Ollama."""

from __future__ import annotations

import argparse
import base64
import re
import select
import sys
import termios
import threading
import time
import tomllib
import tty
from pathlib import Path
from typing import Optional

import ollama
from prompt_toolkit import PromptSession
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.spinner import Spinner
from rich.text import Text

from pylatexenc.latex2text import LatexNodes2Text

from art import show_welcome
from commands import REGISTRY, Session, SlashCompleter, build_profile, prompt_style, trim_history
from tools import get_tool, ollama_tools

CONFIG_PATH = Path(__file__).parent / "config.toml"

console = Console()


def _keyboard_listener(stop_event: threading.Event) -> None:
    """Set stop_event when 'q' is pressed. Runs in a daemon thread."""
    fd = sys.stdin.fileno()
    old = termios.tcgetattr(fd)
    try:
        tty.setcbreak(fd)
        while not stop_event.is_set():
            if select.select([sys.stdin], [], [], 0.05)[0]:
                if sys.stdin.read(1).lower() == "q":
                    stop_event.set()
    except Exception:
        pass
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old)


def load_config() -> dict:
    with open(CONFIG_PATH, "rb") as f:
        return tomllib.load(f)


THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
DISPLAY_MATH_RE = re.compile(r"\$\$(.+?)\$\$", re.DOTALL)
INLINE_MATH_RE = re.compile(r"\$(.+?)\$")
MATRIX_RE = re.compile(r"\\begin\{([pbvB]?matrix)\}(.*?)\\end\{[pbvB]?matrix\}", re.DOTALL)
ENV_RE = re.compile(r"(\\begin\{[^}]+\}.*?\\end\{[^}]+\})", re.DOTALL)
LATEX_FENCE_RE = re.compile(r"```(?:latex|math)\s*\n(.*?)```", re.DOTALL)

_latex = LatexNodes2Text()

# Left and right bracket characters (top, mid, bottom) for each matrix type.
_BRACKETS = {
    "matrix":  ("⎡⎢⎣", "⎤⎥⎦"),
    "pmatrix": ("⎛⎜⎝", "⎞⎟⎠"),
    "bmatrix": ("⎡⎢⎣", "⎤⎥⎦"),
    "vmatrix": ("⎢⎢⎢", "⎥⎥⎥"),
    "Bmatrix": ("⎧⎪⎩", "⎫⎪⎭"),
}

def _render_latex(src: str) -> str:
    try:
        return _latex.latex_to_text(src.strip())
    except Exception:
        return src


def _render_matrix(env: str, content: str) -> str:
    rows = [r.strip() for r in re.split(r"\\\\", content) if r.strip()]
    cells = [[_render_latex(c.strip()) for c in row.split("&")] for row in rows]

    n_cols = max(len(r) for r in cells)
    col_w = [max(len(r[c]) if c < len(r) else 0 for r in cells) for c in range(n_cols)]

    lines = []
    for row in cells:
        padded = [cell.center(col_w[i]) for i, cell in enumerate(row)]
        lines.append("  " + "   ".join(padded) + "  ")

    left, right = _BRACKETS.get(env, ("", ""))
    if not left:
        return "\n".join(lines)

    result = []
    for i, line in enumerate(lines):
        if len(lines) == 1:
            l, r = left[0], right[0]
        elif i == 0:
            l, r = left[0], right[0]
        elif i == len(lines) - 1:
            l, r = left[2], right[2]
        else:
            l, r = left[1], right[1]
        result.append(f"{l}{line}{r}")
    return "\n".join(result)


def convert_latex(text: str) -> str:
    """Convert LaTeX math blocks to Unicode for terminal rendering."""
    text = LATEX_FENCE_RE.sub(lambda m: _render_latex(m.group(1)), text)
    # Matrices must be handled before $$...$$ eats them.
    text = MATRIX_RE.sub(lambda m: "\n" + _render_matrix(m.group(1), m.group(2)) + "\n", text)
    text = DISPLAY_MATH_RE.sub(lambda m: "\n" + _render_latex(m.group(1)) + "\n", text)
    text = INLINE_MATH_RE.sub(lambda m: _render_latex(m.group(1)), text)
    text = ENV_RE.sub(lambda m: _render_latex(m.group(1)), text)
    return text


def run_response(session: Session) -> str:
    """Agentic loop: call Ollama, execute any tool calls, repeat until final answer."""
    messages = list(session.history)

    if session.pending_image is not None:
        image_data = base64.b64encode(session.pending_image.read_bytes()).decode()
        if messages and messages[-1]["role"] == "user":
            messages[-1] = {**messages[-1], "images": [image_data]}

    tools = ollama_tools()
    start = time.monotonic()
    full_response = ""

    with Live(Spinner("dots", style="color(208)"), console=console, refresh_per_second=15, vertical_overflow="visible") as live:
        tools_were_called = False

        # Tool-call phases: non-streaming so the model commits to tool calls
        # before generating any text, avoiding split responses.
        while tools:
            try:
                response = ollama.chat(model=session.model, messages=messages, tools=tools)
            except ollama.ResponseError as e:
                if e.status_code == 400 and "does not support tools" in str(e):
                    tools = None  # fall through to plain streaming
                    break
                raise
            msg = response.message
            tool_calls = list(msg.tool_calls or [])

            if not tool_calls:
                if tools_were_called:
                    # The model has synthesized from tool results — use this response
                    # directly rather than re-requesting without tools (which causes
                    # the model to ignore the tool result messages in history).
                    full_response = msg.content or ""
                    visible = THINK_RE.sub("", full_response).strip()
                    if visible:
                        live.update(Markdown(convert_latex(visible)))
                break

            tools_were_called = True
            messages.append({
                "role": "assistant",
                "content": msg.content or "",
                "tool_calls": tool_calls,
            })

            for tc in tool_calls:
                fn_name = tc.function.name
                fn_args = tc.function.arguments or {}
                tool = get_tool(fn_name)

                if tool and tool.requires_confirmation:
                    live.stop()
                    console.print(Text(f"\n  {fn_name}", style="color(220) bold"))
                    if fn_name == "run_python":
                        for line in fn_args.get("code", "").splitlines():
                            console.print(Text(f"  {line}", style="color(245)"))
                        console.print()
                    elif fn_name == "write_file":
                        console.print(Text(f"  → {fn_args.get('path', '')}", style="color(245)"))
                        for line in fn_args.get("content", "").splitlines()[:20]:
                            console.print(Text(f"  {line}", style="color(245)"))
                        if fn_args.get("content", "").count("\n") > 20:
                            console.print(Text("  [...]", style="color(240)"))
                        console.print()
                    else:
                        console.print(Text(f"  $ {fn_args.get('command', '')}\n", style="color(245)"))
                    answer = console.input(Text("  run this? [y/N] ", style="color(240)"))
                    console.print()
                    if answer.strip().lower() != "y":
                        messages.append({"role": "tool", "content": "User declined to run the command."})
                        live.start()
                        continue
                    live.start()

                live.update(Spinner("dots", text=f" {fn_name}...", style="color(220)"))
                if tool:
                    try:
                        result = tool.fn(**fn_args)
                    except Exception as e:
                        result = f"Error running {fn_name}: {e}"
                else:
                    result = f"Error: unknown tool '{fn_name}'"
                messages.append({"role": "tool", "content": result})

        if not tools_were_called:
            # No tools used — stream the response normally.
            stop_event = threading.Event()
            listener = threading.Thread(target=_keyboard_listener, args=(stop_event,), daemon=True)
            listener.start()
            try:
                for chunk in ollama.chat(model=session.model, messages=messages, stream=True):
                    if stop_event.is_set():
                        break
                    full_response += chunk.message.content or ""

                    thinking = "<think>" in full_response and "</think>" not in full_response
                    visible = THINK_RE.sub("", full_response).strip()

                    if thinking:
                        live.update(Spinner("dots", text=" Thinking...", style="color(240)"))
                    elif visible:
                        live.update(Markdown(convert_latex(visible)))
            finally:
                stop_event.set()
                listener.join(timeout=0.2)

    elapsed = time.monotonic() - start
    console.print(Text(f"  {elapsed:.1f}s", style="color(244)"))

    final = convert_latex(THINK_RE.sub("", full_response).strip())

    # Print bare URLs so terminals without OSC 8 support (e.g. Konsole) can
    # autodetect and make them clickable, while markdown links are preserved above.
    urls = re.findall(r'\[(?:[^\]]+)\]\((https?://[^\)]+)\)', final)
    if urls:
        console.print(Text("  links", style="color(240)"))
        for url in dict.fromkeys(urls):  # deduplicate, preserve order
            console.print(Text(f"  {url}", style="color(244)"))
        console.print()

    return final


def main() -> None:
    parser = argparse.ArgumentParser(prog="nova", description="Nova — local AI assistant")
    parser.add_argument("-p", "--profile", default=None, help="profile to load on startup")
    parser.add_argument("-m", "--model", default=None, help="override the model on startup")
    args = parser.parse_args()

    config = load_config()
    settings = config.get("settings", {})
    max_history = settings.get("max_history_messages", 20)

    profile_name = args.profile or settings.get("default_profile", "default")
    model, history = build_profile(config, profile_name, console)

    if args.model:
        model = args.model

    show_welcome(console, model, profile_name)

    session = Session(
        console=console,
        config=config,
        model=model,
        profile_name=profile_name,
        history=history,
        max_history=max_history,
    )

    prompt_session: PromptSession = PromptSession(
        completer=SlashCompleter(config),
        style=prompt_style,
        complete_while_typing=True,
    )

    while True:
        try:
            user_input = prompt_session.prompt("You  ").strip()
        except (KeyboardInterrupt, EOFError):
            console.print("\n[color(240)]bye[/color(240)]")
            break

        if not user_input:
            continue

        if user_input.startswith("/"):
            parts = user_input.split(maxsplit=1)
            command = parts[0].lower()
            argument = parts[1].strip() if len(parts) > 1 else ""

            if command == "/exit":
                console.print("[color(240)]bye[/color(240)]")
                break

            handler = REGISTRY.get(command)
            if handler:
                handler(session, argument)
            else:
                console.print(f"[red]Unknown command '{command}'. Type /help for the list.[/red]")

            continue

        session.history.append({"role": "user", "content": user_input})
        session.history = trim_history(session.history, session.max_history)

        console.print(Text("Nova  ", style="color(208) bold"))
        try:
            response = run_response(session)
        except ollama.ResponseError as e:
            console.print(f"[red]Ollama error: {e}[/red]")
            session.history.pop()
            continue
        except Exception as e:
            console.print(f"[red]{e}[/red]")
            session.history.pop()
            continue
        finally:
            session.pending_image = None

        console.print()
        session.history.append({"role": "assistant", "content": response})
        session.history = trim_history(session.history, session.max_history)


if __name__ == "__main__":
    main()