Added LLM

app.py

@@ -1,32 +1,33 @@
 """
-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
+Data Analyst Duo MCP with OpenAI Integration
 """
 
 import os
 import json
-import datetime
-import logging
 import uuid
+import logging
+import datetime
 from io import StringIO
 
 import pandas as pd
 import numpy as np
 import requests
 import gradio as gr
+import openai
 
-# Configure logging
+# ——— OpenAI Setup ——————————————————————————————————————
+openai.api_key = os.getenv("OPENAI_API_KEY")
+if not openai.api_key:
+    raise EnvironmentError("Missing OPENAI_API_KEY environment variable")
+
+# ——— Logging setup —————————————————————————————————————
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    format="%(asctime)s %(levelname)s:%(name)s: %(message)s"
 )
-logger = logging.getLogger(__name__)
-
-# ============== MCP Protocol Implementation ==============
+logger = logging.getLogger("DataAnalystDuo")
 
+# ——— MCP Core —————————————————————————————————————
 class MCPMessage:
     def __init__(self, sender, message_type, content):
         self.id = str(uuid.uuid4())
@@ -36,31 +37,27 @@ class MCPMessage:
         self.timestamp = datetime.datetime.now().isoformat()
 
     def to_dict(self):
-        return {
-            "id": self.id,
-            "sender": self.sender,
-            "message_type": self.message_type,
-            "content": self.content,
-            "timestamp": self.timestamp
-        }
+        return {"id": self.id, "sender": self.sender,
+                "message_type": self.message_type,
+                "content": self.content, "timestamp": self.timestamp}
 
 class MCPTool:
-    def __init__(self, name, description, function):
+    def __init__(self, name, description, func):
         self.name = name
         self.description = description
-        self.function = function
+        self.func = func
 
     def execute(self, params):
-        return self.function(params)
+        return self.func(params)
 
 class MCPAgent:
     def __init__(self, name, description):
         self.name = name
         self.description = description
         self.tools = {}
-        self.message_queue = []
         self.peers = {}
-        self.message_history = []
+        self.queue = []
+        self.history = []
 
     def register_tool(self, tool):
         self.tools[tool.name] = tool
@@ -68,220 +65,169 @@ class MCPAgent:
     def connect(self, peer):
         self.peers[peer.name] = peer
 
-    def send_message(self, receiver, message_type, content):
-        if receiver not in self.peers:
-            raise ValueError(f"Peer {receiver} not found")
-        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}")
+    def send_message(self, to, mtype, content):
+        if to not in self.peers:
+            raise ValueError(f"Peer {to} not found")
+        msg = MCPMessage(self.name, mtype, content)
+        self.history.append({"type": "sent", "message": msg.to_dict()})
+        self.peers[to].receive(msg)
+        logger.info(f"{self.name} → {to}: {mtype}")
         return msg.to_dict()
 
-    def receive_message(self, message):
-        self.message_queue.append(message)
-        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):
-        responses = []
-        while self.message_queue:
-            msg = self.message_queue.pop(0)
-            resp = self.handle_message(msg)
-            responses.append(resp)
-        return responses
+    def receive(self, msg):
+        self.queue.append(msg)
+        self.history.append({"type": "received", "message": msg.to_dict()})
+        logger.info(f"{self.name} received {msg.message_type} from {msg.sender}")
 
-    def get_message_history(self):
-        return self.message_history
+    def process(self):
+        while self.queue:
+            msg = self.queue.pop(0)
+            self.handle_message(msg)
 
     def handle_message(self, message):
-        raise NotImplementedError("Override in subclass")
+        raise NotImplementedError
 
-# ============== Compute Agent ==============
+    def get_history(self):
+        return self.history
 
+# ——— ComputeAgent —————————————————————————————————————
 class ComputeAgent(MCPAgent):
-    def __init__(self, name="ComputeAgent"):
-        super().__init__(name, "Loads and computes on datasets")
-        self.dataframe = None
-
-        # 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 __init__(self):
+        super().__init__("ComputeAgent", "Loads & computes data")
+        self.df = None
+        self.register_tool(MCPTool("load_dataset", "Load CSV from URL", self._load))
+        self.register_tool(MCPTool("compute_statistics", "Descriptive stats", self._stats))
+        self.register_tool(MCPTool("compute_correlation", "Correlation matrix", self._corr))
 
