loop.py

import datetime
import json
import logging
import os
import platform
import re
from typing import Any

from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.text import Text

from mellea.backends import ModelOption

from config import MAX_TURNS, MCP_SERVERS, REQUIREMENTS, SYSTEM_PROMPT
import provider
from provider import make_session
from tools import TOOLS
from tools.mcp_tools import load_mcp_tools

logging.getLogger("mellea").setLevel(logging.WARNING)

console = Console()


def _user_context() -> str:
    now = datetime.datetime.now().astimezone()
    tz = now.tzname() or now.strftime("%z")
    return (
        f"User environment: {platform.system()} {platform.release()}, "
        f"local time {now.strftime('%Y-%m-%d %H:%M')} ({tz})"
    )


def _build_system_prompt(mcp_tools: list) -> str:
    prompt = SYSTEM_PROMPT + f"\n\n{_user_context()}"
    if not mcp_tools:
        return prompt
    mcp_lines = "\n".join(
        f"- `{t.name}` — {t.as_json_tool['function'].get('description', '').splitlines()[0]}"
        for t in mcp_tools
    )
    return prompt + f"\n\nMCP server tools also available:\n{mcp_lines}"


def _print_tool_call(name: str, args: Any) -> None:
    args_str = json.dumps(args, ensure_ascii=False, default=str)
    console.print(f"  [bold cyan]·[/bold cyan] [cyan]{name}[/cyan]  [dim]{args_str}[/dim]")


# ---------------------------------------------------------------------------
# LiteRT tool-call parse-error recovery
# ---------------------------------------------------------------------------

_LITERT_FAIL_MARKER = "from code block: call:"


def _extract_litert_call(err: str) -> tuple[str, str] | None:
    """Extract (tool_name, raw_args_str) from a LiteRT parse-error string.

    Normalises <|"|> special-quote tokens to real quotes first, then uses a
    brace-depth + string-state walker to find the matching closing brace even
    when string values contain commas, braces, or escape sequences.
    """
    err = err.replace('<|"|>', '"')
    start = err.find(_LITERT_FAIL_MARKER)
    if start == -1:
        return None
    pos = start + len(_LITERT_FAIL_MARKER)
    brace = err.find('{', pos)
    if brace == -1:
        return None
    name = err[pos:brace].strip()
    if not re.match(r'^\w+$', name):
        return None

    depth, in_str, esc, i = 0, False, False, brace
    while i < len(err):
        c = err[i]
        if esc:
            esc = False
        elif c == '\\' and in_str:
            esc = True
        elif c == '"':
            in_str = not in_str
        elif not in_str:
            if c == '{':
                depth += 1
            elif c == '}':
                depth -= 1
                if depth == 0:
                    return name, err[brace + 1:i]
        i += 1
    return None


def _parse_litert_args(args_str: str) -> dict[str, Any]:
    """Parse LiteRT key:value args, handling double-quoted strings and escape sequences."""
    args: dict[str, Any] = {}
    i, n = 0, len(args_str)
    while i < n:
        while i < n and args_str[i] in ' ,\t\n':
            i += 1
        if i >= n:
            break
        j = i
        while j < n and args_str[j] != ':':
            j += 1
        key = args_str[i:j].strip()
        if not key or j >= n:
            break
        i = j + 1
        while i < n and args_str[i] in ' \t':
            i += 1
        if i >= n:
            break
        if args_str[i] == '"':
            i += 1
            chars: list[str] = []
            while i < n:
                c = args_str[i]
                if c == '\\' and i + 1 < n:
                    nxt = args_str[i + 1]
                    chars.append('\n' if nxt == 'n' else '\t' if nxt == 't' else nxt)
                    i += 2
                elif c == '"':
                    i += 1
                    break
                else:
                    chars.append(c)
                    i += 1
            value: Any = ''.join(chars)
        else:
            j = i
            while j < n and args_str[j] not in ',}':
                j += 1
            value = args_str[i:j].strip()
            i = j
        if key:
            args[key] = value
    return args


def _recover_litert_tool_call(err: str, tool_map: dict) -> tuple[str, str] | None:
    """Parse and execute a tool call extracted from a LiteRT parse-error message.

    Uses mellea's MelleaTool.run() for execution so the same tool implementations
    serve both the API and LiteRT paths. Returns (tool_name, result_str) or None
    if the error cannot be recovered (unknown tool, malformed call, etc.).
    """
    extracted = _extract_litert_call(err)
    if extracted is None:
        return None
    name, args_str = extracted
    tool = tool_map.get(name)
    if tool is None:
        return None
    args = _parse_litert_args(args_str)
    _print_tool_call(name, args)
    try:
        return name, str(tool.run(**args))
    except Exception as exc:
        return name, f"Error: {exc}"


# ---------------------------------------------------------------------------
# Agent loops
# ---------------------------------------------------------------------------

