api/routers/brainstorm.py

import json
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from api.stored_data import stored_data
from pydantic import BaseModel
from typing import Optional
from src.genai.brainstroming_agent.agent import brainstroming_graph
from langgraph.pregel.io import AddableValuesDict

router=APIRouter()
brainstrom_graph = brainstroming_graph()


class BrainstormRequest(BaseModel):
    preferred_topics: Optional[list] = []
    image_base64_list: Optional[list] = []
    thread_id: Optional[str]="default-session"


@router.post("/brainstorm")
def brainstorming_endpoint(request: BrainstormRequest):
    idea = (
        [stored_data['refined_ideation']]
        if stored_data.get('refined_ideation')
        else [str(stored_data['final_ideation'])]
        if stored_data.get('final_ideation')
        else [
            "I don't have any idea right now. Create your own **very creative** and **out of the box** video idea and generate the story for now."
        ]
    )

    def event_stream():

        for stream_mode, chunk in brainstrom_graph.stream(
            {
                "idea": idea,
                "images": request.image_base64_list,
                "latest_preferred_topics": request.preferred_topics,
                "business_details": stored_data["business_details"],
            },
            config={"configurable": {"thread_id": request.thread_id}},
            stream_mode=['messages','values'] 
        ):
            stored_data['brainstorming_response'] = chunk if isinstance(chunk, AddableValuesDict) else stored_data.get('brainstorming_response')
            if isinstance(chunk, tuple): yield chunk[0].content
    return StreamingResponse(event_stream(), media_type="text/event-stream")