mirror of
https://github.com/langgenius/dify.git
synced 2026-09-08 11:04:27 +08:00
fix(agent): meter internal LLM calls through API
This commit is contained in:
parent
cb3357a0e4
commit
e3d5ce63d5
@ -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.
|
||||
|
||||
@ -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),
|
||||
),
|
||||
),
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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",
|
||||
|
||||
122
api/controllers/inner_api/agent/llm.py
Normal file
122
api/controllers/inner_api/agent/llm.py
Normal 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"]
|
||||
@ -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(
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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
|
||||
):
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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")
|
||||
@ -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",
|
||||
|
||||
@ -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."""
|
||||
|
||||
|
||||
23
api/schedule/reconcile_agent_llm_invocations_task.py
Normal file
23
api/schedule/reconcile_agent_llm_invocations_task.py
Normal 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"]
|
||||
474
api/services/agent_llm_inner_service.py
Normal file
474
api/services/agent_llm_inner_service.py
Normal 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"]
|
||||
@ -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:
|
||||
|
||||
63
api/services/entities/agent_llm_inner.py
Normal file
63
api/services/entities/agent_llm_inner.py
Normal 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",
|
||||
]
|
||||
@ -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
|
||||
|
||||
|
||||
140
api/tests/unit_tests/controllers/inner_api/test_agent_llm.py
Normal file
140
api/tests/unit_tests/controllers/inner_api/test_agent_llm.py
Normal 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,
|
||||
}
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
325
api/tests/unit_tests/services/test_agent_llm_inner_service.py
Normal file
325
api/tests/unit_tests/services/test_agent_llm_inner_service.py
Normal 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
|
||||
@ -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),
|
||||
|
||||
@ -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"]
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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())
|
||||
@ -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())
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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.
|
||||
|
||||
Loading…
Reference in New Issue
Block a user