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

@@ -7,10 +7,11 @@ from uuid import UUID
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from app.core.ai_provider import get_ai_provider
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 (
ScriptBuilderMessageResponse,
ScriptBuilderSessionDetail,
@@ -151,8 +152,6 @@ async def create_session(
user_id=user_id,
team_id=team_id,
language=language,
messages=[],
message_count=0,
)
db.add(session)
await db.flush()
@@ -170,30 +169,43 @@ async def send_message(
user_content: str,
) -> ScriptBuilderMessageResponse:
"""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.")
now = datetime.now(timezone.utc)
# Append user message
user_msg = {
"role": "user",
"content": user_content,
"timestamp": now.isoformat(),
}
messages = list(session.messages or [])
messages.append(user_msg)
# Create user message record
user_msg = ScriptBuilderMessage(
session_id=session.id,
role="user",
content=user_content,
created_at=now,
)
db.add(user_msg)
await db.flush()
# Build system prompt
language_prompt = LANGUAGE_PROMPTS.get(session.language, LANGUAGE_PROMPTS["powershell"])
system_prompt = SYSTEM_PROMPT_TEMPLATE.format(language_prompt=language_prompt)
# Build conversation for AI (just role + content)
ai_messages = [{"role": m["role"], "content": m["content"]} for m in messages]
# Context window management: keep last 20 messages (10 exchanges)
if len(ai_messages) > 20:
ai_messages = ai_messages[-20:]
# Build conversation for AI — get last 20 messages for context window
recent_result = await db.execute(
select(ScriptBuilderMessage)
.where(ScriptBuilderMessage.session_id == session.id)
.order_by(ScriptBuilderMessage.created_at.desc())
.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
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)
line_count = len(script.splitlines()) if script else None
# Build assistant message
assistant_msg = {
"role": "assistant",
"content": ai_text,
"timestamp": datetime.now(timezone.utc).isoformat(),
"input_tokens": input_tokens,
"output_tokens": output_tokens,
}
if script:
assistant_msg["script"] = script
assistant_msg["script_filename"] = filename
assistant_msg["line_count"] = line_count
# Create assistant message record
assistant_msg = ScriptBuilderMessage(
session_id=session.id,
role="assistant",
content=ai_text,
script=script,
script_filename=filename,
line_count=line_count,
input_tokens=input_tokens,
output_tokens=output_tokens,
created_at=datetime.now(timezone.utc),
)
db.add(assistant_msg)
messages.append(assistant_msg)
# Update session
session.messages = messages
session.message_count = len([m for m in messages if m["role"] == "user"])
# Update session denormalized fields
if script:
session.latest_script = script
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)
first_user = user_content[:100]
session.title = first_user if len(user_content) <= 100 else first_user + "..."
@@ -254,7 +263,9 @@ async def get_session(
) -> ScriptBuilderSession | None:
"""Get a session by ID, ensuring the user owns it."""
result = await db.execute(
select(ScriptBuilderSession).where(
select(ScriptBuilderSession)
.options(selectinload(ScriptBuilderSession.message_records))
.where(
ScriptBuilderSession.id == session_id,
ScriptBuilderSession.user_id == user_id,
)