-    def _load_dataset(self, params):
+    def _load(self, params):
         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:
-            # fetch via pandas
-            self.dataframe = pd.read_csv(url)
-            return {
-                "status": "success",
-                "rows": self.dataframe.shape[0],
-                "columns": list(self.dataframe.columns),
-                "preview": self.dataframe.head(5).to_dict(orient="records")
-            }
+            self.df = pd.read_csv(url)
+            return {"status":"success","rows":self.df.shape[0],
+                    "columns":list(self.df.columns),
+                    "preview":self.df.head(5).to_dict(orient="records")}
         except Exception as e:
-            logger.exception("Error loading dataset")
-            return {"status": "error", "message": str(e)}
+            logger.exception("Load failed")
+            return {"status":"error","message":str(e)}
 
-    def _compute_statistics(self, params):
-        if self.dataframe is None:
-            return {"status": "error", "message": "No dataset loaded"}
+    def _stats(self, params):
+        if self.df is None:
+            return {"status":"error","message":"No data loaded"}
         try:
-            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}
+            cols = list(self.df.select_dtypes(include=[np.number]).columns)
+            stats = self.df[cols].describe().to_dict()
+            return {"status":"success","statistics":stats}
         except Exception as e:
-            logger.exception("Error computing statistics")
-            return {"status": "error", "message": str(e)}
+            logger.exception("Stats failed")
+            return {"status":"error","message":str(e)}
 
-    def _compute_correlation(self, params):
-        if self.dataframe is None:
-            return {"status": "error", "message": "No dataset loaded"}
+    def _corr(self, params):
+        if self.df is None:
+            return {"status":"error","message":"No data loaded"}
         try:
-            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}
+            cols = list(self.df.select_dtypes(include=[np.number]).columns)
+            corr = self.df[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):
-        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 mtype == "request_statistics":
-            result = self._compute_statistics(content)
-            return self.send_message(message.sender, "statistics_result", result)
-        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 {mtype}"}
-
-# ============== Interpret Agent ==============
-
+            logger.exception("Corr failed")
+            return {"status":"error","message":str(e)}
+
+    def handle_message(self, m):
+        if m.message_type == "request_data_load":
+            res = self._load(m.content);
+            self.send_message(m.sender, "data_load_result", res)
+        elif m.message_type == "request_statistics":
+            res = self._stats(m.content);
+            self.send_message(m.sender, "statistics_result", res)
+        elif m.message_type == "request_correlation":
+            res = self._corr(m.content);
+            self.send_message(m.sender, "correlation_result", res)
+
+# ——— InterpretAgent with LLM —————————————————————————————————
 class InterpretAgent(MCPAgent):
