Aether/_deprecated_py_src/services/provider/adapters/kiro/request.py

"""Helpers for building Kiro generateAssistantResponse requests."""

from __future__ import annotations

from typing import Any

from src.services.provider.adapters.kiro.constants import KIRO_GENERATE_ASSISTANT_PATH
from src.services.provider.adapters.kiro.context import KiroRequestContext
from src.services.provider.adapters.kiro.converter import (
    convert_claude_messages_to_conversation_state,
)
from src.services.provider.adapters.kiro.headers import build_generate_assistant_headers
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
from src.services.provider.adapters.kiro.token_manager import generate_machine_id


def is_kiro_thinking_enabled(request_body: dict[str, Any]) -> bool:
    thinking = request_body.get("thinking")
    if not isinstance(thinking, dict):
        return False
    ttype = str(thinking.get("type") or "").strip().lower()
    return ttype in {"enabled", "adaptive"}


def build_kiro_request_context(
    request_body: dict[str, Any],
    *,
    cfg: KiroAuthConfig,
) -> KiroRequestContext:
    return KiroRequestContext(
        region=cfg.effective_api_region(),
        machine_id=generate_machine_id(cfg),
        kiro_version=cfg.kiro_version,
        system_version=cfg.system_version,
        node_version=cfg.node_version,
        thinking_enabled=is_kiro_thinking_enabled(request_body),
    )


def build_kiro_request_headers(
    cfg: KiroAuthConfig,
    *,
    access_token: str | None = None,
) -> dict[str, str]:
    region = cfg.effective_api_region()
    host = f"q.{region}.amazonaws.com"
    return build_generate_assistant_headers(
        host=host,
        access_token=access_token,
        machine_id=generate_machine_id(cfg),
        kiro_version=cfg.kiro_version,
        system_version=cfg.system_version,
        node_version=cfg.node_version,
    )


def resolve_kiro_base_url(base_url: str, *, cfg: KiroAuthConfig) -> str:
    resolved = str(base_url or "").rstrip("/")
    region = cfg.effective_api_region()
    if "{region}" in resolved:
        resolved = resolved.replace("{region}", region)
    return resolved


def build_kiro_generate_assistant_url(base_url: str, *, cfg: KiroAuthConfig) -> str:
    resolved = resolve_kiro_base_url(base_url, cfg=cfg)
    if resolved.endswith(KIRO_GENERATE_ASSISTANT_PATH):
        return resolved
    return f"{resolved}{KIRO_GENERATE_ASSISTANT_PATH}"


def build_kiro_inference_config(request_body: dict[str, Any]) -> dict[str, Any] | None:
    inference_config: dict[str, Any] = {}

    max_tokens = request_body.get("max_tokens")
    try:
        max_tokens_i = int(max_tokens) if max_tokens is not None else 0
    except Exception:
        max_tokens_i = 0
    if max_tokens_i > 0:
        inference_config["maxTokens"] = max_tokens_i

    temperature = request_body.get("temperature")
    try:
        temperature_f = float(temperature) if temperature is not None else None
    except Exception:
        temperature_f = None
    if temperature_f is not None and temperature_f >= 0:
        inference_config["temperature"] = temperature_f

    top_p = request_body.get("top_p")
    try:
        top_p_f = float(top_p) if top_p is not None else None
    except Exception:
        top_p_f = None
    if top_p_f is not None and top_p_f > 0:
        inference_config["topP"] = top_p_f

    return inference_config or None


def get_profile_arn_for_payload(cfg: KiroAuthConfig) -> str | None:
    profile_arn = str(cfg.profile_arn or "").strip()
    if not profile_arn:
        return None

    from src.services.provider.adapters.kiro.models.credentials import _normalize_auth_method

    if _normalize_auth_method(cfg.auth_method) == "idc":
        return None

    return profile_arn


def build_kiro_request_payload(
    request_body: dict[str, Any],
    *,
    model: str,
    cfg: KiroAuthConfig,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "conversationState": convert_claude_messages_to_conversation_state(
            request_body,
            model=model,
        )
    }

    inference_config = build_kiro_inference_config(request_body)
    if inference_config:
        payload["inferenceConfig"] = inference_config

    profile_arn = get_profile_arn_for_payload(cfg)
    if profile_arn:
        payload["profileArn"] = profile_arn

    return payload


__all__ = [
    "build_kiro_generate_assistant_url",
    "build_kiro_inference_config",
    "build_kiro_request_context",
    "build_kiro_request_headers",
    "build_kiro_request_payload",
    "get_profile_arn_for_payload",
    "is_kiro_thinking_enabled",
    "resolve_kiro_base_url",
]