Skip to content

feat: Add a new video generation AGENT that uses the VEO model #555

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: next
Choose a base branch
from
Open

feat: Add a new video generation AGENT that uses the VEO model #555

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b59c23d
Add a new video generation AGENT that uses the VEO model
dmandar May 15, 2025
636f04c
Merge branch 'next' into md-video
holtskinner May 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
218 changes: 116 additions & 102 deletions demo/ui/components/chat_bubble.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,118 +1,132 @@
import mesop as me

from state.state import AppState, StateMessage


@me.component
def chat_bubble(message: StateMessage, key: str):
"""Chat bubble component"""
app_state = me.state(AppState)
show_progress_bar = (
message.message_id in app_state.background_tasks
or message.message_id in app_state.message_aliases.values()
)
progress_text = ''
if show_progress_bar:
progress_text = app_state.background_tasks[message.message_id]
if not message.content:
print('No message content')
for pair in message.content:
chat_box(
pair[0],
pair[1],
message.role,
key,
progress_bar=show_progress_bar,
progress_text=progress_text,
)


def chat_box(
# Helper function to render individual message content parts (text, image, video)
def _render_message_part(
content: str,
media_type: str,
role: str,
key: str,
progress_bar: bool,
progress_text: str,
):
with me.box(
style=me.Style(
display='flex',
justify_content=('space-between' if role == 'agent' else 'end'),
min_width=500,
"""Render a single part of a message (text, image, video) with bubble styling."""
bubble_style = me.Style(
padding=me.Padding(top=8, left=15, right=15, bottom=8),
margin=me.Margin(top=2, bottom=2),
background=(
me.theme_var('primary-container')
if role == 'user'
else me.theme_var('secondary-container')
),
key=key,
):
with me.box(
style=me.Style(display='flex', flex_direction='column', gap=5)
):
if media_type == 'image/png':
if '/message/file' not in content:
content = 'data:image/png;base64,' + content
me.image(
src=content,
border_radius=15,
box_shadow=(
'0 1px 2px 0 rgba(60, 64, 67, 0.3), '
'0 1px 3px 1px rgba(60, 64, 67, 0.15)'
),
max_width='75%',
word_wrap='break-word',
)

with me.box(style=bubble_style, key=key):
if media_type == 'image/png':
if '/message/file' not in content:
content = 'data:image/png;base64,' + content
me.image(
src=content,
style=me.Style(
max_width='100%',
height='auto',
border_radius=10,
display='block',
),
)
elif media_type and media_type.startswith('video/'): # Handle video/mp4, video/webm, etc.
# Check if the content is a string and looks like a URL or GCS URI
if isinstance(content, str) and (content.startswith('http://') or content.startswith('https://') or content.startswith('gs://')):
me.video(
src=content, # Use the URI as the source
style=me.Style(
width='50%',
object_fit='contain',
width='100%', # Video responsive within the bubble
max_width='480px', # Max width for video elements
border_radius=10,
display='block',
),
)
else:
# Fallback if content is not a valid URL for video
# This might happen if the FilePart had bytes but no URI,
# and extract_content put the bytes (or default) here.
# Or if mimeType was video but content was None.
me.markdown(
content,
style=me.Style(
font_family='Google Sans',
box_shadow=(
'0 1px 2px 0 rgba(60, 64, 67, 0.3), '
'0 1px 3px 1px rgba(60, 64, 67, 0.15)'
),
padding=me.Padding(top=1, left=15, right=15, bottom=1),
margin=me.Margin(top=5, left=0, right=0, bottom=5),
background=(
me.theme_var('primary-container')
if role == 'user'
else me.theme_var('secondary-container')
),
border_radius=15,
),
f"Video content (type: {media_type}) could not be displayed. Content snippet: {str(content)[:100]}...",
style=me.Style(font_family='Google Sans'),
)
if progress_bar:
with me.box(
style=me.Style(
display='flex',
justify_content=('space-between' if role == 'user' else 'end'),
min_width=500,
),
key=key,
):
with me.box(
style=me.Style(display='flex', flex_direction='column', gap=5)
):
with me.box(
style=me.Style(
font_family='Google Sans',
box_shadow=(
'0 1px 2px 0 rgba(60, 64, 67, 0.3), '
'0 1px 3px 1px rgba(60, 64, 67, 0.15)'
),
padding=me.Padding(top=1, left=15, right=15, bottom=1),
margin=me.Margin(top=5, left=0, right=0, bottom=5),
background=(
me.theme_var('primary-container')
if role == 'agent'
else me.theme_var('secondary-container')
),
border_radius=15,
),
):
if not progress_text:
progress_text = 'Working...'
me.text(
progress_text,
style=me.Style(
padding=me.Padding(
top=1, left=15, right=15, bottom=1
),
margin=me.Margin(top=5, left=0, right=0, bottom=5),
),
)
me.progress_bar(color='accent')
else: # Default to markdown for text/plain, application/json, etc.
me.markdown(
content,
style=me.Style(font_family='Google Sans'),
)


# Helper function to render the progress indicator for a message
def _render_message_progress(progress_text: str, role: str, key: str):
"""Renders a progress indicator bubble for a message."""
progress_bubble_style = me.Style(
padding=me.Padding(top=8, left=15, right=15, bottom=8),
margin=me.Margin(top=5),
background=me.theme_var('surface-variant'), # Distinct background for progress
border_radius=15,
box_shadow=(
'0 1px 2px 0 rgba(60, 64, 67, 0.3), '
'0 1px 3px 1px rgba(60, 64, 67, 0.15)'
),
max_width='75%',
font_family='Google Sans',
)
with me.box(style=progress_bubble_style, key=key):
me.text(
progress_text,
style=me.Style(padding=me.Padding(bottom=5)),
)
me.progress_bar(color='accent')


@me.component
def chat_bubble(message: StateMessage, key: str):
"""Chat bubble component. Key should be message.message_id."""
app_state = me.state(AppState)

if not message.content:
# If there's no content, but it's a pending task, we might still show progress.
# This case can be refined if needed. For now, requires content to render bubble.
print(f'No message content for message_id: {message.message_id}')

# Main container for the entire message (all parts + progress bar)
# Aligns the whole message to the left (agent) or right (user)
with me.box(
style=me.Style(
display='flex',
flex_direction='column',
align_items='flex-start' if message.role == 'agent' else 'flex-end',
width='100%', # Take full width to allow alignment
margin=me.Margin(bottom=10), # Space between messages
),
):
# Render each part of the message content
for idx, pair in enumerate(message.content):
content_data, media_type = pair
part_key = f"{key}_part_{idx}" # Unique key for each message part
_render_message_part(content_data, media_type, message.role, part_key)

# Render progress bar for the entire message if it's a background task
show_progress_bar = (
message.message_id in app_state.background_tasks
or message.message_id in app_state.message_aliases.values()
)
if show_progress_bar:
progress_text = app_state.background_tasks.get(message.message_id)
# Ensure progress_text is a string, default to 'Working...'
if not isinstance(progress_text, str) or not progress_text.strip():
progress_text = 'Working...'
progress_bar_key = f"{key}_progress" # Unique key for the progress bar
_render_message_progress(progress_text, message.role, progress_bar_key)
32 changes: 17 additions & 15 deletions demo/ui/components/conversation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,21 +1,9 @@
import mesop as me

import uuid

from state.state import AppState, SettingsState, StateMessage
from state.host_agent_service import (
SendMessage,
ListConversations,
convert_message_to_state,
)
from .chat_bubble import chat_bubble
from .form_render import is_form, render_form, form_sent
import mesop as me
from common.types import Message, TextPart
from state.host_agent_service import (
ListConversations,
SendMessage,
convert_message_to_state,
)
from state.host_agent_service import (ListConversations, SendMessage,
convert_message_to_state)
from state.state import AppState, SettingsState, StateMessage

from .chat_bubble import chat_bubble
Expand Down Expand Up @@ -126,6 +114,20 @@ def conversation():
else:
chat_bubble(message, message.message_id)

# After rendering messages, check for completed tasks associated with the current conversation
# and render their artifacts if they haven't been implicitly shown as a message.
for session_task in app_state.task_list:
if session_task.context_id == page_state.conversation_id and session_task.task.state == "TaskState.COMPLETED":
# Check if this task's result is already part of a message to avoid duplication
# This is a simple check; a more robust way might involve linking task artifacts to specific messages.
task_already_in_message = any(msg.task_id == session_task.task.task_id and msg.content for msg in app_state.messages)

if not task_already_in_message and session_task.task.artifacts:
for artifact_list in session_task.task.artifacts: # artifacts is list[list[Tuple[ContentPart, str]]]
# Create a temporary StateMessage to render the artifact via chat_bubble
artifact_message = StateMessage(message_id=f"task_artifact_{session_task.task.task_id}", role="agent", content=artifact_list)
chat_bubble(artifact_message, artifact_message.message_id)

with me.box(
style=me.Style(
display='flex',
Expand Down
2 changes: 1 addition & 1 deletion demo/ui/pages/settings.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import asyncio

import mesop as me

from components.header import header
from components.page_scaffold import page_frame, page_scaffold
from state.host_agent_service import UpdateApiKey
Expand Down Expand Up @@ -157,6 +156,7 @@ def settings_page_content():
me.SelectOption(
label='Text (Plain)', value='text/plain'
),
me.SelectOption(label='Video', value='video/mp4'),
],
on_selection_change=on_selection_change_output_types,
style=me.Style(width=500),
Expand Down
Loading
Loading