def _litert_agent_loop(user_message: str, tools: list, system_prompt: str) -> None:
    os.environ["GLOG_minloglevel"] = "3"
    os.environ["GLOG_logtostderr"] = "0"
    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

    # Suppress C++ glog/TFLite noise at the fd level for the entire init phase.
    # Stderr is kept closed through Engine.__enter__ and restored only just before
    # send_message_async so real inference errors surface through Python exceptions.
    _devnull_fd = os.open(os.devnull, os.O_WRONLY)
    _saved_stderr = os.dup(2)
    os.dup2(_devnull_fd, 2)
    os.close(_devnull_fd)
    _stderr_restored = False

    def _restore_stderr() -> None:
        nonlocal _stderr_restored
        if not _stderr_restored:
            os.dup2(_saved_stderr, 2)
            os.close(_saved_stderr)
            _stderr_restored = True

    try:
        import litert_lm
    except ImportError:
        _restore_stderr()
        console.print(
            "[bold red][error][/bold red] litert-lm-api not installed. "
            "Run: [dim]uv add litert-lm-api[/dim]  (or uv sync --extra litert)"
        )
        return

    backend_name = os.getenv("LITERT_BACKEND", "cpu").lower()
    if backend_name == "gpu":
        backend = litert_lm.Backend.GPU()
    elif backend_name == "npu":
        backend = litert_lm.Backend.NPU()
    else:
        backend = litert_lm.Backend.CPU()

    sampler = litert_lm.SamplerConfig(temperature=0.0, top_k=1)

    tool_map = {t.name: t for t in tools}

    class _ToolHandler(litert_lm.ToolEventHandler):
        def approve_tool_call(self, tool_call: dict) -> bool:
            fn = tool_call.get("function", {})
            _print_tool_call(fn.get("name", "?"), fn.get("arguments", {}))
            return True

        def process_tool_response(self, result: Any) -> Any:
            return result

    raw_tools = [t._call_tool for t in tools]

    try:
        with litert_lm.Engine(provider.LLM_MODEL, backend=backend, max_num_tokens=4096) as engine:
            messages = [litert_lm.Message.system(system_prompt)]
            with engine.create_conversation(
                messages=messages,
                tools=raw_tools,
                tool_event_handler=_ToolHandler(),
                sampler_config=sampler,
                filter_channel_content_from_kv_cache=True,
            ) as conv:
                current_message: Any = user_message
                while True:
                    buf = Text()
                    try:
                        with Live(buf, console=console, refresh_per_second=20) as live:
                            for chunk in conv.send_message_async(current_message):
                                for item in chunk.get("content", []):
                                    if item.get("type") == "text":
                                        buf.append(item["text"])
                                        live.update(buf)
                        break  # clean completion — no parse error
                    except RuntimeError as e:
                        err = str(e)
                        if "Failed to parse tool calls" not in err:
                            raise
                        recovered = _recover_litert_tool_call(err, tool_map)
                        if recovered is None:
                            raise
                        name, response = recovered
                        current_message = litert_lm.Message.tool(
                            litert_lm.Contents([
                                litert_lm.Content.ToolResponse(
                                    name=name, response=response
                                )
                            ])
                        )
    except Exception as e:
        _restore_stderr()
        console.print(f"[red]error[/red] LiteRT: {e}")
    finally:
        _restore_stderr()


def agent_loop(
    user_message: str,
    history: list[tuple[str, str]] | None = None,
) -> str | None:
    mcp_tools = load_mcp_tools(MCP_SERVERS)
    all_tools = TOOLS + mcp_tools
    system_prompt = _build_system_prompt(mcp_tools)

    if provider.IS_LITERT:
        _litert_agent_loop(user_message, all_tools, system_prompt)
        return None

    history_prefix = ""
    if history:
        turns = "\n\n".join(
            f"User: {u}\nAssistant: {a}" for u, a in history[-10:]
        )
        history_prefix = f"Conversation so far:\n{turns}\n\n"

    m = make_session(system_prompt)
    scratchpad: list[str] = []
    last_content: str = ""
    seen_calls: set[str] = set()

    for turn in range(1, MAX_TURNS + 1):
        # When tool calls return only empty text (no prose), the ChatContext would
        # replay an assistant message with content="" on the next turn, which
        # Cohere (and some other providers) reject. Reset the context before each
        # continuation and rebuild the prompt from the scratchpad instead.
        if scratchpad:
            m.reset()

        base = history_prefix + user_message
        prompt = (
            base
            + "\n\nPrevious tool results:\n"
            + "\n\n".join(scratchpad)
            + "\n\nContinue from here."
            if scratchpad
            else base
        )

        try:
            with console.status("[dim]…[/dim]", spinner="dots"):
                result = m.instruct(
                    prompt,
                    model_options={ModelOption.TOOLS: all_tools},
                    tool_calls=True,
                    requirements=REQUIREMENTS,
                )
        except Exception as e:
            console.print(f"[red]error[/red] {e}")
            try:
                with console.status("[dim]…[/dim]", spinner="dots"):
                    result = m.instruct(
                        prompt,
                        model_options={ModelOption.TOOLS: all_tools},
                        tool_calls=True,
                        requirements=REQUIREMENTS,
                    )
            except Exception as e2:
                console.print(f"[red]retry failed[/red] {e2}")
                return last_content or None

        content = result.value or ""
        if content:
            console.print(Markdown(content))
            last_content = content

        if not result.tool_calls:
            return last_content or None

        any_new = False
        for name, tc in result.tool_calls.items():
            call_key = f"{name}:{json.dumps(tc.args, sort_keys=True, ensure_ascii=False)}"
            if call_key in seen_calls:
                console.print(f"  [yellow]· skip duplicate:[/yellow] {name}")
                continue
            seen_calls.add(call_key)
            any_new = True
            _print_tool_call(name, tc.args)
            output = str(tc.call_func())
            scratchpad.append(
                f"[tool: {name}] args={json.dumps(tc.args, ensure_ascii=False)}\n"
                f"[result]\n{output}"
            )

        if not any_new:
            scratchpad.append(
                "[IMPORTANT] All tool calls this turn were duplicates already executed above. "
                "Do NOT call any more tools. Write your final answer using the information already provided."
            )

    return last_content or None