Safety check implemented incorrectly

[rhiever]

System Info

OS version: Latest MacOS
Python version: 3.9
The current version of pandasai being used: 2.4.0

🐛 Describe the bug

The current safety check seems to be checking if any of these strings are in any of the generated code: "io", "os", "subprocess", "sys", "importlib"

So it's firing off on a matplotlib plt.title command with the word "Distribution" in it, for example.

Please fix!

[dosubot]

Hey @rhiever! I'm here to help you with bug fixes, answer your questions, and guide you on becoming a contributor. Let's squash those bugs together! 🤖

The issue you're encountering is due to the safety check in the BaseRestrictedModule class, which is designed to prevent the use of potentially unsafe modules by checking for specific strings like "io", "os", "subprocess", "sys", and "importlib" in the arguments of functions ^[1]. This check is likely too broad, causing it to incorrectly flag safe code, such as a matplotlib plt.title command containing the word "Distribution".

To address this, you might consider refining the safety check to be more context-aware. For example, you could implement a more sophisticated parsing mechanism that distinguishes between safe and unsafe uses of these keywords, or adjust the logic to only flag these keywords when they appear in specific contexts that are more likely to indicate a security risk.

To continue talking to Dosu, mention @dosu.

[ilyasch199]

i just got the same problem :
pandasai.safe_libs.base_restricted_module.SecurityError: Potential security risk: 'Distribution of Sepal Length in cm' is not allowed

[XJTU-JP]

I also met the same problem. How did you solve it? Hope to hear your suggestions @rhiever @ilyasch199

[rustykhatman]

Did this get solved?
I too get
pandasai.safe_libs.base_restricted_module.SecurityError: Potential security risk: 'Daily Consumption' is not allowed

How did this get solved?! This is ridiculous since Pandas AI creates the column 'Daily Consumption' and then tells me it's a security risk.

What are the work arounds?

[XJTU-JP]

Yes, I have solved this problem. My suggestion is to modify these constraints from the source code

[rustykhatman]

How would I how to modify it? Can you provide an example?

here is my code

Initialize OpenAI

        llm = OpenAI(api_token=os.getenv('OPENAI_API_KEY'))
        
        # Map the actual column names to our expected names
        column_mapping = {
            'ts': 'timestamp',
            'v0': 'consumption',
            'v1': 'demand',
            'v2': 'temperature',
            'site': 'site',
            'status': 'status'
        }
        
        data = data.rename(columns=column_mapping).reset_index()  # Ensure index is a column

        user_defined_path = os.getcwd()

        # Initialize SmartDataframe with OpenAI using GPT-4
        llm = OpenAI(
            api_token=os.getenv('OPENAI_API_KEY'),
            model="gpt-4",
            temperature=0,
            seed=26,
        )
        df = SmartDataframe(
            data,
            name="My Building Data",
            description="This dataframe contains energy consumption, demand, and temperature data for buildings.",
            config={
                "llm": llm,
                "verbose": True,
                "enable_cache": True,
                "save_charts": True,
                "open_charts": True,
                "use_error_correction_framework": True,
                "save_charts_path": "/charts",
                "enable_safe_mode": False,
                "enforce_privacy": False,
                "custom_whitelisted_dependencies": ["scikit-learn", "python-dateutil", "sklearn", "matplotlib"],
                "custom_plot_params": {
                    "title_prefix": "Energy Data: ",
                    "save_path": "/charts"
                }
            }
        )

        

        # Add data preview section
        st.subheader("Data Preview")
        with st.expander("Preview Data"):
            st.dataframe(df.head(), use_container_width=True)
            st.markdown("### Example Prompts")
            st.markdown("""
            Try these example questions:
            - What is the total consumption for each site?
            - Show me the average temperature by site.
            - What is the peak demand time for by site?
            - Create a line chart of consumption over time.
            - Compare the energy usage patterns between different sites
            """)
        st.markdown("---")  # Add a divider

        # Chat input
        if prompt := st.chat_input("Ask about your energy data..."):
            # Display user message
            st.chat_message("user").write(prompt)
            
            # Get response from SmartDataframe
            with st.chat_message("assistant"):
                with st.spinner("Getting AI response..."):
                    try:
                        response = df.chat(prompt)
                        st.write(response)
                        
                        # Check for generated charts
                        charts_path = "./charts"
                        if os.path.exists(charts_path):
                            for file in os.listdir(charts_path):
                                if file.endswith(('.png', '.jpg')):
                                    st.image(os.path.join(charts_path, file))
                                    os.remove(os.path.join(charts_path, file))
                    except Exception as e:
                        st.error(f"Error: {str(e)}")