diff --git a/backend/alembic/env.py b/backend/alembic/env.py index 87f08042..f3a1014d 100644 --- a/backend/alembic/env.py +++ b/backend/alembic/env.py @@ -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.psa_post_log import PsaPostLog # 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 # this is the Alembic Config object diff --git a/backend/alembic/versions/064_normalize_script_builder_messages.py b/backend/alembic/versions/064_normalize_script_builder_messages.py new file mode 100644 index 00000000..224d1d8c --- /dev/null +++ b/backend/alembic/versions/064_normalize_script_builder_messages.py @@ -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) diff --git a/backend/app/api/endpoints/script_builder.py b/backend/app/api/endpoints/script_builder.py index a477d7a9..1ac5851c 100644 --- a/backend/app/api/endpoints/script_builder.py +++ b/backend/app/api/endpoints/script_builder.py @@ -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( diff --git a/backend/app/models/script_builder_session.py b/backend/app/models/script_builder_session.py index fe8009e9..f7075494 100644 --- a/backend/app/models/script_builder_session.py +++ b/backend/app/models/script_builder_session.py @@ -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" + ) diff --git a/backend/app/schemas/script_builder.py b/backend/app/schemas/script_builder.py index 9c74c450..893d16be 100644 --- a/backend/app/schemas/script_builder.py +++ b/backend/app/schemas/script_builder.py @@ -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 diff --git a/backend/app/services/script_builder_service.py b/backend/app/services/script_builder_service.py index ae39fae7..3fdeb3ea 100644 --- a/backend/app/services/script_builder_service.py +++ b/backend/app/services/script_builder_service.py @@ -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, ) diff --git a/frontend/src/pages/ScriptBuilderPage.tsx b/frontend/src/pages/ScriptBuilderPage.tsx index ca567497..d14714ef 100644 --- a/frontend/src/pages/ScriptBuilderPage.tsx +++ b/frontend/src/pages/ScriptBuilderPage.tsx @@ -66,7 +66,7 @@ export default function ScriptBuilderPage() { const userMessage: ScriptBuilderMessage = { role: 'user', content, - timestamp: new Date().toISOString(), + created_at: new Date().toISOString(), } setMessages((prev) => [...prev, userMessage]) setIsLoading(true) @@ -88,7 +88,7 @@ export default function ScriptBuilderPage() { script: response.script, script_filename: response.script_filename, line_count: response.line_count, - timestamp: response.timestamp, + created_at: response.timestamp, } setMessages((prev) => [...prev, assistantMessage]) } catch (err) { @@ -96,7 +96,7 @@ export default function ScriptBuilderPage() { const errorMessage: ScriptBuilderMessage = { role: 'assistant', 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]) } finally { diff --git a/frontend/src/types/ai-session.ts b/frontend/src/types/ai-session.ts index a18e1dfe..34bda117 100644 --- a/frontend/src/types/ai-session.ts +++ b/frontend/src/types/ai-session.ts @@ -56,11 +56,11 @@ export interface ScriptBuilderContent { script_language?: string } -export function isScriptGenerationAction(content: Record): content is ScriptGenerationContent { +export function isScriptGenerationAction(content: Record): content is Record & ScriptGenerationContent { return content.action_type === 'script_generation' } -export function isScriptBuilderAction(content: Record): content is ScriptBuilderContent { +export function isScriptBuilderAction(content: Record): content is Record & ScriptBuilderContent { return content.action_type === 'open_script_builder' } diff --git a/frontend/src/types/script-builder.ts b/frontend/src/types/script-builder.ts index 76dbae02..cfa8165e 100644 --- a/frontend/src/types/script-builder.ts +++ b/frontend/src/types/script-builder.ts @@ -22,12 +22,15 @@ export interface ScriptBuilderSessionDetail { } export interface ScriptBuilderMessage { + id?: string role: 'user' | 'assistant' content: string script?: string | null script_filename?: string | null line_count?: number | null - timestamp: string + input_tokens?: number | null + output_tokens?: number | null + created_at: string } export interface ScriptBuilderMessageResponse {