229 lines
8.1 KiB
Python
229 lines
8.1 KiB
Python
from __future__ import annotations
|
|
|
|
import enum
|
|
from datetime import datetime # noqa: TC003
|
|
from decimal import Decimal # noqa: TC003
|
|
from uuid import UUID # noqa: TC003
|
|
|
|
from sqlalchemy import (
|
|
JSON,
|
|
Boolean,
|
|
DateTime,
|
|
Enum,
|
|
ForeignKey,
|
|
Index,
|
|
Integer,
|
|
LargeBinary,
|
|
Numeric,
|
|
String,
|
|
text,
|
|
)
|
|
from sqlalchemy.dialects.postgresql import UUID as PGUUID
|
|
from sqlalchemy.orm import Mapped, mapped_column
|
|
from sqlalchemy.sql import func
|
|
|
|
from db.base import Base
|
|
|
|
|
|
class SessionStatus(str, enum.Enum): # noqa: UP042
|
|
pending = "pending"
|
|
running = "running"
|
|
waiting_approval = "waiting_approval"
|
|
completed = "completed"
|
|
failed = "failed"
|
|
|
|
|
|
class ApprovalStatus(str, enum.Enum): # noqa: UP042
|
|
pending = "pending"
|
|
approved = "approved"
|
|
rejected = "rejected"
|
|
expired = "expired"
|
|
|
|
|
|
class InvokeMode(str, enum.Enum): # noqa: UP042
|
|
realtime = "realtime"
|
|
cached = "cached"
|
|
batch = "batch"
|
|
|
|
|
|
class TraceSession(Base):
|
|
__tablename__ = "trace_session"
|
|
__table_args__ = (Index("ix_trace_session_tenant_created", "tenant_id", "created_at"),)
|
|
|
|
id: Mapped[UUID] = mapped_column(
|
|
PGUUID(as_uuid=True),
|
|
primary_key=True,
|
|
server_default=text("gen_random_uuid()"),
|
|
)
|
|
trace_id: Mapped[str] = mapped_column(String(128), unique=True, nullable=False)
|
|
session_id: Mapped[str] = mapped_column(String(128), unique=True, nullable=False)
|
|
tenant_id: Mapped[str] = mapped_column(String(128), nullable=False, index=True)
|
|
user_id: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
profile_type: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
status: Mapped[SessionStatus] = mapped_column(
|
|
Enum(SessionStatus, name="session_status", native_enum=False),
|
|
nullable=False,
|
|
server_default=text("'pending'"),
|
|
)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=func.now(),
|
|
)
|
|
updated_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=func.now(),
|
|
onupdate=func.now(),
|
|
)
|
|
completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
final_output_ref: Mapped[str | None] = mapped_column(String(512), nullable=True)
|
|
total_input_tokens: Mapped[int] = mapped_column(
|
|
Integer,
|
|
nullable=False,
|
|
server_default=text("0"),
|
|
)
|
|
total_output_tokens: Mapped[int] = mapped_column(
|
|
Integer,
|
|
nullable=False,
|
|
server_default=text("0"),
|
|
)
|
|
total_cost_cny: Mapped[Decimal] = mapped_column(
|
|
Numeric(12, 4),
|
|
nullable=False,
|
|
server_default=text("0"),
|
|
)
|
|
|
|
|
|
class AuditEvent(Base):
|
|
# APPEND-ONLY: enforced by DB trigger trg_audit_event_no_mutation
|
|
__tablename__ = "audit_event"
|
|
|
|
id: Mapped[UUID] = mapped_column(
|
|
PGUUID(as_uuid=True),
|
|
primary_key=True,
|
|
server_default=text("gen_random_uuid()"),
|
|
)
|
|
trace_id: Mapped[str] = mapped_column(String(128), nullable=False, index=True)
|
|
step_id: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
event_type: Mapped[str] = mapped_column(String(64), nullable=False)
|
|
actor: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
tool_name: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
|
status: Mapped[str] = mapped_column(String(64), nullable=False)
|
|
duration_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
error_message: Mapped[str | None] = mapped_column(String(1000), nullable=True)
|
|
evidence_score: Mapped[Decimal | None] = mapped_column(Numeric(5, 4), nullable=True)
|
|
payload_ref: Mapped[str | None] = mapped_column(String(512), nullable=True)
|
|
mock_flag: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("false"))
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=func.now(),
|
|
)
|
|
tenant_id: Mapped[str] = mapped_column(String(128), nullable=False, index=True)
|
|
|
|
|
|
class AuditPayload(Base):
|
|
__tablename__ = "audit_payload"
|
|
|
|
id: Mapped[UUID] = mapped_column(
|
|
PGUUID(as_uuid=True),
|
|
primary_key=True,
|
|
server_default=text("gen_random_uuid()"),
|
|
)
|
|
audit_event_id: Mapped[UUID] = mapped_column(
|
|
PGUUID(as_uuid=True),
|
|
ForeignKey("audit_event.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
)
|
|
payload_type: Mapped[str] = mapped_column(String(64), nullable=False)
|
|
encrypted_data: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
|
|
encryption_key_id: Mapped[str] = mapped_column(String(64), nullable=False)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=func.now(),
|
|
)
|
|
|
|
|
|
class PendingApproval(Base):
|
|
__tablename__ = "pending_approval"
|
|
__table_args__ = (Index("ix_pending_approval_status_expired", "status", "expired_at"),)
|
|
|
|
id: Mapped[UUID] = mapped_column(
|
|
PGUUID(as_uuid=True),
|
|
primary_key=True,
|
|
server_default=text("gen_random_uuid()"),
|
|
)
|
|
trace_id: Mapped[str] = mapped_column(String(128), unique=True, nullable=False)
|
|
session_id: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
step_id: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
tool_name: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
action_description: Mapped[str] = mapped_column(String(1000), nullable=False)
|
|
risk_level: Mapped[str] = mapped_column(String(64), nullable=False)
|
|
requester_id: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
approver_role: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
params_preview: Mapped[dict[str, object]] = mapped_column(JSON, nullable=False)
|
|
status: Mapped[ApprovalStatus] = mapped_column(
|
|
Enum(ApprovalStatus, name="approval_status", native_enum=False),
|
|
nullable=False,
|
|
server_default=text("'pending'"),
|
|
)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=func.now(),
|
|
)
|
|
expired_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=text("(now() + interval '24 hours')"),
|
|
)
|
|
decided_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
decider_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
|
decision_note: Mapped[str | None] = mapped_column(String(1000), nullable=True)
|
|
|
|
|
|
class CostRecord(Base):
|
|
__tablename__ = "cost_record"
|
|
|
|
id: Mapped[UUID] = mapped_column(
|
|
PGUUID(as_uuid=True),
|
|
primary_key=True,
|
|
server_default=text("gen_random_uuid()"),
|
|
)
|
|
trace_id: Mapped[str] = mapped_column(String(128), nullable=False, index=True)
|
|
tenant_id: Mapped[str] = mapped_column(String(128), nullable=False, index=True)
|
|
model: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
input_tokens: Mapped[int] = mapped_column(Integer, nullable=False)
|
|
output_tokens: Mapped[int] = mapped_column(Integer, nullable=False)
|
|
cache_read_tokens: Mapped[int] = mapped_column(
|
|
Integer,
|
|
nullable=False,
|
|
server_default=text("0"),
|
|
)
|
|
cache_write_tokens: Mapped[int] = mapped_column(
|
|
Integer,
|
|
nullable=False,
|
|
server_default=text("0"),
|
|
)
|
|
cost_usd: Mapped[Decimal] = mapped_column(Numeric(12, 6), nullable=False)
|
|
cost_cny: Mapped[Decimal] = mapped_column(Numeric(12, 4), nullable=False)
|
|
invoke_mode: Mapped[InvokeMode] = mapped_column(
|
|
Enum(InvokeMode, name="invoke_mode", native_enum=False),
|
|
nullable=False,
|
|
server_default=text("'realtime'"),
|
|
)
|
|
recorded_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=func.now(),
|
|
)
|
|
|
|
|
|
TraceSession.table = TraceSession.__table__
|
|
AuditEvent.table = AuditEvent.__table__
|
|
AuditPayload.table = AuditPayload.__table__
|
|
PendingApproval.table = PendingApproval.__table__
|
|
CostRecord.table = CostRecord.__table__
|