import gradio as gr from gradio import ChatMessage import asyncio import json import hashlib from datetime import datetime from agent import LlamaIndexReportAgent from tools.simple_tools import get_workflow_state from llama_index.core.agent.workflow import ( AgentInput, AgentOutput, ToolCall, ToolCallResult, AgentStream, ) from llama_index.core.workflow import Context # Initialize the agent workflow agent_workflow = None def get_agent_workflow(): global agent_workflow if agent_workflow is None: agent_workflow = LlamaIndexReportAgent() return agent_workflow async def chat_with_agent(message, history): """ Async chat function that runs the agent workflow and streams each step. """ history = history or [] history.append(ChatMessage(role="user", content=message)) # Initial yield to show user message immediately yield history, None, None, gr.update(value="", interactive=False) final_report_content = None structured_report_data = None displayed_tool_calls = set() try: workflow = get_agent_workflow() # Create context and initialize state properly ctx = Context(workflow.agent_workflow) await ctx.set("state", { "research_notes": {}, "report_content": "Not written yet.", "review": "Review required.", }) handler = workflow.agent_workflow.run(user_msg=message, ctx=ctx) current_agent = None async for event in handler.stream_events(): print(f"DEBUG: Event type: {type(event).__name__}") if hasattr(event, "current_agent_name") and event.current_agent_name != current_agent: current_agent = event.current_agent_name history.append(ChatMessage( role="assistant", content=f"**🤖 Agent: {current_agent}**", metadata={"title": f"Agent: {current_agent}"} )) yield history, final_report_content, structured_report_data, gr.update(interactive=False) if isinstance(event, ToolCall): tool_call_kwargs_str = json.dumps(getattr(event, 'tool_kwargs', {}), sort_keys=True) tool_call_key = f"{current_agent}:{event.tool_name}:{hashlib.md5(tool_call_kwargs_str.encode()).hexdigest()[:8]}" print(f"DEBUG: ToolCall detected - Agent: {current_agent}, Tool: {event.tool_name}, Args: {getattr(event, 'tool_kwargs', {})}") if tool_call_key not in displayed_tool_calls: args_preview = str(getattr(event, 'tool_kwargs', {}))[:100] + "..." if len(str(getattr(event, 'tool_kwargs', {}))) > 100 else str(getattr(event, 'tool_kwargs', {})) history.append(ChatMessage( role="assistant", content=f"**🔨 Calling Tool:** `{event.tool_name}`\n**Arguments:** {args_preview}", metadata={"title": f"{current_agent} - Tool Call"} )) displayed_tool_calls.add(tool_call_key) yield history, final_report_content, structured_report_data, gr.update(interactive=False) elif isinstance(event, ToolCallResult): print(f"DEBUG: ToolCallResult - Tool: {getattr(event, 'tool_name', 'unknown')}, Output: {getattr(event, 'tool_output', 'no output')}") # Show tool result in UI tool_output = getattr(event, 'tool_output', 'No output') tool_name = getattr(event, 'tool_name', 'unknown') output_preview = str(tool_output)[:200] + "..." if len(str(tool_output)) > 200 else str(tool_output) history.append(ChatMessage( role="assistant", content=f"**🔧 Tool Result ({tool_name}):**\n{output_preview}", metadata={"title": f"{current_agent} - Tool Result"} )) yield history, final_report_content, structured_report_data, gr.update(interactive=False) elif isinstance(event, AgentOutput) and event.response.content: print(f"DEBUG: AgentOutput from {current_agent}: {event.response.content}") # This is the agent's final thought or handoff message history.append(ChatMessage( role="assistant", content=f"**📤 Thought:** {event.response.content}", metadata={"title": f"{current_agent} - Output"} )) yield history, final_report_content, structured_report_data, gr.update(interactive=False) # Final state extraction - use the simple tools state print("DEBUG: Workflow completed, extracting final state...") final_state = get_workflow_state() print(f"DEBUG: Final state keys: {final_state.keys() if final_state else 'None'}") if final_state: print(f"DEBUG: Final state content: {json.dumps(final_state, indent=2, default=str)}") # Check for research notes research_notes = final_state.get("research_notes", {}) print(f"DEBUG: Research notes found: {len(research_notes)} items") for title, content in research_notes.items(): print(f"DEBUG: Research note '{title}': {content[:100]}..." if len(content) > 100 else f"DEBUG: Research note '{title}': {content}") # Check if we have a structured report if final_state.get("structured_report"): structured_report_data = final_state["structured_report"] final_report_content = structured_report_data.get("content", "*Report content not found in structured report.*") print(f"DEBUG: Found structured report with content length: {len(final_report_content) if final_report_content else 0}") else: # Fallback: try to get report_content directly from state final_report_content = final_state.get("report_content", None) if final_report_content and final_report_content != "Not written yet.": print(f"DEBUG: Found report_content directly in state with length: {len(final_report_content)}") # Create minimal structured data for JSON display structured_report_data = { "title": "Generated Report", "content": final_report_content, "word_count": len(final_report_content.split()), "generated_at": datetime.now().isoformat(), "research_notes_count": len(final_state.get("research_notes", {})) } else: print("DEBUG: No valid report content found in final state") print(f"DEBUG: report_content value: '{final_report_content}'") # If we have research notes but no report, show that as partial success if research_notes: final_report_content = f"**Research completed but report not written.**\n\n**Research Notes:**\n\n" for title, content in research_notes.items(): final_report_content += f"### {title}\n{content}\n\n" structured_report_data = { "title": "Research Notes (Report Incomplete)", "content": final_report_content, "word_count": len(final_report_content.split()), "generated_at": datetime.now().isoformat(), "research_notes_count": len(research_notes), "status": "incomplete" } print(f"DEBUG: Created fallback report from research notes") else: final_report_content = None structured_report_data = None else: print("DEBUG: No final state retrieved") final_report_content = None structured_report_data = None history.append(ChatMessage( role="assistant", content="✅ **Workflow completed!**", metadata={"title": "Workflow Complete"} )) if final_report_content: final_report_update = gr.update(value=final_report_content, visible=True) json_report_update = gr.update(value=structured_report_data, visible=True) if structured_report_data else gr.update(visible=False) else: final_report_update = gr.update(value="*No final report was generated. Check the workflow execution above.*", visible=True) json_report_update = gr.update(visible=False) yield history, final_report_update, json_report_update, gr.update(interactive=True, placeholder="Enter your next request...") except Exception as e: print(f"ERROR in chat_with_agent: {e}") import traceback traceback.print_exc() history.append(ChatMessage(role="assistant", content=f"❌ **Error:** {str(e)}", metadata={"title": "Error"})) yield history, gr.update(visible=False), gr.update(visible=False), gr.update(interactive=True) def like_feedback(evt: gr.LikeData): """Handle user feedback on messages.""" print(f"User feedback - Index: {evt.index}, Liked: {evt.liked}, Value: {evt.value}") def format_structured_report_display(structured_report_data): """Format structured report data for JSON display component.""" if not structured_report_data: return gr.JSON(visible=False) return gr.JSON( value=structured_report_data, visible=True ) # Create the Gradio interface with gr.Blocks(title="LlamaIndex Report Generation Agent", theme=gr.themes.Soft()) as demo: gr.Markdown(""" # 🤖 LlamaIndex Report Generation Agent A multi-agent workflow built with LlamaIndex that uses teacher-student methodology to generate comprehensive reports. The system employs three specialized agents that collaborate step by step: - **ResearchAgent**: Searches the web and records research notes - **WriteAgent**: Creates structured reports based on research findings - **ReviewAgent**: Reviews reports and provides iterative feedback for improvement Enter any topic below to see the LlamaIndex agents collaborate using teacher-student methodology! """) chatbot = gr.Chatbot( label="Agent Workflow", type="messages", height=600, show_copy_button=True, placeholder="Ask me to write a report on any topic...", render_markdown=True ) with gr.Row(): textbox = gr.Textbox( placeholder="Enter your request...", container=False, scale=7 ) submit_btn = gr.Button("Submit", variant="primary", scale=1) with gr.Row(): with gr.Column(scale=2): final_report_output = gr.Textbox( label="📄 Final Report", interactive=False, lines=20, show_copy_button=True, visible=False ) with gr.Column(scale=1): structured_report_json = gr.JSON(label="📊 Report Metadata", visible=False) gr.Examples( examples=[ "Write a report on the history of artificial intelligence", "Create a report about renewable energy technologies", "Write a report on the impact of social media on society", ], inputs=textbox, ) gr.Markdown(""" ### How the LlamaIndex Teacher-Student Agent Works: 1. **ResearchAgent** searches for information and takes comprehensive notes 2. **WriteAgent** creates a structured report based on the research findings 3. **ReviewAgent** reviews the report and provides constructive feedback 4. The process iterates until the report meets quality standards Watch the real-time collaboration between LlamaIndex agents as they employ teacher-student methodology! """) # Event handlers submit_btn.click( chat_with_agent, inputs=[textbox, chatbot], outputs=[chatbot, final_report_output, structured_report_json, textbox], queue=True ) textbox.submit( chat_with_agent, inputs=[textbox, chatbot], outputs=[chatbot, final_report_output, structured_report_json, textbox], queue=True ) if __name__ == "__main__": demo.launch()