o2o-negosium-original/agent/common/database/model/models.py

from sqlalchemy.orm import declarative_base
from sqlalchemy import (
    BigInteger,
    Boolean,
    Column,
    DateTime,
    Integer,
    Numeric,
    SmallInteger,
    String,
    Float,
)
from sqlalchemy.dialects.postgresql import JSONB, UUID
from sqlalchemy.sql import text

from common.enums import DBType

# 모든 ORM 모델의 베이스. insert 시 isinstance 체크에도 사용된다. (backend 와 동일 패턴)
#
# agent 는 negosium_db 를 backend 와 공유한다.
# - 기존 스키마(company/card/negotiation 등)는 postgres-init/01-schema*.sql 이 소유 → 읽기/쓰기만.
# - RL 학습 자산은 신설 `learning` 스키마(postgres-init/02-learning-schema.sql)에 둔다.
MAIN_BASE = declarative_base()

LEARNING_SCHEMA = "learning"

# 공유 베이스 정책의 예약 테넌트 키 (company_id 컬럼에 저장). company uuid 와 혼용되므로 VARCHAR.
BASE_COMPANY_ID = "_base"


class _DBTypeMixin:
    """모델이 자신이 속한 논리 DB 를 알려준다 (람다 실행 시 DBType 으로 세션 선택)."""

    @staticmethod
    def DBType():
        return DBType.MAIN.value


# ============================================================
# learning 스키마 ORM (02-learning-schema.sql 과 1:1)
# 모든 테이블에 company_id(테넌트 키) — 논리 격리.
# ============================================================
class QTableVersion(_DBTypeMixin, MAIN_BASE):
    __tablename__ = "q_table_versions"
    __table_args__ = {"schema": LEARNING_SCHEMA}

    version_id = Column(UUID(as_uuid=True), primary_key=True, server_default=text("gen_random_uuid()"))
    company_id = Column(String(64), nullable=False)
    version_name = Column(String(50), nullable=False)
    scope = Column(SmallInteger, nullable=False, default=2)  # 1=base, 2=tenant
    base_version_id = Column(UUID(as_uuid=True), nullable=True)
    state_space_size = Column(Integer, nullable=False)
    action_space_size = Column(Integer, nullable=False)
    learning_rate = Column(Numeric(6, 4), nullable=False, default=0.1)
    discount_factor = Column(Numeric(6, 4), nullable=False, default=0.95)
    epochs = Column(Integer, nullable=False, default=0)
    is_active = Column(Boolean, nullable=False, default=False)
    created_at = Column(DateTime(timezone=True), server_default=text("now()"))
    deleted = Column(Boolean, nullable=False, default=False)


class QValue(_DBTypeMixin, MAIN_BASE):
    __tablename__ = "q_values"
    __table_args__ = {"schema": LEARNING_SCHEMA}

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    company_id = Column(String(64), nullable=False)
    version_id = Column(UUID(as_uuid=True), nullable=False)
    state_index = Column(Integer, nullable=False)
    action_id = Column(Integer, nullable=False)
    q_value = Column(Float, nullable=False, default=0.0)


class VisitCount(_DBTypeMixin, MAIN_BASE):
    __tablename__ = "visit_counts"
    __table_args__ = {"schema": LEARNING_SCHEMA}

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    company_id = Column(String(64), nullable=False)
    version_id = Column(UUID(as_uuid=True), nullable=False)
    state_index = Column(Integer, nullable=False)
    action_id = Column(Integer, nullable=False)
    count = Column(BigInteger, nullable=False, default=0)


class ExperienceLog(_DBTypeMixin, MAIN_BASE):
    __tablename__ = "experience_logs"
    __table_args__ = {"schema": LEARNING_SCHEMA}

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    company_id = Column(String(64), nullable=False)
    transition_id = Column(UUID(as_uuid=True), nullable=False, server_default=text("gen_random_uuid()"))
    session_id = Column(UUID(as_uuid=True), nullable=True)
    state_index = Column(Integer, nullable=False)
    action_id = Column(Integer, nullable=False)
    card_id = Column(String(40), nullable=True)
    q_value_at_selection = Column(Float, nullable=True)
    reward = Column(Float, nullable=True)
    next_state_index = Column(Integer, nullable=True)
    done = Column(Boolean, nullable=False, default=False)
    snapshot = Column(JSONB, nullable=True)
    propensity = Column(Float, nullable=True)            # OPE 필수
    turn = Column(Integer, nullable=True)
    available_actions = Column(JSONB, nullable=True)
    settled_price = Column(BigInteger, nullable=True)
    visit_count_at_selection = Column(BigInteger, nullable=True)
    total_visits_at_selection = Column(BigInteger, nullable=True)
    ucb_score_at_selection = Column(Float, nullable=True)
    is_new_quote = Column(Boolean, nullable=False, default=False)
    is_invalidated = Column(Boolean, nullable=False, default=False)
    invalidated_reason = Column(String(255), nullable=True)
    created_at = Column(DateTime(timezone=True), server_default=text("now()"))


class TenantActionCard(_DBTypeMixin, MAIN_BASE):
    __tablename__ = "tenant_action_cards"
    __table_args__ = {"schema": LEARNING_SCHEMA}

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    company_id = Column(String(64), nullable=False)
    action_id = Column(Integer, nullable=False)
    card_id = Column(String(40), nullable=False)
    created_at = Column(DateTime(timezone=True), server_default=text("now()"))
    updated_at = Column(DateTime(timezone=True), server_default=text("now()"))
    deleted = Column(Boolean, nullable=False, default=False)


class ChatSessionRow(_DBTypeMixin, MAIN_BASE):
    """/chat 진행 상태 영속화 (P8-A). 채팅 메시지 로그가 아니라 step 머신 상태."""

    __tablename__ = "chat_sessions"
    __table_args__ = {"schema": LEARNING_SCHEMA}

    session_id = Column(UUID(as_uuid=True), primary_key=True)
    company_id = Column(String(64), nullable=False)
    tenant_id = Column(String(64), nullable=False)
    rq_type = Column(String(10), nullable=False, default="재협상")
    step = Column(String(40), nullable=False, default="시작")
    context = Column(JSONB, nullable=False, default=dict)
    used_action_ids = Column(JSONB, nullable=False, default=list)
    action_space_size = Column(Integer, nullable=False, default=0)
    ended = Column(Boolean, nullable=False, default=False)
    created_at = Column(DateTime(timezone=True), server_default=text("now()"))
    updated_at = Column(DateTime(timezone=True), server_default=text("now()"))