simbiot-harness/harness/_llm_worker.py

"""Воркер одного 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,
                             model: str = None):
    """Записать BudgetSpent за один модельный вызов.

    Воркер — единственный автор spend-событий, но scheduler может убить
    его по таймауту в любой момент (SIGKILL). Поэтому usage пишем
    ИНКРЕМЕНТАЛЬНО: каждый завершённый модельный вызов — своё событие.
    Хвост недописанного вызова теряется, но весь тик — никогда.

    model (e-260930062120): запрошенный model_id уровня пишется в каждое
    событие — в кредитах роутера видно, ЧЬИ это траты; "auto" резолвится
    роутером, фактическую модель до usage strands не выносит.
    """
    if not usage_delta:
        return
    try:
        from harness import es
        es.append_event("harness", "budget", "BudgetSpent", {
            "level": level, "trace_id": trace_id, "source": "worker",
            **({"model": model} if model else {}),
            **usage_delta,
        })
    except Exception:
        pass  # трекинг не должен валить тик


def _wire_budget_hooks(agent, level: str, trace_id: str, model: str = None):
    """После каждого модельного вызова — 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,
        }, model=model)

    agent.add_hook(on_message_added, MessageAddedEvent)
    return agent


def resolve_tool_budget(level: str) -> int:
    """Бюджет тул-коллов тика по уровню (e-260930071900, просьба Alex:
    часовому уровню нужен большой цикл — «хоть 300»).

    Приоритет: SIMBIOT_MAX_TOOL_CALLS_<LEVEL> > SIMBIOT_MAX_TOOL_CALLS
    > per-level дефолт. Минута остаётся компактной (30) — жёсткий кап
    и спасает её от таймаута (e-260930031800: 90+ шагов → 300с → смерть).
    """
    env = (os.environ.get(f"SIMBIOT_MAX_TOOL_CALLS_{level.upper()}")
           or os.environ.get("SIMBIOT_MAX_TOOL_CALLS"))
    if env:
        return int(env)
    return {"min": 30, "hour": 300}.get(level, 30)


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: вендор не должен валить тик
        # stdout воркера не читается (subprocess.run с capture_output и
        # длинным выводом переполняет pipe → BrokenPipeError в печати,
        # результат тика терялся, дайджесты 04:11-04:17 ушли пустыми).
        # Результат живёт в JSON-файле (result_path), печать не нужна.
        from strands.handlers.callback_handler import null_callback_handler
        agent = create_harness(
            model=build_level_model_from_payload(payload),
            effort="auto",  # effort уже в параметрах модели
            instructions=payload.get("instructions", ""),
            session=None,
            tools=tools,
            callback_handler=null_callback_handler,
            builtin_tools=("read", "write", "edit", "shell") if payload.get("can_shell") else [],
            memory=False,
            skills=False,
        )
        agent = _wire_trace_hooks(agent, es, trace_id, model=payload.get("model"))
        agent = _wire_budget_hooks(agent, level, trace_id, model=payload.get("model"))
        # хард-кап шагов (e-260930031800): cancel_tool + end_turn.
        # ВАЖНО: НЕ используем limits={"turns": N} — нативный стоп SDK
        # обрывает цикл на границе хода, и result.message остаётся
        # toolResult-сообщением → str(result)='' → пустой дайджест
        # (репродукция: stop_reason=limit_turns, text='', тики 04:18+).
        # Наш end_turn (AfterToolsEvent) кладёт в историю ассистентский
        # текст с итогом — дайджест не пустеет.
        agent = _wire_tool_budget(agent, resolve_tool_budget(level))
        result = agent(payload.get("prompt", ""))
        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()