34df7f3d72
RAG (t-260929212316): каждому LLM-тику доступен rag_lookup — семантический поиск по базе кооператива (/root/rag, fastembed+sqlite). Лимиты зашиты в инструмент: 1 запрос/тик, k<=5, таймаут 20с, сабпроцесс rag-venv, мёртвый вендор не валит тик. Трейс RagQueried в сторе. CLI: rag-query. Хард-кап шагов (e-260930031800): мин-воркер тонул в 90+ tool-вызовах и умирал по таймауту 300с. До: 3 CycleError за 2 часа, эскалации min→hour. Фикс: cancel_tool на вызовы сверх лимита + end_turn (AfterToolsEvent) останавливает цикл после партии + limits.turns (нативный механизм SDK). e2e: 20 параллельных вызовов при лимите 3 → ровно 3 исполнения, цикл остановлен, тик завершён текстом за ~12с (было: таймаут 300с). Мандат мин-уровня: про rag_lookup и приоритет RAG для вопросов про кооператив. Тесты: tests/test_rag_lookup.py (29 PASS), tests/test_tool_budget.py (12 PASS).
254 lines
No EOL
12 KiB
Python
254 lines
No EOL
12 KiB
Python
"""Воркер одного LLM-тика (запускается сабпроцессом из isolated_runner).
|
||
|
||
Читает payload из stdin, гоняет create_harness() один раз и умирает —
|
||
каждый тик получает чистый asyncio-рантайм Strands.
|
||
Трасса пишется напрямую в event store (es доступен).
|
||
Результат — в JSON-файл (path в payload['result_path']); stdout — только лог.
|
||
"""
|
||
import json
|
||
import os
|
||
import sys
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
|
||
def _rag_tool_for(level: str, trace_id: str):
|
||
"""RAG-инструмент тика (t-260929212316): семантический поиск по базе
|
||
кооператива (вендор /root/rag, fastembed bge-small + sqlite).
|
||
|
||
Ограничения зашиты в сам инструмент (не в промпт):
|
||
1 запрос на тик, k<=5, таймаут 20с, сабпроцесс rag-venv — упавший
|
||
вендор не валит тик. RAG_DISABLED=1 выключает.
|
||
"""
|
||
from strands import tool as strands_tool
|
||
from harness import rag_lookup as _rag
|
||
state = {"calls": 0}
|
||
|
||
@strands_tool
|
||
def rag_lookup(question: str) -> str:
|
||
"""Семантический поиск по базе знаний кооператива: контекст, проекты,
|
||
порталы, встречи, документы. Возвращает релевантные фрагменты.
|
||
Один вызов на тик — формулируй запрос целиком."""
|
||
if state["calls"] >= 1:
|
||
return "лимит: 1 RAG-запрос на тик (использован)"
|
||
state["calls"] += 1
|
||
chunks = _rag.rag_query(question, k=_rag.MAX_K,
|
||
level=level, trace_id=trace_id)
|
||
if not chunks:
|
||
return "(поиск ничего не вернул — отвечай из контекста тика)"
|
||
return _rag.format_chunks(chunks)
|
||
|
||
return rag_lookup
|
||
|
||
|
||
def _wire_tool_budget(agent, max_calls: int = 30):
|
||
"""Хард-кап тул-коллов на тик (e-260930031800: мин-воркер уходил в
|
||
90+ шагов исследования и умирал по таймауту 300с).
|
||
|
||
Механизм (e2e-проверено): BeforeToolCallEvent.cancel_tool отменяет
|
||
каждый вызов сверх лимита; AfterToolsEvent.end_turn останавливает
|
||
цикл агента ПОСЛЕ партии инструментов — модель больше не
|
||
вызывается, финальный текст тика — наш отчёт о недовыполненном.
|
||
Без end_turn модель «настаивает»: в тесте слала 20 шагов подряд,
|
||
получая отказы (дёшево, но цикл не кончается).
|
||
|
||
Партия: 20 параллельных tool-вызовов за один ход. Лимит считается
|
||
по РЕАЛЬНЫМ исполнениям (отменённые не съедают бюджет — модель
|
||
не может заблокировать тик потоком отмен), поэтому первые
|
||
max_calls вызова партии выполняются, хвост отменяется, затем
|
||
цикл останавливается с итоговым текстом.
|
||
"""
|
||
from strands.hooks import AfterToolsEvent, BeforeToolCallEvent
|
||
state = {"executed": 0, "halted": False}
|
||
|
||
def on_before_tool(event):
|
||
if state["halted"] or state["executed"] >= max_calls:
|
||
event.cancel_tool = (
|
||
f"тиковый бюджет шагов исчерпан ({max_calls} tool-вызовов): "
|
||
f"вызов отменён; новые вызовы инструментов отключены до конца тика"
|
||
)
|
||
return
|
||
state["executed"] += 1
|
||
|
||
def on_after_tools(event):
|
||
# партия инструментов закончилась; лимит пробит → стоп-цикл
|
||
# (end_turn: цикл агента не зовёт модель повторно)
|
||
if not state["halted"] and state["executed"] >= max_calls:
|
||
state["halted"] = True
|
||
event.end_turn = (
|
||
f"тик завершён: бюджет шагов исчерпан ({max_calls} tool-вызовов). "
|
||
f"Выше — результат того, что успело выполниться. Не успевшее "
|
||
f"оформи задачей (cli.py task) или продолжи в следующем тике."
|
||
)
|
||
|
||
agent.add_hook(on_before_tool, BeforeToolCallEvent)
|
||
agent.add_hook(on_after_tools, AfterToolsEvent)
|
||
return agent
|
||
|
||
|
||
def build_level_model_from_payload(payload: dict):
|
||
"""Модель для воркера — та же сборка, что и в runtime.build_level_model:
|
||
9router понимает chat/completions (usage там же), Responses API — шим без usage."""
|
||
from strands.models.openai import OpenAIModel
|
||
import os
|
||
params = {}
|
||
effort = payload.get("effort") or "auto"
|
||
if effort and effort != "auto":
|
||
params["reasoning_effort"] = effort
|
||
model_id = payload.get("model") or "auto"
|
||
if model_id == "none":
|
||
model_id = "auto"
|
||
return OpenAIModel(
|
||
client_args={"api_key": os.environ.get("OPENAI_API_KEY", ""),
|
||
"base_url": os.environ.get("OPENAI_BASE_URL", "http://127.0.0.1:20128/v1")},
|
||
model_id=model_id,
|
||
**({"params": params} if params else {}),
|
||
)
|
||
|
||
|
||
def _emit(payload, obj):
|
||
rp = payload.get("result_path")
|
||
if rp:
|
||
try:
|
||
with open(rp, "w") as f:
|
||
json.dump(obj, f, ensure_ascii=False)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _extract_usage(agent) -> dict:
|
||
"""Usage из истории сообщений агента.
|
||
|
||
Strands прикрепляет usage к каждой assistant-сообщении в
|
||
message["metadata"]["usage"] (event_loop.py) — по одному словарю
|
||
на модельный вызов (шаг тика). Накапливаем по шагам: тик может
|
||
состоять из многих вызовов, бюджет интересует суммарный.
|
||
|
||
Старый разбор m["message"]["usage"] не находил ничего (структура
|
||
другая) — usage всегда возвращался пустым, spend не рос.
|
||
"""
|
||
usage = {"prompt_tokens": 0, "completion_tokens": 0, "model_calls": 0}
|
||
try:
|
||
msgs = getattr(agent, "messages", None) or []
|
||
for m in msgs:
|
||
u = ((m or {}).get("metadata") or {}).get("usage") or {}
|
||
p = u.get("inputTokens") or 0
|
||
c = u.get("outputTokens") or 0
|
||
if isinstance(p, (int, float)):
|
||
usage["prompt_tokens"] += p
|
||
if isinstance(c, (int, float)):
|
||
usage["completion_tokens"] += c
|
||
if p or c:
|
||
usage["model_calls"] += 1
|
||
except Exception:
|
||
pass
|
||
if not usage["model_calls"]:
|
||
return {}
|
||
return usage
|
||
|
||
|
||
def _emit_budget_incremental(level: str, trace_id: str, usage_delta: dict):
|
||
"""Записать BudgetSpent за один модельный вызов.
|
||
|
||
Воркер — единственный автор spend-событий, но scheduler может убить
|
||
его по таймауту в любой момент (SIGKILL). Поэтому usage пишем
|
||
ИНКРЕМЕНТАЛЬНО: каждый завершённый модельный вызов — своё событие.
|
||
Хвост недописанного вызова теряется, но весь тик — никогда.
|
||
"""
|
||
if not usage_delta:
|
||
return
|
||
try:
|
||
from harness import es
|
||
es.append_event("harness", "budget", "BudgetSpent", {
|
||
"level": level, "trace_id": trace_id, "source": "worker",
|
||
**usage_delta,
|
||
})
|
||
except Exception:
|
||
pass # трекинг не должен валить тик
|
||
|
||
|
||
def _wire_budget_hooks(agent, level: str, trace_id: str):
|
||
"""После каждого модельного вызова — usage-дельта в стор.
|
||
|
||
Strands вешает usage на assistant-сообщение в
|
||
message["metadata"]["usage"] (одна запись на модельный вызов).
|
||
MessageAddedEvent стреляет на добавление КАЖДОЙ записи истории,
|
||
нас интересуют только assistant с usage.
|
||
"""
|
||
from strands.hooks import MessageAddedEvent
|
||
|
||
seen = set()
|
||
|
||
def on_message_added(event):
|
||
msg = getattr(event, "message", None) or {}
|
||
if msg.get("role") != "assistant":
|
||
return
|
||
usage = ((msg.get("metadata") or {}).get("usage")) or {}
|
||
p = usage.get("inputTokens") or 0
|
||
c = usage.get("outputTokens") or 0
|
||
if not (p or c):
|
||
return
|
||
key = id(msg)
|
||
if key in seen: # MessageUpdatedEvent может переигрывать — дедуп по объекту
|
||
return
|
||
seen.add(key)
|
||
_emit_budget_incremental(level, trace_id, {
|
||
"prompt_tokens": p, "completion_tokens": c,
|
||
})
|
||
|
||
agent.add_hook(on_message_added, MessageAddedEvent)
|
||
return agent
|
||
|
||
|
||
def main():
|
||
payload = json.loads(sys.stdin.read())
|
||
|
||
from harness import es
|
||
from harness.runtime import setup_provider, _wire_trace_hooks
|
||
|
||
setup_provider()
|
||
from strands_harness import create_harness
|
||
|
||
trace_id = payload.get("trace_prefix") or "tick"
|
||
level = payload.get("level") or "?"
|
||
try:
|
||
# RAG-инструмент тика (t-260929212316): у LLM появляется rag_lookup —
|
||
# поиск по базе кооператива; лимиты (1 запрос/tick, k<=5, таймаут)
|
||
# зашиты в инструмент. Уровень без shell получает только его.
|
||
tools = []
|
||
try:
|
||
tools.append(_rag_tool_for(level, trace_id))
|
||
except Exception:
|
||
pass # RAG — best effort: вендор не должен валить тик
|
||
agent = create_harness(
|
||
model=build_level_model_from_payload(payload),
|
||
effort="auto", # effort уже в параметрах модели
|
||
instructions=payload.get("instructions", ""),
|
||
session=None,
|
||
tools=tools,
|
||
builtin_tools=("read", "write", "edit", "shell") if payload.get("can_shell") else [],
|
||
memory=False,
|
||
skills=False,
|
||
)
|
||
agent = _wire_trace_hooks(agent, es, trace_id)
|
||
agent = _wire_budget_hooks(agent, level, trace_id)
|
||
# хард-кап шагов (e-260930031800): двойная защита.
|
||
# 1) limits={"turns": N} — нативный механизм SDK: цикл агента
|
||
# останавливается на границе хода (stop_reason="limit_turns"),
|
||
# модель не может уйти в бесконечную цепочку tool-вызовов.
|
||
# 2) хук cancel_tool — страховка на случай параллельного блока
|
||
# (20 вызовов одним блоком: turns их не режет, хук отменяет).
|
||
max_turns = int(os.environ.get("SIMBIOT_MAX_TURNS") or 15)
|
||
agent = _wire_tool_budget(
|
||
agent, int(os.environ.get("SIMBIOT_MAX_TOOL_CALLS") or 30))
|
||
result = agent(payload.get("prompt", ""), limits={"turns": max_turns})
|
||
usage = _extract_usage(agent)
|
||
_emit(payload, {"ok": True, "text": str(result), "usage": usage, "trace_id": trace_id})
|
||
except Exception as e:
|
||
import traceback
|
||
_emit(payload, {"ok": False, "error": f"{e}\n{traceback.format_exc()[-800:]}"})
|
||
# орфан-воркер после убийства родителя может умереть на печати — не критично
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |