Issue with state management , how to prevent structured output from being streamed

[AdityaVernekar]

I am building ai ai agent that generates content for you based on some predefined steps

Here is my code

States

class LinkedInState(TypedDict):
    """State for the LinkedIn agent."""
    messages: Annotated[list[AnyMessage], add_messages]
    context: Optional[Dict[str, RunningSummary]]  # Running summary context
    summarized_messages: Optional[list] = None
    topic: str
    tone: str
    audience: str
    goal: str
    is_websearch_required: bool
    is_context_valid: bool
    content: str
    follow_up_response: str

class InputState(TypedDict):
    """State for the LinkedIn agent."""
    messages: Annotated[list[AnyMessage], add_messages]
    user_id: str 
    space_id: str 


class OutputState(TypedDict):
    """State for the LinkedIn agent."""
    messages: Annotated[list[AnyMessage], add_messages]
    content: str 
    is_context_valid: bool 

class ContextValidatorOutput(BaseModel):
    """Output for the context validator agent."""
    is_context_valid: bool = Field(default=False,description="Whether the context is valid or not.")
    topic: str = Field(default="",description="The topic for the post.")
    tone: str = Field(default="",description="The tone for the post.")
    audience: str = Field(default="",description="The audience for the post.")
    goal: str = Field(default="",description="The goal for the post.")
    is_websearch_required: bool = Field(default=False,description="Whether web search is required or not.")
    follow_up_response: str = Field(default="",description="The response contains the follow up questions to the user to get the missing information.")
   

model = ChatOpenAI(model="gpt-4o-mini")
async def context_validator(state: LinkedInState) -> LinkedInState:
    messages = [
        SystemMessage(content=prompt),
        *state["messages"]
    ]

    response = await model.with_structured_output(ContextValidatorOutput).ainvoke(messages)
    # Return a new state with updated values using TypedDict syntax
    return {
        "topic": response.topic,
        "tone": response.tone,
        "audience": response.audience,
        "goal": response.goal,
        "is_websearch_required": response.is_websearch_required,
        "is_context_valid": response.is_context_valid,
        "follow_up_response": response.follow_up_response
    }


async def ask_follow_up_questions(state: LinkedInState) -> LinkedInState:
    # Create a new message with the follow-up response
    new_message = AIMessage(content=state["follow_up_response"])
    
    # Return updated state with new message using TypedDict syntax
    return {
        "messages": [*state["messages"], new_message],
        "topic": state["topic"],
        "tone": state["tone"],
        "audience": state["audience"],
        "goal": state["goal"],
        "is_websearch_required": state["is_websearch_required"],
        "is_context_valid": state["is_context_valid"],
        "content":"",  # Provide default value if content doesn't exist
        "follow_up_response": state["follow_up_response"]
    }

