feat: add ai cost tracking foundation

2026-07-31 16:24:35 +08:00
parent ec9f8a015a
commit a84650bb9e
14 changed files with 327 additions and 2 deletions
--- a/ai_knowledge_base_v2/apps/backend/alembic/versions/0018_ai_cost_tracking.py
+++ b/ai_knowledge_base_v2/apps/backend/alembic/versions/0018_ai_cost_tracking.py
@@ -0,0 +1,66 @@
+"""add ai cost tracking fields
+
+Revision ID: 0018_ai_cost_tracking
+Revises: 0017_share_drafts
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+revision = "0018_ai_cost_tracking"
+down_revision = "0017_share_drafts"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    inspector = sa.inspect(op.get_bind())
+    model_columns = {column["name"] for column in inspector.get_columns("sys_model")}
+    _add_column_if_missing(model_columns, "sys_model", sa.Column("input_price_per_1k", sa.Numeric(12, 6), nullable=True))
+    _add_column_if_missing(model_columns, "sys_model", sa.Column("output_price_per_1k", sa.Numeric(12, 6), nullable=True))
+    _add_column_if_missing(model_columns, "sys_model", sa.Column("currency", sa.String(10), nullable=False, server_default="CNY"))
+    _add_column_if_missing(model_columns, "sys_model", sa.Column("usage_scenarios", sa.String(255), nullable=True))
+    _add_column_if_missing(model_columns, "sys_model", sa.Column("allow_summary", sa.Integer(), nullable=False, server_default="1"))
+    _add_column_if_missing(model_columns, "sys_model", sa.Column("allow_report", sa.Integer(), nullable=False, server_default="1"))
+    _add_column_if_missing(model_columns, "sys_model", sa.Column("allow_fixed_info", sa.Integer(), nullable=False, server_default="1"))
+    _add_column_if_missing(model_columns, "sys_model", sa.Column("allow_deep_chat", sa.Integer(), nullable=False, server_default="1"))
+
+    log_columns = {column["name"] for column in inspector.get_columns("sys_ai_request_log")}
+    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("model_id", sa.BigInteger(), nullable=True))
+    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("estimated_cost", sa.Numeric(18, 6), nullable=True))
+    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("currency", sa.String(10), nullable=True))
+    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("route_reason", sa.String(255), nullable=True))
+    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("question_type", sa.String(50), nullable=True))
+    _add_column_if_missing(log_columns, "sys_ai_request_log", sa.Column("knowledge_hit", sa.Integer(), nullable=False, server_default="0"))
+    if "model_id" not in log_columns:
+        op.create_index("ix_sys_ai_request_log_model_id", "sys_ai_request_log", ["model_id"])
+
+
+def downgrade() -> None:
+    inspector = sa.inspect(op.get_bind())
+    log_columns = {column["name"] for column in inspector.get_columns("sys_ai_request_log")}
+    if "model_id" in log_columns:
+        op.drop_index("ix_sys_ai_request_log_model_id", table_name="sys_ai_request_log")
+    for column in ["knowledge_hit", "question_type", "route_reason", "currency", "estimated_cost", "model_id"]:
+        if column in log_columns:
+            op.drop_column("sys_ai_request_log", column)
+
+    model_columns = {column["name"] for column in inspector.get_columns("sys_model")}
+    for column in [
+        "allow_deep_chat",
+        "allow_fixed_info",
+        "allow_report",
+        "allow_summary",
+        "usage_scenarios",
+        "currency",
+        "output_price_per_1k",
+        "input_price_per_1k",
+    ]:
+        if column in model_columns:
+            op.drop_column("sys_model", column)
+
+
+def _add_column_if_missing(columns: set[str], table: str, column: sa.Column) -> None:
+    if column.name not in columns:
+        op.add_column(table, column)