api/core/telemetry/events.py

Rama main · réplica de lectura del repositorio original.

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Any, NotRequired, Protocol, TypedDict, runtime_checkable

from dify_trace_sdk.events import DraftNodeExecuted, PromptGenerated, TracePayload, json_safe

from enterprise.telemetry.contracts import SignalType, TelemetryCase

# ---------------------------------------------------------------------------
# Context shapes
# ---------------------------------------------------------------------------


class TraceContext(TypedDict):
    """Common ``context`` shape for TRACE-routed cases."""

    tenant_id: str | None
    user_id: NotRequired[str | None]
    app_id: NotRequired[str | None]


class MetricLogContext(TypedDict):
    """Common ``context`` shape for METRIC_LOG-routed cases."""

    tenant_id: str | None


# ---------------------------------------------------------------------------
# Payload shapes
# ---------------------------------------------------------------------------


class NodeExecutionPayload(TypedDict):
    """``payload`` shape for ``TelemetryCase.NODE_EXECUTION`` and ``DRAFT_NODE_EXECUTION``."""

    node_execution_data: dict[str, Any]


class AppCreatedPayload(TypedDict):
    """``payload`` shape for ``TelemetryCase.APP_CREATED``."""

    app_id: str | None
    mode: NotRequired[str | None]


class AppUpdatedPayload(TypedDict):
    """``payload`` shape for ``TelemetryCase.APP_UPDATED``."""

    app_id: str | None


class AppDeletedPayload(TypedDict):
    """``payload`` shape for ``TelemetryCase.APP_DELETED``."""

    app_id: str | None


class PromptGenerationPayload(TypedDict):
    """``payload`` shape for ``TelemetryCase.PROMPT_GENERATION``."""

    tenant_id: str
    app_id: NotRequired[str | None]
    operation_type: str
    instruction: str
    generated_output: str
    model_provider: str
    model_name: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    latency: float
    total_price: float | None
    currency: str | None
    error: str | None
    started_at: datetime
    ended_at: datetime


class FeedbackCreatedPayload(TypedDict):
    """``payload`` shape for ``TelemetryCase.FEEDBACK_CREATED``."""

    message_id: str
    app_id: str | None
    conversation_id: NotRequired[str | None]
    from_end_user_id: NotRequired[str | None]
    from_account_id: NotRequired[str | None]
    rating: NotRequired[str | None]
    from_source: NotRequired[str | None]
    content: NotRequired[str | None]


# ---------------------------------------------------------------------------
# Shared context dataclass
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class TelemetryContext:
    tenant_id: str | None = None
    user_id: str | None = None
    app_id: str | None = None


# ---------------------------------------------------------------------------
# Event protocol — every concrete event must declare these routing fields
# so the gateway can dispatch without external mapping tables.
# ---------------------------------------------------------------------------


@runtime_checkable
class TelemetryEvent(Protocol):
    """Structural contract for all telemetry events.

    All attributes are declared as read-only ``@property`` so that
    frozen dataclasses satisfy the protocol (a bare attribute
    annotation would require writability).
    """

    @property
    def context(self) -> TelemetryContext: ...
    @property
    def payload(self) -> Mapping[str, Any]: ...
    @property
    def case(self) -> TelemetryCase: ...
    @property
    def signal_type(self) -> SignalType: ...
    @property
    def ce_eligible(self) -> bool: ...


@runtime_checkable
class TraceSignalEvent(TelemetryEvent, Protocol):
    """A ``SignalType.TRACE`` event, delivered through the trace pipeline."""

    def trace_payload(self) -> TracePayload: ...


# ---------------------------------------------------------------------------
# Concrete event classes
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class DraftNodeExecutionTraceEvent:
    context: TelemetryContext
    payload: NodeExecutionPayload
    case: TelemetryCase = TelemetryCase.DRAFT_NODE_EXECUTION
    signal_type: SignalType = SignalType.TRACE
    ce_eligible: bool = False

    def trace_payload(self) -> TracePayload:
        node_execution = self.payload["node_execution_data"]
        return DraftNodeExecuted(node_execution={key: json_safe(value) for key, value in node_execution.items()})


@dataclass(frozen=True)
class PromptGenerationEvent:
    context: TelemetryContext
    payload: PromptGenerationPayload
    case: TelemetryCase = TelemetryCase.PROMPT_GENERATION
    signal_type: SignalType = SignalType.TRACE
    ce_eligible: bool = False

    def trace_payload(self) -> TracePayload:
        payload = self.payload
        return PromptGenerated(
            operation_type=payload["operation_type"],
            instruction=payload["instruction"],
            generated_output=payload["generated_output"],
            model_provider=payload["model_provider"],
            model_name=payload["model_name"],
            prompt_tokens=payload["prompt_tokens"],
            completion_tokens=payload["completion_tokens"],
            total_tokens=payload["total_tokens"],
            latency=payload["latency"],
            total_price=payload["total_price"],
            currency=payload["currency"],
            error=payload["error"],
            started_at=payload["started_at"],
            ended_at=payload["ended_at"],
        )


@dataclass(frozen=True)
class AppCreatedEvent:
    context: TelemetryContext
    payload: AppCreatedPayload
    case: TelemetryCase = TelemetryCase.APP_CREATED
    signal_type: SignalType = SignalType.METRIC_LOG
    ce_eligible: bool = False


@dataclass(frozen=True)
class AppUpdatedEvent:
    context: TelemetryContext
    payload: AppUpdatedPayload
    case: TelemetryCase = TelemetryCase.APP_UPDATED
    signal_type: SignalType = SignalType.METRIC_LOG
    ce_eligible: bool = False


@dataclass(frozen=True)
class AppDeletedEvent:
    context: TelemetryContext
    payload: AppDeletedPayload
    case: TelemetryCase = TelemetryCase.APP_DELETED
    signal_type: SignalType = SignalType.METRIC_LOG
    ce_eligible: bool = False


@dataclass(frozen=True)
class FeedbackCreatedEvent:
    context: TelemetryContext
    payload: FeedbackCreatedPayload
    case: TelemetryCase = TelemetryCase.FEEDBACK_CREATED
    signal_type: SignalType = SignalType.METRIC_LOG
    ce_eligible: bool = False