Added url field for dataset

app.py

@@ -1,27 +1,33 @@
 """
-Data Analyst Duo MCP Implementation - Simplified version
+Data Analyst Duo MCP Implementation - Full Working Version
+Supports loading any CSV over HTTP(S), including:
+  - Default cereal dataset
+  - Seaborn diamonds.csv
+  - FiveThirtyEight candy-data.csv
 """
 
 import os
 import json
 import datetime
-import gradio as gr
+import logging
+import uuid
+from io import StringIO
+
 import pandas as pd
 import numpy as np
 import requests
-from io import StringIO
-import logging
-import uuid
+import gradio as gr
 
 # Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
 # ============== MCP Protocol Implementation ==============
 
 class MCPMessage:
-    """Base class for MCP messages that agents exchange"""
-
     def __init__(self, sender, message_type, content):
         self.id = str(uuid.uuid4())
         self.sender = sender
@@ -38,10 +44,7 @@ class MCPMessage:
             "timestamp": self.timestamp
         }
 
-
 class MCPTool:
-    """Defines a tool that can be used by agents through the MCP protocol"""
-
     def __init__(self, name, description, function):
         self.name = name
         self.description = description
@@ -50,10 +53,7 @@ class MCPTool:
     def execute(self, params):
         return self.function(params)
 
-
 class MCPAgent:
-    """Base agent class implementing MCP protocol"""
-
     def __init__(self, name, description):
         self.name = name
         self.description = description
@@ -63,427 +63,225 @@ class MCPAgent:
         self.message_history = []
 
     def register_tool(self, tool):
-        """Register a tool that this agent can use"""
         self.tools[tool.name] = tool
 
     def connect(self, peer):
-        """Connect to another agent as a peer"""
         self.peers[peer.name] = peer
 
     def send_message(self, receiver, message_type, content):
-        """Send a message to a peer agent"""
         if receiver not in self.peers:
             raise ValueError(f"Peer {receiver} not found")
-
-        message = MCPMessage(self.name, message_type, content)
-        message_dict = message.to_dict()
-
-        # Save to message history
-        self.message_history.append({
-            "type": "sent",
-            "message": message_dict
-        })
-
-        # Send to receiver
-        self.peers[receiver].receive_message(message)
-        logger.info(f"Agent {self.name} sent {message_type} to {receiver}")
-        return message_dict
+        msg = MCPMessage(self.name, message_type, content)
+        self.message_history.append({"type": "sent", "message": msg.to_dict()})
+        self.peers[receiver].receive_message(msg)
+        logger.info(f"{self.name} → {receiver}: {message_type}")
+        return msg.to_dict()
 
     def receive_message(self, message):
-        """Receive a message from a peer agent"""
         self.message_queue.append(message)
-
-        # Save to message history
-        self.message_history.append({
-            "type": "received",
-            "message": message.to_dict()
-        })
-
-        logger.info(f"Agent {self.name} received {message.message_type} from {message.sender}")
+        self.message_history.append({"type": "received", "message": message.to_dict()})
+        logger.info(f"{self.name} received {message.message_type} from {message.sender}")
 
     def process_messages(self):
-        """Process all messages in the queue"""
-        processed = []
+        responses = []
         while self.message_queue:
-            message = self.message_queue.pop(0)
-            response = self.handle_message(message)
-            processed.append(response)
-        return processed
-
-    def handle_message(self, message):
-        """Handle a message - to be implemented by subclasses"""
-        raise NotImplementedError("Subclasses must implement handle_message")
+            msg = self.message_queue.pop(0)
+            resp = self.handle_message(msg)
+            responses.append(resp)
+        return responses
 
     def get_message_history(self):
-        """Get the agent's message history"""
         return self.message_history
 
+    def handle_message(self, message):
+        raise NotImplementedError("Override in subclass")
 
-# ============== Compute Agent Implementation ==============
+# ============== Compute Agent ==============
 
 class ComputeAgent(MCPAgent):
-    """Agent responsible for data loading, cleaning, and computation"""
-
     def __init__(self, name="ComputeAgent"):
-        super().__init__(name, "Agent responsible for data loading, cleaning and computation")
+        super().__init__(name, "Loads and computes on datasets")
         self.dataframe = None
 
