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".*?", 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))