Merge pull request #1041 from IlumCI/format
[FEAT-MD_OUTPUT][Added optional markdown output with proper table and code block support]pull/1051/head
Kye Gomez
1 week ago
committed by
GitHub
commit
e8f161beea
@ -0,0 +1,519 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script demonstrating markdown output functionality with a real swarm
|
||||
Uses the current state of formatter.py to show agent markdown output capabilities
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
from typing import List, Dict, Any
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
from swarms import Agent
|
||||
from swarms.structs import SequentialWorkflow, ConcurrentWorkflow, GroupChat
|
||||
from swarms.utils.formatter import Formatter
|
||||
|
||||
class MarkdownTestSwarm:
|
||||
"""A test swarm that demonstrates markdown output capabilities"""
|
||||
|
||||
def __init__(self):
|
||||
self.formatter = Formatter(markdown=True)
|
||||
self.setup_agents()
|
||||
self.setup_swarm()
|
||||
|
||||
def setup_agents(self):
|
||||
"""Setup specialized agents for markdown testing"""
|
||||
|
||||
# Research Agent - Generates structured markdown reports
|
||||
self.research_agent = Agent(
|
||||
agent_name="Research Agent",
|
||||
system_prompt="""You are a research specialist. When given a topic, create a comprehensive markdown report with:
|
||||
- Clear headers and subheaders
|
||||
- Code examples when relevant
|
||||
- Bullet points and numbered lists
|
||||
- Bold and italic text for emphasis
|
||||
- Tables for data comparison
|
||||
- Code blocks with syntax highlighting
|
||||
|
||||
Always format your response as clean markdown with proper structure.""",
|
||||
model_name="gpt-4o-mini", # Use a more capable model
|
||||
temperature=0.7,
|
||||
max_tokens=4000,
|
||||
max_loops=1,
|
||||
context_length=8000, # Limit context to prevent overflow
|
||||
return_history=False # Don't return history to reduce context
|
||||
)
|
||||
|
||||
# Code Analysis Agent - Generates code-heavy markdown
|
||||
self.code_agent = Agent(
|
||||
agent_name="Code Analysis Agent",
|
||||
system_prompt="""You are a code analysis specialist. When given code or programming concepts, create markdown documentation with:
|
||||
- Syntax-highlighted code blocks
|
||||
- Function documentation
|
||||
- Code examples
|
||||
- Performance analysis
|
||||
- Best practices
|
||||
|
||||
Use proper markdown formatting with code blocks, inline code, and structured content.""",
|
||||
model_name="gpt-4o-mini", # Use a more capable model
|
||||
temperature=0.5,
|
||||
max_tokens=4000,
|
||||
max_loops=1,
|
||||
context_length=8000, # Limit context to prevent overflow
|
||||
return_history=False # Don't return history to reduce context
|
||||
)
|
||||
|
||||
# Data Visualization Agent - Creates data-focused markdown
|
||||
self.data_agent = Agent(
|
||||
agent_name="Data Visualization Agent",
|
||||
system_prompt="""You are a data visualization specialist. When given data or analysis requests, create markdown reports with:
|
||||
- Data tables
|
||||
- Statistical analysis
|
||||
- Charts and graphs descriptions
|
||||
- Key insights with bold formatting
|
||||
- Recommendations in structured lists
|
||||
|
||||
Format everything as clean, readable markdown.""",
|
||||
model_name="gpt-4o-mini", # Use a more capable model
|
||||
temperature=0.6,
|
||||
max_tokens=4000,
|
||||
max_loops=1,
|
||||
context_length=8000, # Limit context to prevent overflow
|
||||
return_history=False # Don't return history to reduce context
|
||||
)
|
||||
|
||||
def setup_swarm(self):
|
||||
"""Setup the swarm with the agents"""
|
||||
# Create different swarm types for testing
|
||||
self.sequential_swarm = SequentialWorkflow(
|
||||
name="Markdown Test Sequential",
|
||||
description="Sequential workflow for markdown testing",
|
||||
agents=[self.research_agent, self.code_agent, self.data_agent],
|
||||
max_loops=1 # Reduce loops to prevent context overflow
|
||||
)
|
||||
|
||||
self.concurrent_swarm = ConcurrentWorkflow(
|
||||
name="Markdown Test Concurrent",
|
||||
description="Concurrent workflow for markdown testing",
|
||||
agents=[self.research_agent, self.code_agent, self.data_agent],
|
||||
max_loops=1 # Reduce loops to prevent context overflow
|
||||
)
|
||||
|
||||
self.groupchat_swarm = GroupChat(
|
||||
name="Markdown Test Group Chat",
|
||||
description="A group chat for testing markdown output",
|
||||
agents=[self.research_agent, self.code_agent, self.data_agent],
|
||||
max_loops=1 # Reduce loops to prevent context overflow
|
||||
)
|
||||
|
||||
# Default swarm for main tests
|
||||
self.swarm = self.sequential_swarm
|
||||
|
||||
def test_basic_markdown_output(self):
|
||||
"""Test basic markdown output with a simple topic"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 1: Basic Markdown Output")
|
||||
print("="*60)
|
||||
|
||||
topic = "Python Web Development with FastAPI"
|
||||
|
||||
self.formatter.print_panel(
|
||||
f"Starting research on: {topic}",
|
||||
title="Research Topic",
|
||||
style="bold blue"
|
||||
)
|
||||
|
||||
# Run the research agent
|
||||
result = self.research_agent.run(topic)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result,
|
||||
title="Research Report",
|
||||
border_style="green"
|
||||
)
|
||||
|
||||
def test_code_analysis_markdown(self):
|
||||
"""Test markdown output with code analysis"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 2: Code Analysis Markdown")
|
||||
print("="*60)
|
||||
|
||||
code_sample = """
|
||||
def fibonacci(n):
|
||||
if n <= 1:
|
||||
return n
|
||||
return fibonacci(n-1) + fibonacci(n-2)
|
||||
|
||||
# Test the function
|
||||
for i in range(10):
|
||||
print(fibonacci(i))
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Analyzing Python code sample",
|
||||
title="Code Analysis",
|
||||
style="bold cyan"
|
||||
)
|
||||
|
||||
# Run the code analysis agent
|
||||
result = self.code_agent.run(f"Analyze this Python code and provide improvements:\n\n{code_sample}")
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result,
|
||||
title="Code Analysis Report",
|
||||
border_style="yellow"
|
||||
)
|
||||
|
||||
def test_data_analysis_markdown(self):
|
||||
"""Test markdown output with data analysis"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 3: Data Analysis Markdown")
|
||||
print("="*60)
|
||||
|
||||
data_request = """
|
||||
Analyze the following dataset:
|
||||
- Sales: $1.2M (Q1), $1.5M (Q2), $1.8M (Q3), $2.1M (Q4)
|
||||
- Growth Rate: 8%, 12%, 15%, 18%
|
||||
- Customer Count: 1000, 1200, 1400, 1600
|
||||
|
||||
Provide insights and recommendations in markdown format.
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Analyzing quarterly business data",
|
||||
title="Data Analysis",
|
||||
style="bold magenta"
|
||||
)
|
||||
|
||||
# Run the data analysis agent
|
||||
result = self.data_agent.run(data_request)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result,
|
||||
title="Data Analysis Report",
|
||||
border_style="red"
|
||||
)
|
||||
|
||||
def test_swarm_collaboration_markdown(self):
|
||||
"""Test markdown output with swarm collaboration"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 4: Swarm Collaboration Markdown")
|
||||
print("="*60)
|
||||
|
||||
complex_topic = """
|
||||
Create a comprehensive guide on building a machine learning pipeline that includes:
|
||||
1. Data preprocessing techniques
|
||||
2. Model selection strategies
|
||||
3. Performance evaluation metrics
|
||||
4. Deployment considerations
|
||||
|
||||
Each agent should contribute their expertise and the final output should be well-formatted markdown.
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Swarm collaboration on ML pipeline guide",
|
||||
title="Swarm Task",
|
||||
style="bold green"
|
||||
)
|
||||
|
||||
# Run the swarm
|
||||
results = self.swarm.run(complex_topic)
|
||||
|
||||
# Display individual agent results
|
||||
# SequentialWorkflow returns a list of results, not a dict
|
||||
for i, result in enumerate(results, 1):
|
||||
agent_name = f"Agent {i}"
|
||||
|
||||
# Handle different result types
|
||||
if isinstance(result, dict):
|
||||
# Extract the output from dict result
|
||||
result_content = result.get('output', str(result))
|
||||
else:
|
||||
result_content = str(result)
|
||||
self.formatter.print_markdown(
|
||||
result_content,
|
||||
title=f"Agent {i}: {agent_name}",
|
||||
border_style="blue"
|
||||
)
|
||||
|
||||
def test_markdown_toggle_functionality(self):
|
||||
"""Test the markdown enable/disable functionality"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 5: Markdown Toggle Functionality")
|
||||
print("="*60)
|
||||
|
||||
test_content = """
|
||||
# Test Content
|
||||
|
||||
This is a **bold** test with `inline code`.
|
||||
|
||||
## Code Block
|
||||
```python
|
||||
def test_function():
|
||||
return "Hello, World!"
|
||||
```
|
||||
|
||||
## List
|
||||
- Item 1
|
||||
- Item 2
|
||||
- Item 3
|
||||
"""
|
||||
|
||||
# Test with markdown enabled
|
||||
self.formatter.print_panel(
|
||||
"Testing with markdown ENABLED",
|
||||
title="Markdown Enabled",
|
||||
style="bold green"
|
||||
)
|
||||
self.formatter.print_markdown(test_content, "Markdown Output")
|
||||
|
||||
# Disable markdown
|
||||
self.formatter.disable_markdown()
|
||||
self.formatter.print_panel(
|
||||
"Testing with markdown DISABLED",
|
||||
title="Markdown Disabled",
|
||||
style="bold red"
|
||||
)
|
||||
self.formatter.print_panel(test_content, "Plain Text Output")
|
||||
|
||||
# Re-enable markdown
|
||||
self.formatter.enable_markdown()
|
||||
self.formatter.print_panel(
|
||||
"Testing with markdown RE-ENABLED",
|
||||
title="Markdown Re-enabled",
|
||||
style="bold blue"
|
||||
)
|
||||
self.formatter.print_markdown(test_content, "Markdown Output Again")
|
||||
|
||||
def test_different_swarm_types(self):
|
||||
"""Test markdown output with different swarm types"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 6: Different Swarm Types")
|
||||
print("="*60)
|
||||
|
||||
simple_topic = "Explain the benefits of using Python for data science"
|
||||
|
||||
# Test Sequential Workflow
|
||||
print("\n--- Sequential Workflow ---")
|
||||
self.formatter.print_panel(
|
||||
"Testing Sequential Workflow (agents work in sequence)",
|
||||
title="Swarm Type Test",
|
||||
style="bold blue"
|
||||
)
|
||||
sequential_results = self.sequential_swarm.run(simple_topic)
|
||||
for i, result in enumerate(sequential_results, 1):
|
||||
# Handle different result types
|
||||
if isinstance(result, dict):
|
||||
result_content = result.get('output', str(result))
|
||||
else:
|
||||
result_content = str(result)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result_content,
|
||||
title=f"Sequential Agent {i}",
|
||||
border_style="blue"
|
||||
)
|
||||
|
||||
# Test Concurrent Workflow
|
||||
print("\n--- Concurrent Workflow ---")
|
||||
self.formatter.print_panel(
|
||||
"Testing Concurrent Workflow (agents work in parallel)",
|
||||
title="Swarm Type Test",
|
||||
style="bold green"
|
||||
)
|
||||
concurrent_results = self.concurrent_swarm.run(simple_topic)
|
||||
for i, result in enumerate(concurrent_results, 1):
|
||||
# Handle different result types
|
||||
if isinstance(result, dict):
|
||||
result_content = result.get('output', str(result))
|
||||
else:
|
||||
result_content = str(result)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result_content,
|
||||
title=f"Concurrent Agent {i}",
|
||||
border_style="green"
|
||||
)
|
||||
|
||||
# Test Group Chat
|
||||
print("\n--- Group Chat ---")
|
||||
self.formatter.print_panel(
|
||||
"Testing Group Chat (agents collaborate in conversation)",
|
||||
title="Swarm Type Test",
|
||||
style="bold magenta"
|
||||
)
|
||||
groupchat_results = self.groupchat_swarm.run(simple_topic)
|
||||
|
||||
# Handle different result types for GroupChat
|
||||
if isinstance(groupchat_results, dict):
|
||||
result_content = groupchat_results.get('output', str(groupchat_results))
|
||||
else:
|
||||
result_content = str(groupchat_results)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
result_content,
|
||||
title="Group Chat Result",
|
||||
border_style="magenta"
|
||||
)
|
||||
|
||||
def test_simple_formatter_only(self):
|
||||
"""Test just the formatter functionality without agents"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 7: Simple Formatter Test (No Agents)")
|
||||
print("="*60)
|
||||
|
||||
# Test basic markdown rendering
|
||||
simple_markdown = """
|
||||
# Simple Test
|
||||
|
||||
This is a **bold** test with `inline code`.
|
||||
|
||||
## Code Block
|
||||
```python
|
||||
def hello_world():
|
||||
print("Hello, World!")
|
||||
return "Success"
|
||||
```
|
||||
|
||||
## List
|
||||
- Item 1
|
||||
- Item 2
|
||||
- Item 3
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Testing formatter without agents",
|
||||
title="Formatter Test",
|
||||
style="bold cyan"
|
||||
)
|
||||
|
||||
self.formatter.print_markdown(
|
||||
simple_markdown,
|
||||
title="Simple Markdown Test",
|
||||
border_style="green"
|
||||
)
|
||||
|
||||
# Test toggle functionality
|
||||
self.formatter.disable_markdown()
|
||||
self.formatter.print_panel(
|
||||
"Markdown disabled - this should be plain text",
|
||||
title="Plain Text Test",
|
||||
style="bold red"
|
||||
)
|
||||
self.formatter.enable_markdown()
|
||||
|
||||
def test_error_handling_markdown(self):
|
||||
"""Test markdown output with error handling"""
|
||||
print("\n" + "="*60)
|
||||
print("TEST 8: Error Handling in Markdown")
|
||||
print("="*60)
|
||||
|
||||
# Test with malformed markdown
|
||||
malformed_content = """
|
||||
# Incomplete header
|
||||
**Unclosed bold
|
||||
```python
|
||||
def incomplete_code():
|
||||
# Missing closing backticks
|
||||
"""
|
||||
|
||||
self.formatter.print_panel(
|
||||
"Testing error handling with malformed markdown",
|
||||
title="Error Handling Test",
|
||||
style="bold yellow"
|
||||
)
|
||||
|
||||
# This should handle the error gracefully
|
||||
self.formatter.print_markdown(
|
||||
malformed_content,
|
||||
title="Malformed Markdown Test",
|
||||
border_style="yellow"
|
||||
)
|
||||
|
||||
# Test with empty content
|
||||
self.formatter.print_markdown(
|
||||
"",
|
||||
title="Empty Content Test",
|
||||
border_style="cyan"
|
||||
)
|
||||
|
||||
# Test with None content
|
||||
self.formatter.print_markdown(
|
||||
None,
|
||||
title="None Content Test",
|
||||
border_style="magenta"
|
||||
)
|
||||
|
||||
def run_all_tests(self):
|
||||
"""Run all markdown output tests"""
|
||||
print(" Starting Swarm Markdown Output Tests")
|
||||
print("="*60)
|
||||
|
||||
try:
|
||||
# Test 1: Basic markdown output
|
||||
self.test_basic_markdown_output()
|
||||
|
||||
# Test 2: Code analysis markdown
|
||||
self.test_code_analysis_markdown()
|
||||
|
||||
# Test 3: Data analysis markdown
|
||||
self.test_data_analysis_markdown()
|
||||
|
||||
# Test 4: Swarm collaboration
|
||||
self.test_swarm_collaboration_markdown()
|
||||
|
||||
# Test 5: Markdown toggle functionality
|
||||
self.test_markdown_toggle_functionality()
|
||||
|
||||
# Test 6: Different swarm types
|
||||
self.test_different_swarm_types()
|
||||
|
||||
# Test 7: Simple formatter test (no agents)
|
||||
self.test_simple_formatter_only()
|
||||
|
||||
# Test 8: Error handling
|
||||
self.test_error_handling_markdown()
|
||||
|
||||
print("\n" + "="*60)
|
||||
print(" All tests completed successfully!")
|
||||
print("="*60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n Test failed with error: {str(e)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
def main():
|
||||
"""Main function to run the markdown output tests"""
|
||||
print("Swarms Markdown Output Test Suite")
|
||||
print("Testing the current state of formatter.py with real swarm agents")
|
||||
print("="*60)
|
||||
|
||||
# Check environment setup
|
||||
api_key = os.getenv("OPENAI_API_KEY") or os.getenv("SWARMS_API_KEY")
|
||||
if not api_key:
|
||||
print("⚠ Warning: No API key found. Please set OPENAI_API_KEY or SWARMS_API_KEY environment variable.")
|
||||
print(" You can create a .env file with: OPENAI_API_KEY=your_api_key_here")
|
||||
print(" Or set it in your environment: export OPENAI_API_KEY=your_api_key_here")
|
||||
print()
|
||||
|
||||
try:
|
||||
# Create and run the test swarm
|
||||
test_swarm = MarkdownTestSwarm()
|
||||
test_swarm.run_all_tests()
|
||||
except Exception as e:
|
||||
print(f"\n Test failed with error: {str(e)}")
|
||||
print("\n Troubleshooting tips:")
|
||||
print("1. Make sure you have set your API key (OPENAI_API_KEY or SWARMS_API_KEY)")
|
||||
print("2. Check your internet connection")
|
||||
print("3. Verify you have sufficient API credits")
|
||||
print("4. Try running with a simpler test first")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in new issue