output shown properly using pydantic

agent.py

@@ -1,5 +1,5 @@
 from llama_index.llms.openai import OpenAI
-from tools.tavily_search_tool import search_web, record_notes, write_report, review_report
+from tools.tavily_search_tool import search_web, record_notes, write_report, review_report, ReportOutput, get_structured_report_from_state
 from dotenv import load_dotenv
 import os
 from llama_index.core.agent.workflow import FunctionAgent
@@ -36,11 +36,13 @@ class TeacherStudentAgentWorkflow:
 
         self.write_agent = FunctionAgent(
             name="WriteAgent",
-            description="Useful for writing a report on a given topic.",
+            description="Useful for writing a structured report on a given topic.",
             system_prompt=(
-                "You are the WriteAgent that can write a report on a given topic. "
+                "You are the WriteAgent that can write a structured report on a given topic. "
                 "IMPORTANT: Never make duplicate tool calls. Write the report only ONCE with all available research. "
-                "Your report should be in a markdown format. The content should be grounded in the research notes. "
+                "Your report should be in markdown format and include a descriptive title. "
+                "When calling write_report, provide both the markdown content AND a clear, descriptive title. "
+                "The content should be grounded in the research notes and well-structured with clear sections. "
                 "Once the report is written ONCE, immediately hand off control to the ReviewAgent for feedback."
             ),
             llm=self.llm,
@@ -54,7 +56,9 @@ class TeacherStudentAgentWorkflow:
             system_prompt=(
                 "You are the ReviewAgent that can review the report and provide feedback. "
                 "IMPORTANT: Never make duplicate tool calls. Review the report only ONCE and provide clear feedback. "
-                "Your review should either approve the current report or request specific changes for the WriteAgent to implement. "
+                "Your review should either APPROVE the current report or request specific changes for the WriteAgent to implement. "
+                "When APPROVING a report, use clear approval language like 'APPROVED', 'READY', 'EXCELLENT', 'SATISFACTORY', or 'COMPLETE'. "
+                "If the report meets the requirements and is well-written, APPROVE it to complete the workflow. "
                 "If you have feedback that requires changes, hand off control to the WriteAgent to implement the changes after submitting the review ONCE."
             ),
             llm=self.llm,
@@ -72,6 +76,15 @@ class TeacherStudentAgentWorkflow:
             },
         )
 
+    async def get_structured_report(self, handler) -> ReportOutput:
+        """Extract the structured report from the workflow state."""
+        try:
+            final_state = await handler.ctx.get("state")
+            return get_structured_report_from_state(final_state)
+        except Exception as e:
+            print(f"Error getting structured report: {e}")
+            return None
+
     async def run_workflow(self, user_msg=None):
         if user_msg is None:
             user_msg = (

app.py

@@ -4,6 +4,7 @@ import asyncio
 import json
 import hashlib
 from agent import TeacherStudentAgentWorkflow
+from tools.tavily_search_tool import ReportOutput, get_structured_report_from_state
 from llama_index.core.agent.workflow import (
     AgentInput,
     AgentOutput,
@@ -11,6 +12,7 @@ from llama_index.core.agent.workflow import (
     ToolCallResult,
     AgentStream,
 )
+from datetime import datetime
 
 # Initialize the agent workflow
 agent_workflow = None
@@ -24,14 +26,15 @@ def get_agent_workflow():
 async def chat_with_agent(message, history):
     """
     Async chat function that runs the agent workflow and streams each step.
+    Returns structured report data for separate display.
     """
     if not message.strip():
-        yield history, None
+        yield history, None, gr.JSON(visible=False)
         return
     
     # Add user message to history
     history.append(ChatMessage(role="user", content=message))
-    yield history, None
+    yield history, None, gr.JSON(visible=False)
     
     try:
         # Get the agent workflow
@@ -43,6 +46,7 @@ async def chat_with_agent(message, history):
         current_agent = None
         current_step_messages = []
         final_report = None
+        structured_report_data = None
         workflow_state = {}
         
         # Track recent tool calls to prevent UI duplicates
@@ -67,7 +71,7 @@ async def chat_with_agent(message, history):
                     metadata={"title": f"Agent: {current_agent}"}
                 )
                 history.append(agent_header)
-                yield history, final_report
+                yield history, final_report, gr.JSON(visible=False)
             
             # Handle different event types
             if isinstance(event, AgentOutput):
@@ -79,7 +83,7 @@ async def chat_with_agent(message, history):
                         metadata={"title": f"{current_agent} - Output"}
                     )
                     history.append(output_msg)
