fix(agent): meter internal LLM calls through API

This commit is contained in:
Yansong Zhang 2026-08-12 17:57:44 +08:00
parent cb3357a0e4
commit e3d5ce63d5
32 changed files with 1855 additions and 62 deletions

View File

@ -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.

View File

@ -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),
),
),

View File

@ -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",

View File

@ -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",

View File

@ -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"]

View File

@ -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(

View File

@ -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):

View File

@ -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
):

View File

@ -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

View File

@ -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")

View File

@ -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",

View File

@ -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."""

View File

@ -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"]

View File

@ -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"]

View File

@ -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:

View File

@ -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",
]

View File

@ -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

View File

@ -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,
}

View File

@ -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:

View File

@ -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

View File

@ -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

View File

@ -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),

View File

@ -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"]

View File

@ -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 = "<think>\n"
_THINK_END = "\n</think>"
@ -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

View File

@ -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,

View File

@ -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,
)

View File

@ -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(

View File

@ -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,

View File

@ -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())

View File

@ -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())

View File

@ -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

View File

@ -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.