Update mylab/gpt4o_based_sol.py

mylab/gpt4o_based_sol.py

@@ -0,0 +1,75 @@
+import streamlit as st
+import pandas as pd
+import plotly.express as px
+import json
+
+
+# GPT-4o powered function to generate visualization
+def ask_gpt4o_for_viz(query, df, llm):
+    columns = ', '.join(df.columns)
+    prompt = f"""
+    Analyze the user query and suggest the best way to visualize the data.
+    Query: "{query}"
+    Available Columns: {columns}
+    Respond in this JSON format:
+    {{
+      "chart_type": "bar/box/line/scatter",
+      "x_axis": "column_name",
+      "y_axis": "column_name",
+      "group_by": "optional_column_name"
+    }}
+    """
+    
+    response = llm.generate(prompt)
+    try:
+        suggestion = json.loads(response)
+        return suggestion
+    except json.JSONDecodeError:
+        st.error("⚠️ Failed to interpret AI response. Please refine your query.")
+        return None
+
+# Visualization generator
+def generate_viz(suggestion, df):
+    chart_type = suggestion.get("chart_type", "bar")
+    x_axis = suggestion.get("x_axis")
+    y_axis = suggestion.get("y_axis", "salary_in_usd")
+    group_by = suggestion.get("group_by")
+    
+    if not x_axis or not y_axis:
+        st.warning("⚠️ Could not identify required columns.")
+        return None
+    
+    # Generate the specified chart
+    if chart_type == "bar":
+        fig = px.bar(df, x=x_axis, y=y_axis, color=group_by)
+    elif chart_type == "box":
+        fig = px.box(df, x=x_axis, y=y_axis, color=group_by)
+    elif chart_type == "line":
+        fig = px.line(df, x=x_axis, y=y_axis, color=group_by)
+    elif chart_type == "scatter":
+        fig = px.scatter(df, x=x_axis, y=y_axis, color=group_by)
+    else:
+        st.warning("⚠️ Unsupported chart type suggested.")
+        return None
+    
+    fig.update_layout(title=f"{chart_type.title()} Plot of {y_axis} by {x_axis}")
+    return fig
+
+# Streamlit App
+st.title("📊 GPT-4o Powered Data Visualization")
+uploaded_file = st.file_uploader("Upload CSV File", type=["csv"])
+query = st.text_input("Ask a question about the data:")
+
+if uploaded_file:
+    df = load_data(uploaded_file)
+    st.write("### Dataset Preview", df.head())
+    
+    if query and st.button("Generate Visualization"):
+        llm = ChatOpenAI(api_key=os.getenv("OPENAI_API_KEY"), model="gpt-4o")  # Initialize GPT-4o
+        suggestion = ask_gpt4o_for_viz(query, df, llm)
+        if suggestion:
+            fig = generate_viz(suggestion, df)
+            if fig:
+                st.plotly_chart(fig, use_container_width=True)
+else:
+    st.info("Upload a CSV file to get started.")