src/genai/brainstroming_agent/utils/state.py

from pydantic import BaseModel, ConfigDict , Field
from typing import Optional
import pandas as pd

class State(BaseModel):
    carry_on: Optional[bool]=False
    messages: Optional[str] = None
    idea: list
    brainstroming_topics: Optional[list] = []
    preferred_topics: Optional[list] = []
    stories : Optional[list]=[]
    final_story: Optional[list]=[]
    retrievals : Optional[list]=[]
    business_details : Optional[dict]={}
    latest_preferred_topics: Optional[list] = []
    images: Optional[list] = []  # Base64-encoded strings of images
    image_captions: Optional[list] = []
    model_config = ConfigDict(arbitrary_types_allowed=True)


class BrainstromTopicFormatter(BaseModel):
    topic1:str=Field(description="First brainstorming topic of the story")
    topic2:str=Field(description="Second brainstorming topic of the story")
    topic3:str=Field(description="Third brainstorming topic of the story")
    topic4:str=Field(description="Fourth brainstorming topic of the story")

class QueryFormatter(BaseModel):
    video_topic: str = Field(description="The video topic that user passes to the agent")
    # business_details: dict = Field (description="The business details that user passes to the agent")