logger = setup_logger() async def generate_sse_event(event: str, data: Any, step: str = None) -> str: """Generate SSE formatted message with timestamp and consistent structure.""" timestamp = time.time() event_data = { "event": event, "step": step or event, "timestamp": timestamp, "status": "success" } if isinstance(data, (dict, list)): event_data["data"] = data else: event_data["data"] = str(data) if event == "error": event_data["status"] = "error" event_data["error_details"] = { "message": str(data), "type": type(data).__name__ } return f"data: {json.dumps(event_data)}\n\n" async def rag_node(state: AgentState) -> AgentState: """Wrapper function for RAG processing with proper state handling.""" result = await rag_tool.ainvoke({"state": state.model_dump()}) return AgentState(**{**state.model_dump(), **result}) async def stream_rag_events(state: AgentState) -> AsyncGenerator[str, None]: """Stream RAG events as SSE.""" try: yield await generate_sse_event("rag_start", { "message": "Starting RAG processing", "query": state.query }) # Initial processing event yield await generate_sse_event("status", {"message": "Starting RAG processing", "query": state.query}) # Cache check event yield await generate_sse_event("cache_check", {"message": "Checking cache for query"}) # Get chat history if available if state.session_id: yield await generate_sse_event("history_retrieval", { "message": f"Retrieving chat history for session {state.session_id}" }) # Personalization check yield await generate_sse_event("personalization_check", { "is_personalized": state.is_personalized, "entity_id": state.entity_id }) # Query refinement yield await generate_sse_event("query_refinement_start", {"message": "Starting query refinement"}) # Parallel processing events yield await generate_sse_event("parallel_processing", { "message": "Starting parallel processing of intent analysis and embeddings" }) # Process through RAG pipeline result = await rag_node(state) result_dict = result.model_dump() # Convert AgentState to dictionary # Emit events for each major step yield await generate_sse_event("query_refinement_complete", { "refined_query": result_dict.get("refined_query", state.query), "original_query": state.query, "personalization_details": result_dict.get("personalization_details", {}) }) yield await generate_sse_event("intent_analysis_complete", { "intents": result_dict.get("intent_data", []), "entities": result_dict.get("entity_data", []) }) yield await generate_sse_event("embedding_generation_complete", { "message": "Embeddings generated successfully" }) yield await generate_sse_event("search_execution", { "filter_expression": result_dict.get("search_filters"), "message": "Executing vector search" }) yield await generate_sse_event("search_results", { "count": len(result_dict.get("search_results", [])), "top_results": result_dict.get("search_results", [])[:3] }) yield await generate_sse_event("response_generation", { "message": "Generating final response" }) # Final result with all metadata yield await generate_sse_event("completion", { "response": result_dict.get("response", {}).get("message"), "final_response": { "session_id": result_dict.get("session_id"), "processing_time": result_dict.get("processing_details", {}).get("processing_time"), "refined_query": result_dict.get("refined_query"), "search_results_count": len(result_dict.get("search_results", [])) } }) yield await generate_sse_event("rag_complete", { "message": "RAG processing completed", "response": result.response, "session_id": result.session_id }) except Exception as e: logger.error(f"RAG processing error: {str(e)}", exc_info=True) yield await generate_sse_event("error", { "message": f"RAG processing failed: {str(e)}", "stage": "rag_processing" }) class EnhancedStateGraph(StateGraph): async def ainvoke(self, state: AgentState) -> AsyncGenerator[str, None]: """Execute the graph with SSE streaming.""" try: yield await generate_sse_event("connected", "Connection established") # Track processed nodes to prevent duplicates processed_nodes = set() async for node_name, node_result in self.astream(state): if node_name == "supervisor": yield await generate_sse_event("supervisor_decision", { "next_step": node_result.next_step, "reason": getattr(node_result, "reason", "No reason provided") }) elif node_name == "small_talk" and "small_talk" not in processed_nodes: processed_nodes.add("small_talk") yield await generate_sse_event("small_talk_start", { "message": "Processing small talk request" }) # Get the final response from small talk agent response = await small_talk_agent_node(state) yield await generate_sse_event("small_talk_response", { "response": response, "session_id": state.session_id }) elif node_name == "rag" and "rag" not in processed_nodes: processed_nodes.add("rag") async for event in stream_rag_events(state): yield event yield await generate_sse_event("workflow_complete", { "message": "Workflow completed successfully", "processed_nodes": list(processed_nodes) }) except Exception as e: logger.error(f"Workflow error: {str(e)}") yield await generate_sse_event("error", { "message": f"Workflow failed: {str(e)}", "stage": "workflow_execution" }) def create_main_flow() -> EnhancedStateGraph: """Create the workflow graph with proper node connections.""" builder = EnhancedStateGraph(AgentState) # Add nodes builder.add_node("supervisor", supervisor_node) builder.add_node("small_talk", lambda state: state) # Pass-through, actual processing in ainvoke builder.add_node("rag", lambda state: state) # Pass-through, actual processing in ainvoke # Set entry point builder.set_entry_point("supervisor") # Add conditional edges builder.add_conditional_edges( "supervisor", lambda state: state.next_step, { "small_talk": "small_talk", "rag": "rag" } ) # Add edges to END builder.add_edge("small_talk", END) builder.add_edge("rag", END) return builder.compile() async def execute_with_sse(state: AgentState) -> AsyncGenerator[str, None]: """Main entry point for executing the workflow with SSE.""" workflow = create_main_flow() async for event in workflow.ainvoke(state): yield event Continue reading...
