import os import asyncio from textwrap import dedent from dotenv import load_dotenv from crewai import Agent, Crew, Process, Task from langchain_openai import ChatOpenAI import gradio as gr # Load the .env file to get the API key load_dotenv() # Verify API key exists if not os.getenv("OPENAI_API_KEY"): raise ValueError("OPENAI_API_KEY not found in environment variables. Please add it to your .env file.") # Initialize LLM llm = ChatOpenAI( model="gpt-4-turbo-preview", temperature=0.1, api_key=os.getenv("OPENAI_API_KEY") ) # --- 1. Agent Definitions --- lead_architect = Agent( role="Lead Network Architect", goal=dedent(""" Translate high-level business and application needs into a detailed list of specific, actionable technical requirements for a new data center network. """), backstory=dedent(""" You are a seasoned Lead Architect who excels at interfacing with business stakeholders. Your primary skill is listening to business goals and application needs, then distilling them into a clear and comprehensive list of technical specifications that engineers can build from. You think in terms of availability, security, scalability, and performance metrics. """), llm=llm, allow_delegation=False, verbose=True, ) senior_network_architect = Agent( role="Senior Network Architect", goal=dedent(""" Take a detailed list of technical requirements and create three distinct, high-level data center network designs. Each design must be for a different major vendor (Cisco, Arista, Juniper) and must meet all specified technical requirements. """), backstory=dedent(""" You are a hands-on Senior Network Architect with deep product knowledge across Cisco, Arista, and Juniper. You are an expert in creating detailed High-Level Designs (HLDs) based on a given set of technical specifications. You focus purely on the technical implementation. """), llm=llm, allow_delegation=False, verbose=True, ) peer_review_architect = Agent( role="Peer Review Architect", goal=dedent(""" Critically review three network designs against the original technical requirements to validate their feasibility, compare their trade-offs, and select the optimal solution. """), backstory=dedent(""" You are an impartial and meticulous architect responsible for quality control. Your job is to act as a peer reviewer. You take a set of designs and the original requirements, and your task is to find any gaps, inconsistencies, or trade-offs. You provide the final validation and reasoned recommendation on which design to proceed with. """), llm=llm, allow_delegation=False, verbose=True, ) # --- 2. Task Definitions --- def create_tasks(business_requirements, application_requirements): """Create tasks with proper input context""" initial_requirements = f"Business Requirements: {business_requirements}\nApplication Requirements: {application_requirements}" technical_translation_task = Task( description=dedent(f""" Analyze the following business and application requirements and translate them into a detailed list of technical network requirements. Create a comprehensive table with three columns: 'Category', 'Requirement', and 'Justification'. Categories should include: Availability, Performance, Security, Scalability, Management, and Budget. Requirements to analyze: {initial_requirements} Focus on creating specific, measurable technical requirements that can guide network design decisions. """), expected_output="A detailed markdown table with three columns: 'Category', 'Requirement', and 'Justification'.", agent=lead_architect, ) design_creation_task = Task( description=dedent(""" Based on the technical requirements from the Lead Architect, create three distinct high-level data center designs using a spine-leaf architecture. **CRITICAL:** Prioritize the original business requirements when making design trade-offs. If 'Strict Budget' was selected, choose cost-effective hardware even if it means compromising on supporting requirements. Create one design for each vendor: - Cisco design with specific switch models and NX-OS - Juniper design with specific switch models and Junos - Arista design with specific switch models and EOS For each design, specify: 1. Recommended spine switch models and quantities 2. Recommended leaf switch models and quantities 3. Network operating system version 4. Brief rationale explaining how it meets the technical requirements 5. Estimated cost range if budget was a consideration """), expected_output="Three complete high-level designs, one for each vendor (Cisco, Juniper, Arista).", agent=senior_network_architect, context=[technical_translation_task] ) validation_task = Task( description=dedent(""" Review the technical requirements and the three vendor designs (Cisco, Juniper, Arista) provided by previous team members. Create a comprehensive validation report with these sections: 1. **Requirements Validation**: Check if each vendor design meets all technical requirements 2. **Comparative Analysis**: Compare pros/cons of Cisco vs Juniper vs Arista solutions 3. **Trade-off Analysis**: Identify any compromises made in each design 4. **Final Recommendation**: Select the best vendor solution with clear justification 5. **Implementation Considerations**: Next steps and potential risks Base your analysis solely on the specific vendor designs provided, not generic concepts. """), expected_output="A comprehensive validation report with requirements check, vendor comparison, and final recommendation.", agent=peer_review_architect, context=[technical_translation_task, design_creation_task] ) return [technical_translation_task, design_creation_task, validation_task] # --- 3. CrewAI Execution Function --- def run_design_crew(business_reqs_list, app_reqs_list, progress=gr.Progress()): """Run the CrewAI design crew with proper execution flow""" # Input validation if not business_reqs_list and not app_reqs_list: return ( "❌ **Error**: Please select at least one requirement.", "", "" ) # Verify API key if not os.getenv("OPENAI_API_KEY"): return ( "❌ **Error**: OPENAI_API_KEY not found. Please add it to your .env file.", "", "" ) try: # Prepare requirements strings business_reqs_str = ", ".join(business_reqs_list) if business_reqs_list else "None specified" app_reqs_str = ", ".join(app_reqs_list) if app_reqs_list else "None specified" progress(0.1, desc="Initializing AI agents...") # Create tasks tasks = create_tasks(business_reqs_str, app_reqs_str) # Create crew crew = Crew( agents=[lead_architect, senior_network_architect, peer_review_architect], tasks=tasks, process=Process.sequential, verbose=True, memory=True, # Enable memory for better context retention ) progress(0.2, desc="Starting design process...") # Prepare inputs inputs = { 'business_requirements': business_reqs_str, 'application_requirements': app_reqs_str, 'requirements_context': f"Business: {business_reqs_str}\nApplications: {app_reqs_str}" } progress(0.3, desc="Lead Architect analyzing requirements...") # Execute crew result = crew.kickoff(inputs=inputs) progress(0.9, desc="Finalizing reports...") # Extract individual task outputs task_outputs = result.tasks_output if hasattr(result, 'tasks_output') else [] if len(task_outputs) >= 3: lead_output = task_outputs[0].raw if hasattr(task_outputs[0], 'raw') else str(task_outputs[0]) senior_output = task_outputs[1].raw if hasattr(task_outputs[1], 'raw') else str(task_outputs[1]) peer_output = task_outputs[2].raw if hasattr(task_outputs[2], 'raw') else str(task_outputs[2]) else: # Fallback if task outputs aren't properly separated full_output = str(result) # Try to split the output into sections sections = full_output.split("# Task") if len(sections) >= 3: lead_output = f"# Technical Requirements\n{sections[1]}" if len(sections) > 1 else "Output processing..." senior_output = f"# Network Designs\n{sections[2]}" if len(sections) > 2 else "Processing designs..." peer_output = f"# Validation Report\n{sections[3]}" if len(sections) > 3 else "Finalizing validation..." else: lead_output = "✅ Requirements analysis complete" senior_output = "✅ Designs generated" peer_output = full_output progress(1.0, desc="Complete!") return lead_output, senior_output, peer_output except Exception as e: error_msg = f"❌ **Error occurred**: {str(e)}\n\n**Troubleshooting:**\n- Check your OpenAI API key\n- Ensure stable internet connection\n- Try selecting fewer requirements if the request is too complex" return error_msg, error_msg, error_msg # --- 4. Enhanced CSS --- custom_css = """ """ # --- 5. Gradio Interface --- with gr.Blocks( theme="soft", css=custom_css, title="AI Co-Designer for Data Center Networks" ) as demo: # Header gr.HTML("""
AI Co-Designer for Data Center Networks

