mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-08-02 21:22:51 +02:00
108 lines
3.4 KiB
Python
108 lines
3.4 KiB
Python
"""agent_tracking
|
|
|
|
Revision ID: 98a5008d8711
|
|
Revises: 2f80c6a2550f
|
|
Create Date: 2025-01-29 17:00:00.000001
|
|
|
|
"""
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
from sqlalchemy.dialects.postgresql import UUID
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision = "98a5008d8711"
|
|
down_revision = "2f80c6a2550f"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
op.create_table(
|
|
"agent__search_metrics",
|
|
sa.Column("id", sa.Integer(), nullable=False),
|
|
sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=True),
|
|
sa.Column("persona_id", sa.Integer(), nullable=True),
|
|
sa.Column("agent_type", sa.String(), nullable=False),
|
|
sa.Column("start_time", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("base_duration_s", sa.Float(), nullable=False),
|
|
sa.Column("full_duration_s", sa.Float(), nullable=False),
|
|
sa.Column("base_metrics", postgresql.JSONB(), nullable=True),
|
|
sa.Column("refined_metrics", postgresql.JSONB(), nullable=True),
|
|
sa.Column("all_metrics", postgresql.JSONB(), nullable=True),
|
|
sa.ForeignKeyConstraint(
|
|
["persona_id"],
|
|
["persona.id"],
|
|
),
|
|
sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
|
|
sa.PrimaryKeyConstraint("id"),
|
|
)
|
|
|
|
# Create sub_question table
|
|
op.create_table(
|
|
"agent__sub_question",
|
|
sa.Column("id", sa.Integer, primary_key=True),
|
|
sa.Column("primary_question_id", sa.Integer, sa.ForeignKey("chat_message.id")),
|
|
sa.Column(
|
|
"chat_session_id", UUID(as_uuid=True), sa.ForeignKey("chat_session.id")
|
|
),
|
|
sa.Column("sub_question", sa.Text),
|
|
sa.Column(
|
|
"time_created", sa.DateTime(timezone=True), server_default=sa.func.now()
|
|
),
|
|
sa.Column("sub_answer", sa.Text),
|
|
sa.Column("sub_question_doc_results", postgresql.JSONB(), nullable=True),
|
|
sa.Column("level", sa.Integer(), nullable=False),
|
|
sa.Column("level_question_num", sa.Integer(), nullable=False),
|
|
)
|
|
|
|
# Create sub_query table
|
|
op.create_table(
|
|
"agent__sub_query",
|
|
sa.Column("id", sa.Integer, primary_key=True),
|
|
sa.Column(
|
|
"parent_question_id", sa.Integer, sa.ForeignKey("agent__sub_question.id")
|
|
),
|
|
sa.Column(
|
|
"chat_session_id", UUID(as_uuid=True), sa.ForeignKey("chat_session.id")
|
|
),
|
|
sa.Column("sub_query", sa.Text),
|
|
sa.Column(
|
|
"time_created", sa.DateTime(timezone=True), server_default=sa.func.now()
|
|
),
|
|
)
|
|
|
|
# Create sub_query__search_doc association table
|
|
op.create_table(
|
|
"agent__sub_query__search_doc",
|
|
sa.Column(
|
|
"sub_query_id",
|
|
sa.Integer,
|
|
sa.ForeignKey("agent__sub_query.id"),
|
|
primary_key=True,
|
|
),
|
|
sa.Column(
|
|
"search_doc_id",
|
|
sa.Integer,
|
|
sa.ForeignKey("search_doc.id"),
|
|
primary_key=True,
|
|
),
|
|
)
|
|
|
|
op.add_column(
|
|
"chat_message",
|
|
sa.Column(
|
|
"refined_answer_improvement",
|
|
sa.Boolean(),
|
|
nullable=True,
|
|
),
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.drop_column("chat_message", "refined_answer_improvement")
|
|
op.drop_table("agent__sub_query__search_doc")
|
|
op.drop_table("agent__sub_query")
|
|
op.drop_table("agent__sub_question")
|
|
op.drop_table("agent__search_metrics")
|