checkbot/app/agent/loop.py

import asyncio
import inspect
import json
import logging
import re
import time
from typing import Callable

from .constants import (ABORT_MODEL_ERROR, ABORT_STEP_LIMIT, ABORT_TIMEOUT, ANSWER_MAX, FALLBACK_SNIPPET_MAX,
                        FINISH_RESERVE, KEEP_FULL_TOOL_RESULTS, MAX_STEPS, MAX_TOOL_CALLS_PER_STEP, ROLE_SYSTEM,
                        ROLE_TOOL, ROLE_USER, STEP_ARG_LOG_MAX, TOTAL_TIMEOUT)
from .executor import ToolExecutor
from .prompts import SUMMARY_REQUEST, SYSTEM_PROMPT, ask_prompt, explain_prompt
from .sanitize import sanitize, wrap_data
from .tools import TOOL_SCHEMAS
from .types import AgentResult, AgentStep, ChatModel, Recorder

log = logging.getLogger(__name__)

THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
NO_ANSWER = "Keine Antwort erhalten."
NOT_EXECUTED = "Nicht ausgeführt."
COMPACTED_RESULT = "[gekürzt] Ergebnis früherer Abfrage"
MESSAGES = {
    ABORT_TIMEOUT: "Abgebrochen: Zeitlimit erreicht.",
    ABORT_STEP_LIMIT: "Abgebrochen: Schrittlimit erreicht.",
    ABORT_MODEL_ERROR: "Die KI ist gerade nicht verfügbar.",
}


def _arguments(call: dict) -> tuple[str, object]:
    function = call.get("function") if isinstance(call, dict) else None
    if not isinstance(function, dict):
        return "", {}
    name, arguments = function.get("name"), function.get("arguments", {})
    if isinstance(arguments, str):
        try:
            arguments = json.loads(arguments or "{}")
        except ValueError:
            arguments = None
    return (name if isinstance(name, str) else ""), arguments


def _clean_answer(text: object) -> str:
    return sanitize(THINK_RE.sub("", str(text or "")), ANSWER_MAX) or NO_ANSWER


def _not_executed(calls: list[dict]) -> list[dict]:
    return [{"role": ROLE_TOOL, "tool_name": _arguments(call)[0], "content": NOT_EXECUTED} for call in calls]


def _compact_old_tool_results(messages: list[dict], keep: int = KEEP_FULL_TOOL_RESULTS) -> None:
    tool_indexes = [i for i, m in enumerate(messages) if m.get("role") == ROLE_TOOL]
    for i in tool_indexes[:-keep] if keep > 0 else tool_indexes:
        if messages[i]["content"] not in (COMPACTED_RESULT, NOT_EXECUTED):
            messages[i] = {**messages[i], "content": COMPACTED_RESULT}


class Agent:
    def __init__(self, chat: ChatModel, model: str, executor: ToolExecutor, recorder: Recorder | None = None,
                 clock: Callable[[], float] = time.monotonic, max_steps: int = MAX_STEPS,
                 timeout: float = TOTAL_TIMEOUT, system_prompt: str = SYSTEM_PROMPT):
        self._system_prompt = system_prompt
        self._chat = chat
        self._model = model
        self._executor = executor
        self._recorder = recorder
        self._clock = clock
        self._max_steps = max_steps
        self._timeout = timeout

    async def ask(self, question: str) -> AgentResult:
        return await self.run(ask_prompt(question))

    async def explain(self, host: str, service: str) -> AgentResult:
        return await self.run(explain_prompt(host, service))

    async def _record(self, step: AgentStep) -> None:
        if self._recorder is None:
            return
        try:
            outcome = self._recorder(step)
            if inspect.isawaitable(outcome):
                await outcome
        except Exception:
            log.exception("agent recorder failed")

    async def _call(self, messages: list[dict], deadline: float, tools: list[dict] | None) -> dict:
        remaining = deadline - self._clock()
        if remaining <= 0:
            raise asyncio.TimeoutError
        return await asyncio.wait_for(self._chat.chat(self._model, messages, tools=tools), remaining)

    async def _run_tool(self, number: int, call: dict, deadline: float) -> tuple[AgentStep, str]:
        name, arguments = _arguments(call)
        started = self._clock()
        try:
            result, error = await asyncio.wait_for(self._executor.execute(name, arguments),
                                                   max(deadline - started, 0))
        except asyncio.TimeoutError:
            result, error = "Fehler: Zeitlimit erreicht.", ABORT_TIMEOUT
        shown = arguments if isinstance(arguments, dict) else {}
        step = AgentStep(number, name, {str(k): sanitize(v, STEP_ARG_LOG_MAX, True) for k, v in shown.items()},
                         result, self._clock() - started, error)
        await self._record(step)
        return step, result

    def _fallback(self, reason: str, steps: list[AgentStep]) -> str:
        lines = [MESSAGES[reason]]
        known = [s for s in steps if not s.error]
        if known:
            lines.append("Bisher bekannt:")
            lines += [f"- {s.tool}: {sanitize(s.result, FALLBACK_SNIPPET_MAX, True)}" for s in known]
        return "\n".join(lines)

    async def _finish(self, reason: str, messages: list[dict], steps: list[AgentStep],
                      final_deadline: float) -> AgentResult:
        try:
            messages.append({"role": ROLE_USER, "content": SUMMARY_REQUEST})
            reply = await self._call(messages, final_deadline, None)
            text = _clean_answer(reply.get("content"))
            if text != NO_ANSWER:
                return AgentResult(f"{text}\n({MESSAGES[reason]})", tuple(steps), reason)
        except Exception:
            log.warning("agent summary after %s failed", reason)
        return AgentResult(self._fallback(reason, steps), tuple(steps), reason)

    async def run(self, prompt: str) -> AgentResult:
        start = self._clock()
        final_deadline = start + self._timeout
        deadline = final_deadline - FINISH_RESERVE
        messages: list[dict] = [{"role": ROLE_SYSTEM, "content": self._system_prompt},
                                {"role": ROLE_USER, "content": prompt}]
        steps: list[AgentStep] = []
        while True:
            last_step = len(steps) >= self._max_steps - 1
            tools = None if last_step else TOOL_SCHEMAS
            try:
                reply = await self._call(messages, deadline, tools)
            except asyncio.TimeoutError:
                return await self._finish(ABORT_TIMEOUT, messages, steps, final_deadline)
            except Exception:
                log.warning("agent model call failed", exc_info=True)
                return await self._finish(ABORT_MODEL_ERROR, messages, steps, final_deadline)
            calls = reply.get("tool_calls") or []
            if not calls:
                return AgentResult(_clean_answer(reply.get("content")), tuple(steps))
            allowed, overflow = calls[:MAX_TOOL_CALLS_PER_STEP], calls[MAX_TOOL_CALLS_PER_STEP:]
            messages.append({"role": "assistant", "content": reply.get("content") or "", "tool_calls": calls})
            for index, call in enumerate(allowed):
                if len(steps) >= self._max_steps or self._clock() >= deadline:
                    messages.extend(_not_executed(allowed[index:] + overflow))
                    reason = ABORT_STEP_LIMIT if len(steps) >= self._max_steps else ABORT_TIMEOUT
                    return await self._finish(reason, messages, steps, final_deadline)
                step, result = await self._run_tool(len(steps) + 1, call, deadline)
                steps.append(step)
                messages.append({"role": ROLE_TOOL, "tool_name": step.tool, "content": wrap_data(result)})
                _compact_old_tool_results(messages)
            messages.extend(_not_executed(overflow))