ElebobaandClaude Opus 5.5
65acf3abf1
CI / ruff · mypy · pytest (pull_request) Successful in 39m13s
- prefilter: правила из config/agent_rules.yaml (длина, команды, shell,
секреты и системные файлы, prompt injection) — без LLM
- guard: отдельный вызов, JSON {verdict, reason}; невалидный ответ = attack,
недоступный LLM — «попробуйте позже», а не атака
- planner: только AgentLogQuery (контейнер из белого списка, ≤ 7 д, шаблон
≤ 200 без вложенных квантификаторов, лимит ≤ 500); логи в LLM не идут
- бот: «🤖 Спросить агента» → показывает построенный grep и строки;
атака — вежливый отказ и 🛡 администраторам; квота 30 в сутки
- журнал agent_requests + log_queries (mode=agent)
- при LLM_PROVIDER=mock агент выключен
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
"""The agent pipeline end to end, with the daily quota and the `agent_requests` journal (SPEC L-10…L-17)."""
|
|
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta
|
|
from enum import StrEnum
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import func, select
|
|
|
|
from app.agent import guard, planner
|
|
from app.agent.log_query import AgentLogQuery
|
|
from app.agent.prefilter import Prefilter, load_rules
|
|
from app.db.models import AgentRequest, Verdict, utcnow
|
|
from app.db.session import SessionMaker
|
|
from app.llm.provider import LLMError, LLMProvider
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
DAILY_QUOTA = 30 # L-17
|
|
|
|
|
|
class Outcome(StrEnum):
|
|
OK = "ok"
|
|
ATTACK = "attack" # prefilter or guard: refuse politely, alert the admins
|
|
QUOTA = "quota"
|
|
UNAVAILABLE = "unavailable" # the LLM is down
|
|
NOT_UNDERSTOOD = "not_understood" # the planner's reply failed validation
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Answer:
|
|
outcome: Outcome
|
|
query: AgentLogQuery | None = None
|
|
reason: str = "" # for the admins' alert and the journal; never shown to the asker
|
|
stage: str = "" # prefilter | guard
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Prompts:
|
|
guard: str
|
|
planner: str
|
|
|
|
|
|
def load_prompts(config_dir: Path) -> Prompts:
|
|
d = config_dir / "prompts"
|
|
return Prompts(
|
|
guard=(d / "guard.md").read_text(encoding="utf-8"), planner=(d / "planner.md").read_text(encoding="utf-8")
|
|
)
|
|
|
|
|
|
class Agent:
|
|
def __init__(
|
|
self, llm: LLMProvider, sessions: SessionMaker, prefilter: Prefilter, prompts: Prompts, quota: int = DAILY_QUOTA
|
|
) -> None:
|
|
self._llm = llm
|
|
self._sessions = sessions
|
|
self._prefilter = prefilter
|
|
self._prompts = prompts
|
|
self.quota = quota
|
|
|
|
@classmethod
|
|
def from_config(cls, llm: LLMProvider, sessions: SessionMaker, config_dir: Path) -> "Agent":
|
|
return cls(llm, sessions, Prefilter(load_rules(config_dir)), load_prompts(config_dir))
|
|
|
|
async def used_today(self, user_id: int, now: datetime) -> int:
|
|
async with self._sessions() as session:
|
|
count = await session.scalar(
|
|
select(func.count())
|
|
.select_from(AgentRequest)
|
|
.where(AgentRequest.user_id == user_id, AgentRequest.at >= now - timedelta(days=1))
|
|
)
|
|
return int(count or 0)
|
|
|
|
async def ask(self, user_id: int, prompt: str, containers: list[str], default: str) -> Answer:
|
|
now = utcnow()
|
|
prompt = prompt.strip()
|
|
if await self.used_today(user_id, now) >= self.quota:
|
|
return Answer(Outcome.QUOTA)
|
|
|
|
reason = self._prefilter.check(prompt)
|
|
if reason is not None:
|
|
await self._journal(user_id, prompt, Verdict.ATTACK, f"prefilter: {reason}", None, 0)
|
|
return Answer(Outcome.ATTACK, reason=reason, stage="prefilter")
|
|
|
|
tokens = 0
|
|
try:
|
|
verdict = await guard.classify(self._llm, self._prompts.guard, prompt)
|
|
tokens += verdict.tokens
|
|
if verdict.verdict is not Verdict.SAFE:
|
|
await self._journal(user_id, prompt, Verdict.ATTACK, f"guard: {verdict.reason}", None, tokens)
|
|
return Answer(Outcome.ATTACK, reason=verdict.reason, stage="guard")
|
|
result = await planner.plan(self._llm, self._prompts.planner, prompt, containers, default, now)
|
|
tokens += result.tokens
|
|
except LLMError as e:
|
|
log.warning("agent: llm unavailable: %s", e)
|
|
await self._journal(user_id, prompt, Verdict.ERROR, f"llm: {str(e)[:200]}", None, tokens)
|
|
return Answer(Outcome.UNAVAILABLE)
|
|
|
|
if result.query is None:
|
|
await self._journal(user_id, prompt, Verdict.ERROR, f"planner: {result.error}", None, tokens)
|
|
return Answer(Outcome.NOT_UNDERSTOOD, reason=result.error)
|
|
await self._journal(user_id, prompt, Verdict.SAFE, None, result.query.model_dump(), tokens)
|
|
return Answer(Outcome.OK, query=result.query)
|
|
|
|
async def _journal(
|
|
self,
|
|
user_id: int,
|
|
prompt: str,
|
|
verdict: Verdict,
|
|
reason: str | None,
|
|
query: dict[str, object] | None,
|
|
tokens: int,
|
|
) -> None:
|
|
async with self._sessions.begin() as session:
|
|
session.add(
|
|
AgentRequest(
|
|
user_id=user_id, prompt=prompt, verdict=verdict, reason=reason, log_query=query, tokens=tokens
|
|
)
|
|
)
|