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:
@@ -9,10 +9,12 @@ from app.core.database import get_db
|
||||
from app.core.rate_limit import limiter
|
||||
from app.api.deps import get_current_active_user
|
||||
from app.models.user import User
|
||||
from app.models.script_builder_session import ScriptBuilderSession
|
||||
from app.schemas.script_builder import (
|
||||
ScriptBuilderCreateRequest,
|
||||
ScriptBuilderMessageRequest,
|
||||
ScriptBuilderMessageResponse,
|
||||
ScriptBuilderMessageSchema,
|
||||
ScriptBuilderSessionDetail,
|
||||
ScriptBuilderSessionSummary,
|
||||
SaveToLibraryRequest,
|
||||
@@ -25,6 +27,39 @@ router = APIRouter(prefix="/scripts/builder", tags=["script-builder"])
|
||||
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)
|
||||
async def create_session(
|
||||
data: ScriptBuilderCreateRequest,
|
||||
@@ -47,7 +82,9 @@ async def create_session(
|
||||
language=data.language,
|
||||
)
|
||||
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])
|
||||
@@ -61,7 +98,7 @@ async def list_sessions(
|
||||
sessions = await script_builder_service.list_sessions(
|
||||
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)
|
||||
@@ -74,7 +111,7 @@ async def get_session(
|
||||
session = await script_builder_service.get_session(db, session_id, current_user.id)
|
||||
if not session:
|
||||
raise HTTPException(status_code=404, detail="Session not found")
|
||||
return ScriptBuilderSessionDetail.model_validate(session)
|
||||
return _session_to_detail(session)
|
||||
|
||||
|
||||
@router.post(
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
"""Script Builder session model.
|
||||
"""Script Builder session and message models.
|
||||
|
||||
Tracks AI-powered script generation conversations.
|
||||
"""
|
||||
import uuid
|
||||
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.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
|
||||
|
||||
@@ -42,10 +42,6 @@ class ScriptBuilderSession(Base):
|
||||
String(200), nullable=True,
|
||||
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(
|
||||
Text, nullable=True,
|
||||
comment="Most recent generated script for quick access",
|
||||
@@ -54,9 +50,6 @@ class ScriptBuilderSession(Base):
|
||||
String(200), nullable=True,
|
||||
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(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("ai_sessions.id", ondelete="SET NULL"),
|
||||
@@ -74,3 +67,45 @@ class ScriptBuilderSession(Base):
|
||||
|
||||
# Relationships
|
||||
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"
|
||||
)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Pydantic schemas for the AI Script Builder."""
|
||||
from datetime import datetime
|
||||
from typing import Any, Optional
|
||||
from typing import Optional
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
@@ -20,6 +20,21 @@ class ScriptBuilderMessageRequest(BaseModel):
|
||||
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):
|
||||
"""AI response to a builder message."""
|
||||
role: str = "assistant"
|
||||
@@ -37,7 +52,7 @@ class ScriptBuilderSessionSummary(BaseModel):
|
||||
id: UUID
|
||||
language: str
|
||||
title: str | None = None
|
||||
message_count: int
|
||||
message_count: int = 0
|
||||
latest_script_filename: str | None = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
@@ -50,10 +65,10 @@ class ScriptBuilderSessionDetail(BaseModel):
|
||||
id: UUID
|
||||
language: str
|
||||
title: str | None = None
|
||||
messages: list[dict[str, Any]]
|
||||
messages: list[ScriptBuilderMessageSchema] = []
|
||||
latest_script: str | None = None
|
||||
latest_script_filename: str | None = None
|
||||
message_count: int
|
||||
message_count: int = 0
|
||||
ai_session_id: UUID | None = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user