dify/api/controllers/inner_api/agent/llm.py

"""Trusted streaming LLM gateway for dify-agent runs."""

from __future__ import annotations

import json
from collections.abc import Generator

from dify_agent.adapters.llm.provider import model_chunk_to_wire
from flask import Response, stream_with_context
from flask_restx import Resource
from pydantic import ValidationError

from controllers.common.schema import register_schema_models
from controllers.inner_api import inner_api_ns
from controllers.inner_api.wraps import agent_endpoint_api_only
from core.errors.error import QuotaExceededError
from libs.exception import BaseHTTPException
from runtime_contracts.builder.execution import BuilderExecutionPolicyError
from runtime_contracts.builder.native.observation import format_native_model_failure
from services.agent_llm_inner_service import (
    AgentLLMInnerService,
    AgentLLMInnerServiceError,
    NativeAgentModelDispatchError,
)
from services.dify_builder.native_testing.agent_bridge import claim_agent_endpoint
from services.entities.agent_llm_inner import AgentLLMInvokeRequest


class AgentLLMInvokeHttpError(BaseHTTPException):
    error_code = "agent_llm_invoke_failed"
    description = "Agent LLM invocation failed."
    code = 500

    def __init__(self, *, error_code: str, description: str, status_code: int) -> None:
        self.error_code = error_code
        self.description = description
        self.code = status_code
        super().__init__(description)


register_schema_models(inner_api_ns, AgentLLMInvokeRequest)


@inner_api_ns.route("/agent/llm/invoke")
class AgentLLMInvokeApi(Resource):
    """Resolve one dify-agent model request before proxying it."""

    @agent_endpoint_api_only
    @inner_api_ns.doc("inner_agent_llm_invoke")
    @inner_api_ns.expect(inner_api_ns.models[AgentLLMInvokeRequest.__name__])
    @inner_api_ns.produces(["text/event-stream"])
    def post(self) -> Response:
        try:
            payload = AgentLLMInvokeRequest.model_validate(inner_api_ns.payload or {})
        except ValidationError as exc:
            raise AgentLLMInvokeHttpError(
                error_code="invalid_request",
                description=str(exc),
                status_code=400,
            ) from exc

        try:
            claim_agent_endpoint(
                boundary="model",
                caller=payload.caller.model_dump(mode="json"),
                envelope=payload.native_dispatch,
                target={
                    "plugin_id": "/".join(payload.target.provider.split("/")[:2]),
                    "provider": payload.target.provider,
                    "model": payload.target.model,
                },
                parameters={**payload.target.model_dump(mode="json", exclude={"provider", "model"}), "stream": True},
                kind="model",
            )
        except BuilderExecutionPolicyError:
            raise AgentLLMInvokeHttpError(
                error_code="native_agent_control_refused", description="Native Agent dispatch refused.", status_code=409
            ) from None
        service = AgentLLMInnerService()
        try:
            prepared = service.prepare(payload)
        except BuilderExecutionPolicyError:
            raise AgentLLMInvokeHttpError(
                error_code="native_agent_control_refused", description="Native Agent dispatch refused.", status_code=409
            ) from None
        except AgentLLMInnerServiceError as exc:
            raise AgentLLMInvokeHttpError(
                error_code=exc.error_code,
                description=exc.description,
                status_code=exc.status_code,
            ) from exc
        except QuotaExceededError as exc:
            raise AgentLLMInvokeHttpError(
                error_code="agent_llm_quota_exceeded",
                description=str(exc) or "Insufficient Message Credits.",
                status_code=429,
            ) from exc
        except ValueError as exc:
            raise AgentLLMInvokeHttpError(
                error_code="invalid_model_request",
                description=str(exc),
                status_code=400,
            ) from exc

        def generate() -> Generator[str]:
            try:
                for chunk in service.invoke(prepared):
                    # Agent owns its history. Do not reserialize opaque request
                    # messages through the runtime's typed PromptMessage field.
                    envelope = {
                        "code": 0,
                        "message": "",
                        "data": model_chunk_to_wire(chunk),
                    }
                    yield f"data: {json.dumps(envelope, ensure_ascii=False, separators=(',', ':'))}\n\n"
            except GeneratorExit:
                raise
            except QuotaExceededError as exc:
                error = {
                    "error_type": "AgentLLMQuotaExceededError",
                    "message": (
                        "Insufficient hosted model quota remaining."
                        if payload.native_dispatch is not None
                        else str(exc) or "Insufficient hosted model quota remaining."
                    ),
                }
                envelope = {
                    "code": -429,
                    "message": json.dumps(error, ensure_ascii=False, separators=(",", ":")),
                    "data": None,
                }
                yield f"data: {json.dumps(envelope, ensure_ascii=False, separators=(',', ':'))}\n\n"
            except Exception as exc:
                error = {
                    "error_type": "NativeAgentModelDispatchFailed"
                    if payload.native_dispatch is not None
                    else type(exc).__name__,
                    "message": (
                        format_native_model_failure(exc.diagnostic)
                        if isinstance(exc, NativeAgentModelDispatchError)
                        else "Native Agent model invocation failed."
                    )
                    if payload.native_dispatch is not None
                    else str(exc) or "Agent LLM invocation failed.",
                }
                envelope = {
                    "code": -500,
                    "message": json.dumps(error, ensure_ascii=False, separators=(",", ":")),
                    "data": None,
                }
                yield f"data: {json.dumps(envelope, ensure_ascii=False, separators=(',', ':'))}\n\n"

        return Response(
            stream_with_context(generate()),  # pyrefly: ignore[no-matching-overload]
            content_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )


__all__ = ["AgentLLMInvokeApi", "AgentLLMInvokeHttpError"]