from common.extensions import db
from sqlalchemy.dialects.postgresql import JSONB
import sqlalchemy as sa


class LLMUsageMetric(db.Model):
    __bind_key__ = 'public'
    __table_args__ = {'schema': 'public'}

    id = db.Column(db.Integer, primary_key=True)
    tenant_id = db.Column(db.Integer, nullable=False)
    environment = db.Column(db.String(20), nullable=False)
    activity = db.Column(db.String(20), nullable=False)
    sub_activity = db.Column(db.String(20), nullable=False)
    activity_detail = db.Column(db.String(50), nullable=True)
    session_id = db.Column(db.String(50), nullable=True)        # Chat Session ID
    interaction_id = db.Column(db.Integer, nullable=True)       # Chat Interaction ID
    document_version_id = db.Column(db.Integer, nullable=True)
    prompt_tokens = db.Column(db.Integer, nullable=True)
    completion_tokens = db.Column(db.Integer, nullable=True)
    total_tokens = db.Column(db.Integer, nullable=True)
    cost = db.Column(db.Float, nullable=True)
    latency = db.Column(db.Float, nullable=True)
    model_name = db.Column(db.String(50), nullable=False)
    timestamp = db.Column(db.DateTime, nullable=False)
    additional_info = db.Column(JSONB, nullable=True)

    # Add any additional fields or methods as needed