Select your business and application requirements to generate comprehensive network designs with AI-powered validation.

""") # Main layout with gr.Row(): # Left column - Requirements with gr.Column(scale=1): # Business Requirements gr.HTML('
Business Requirements
') business_checkboxes = gr.CheckboxGroup( choices=[ "Strict Budget (CAPEX)", "High Scalability", "High Automation (OPEX)", "High Availability", "Regulatory Compliance" ], value=[], container=False, show_label=False ) # Application Requirements gr.HTML('
Application Requirements
') application_checkboxes = gr.CheckboxGroup( choices=[ "Large-Scale Virtualization", "AI/ML Workloads", "Low-Latency / Real-time", "Big Data Clusters", "High East-West Traffic" ], value=[], container=False, show_label=False ) # Generate Button submit_button = gr.Button( "🚀 Generate & Validate Design", elem_classes=["generate-btn"] ) # Right column - Agent outputs with gr.Column(scale=2): # Lead Architect Output gr.HTML('
💡 Lead Architect: Technical Requirements
') lead_architect_output = gr.Markdown( value="""

Requirements analysis will appear here...

""", container=False, show_label=False ) # Senior Architect Output gr.HTML('
🏗️ Senior Architect: Network Designs
') senior_architect_output = gr.Markdown( value="""

Network designs will appear here...

""", container=False, show_label=False ) # Peer Review Output gr.HTML('
✅ Peer Review: Validation & Recommendation
') peer_review_output = gr.Markdown( value="""

Validation report will appear here...

""", container=False, show_label=False ) # Connect the functionality submit_button.click( fn=run_design_crew, inputs=[business_checkboxes, application_checkboxes], outputs=[lead_architect_output, senior_architect_output, peer_review_output], show_progress=True ) # Launch the app if __name__ == "__main__": demo.launch( share=True, server_name="0.0.0.0", server_port=7860, show_error=True )