api/routers/brainstorm.py

import logging
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from api.stored_data import stored_data
from src.genai.brainstroming_agent.agent import BrainstormingAgent
from api.schemas.brainstorming import BrainstormRequest
from api.backup_prompts import brainstorm_backup_prompt

router=APIRouter()
agent = BrainstormingAgent()
brainstorm_graph = agent.brainstorming_graph()
app_logger = logging.getLogger("app_logger")
error_logger = logging.getLogger("error_logger")


@router.post("/brainstorm")
def brainstorming_endpoint(request: BrainstormRequest):
    try:
        idea = (
            [stored_data['refined_ideation']]
            if stored_data.get('refined_ideation')
            else [str(stored_data['final_ideation'])]
            if stored_data.get('final_ideation')
            else [brainstorm_backup_prompt]
        )

        def event_stream():
            for stream_mode, chunk in brainstorm_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, dict) else stored_data.get('brainstorming_response')
                # if isinstance(chunk, tuple): yield f"data: {json.dumps(chunk[0].content)}"
                if isinstance(chunk, tuple): yield chunk[0].content
        app_logger.info('Executed brainstorming agent.')
        return StreamingResponse(event_stream(), media_type="text/event-stream")
    
    except Exception as e:
        error_logger.error('Unable to execute brainstorm agent.')
        raise