139 lines
6.1 KiB
Python
139 lines
6.1 KiB
Python
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,
|
|
MAX_STEPS, 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."
|
|
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
|
|
|
|
|
|
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):
|
|
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],
|
|
deadline: float) -> AgentResult:
|
|
if reason == ABORT_STEP_LIMIT:
|
|
try:
|
|
messages.append({"role": ROLE_USER, "content": SUMMARY_REQUEST})
|
|
reply = await self._call(messages, 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 step limit failed")
|
|
return AgentResult(self._fallback(reason, steps), tuple(steps), reason)
|
|
|
|
async def run(self, prompt: str) -> AgentResult:
|
|
deadline = self._clock() + self._timeout
|
|
messages: list[dict] = [{"role": ROLE_SYSTEM, "content": SYSTEM_PROMPT},
|
|
{"role": ROLE_USER, "content": prompt}]
|
|
steps: list[AgentStep] = []
|
|
while True:
|
|
try:
|
|
reply = await self._call(messages, deadline, TOOL_SCHEMAS)
|
|
except asyncio.TimeoutError:
|
|
return AgentResult(self._fallback(ABORT_TIMEOUT, steps), tuple(steps), ABORT_TIMEOUT)
|
|
except Exception:
|
|
log.warning("agent model call failed", exc_info=True)
|
|
return AgentResult(self._fallback(ABORT_MODEL_ERROR, steps), tuple(steps), ABORT_MODEL_ERROR)
|
|
calls = reply.get("tool_calls") or []
|
|
if not calls:
|
|
return AgentResult(_clean_answer(reply.get("content")), tuple(steps))
|
|
messages.append({"role": "assistant", "content": reply.get("content") or "", "tool_calls": calls})
|
|
for call in calls:
|
|
if len(steps) >= self._max_steps:
|
|
return await self._finish(ABORT_STEP_LIMIT, messages, steps, deadline)
|
|
if self._clock() >= deadline:
|
|
return AgentResult(self._fallback(ABORT_TIMEOUT, steps), tuple(steps), ABORT_TIMEOUT)
|
|
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)})
|