class LinkedInGraph:    
    def __init__(self, llm: Optional[ChatOpenAI] = None, batch_interval: float = 4.0, min_content_threshold: int = 10):
        self.llm = llm or ChatOpenAI(model="gpt-4o-mini")
        self.logger = Logger.get_logger(__name__)
        
    def _create_graph(self):

        graph = StateGraph(state_schema=LinkedInState,input_schema=InputState,output_schema=OutputState)


        def is_context_validated(state: LinkedInState):
            if state["is_context_valid"]:
                return "end"
            else:
                return "ask_follow_up_questions"

        graph.add_node("summarize",self.summarization_node)
        graph.add_node("context_validator",context_validator)
        graph.add_node("ask_follow_up_questions",ask_follow_up_questions)
        graph.add_edge(START, "summarize")

        graph.add_edge("summarize", "context_validator")
        graph.add_conditional_edges("context_validator",is_context_validated,{
            "end": END,
            "ask_follow_up_questions": "ask_follow_up_questions"
        })
        graph.add_edge("ask_follow_up_questions", END)


        return graph

    async def stream_request(self, request: SupervisorRequest) -> AsyncGenerator[StreamingEvent, None]:
        """Stream process a request using the handoff-based multi-agent system with memory and time-based batching."""
        try:
            start_time = time.time()
            last_progress_update = start_time
            progress_update_interval = 5.0  # Send progress updates every 5 seconds
            
            # Initialize utilities with time-based batching (4 second intervals)
            message_manager = MessageGroupingManager(
                batch_interval=self.batch_interval, 
                min_content_threshold=self.min_content_threshold
            )
            
            # Initialize message ID manager for conversation flow
            message_id_manager = MessageIDManager()
            initial_message_id = message_id_manager.start_conversation()
            
            # Send start event with initial message ID
            yield StartEvent(
                message="Thinking..",
                task_name="marketing_supervisor",
                message_id=initial_message_id,
                metadata=StreamingEventFormatter.clean_metadata(
                    use_case_id=request.use_case_id,
                    space_id=request.space_id
                )
            )
            
            from langchain_core.messages import HumanMessage
            
            initial_state = LinkedInState(
                messages=[
                    HumanMessage(content=request.query)
                ],
                context={},
                summarized_messages=[],
                topic="",
                tone="",
                audience="",
                goal="",
                is_websearch_required=False,
                is_context_valid=False,
                content="",
                follow_up_response=""
            )
            
            # Get or create thread ID for memory persistence
            thread_id = await self._get_or_create_thread_id(request.space_id)
            
            # Create runtime context with user information
            user_data = await get_user_data(request.user_id)
            runtime_context = user_data

            mongodb_url = sttgs.get("MONGODB_URI")
            async with AsyncMongoDBSaver.from_conn_string(mongodb_url) as checkpointer:
                # await checkpointer.asetup()
                
            # Compile graph with in-memory checkpointer
                graph = self._create_graph().compile(checkpointer=checkpointer)
            
            # Stream the request and yield responses
                async for item in graph.astream(
                        initial_state,
                        config={"configurable": {"thread_id": thread_id,"user_id":request.user_id}},
                        context=runtime_context,
                        subgraphs=True,
                        stream_mode=["messages","custom"]
                    ):
                        event_id = str(uuid.uuid4())[:8]  # Generate one UUID per event for internal tracking

                        # Handle different types of LangGraph events
                        if isinstance(item, tuple) and len(item) == 3:
                            node_name, event_type, event_data = item
                            
                            # Extract agent name from node
                            agent_name = str(node_name[0]) if node_name else 'unknown'
                            
                            if event_type == 'custom':
                                # Handle custom events (tool updates, agent messages) using utility
                                parsed_custom = LangGraphStreamParser.parse_custom_event(item)
                                
                                if parsed_custom:
                                    agent_name, message, metadata = parsed_custom
                                    
                                    # Check if this is a structured tool update from our wrapper
                                    if isinstance(metadata.get("raw_data"), dict) and "tool_name" in metadata["raw_data"]:
                                        # Use the structured data from our wrapper
                                        raw_data = metadata["raw_data"]
                                        tool_data = ToolUpdateData(
                                            tool_name=raw_data.get("tool_name", agent_name),
                                            message=raw_data.get("message", message),
                                            data=raw_data.get("data"),
                                            metadata=raw_data.get("metadata")
                                        )
                                    else:
                                        # Fallback to basic parsing
                                        tool_data = ToolUpdateData(
                                            tool_name=agent_name,
                                            message=message,
                                            data=metadata.get("raw_data"),
                                            metadata=metadata
                                        )
                                    
                                    # Use message ID manager for tool updates
                                    yield ToolUpdateEvent(
                                        tool_data=tool_data,
                                        message_id=message_id_manager.get_message_id_for_event("tool_update"),
                                        metadata=StreamingEventFormatter.clean_metadata(
                                            agent_name=agent_name,
                                            event_type="custom"
                                        )
                                    )
                                    
                                    # Tool updates are now handled by MessageIDManager automatically
                            
                            elif event_type == 'messages':
                                # Handle message chunks with time-based batching
                                parsed_data = LangGraphStreamParser.parse_message_chunk(item)

                                if parsed_data:
                                    agent_name, content, is_final_chunk, metadata = parsed_data
                                  
                                    # Check if content chunk should be sent
                                    if LangGraphStreamParser.is_content_chunk_worth_sending(content, is_final_chunk):
                                        # Get message ID from message ID manager
                                        message_id = message_id_manager.get_message_id_for_event("content_chunk")
                                        
                                        # Add content to buffer and check if we should send a batch
                                        batch_result = message_manager.add_content_chunk(
                                            agent_name, content, is_final_chunk
                                        )
                                        
                                        if batch_result:
                                            should_send, accumulated_content, is_progressive = batch_result
                                            
                                            if should_send:
                                                # Send the accumulated content with current message ID
                                                yield ContentChunkEvent(
                                                    content=accumulated_content,
                                                    agent_name=agent_name,
                                                    chunk_index=event_id,
                                                    is_final_chunk=is_final_chunk,
                                                    message_id=message_id,
                                                    metadata={
                                                        **metadata,
                                                        "is_progressive": is_progressive,
                                                        "batch_size": len(accumulated_content),
                                                        "original_chunk_size": len(content)
                                                    }
                                                )
                                                
                                                # Reset message group if this was the final chunk
                                                if is_final_chunk:
                                                    message_manager.reset_message_group()
                                            else:
                                                # Content is buffered, no send needed
                                                pass
                                    else:
                                        # Debug: Check if this was a ToolMessage that got filtered
                                        message_chunk = item[2][0]  # Extract the message chunk from the tuple
                                        if hasattr(message_chunk, 'name') and message_chunk.name:
                                            self.logger.info(f"🔧 ToolMessage '{message_chunk.name}' filtered out to prevent duplication")
                        else:
                            # Handle other types of events or log them
                            self.logger.debug(f"Unhandled event type: {type(item)} - {item}")
                    
                self.logger.info(f"=== STREAMING COMPLETED ===")
                
                # Get the final state using the graph's get_state method (async version)
                try:
                    final_state = await graph.aget_state(config={"configurable": {"thread_id": thread_id}})
                    self.logger.info(f"Final Graph State: {final_state.values}")
                except Exception as e:
                    self.logger.warning(f"Could not get final state: {e}")
               
        except Exception as e:
            execution_time = time.time() - start_time
            self.logger.error(f"Error in streaming request: {e}")
            
            # Use message ID manager for error events as well
            yield ErrorEvent(
                error=f"Streaming error: {str(e)}",
                details="An error occurred while processing your request",
                message_id=message_id_manager.get_message_id_for_event("error"),
                metadata=StreamingEventFormatter.clean_metadata(error_type=type(e).__name__)
            )
    
    

Langraph version -
Name: langgraph
Version: 0.6.4
Summary: Building stateful, multi-actor applications with LLMs
Home-page:
Author:
Author-email:
License-Expression: MIT
Location: /Users/adityavernekar/Documents/insightx/insightx-backend/venv/lib/python3.12/site-packages
Requires: langchain-core, langgraph-checkpoint, langgraph-prebuilt, langgraph-sdk, pydantic, xxhash
Required-by: insightx-backend, langgraph-supervisor, langmem, trustcall

So the graph is working but how i want the logic to work is it should stream the followup response , but now its streaming the entire json