services/gr-processor/src/models.py - develop - laritsiuriumov/aef-manager-python-services

services/gr-processor/src/models.py
147 строк4 KB
Цюрюмов Лари Валерьевич [B]
Fix contracts
13 июл 2026, 13:47
13 июл 2026, 13:478b6f2eb
from datetime import datetime
from enum import StrEnum, auto

from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    StrictInt,
    model_validator,
)
from pydantic.alias_generators import to_camel


class EventType(StrEnum):
    EGRESS = "EGRESS"
    GR_PROXY_SERVICE_TRACE = "GR_PROXY_SERVICE_TRACE"


class StatisticType(StrEnum):
    GR_CORE_TOTAL = auto()
    GR_CORE_FAILED = auto()
    GIGACHAT_TOTAL = auto()
    GIGACHAT_GR_CORE = auto()
    GR_CORE_ALERT = auto()
    GR_CORE_INCOMPLETE_ALERT = auto()
    GR_CORE_INCOMPLETE_NO_ALERT = auto()


class GuardCheck(BaseModel):
    model_config = ConfigDict(extra="ignore")

    guard_id: str
    duration_ms: int = Field(alias="duration")
    decision: str
    fail_message: str


class GuardrailResultDetail(BaseModel):
    decision: str
    policy_code: str | None
    description: str | None
    subcode: str | None = None
    duration: StrictInt
    input: str


class GuardrailsTrace(BaseModel):
    timestamp: datetime
    start_time: datetime = Field(alias="start-time")
    duration: StrictInt
    agent_id: str = Field(alias="x-agent-id")
    contract_version: str | None = None
    event_type: str = Field(alias="event-type")
    x_trace_id: str = Field(alias="x-trace-id")
    direction: str
    guardrails_result: dict[str, GuardrailResultDetail] = Field(
        alias="guardrails-result"
    )
    decision: str

    @property
    def x_agent_id(self) -> str:
        return self.agent_id


class SourceLog(BaseModel):
    model_config = ConfigDict(extra="ignore")

    duration: int = Field(..., description="Длительность запроса в милисекундах")
    event_type: EventType = Field(
        alias="event-type",
        examples=["EGRESS", "GR_PROXY_SERVICE_TRACE"],
    )
    trace_id: str = Field(..., alias="x-trace-id")
    agent_id: str = Field(..., alias="x-agent-id")

    distributive: str | None = Field(None, description="Приходит только от gr-proxy")
    egress_timestamp: datetime | None = Field(
        None,
        description="Временная метка, проставляется автоматически только от EGRESS",
        alias="timestamp",
    )
    start_time: datetime | None = Field(
        None,
        description="Старт запроса, проставляется для GR_PROXY_SERVICE_TRACE",
        alias="start-time",
    )
    upstream_cluster: str | None = Field(
        None,
        description="Название сервиса с которого идет запрос. Приходит только от EGRESS",
        examples=["gr-core-service", "some-agent"],
        alias="upstream-cluster",
    )
    path: str | None = Field(
        None, description="Путь запроса, приходит только от EGRESS"
    )
    http_code: int | None = Field(
        None,
        alias="http-code",
        description="Приходит только с EGRESS. Код ответа запроса.",
    )
    decision: str | None = Field(
        None,
        description="Decision проверки. Используется для GR_PROXY_SERVICE_TRACE.",
    )
    guard_checks_data: list[list[GuardCheck]] = Field(
        default_factory=list,
        alias="guard-checks-data",
        description="Данные о проверках, если присутствуют в trace payload",
    )

    @model_validator(mode="after")
    def validate_timestamp_source(self) -> "SourceLog":
        if self.event_type is EventType.EGRESS and self.egress_timestamp is None:
            raise ValueError("EGRESS payload requires timestamp")

        if (
            self.event_type is EventType.GR_PROXY_SERVICE_TRACE
            and self.start_time is None
        ):
            raise ValueError("GR_PROXY_SERVICE_TRACE payload requires start-time")

        return self

    @property
    def timestamp(self) -> datetime:
        if self.event_type is EventType.EGRESS:
            assert self.egress_timestamp is not None
            return self.egress_timestamp

        assert self.start_time is not None
        return self.start_time


class Key(BaseModel):
    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)

    agent_id: str
    stat_type: StatisticType
    start_time: datetime


class Metric(BaseModel):
    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)

    key: Key
    count: int