Create app.py

app.py

@@ -0,0 +1,216 @@
+import streamlit as st
+from dotenv import load_dotenv
+import os 
+import pandas as pd
+from pandasai import SmartDataframe, SmartDatalake
+from pandasai.responses.response_parser import ResponseParser
+from pandasai.llm import GoogleGemini
+import plotly.graph_objects as go
+from PIL import Image
+import io
+import base64
+import requests
+
+# API Endpoint
+API_URL = "https://irisplus.elixir.co.zw/public/api/profile/reporting/stock-card/genericReports"
+PAYLOAD = {
+    "stock_card_report_id": "d2f1a0e1-7be1-472c-9610-94287154e544"
+}
+
+def fetch_data():
+    """Fetch stock card report data from API, return cleaned DataFrame"""
+    response = requests.post(API_URL, data=PAYLOAD)
+    if response.status_code == 200:
+        try:
+            data = response.json()
+            if isinstance(data, dict) and 'actual_report' in data and isinstance(data['actual_report'], list):
+                df = pd.DataFrame(data['actual_report'])  # Convert list to DataFrame
+                
+                # Remove columns where all values are None
+                df.dropna(axis=1, how='all', inplace=True)
+                
+                return df
+            else:
+                st.error("Unexpected response format from API.")
+                return None
+        except ValueError:
+            st.error("Error: Response is not valid JSON.")
+            return None
+    else:
+        st.error(f"Error fetching data: {response.status_code} - {response.text}")
+        return None
+
+
+class StreamLitResponse(ResponseParser):
+    def __init__(self, context):
+        super().__init__(context)
+
+    def format_dataframe(self, result):
+        """Enhanced DataFrame rendering with type identifier"""
+        return {
+            'type': 'dataframe',
+            'value': result['value']
+        }
+
+    def format_plot(self, result):
+        """Enhanced plot rendering with type identifier"""
+        try:
+            image = result['value']
+            
+            # Convert image to base64 for consistent storage
+            if isinstance(image, Image.Image):
+                buffered = io.BytesIO()
+                image.save(buffered, format="PNG")
+                base64_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
+            elif isinstance(image, bytes):
+                base64_image = base64.b64encode(image).decode('utf-8')
+            elif isinstance(image, str) and os.path.exists(image):
+                with open(image, "rb") as f:
+                    base64_image = base64.b64encode(f.read()).decode('utf-8')
+            else:
+                return {'type': 'text', 'value': "Unsupported image format"}
+
+            return {
+                'type': 'plot',
+                'value': base64_image
+            }
+        except Exception as e:
+            return {'type': 'text', 'value': f"Error processing plot: {e}"}
+
+    def format_other(self, result):
+        """Handle other types of responses"""
+        return {
+            'type': 'text',
+            'value': str(result['value'])
+        }
+
+# Load environment variables
+load_dotenv()
+GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
+
+if not GOOGLE_API_KEY:
+    st.error("GOOGLE_API_KEY environment variable not set.")
+    st.stop()
+
+def generateResponse(prompt, dfs):
+    """Generate response using PandasAI"""
+    llm = GoogleGemini(api_key=GOOGLE_API_KEY)
+    pandas_agent = SmartDatalake(dfs, config={
+        "llm": llm, 
+        "response_parser": StreamLitResponse
+    })
+    return pandas_agent.chat(prompt)
+
+def render_chat_message(message):
+    """Render different types of chat messages"""
+    if "dataframe" in message:
+        st.dataframe(message["dataframe"])
+    elif "plot" in message:
+        try:
+            # Handle base64 encoded images
+            plot_data = message["plot"]
+            if isinstance(plot_data, str):
+                st.image(f"data:image/png;base64,{plot_data}")
+            elif isinstance(plot_data, Image.Image):
+                st.image(plot_data)
+            elif isinstance(plot_data, go.Figure):
+                st.plotly_chart(plot_data)
+            elif isinstance(plot_data, bytes):
+                image = Image.open(io.BytesIO(plot_data))
+                st.image(image)
+            else:
+                st.write("Unsupported plot format")
+        except Exception as e:
+            st.error(f"Error rendering plot: {e}")
+    
+    # Always render text content
+    if "content" in message:
+        st.markdown(message["content"])
+
+
+def handle_userinput(question, dfs):
+    """Enhanced input handling with robust content processing"""
+    try:
+        # Ensure data is loaded
+        if dfs:
+            # Append user input to chat history
+            st.session_state.chat_history.append({
+                "role": "user",
+                "content": question
+            })
+
+            # Generate response with PandasAI
+            result = generateResponse(question, dfs)
+
+            # Handle different response types
+            if isinstance(result, dict):
+                response_type = result.get('type', 'text')
+                response_value = result.get('value')
+
+                if response_type == 'dataframe':
+                    st.session_state.chat_history.append({
+                        "role": "assistant",
+                        "content": "Here's the table:",
+                        "dataframe": response_value
+                    })
+                elif response_type == 'plot':
+                    st.session_state.chat_history.append({
+                        "role": "assistant",
+                        "content": "Here's the chart:",
+                        "plot": response_value
+                    })
+                else:
+                    st.session_state.chat_history.append({
+                        "role": "assistant",
+                        "content": str(response_value)
+                    })
+            else:
+                st.session_state.chat_history.append({
+                    "role": "assistant",
+                    "content": str(result)
+                })
+        else:
+            st.write("No data loaded.")
+
+    except Exception as e:
+        st.error(f"Error processing input: {e}")
+
+def main():
+    st.set_page_config(page_title="AI Chat with Your Data", page_icon="📊")
+
+    # Initialize session state variables
+    if "chat_history" not in st.session_state:
+        st.session_state.chat_history = []
+    if "dfs" not in st.session_state:
+        st.session_state.dfs = fetch_data()  # Load DataFrame at startup
+
+    st.title("AI Chat with Your Data 📊")
+
+    # Display chat history
+    for message in st.session_state.chat_history:
+        with st.chat_message(message["role"]):
+            render_chat_message(message)
+
+    # Chat input
+    user_question = st.chat_input("Ask a question about your data:")
+
+    if user_question:
+        handle_userinput(user_question, st.session_state.dfs)
+
+    # Sidebar with options
+    with st.sidebar:
+        st.subheader("Options")
+
+        if st.button("Reload Data"):
+            with st.spinner("Fetching latest data..."):
+                st.session_state.dfs = fetch_data()
+            st.success("Data refreshed!")
+
+        if st.button("Clear Chat"):
+            st.session_state.chat_history = []
+            st.rerun()
+
+if __name__ == "__main__":
+    main()
+
+