-        # Register tools
-        self.register_tool(MCPTool(
-            "load_dataset",
-            "Load a dataset from URL",
-            self._load_dataset
-        ))
-
-        self.register_tool(MCPTool(
-            "compute_statistics",
-            "Compute basic statistics on the dataset",
-            self._compute_statistics
-        ))
-
-        self.register_tool(MCPTool(
-            "compute_correlation",
-            "Compute correlation between columns",
-            self._compute_correlation
-        ))
+        # Tools
+        self.register_tool(MCPTool("load_dataset", "Load a dataset from URL", self._load_dataset))
+        self.register_tool(MCPTool("compute_statistics", "Compute basic statistics", self._compute_statistics))
+        self.register_tool(MCPTool("compute_correlation", "Compute correlation matrix", self._compute_correlation))
 
     def _load_dataset(self, params):
-        """Load a dataset from URL"""
-        dataset_url = params.get("url")
-
+        url = params.get("url", "").strip()
+        # default cereal dataset
+        if not url or url.lower() == "default":
+            url = "https://raw.githubusercontent.com/datasciencedojo/datasets/master/cereal.csv"
         try:
-            # Use default cereals dataset if not specified
-            if not dataset_url or dataset_url == "default":
-                dataset_url = "https://raw.githubusercontent.com/datasciencedojo/datasets/master/cereal.csv"
-
-            # Load the dataset
-            response = requests.get(dataset_url)
-            content = response.content.decode('utf-8')
-            self.dataframe = pd.read_csv(StringIO(content))
-
-            # Basic info about the dataset
-            info = {
+            # fetch via pandas
+            self.dataframe = pd.read_csv(url)
+            return {
                 "status": "success",
-                "rows": len(self.dataframe),
+                "rows": self.dataframe.shape[0],
                 "columns": list(self.dataframe.columns),
                 "preview": self.dataframe.head(5).to_dict(orient="records")
             }
-
-            return info
-
         except Exception as e:
+            logger.exception("Error loading dataset")
             return {"status": "error", "message": str(e)}
 
     def _compute_statistics(self, params):
-        """Compute basic statistics on the dataset"""
         if self.dataframe is None:
             return {"status": "error", "message": "No dataset loaded"}
-
         try:
-            # Get columns to compute stats for
-            columns = params.get("columns", list(self.dataframe.select_dtypes(include=[np.number]).columns))
-
-            # Basic descriptive statistics
-            stats = self.dataframe[columns].describe().to_dict()
-
-            return {
-                "status": "success",
-                "statistics": stats
-            }
-
+            cols = params.get("columns", list(self.dataframe.select_dtypes(include=[np.number]).columns))
+            stats = self.dataframe[cols].describe().to_dict()
+            return {"status": "success", "statistics": stats}
         except Exception as e:
+            logger.exception("Error computing statistics")
             return {"status": "error", "message": str(e)}
 
     def _compute_correlation(self, params):
-        """Compute correlation between columns"""
         if self.dataframe is None:
             return {"status": "error", "message": "No dataset loaded"}
-
         try:
-            # Get columns to compute correlation for
-            columns = params.get("columns", list(self.dataframe.select_dtypes(include=[np.number]).columns))
-
-            corr_matrix = self.dataframe[columns].corr().to_dict()
-
-            return {
-                "status": "success",
-                "correlation_matrix": corr_matrix
-            }
-
+            cols = params.get("columns", list(self.dataframe.select_dtypes(include=[np.number]).columns))
+            corr = self.dataframe[cols].corr().to_dict()
+            return {"status": "success", "correlation_matrix": corr}
         except Exception as e:
+            logger.exception("Error computing correlation")
             return {"status": "error", "message": str(e)}
 
     def handle_message(self, message):
-        """Handle incoming messages from other agents"""
-        if message.message_type == "request_data_load":
-            result = self._load_dataset(message.content)
+        mtype = message.message_type
+        content = message.content
+        if mtype == "request_data_load":
+            result = self._load_dataset(content)
             return self.send_message(message.sender, "data_load_result", result)
-
-        elif message.message_type == "request_statistics":
-            result = self._compute_statistics(message.content)
+        elif mtype == "request_statistics":
+            result = self._compute_statistics(content)
             return self.send_message(message.sender, "statistics_result", result)
-
-        elif message.message_type == "request_correlation":
-            result = self._compute_correlation(message.content)
+        elif mtype == "request_correlation":
+            result = self._compute_correlation(content)
             return self.send_message(message.sender, "correlation_result", result)
-
         else:
-            return {"status": "error", "message": f"Unknown message type: {message.message_type}"}
-
+            return {"status": "error", "message": f"Unknown message type {mtype}"}
 
-# ============== Interpret Agent Implementation ==============
+# ============== Interpret Agent ==============
 
 class InterpretAgent(MCPAgent):
-    """Agent responsible for interpreting results and visualizing data"""
-
     def __init__(self, name="InterpretAgent"):
-        super().__init__(name, "Agent responsible for interpreting results and visualizing data")
+        super().__init__(name, "Interprets and reports on results")
         self.dataset_info = None
         self.statistics = None
-        self.correlation_data = None
-
-        # Register tools
-        self.register_tool(MCPTool(
-            "interpret_statistics",
-            "Interpret statistical results and provide insights",
-            self._interpret_statistics
-        ))
-
-        self.register_tool(MCPTool(
-            "interpret_correlation",
-            "Interpret correlation results and provide insights",
-            self._interpret_correlation
-        ))
-
-        self.register_tool(MCPTool(
-            "generate_report",
-            "Generate a report with key findings",
-            self._generate_report
-        ))
-
-    def _interpret_statistics(self, params):
-        """Interpret statistical results and provide insights"""
-        if not self.statistics:
-            return {"status": "error", "message": "No statistics data available"}
-
-        try:
-            insights = []
-            stats = self.statistics.get("statistics", {})
-
-            # Simple rule-based insights
-            for col, col_stats in stats.items():
-                # Add a simple insight about the mean value
-                if "mean" in col_stats:
-                    insights.append(f"The average {col} is {col_stats['mean']:.2f}")
+        self.correlation = None
 
-                # Add insight about range
-                if "min" in col_stats and "max" in col_stats:
-                    insights.append(f"{col} ranges from {col_stats['min']:.2f} to {col_stats['max']:.2f}")
+        self.register_tool(MCPTool("interpret_statistics", "", self._interpret_statistics))
+        self.register_tool(MCPTool("interpret_correlation", "", self._interpret_correlation))
+        self.register_tool(MCPTool("generate_report", "", self._generate_report))
 
-            return {
-                "status": "success",
-                "insights": insights[:3],  # Limit to top 3 insights
-                "summary": "Statistical analysis complete."
-            }
-
-        except Exception as e:
-            return {"status": "error", "message": str(e)}
+    def _interpret_statistics(self, params):
+        stats = self.statistics.get("statistics", {})
+        insights = []
+        for col, vals in stats.items():
+            if "mean" in vals:
+                insights.append(f"{col} avg = {vals['mean']:.2f}")
+            if "min" in vals and "max" in vals:
+                insights.append(f"{col} ranges {vals['min']:.2f}–{vals['max']:.2f}")
+        return {"status": "success", "insights": insights[:3], "summary": "Stats interpreted"}
 
     def _interpret_correlation(self, params):
-        """Interpret correlation results and provide insights"""
-        if not self.correlation_data:
-            return {"status": "error", "message": "No correlation data available"}
-
-        try:
-            insights = ["Correlation analysis complete."]
-
-            return {
-                "status": "success",
-                "insights": insights,
-                "summary": "Correlation analysis complete."
-            }
-
-        except Exception as e:
-            return {"status": "error", "message": str(e)}
+        return {"status": "success", "insights": ["Correlation matrix computed"], "summary": ""}
 
     def _generate_report(self, params):
-        """Generate a report with key findings"""
-        try:
-            report_sections = []
-
-            # Dataset overview
-            if self.dataset_info:
-                report_sections.append({
-                    "title": "Dataset Overview",
-                    "content": f"The dataset contains {self.dataset_info.get('rows', 0)} rows and {len(self.dataset_info.get('columns', []))} columns."
-                })
-
-            # Simple conclusion
-            report_sections.append({
-                "title": "Conclusions",
-                "content": "Analysis complete."
+        sections = []
+        if self.dataset_info:
+            sections.append({
+                "title": "Overview",
+                "content": f"{self.dataset_info['rows']} rows × {len(self.dataset_info['columns'])} cols"
             })
-
-            return {
-                "status": "success",
-                "report": {
-                    "title": params.get("report_title", "Data Analysis Report"),
-                    "sections": report_sections
-                }
+        sections.append({"title": "Conclusion", "content": "Analysis complete."})
+        return {
+            "status": "success",
+            "report": {
+                "title": params.get("report_title", "Report"),
+                "sections": sections
             }
-
-        except Exception as e:
-            return {"status": "error", "message": str(e)}
+        }
 
     def handle_message(self, message):
-        """Handle incoming messages from other agents"""
-        if message.message_type == "data_load_result":
-            self.dataset_info = message.content
-            return self.send_message(message.sender, "acknowledge", {"status": "received", "message": "Dataset info received"})
-
-        elif message.message_type == "statistics_result":
-            self.statistics = message.content
-            insights = self._interpret_statistics({})
-            return self.send_message(message.sender, "statistics_interpretation", insights)
-
-        elif message.message_type == "correlation_result":
-            self.correlation_data = message.content
-            insights = self._interpret_correlation({})
-            return self.send_message(message.sender, "correlation_interpretation", insights)
-
-        elif message.message_type == "request_report":
-            report = self._generate_report(message.content)
+        mtype = message.message_type
+        content = message.content
+        if mtype == "data_load_result":
+            self.dataset_info = content
+            return self.send_message(message.sender, "ack", {"status": "loaded"})
+        elif mtype == "statistics_result":
+            self.statistics = content
+            interp = self._interpret_statistics({})
+            return self.send_message(message.sender, "statistics_interpretation", interp)
+        elif mtype == "correlation_result":
+            self.correlation = content
+            interp = self._interpret_correlation({})
+            return self.send_message(message.sender, "correlation_interpretation", interp)
+        elif mtype == "request_report":
+            report = self._generate_report(content)
             return self.send_message(message.sender, "report_result", report)
-
         else:
-            return {"status": "error", "message": f"Unknown message type: {message.message_type}"}
+            return {"status": "error", "message": f"Unknown message type {mtype}"}
 
-
-# ============== Main Analysis Workflow ==============
+# ============== Main Workflow ==============
 
 class DataAnalystDuo:
-    """Main class for the Data Analyst Duo MCP implementation"""
-
     def __init__(self):
         self.compute_agent = ComputeAgent()
         self.interpret_agent = InterpretAgent()
-
-        # Connect the agents as peers
         self.compute_agent.connect(self.interpret_agent)
         self.interpret_agent.connect(self.compute_agent)
 
     def run_analysis(self, dataset_url="default"):
-        """Run the complete analysis workflow"""
-
-        # 1. Load dataset
+        # 1. Load
         self.interpret_agent.send_message("ComputeAgent", "request_data_load", {"url": dataset_url})
-        self.compute_agent.process_messages()
-        self.interpret_agent.process_messages()
-
-        # 2. Compute statistics
-        self.interpret_agent.send_message("ComputeAgent", "request_statistics", {"descriptive": True})
-        self.compute_agent.process_messages()
-        self.interpret_agent.process_messages()
-
-        # 3. Compute correlation
-        self.interpret_agent.send_message("ComputeAgent", "request_correlation", {"method": "pearson"})
-        self.compute_agent.process_messages()
-        self.interpret_agent.process_messages()
-
-        # 4. Generate final report
-        self.compute_agent.send_message("InterpretAgent", "request_report", {"report_title": "Data Analysis Report"})
-        self.interpret_agent.process_messages()
-        self.compute_agent.process_messages()
-
-        # Collect results
-        results = {
-            "compute_agent_messages": self.compute_agent.get_message_history(),
-            "interpret_agent_messages": self.interpret_agent.get_message_history()
-        }
-
-        return results
+        self.compute_agent.process_messages(); self.interpret_agent.process_messages()
+        # 2. Stats
+        self.interpret_agent.send_message("ComputeAgent", "request_statistics", {})
+        self.compute_agent.process_messages(); self.interpret_agent.process_messages()
+        # 3. Corr
+        self.interpret_agent.send_message("ComputeAgent", "request_correlation", {})
+        self.compute_agent.process_messages(); self.interpret_agent.process_messages()
+        # 4. Report
+        self.compute_agent.send_message("InterpretAgent", "request_report", {"report_title": "Analysis Report"})
+        self.interpret_agent.process_messages(); self.compute_agent.process_messages()
 
+        return {
+            "compute": self.compute_agent.get_message_history(),
+            "interpret": self.interpret_agent.get_message_history()
+        }
 
 # ============== Gradio Interface ==============
 
-def format_json(json_data):
-    """Format JSON data for display"""
-    if isinstance(json_data, dict) or isinstance(json_data, list):
-        return json.dumps(json_data, indent=2)
-    return str(json_data)
+def format_json(data):
+    return json.dumps(data, indent=2) if isinstance(data, (dict, list)) else str(data)
 
 def run_analysis(dataset_url):
-    """Run the data analysis workflow and return formatted messages"""
-    try:
-        # Use default cereals dataset if not specified
-        if not dataset_url:
-            dataset_url = "default"
-
-        # Create and run the analyst duo
-        duo = DataAnalystDuo()
-        results = duo.run_analysis(dataset_url)
-
-        # Format messages for display
-        all_messages = []
-
-        # Add compute agent messages
-        for msg in results["compute_agent_messages"]:
-            formatted_msg = f"[{msg['message']['timestamp']}] ComputeAgent {msg['type'].upper()} - Type: {msg['message']['message_type']}\n"
-            formatted_msg += format_json(msg['message']['content'])
-            formatted_msg += "\n\n" + "-"*80 + "\n\n"
-            all_messages.append((msg['message']['timestamp'], formatted_msg))
-
-        # Add interpret agent messages
-        for msg in results["interpret_agent_messages"]:
-            formatted_msg = f"[{msg['message']['timestamp']}] InterpretAgent {msg['type'].upper()} - Type: {msg['message']['message_type']}\n"
-            formatted_msg += format_json(msg['message']['content'])
-            formatted_msg += "\n\n" + "-"*80 + "\n\n"
-            all_messages.append((msg['message']['timestamp'], formatted_msg))
-
-        # Sort messages by timestamp
-        all_messages.sort(key=lambda x: x[0])
-
-        # Join messages
-        formatted_output = "\n".join([msg[1] for msg in all_messages])
-
-        return formatted_output
-
-    except Exception as e:
-        import traceback
-        return f"Error: {str(e)}\n\n{traceback.format_exc()}"
-
-# Define the Gradio interface
-with gr.Blocks(title="Data Analyst Duo - MCP Communication") as app:
-    gr.Markdown("""
-    # Data Analyst Duo - Model Context Protocol (MCP) Implementation
-
-    This application demonstrates a multi-agent system using the Model Context Protocol (MCP).
-    It consists of two agents:
-
-    1. **ComputeAgent**: Responsible for data loading, cleaning, and computation
-    2. **InterpretAgent**: Responsible for interpreting results
-
-    The agents communicate directly using standardized MCP messages, showcasing agent-to-agent communication.
-    """)
-
-    dataset_url = gr.Textbox(label="Dataset URL (leave empty for default cereals dataset)", placeholder="Enter dataset URL or leave empty for default")
-    run_button = gr.Button("Run Analysis")
-    mcp_messages = gr.Textbox(label="MCP Message Flow", lines=30)
-
-    run_button.click(fn=run_analysis, inputs=dataset_url, outputs=mcp_messages)
-
-    gr.Markdown("""
-    ## How This Demonstrates MCP
-
-    This application shows the Model Context Protocol in action:
-
-    1. **Standardized Message Structure**: All communication between agents follows a consistent format
-    2. **Direct Peer Communication**: Agents communicate directly with structured messages
-    3. **Asynchronous Processing**: Each agent processes messages independently
-
-    The message flow display shows the exact JSON messages exchanged between agents, demonstrating the protocol in action.
-    """)
+    duo = DataAnalystDuo()
+    histories = duo.run_analysis(dataset_url.strip())
+
+    all_msgs = []
+    for side in ["compute", "interpret"]:
+        for entry in histories[side]:
+            msg = entry["message"]
+            line = (f"[{msg['timestamp']}] {msg['sender']} "
+                    f"{entry['type'].upper()} {msg['message_type']}\n"
+                    f"{format_json(msg['content'])}\n\n" + "-"*60 + "\n")
+            all_msgs.append((msg['timestamp'], line))
+    all_msgs.sort(key=lambda x: x[0])
+    return "\n".join(line for _, line in all_msgs)
+
+with gr.Blocks(title="Data Analyst Duo MCP") as app:
+    gr.Markdown("## Data Analyst Duo — Load any CSV URL")
+    input_box = gr.Textbox(
+        label="Dataset URL",
+        placeholder="e.g. https://raw.githubusercontent.com/.../diamonds.csv"
+    )
+    run_btn = gr.Button("Run")
+    output_box = gr.Textbox(label="MCP Flow", lines=25)
+    run_btn.click(fn=run_analysis, inputs=input_box, outputs=output_box)
 
-# Launch the app
 if __name__ == "__main__":
-    app.launch()
+    app.launch()