-                    yield history, final_report
+                    yield history, final_report, gr.JSON(visible=False)
                 
                 if event.tool_calls:
                     # Show planned tools
@@ -90,7 +94,7 @@ async def chat_with_agent(message, history):
                         metadata={"title": f"{current_agent} - Tool Planning"}
                     )
                     history.append(tools_msg)
-                    yield history, final_report
+                    yield history, final_report, gr.JSON(visible=False)
                     
             elif isinstance(event, ToolCall):
                 # Create a unique identifier for this tool call using a more robust approach
@@ -118,7 +122,7 @@ async def chat_with_agent(message, history):
                         metadata={"title": f"{current_agent} - Tool Call"}
                     )
                     history.append(tool_msg)
-                    yield history, final_report
+                    yield history, final_report, gr.JSON(visible=False)
                 else:
                     # Debug: Log duplicate detection (remove this in production)
                     print(f"🚫 Duplicate tool call detected and skipped: {event.tool_name} with args {event.tool_kwargs}")
@@ -154,24 +158,62 @@ async def chat_with_agent(message, history):
                         workflow_state["has_report"] = True
                     elif event.tool_name == "review_report" and current_agent == "ReviewAgent":
                         workflow_state["has_review"] = True
-                        # Check if review indicates approval
-                        if any(word in result_content.lower() for word in ["approved", "ready", "good", "excellent"]):
+                        # Check if review indicates approval (expanded keywords)
+                        approval_keywords = [
+                            "approved", "ready", "good", "excellent", "satisfactory", 
+                            "complete", "accept", "final", "publish", "meets", "solid",
+                            "well-written", "comprehensive", "thorough"
+                        ]
+                        if any(word in result_content.lower() for word in approval_keywords):
                             workflow_state["review_approved"] = True
                 
-                yield history, final_report
+                yield history, final_report, gr.JSON(visible=False)
         
-        # Get the final state to extract the report
+        # Get the final state to extract the structured report
         try:
             final_state = await handler.ctx.get("state")
-            if final_state and "report_content" in final_state:
-                report_content = final_state["report_content"]
-                if report_content and report_content != "Not written yet.":
-                    # Show the final report if we have one and it's been reviewed
-                    if workflow_state.get("has_report") and workflow_state.get("has_review"):
-                        final_report = gr.Markdown(report_content, visible=True)
+            if final_state:
+                # Get structured report data
+                structured_report = get_structured_report_from_state(final_state)
+                if structured_report:
+                    # Include ALL fields from the Pydantic model, including content
+                    structured_report_data = {
+                        "title": structured_report.title,
+                        "abstract": structured_report.abstract,
+                        "content": structured_report.content,
+                        "sections": structured_report.sections,
+                        "word_count": structured_report.word_count,
+                        "generated_at": structured_report.generated_at.strftime("%Y-%m-%d %H:%M:%S"),
+                        "sources_used": structured_report.sources_used or []
+                    }
+                    
+                    # Show the final report if we have one
+                    # Check if review indicates approval OR if we just have a completed report
+                    if (workflow_state.get("has_report") and 
+                        (workflow_state.get("review_approved") or workflow_state.get("has_review"))):
+                        final_report = gr.Markdown(structured_report.content, visible=True)
                     elif workflow_state.get("has_report"):
-                        # Show report even if not reviewed yet
-                        final_report = gr.Markdown(report_content, visible=True)
+                        # Show report even if not reviewed yet, but mark it as preliminary
+                        final_report = gr.Markdown(
+                            f"## 📝 Preliminary Report (Pending Review)\n\n{structured_report.content}", 
+                            visible=True
+                        )
+                        
+                # Fallback to regular content if structured report is not available
+                elif "report_content" in final_state:
+                    report_content = final_state["report_content"]
+                    if report_content and report_content != "Not written yet.":
+                        # Create basic structured data from the raw content
+                        structured_report_data = {
+                            "title": "Generated Report",
+                            "content": report_content,
+                            "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                            "word_count": len(report_content.split()),
+                            "sources_used": list(final_state.get("research_notes", {}).keys()) if "research_notes" in final_state else []
+                        }
+                        
+                        if workflow_state.get("has_report"):
+                            final_report = gr.Markdown(report_content, visible=True)
         except Exception as state_error:
             print(f"Could not extract final state: {state_error}")
             # Try to show any report that was generated during the conversation
