simbiot-harness/harness/runtime.py

"""Simbiot Harness — Strands runtime-адаптер.

create_harness() на тик: модель/effort из конфига уровня.
Hooks → трасса в event store (агрегат trace).
"""
import json
import os
import uuid
from datetime import datetime, timezone

_HARNESS_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def _router_key():
    env = os.environ.get("ROUTER_KEY", "")
    if not env:
        path = os.path.join(_HARNESS_DIR, ".env")
        if os.path.exists(path):
            for line in open(path):
                if line.startswith("ROUTER_KEY="):
                    env = line.split("=", 1)[1].strip()
                    break
    return env


def _router_base_url():
    return os.environ.get("ROUTER_BASE_URL", "http://127.0.0.1:20128/v1")


def build_level_model(level_cfg):
    """Модель уровня для 9router: CHAT COMPLETIONS, не Responses API.

    strands_harness для провайдера openai строит OpenAIResponsesModel,
    но у 9router /responses — шим: usage в response.completed отсутствует
    (проверено: usage:{} и в стриме, и в non-stream). Токены видны только
    в /chat/completions (OpenAIModel, stream_options include_usage) —
    а без них BudgetSpent бесполезен. Проверено: OpenAIModel + 9router
    отдаёт usage и с инструментами, и в reasoning-режиме.

    Сессии выключены, поэтому server-side state Responses API не нужен;
    chat/completions stateless и полностью покрывает потребности тика.
    """
    from strands.models.openai import OpenAIModel
    params = {}
    if level_cfg.effort and level_cfg.effort != "auto":
        params["reasoning_effort"] = level_cfg.effort
    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=level_cfg.model if level_cfg.model != "none" else "auto",
        **({"params": params} if params else {}),
    )


def setup_provider():
    """9router как openai-совместимый провайдер strands."""
    key = _router_key()
    if not key:
        raise RuntimeError("ROUTER_KEY не найден (.env)")
    os.environ["OPENAI_API_KEY"] = key
    os.environ["OPENAI_BASE_URL"] = _router_base_url()


def _wire_trace_hooks(agent, es, trace_id: str):
    """Каждая стадия агентного цикла → событие агрегата trace."""
    from strands import Agent
    from strands.hooks import BeforeInvocationEvent, BeforeToolCallEvent, AfterToolCallEvent

    seq = {"n": 0}

    def _trace(kind: str, payload: dict):
        seq["n"] += 1
        es.append_event("trace", f"{trace_id}:{seq['n']}", kind, payload)

    def on_before_invocation(event):
        _trace("InvocationStarted", {"trace_id": trace_id, "ts": _now()})

    def on_before_tool(event):
        tool = getattr(event, "tool_use", None) or {}
        _trace("ToolCalled", {
            "tool": tool.get("name", "?"),
            "input": json.dumps(tool.get("input", {}), ensure_ascii=False, default=str)[:4000],
        })

    def on_after_tool(event):
        try:
            res = str(getattr(event, "result", ""))[:4000]
        except Exception:
            res = "<unreadable>"
        _trace("ToolResult", {"tool": "?", "result": res})

    agent.add_hook(on_before_invocation, BeforeInvocationEvent)
    agent.add_hook(on_before_tool, BeforeToolCallEvent)
    agent.add_hook(on_after_tool, AfterToolCallEvent)
    return agent


def _now():
    return datetime.now(timezone.utc).isoformat()


def build_agent(level_cfg, session_id=None, extra_instructions=""):
    """Собирает Strands-агента для тика уровня. None если уровень без LLM."""
    if not level_cfg.uses_llm:
        return None
    setup_provider()
    from strands_harness import create_harness

    instructions = level_cfg.mandate + ("\n\n" + extra_instructions if extra_instructions else "")
    agent = create_harness(
        model=build_level_model(level_cfg),
        effort="auto",  # effort уже в параметрах модели (build_level_model)
        instructions=instructions,
        session={"id": session_id} if session_id else None,
        builtin_tools=("read", "write", "edit", "shell") if level_cfg.can_shell else [],
        memory=False,
        skills=False,
    )
    return agent


def run_cycle(level_cfg, es, prompt: str, session_id=None, trace_prefix=None):
    """Прогон тика уровня с LLM и трассой. Возвращает (text, usage_dict, trace_id)."""
    setup_provider()
    from strands_harness import create_harness

    trace_id = trace_prefix or uuid.uuid4().hex[:12]
    instructions = level_cfg.mandate
    agent = create_harness(
        model=build_level_model(level_cfg),
        effort="auto",  # effort уже в параметрах модели (build_level_model)
        instructions=instructions,
        session={"id": session_id} if session_id else None,
        builtin_tools=("read", "write", "edit", "shell") if level_cfg.can_shell else [],
        memory=False,
        skills=False,
    )
    agent = _wire_trace_hooks(agent, es, trace_id)
    result = agent(prompt)
    # AgentResult: текст из message, usage из metrics.accumulated_usage
    # (старый разбор agent.messages молчал → BudgetSpent всегда был пуст).
    text = str(result)
    usage = _extract_usage(result)
    return text, usage, trace_id


def _extract_usage(result):
    """Usage из AgentResult.metrics.accumulated_usage (input/output/total)."""
    try:
        u = getattr(getattr(result, "metrics", None), "accumulated_usage", None)
        if u:
            return {
                "prompt_tokens": u.get("inputTokens"),
                "completion_tokens": u.get("outputTokens"),
                "total_tokens": u.get("totalTokens"),
            }
    except Exception:
        pass
    return {}