Setup check

MODEL_ID = os.getenv("MODEL_ID", "gpt-6-luna")


def make_model() -> OpenAIModel:
    """Every task talks to the model through the Opsfleet LiteLLM gateway (keyless, budgeted per participant)."""
    return litellm_gateway.make_model(MODEL_ID)

class StsIdentityAuth(httpx.Auth):
    """Adds a fresh presigned STS identity URL to each request (cached until shortly before it expires).

    Short-lived (2 min) because the URL is a bearer credential for your LLM budget until it expires.
    The URL must be signed for GET: STS is a query-protocol API, so boto3 presigns for POST by default and the
    gateway's GET replay would fail with SignatureDoesNotMatch.
    """

    def __init__(self, session: boto3.Session | None = None, region: str | None = None) -> None:
        self._session = session or boto3.Session()
        self._region = region or self._session.region_name or os.getenv("AWS_REGION", "us-east-1")
        self._url = ""
        self._expires_at = 0.0

    def identity_url(self) -> str:
        if time.time() > self._expires_at - REFRESH_MARGIN_SECONDS:
            sts = self._session.client("sts", region_name=self._region)
            self._url = sts.generate_presigned_url("get_caller_identity", ExpiresIn=PRESIGN_SECONDS, HttpMethod="GET")
            self._expires_at = time.time() + PRESIGN_SECONDS
        return self._url

    def auth_flow(self, request: httpx.Request):
        request.headers[IDENTITY_HEADER] = self.identity_url()
        yield request


class GatewayModel(OpenAIModel):
    """Strands OpenAIModel bound to the gateway.

    Strands opens (and closes) a fresh OpenAI client per request from `_resolve_client_args()`, so the keyless
    httpx client must be created per request too; the STS signer (and its cached URL) is shared.
    """

    def __init__(self, base_url: str, api_key: str | None = None, session: boto3.Session | None = None, **config):
        super().__init__(
            client_args={"base_url": f"{base_url.rstrip('/')}/v1", "api_key": api_key or KEYLESS_PLACEHOLDER}, **config
        )
        self._auth = None if api_key else StsIdentityAuth(session)

    def _resolve_client_args(self) -> dict:
        args = dict(self.client_args)
        if self._auth:
            args["http_client"] = httpx.AsyncClient(auth=self._auth, timeout=600)
        return args

    async def stream(self, *args: Any, **kwargs: Any) -> AsyncGenerator[StreamEvent]:
        """Stream like OpenAIModel, but raise `EmptyModelResponseError` if the reply carried no output at all."""
        produced = False
        async for event in super().stream(*args, **kwargs):
            produced = produced or has_output(event)
            yield event
        if not produced:
            raise EmptyModelResponseError("the model returned an empty response (no text, tool call or tokens)")


def make_model(model_id: str | None = None, base_url: str | None = None, **params) -> OpenAIModel:
    """Strands model routed through the gateway; reads MODEL_ID / LITELLM_BASE_URL / LITELLM_API_KEY from env."""
    return GatewayModel(
        base_url or os.environ["LITELLM_BASE_URL"],
        os.getenv("LITELLM_API_KEY"),
        model_id=model_id or os.getenv("MODEL_ID", "gpt-6-luna"),
        params=params or None,
    )