SupplyAgent/find_agent_v2/models.py

"""Completely isolated persistence models for ``find_agent_v2``."""

from __future__ import annotations

from datetime import datetime
from decimal import Decimal

from sqlalchemy import BigInteger, Index, Integer, Numeric, String, Text, UniqueConstraint, func
from sqlalchemy.dialects.mysql import LONGTEXT
from sqlalchemy.orm import Mapped, mapped_column

from supply_infra.db.base import Base

_LONG_TEXT = Text().with_variant(LONGTEXT(), "mysql")


class FindAgentV2Run(Base):
    __tablename__ = "find_agent_v2_run"
    __table_args__ = (
        UniqueConstraint("run_id", name="uk_find_agent_v2_run_id"),
        Index("idx_find_agent_v2_run_status", "status"),
        Index("idx_find_agent_v2_run_demand", "demand_grade_id"),
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    run_id: Mapped[str] = mapped_column(String(64), nullable=False)
    demand_grade_id: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    demand_word: Mapped[str] = mapped_column(String(256), nullable=False)
    triggered_by_user_id: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    triggered_by_username: Mapped[str | None] = mapped_column(String(64), nullable=True)
    trigger_source: Mapped[str] = mapped_column(String(32), nullable=False, default="cli")
    input_json: Mapped[str] = mapped_column(_LONG_TEXT, nullable=False)
    rule_config_json: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(String(24), nullable=False, default="running")
    outcome_status: Mapped[str | None] = mapped_column(String(24), nullable=True)
    current_round: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    search_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    candidate_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    valid_primary_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    input_tokens: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
    output_tokens: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
    total_tokens: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
    cost_usd: Mapped[Decimal] = mapped_column(Numeric(18, 8), nullable=False, default=0)
    intent_summary: Mapped[str | None] = mapped_column(Text, nullable=True)
    stop_reason: Mapped[str | None] = mapped_column(Text, nullable=True)
    obagent_run_uid: Mapped[str | None] = mapped_column(String(64), nullable=True)
    create_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now(), onupdate=func.now())


class FindAgentV2Round(Base):
    __tablename__ = "find_agent_v2_round"
    __table_args__ = (
        UniqueConstraint("run_id", "round_index", name="uk_find_agent_v2_round"),
        Index("idx_find_agent_v2_round_run", "run_id", "round_index"),
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    run_id: Mapped[str] = mapped_column(String(64), nullable=False)
    round_index: Mapped[int] = mapped_column(Integer, nullable=False)
    phase: Mapped[str] = mapped_column(String(24), nullable=False, default="planning")
    status: Mapped[str] = mapped_column(String(16), nullable=False, default="open")
    plan_json: Mapped[str | None] = mapped_column(_LONG_TEXT, nullable=True)
    start_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    end_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    create_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now(), onupdate=func.now())


class FindAgentV2Search(Base):
    __tablename__ = "find_agent_v2_search"
    __table_args__ = (Index("idx_find_agent_v2_search_run", "run_id", "id"),)

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    run_id: Mapped[str] = mapped_column(String(64), nullable=False)
    round_index: Mapped[int] = mapped_column(Integer, nullable=False)
    keyword: Mapped[str] = mapped_column(String(256), nullable=False)
    query_reason: Mapped[str] = mapped_column(Text, nullable=False)
    source_type: Mapped[str] = mapped_column(String(32), nullable=False)
    provider: Mapped[str] = mapped_column(String(32), nullable=False)
    cursor: Mapped[str] = mapped_column(String(128), nullable=False, default="0")
    page_no: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
    has_more: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    next_cursor: Mapped[str | None] = mapped_column(String(128), nullable=True)
    provider_state_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    result_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    status: Mapped[str] = mapped_column(String(16), nullable=False)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    raw_response_json: Mapped[str | None] = mapped_column(_LONG_TEXT, nullable=True)
    create_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())


class FindAgentV2Candidate(Base):
    __tablename__ = "find_agent_v2_candidate"
    __table_args__ = (
        UniqueConstraint("run_id", "aweme_id", name="uk_find_agent_v2_candidate_aweme"),
        Index("idx_find_agent_v2_candidate_bucket", "run_id", "decision_bucket"),
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    run_id: Mapped[str] = mapped_column(String(64), nullable=False)
    # Logical reference only. find_agent_v2 deliberately uses no database FKs.
    first_search_id: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    aweme_id: Mapped[str] = mapped_column(String(64), nullable=False)
    title: Mapped[str | None] = mapped_column(String(512), nullable=True)
    content_link: Mapped[str | None] = mapped_column(String(1024), nullable=True)
    author_name: Mapped[str | None] = mapped_column(String(256), nullable=True)
    author_sec_uid: Mapped[str | None] = mapped_column(String(256), nullable=True)
    source_keywords_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    tags_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    publish_at: Mapped[datetime | None] = mapped_column(nullable=True)
    duration_seconds: Mapped[Decimal | None] = mapped_column(Numeric(10, 3), nullable=True)
    play_count: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    like_count: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    comment_count: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    collect_count: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    share_count: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    detail_json: Mapped[str | None] = mapped_column(_LONG_TEXT, nullable=True)
    portrait_json: Mapped[str | None] = mapped_column(_LONG_TEXT, nullable=True)
    detail_status: Mapped[str] = mapped_column(String(16), nullable=False, default="pending")
    portrait_status: Mapped[str] = mapped_column(String(16), nullable=False, default="pending")
    content_50_plus_ratio: Mapped[Decimal | None] = mapped_column(Numeric(8, 6), nullable=True)
    account_50_plus_ratio: Mapped[Decimal | None] = mapped_column(Numeric(8, 6), nullable=True)
    relevance_score: Mapped[Decimal | None] = mapped_column(Numeric(8, 6), nullable=True)
    elder_score: Mapped[Decimal | None] = mapped_column(Numeric(8, 6), nullable=True)
    share_score: Mapped[Decimal | None] = mapped_column(Numeric(8, 6), nullable=True)
    value_score: Mapped[Decimal | None] = mapped_column(Numeric(8, 6), nullable=True)
    gate_status: Mapped[str | None] = mapped_column(String(16), nullable=True)
    gate_result_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    decision_bucket: Mapped[str] = mapped_column(String(24), nullable=False, default="pending_evaluation")
    decision_reason: Mapped[str | None] = mapped_column(Text, nullable=True)
    reject_reason_code: Mapped[str | None] = mapped_column(String(64), nullable=True)
    create_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
    update_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now(), onupdate=func.now())


class FindAgentV2Evidence(Base):
    __tablename__ = "find_agent_v2_evidence"
    __table_args__ = (Index("idx_find_agent_v2_evidence_subject", "run_id", "candidate_id", "id"),)

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    run_id: Mapped[str] = mapped_column(String(64), nullable=False)
    candidate_id: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
    evidence_type: Mapped[str] = mapped_column(String(32), nullable=False)
    provider: Mapped[str] = mapped_column(String(32), nullable=False)
    status: Mapped[str] = mapped_column(String(16), nullable=False)
    raw_json: Mapped[str | None] = mapped_column(_LONG_TEXT, nullable=True)
    normalized_json: Mapped[str | None] = mapped_column(_LONG_TEXT, nullable=True)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    create_time: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())