src/genai/ideation_agent/utils/state.py

from pydantic import BaseModel, Field
from typing import Optional
import operator
from typing import Annotated

class State(BaseModel):
  business_details : Annotated[list[dict],operator.add] = []
  ideator_response: Annotated[list[str],operator.add] = []
  critic_response: Annotated[list[str],operator.add] = []
  critic_fault: Annotated[list[str],operator.add] = []
  improver_response: Annotated[list[str],operator.add] = []
  validator1_response: Annotated[list[str],operator.add] = []
  validator2_response: Annotated[list[str],operator.add] = []
  validator3_response: Annotated[list[str],operator.add] = []
  validator4_response: Annotated[list[str],operator.add] = []
  disagreement_reason: Annotated[list[str],operator.add] = []
  image_caption: Annotated[list[str],operator.add] = []

class QueryFormatter(BaseModel):
    business_details: str = Field(description="The details of the business that user passes to the agent")

class ImproverResponseFormatter(BaseModel):
   faults: str= Field(description="Returns the faults of the critic's response if any.")
   improved_idea1: str = Field(description="Returns the first improved idea of exactly 40 words. If no faults found, repeat the critic's version.")
   improved_idea2: str = Field(description="Returns the second improved idea of exactly 40 words. If no faults found, repeat the critic's version.")
   improved_idea3: str = Field(description="Returns the third improved idea of exactly 40 words. If no faults found, repeat the critic's version.")
   improved_idea4: str = Field(description="Returns the fourth improved idea of exactly 40 words. If no faults found, repeat the critic's version.")

class ValidationFormatter(BaseModel):
  result: str = Field(description="Returns **validated** if the idea is validated. Returns **not validated** if story is not validated.")
  reason: str = Field(description="If the idea is **not validated**, it returns the reason why it is not validated.")