diff --git a/api/.env.example b/api/.env.example index 9d0f17fe664..b12b2a208d8 100644 --- a/api/.env.example +++ b/api/.env.example @@ -830,6 +830,12 @@ ENABLE_HUMAN_INPUT_TIMEOUT_TASK=true # Human input timeout check interval in minutes HUMAN_INPUT_TIMEOUT_TASK_INTERVAL=1 +# Reconcile Agent LLM gateway calls interrupted before a terminal ledger update. +ENABLE_AGENT_LLM_INVOCATION_RECONCILIATION_TASK=true +AGENT_LLM_INVOCATION_RECONCILIATION_INTERVAL=5 +AGENT_LLM_INVOCATION_STALE_AFTER_SECONDS=900 +AGENT_LLM_INVOCATION_RECONCILIATION_BATCH_SIZE=100 + # Nacos remote settings source HTTP timeouts (seconds). # Bound how long requests to the Nacos endpoint wait before failing, so a slow or # unresponsive Nacos server cannot stall API startup or token refresh. diff --git a/api/clients/agent_backend/request_builder.py b/api/clients/agent_backend/request_builder.py index 61d74eb4a4b..9934a7b647f 100644 --- a/api/clients/agent_backend/request_builder.py +++ b/api/clients/agent_backend/request_builder.py @@ -414,7 +414,9 @@ class AgentBackendRunRequestBuilder: plugin_id=run_input.model.plugin_id, model_provider=run_input.model.model_provider, model=run_input.model.model, - credentials=run_input.model.credentials, + # The API gateway resolves live model credentials for every + # invocation; do not transmit them to the Agent runtime. + credentials={}, model_settings=_agent_model_settings(run_input.model.model_settings), ), ) @@ -607,7 +609,7 @@ class AgentBackendRunRequestBuilder: plugin_id=run_input.model.plugin_id, model_provider=run_input.model.model_provider, model=run_input.model.model, - credentials=run_input.model.credentials, + credentials={}, model_settings=_agent_model_settings(run_input.model.model_settings), ), ), diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 82c46cf8d26..dff3685a3a1 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1434,6 +1434,23 @@ class CeleryScheduleTasksConfig(BaseSettings): default=30, ) + ENABLE_AGENT_LLM_INVOCATION_RECONCILIATION_TASK: bool = Field( + description="Enable reconciliation for interrupted Agent LLM gateway invocations", + default=True, + ) + AGENT_LLM_INVOCATION_RECONCILIATION_INTERVAL: PositiveInt = Field( + description="Agent LLM invocation reconciliation interval in minutes", + default=5, + ) + AGENT_LLM_INVOCATION_STALE_AFTER_SECONDS: PositiveInt = Field( + description="Age in seconds after which a non-terminal Agent LLM invocation is considered interrupted", + default=900, + ) + AGENT_LLM_INVOCATION_RECONCILIATION_BATCH_SIZE: PositiveInt = Field( + description="Maximum interrupted Agent LLM invocations reconciled per task run", + default=100, + ) + # Trigger provider refresh (simple version) ENABLE_TRIGGER_PROVIDER_REFRESH_TASK: bool = Field( description="Enable trigger provider refresh poller", diff --git a/api/controllers/inner_api/__init__.py b/api/controllers/inner_api/__init__.py index 277e0f5ec0e..32dabf5bb01 100644 --- a/api/controllers/inner_api/__init__.py +++ b/api/controllers/inner_api/__init__.py @@ -18,6 +18,7 @@ inner_api_ns = Namespace("inner_api", description="Internal API operations", pat from . import mail as _mail from . import runtime_credentials as _runtime_credentials from .agent import files as _agent_files +from .agent import llm as _agent_llm from .agent import tools as _agent_tools from .app import dsl as _app_dsl from .knowledge import retrieval as _knowledge_retrieval @@ -32,6 +33,7 @@ __all__ = [ "_agent_config", "_agent_drive", "_agent_files", + "_agent_llm", "_agent_tools", "_app_dsl", "_knowledge_retrieval", diff --git a/api/controllers/inner_api/agent/llm.py b/api/controllers/inner_api/agent/llm.py new file mode 100644 index 00000000000..b611c99b335 --- /dev/null +++ b/api/controllers/inner_api/agent/llm.py @@ -0,0 +1,122 @@ +"""Trusted streaming LLM gateway for dify-agent runs.""" + +from __future__ import annotations + +import json +import logging +import time +from collections.abc import Generator + +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_inner_api_only +from core.errors.error import QuotaExceededError +from libs.exception import BaseHTTPException +from services.agent_llm_inner_service import AgentLLMInnerService, AgentLLMInnerServiceError +from services.entities.agent_llm_inner import AgentLLMInvokeRequest + +logger = logging.getLogger(__name__) +_HEARTBEAT_INTERVAL_SECONDS = 60 + + +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 and meter one dify-agent model request before proxying it.""" + + @agent_inner_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 + + service = AgentLLMInnerService() + try: + prepared = service.prepare(payload) + service.mark_running(prepared.invocation_id) + 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, None, None]: + usage = None + last_heartbeat = time.monotonic() + try: + for chunk in service.invoke(prepared): + now = time.monotonic() + if now - last_heartbeat >= _HEARTBEAT_INTERVAL_SECONDS: + try: + service.heartbeat(prepared.invocation_id) + except Exception: + logger.exception("Failed to heartbeat Agent LLM invocation %s", prepared.invocation_id) + last_heartbeat = now + if chunk.delta.usage is not None: + usage = chunk.delta.usage + envelope = {"code": 0, "message": "", "data": chunk.model_dump(mode="json")} + yield f"data: {json.dumps(envelope, ensure_ascii=False, separators=(',', ':'))}\n\n" + service.mark_succeeded(prepared.invocation_id, usage) + except GeneratorExit as exc: + service.mark_failed(prepared.invocation_id, exc, usage) + raise + except Exception as exc: + service.mark_failed(prepared.invocation_id, exc, usage) + error = { + "error_type": type(exc).__name__, + "message": 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"] diff --git a/api/core/app/apps/agent_app/app_generator.py b/api/core/app/apps/agent_app/app_generator.py index a291b84b39b..57d6575e878 100644 --- a/api/core/app/apps/agent_app/app_generator.py +++ b/api/core/app/apps/agent_app/app_generator.py @@ -219,6 +219,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): agent_config_snapshot_id=agent_config_id, agent_config_version_kind=agent_config_version_kind, agent_session_scope_config_version_id=session_scope_config_version_id, + agent_llm_gateway_enabled=True, ) conversation, message = self._init_generate_records( @@ -349,6 +350,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): agent_id=agent.id, agent_config_snapshot_id=agent_config_id, agent_config_version_kind=agent_config_version_kind, + agent_llm_gateway_enabled=True, ) conversation, message = self._init_generate_records( diff --git a/api/core/app/entities/app_invoke_entities.py b/api/core/app/entities/app_invoke_entities.py index b4e2ced08aa..c59208c4aaf 100644 --- a/api/core/app/entities/app_invoke_entities.py +++ b/api/core/app/entities/app_invoke_entities.py @@ -238,6 +238,7 @@ class AgentAppGenerateEntity(ChatAppGenerateEntity): agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot" agent_session_scope_config_version_id: str | None = None prompt_file_mappings: Sequence[JsonValue] = Field(default_factory=list) + agent_llm_gateway_enabled: bool = False class AdvancedChatAppGenerateEntity(ConversationAppGenerateEntity): diff --git a/api/events/event_handlers/update_provider_when_message_created.py b/api/events/event_handlers/update_provider_when_message_created.py index 15b40afdbf2..a2ed5568c99 100644 --- a/api/events/event_handlers/update_provider_when_message_created.py +++ b/api/events/event_handlers/update_provider_when_message_created.py @@ -9,7 +9,11 @@ from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session from configs import dify_config -from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity +from core.app.entities.app_invoke_entities import ( + AgentAppGenerateEntity, + AgentChatAppGenerateEntity, + ChatAppGenerateEntity, +) from core.entities.provider_entities import ProviderQuotaType, QuotaUnit, SystemConfiguration from events.message_event import message_was_created from extensions.ext_database import db @@ -121,8 +125,13 @@ def handle(sender: Message, **kwargs): provider_model_bundle = model_config.provider_model_bundle provider_configuration = provider_model_bundle.configuration + agent_gateway_metered = ( + isinstance(application_generate_entity, AgentAppGenerateEntity) + and application_generate_entity.agent_llm_gateway_enabled + ) if ( - provider_configuration.using_provider_type == ProviderType.SYSTEM + not agent_gateway_metered + and provider_configuration.using_provider_type == ProviderType.SYSTEM and provider_configuration.system_configuration and provider_configuration.system_configuration.current_quota_type is not None ): diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index a96d50e1caf..b267e380bb8 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -275,6 +275,13 @@ def init_app(app: DifyApp) -> Celery: "schedule": timedelta(minutes=dify_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL), } + if dify_config.ENABLE_AGENT_LLM_INVOCATION_RECONCILIATION_TASK: + imports.append("schedule.reconcile_agent_llm_invocations_task") + beat_schedule["reconcile_agent_llm_invocations"] = { + "task": "schedule.reconcile_agent_llm_invocations_task.reconcile_agent_llm_invocations", + "schedule": timedelta(minutes=dify_config.AGENT_LLM_INVOCATION_RECONCILIATION_INTERVAL), + } + if ( dify_config.EDITION == "SELF_HOSTED" and not dify_config.ENTERPRISE_ENABLED diff --git a/api/migrations/versions/2026_08_12_1200-39f8d6c14a21_add_agent_llm_invocation_ledger.py b/api/migrations/versions/2026_08_12_1200-39f8d6c14a21_add_agent_llm_invocation_ledger.py new file mode 100644 index 00000000000..4bc23c9d38b --- /dev/null +++ b/api/migrations/versions/2026_08_12_1200-39f8d6c14a21_add_agent_llm_invocation_ledger.py @@ -0,0 +1,74 @@ +"""add agent llm invocation ledger + +Revision ID: 39f8d6c14a21 +Revises: e4708db55c1d +Create Date: 2026-08-12 12:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +import models + +revision = "39f8d6c14a21" +down_revision = "e4708db55c1d" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "agent_llm_invocations", + sa.Column("invocation_id", models.types.StringUUID(), nullable=False), + sa.Column("tenant_id", models.types.StringUUID(), nullable=False), + sa.Column("agent_run_id", models.types.StringUUID(), nullable=False), + sa.Column("call_index", sa.Integer(), nullable=False), + sa.Column("agent_mode", sa.String(length=32), nullable=False), + sa.Column("invoke_from", sa.String(length=32), nullable=False), + sa.Column("user_id", models.types.StringUUID(), nullable=False), + sa.Column("user_from", sa.String(length=16), nullable=False), + sa.Column("app_id", models.types.StringUUID(), nullable=False), + sa.Column("workflow_id", models.types.StringUUID(), nullable=True), + sa.Column("workflow_run_id", models.types.StringUUID(), nullable=True), + sa.Column("node_id", sa.String(length=255), nullable=True), + sa.Column("node_execution_id", sa.String(length=255), nullable=True), + sa.Column("conversation_id", models.types.StringUUID(), nullable=True), + sa.Column("agent_id", models.types.StringUUID(), nullable=True), + sa.Column("agent_config_version_id", models.types.StringUUID(), nullable=True), + sa.Column("agent_config_version_kind", sa.String(length=16), nullable=True), + sa.Column("trace_id", sa.String(length=255), nullable=True), + sa.Column("provider", sa.String(length=255), nullable=False), + sa.Column("model", sa.String(length=255), nullable=False), + sa.Column("credential_source", sa.String(length=16), nullable=False), + sa.Column("quota_type", sa.String(length=32), nullable=True), + sa.Column("pool_type", sa.String(length=32), nullable=True), + sa.Column("credits", sa.Integer(), server_default="0", nullable=False), + sa.Column("billing_status", sa.String(length=24), server_default="pending", nullable=False), + sa.Column("execution_status", sa.String(length=16), server_default="prepared", nullable=False), + sa.Column("usage", models.types.LongText(), nullable=True), + sa.Column("error_type", sa.String(length=255), nullable=True), + sa.Column("error_message", models.types.LongText(), nullable=True), + sa.Column("started_at", sa.DateTime(), nullable=True), + sa.Column("finished_at", sa.DateTime(), nullable=True), + sa.Column("id", models.types.StringUUID(), nullable=False), + sa.Column("created_at", sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False), + sa.Column("updated_at", sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("agent_llm_invocation_pkey")), + sa.UniqueConstraint("invocation_id", name=op.f("agent_llm_invocation_id_unique")), + sa.UniqueConstraint("agent_run_id", "call_index", name=op.f("agent_llm_invocation_run_call_unique")), + ) + with op.batch_alter_table("agent_llm_invocations", schema=None) as batch_op: + batch_op.create_index("agent_llm_invocation_billing_status_idx", ["billing_status", "updated_at"], unique=False) + batch_op.create_index( + "agent_llm_invocation_execution_status_idx", ["execution_status", "updated_at"], unique=False + ) + batch_op.create_index("agent_llm_invocation_tenant_created_idx", ["tenant_id", "created_at"], unique=False) + + +def downgrade() -> None: + with op.batch_alter_table("agent_llm_invocations", schema=None) as batch_op: + batch_op.drop_index("agent_llm_invocation_tenant_created_idx") + batch_op.drop_index("agent_llm_invocation_execution_status_idx") + batch_op.drop_index("agent_llm_invocation_billing_status_idx") + op.drop_table("agent_llm_invocations") diff --git a/api/models/__init__.py b/api/models/__init__.py index 387a6761820..8f3ba618dc9 100644 --- a/api/models/__init__.py +++ b/api/models/__init__.py @@ -22,6 +22,10 @@ from .agent import ( AgentHomeSnapshot, AgentIconType, AgentKind, + AgentLLMBillingStatus, + AgentLLMCredentialSource, + AgentLLMExecutionStatus, + AgentLLMInvocation, AgentScope, AgentSource, AgentStatus, @@ -172,6 +176,10 @@ __all__ = [ "AgentHomeSnapshot", "AgentIconType", "AgentKind", + "AgentLLMBillingStatus", + "AgentLLMCredentialSource", + "AgentLLMExecutionStatus", + "AgentLLMInvocation", "AgentScope", "AgentSource", "AgentStatus", diff --git a/api/models/agent.py b/api/models/agent.py index a5863dd1259..fed8cdb171d 100644 --- a/api/models/agent.py +++ b/api/models/agent.py @@ -7,6 +7,7 @@ import sqlalchemy as sa from sqlalchemy import DateTime, Index, String, UniqueConstraint, func from sqlalchemy.orm import Mapped, mapped_column +from graphon.model_runtime.entities.llm_entities import LLMUsage from libs.datetime_utils import naive_utc_now from libs.uuid_utils import uuidv7 @@ -137,6 +138,26 @@ class AgentConfigVersionKind(StrEnum): BUILD_DRAFT = "build_draft" +class AgentLLMCredentialSource(StrEnum): + SYSTEM = "system" + CUSTOM = "custom" + + +class AgentLLMBillingStatus(StrEnum): + PENDING = "pending" + REJECTED = "rejected" + INDETERMINATE = "indeterminate" + NOT_BILLABLE = "not_billable" + CHARGED = "charged" + + +class AgentLLMExecutionStatus(StrEnum): + PREPARED = "prepared" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + class Agent(DefaultFieldsMixin, Base): """Agent Soul and source lineage; ``AgentWorkspaceBinding.id`` identifies each materialized participant.""" @@ -538,6 +559,66 @@ class AgentWorkspaceBinding(DefaultFieldsMixin, Base): pending_tool_call_id: Mapped[str | None] = mapped_column(String(255), nullable=True) +class AgentLLMInvocation(DefaultFieldsMixin, Base): + """API-owned ledger for one model request made by a dify-agent run.""" + + __tablename__ = "agent_llm_invocations" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="agent_llm_invocation_pkey"), + UniqueConstraint("invocation_id", name="agent_llm_invocation_id_unique"), + UniqueConstraint("agent_run_id", "call_index", name="agent_llm_invocation_run_call_unique"), + Index("agent_llm_invocation_tenant_created_idx", "tenant_id", "created_at"), + Index("agent_llm_invocation_billing_status_idx", "billing_status", "updated_at"), + Index("agent_llm_invocation_execution_status_idx", "execution_status", "updated_at"), + ) + + invocation_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + agent_run_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + call_index: Mapped[int] = mapped_column(sa.Integer, nullable=False) + agent_mode: Mapped[str] = mapped_column(String(32), nullable=False) + invoke_from: Mapped[str] = mapped_column(String(32), nullable=False) + + user_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + user_from: Mapped[str] = mapped_column(String(16), nullable=False) + app_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + workflow_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) + workflow_run_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) + node_id: Mapped[str | None] = mapped_column(String(255), nullable=True) + node_execution_id: Mapped[str | None] = mapped_column(String(255), nullable=True) + conversation_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) + agent_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) + agent_config_version_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) + agent_config_version_kind: Mapped[str | None] = mapped_column(String(16), nullable=True) + trace_id: Mapped[str | None] = mapped_column(String(255), nullable=True) + + provider: Mapped[str] = mapped_column(String(255), nullable=False) + model: Mapped[str] = mapped_column(String(255), nullable=False) + credential_source: Mapped[AgentLLMCredentialSource] = mapped_column( + EnumText(AgentLLMCredentialSource, length=16), nullable=False + ) + quota_type: Mapped[str | None] = mapped_column(String(32), nullable=True) + pool_type: Mapped[str | None] = mapped_column(String(32), nullable=True) + credits: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0, server_default="0") + billing_status: Mapped[AgentLLMBillingStatus] = mapped_column( + EnumText(AgentLLMBillingStatus, length=24), + nullable=False, + default=AgentLLMBillingStatus.PENDING, + server_default=AgentLLMBillingStatus.PENDING.value, + ) + execution_status: Mapped[AgentLLMExecutionStatus] = mapped_column( + EnumText(AgentLLMExecutionStatus, length=16), + nullable=False, + default=AgentLLMExecutionStatus.PREPARED, + server_default=AgentLLMExecutionStatus.PREPARED.value, + ) + usage: Mapped[LLMUsage | None] = mapped_column(JSONModelColumn(LLMUsage), nullable=True) + error_type: Mapped[str | None] = mapped_column(String(255), nullable=True) + error_message: Mapped[str | None] = mapped_column(LongText, nullable=True) + started_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) + finished_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) + + class AgentDriveFileKind(StrEnum): """Kind of existing file record an agent-drive KV entry points at.""" diff --git a/api/schedule/reconcile_agent_llm_invocations_task.py b/api/schedule/reconcile_agent_llm_invocations_task.py new file mode 100644 index 00000000000..3a747a131a6 --- /dev/null +++ b/api/schedule/reconcile_agent_llm_invocations_task.py @@ -0,0 +1,23 @@ +"""Reconcile Agent LLM gateway calls left non-terminal by interrupted API workers.""" + +import logging +from datetime import timedelta + +import app +from configs import dify_config +from services.agent_llm_inner_service import AgentLLMInnerService + +logger = logging.getLogger(__name__) + + +@app.celery.task(queue="retention") +def reconcile_agent_llm_invocations() -> None: + reconciled = AgentLLMInnerService.reconcile_stale( + stale_after=timedelta(seconds=dify_config.AGENT_LLM_INVOCATION_STALE_AFTER_SECONDS), + limit=dify_config.AGENT_LLM_INVOCATION_RECONCILIATION_BATCH_SIZE, + ) + if reconciled: + logger.info("Reconciled %d interrupted Agent LLM invocation(s)", reconciled) + + +__all__ = ["reconcile_agent_llm_invocations"] diff --git a/api/services/agent_llm_inner_service.py b/api/services/agent_llm_inner_service.py new file mode 100644 index 00000000000..dc3f26ed719 --- /dev/null +++ b/api/services/agent_llm_inner_service.py @@ -0,0 +1,474 @@ +"""API-owned model invocation and Message Credits accounting for dify-agent.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Generator +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import cast + +from sqlalchemy import or_, select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session + +from configs import dify_config +from core.db.session_factory import session_factory as default_session_factory +from core.entities.provider_entities import ProviderQuotaType, QuotaUnit +from core.errors.error import QuotaExceededError +from core.model_manager import ModelInstance, ModelManager +from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager +from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMUsage +from graphon.model_runtime.entities.model_entities import ModelType +from libs.datetime_utils import naive_utc_now +from models import ( + AgentLLMBillingStatus, + AgentLLMCredentialSource, + AgentLLMExecutionStatus, + AgentLLMInvocation, + Provider, + ProviderType, +) +from models.model import App +from models.provider_ids import ModelProviderID +from services.credit_pool_service import CreditPoolService +from services.entities.agent_llm_inner import AgentLLMInvokeRequest + +logger = logging.getLogger(__name__) + + +class AgentLLMInnerServiceError(RuntimeError): + def __init__(self, error_code: str, description: str, *, status_code: int = 500) -> None: + self.error_code = error_code + self.description = description + self.status_code = status_code + super().__init__(description) + + +@dataclass(frozen=True) +class _BillingPlan: + credential_source: AgentLLMCredentialSource + quota_type: str | None = None + pool_type: str | None = None + credits: int = 0 + + @property + def billable(self) -> bool: + return self.credits > 0 and self.quota_type is not None + + +@dataclass(frozen=True) +class PreparedAgentLLMInvocation: + request: AgentLLMInvokeRequest + model_instance: ModelInstance + + @property + def invocation_id(self) -> str: + return self.request.caller.invocation_id + + +class AgentLLMInnerService: + """Resolve credentials, charge one invocation, and proxy it through the API model runtime.""" + + def __init__(self, *, session_factory: Callable[[], Session] | None = None) -> None: + self._session_factory = session_factory or default_session_factory.create_session + + def prepare(self, request: AgentLLMInvokeRequest) -> PreparedAgentLLMInvocation: + caller = request.caller + target = request.target + self._validate_app_tenant(app_id=caller.app_id, tenant_id=caller.tenant_id) + provider_manager = create_plugin_provider_manager(tenant_id=caller.tenant_id, user_id=caller.user_id) + model_manager = ModelManager(provider_manager=provider_manager) + model_instance = model_manager.get_model_instance( + tenant_id=caller.tenant_id, + provider=target.provider, + model_type=ModelType.LLM, + model=target.model, + ) + + provider_configuration = model_instance.provider_model_bundle.configuration + provider_model = provider_configuration.get_provider_model(model_type=ModelType.LLM, model=target.model) + if provider_model is None: + raise AgentLLMInnerServiceError( + "model_not_found", + f"Model {target.model} does not exist for provider {target.provider}.", + status_code=404, + ) + provider_model.raise_for_status() + + plan = self._build_billing_plan(model_instance) + created = self._create_ledger(request=request, plan=plan) + if not created: + raise AgentLLMInnerServiceError( + "duplicate_invocation", + f"Agent LLM invocation {caller.invocation_id} has already been accepted.", + status_code=409, + ) + + if plan.billable: + try: + charged = self._charge_pending_invocation(caller.invocation_id) + except QuotaExceededError as exc: + self._mark_billing_rejected(caller.invocation_id, exc) + raise AgentLLMInnerServiceError( + "agent_llm_quota_exceeded", + str(exc) or "Insufficient Message Credits.", + status_code=429, + ) from exc + if not charged: + raise AgentLLMInnerServiceError( + "duplicate_invocation", + f"Agent LLM invocation {caller.invocation_id} has already been charged.", + status_code=409, + ) + + return PreparedAgentLLMInvocation(request=request, model_instance=model_instance) + + def invoke(self, prepared: PreparedAgentLLMInvocation) -> Generator[LLMResultChunk, None, None]: + request = prepared.request + caller = request.caller + target = request.target + result = prepared.model_instance.invoke_llm( + prompt_messages=target.prompt_messages, + model_parameters=target.model_parameters, + tools=target.tools, + stop=target.stop, + # The gateway transport is always streamed, including Pydantic AI's + # non-streaming request path, so one response protocol is sufficient. + stream=True, + request_metadata={ + "source": "agent_llm_gateway", + "invocation_id": caller.invocation_id, + "agent_run_id": caller.agent_run_id, + "agent_mode": caller.agent_mode, + "call_index": caller.call_index, + "app_id": caller.app_id, + "workflow_run_id": caller.workflow_run_id, + "node_execution_id": caller.node_execution_id, + "trace_id": caller.trace_id, + }, + ) + yield from cast(Generator[LLMResultChunk, None, None], result) + + def mark_running(self, invocation_id: str) -> None: + with self._session_factory() as session: + invocation = self._get_invocation_for_update(session, invocation_id) + if invocation.execution_status == AgentLLMExecutionStatus.PREPARED: + invocation.execution_status = AgentLLMExecutionStatus.RUNNING + invocation.started_at = naive_utc_now() + session.commit() + + def mark_succeeded(self, invocation_id: str, usage: LLMUsage | None) -> None: + with self._session_factory() as session: + invocation = self._get_invocation_for_update(session, invocation_id) + if invocation.execution_status in { + AgentLLMExecutionStatus.SUCCEEDED, + AgentLLMExecutionStatus.FAILED, + }: + return + invocation.execution_status = AgentLLMExecutionStatus.SUCCEEDED + invocation.usage = usage + invocation.finished_at = naive_utc_now() + invocation.error_type = None + invocation.error_message = None + session.commit() + + def mark_failed(self, invocation_id: str, error: BaseException, usage: LLMUsage | None = None) -> None: + with self._session_factory() as session: + invocation = self._get_invocation_for_update(session, invocation_id) + if invocation.execution_status == AgentLLMExecutionStatus.SUCCEEDED: + return + invocation.execution_status = AgentLLMExecutionStatus.FAILED + invocation.usage = usage + invocation.finished_at = naive_utc_now() + invocation.error_type = type(error).__name__ + invocation.error_message = str(error) + session.commit() + + def heartbeat(self, invocation_id: str) -> None: + with self._session_factory() as session: + invocation = self._get_invocation_for_update(session, invocation_id) + if invocation.execution_status == AgentLLMExecutionStatus.RUNNING: + invocation.updated_at = naive_utc_now() + session.commit() + + @classmethod + def reconcile_stale(cls, *, stale_after: timedelta, limit: int = 100) -> int: + """Close interrupted invocations without replaying uncertain charges.""" + cutoff = naive_utc_now() - stale_after + service = cls() + with service._session_factory() as session: + invocation_ids = list( + session.scalars( + select(AgentLLMInvocation.invocation_id) + .where( + AgentLLMInvocation.updated_at < cutoff, + AgentLLMInvocation.execution_status.in_( + [AgentLLMExecutionStatus.PREPARED, AgentLLMExecutionStatus.RUNNING] + ), + ) + .order_by(AgentLLMInvocation.updated_at) + .limit(limit) + ) + ) + + reconciled = 0 + for invocation_id in invocation_ids: + try: + if service._reconcile_one(invocation_id, cutoff=cutoff): + reconciled += 1 + except Exception: + logger.exception("Failed to reconcile Agent LLM invocation %s", invocation_id) + return reconciled + + def _reconcile_one(self, invocation_id: str, *, cutoff: datetime) -> bool: + with self._session_factory() as session: + invocation = self._get_invocation_for_update(session, invocation_id) + if invocation.updated_at >= cutoff or invocation.execution_status not in { + AgentLLMExecutionStatus.PREPARED, + AgentLLMExecutionStatus.RUNNING, + }: + return False + if invocation.billing_status == AgentLLMBillingStatus.PENDING: + # No model call starts before charging returns successfully. A + # stale PENDING row is therefore never replayed into a new + # charge; an external billing timeout may have an uncertain + # outcome, which is recorded explicitly for reconciliation. + invocation.billing_status = AgentLLMBillingStatus.INDETERMINATE + invocation.execution_status = AgentLLMExecutionStatus.FAILED + invocation.finished_at = naive_utc_now() + invocation.error_type = "InterruptedAgentLLMInvocation" + invocation.error_message = "Agent LLM invocation was interrupted before a terminal event." + session.commit() + return True + + def _validate_app_tenant(self, *, app_id: str, tenant_id: str) -> None: + with self._session_factory() as session: + app = session.get(App, app_id) + if app is None: + raise AgentLLMInnerServiceError( + "app_not_found", + "App not found.", + status_code=404, + ) + if app.tenant_id != tenant_id: + raise AgentLLMInnerServiceError( + "app_tenant_mismatch", + "App does not belong to the caller tenant.", + status_code=403, + ) + + def _create_ledger(self, *, request: AgentLLMInvokeRequest, plan: _BillingPlan) -> bool: + caller = request.caller + target = request.target + billing_status = AgentLLMBillingStatus.PENDING if plan.billable else AgentLLMBillingStatus.NOT_BILLABLE + invocation = AgentLLMInvocation( + invocation_id=caller.invocation_id, + tenant_id=caller.tenant_id, + agent_run_id=caller.agent_run_id, + call_index=caller.call_index, + agent_mode=caller.agent_mode, + invoke_from=caller.invoke_from, + user_id=caller.user_id, + user_from=caller.user_from, + app_id=caller.app_id, + workflow_id=caller.workflow_id, + workflow_run_id=caller.workflow_run_id, + node_id=caller.node_id, + node_execution_id=caller.node_execution_id, + conversation_id=caller.conversation_id, + agent_id=caller.agent_id, + agent_config_version_id=caller.agent_config_version_id, + agent_config_version_kind=caller.agent_config_version_kind, + trace_id=caller.trace_id, + provider=target.provider, + model=target.model, + credential_source=plan.credential_source, + quota_type=plan.quota_type, + pool_type=plan.pool_type, + credits=plan.credits, + billing_status=billing_status, + execution_status=AgentLLMExecutionStatus.PREPARED, + ) + with self._session_factory() as session: + session.add(invocation) + try: + session.commit() + return True + except IntegrityError: + session.rollback() + existing = session.scalar( + select(AgentLLMInvocation).where( + or_( + AgentLLMInvocation.invocation_id == caller.invocation_id, + ( + (AgentLLMInvocation.agent_run_id == caller.agent_run_id) + & (AgentLLMInvocation.call_index == caller.call_index) + ), + ) + ) + ) + if existing is None: + raise + self._validate_existing_identity(existing, request) + return False + + def _charge_pending_invocation(self, invocation_id: str) -> bool: + with self._session_factory() as session: + invocation = self._get_invocation_for_update(session, invocation_id) + if invocation.billing_status != AgentLLMBillingStatus.PENDING: + return False + if invocation.credits <= 0 or invocation.quota_type is None: + invocation.billing_status = AgentLLMBillingStatus.NOT_BILLABLE + session.commit() + return True + + invocation.billing_status = AgentLLMBillingStatus.CHARGED + if invocation.quota_type in {ProviderQuotaType.TRIAL.value, ProviderQuotaType.PAID.value}: + CreditPoolService.check_and_deduct_credits( + tenant_id=invocation.tenant_id, + credits_required=invocation.credits, + pool_type=invocation.pool_type or invocation.quota_type, + request_id=invocation.invocation_id, + metadata=self._billing_metadata(invocation), + session=session, + ) + elif invocation.quota_type == ProviderQuotaType.FREE.value: + self._deduct_free_quota(session, invocation) + else: + raise AgentLLMInnerServiceError( + "unsupported_quota_type", + f"Unsupported Agent LLM quota type: {invocation.quota_type}", + ) + session.commit() + return True + + @staticmethod + def _build_billing_plan(model_instance: ModelInstance) -> _BillingPlan: + configuration = model_instance.provider_model_bundle.configuration + if configuration.using_provider_type != ProviderType.SYSTEM: + return _BillingPlan(credential_source=AgentLLMCredentialSource.CUSTOM) + + system_configuration = configuration.system_configuration + quota_type = system_configuration.current_quota_type + if quota_type is None: + return _BillingPlan(credential_source=AgentLLMCredentialSource.SYSTEM) + + quota_configuration = next( + (item for item in system_configuration.quota_configurations if item.quota_type == quota_type), + None, + ) + if quota_configuration is None or quota_configuration.quota_limit == -1: + return _BillingPlan( + credential_source=AgentLLMCredentialSource.SYSTEM, + quota_type=quota_type.value, + ) + + if quota_configuration.quota_unit == QuotaUnit.CREDITS: + credits = dify_config.get_model_credits(model_instance.model_name) + elif quota_configuration.quota_unit == QuotaUnit.TIMES: + credits = 1 + else: + raise AgentLLMInnerServiceError( + "unsupported_quota_unit", + "Agent LLM Gateway supports fixed Message Credits quotas only.", + status_code=422, + ) + return _BillingPlan( + credential_source=AgentLLMCredentialSource.SYSTEM, + quota_type=quota_type.value, + pool_type=quota_type.value if quota_type in {ProviderQuotaType.TRIAL, ProviderQuotaType.PAID} else None, + credits=credits, + ) + + @staticmethod + def _deduct_free_quota(session: Session, invocation: AgentLLMInvocation) -> None: + provider_record = session.scalar( + select(Provider) + .where( + Provider.tenant_id == invocation.tenant_id, + Provider.provider_name == ModelProviderID(invocation.provider).provider_name, + Provider.provider_type == ProviderType.SYSTEM.value, + Provider.quota_type == ProviderQuotaType.FREE, + ) + .limit(1) + .with_for_update() + ) + if ( + provider_record is None + or provider_record.quota_limit is None + or provider_record.quota_used is None + or provider_record.quota_limit - provider_record.quota_used < invocation.credits + ): + raise QuotaExceededError("Insufficient hosted model quota remaining") + provider_record.quota_used += invocation.credits + provider_record.last_used = naive_utc_now() + + @staticmethod + def _billing_metadata(invocation: AgentLLMInvocation) -> dict[str, str]: + return { + "source": "agent_llm_gateway", + "invocation_id": invocation.invocation_id, + "agent_run_id": invocation.agent_run_id, + "agent_mode": invocation.agent_mode, + "call_index": str(invocation.call_index), + "provider": invocation.provider, + "model": invocation.model, + } + + @staticmethod + def _validate_existing_identity(existing: AgentLLMInvocation, request: AgentLLMInvokeRequest) -> None: + caller = request.caller + target = request.target + expected = ( + caller.tenant_id, + caller.agent_run_id, + caller.call_index, + caller.agent_mode, + caller.user_id, + caller.app_id, + target.provider, + target.model, + ) + actual = ( + existing.tenant_id, + existing.agent_run_id, + existing.call_index, + existing.agent_mode, + existing.user_id, + existing.app_id, + existing.provider, + existing.model, + ) + if actual != expected: + raise AgentLLMInnerServiceError( + "invocation_identity_conflict", + "The invocation_id is already bound to a different Agent LLM request.", + status_code=409, + ) + + @staticmethod + def _get_invocation_for_update(session: Session, invocation_id: str) -> AgentLLMInvocation: + invocation = session.scalar( + select(AgentLLMInvocation).where(AgentLLMInvocation.invocation_id == invocation_id).with_for_update() + ) + if invocation is None: + raise AgentLLMInnerServiceError( + "invocation_not_found", + f"Agent LLM invocation {invocation_id} was not found.", + status_code=404, + ) + return invocation + + def _mark_billing_rejected(self, invocation_id: str, error: BaseException) -> None: + with self._session_factory() as session: + invocation = self._get_invocation_for_update(session, invocation_id) + invocation.billing_status = AgentLLMBillingStatus.REJECTED + invocation.execution_status = AgentLLMExecutionStatus.FAILED + invocation.finished_at = naive_utc_now() + invocation.error_type = type(error).__name__ + invocation.error_message = str(error) + session.commit() + + +__all__ = ["AgentLLMInnerService", "AgentLLMInnerServiceError", "PreparedAgentLLMInvocation"] diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py index f398ed6e6e4..deea8245ffc 100644 --- a/api/services/credit_pool_service.py +++ b/api/services/credit_pool_service.py @@ -6,7 +6,7 @@ from piling up database transactions while preserving cross-tenant concurrency. """ import logging -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from uuid import uuid4 @@ -169,6 +169,8 @@ class CreditPoolService: credits_required: int, pool_type: str | ProviderQuotaType = "trial", *, + request_id: str | None = None, + metadata: Mapping[str, str] | None = None, session: Session | None = None, ) -> int: """Deduct exactly the requested credits or raise without mutating the pool.""" @@ -179,14 +181,15 @@ class CreditPoolService: if cls._use_billing_quota(): from services.billing_service import BillingService - request_id = str(uuid4()) + resolved_request_id = request_id or str(uuid4()) + billing_metadata = {"source": "credit_pool.check_and_deduct", **dict(metadata or {})} result = BillingService.quota_reserve( tenant_id=tenant_id, feature_key=FEATURE_KEY_CREDIT_POOL, bucket=normalized_pool_type, - request_id=request_id, + request_id=resolved_request_id, amount=credits_required, - meta={"source": "credit_pool.check_and_deduct"}, + meta=billing_metadata, ) reservation_id = result.get("reservation_id", "") if not reservation_id: @@ -198,7 +201,7 @@ class CreditPoolService: bucket=normalized_pool_type, reservation_id=reservation_id, actual_amount=credits_required, - meta={"source": "credit_pool.check_and_deduct"}, + meta=billing_metadata, ) except Exception: try: diff --git a/api/services/entities/agent_llm_inner.py b/api/services/entities/agent_llm_inner.py new file mode 100644 index 00000000000..918d2d0c120 --- /dev/null +++ b/api/services/entities/agent_llm_inner.py @@ -0,0 +1,63 @@ +"""DTOs for the API-owned Agent LLM gateway.""" + +from __future__ import annotations + +from typing import ClassVar, Literal + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from graphon.model_runtime.entities.message_entities import PromptMessage, PromptMessageTool + +type AgentLLMMode = Literal["workflow_run", "single_step", "agent_app", "babysit", "fasten"] +type AgentConfigVersionKind = Literal["snapshot", "draft", "build_draft"] + + +class AgentLLMInvokeCaller(BaseModel): + invocation_id: str + agent_run_id: str + call_index: int = Field(ge=1) + tenant_id: str + user_id: str + user_from: Literal["account", "end-user"] + app_id: str + invoke_from: str + agent_mode: AgentLLMMode + conversation_id: str | None = None + workflow_id: str | None = None + workflow_run_id: str | None = None + node_id: str | None = None + node_execution_id: str | None = None + agent_id: str | None = None + agent_config_version_id: str | None = None + agent_config_version_kind: AgentConfigVersionKind | None = None + trace_id: str | None = None + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +class AgentLLMInvokeTarget(BaseModel): + provider: str + model: str + prompt_messages: list[PromptMessage] + model_parameters: dict[str, JsonValue] = Field(default_factory=dict) + tools: list[PromptMessageTool] | None = None + stop: list[str] | None = None + stream: bool = True + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + +class AgentLLMInvokeRequest(BaseModel): + caller: AgentLLMInvokeCaller + target: AgentLLMInvokeTarget + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + +__all__ = [ + "AgentConfigVersionKind", + "AgentLLMInvokeCaller", + "AgentLLMInvokeRequest", + "AgentLLMInvokeTarget", + "AgentLLMMode", +] diff --git a/api/tests/unit_tests/clients/agent_backend/test_request_builder.py b/api/tests/unit_tests/clients/agent_backend/test_request_builder.py index 6f2a10452c5..c504e6a3ba0 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_request_builder.py +++ b/api/tests/unit_tests/clients/agent_backend/test_request_builder.py @@ -207,7 +207,9 @@ def test_request_builder_sets_model_and_output_layer_contract_ids(): assert execution_context_config.invoke_from == "debugger" assert layers[DIFY_AGENT_HISTORY_LAYER_ID].type == PYDANTIC_AI_HISTORY_LAYER_TYPE_ID assert layers[DIFY_AGENT_MODEL_LAYER_ID].type == DIFY_PLUGIN_LLM_LAYER_TYPE_ID - assert cast(DifyPluginLLMLayerConfig, layers[DIFY_AGENT_MODEL_LAYER_ID].config).plugin_id == "langgenius/openai" + model_config = cast(DifyPluginLLMLayerConfig, layers[DIFY_AGENT_MODEL_LAYER_ID].config) + assert model_config.plugin_id == "langgenius/openai" + assert model_config.credentials == {} assert layers[DIFY_AGENT_MODEL_LAYER_ID].deps == {"execution_context": DIFY_EXECUTION_CONTEXT_LAYER_ID} assert layers[DIFY_AGENT_OUTPUT_LAYER_ID].type == DIFY_OUTPUT_LAYER_TYPE_ID diff --git a/api/tests/unit_tests/controllers/inner_api/test_agent_llm.py b/api/tests/unit_tests/controllers/inner_api/test_agent_llm.py new file mode 100644 index 00000000000..59650b6f66f --- /dev/null +++ b/api/tests/unit_tests/controllers/inner_api/test_agent_llm.py @@ -0,0 +1,140 @@ +"""Contract tests for the trusted Agent LLM streaming gateway.""" + +import json +from collections.abc import Generator +from contextlib import contextmanager +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +from flask import Flask + +from controllers.inner_api import bp as inner_api_bp +from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta +from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage +from services.agent_llm_inner_service import AgentLLMInnerServiceError, PreparedAgentLLMInvocation +from services.entities.agent_llm_inner import AgentLLMInvokeRequest + + +def _payload() -> dict[str, object]: + return { + "caller": { + "invocation_id": str(uuid4()), + "agent_run_id": str(uuid4()), + "call_index": 1, + "tenant_id": str(uuid4()), + "user_id": str(uuid4()), + "user_from": "account", + "app_id": str(uuid4()), + "invoke_from": "debugger", + "agent_mode": "workflow_run", + "agent_config_version_kind": "draft", + }, + "target": { + "provider": "openai", + "model": "gpt-test", + "prompt_messages": [UserPromptMessage(content="hello").model_dump(mode="json")], + "model_parameters": {"temperature": 0.2}, + "stream": True, + }, + } + + +@contextmanager +def _agent_inner_auth() -> Generator[None]: + with ( + patch("configs.dify_config.PLUGIN_DAEMON_KEY", "plugin-daemon-key"), + patch("configs.dify_config.INNER_API_KEY_FOR_PLUGIN", "inner-key"), + ): + yield + + +def _app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(inner_api_bp) + return app + + +def test_post_streams_plugin_compatible_envelope_and_marks_terminal_status() -> None: + payload = _payload() + request = AgentLLMInvokeRequest.model_validate(payload) + prepared = PreparedAgentLLMInvocation(request=request, model_instance=MagicMock()) + chunk = LLMResultChunk( + model="gpt-test", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content="done", tool_calls=[]), + ), + ) + + with ( + _agent_inner_auth(), + patch("controllers.inner_api.agent.llm.AgentLLMInnerService.prepare", return_value=prepared) as prepare, + patch("controllers.inner_api.agent.llm.AgentLLMInnerService.mark_running") as mark_running, + patch("controllers.inner_api.agent.llm.AgentLLMInnerService.invoke", return_value=iter([chunk])), + patch("controllers.inner_api.agent.llm.AgentLLMInnerService.mark_succeeded") as mark_succeeded, + ): + response = ( + _app() + .test_client() + .post( + "/inner/api/agent/llm/invoke", + json=payload, + headers={"X-Inner-Api-Key": "inner-key"}, + ) + ) + assert response.status_code == 200 + assert response.content_type == "text/event-stream" + data_line = response.get_data(as_text=True).strip().removeprefix("data: ") + envelope = json.loads(data_line) + assert envelope["code"] == 0 + assert envelope["data"]["delta"]["message"]["content"] == "done" + prepare.assert_called_once() + mark_running.assert_called_once_with(request.caller.invocation_id) + mark_succeeded.assert_called_once_with(request.caller.invocation_id, None) + + +def test_post_rejects_invalid_body_before_model_resolution() -> None: + with _agent_inner_auth(): + response = ( + _app() + .test_client() + .post( + "/inner/api/agent/llm/invoke", + json={"caller": {}}, + headers={"X-Inner-Api-Key": "inner-key"}, + ) + ) + + assert response.status_code == 400 + assert response.get_json()["code"] == "invalid_request" + + +def test_post_preserves_preflight_quota_failure() -> None: + with ( + _agent_inner_auth(), + patch( + "controllers.inner_api.agent.llm.AgentLLMInnerService.prepare", + side_effect=AgentLLMInnerServiceError( + "agent_llm_quota_exceeded", + "Insufficient Message Credits.", + status_code=429, + ), + ), + ): + response = ( + _app() + .test_client() + .post( + "/inner/api/agent/llm/invoke", + json=_payload(), + headers={"X-Inner-Api-Key": "inner-key"}, + ) + ) + + assert response.status_code == 429 + assert response.get_json() == { + "code": "agent_llm_quota_exceeded", + "message": "Insufficient Message Credits.", + "status": 429, + } diff --git a/api/tests/unit_tests/events/test_update_provider_when_message_created.py b/api/tests/unit_tests/events/test_update_provider_when_message_created.py index a700783b31c..7c55c949fab 100644 --- a/api/tests/unit_tests/events/test_update_provider_when_message_created.py +++ b/api/tests/unit_tests/events/test_update_provider_when_message_created.py @@ -8,7 +8,7 @@ from sqlalchemy import select from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker -from core.app.entities.app_invoke_entities import ChatAppGenerateEntity +from core.app.entities.app_invoke_entities import AgentAppGenerateEntity, ChatAppGenerateEntity from core.entities.provider_entities import ProviderQuotaType, QuotaUnit from events.event_handlers import update_provider_when_message_created from models import Message, TenantCreditPool @@ -121,6 +121,48 @@ def test_message_created_paid_credit_accounting_uses_paid_pool() -> None: ) +def test_agent_app_gateway_accounting_skips_legacy_message_charge() -> None: + tenant_id = str(uuid4()) + system_configuration = SimpleNamespace( + current_quota_type=ProviderQuotaType.TRIAL, + quota_configurations=[ + SimpleNamespace( + quota_type=ProviderQuotaType.TRIAL, + quota_unit=QuotaUnit.CREDITS, + quota_limit=10, + ) + ], + ) + application_generate_entity = AgentAppGenerateEntity.model_construct( + app_config=SimpleNamespace(tenant_id=tenant_id), + model_conf=SimpleNamespace( + provider="openai", + model="gpt-4o", + provider_model_bundle=SimpleNamespace( + configuration=SimpleNamespace( + using_provider_type=ProviderType.SYSTEM, + system_configuration=system_configuration, + ) + ), + ), + agent_llm_gateway_enabled=True, + ) + + with ( + patch.object(update_provider_when_message_created, "_deduct_credit_pool_quota_capped") as mock_deduct, + patch.object(update_provider_when_message_created, "_execute_provider_updates") as mock_updates, + ): + update_provider_when_message_created.handle( + sender=Message(message_tokens=2, answer_tokens=1), + application_generate_entity=application_generate_entity, + ) + + mock_deduct.assert_not_called() + mock_updates.assert_called_once() + assert len(mock_updates.call_args.args[0]) == 1 + assert mock_updates.call_args.args[0][0].description == "basic_last_used_update" + + def test_capped_credit_pool_accounting_skips_exhaustion_warning_when_full_amount_is_deducted( caplog: pytest.LogCaptureFixture, ) -> None: diff --git a/api/tests/unit_tests/extensions/test_celery_ssl.py b/api/tests/unit_tests/extensions/test_celery_ssl.py index 784f6b6417f..58511df49f9 100644 --- a/api/tests/unit_tests/extensions/test_celery_ssl.py +++ b/api/tests/unit_tests/extensions/test_celery_ssl.py @@ -176,6 +176,7 @@ class TestCelerySSLConfiguration: mock_config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15 mock_config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False mock_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30 + mock_config.ENABLE_AGENT_LLM_INVOCATION_RECONCILIATION_TASK = False with patch("extensions.ext_celery.dify_config", mock_config): from dify_app import DifyApp @@ -226,6 +227,7 @@ class TestCelerySSLConfiguration: mock_config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15 mock_config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False mock_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30 + mock_config.ENABLE_AGENT_LLM_INVOCATION_RECONCILIATION_TASK = False mock_config.ENTERPRISE_ENABLED = False mock_config.ENTERPRISE_TELEMETRY_ENABLED = False diff --git a/api/tests/unit_tests/services/test_agent_llm_inner_service.py b/api/tests/unit_tests/services/test_agent_llm_inner_service.py new file mode 100644 index 00000000000..30629e51258 --- /dev/null +++ b/api/tests/unit_tests/services/test_agent_llm_inner_service.py @@ -0,0 +1,325 @@ +"""Agent LLM gateway accounting tests backed by the unit-test SQLite database.""" + +from datetime import timedelta +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest +from sqlalchemy import select, update +from sqlalchemy.orm import Session, sessionmaker + +from graphon.model_runtime.entities.message_entities import UserPromptMessage +from libs.datetime_utils import naive_utc_now +from models import ( + AgentLLMBillingStatus, + AgentLLMCredentialSource, + AgentLLMExecutionStatus, + AgentLLMInvocation, + TenantCreditPool, +) +from models.enums import ProviderQuotaType +from models.model import App, AppMode +from services.agent_llm_inner_service import AgentLLMInnerService, AgentLLMInnerServiceError, _BillingPlan +from services.entities.agent_llm_inner import AgentLLMInvokeCaller, AgentLLMInvokeRequest, AgentLLMInvokeTarget + + +def _request(*, invocation_id: str | None = None, agent_run_id: str | None = None) -> AgentLLMInvokeRequest: + return AgentLLMInvokeRequest( + caller=AgentLLMInvokeCaller( + invocation_id=invocation_id or str(uuid4()), + agent_run_id=agent_run_id or str(uuid4()), + call_index=1, + tenant_id=str(uuid4()), + user_id=str(uuid4()), + user_from="account", + app_id=str(uuid4()), + invoke_from="debugger", + agent_mode="workflow_run", + agent_config_version_kind="draft", + trace_id="trace-1", + ), + target=AgentLLMInvokeTarget( + provider="openai", + model="gpt-test", + prompt_messages=[UserPromptMessage(content="hello")], + ), + ) + + +def _model_instance() -> MagicMock: + provider_model = MagicMock() + configuration = MagicMock() + configuration.get_provider_model.return_value = provider_model + model_instance = MagicMock() + model_instance.provider_model_bundle.configuration = configuration + return model_instance + + +def _prepare_with_plan( + service: AgentLLMInnerService, + request: AgentLLMInvokeRequest, + plan: _BillingPlan, + *, + session: Session, +) -> None: + _persist_app(session, request=request) + manager = MagicMock() + manager.get_model_instance.return_value = _model_instance() + with ( + patch("services.agent_llm_inner_service.create_plugin_provider_manager"), + patch("services.agent_llm_inner_service.ModelManager", return_value=manager), + patch.object(service, "_build_billing_plan", return_value=plan), + ): + service.prepare(request) + + +def _persist_app( + session: Session, + *, + request: AgentLLMInvokeRequest, + tenant_id: str | None = None, +) -> App: + existing = session.get(App, request.caller.app_id) + if existing is not None: + return existing + app = App( + id=request.caller.app_id, + tenant_id=tenant_id or request.caller.tenant_id, + name="Agent LLM gateway test app", + description="", + mode=AppMode.CHAT, + enable_site=False, + enable_api=False, + max_active_requests=None, + ) + session.add(app) + session.commit() + return app + + +def _create_pool(session: Session, *, tenant_id: str, quota_limit: int, quota_used: int = 0) -> TenantCreditPool: + pool = TenantCreditPool( + tenant_id=tenant_id, + pool_type=ProviderQuotaType.TRIAL, + quota_limit=quota_limit, + quota_used=quota_used, + ) + session.add(pool) + session.commit() + return pool + + +def test_prepare_rejects_missing_app(sqlite_session_factory: sessionmaker[Session]) -> None: + request = _request() + service = AgentLLMInnerService(session_factory=sqlite_session_factory) + + with pytest.raises(AgentLLMInnerServiceError) as exc_info: + service.prepare(request) + + assert exc_info.value.error_code == "app_not_found" + assert exc_info.value.status_code == 404 + + +def test_prepare_rejects_cross_tenant_app( + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, +) -> None: + request = _request() + _persist_app(sqlite_session, request=request, tenant_id=str(uuid4())) + service = AgentLLMInnerService(session_factory=sqlite_session_factory) + + with pytest.raises(AgentLLMInnerServiceError) as exc_info: + service.prepare(request) + + assert exc_info.value.error_code == "app_tenant_mismatch" + assert exc_info.value.status_code == 403 + + +def test_system_invocation_is_charged_once_per_run_call( + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, +) -> None: + request = _request() + pool = _create_pool(sqlite_session, tenant_id=request.caller.tenant_id, quota_limit=10) + service = AgentLLMInnerService(session_factory=sqlite_session_factory) + plan = _BillingPlan( + credential_source=AgentLLMCredentialSource.SYSTEM, + quota_type="trial", + pool_type="trial", + credits=3, + ) + + _prepare_with_plan(service, request, plan, session=sqlite_session) + + sqlite_session.expire_all() + invocation = sqlite_session.scalar( + select(AgentLLMInvocation).where(AgentLLMInvocation.invocation_id == request.caller.invocation_id) + ) + assert invocation is not None + assert invocation.billing_status == AgentLLMBillingStatus.CHARGED + assert invocation.execution_status == AgentLLMExecutionStatus.PREPARED + assert invocation.user_from == "account" + assert invocation.agent_config_version_kind == "draft" + assert invocation.trace_id == "trace-1" + assert sqlite_session.get(TenantCreditPool, pool.id).quota_used == 3 + + with pytest.raises(AgentLLMInnerServiceError, match="already been accepted"): + _prepare_with_plan(service, request, plan, session=sqlite_session) + + changed_id_request = request.model_copy( + update={"caller": request.caller.model_copy(update={"invocation_id": str(uuid4())})} + ) + with pytest.raises(AgentLLMInnerServiceError, match="already been accepted"): + _prepare_with_plan(service, changed_id_request, plan, session=sqlite_session) + + sqlite_session.expire_all() + assert sqlite_session.get(TenantCreditPool, pool.id).quota_used == 3 + + +def test_custom_credentials_create_non_billable_ledger( + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, +) -> None: + request = _request() + service = AgentLLMInnerService(session_factory=sqlite_session_factory) + + _prepare_with_plan( + service, + request, + _BillingPlan(credential_source=AgentLLMCredentialSource.CUSTOM), + session=sqlite_session, + ) + + invocation = sqlite_session.scalar( + select(AgentLLMInvocation).where(AgentLLMInvocation.invocation_id == request.caller.invocation_id) + ) + assert invocation is not None + assert invocation.credential_source == AgentLLMCredentialSource.CUSTOM + assert invocation.billing_status == AgentLLMBillingStatus.NOT_BILLABLE + assert invocation.credits == 0 + + +def test_insufficient_credits_reject_before_model_invocation( + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, +) -> None: + request = _request() + pool = _create_pool(sqlite_session, tenant_id=request.caller.tenant_id, quota_limit=2) + service = AgentLLMInnerService(session_factory=sqlite_session_factory) + plan = _BillingPlan( + credential_source=AgentLLMCredentialSource.SYSTEM, + quota_type="trial", + pool_type="trial", + credits=3, + ) + + with pytest.raises(AgentLLMInnerServiceError) as exc_info: + _prepare_with_plan(service, request, plan, session=sqlite_session) + + assert exc_info.value.error_code == "agent_llm_quota_exceeded" + assert exc_info.value.status_code == 429 + sqlite_session.expire_all() + invocation = sqlite_session.scalar( + select(AgentLLMInvocation).where(AgentLLMInvocation.invocation_id == request.caller.invocation_id) + ) + assert invocation is not None + assert invocation.billing_status == AgentLLMBillingStatus.REJECTED + assert invocation.execution_status == AgentLLMExecutionStatus.FAILED + assert sqlite_session.get(TenantCreditPool, pool.id).quota_used == 0 + + +def test_terminal_status_does_not_regress_from_succeeded( + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, +) -> None: + request = _request() + service = AgentLLMInnerService(session_factory=sqlite_session_factory) + _prepare_with_plan( + service, + request, + _BillingPlan(credential_source=AgentLLMCredentialSource.CUSTOM), + session=sqlite_session, + ) + + service.mark_running(request.caller.invocation_id) + service.mark_succeeded(request.caller.invocation_id, usage=None) + service.mark_failed(request.caller.invocation_id, RuntimeError("late disconnect")) + + sqlite_session.expire_all() + invocation = sqlite_session.scalar( + select(AgentLLMInvocation).where(AgentLLMInvocation.invocation_id == request.caller.invocation_id) + ) + assert invocation is not None + assert invocation.execution_status == AgentLLMExecutionStatus.SUCCEEDED + assert invocation.error_message is None + + +def test_reconciliation_does_not_create_a_charge_for_unconfirmed_pending_invocation( + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, +) -> None: + request = _request() + pool = _create_pool(sqlite_session, tenant_id=request.caller.tenant_id, quota_limit=10) + service = AgentLLMInnerService(session_factory=sqlite_session_factory) + plan = _BillingPlan( + credential_source=AgentLLMCredentialSource.SYSTEM, + quota_type="trial", + pool_type="trial", + credits=2, + ) + assert service._create_ledger(request=request, plan=plan) + sqlite_session.execute( + update(AgentLLMInvocation) + .where(AgentLLMInvocation.invocation_id == request.caller.invocation_id) + .values(updated_at=naive_utc_now() - timedelta(hours=1)) + ) + sqlite_session.commit() + + assert AgentLLMInnerService.reconcile_stale(stale_after=timedelta(minutes=15)) == 1 + + sqlite_session.expire_all() + invocation = sqlite_session.scalar( + select(AgentLLMInvocation).where(AgentLLMInvocation.invocation_id == request.caller.invocation_id) + ) + assert invocation is not None + assert invocation.billing_status == AgentLLMBillingStatus.INDETERMINATE + assert invocation.execution_status == AgentLLMExecutionStatus.FAILED + assert sqlite_session.get(TenantCreditPool, pool.id).quota_used == 0 + + assert AgentLLMInnerService.reconcile_stale(stale_after=timedelta(minutes=15)) == 0 + sqlite_session.expire_all() + assert sqlite_session.get(TenantCreditPool, pool.id).quota_used == 0 + + +def test_reconciliation_closes_confirmed_charge_without_deducting_again( + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, +) -> None: + request = _request() + pool = _create_pool(sqlite_session, tenant_id=request.caller.tenant_id, quota_limit=10) + service = AgentLLMInnerService(session_factory=sqlite_session_factory) + plan = _BillingPlan( + credential_source=AgentLLMCredentialSource.SYSTEM, + quota_type="trial", + pool_type="trial", + credits=2, + ) + _prepare_with_plan(service, request, plan, session=sqlite_session) + service.mark_running(request.caller.invocation_id) + sqlite_session.execute( + update(AgentLLMInvocation) + .where(AgentLLMInvocation.invocation_id == request.caller.invocation_id) + .values(updated_at=naive_utc_now() - timedelta(hours=1)) + ) + sqlite_session.commit() + + assert AgentLLMInnerService.reconcile_stale(stale_after=timedelta(minutes=15)) == 1 + + sqlite_session.expire_all() + invocation = sqlite_session.scalar( + select(AgentLLMInvocation).where(AgentLLMInvocation.invocation_id == request.caller.invocation_id) + ) + assert invocation is not None + assert invocation.billing_status == AgentLLMBillingStatus.CHARGED + assert invocation.execution_status == AgentLLMExecutionStatus.FAILED + assert sqlite_session.get(TenantCreditPool, pool.id).quota_used == 2 diff --git a/api/tests/unit_tests/services/test_credit_pool_service.py b/api/tests/unit_tests/services/test_credit_pool_service.py index f3b97e8bbd6..09008f60141 100644 --- a/api/tests/unit_tests/services/test_credit_pool_service.py +++ b/api/tests/unit_tests/services/test_credit_pool_service.py @@ -308,6 +308,45 @@ def test_check_and_deduct_credits_uses_billing_reserve_and_commit_when_enabled() quota_release.assert_not_called() +def test_check_and_deduct_credits_forwards_deterministic_billing_identity() -> None: + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + patch("services.billing_service.BillingService.quota_commit") as quota_commit, + ): + quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3} + + result = CreditPoolService.check_and_deduct_credits( + tenant_id="tenant-1", + credits_required=3, + pool_type="trial", + request_id="invocation-1", + metadata={"agent_run_id": "run-1"}, + ) + + assert result == 3 + expected_metadata = { + "source": "credit_pool.check_and_deduct", + "agent_run_id": "run-1", + } + quota_reserve.assert_called_once_with( + tenant_id="tenant-1", + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + request_id="invocation-1", + amount=3, + meta=expected_metadata, + ) + quota_commit.assert_called_once_with( + tenant_id="tenant-1", + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + reservation_id="reservation-1", + actual_amount=3, + meta=expected_metadata, + ) + + def test_check_and_deduct_credits_raises_when_billing_reserve_is_insufficient() -> None: with ( patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), diff --git a/dify-agent/src/dify_agent/adapters/llm/__init__.py b/dify-agent/src/dify_agent/adapters/llm/__init__.py index b63771ed616..28705650747 100644 --- a/dify-agent/src/dify_agent/adapters/llm/__init__.py +++ b/dify-agent/src/dify_agent/adapters/llm/__init__.py @@ -1,6 +1,6 @@ """LLM adapters for Dify plugin-daemon integrations.""" from .model import DifyLLMAdapterModel -from .provider import DifyPluginDaemonProvider +from .provider import DifyApiLLMProvider, DifyPluginDaemonProvider -__all__ = ["DifyLLMAdapterModel", "DifyPluginDaemonProvider"] +__all__ = ["DifyApiLLMProvider", "DifyLLMAdapterModel", "DifyPluginDaemonProvider"] diff --git a/dify-agent/src/dify_agent/adapters/llm/model.py b/dify-agent/src/dify_agent/adapters/llm/model.py index 3ffa4efd4e3..b9915c0e597 100644 --- a/dify-agent/src/dify_agent/adapters/llm/model.py +++ b/dify-agent/src/dify_agent/adapters/llm/model.py @@ -65,10 +65,11 @@ from pydantic_ai.messages import ( ) from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse from pydantic_ai.profiles import ModelProfileSpec +from pydantic_ai.providers import Provider from pydantic_ai.settings import ModelSettings from pydantic_ai.usage import RequestUsage -from .provider import DifyPluginDaemonLLMClient, DifyPluginDaemonProvider +from .provider import DifyLLMClient _THINK_START = "\n" _THINK_END = "\n" @@ -88,7 +89,7 @@ class _DifyRequestInput: @dataclass(slots=True) -class DifyLLMAdapterModel(Model[DifyPluginDaemonLLMClient]): +class DifyLLMAdapterModel(Model[DifyLLMClient]): """Use a Dify plugin-daemon transport and retain complete usage for one Agent run. A model instance belongs to one runner invocation. Pydantic AI may call it repeatedly while @@ -97,7 +98,7 @@ class DifyLLMAdapterModel(Model[DifyPluginDaemonLLMClient]): """ model: str - daemon_provider: DifyPluginDaemonProvider + daemon_provider: Provider[DifyLLMClient] _: KW_ONLY model_provider: str credentials: dict[str, object] = field(default_factory=dict, repr=False) @@ -118,7 +119,7 @@ class DifyLLMAdapterModel(Model[DifyPluginDaemonLLMClient]): @property @override - def provider(self) -> DifyPluginDaemonProvider: + def provider(self) -> Provider[DifyLLMClient]: return self.daemon_provider @property diff --git a/dify-agent/src/dify_agent/adapters/llm/provider.py b/dify-agent/src/dify_agent/adapters/llm/provider.py index a210cce1e32..5ae7aeebb7d 100644 --- a/dify-agent/src/dify_agent/adapters/llm/provider.py +++ b/dify-agent/src/dify_agent/adapters/llm/provider.py @@ -8,9 +8,12 @@ this provider. from __future__ import annotations +import json from collections.abc import AsyncIterator, Callable, Mapping from dataclasses import dataclass, field -from typing import NoReturn +from itertools import count +from typing import NoReturn, Protocol, cast +from uuid import NAMESPACE_URL, uuid5 import httpx from graphon.model_runtime.entities.llm_entities import LLMResultChunk @@ -21,6 +24,7 @@ from typing_extensions import override from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError, UnexpectedModelBehavior, UserError from pydantic_ai.providers import Provider +from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.plugin_daemon_transport import ( decode_plugin_daemon_error_payload, to_plugin_daemon_jsonable, @@ -30,12 +34,139 @@ from dify_agent.plugin_daemon_transport import ( _DEFAULT_DAEMON_TIMEOUT: float | httpx.Timeout | None = 600.0 +class DifyLLMClient(Protocol): + """Transport contract consumed by the Pydantic AI model adapter.""" + + http_client: httpx.AsyncClient + + def iter_llm_result_chunks( + self, + *, + provider: str, + model: str, + credentials: dict[str, object], + prompt_messages: list[PromptMessage], + model_parameters: dict[str, object], + tools: list[PromptMessageTool] | None, + stop: list[str] | None, + stream: bool, + ) -> AsyncIterator[LLMResultChunk]: ... + + class PluginDaemonBasicResponse(BaseModel): code: int message: str data: object | None = None +@dataclass(slots=True) +class DifyApiLLMClient: + """HTTP client for the API-owned Agent LLM metering gateway.""" + + plugin_id: str + inner_api_url: str + inner_api_key: str = field(repr=False) + execution_context: DifyExecutionContextLayerConfig + agent_run_id: str + http_client: httpx.AsyncClient = field(repr=False) + _call_counter: count[int] = field(default_factory=lambda: count(1), init=False, repr=False) + + def __post_init__(self) -> None: + self.inner_api_url = self.inner_api_url.rstrip("/") + + async def iter_llm_result_chunks( + self, + *, + provider: str, + model: str, + credentials: dict[str, object], + prompt_messages: list[PromptMessage], + model_parameters: dict[str, object], + tools: list[PromptMessageTool] | None, + stop: list[str] | None, + stream: bool, + ) -> AsyncIterator[LLMResultChunk]: + del credentials + call_index = next(self._call_counter) + invocation_id = str(uuid5(NAMESPACE_URL, f"dify-agent:{self.agent_run_id}:llm:{call_index}")) + context = self.execution_context + missing = [ + field_name for field_name in ("user_id", "user_from", "app_id") if getattr(context, field_name) is None + ] + if missing: + raise UserError(f"Agent LLM Gateway requires execution context fields: {', '.join(missing)}") + + caller = context.model_dump(mode="json") + caller.update( + { + "invocation_id": invocation_id, + "agent_run_id": self.agent_run_id, + "call_index": call_index, + } + ) + provider_id = provider if provider.count("/") == 2 else f"{self.plugin_id}/{provider}" + payload = to_plugin_daemon_jsonable( + { + "caller": caller, + "target": { + "provider": provider_id, + "model": model, + "prompt_messages": prompt_messages, + "model_parameters": model_parameters, + "tools": tools, + "stop": stop, + "stream": stream, + }, + } + ) + url = f"{self.inner_api_url}/inner/api/agent/llm/invoke" + headers = { + "X-Inner-Api-Key": self.inner_api_key, + "Content-Type": "application/json", + } + + try: + async with self.http_client.stream("POST", url, headers=headers, json=payload) as response: + if response.is_error: + body = (await response.aread()).decode("utf-8", errors="replace") + _raise_agent_gateway_http_error( + model_name=model, + status_code=response.status_code, + body=body, + ) + + async for raw_line in response.aiter_lines(): + line = raw_line.strip() + if not line: + continue + if line.startswith("data:"): + line = line[5:].strip() + wrapped = PluginDaemonBasicResponse.model_validate_json(line) + if wrapped.code != 0: + error = decode_plugin_daemon_error_payload(wrapped.message) + if error is not None: + resolved_error = unwrap_plugin_daemon_error( + error_type=error["error_type"], + message=error["message"], + ) + _raise_plugin_daemon_error( + model_name=model, + error_type=resolved_error["error_type"], + message=resolved_error["message"], + body=resolved_error, + ) + raise ModelAPIError(model, wrapped.message) + if wrapped.data is None: + raise UnexpectedModelBehavior("Agent LLM Gateway returned an empty stream item") + yield LLMResultChunk.model_validate(wrapped.data) + except (httpx.InvalidURL, httpx.UnsupportedProtocol) as exc: + raise UserError(f"Agent LLM Gateway is misconfigured: {exc}") from exc + except httpx.TimeoutException as exc: + raise ModelHTTPError(504, model, "Agent LLM Gateway timed out") from exc + except httpx.RequestError as exc: + raise ModelHTTPError(503, model, f"Agent LLM Gateway request failed: {exc}") from exc + + @dataclass(slots=True) class DifyPluginDaemonLLMClient: """HTTP client wrapper for plugin-daemon LLM dispatch requests.""" @@ -148,7 +279,7 @@ class DifyPluginDaemonLLMClient: @dataclass(slots=True, kw_only=True) -class DifyPluginDaemonProvider(Provider[DifyPluginDaemonLLMClient]): +class DifyPluginDaemonProvider(Provider[DifyLLMClient]): """Pydantic AI provider for Dify plugin-daemon dispatch requests. The provider ``name`` identifies the daemon/plugin context. The business LLM @@ -166,7 +297,7 @@ class DifyPluginDaemonProvider(Provider[DifyPluginDaemonLLMClient]): user_id: str | None = None timeout: float | httpx.Timeout | None = _DEFAULT_DAEMON_TIMEOUT http_client: httpx.AsyncClient | None = field(default=None, repr=False) - _client: DifyPluginDaemonLLMClient = field(init=False, repr=False) + _client: DifyLLMClient = field(init=False, repr=False) _own_http_client: httpx.AsyncClient | None = field(init=False, default=None, repr=False) _http_client_factory: Callable[[], httpx.AsyncClient] | None = field(init=False, default=None, repr=False) @@ -208,10 +339,67 @@ class DifyPluginDaemonProvider(Provider[DifyPluginDaemonLLMClient]): @property @override - def client(self) -> DifyPluginDaemonLLMClient: + def client(self) -> DifyLLMClient: return self._client +@dataclass(slots=True, kw_only=True) +class DifyApiLLMProvider(Provider[DifyLLMClient]): + """Pydantic AI provider backed by Dify API's metered model gateway.""" + + plugin_id: str + inner_api_url: str + inner_api_key: str = field(repr=False) + execution_context: DifyExecutionContextLayerConfig + agent_run_id: str + http_client: httpx.AsyncClient = field(repr=False) + _client: DifyLLMClient = field(init=False, repr=False) + + def __post_init__(self) -> None: + self.inner_api_url = self.inner_api_url.rstrip("/") + self._client = DifyApiLLMClient( + plugin_id=self.plugin_id, + inner_api_url=self.inner_api_url, + inner_api_key=self.inner_api_key, + execution_context=self.execution_context, + agent_run_id=self.agent_run_id, + http_client=self.http_client, + ) + + @override + def _set_http_client(self, http_client: httpx.AsyncClient) -> None: + self._client.http_client = http_client + + @property + @override + def name(self) -> str: + return f"DifyAPI/{self.plugin_id}" + + @property + @override + def base_url(self) -> str: + return self.inner_api_url + + @property + @override + def client(self) -> DifyLLMClient: + return self._client + + +def _raise_agent_gateway_http_error(*, model_name: str, status_code: int, body: str) -> NoReturn: + message = body or f"Agent LLM Gateway returned HTTP {status_code}" + try: + payload = cast(object, json.loads(body)) + except json.JSONDecodeError: + payload = None + if isinstance(payload, dict): + typed_payload = cast(dict[str, object], payload) + candidate = typed_payload.get("message") or typed_payload.get("description") + if isinstance(candidate, str) and candidate: + message = candidate + raise ModelHTTPError(status_code, model_name, message) + + def _raise_plugin_daemon_error( *, model_name: str, diff --git a/dify-agent/src/dify_agent/layers/dify_plugin/llm_layer.py b/dify-agent/src/dify_agent/layers/dify_plugin/llm_layer.py index bdd9d947565..c7321352c5f 100644 --- a/dify-agent/src/dify_agent/layers/dify_plugin/llm_layer.py +++ b/dify-agent/src/dify_agent/layers/dify_plugin/llm_layer.py @@ -10,14 +10,14 @@ while the DTO's ``model_provider`` is passed to the adapter as request-level model identity. """ -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import ClassVar import httpx from typing_extensions import Self, override from agenton.layers import LayerDeps, PlainLayer -from dify_agent.adapters.llm import DifyLLMAdapterModel +from dify_agent.adapters.llm import DifyApiLLMProvider, DifyLLMAdapterModel from dify_agent.layers.dify_plugin.configs import DIFY_PLUGIN_LLM_LAYER_TYPE_ID, DifyPluginLLMLayerConfig from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer @@ -35,6 +35,8 @@ class DifyPluginLLMLayer(PlainLayer[DifyPluginLLMDeps, DifyPluginLLMLayerConfig] type_id: ClassVar[str | None] = DIFY_PLUGIN_LLM_LAYER_TYPE_ID config: DifyPluginLLMLayerConfig + inner_api_url: str = "http://localhost:5001" + inner_api_key: str = field(default="", repr=False) @classmethod @override @@ -42,17 +44,35 @@ class DifyPluginLLMLayer(PlainLayer[DifyPluginLLMDeps, DifyPluginLLMLayerConfig] """Create the LLM layer from validated public config.""" return cls(config=config) - def get_model(self, *, http_client: httpx.AsyncClient) -> DifyLLMAdapterModel: - """Return the configured model using the directly bound execution context.""" - provider = self.deps.execution_context.create_daemon_provider( + @classmethod + def from_config_with_settings( + cls, + config: DifyPluginLLMLayerConfig, + *, + inner_api_url: str, + inner_api_key: str, + ) -> Self: + return cls(config=config, inner_api_url=inner_api_url.rstrip("/"), inner_api_key=inner_api_key) + + def get_model(self, *, http_client: httpx.AsyncClient, agent_run_id: str) -> DifyLLMAdapterModel: + """Return the configured model through the API-owned metered gateway.""" + if http_client.is_closed: + raise RuntimeError("DifyPluginLLMLayer.get_model() requires an open Dify API HTTP client.") + provider = DifyApiLLMProvider( plugin_id=self.config.plugin_id, + inner_api_url=self.inner_api_url, + inner_api_key=self.inner_api_key, + execution_context=self.deps.execution_context.config, + agent_run_id=agent_run_id, http_client=http_client, ) return DifyLLMAdapterModel( model=self.config.model, daemon_provider=provider, model_provider=self.config.model_provider, - credentials=dict(self.config.credentials), + # Older run snapshots may still contain credentials. The API owns + # credential resolution, so never forward or retain them here. + credentials={}, model_settings=self.config.model_settings, ) diff --git a/dify-agent/src/dify_agent/runtime/compositor_factory.py b/dify-agent/src/dify_agent/runtime/compositor_factory.py index 429c7fce57a..771a0ecb473 100644 --- a/dify-agent/src/dify_agent/runtime/compositor_factory.py +++ b/dify-agent/src/dify_agent/runtime/compositor_factory.py @@ -43,8 +43,8 @@ from dify_agent.layers.ask_human.layer import DifyAskHumanLayer from dify_agent.layers.config.layer import DifyConfigLayer from dify_agent.layers.dify_core_tools.configs import DifyCoreToolsLayerConfig from dify_agent.layers.dify_core_tools.layer import DifyCoreToolsLayer +from dify_agent.layers.dify_plugin.configs import DifyPluginLLMLayerConfig, DifyPluginToolsLayerConfig from dify_agent.layers.dify_plugin.llm_layer import DifyPluginLLMLayer -from dify_agent.layers.dify_plugin.configs import DifyPluginToolsLayerConfig from dify_agent.layers.dify_plugin.tools_layer import DifyPluginToolsLayer from dify_agent.layers.drive.layer import DifyDriveLayer from dify_agent.layers.execution_context.configs import DifyExecutionContextLayerConfig @@ -97,7 +97,14 @@ def create_default_layer_providers( agent_stub_token_factory=agent_stub_token_factory, ), ), - LayerProvider.from_layer_type(DifyPluginLLMLayer), + LayerProvider.from_factory( + layer_type=DifyPluginLLMLayer, + create=lambda config: DifyPluginLLMLayer.from_config_with_settings( + DifyPluginLLMLayerConfig.model_validate(config), + inner_api_url=inner_api_url, + inner_api_key=inner_api_key, + ), + ), LayerProvider.from_factory( layer_type=DifyPluginToolsLayer, create=lambda config: DifyPluginToolsLayer.from_config_with_settings( diff --git a/dify-agent/src/dify_agent/runtime/runner.py b/dify-agent/src/dify_agent/runtime/runner.py index 814447871e5..1b76c561cfa 100644 --- a/dify-agent/src/dify_agent/runtime/runner.py +++ b/dify-agent/src/dify_agent/runtime/runner.py @@ -312,7 +312,10 @@ class AgentRunRunner: ) ask_human_layer = get_ask_human_layer(run) llm_layer = run.get_layer(DIFY_AGENT_MODEL_LAYER_ID, DifyPluginLLMLayer) - model = llm_layer.get_model(http_client=self.plugin_daemon_http_client) + model = llm_layer.get_model( + http_client=self.dify_api_http_client, + agent_run_id=self.run_id, + ) tools = await _resolve_run_tools( run, plugin_daemon_http_client=self.plugin_daemon_http_client, diff --git a/dify-agent/tests/local/dify_agent/adapters/llm/test_api_provider.py b/dify-agent/tests/local/dify_agent/adapters/llm/test_api_provider.py new file mode 100644 index 00000000000..59e924ccbc0 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/adapters/llm/test_api_provider.py @@ -0,0 +1,124 @@ +import asyncio +import json +from uuid import NAMESPACE_URL, uuid5 + +import httpx +import pytest +from graphon.model_runtime.entities.message_entities import UserPromptMessage +from pydantic_ai.exceptions import ModelHTTPError + +from dify_agent.adapters.llm.provider import DifyApiLLMClient +from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig + +from ._test_support import build_stream_response, single_text_chunk + + +def _execution_context() -> DifyExecutionContextLayerConfig: + return DifyExecutionContextLayerConfig( + tenant_id="tenant-1", + user_id="user-1", + user_from="account", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="workflow-run-1", + node_id="node-1", + node_execution_id="execution-1", + agent_config_version_kind="draft", + agent_mode="workflow_run", + invoke_from="debugger", + trace_id="trace-1", + ) + + +def test_api_client_uses_stable_per_run_call_identity_and_omits_credentials() -> None: + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return build_stream_response(*single_text_chunk("done")) + + async def scenario() -> None: + async with httpx.AsyncClient(transport=httpx.MockTransport(handler), trust_env=False) as http_client: + client = DifyApiLLMClient( + plugin_id="acme/custom-model", + inner_api_url="http://dify-api/", + inner_api_key="inner-secret", + execution_context=_execution_context(), + agent_run_id=run_id, + http_client=http_client, + ) + + for _ in range(2): + chunks = [ + chunk + async for chunk in client.iter_llm_result_chunks( + provider="openai", + model="gpt-test", + credentials={"api_key": "must-not-leave-agent"}, + prompt_messages=[UserPromptMessage(content="hello")], + model_parameters={"temperature": 0.2}, + tools=None, + stop=None, + stream=True, + ) + ] + assert chunks[0].delta.message.content == "done" + + run_id = "00000000-0000-0000-0000-000000000001" + asyncio.run(scenario()) + + assert len(requests) == 2 + first_payload = json.loads(requests[0].content) + second_payload = json.loads(requests[1].content) + assert requests[0].url == "http://dify-api/inner/api/agent/llm/invoke" + assert requests[0].headers["X-Inner-Api-Key"] == "inner-secret" + assert first_payload["caller"]["call_index"] == 1 + assert first_payload["caller"]["invocation_id"] == str(uuid5(NAMESPACE_URL, f"dify-agent:{run_id}:llm:1")) + assert second_payload["caller"]["call_index"] == 2 + assert second_payload["caller"]["invocation_id"] == str(uuid5(NAMESPACE_URL, f"dify-agent:{run_id}:llm:2")) + assert first_payload["caller"]["agent_config_version_kind"] == "draft" + assert first_payload["target"]["provider"] == "acme/custom-model/openai" + assert "credentials" not in first_payload["target"] + assert b"must-not-leave-agent" not in requests[0].content + + +def test_api_client_propagates_gateway_quota_error() -> None: + async def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 429, + json={ + "code": "agent_llm_quota_exceeded", + "message": "Insufficient Message Credits.", + "status": 429, + }, + ) + + async def scenario() -> None: + async with httpx.AsyncClient(transport=httpx.MockTransport(handler), trust_env=False) as http_client: + client = DifyApiLLMClient( + plugin_id="langgenius/openai", + inner_api_url="http://dify-api", + inner_api_key="inner-secret", + execution_context=_execution_context(), + agent_run_id="00000000-0000-0000-0000-000000000001", + http_client=http_client, + ) + + with pytest.raises(ModelHTTPError) as exc_info: + _ = [ + chunk + async for chunk in client.iter_llm_result_chunks( + provider="openai", + model="gpt-test", + credentials={}, + prompt_messages=[UserPromptMessage(content="hello")], + model_parameters={}, + tools=None, + stop=None, + stream=True, + ) + ] + assert exc_info.value.status_code == 429 + assert "Insufficient Message Credits" in str(exc_info.value) + + asyncio.run(scenario()) diff --git a/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py b/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py index 995fea0120c..b7aaf42c21b 100644 --- a/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py +++ b/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py @@ -297,14 +297,14 @@ def test_dify_plugin_llm_layer_builds_adapter_model_from_direct_dependency() -> execution_context = run.get_layer("renamed-execution-context", DifyExecutionContextLayer) llm = run.get_layer("llm", DifyPluginLLMLayer) - model = llm.get_model(http_client=client) + model = llm.get_model(http_client=client, agent_run_id="00000000-0000-0000-0000-000000000001") assert llm.deps.execution_context is execution_context assert isinstance(model, DifyLLMAdapterModel) assert model.model_name == "demo-model" assert model.model_provider == "openai" - assert model.credentials == {"api_key": "secret"} - assert model.provider.name == "DifyPlugin/langgenius/openai" + assert model.credentials == {} + assert model.provider.name == "DifyAPI/langgenius/openai" assert model.provider.client.http_client is client asyncio.run(scenario()) diff --git a/dify-agent/tests/local/dify_agent/runtime/test_runner.py b/dify-agent/tests/local/dify_agent/runtime/test_runner.py index e2fd7a4099c..3002e286b91 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_runner.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_runner.py @@ -485,7 +485,7 @@ class FakeAgentRunResult: def test_runner_emits_terminal_success_and_snapshot(monkeypatch: pytest.MonkeyPatch) -> None: seen_clients: list[httpx.AsyncClient] = [] - def fake_get_model(self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert self.config.model == "demo-model" assert self.config.plugin_id == "langgenius/openai" seen_clients.append(http_client) @@ -558,7 +558,7 @@ def test_runner_emits_complete_plugin_usage_in_terminal_event(monkeypatch: pytes def accumulated_usage(self) -> LLMUsage: return complete_usage - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return PricedTestModel(custom_output_text="done") # pyright: ignore[reportReturnType] @@ -592,7 +592,7 @@ def test_runner_emits_complete_plugin_usage_in_terminal_event(monkeypatch: pytes def test_runner_preserves_explicit_json_null_output(monkeypatch: pytest.MonkeyPatch) -> None: - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -625,7 +625,7 @@ def test_runner_preserves_explicit_json_null_output(monkeypatch: pytest.MonkeyPa def test_runner_passes_explicit_step_limit_to_agent(monkeypatch: pytest.MonkeyPatch) -> None: - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -666,7 +666,7 @@ def test_runner_emits_deferred_tool_call_and_persists_pending_history(monkeypatc tool_call_id="tool-call-1", ) - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -743,7 +743,7 @@ def test_runner_resumes_with_deferred_tool_results_and_no_user_prompt(monkeypatc tool_call_id="tool-call-1", ) - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -848,7 +848,7 @@ def test_runner_can_emit_second_deferred_tool_call_after_resume(monkeypatch: pyt tool_call_id="tool-call-2", ) - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -947,7 +947,7 @@ def test_runner_can_emit_second_deferred_tool_call_after_resume(monkeypatch: pyt def test_runner_rejects_deferred_tool_call_without_history_layer(monkeypatch: pytest.MonkeyPatch) -> None: - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -991,7 +991,7 @@ def test_runner_rejects_resume_with_deferred_tool_results_without_history_layer( ) -> None: agent_run_called = False - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -1038,7 +1038,7 @@ def test_runner_rejects_resume_with_deferred_tool_results_without_history_layer( def test_runner_rejects_multiple_deferred_tool_calls(monkeypatch: pytest.MonkeyPatch) -> None: - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -1076,7 +1076,7 @@ def test_runner_rejects_multiple_deferred_tool_calls(monkeypatch: pytest.MonkeyP def test_runner_rejects_deferred_approval_requests(monkeypatch: pytest.MonkeyPatch) -> None: - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] @@ -1120,7 +1120,7 @@ def test_runner_passes_dynamic_dify_plugin_tools_to_agent(monkeypatch: pytest.Mo async def plugin_tool() -> str: return "tool" - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] @@ -1229,7 +1229,7 @@ def test_runner_passes_dynamic_dify_knowledge_tools_to_agent(monkeypatch: pytest async def knowledge_tool() -> str: return "knowledge" - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] @@ -1335,7 +1335,7 @@ def test_runner_passes_dynamic_dify_core_tools_to_agent(monkeypatch: pytest.Monk async def core_tool() -> str: return "core" - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] @@ -1442,7 +1442,7 @@ def test_runner_rejects_duplicate_tool_names_across_dynamic_tool_layers( async def duplicate_tool() -> str: return "tool" - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] @@ -1566,7 +1566,7 @@ def test_runner_rejects_duplicate_tool_names_between_static_and_dynamic_tools( async def dynamic_duplicate_tool() -> str: return "tool" - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] @@ -1676,7 +1676,7 @@ def test_runner_rejects_duplicate_tool_names_between_shell_and_other_layers( create_agent_called = False shell_client = FakeRunnerShellctlClient() - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] @@ -1801,7 +1801,7 @@ def test_runner_rejects_duplicate_tool_names_between_shell_and_other_layers( def test_runner_passes_temporary_system_prompt_prefix_without_history_layer(monkeypatch: pytest.MonkeyPatch) -> None: model = RecordingTestModel(custom_output_text="done") - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return model # pyright: ignore[reportReturnType] @@ -1843,7 +1843,7 @@ def test_runner_prepends_current_system_prompt_to_stored_history_and_appends_onl ModelResponse(parts=[TextPart(content="old assistant")]), ] - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return model # pyright: ignore[reportReturnType] @@ -1894,7 +1894,7 @@ def test_runner_with_empty_history_layer_still_sends_system_prompt_and_saves_onl ) -> None: model = RecordingTestModel(custom_output_text="done") - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return model # pyright: ignore[reportReturnType] @@ -1944,7 +1944,7 @@ def test_runner_failure_with_history_layer_emits_failed_terminal_event_without_s ModelResponse(parts=[TextPart(content="old assistant")]), ] - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return model # pyright: ignore[reportReturnType] @@ -2004,7 +2004,7 @@ def test_runner_persists_usage_limit_failure_type_in_event_and_status( def test_runner_applies_on_exit_overrides_to_success_snapshot(monkeypatch: pytest.MonkeyPatch) -> None: - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] @@ -2070,7 +2070,7 @@ def test_runner_passes_output_layer_spec_to_agent_and_serializes_structured_resu } ) - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return model # pyright: ignore[reportReturnType] @@ -2165,7 +2165,7 @@ def test_runner_retries_invalid_structured_output_and_eventually_succeeds(monkey ] ) - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return model # pyright: ignore[reportReturnType] @@ -2218,7 +2218,7 @@ def test_runner_fails_when_invalid_structured_output_exhausts_retries(monkeypatc } ) - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): assert http_client.is_closed is False return model # pyright: ignore[reportReturnType] @@ -2263,7 +2263,7 @@ def test_runner_fails_when_invalid_structured_output_exhausts_retries(monkeypatc def test_runner_rejects_invalid_output_layer_before_model_resolution(monkeypatch: pytest.MonkeyPatch) -> None: model_requested = False - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): del http_client nonlocal model_requested model_requested = True @@ -2344,7 +2344,7 @@ def test_runner_rejects_misnamed_output_layer_before_model_resolution(monkeypatc ) sink = InMemoryRunEventSink() - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): del http_client nonlocal model_requested model_requested = True @@ -2431,7 +2431,7 @@ def test_runner_rejects_multiple_output_layers_before_model_resolution(monkeypat ) sink = InMemoryRunEventSink() - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): del http_client nonlocal model_requested model_requested = True @@ -2462,7 +2462,7 @@ def test_runner_rejects_reserved_output_name_with_wrong_layer_type_before_model_ ) -> None: model_requested = False - def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient): + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): del http_client nonlocal model_requested model_requested = True diff --git a/docker/envs/core-services/shared.env.example b/docker/envs/core-services/shared.env.example index 425ff1b8028..4715a137944 100644 --- a/docker/envs/core-services/shared.env.example +++ b/docker/envs/core-services/shared.env.example @@ -512,6 +512,12 @@ MILVUS_ENABLE_HYBRID_SEARCH=False ENABLE_HUMAN_INPUT_TIMEOUT_TASK=true HUMAN_INPUT_TIMEOUT_TASK_INTERVAL=1 +# Agent LLM invocation ledger reconciliation +ENABLE_AGENT_LLM_INVOCATION_RECONCILIATION_TASK=true +AGENT_LLM_INVOCATION_RECONCILIATION_INTERVAL=5 +AGENT_LLM_INVOCATION_STALE_AFTER_SECONDS=900 +AGENT_LLM_INVOCATION_RECONCILIATION_BATCH_SIZE=100 + # Nacos remote settings source HTTP timeouts (seconds). # Bound how long requests to the Nacos endpoint wait before failing, so a slow or # unresponsive Nacos server cannot stall API startup or token refresh.