refactor: normalize script_builder_messages into separate table

Extract JSONB messages array from script_builder_sessions into a proper
script_builder_messages table with individual columns for role, content,
script, tokens, etc. Migration handles data migration from JSONB to rows.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Michael Chihlas
2026-03-21 21:06:58 -04:00
parent 0b261ee21a
commit b801f6cac5
9 changed files with 290 additions and 61 deletions

View File

@@ -24,7 +24,7 @@ from app.models.ai_session import AISession # noqa: F401
from app.models.ai_session_step import AISessionStep # noqa: F401 from app.models.ai_session_step import AISessionStep # noqa: F401
from app.models.psa_post_log import PsaPostLog # noqa: F401 from app.models.psa_post_log import PsaPostLog # noqa: F401
from app.models.psa_member_mapping import PsaMemberMapping # noqa: F401 from app.models.psa_member_mapping import PsaMemberMapping # noqa: F401
from app.models.script_builder_session import ScriptBuilderSession # noqa: F401 from app.models.script_builder_session import ScriptBuilderSession, ScriptBuilderMessage # noqa: F401
from app.core.config import settings from app.core.config import settings
# this is the Alembic Config object # this is the Alembic Config object

View File

@@ -0,0 +1,128 @@
"""normalize script builder messages
Revision ID: 064
Revises: 063
Create Date: 2026-03-21 22:00:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID, JSONB
# revision identifiers, used by Alembic.
revision: str = '064'
down_revision: Union[str, None] = '063'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# 1. Create the new script_builder_messages table
op.create_table(
'script_builder_messages',
sa.Column('id', UUID(as_uuid=True), primary_key=True),
sa.Column('session_id', UUID(as_uuid=True), sa.ForeignKey('script_builder_sessions.id', ondelete='CASCADE'), nullable=False, index=True),
sa.Column('role', sa.String(20), nullable=False, comment='user or assistant'),
sa.Column('content', sa.Text(), nullable=False),
sa.Column('script', sa.Text(), nullable=True, comment='Extracted script from AI response'),
sa.Column('script_filename', sa.String(200), nullable=True),
sa.Column('line_count', sa.Integer(), nullable=True),
sa.Column('input_tokens', sa.Integer(), nullable=True),
sa.Column('output_tokens', sa.Integer(), nullable=True),
sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
)
# 2. Migrate existing JSONB messages to the new table (if any exist)
conn = op.get_bind()
sessions = conn.execute(
sa.text("SELECT id, messages FROM script_builder_sessions WHERE messages IS NOT NULL AND messages != '[]'::jsonb")
).fetchall()
for session_row in sessions:
session_id = session_row[0]
messages = session_row[1]
if not messages:
continue
for msg in messages:
conn.execute(
sa.text("""
INSERT INTO script_builder_messages (id, session_id, role, content, script, script_filename, line_count, input_tokens, output_tokens, created_at)
VALUES (gen_random_uuid(), :session_id, :role, :content, :script, :script_filename, :line_count, :input_tokens, :output_tokens, COALESCE(:timestamp::timestamptz, NOW()))
"""),
{
"session_id": session_id,
"role": msg.get("role", "user"),
"content": msg.get("content", ""),
"script": msg.get("script"),
"script_filename": msg.get("script_filename"),
"line_count": msg.get("line_count"),
"input_tokens": msg.get("input_tokens"),
"output_tokens": msg.get("output_tokens"),
"timestamp": msg.get("timestamp"),
},
)
# 3. Drop the old columns
op.drop_column('script_builder_sessions', 'messages')
op.drop_column('script_builder_sessions', 'message_count')
def downgrade() -> None:
# 1. Re-add the JSONB columns
op.add_column('script_builder_sessions', sa.Column('messages', JSONB, nullable=False, server_default='[]'))
op.add_column('script_builder_sessions', sa.Column('message_count', sa.Integer(), nullable=False, server_default='0'))
# 2. Migrate data back from the messages table to JSONB
conn = op.get_bind()
sessions = conn.execute(
sa.text("SELECT DISTINCT session_id FROM script_builder_messages")
).fetchall()
for (session_id,) in sessions:
messages = conn.execute(
sa.text("""
SELECT role, content, script, script_filename, line_count, input_tokens, output_tokens, created_at
FROM script_builder_messages
WHERE session_id = :session_id
ORDER BY created_at ASC
"""),
{"session_id": session_id},
).fetchall()
json_messages = []
user_count = 0
for msg in messages:
entry = {
"role": msg[0],
"content": msg[1],
"timestamp": msg[7].isoformat() if msg[7] else None,
}
if msg[2]: # script
entry["script"] = msg[2]
if msg[3]: # script_filename
entry["script_filename"] = msg[3]
if msg[4] is not None: # line_count
entry["line_count"] = msg[4]
if msg[5] is not None: # input_tokens
entry["input_tokens"] = msg[5]
if msg[6] is not None: # output_tokens
entry["output_tokens"] = msg[6]
json_messages.append(entry)
if msg[0] == "user":
user_count += 1
import json
conn.execute(
sa.text("UPDATE script_builder_sessions SET messages = :messages::jsonb, message_count = :count WHERE id = :id"),
{"messages": json.dumps(json_messages), "count": user_count, "id": session_id},
)
# 3. Drop the messages table
op.drop_table('script_builder_messages')
# 4. Remove server defaults added for downgrade
op.alter_column('script_builder_sessions', 'messages', server_default=None)
op.alter_column('script_builder_sessions', 'message_count', server_default=None)

View File

@@ -9,10 +9,12 @@ from app.core.database import get_db
from app.core.rate_limit import limiter from app.core.rate_limit import limiter
from app.api.deps import get_current_active_user from app.api.deps import get_current_active_user
from app.models.user import User from app.models.user import User
from app.models.script_builder_session import ScriptBuilderSession
from app.schemas.script_builder import ( from app.schemas.script_builder import (
ScriptBuilderCreateRequest, ScriptBuilderCreateRequest,
ScriptBuilderMessageRequest, ScriptBuilderMessageRequest,
ScriptBuilderMessageResponse, ScriptBuilderMessageResponse,
ScriptBuilderMessageSchema,
ScriptBuilderSessionDetail, ScriptBuilderSessionDetail,
ScriptBuilderSessionSummary, ScriptBuilderSessionSummary,
SaveToLibraryRequest, SaveToLibraryRequest,
@@ -25,6 +27,39 @@ router = APIRouter(prefix="/scripts/builder", tags=["script-builder"])
MAX_SESSIONS_PER_USER = 5 MAX_SESSIONS_PER_USER = 5
def _session_to_detail(session: ScriptBuilderSession) -> ScriptBuilderSessionDetail:
"""Convert a session ORM object (with message_records loaded) to detail schema."""
messages = [
ScriptBuilderMessageSchema.model_validate(m)
for m in session.message_records
]
return ScriptBuilderSessionDetail(
id=session.id,
language=session.language,
title=session.title,
messages=messages,
latest_script=session.latest_script,
latest_script_filename=session.latest_script_filename,
message_count=len([m for m in messages if m.role == "user"]),
ai_session_id=session.ai_session_id,
created_at=session.created_at,
updated_at=session.updated_at,
)
def _session_to_summary(session: ScriptBuilderSession) -> ScriptBuilderSessionSummary:
"""Convert a session ORM object to summary schema (no messages needed)."""
return ScriptBuilderSessionSummary(
id=session.id,
language=session.language,
title=session.title,
message_count=0, # Summary doesn't eagerly load messages
latest_script_filename=session.latest_script_filename,
created_at=session.created_at,
updated_at=session.updated_at,
)
@router.post("/sessions", response_model=ScriptBuilderSessionDetail, status_code=201) @router.post("/sessions", response_model=ScriptBuilderSessionDetail, status_code=201)
async def create_session( async def create_session(
data: ScriptBuilderCreateRequest, data: ScriptBuilderCreateRequest,
@@ -47,7 +82,9 @@ async def create_session(
language=data.language, language=data.language,
) )
await db.commit() await db.commit()
return ScriptBuilderSessionDetail.model_validate(session) # Re-fetch with message_records loaded
session = await script_builder_service.get_session(db, session.id, current_user.id)
return _session_to_detail(session)
@router.get("/sessions", response_model=list[ScriptBuilderSessionSummary]) @router.get("/sessions", response_model=list[ScriptBuilderSessionSummary])
@@ -61,7 +98,7 @@ async def list_sessions(
sessions = await script_builder_service.list_sessions( sessions = await script_builder_service.list_sessions(
db=db, user_id=current_user.id, limit=limit, offset=offset db=db, user_id=current_user.id, limit=limit, offset=offset
) )
return [ScriptBuilderSessionSummary.model_validate(s) for s in sessions] return [_session_to_summary(s) for s in sessions]
@router.get("/sessions/{session_id}", response_model=ScriptBuilderSessionDetail) @router.get("/sessions/{session_id}", response_model=ScriptBuilderSessionDetail)
@@ -74,7 +111,7 @@ async def get_session(
session = await script_builder_service.get_session(db, session_id, current_user.id) session = await script_builder_service.get_session(db, session_id, current_user.id)
if not session: if not session:
raise HTTPException(status_code=404, detail="Session not found") raise HTTPException(status_code=404, detail="Session not found")
return ScriptBuilderSessionDetail.model_validate(session) return _session_to_detail(session)
@router.post( @router.post(

View File

@@ -1,14 +1,14 @@
"""Script Builder session model. """Script Builder session and message models.
Tracks AI-powered script generation conversations. Tracks AI-powered script generation conversations.
""" """
import uuid import uuid
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import Optional, Any, TYPE_CHECKING from typing import Optional, TYPE_CHECKING
from sqlalchemy import String, Text, DateTime, ForeignKey, Integer from sqlalchemy import String, Text, DateTime, ForeignKey, Integer
from sqlalchemy.orm import Mapped, mapped_column, relationship from sqlalchemy.orm import Mapped, mapped_column, relationship
from sqlalchemy.dialects.postgresql import UUID, JSONB from sqlalchemy.dialects.postgresql import UUID
from app.core.database import Base from app.core.database import Base
@@ -42,10 +42,6 @@ class ScriptBuilderSession(Base):
String(200), nullable=True, String(200), nullable=True,
comment="Auto-generated from first AI response", comment="Auto-generated from first AI response",
) )
messages: Mapped[list[dict[str, Any]]] = mapped_column(
JSONB, nullable=False, default=list,
comment="Array of {role, content, script?, script_filename?, timestamp}",
)
latest_script: Mapped[Optional[str]] = mapped_column( latest_script: Mapped[Optional[str]] = mapped_column(
Text, nullable=True, Text, nullable=True,
comment="Most recent generated script for quick access", comment="Most recent generated script for quick access",
@@ -54,9 +50,6 @@ class ScriptBuilderSession(Base):
String(200), nullable=True, String(200), nullable=True,
comment="Filename of the latest generated script", comment="Filename of the latest generated script",
) )
message_count: Mapped[int] = mapped_column(
Integer, nullable=False, default=0,
)
ai_session_id: Mapped[Optional[uuid.UUID]] = mapped_column( ai_session_id: Mapped[Optional[uuid.UUID]] = mapped_column(
UUID(as_uuid=True), UUID(as_uuid=True),
ForeignKey("ai_sessions.id", ondelete="SET NULL"), ForeignKey("ai_sessions.id", ondelete="SET NULL"),
@@ -74,3 +67,45 @@ class ScriptBuilderSession(Base):
# Relationships # Relationships
user: Mapped["User"] = relationship("User") user: Mapped["User"] = relationship("User")
message_records: Mapped[list["ScriptBuilderMessage"]] = relationship(
"ScriptBuilderMessage",
back_populates="session",
order_by="ScriptBuilderMessage.created_at",
cascade="all, delete-orphan",
)
class ScriptBuilderMessage(Base):
"""A single message in a Script Builder conversation."""
__tablename__ = "script_builder_messages"
id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
)
session_id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True),
ForeignKey("script_builder_sessions.id", ondelete="CASCADE"),
nullable=False,
index=True,
)
role: Mapped[str] = mapped_column(
String(20), nullable=False, comment="user or assistant"
)
content: Mapped[str] = mapped_column(Text, nullable=False)
script: Mapped[Optional[str]] = mapped_column(
Text, nullable=True, comment="Extracted script from AI response"
)
script_filename: Mapped[Optional[str]] = mapped_column(
String(200), nullable=True
)
line_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
input_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
output_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
)
# Relationship
session: Mapped["ScriptBuilderSession"] = relationship(
"ScriptBuilderSession", back_populates="message_records"
)

View File

@@ -1,6 +1,6 @@
"""Pydantic schemas for the AI Script Builder.""" """Pydantic schemas for the AI Script Builder."""
from datetime import datetime from datetime import datetime
from typing import Any, Optional from typing import Optional
from uuid import UUID from uuid import UUID
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
@@ -20,6 +20,21 @@ class ScriptBuilderMessageRequest(BaseModel):
content: str = Field(min_length=1, max_length=5000) content: str = Field(min_length=1, max_length=5000)
class ScriptBuilderMessageSchema(BaseModel):
"""A single message in a builder session."""
id: UUID
role: str
content: str
script: str | None = None
script_filename: str | None = None
line_count: int | None = None
input_tokens: int | None = None
output_tokens: int | None = None
created_at: datetime
model_config = {"from_attributes": True}
class ScriptBuilderMessageResponse(BaseModel): class ScriptBuilderMessageResponse(BaseModel):
"""AI response to a builder message.""" """AI response to a builder message."""
role: str = "assistant" role: str = "assistant"
@@ -37,7 +52,7 @@ class ScriptBuilderSessionSummary(BaseModel):
id: UUID id: UUID
language: str language: str
title: str | None = None title: str | None = None
message_count: int message_count: int = 0
latest_script_filename: str | None = None latest_script_filename: str | None = None
created_at: datetime created_at: datetime
updated_at: datetime updated_at: datetime
@@ -50,10 +65,10 @@ class ScriptBuilderSessionDetail(BaseModel):
id: UUID id: UUID
language: str language: str
title: str | None = None title: str | None = None
messages: list[dict[str, Any]] messages: list[ScriptBuilderMessageSchema] = []
latest_script: str | None = None latest_script: str | None = None
latest_script_filename: str | None = None latest_script_filename: str | None = None
message_count: int message_count: int = 0
ai_session_id: UUID | None = None ai_session_id: UUID | None = None
created_at: datetime created_at: datetime
updated_at: datetime updated_at: datetime

View File

@@ -7,10 +7,11 @@ from uuid import UUID
from sqlalchemy import select, func from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from app.core.ai_provider import get_ai_provider from app.core.ai_provider import get_ai_provider
from app.core.config import settings from app.core.config import settings
from app.models.script_builder_session import ScriptBuilderSession from app.models.script_builder_session import ScriptBuilderSession, ScriptBuilderMessage
from app.schemas.script_builder import ( from app.schemas.script_builder import (
ScriptBuilderMessageResponse, ScriptBuilderMessageResponse,
ScriptBuilderSessionDetail, ScriptBuilderSessionDetail,
@@ -151,8 +152,6 @@ async def create_session(
user_id=user_id, user_id=user_id,
team_id=team_id, team_id=team_id,
language=language, language=language,
messages=[],
message_count=0,
) )
db.add(session) db.add(session)
await db.flush() await db.flush()
@@ -170,30 +169,43 @@ async def send_message(
user_content: str, user_content: str,
) -> ScriptBuilderMessageResponse: ) -> ScriptBuilderMessageResponse:
"""Send a user message and get AI response with generated script.""" """Send a user message and get AI response with generated script."""
if session.message_count >= MAX_MESSAGES_PER_SESSION: # Count existing messages for the session
msg_count_result = await db.execute(
select(func.count(ScriptBuilderMessage.id)).where(
ScriptBuilderMessage.session_id == session.id,
ScriptBuilderMessage.role == "user",
)
)
user_msg_count = msg_count_result.scalar_one()
if user_msg_count >= MAX_MESSAGES_PER_SESSION:
raise ValueError(f"Session has reached the maximum of {MAX_MESSAGES_PER_SESSION} messages.") raise ValueError(f"Session has reached the maximum of {MAX_MESSAGES_PER_SESSION} messages.")
now = datetime.now(timezone.utc) now = datetime.now(timezone.utc)
# Append user message # Create user message record
user_msg = { user_msg = ScriptBuilderMessage(
"role": "user", session_id=session.id,
"content": user_content, role="user",
"timestamp": now.isoformat(), content=user_content,
} created_at=now,
messages = list(session.messages or []) )
messages.append(user_msg) db.add(user_msg)
await db.flush()
# Build system prompt # Build system prompt
language_prompt = LANGUAGE_PROMPTS.get(session.language, LANGUAGE_PROMPTS["powershell"]) language_prompt = LANGUAGE_PROMPTS.get(session.language, LANGUAGE_PROMPTS["powershell"])
system_prompt = SYSTEM_PROMPT_TEMPLATE.format(language_prompt=language_prompt) system_prompt = SYSTEM_PROMPT_TEMPLATE.format(language_prompt=language_prompt)
# Build conversation for AI (just role + content) # Build conversation for AI — get last 20 messages for context window
ai_messages = [{"role": m["role"], "content": m["content"]} for m in messages] recent_result = await db.execute(
select(ScriptBuilderMessage)
# Context window management: keep last 20 messages (10 exchanges) .where(ScriptBuilderMessage.session_id == session.id)
if len(ai_messages) > 20: .order_by(ScriptBuilderMessage.created_at.desc())
ai_messages = ai_messages[-20:] .limit(20)
)
recent_msgs = list(reversed(recent_result.scalars().all()))
ai_messages = [{"role": m.role, "content": m.content} for m in recent_msgs]
# Call AI # Call AI
model = settings.get_model_for_action("script_build") model = settings.get_model_for_action("script_build")
@@ -208,28 +220,25 @@ async def send_message(
script, filename = _extract_script_from_response(ai_text, session.language) script, filename = _extract_script_from_response(ai_text, session.language)
line_count = len(script.splitlines()) if script else None line_count = len(script.splitlines()) if script else None
# Build assistant message # Create assistant message record
assistant_msg = { assistant_msg = ScriptBuilderMessage(
"role": "assistant", session_id=session.id,
"content": ai_text, role="assistant",
"timestamp": datetime.now(timezone.utc).isoformat(), content=ai_text,
"input_tokens": input_tokens, script=script,
"output_tokens": output_tokens, script_filename=filename,
} line_count=line_count,
if script: input_tokens=input_tokens,
assistant_msg["script"] = script output_tokens=output_tokens,
assistant_msg["script_filename"] = filename created_at=datetime.now(timezone.utc),
assistant_msg["line_count"] = line_count )
db.add(assistant_msg)
messages.append(assistant_msg) # Update session denormalized fields
# Update session
session.messages = messages
session.message_count = len([m for m in messages if m["role"] == "user"])
if script: if script:
session.latest_script = script session.latest_script = script
session.latest_script_filename = filename session.latest_script_filename = filename
if not session.title and len(messages) >= 2: if not session.title:
# Auto-generate title from first user message (truncate) # Auto-generate title from first user message (truncate)
first_user = user_content[:100] first_user = user_content[:100]
session.title = first_user if len(user_content) <= 100 else first_user + "..." session.title = first_user if len(user_content) <= 100 else first_user + "..."
@@ -254,7 +263,9 @@ async def get_session(
) -> ScriptBuilderSession | None: ) -> ScriptBuilderSession | None:
"""Get a session by ID, ensuring the user owns it.""" """Get a session by ID, ensuring the user owns it."""
result = await db.execute( result = await db.execute(
select(ScriptBuilderSession).where( select(ScriptBuilderSession)
.options(selectinload(ScriptBuilderSession.message_records))
.where(
ScriptBuilderSession.id == session_id, ScriptBuilderSession.id == session_id,
ScriptBuilderSession.user_id == user_id, ScriptBuilderSession.user_id == user_id,
) )

View File

@@ -66,7 +66,7 @@ export default function ScriptBuilderPage() {
const userMessage: ScriptBuilderMessage = { const userMessage: ScriptBuilderMessage = {
role: 'user', role: 'user',
content, content,
timestamp: new Date().toISOString(), created_at: new Date().toISOString(),
} }
setMessages((prev) => [...prev, userMessage]) setMessages((prev) => [...prev, userMessage])
setIsLoading(true) setIsLoading(true)
@@ -88,7 +88,7 @@ export default function ScriptBuilderPage() {
script: response.script, script: response.script,
script_filename: response.script_filename, script_filename: response.script_filename,
line_count: response.line_count, line_count: response.line_count,
timestamp: response.timestamp, created_at: response.timestamp,
} }
setMessages((prev) => [...prev, assistantMessage]) setMessages((prev) => [...prev, assistantMessage])
} catch (err) { } catch (err) {
@@ -96,7 +96,7 @@ export default function ScriptBuilderPage() {
const errorMessage: ScriptBuilderMessage = { const errorMessage: ScriptBuilderMessage = {
role: 'assistant', role: 'assistant',
content: `An error occurred: ${err instanceof Error ? err.message : 'Failed to generate response. Please try again.'}`, content: `An error occurred: ${err instanceof Error ? err.message : 'Failed to generate response. Please try again.'}`,
timestamp: new Date().toISOString(), created_at: new Date().toISOString(),
} }
setMessages((prev) => [...prev, errorMessage]) setMessages((prev) => [...prev, errorMessage])
} finally { } finally {

View File

@@ -56,11 +56,11 @@ export interface ScriptBuilderContent {
script_language?: string script_language?: string
} }
export function isScriptGenerationAction(content: Record<string, unknown>): content is ScriptGenerationContent { export function isScriptGenerationAction(content: Record<string, unknown>): content is Record<string, unknown> & ScriptGenerationContent {
return content.action_type === 'script_generation' return content.action_type === 'script_generation'
} }
export function isScriptBuilderAction(content: Record<string, unknown>): content is ScriptBuilderContent { export function isScriptBuilderAction(content: Record<string, unknown>): content is Record<string, unknown> & ScriptBuilderContent {
return content.action_type === 'open_script_builder' return content.action_type === 'open_script_builder'
} }

View File

@@ -22,12 +22,15 @@ export interface ScriptBuilderSessionDetail {
} }
export interface ScriptBuilderMessage { export interface ScriptBuilderMessage {
id?: string
role: 'user' | 'assistant' role: 'user' | 'assistant'
content: string content: string
script?: string | null script?: string | null
script_filename?: string | null script_filename?: string | null
line_count?: number | null line_count?: number | null
timestamp: string input_tokens?: number | null
output_tokens?: number | null
created_at: string
} }
export interface ScriptBuilderMessageResponse { export interface ScriptBuilderMessageResponse {