@@ -185,7 +227,19 @@ async def chat_with_agent(message, history):
             metadata={"title": "Workflow Complete"}
         )
         history.append(completion_msg)
-        yield history, final_report
+        
+        # Ensure we show the final report if we have structured data but no report was set
+        if structured_report_data and final_report is None:
+            if "content" in structured_report_data:
+                final_report = gr.Markdown(structured_report_data["content"], visible=True)
+        
+        # Create the structured report display component
+        structured_report_display = gr.JSON(
+            value=structured_report_data,
+            visible=bool(structured_report_data)
+        )
+        
+        yield history, final_report, structured_report_display
         
     except Exception as e:
         # Handle errors gracefully
@@ -195,12 +249,22 @@ async def chat_with_agent(message, history):
             metadata={"title": "Error"}
         )
         history.append(error_msg)
-        yield history, None
+        yield history, None, gr.JSON(visible=False)
 
 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="Teacher-Student Agent Workflow", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
@@ -235,6 +299,13 @@ with gr.Blocks(title="Teacher-Student Agent Workflow", theme=gr.themes.Soft()) a
         render=False
     )
     
+    # Create structured report metadata component
+    structured_report_json = gr.JSON(
+        label="📊 Report Metadata",
+        visible=False,
+        render=False
+    )
+    
     # Set up the chat interface with additional outputs
     chat_interface = gr.ChatInterface(
         fn=chat_with_agent,
@@ -254,7 +325,7 @@ with gr.Blocks(title="Teacher-Student Agent Workflow", theme=gr.themes.Soft()) a
             "Space Exploration Report"
         ],
         cache_examples=False,
-        additional_outputs=[final_report_output]
+        additional_outputs=[final_report_output, structured_report_json]
     )
     
     # Add feedback handling
@@ -262,9 +333,12 @@ with gr.Blocks(title="Teacher-Student Agent Workflow", theme=gr.themes.Soft()) a
     
     # Render the final report output in a separate section
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=2):
             gr.Markdown("### 📋 Final Report")
             final_report_output.render()