-    def __init__(self, name="InterpretAgent"):
-        super().__init__(name, "Interprets and reports on results")
-        self.dataset_info = None
-        self.statistics = None
-        self.correlation = None
-
-        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))
-
-    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}")
+    def __init__(self):
+        super().__init__("InterpretAgent","Interprets & reports via LLM")
+        self.data_info = None
+        self.stats = None
+        self.corr = None
+        # tools for rule-based fallback
+        self.register_tool(MCPTool("interpret_statistics","",self._int_stats))
+        self.register_tool(MCPTool("interpret_correlation","",self._int_corr))
+        # LLM-powered tools
+        self.register_tool(MCPTool("llm_interpret","GPT-4 insights",self._llm_interpret))
+        self.register_tool(MCPTool("llm_report","GPT-4 report",self._llm_report))
+
+    def _int_stats(self, params):
+        ins=[]
+        for col,vals in self.stats.get("statistics",{}).items():
+            if "mean" in vals: ins.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):
-        return {"status": "success", "insights": ["Correlation matrix computed"], "summary": ""}
-
-    def _generate_report(self, params):
-        sections = []
-        if self.dataset_info:
-            sections.append({
-                "title": "Overview",
-                "content": f"{self.dataset_info['rows']} rows × {len(self.dataset_info['columns'])} cols"
-            })
-        sections.append({"title": "Conclusion", "content": "Analysis complete."})
-        return {
-            "status": "success",
-            "report": {
-                "title": params.get("report_title", "Report"),
-                "sections": sections
-            }
-        }
-
-    def handle_message(self, message):
-        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 {mtype}"}
-
-# ============== Main Workflow ==============
-
+                ins.append(f"{col} ∈ [{vals['min']:.2f},{vals['max']:.2f}]")
+        return {"status":"success","insights":ins[:3],"summary":"Rule-based insights"}
+
+    def _int_corr(self, params):
+        return {"status":"success","insights":["Correlation computed"],"summary":"Rule-based corr"}
+
+    def _llm_interpret(self, params):
+        prompt = (
+            "I have these statistics in JSON format:\n" +
+            json.dumps(self.stats, indent=2) +
+            "\nPlease summarize the top 3 insights in plain English." )
+        resp = openai.ChatCompletion.create(
+            model="gpt-4",
+            messages=[{"role":"user","content":prompt}]
+        )
+        text = resp.choices[0].message.content.strip()
+        lines = [l for l in text.split("\n") if l.strip()]
+        return {"status":"success","insights":lines,"summary":text}
+
+    def _llm_report(self, params):
+        prompt = (
+            "Generate a concise analysis report in Markdown format based on the following:\n" +
+            f"Dataset info: rows={self.data_info['rows']}, cols={len(self.data_info['columns'])}\n" +
+            json.dumps(self.stats, indent=2) + "\n" +
+            json.dumps(self.corr, indent=2)
+        )
+        resp = openai.ChatCompletion.create(
+            model="gpt-4",
+            messages=[{"role":"user","content":prompt}]
+        )
+        return {"status":"success","report_md":resp.choices[0].message.content.strip()}
+
+    def handle_message(self, m):
+        if m.message_type == "data_load_result":
+            self.data_info = m.content
+            self.send_message(m.sender,"ack",{"status":"loaded"})
+        elif m.message_type == "statistics_result":
+            self.stats = m.content
+            # LLM-driven interpretation
+            llm_res = self._llm_interpret({})
+            self.send_message(m.sender,"llm_statistics_interpretation",llm_res)
+        elif m.message_type == "correlation_result":
+            self.corr = m.content
+            llm_res = self._llm_interpret({})
+            self.send_message(m.sender,"llm_correlation_interpretation",llm_res)
+        elif m.message_type == "request_report":
+            llm_res = self._llm_report({})
+            self.send_message(m.sender,"report_result",llm_res)
+
+# ——— Orchestration & Gradio —————————————————————————————————
 class DataAnalystDuo:
     def __init__(self):
-        self.compute_agent = ComputeAgent()
-        self.interpret_agent = InterpretAgent()
-        self.compute_agent.connect(self.interpret_agent)
-        self.interpret_agent.connect(self.compute_agent)
-
-    def run_analysis(self, dataset_url="default"):
-        # 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. 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(data):
-    return json.dumps(data, indent=2) if isinstance(data, (dict, list)) else str(data)
-
-def run_analysis(dataset_url):
-    duo = DataAnalystDuo()
-    histories = duo.run_analysis(dataset_url.strip())
+        self.C=ComputeAgent(); self.I=InterpretAgent()
+        self.C.connect(self.I); self.I.connect(self.C)
 
-    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)
+    def run(self,url):
+        self.I.send_message("ComputeAgent","request_data_load",{"url":url});
+        self.C.process(); self.I.process()
+        self.I.send_message("ComputeAgent","request_statistics",{});
+        self.C.process(); self.I.process()
+        self.I.send_message("ComputeAgent","request_correlation",{});
+        self.C.process(); self.I.process()
+        self.C.send_message("InterpretAgent","request_report",{"report_title":"Analysis Report"});
+        self.I.process(); self.C.process()
+        return self.C.get_history(), self.I.get_history()
 
-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)
 
-if __name__ == "__main__":
-    app.launch()