+        with gr.Column(scale=1):
+            gr.Markdown("### 📊 Report Metadata")
+            structured_report_json.render()
     
     gr.Markdown("""
     ### How it works:

test_structured_output.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""
+Test script to verify the Pydantic structured output functionality.
+"""
+
+import asyncio
+import sys
+import os
+
+# Add the current directory to the path so we can import modules
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+from tools.tavily_search_tool import ReportOutput, get_structured_report_from_state
+from datetime import datetime
+
+async def test_structured_output():
+    """Test the structured report output functionality."""
+    
+    print("🧪 Testing Pydantic ReportOutput model...")
+    
+    # Test creating a ReportOutput instance
+    test_report = ReportOutput(
+        title="Test Report",
+        abstract="This is a test abstract for our report.",
+        content="""# Test Report
+
+## Introduction
+This is the introduction section.
+
+## Main Content
+This is the main content section with some details.
+
+## Conclusion
+This is the conclusion section.
+""",
+        sections=["Introduction", "Main Content", "Conclusion"],
+        word_count=25,
+        sources_used=["Source 1", "Source 2"]
+    )
+    
+    print(f"✅ Created ReportOutput successfully!")
+    print(f"📋 Title: {test_report.title}")
+    print(f"📝 Abstract: {test_report.abstract}")
+    print(f"📊 Word Count: {test_report.word_count}")
+    print(f"🗂️ Sections: {test_report.sections}")
+    print(f"📚 Sources: {test_report.sources_used}")
+    print(f"⏰ Generated at: {test_report.generated_at}")
+    
+    # Test serialization
+    print("\n🔄 Testing serialization...")
+    serialized = test_report.model_dump()
+    print(f"✅ Serialized data keys: {list(serialized.keys())}")
+    
+    # Test state extraction
+    print("\n🔍 Testing state extraction...")
+    mock_state = {
+        "structured_report": serialized,
+        "other_data": "some value"
+    }
+    
+    extracted_report = get_structured_report_from_state(mock_state)
+    if extracted_report:
+        print(f"✅ Successfully extracted report from state!")
+        print(f"📋 Extracted title: {extracted_report.title}")
+        print(f"📊 Extracted word count: {extracted_report.word_count}")
+    else:
+        print("❌ Failed to extract report from state")
+    
+    # Test with empty state
+    empty_report = get_structured_report_from_state({})
+    if empty_report is None:
+        print("✅ Correctly returned None for empty state")
+    else:
+        print("❌ Should have returned None for empty state")
+    
+    print("\n🎉 All tests completed!")
+
+if __name__ == "__main__":
+    asyncio.run(test_structured_output()) 

tools/tavily_search_tool.py

@@ -5,6 +5,9 @@ import os
 import time
 import hashlib
 import json
+from pydantic import BaseModel, Field
+from typing import List, Optional
+from datetime import datetime
 
 load_dotenv(os.path.join(os.path.dirname(__file__), '../env.local'))
 
@@ -61,16 +64,58 @@ async def record_notes(ctx: Context, notes: str, notes_title: str) -> str:
     return "Notes recorded."
 
 
-async def write_report(ctx: Context, report_content: str) -> str:
-    """Useful for writing a report on a given topic. Your input should be a markdown formatted report."""
+class ReportOutput(BaseModel):
+    """Structured output for the writer agent's report."""
+    title: str = Field(description="The title of the report")
+    abstract: str = Field(description="A brief abstract or summary of the report")
+    content: str = Field(description="The full markdown content of the report")
+    sections: List[str] = Field(description="List of main section titles in the report")
+    word_count: int = Field(description="Approximate word count of the report")
+    generated_at: datetime = Field(default_factory=datetime.now, description="Timestamp when the report was generated")
+    sources_used: Optional[List[str]] = Field(default=None, description="List of sources or research notes used")
+
+async def write_report(ctx: Context, report_content: str, title: str = "Research Report") -> str:
+    """Useful for writing a report on a given topic. Your input should be a markdown formatted report with a title."""
     # Check for duplicate calls
-    if not _should_execute_call("write_report", report_content=report_content):
+    if not _should_execute_call("write_report", report_content=report_content, title=title):
         return "Duplicate report writing detected. Skipping to avoid redundant report generation."
     
     current_state = await ctx.get("state")
+    
+    # Extract sections from markdown content (look for ## headers)
+    import re
+    sections = re.findall(r'^#{1,3}\s+(.+)$', report_content, re.MULTILINE)
+    
+    # Calculate word count (approximate)
+    word_count = len(report_content.split())
+    
+    # Extract abstract (first paragraph after title)
+    lines = report_content.split('\n')
+    abstract = ""
+    for line in lines:
+        if line.strip() and not line.startswith('#'):
+            abstract = line.strip()
+            break
+    
+    # Get sources from research notes
+    sources_used = list(current_state.get("research_notes", {}).keys()) if "research_notes" in current_state else None
+    
+    # Create structured report output
+    structured_report = ReportOutput(
+        title=title,
+        abstract=abstract[:200] + "..." if len(abstract) > 200 else abstract,
+        content=report_content,
+        sections=sections,
+        word_count=word_count,
+        sources_used=sources_used
+    )
+    
+    # Store both the original content and structured output
     current_state["report_content"] = report_content
+    current_state["structured_report"] = structured_report.model_dump()
     await ctx.set("state", current_state)
-    return "Report written."
+    
+    return f"Report written successfully. Title: '{title}', Word count: {word_count}, Sections: {len(sections)}"
 
 
 async def review_report(ctx: Context, review: str) -> str:
@@ -82,4 +127,11 @@ async def review_report(ctx: Context, review: str) -> str:
     current_state = await ctx.get("state")
     current_state["review"] = review
     await ctx.set("state", current_state)
-    return "Report reviewed."
+    return "Report reviewed."
+
+
+def get_structured_report_from_state(state: dict) -> Optional[ReportOutput]:
+    """Helper function to extract structured report from workflow state."""
+    if "structured_report" in state:
+        return ReportOutput(**state["structured_report"])
+    return None