Data Analysis via AutoGen Agents

Below is a sample AutoGen app via Gemini 2.5

import autogen
import pandas as pd
import os
import traceback

# --- Configuration ---
# Configure the LLM provider. Uses environment variable OPENAI_API_KEY by default.
# You can change model names if needed, e.g., "gpt-3.5-turbo"
config_list = autogen.config_list_from_json(
    "OAI_CONFIG_LIST", # Reads environment variable OAI_CONFIG_LIST if available
                       # If not, searches for 'OAI_CONFIG_LIST.json' file in the current directory
                       # If neither is found, uses OPENAI_API_KEY environment variable
    filter_dict={
        "model": ["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"] # Specify desired models
    }
)

# Check if config_list is empty (meaning no valid configuration found)
if not config_list:
    print("Error: LLM configuration not found.")
    print("Please set the OPENAI_API_KEY environment variable or provide an OAI_CONFIG_LIST file.")
    exit()

llm_config = {
    "config_list": config_list,
    "cache_seed": 42, # Use None for reproducibility. Use an int for caching.
    "timeout": 120,
}

# Directory to store generated code (optional, for inspection)
CODE_DIR = "coding"
os.makedirs(CODE_DIR, exist_ok=True)

# Global variable to hold the DataFrame
df_global = None

# --- Helper Functions ---

def load_data(file_path):
    """Loads data from CSV or Excel into a pandas DataFrame."""
    global df_global
    try:
        if file_path.endswith('.csv'):
            df_global = pd.read_csv(file_path)
        elif file_path.endswith('.xlsx') or file_path.endswith('.xls'):
            df_global = pd.read_excel(file_path)
        else:
            print(f"Error: Unsupported file format: {file_path}")
            return False
        print(f"Successfully loaded data from {file_path}")
        return True
    except FileNotFoundError:
        print(f"Error: File not found at {file_path}")
        return False
    except Exception as e:
        print(f"Error loading data: {e}")
        traceback.print_exc()
        return False

def get_data_summary():
    """Returns a string summary of the loaded DataFrame (column names and types)."""
    if df_global is None:
        return "No data loaded."
    summary = "DataFrame Summary:\n"
    summary += f"- Shape: {df_global.shape}\n"
    summary += "- Columns and Types:\n"
    for col, dtype in df_global.dtypes.items():
        summary += f"  - {col}: {dtype}\n"
    # Optionally add head() for a small sample, but be mindful of data privacy
    # summary += f"\n- First 5 rows:\n{df_global.head().to_string()}"
    return summary

def execute_reviewed_code(code_to_execute):
    """
    Executes the provided Python code string after user confirmation.
    The code has access to the global DataFrame 'df' and pandas 'pd'.
    """
    global df_global
    if df_global is None:
        return "Error: No data loaded to operate on."

    print("\n--- Code Review ---")
    print("The following code is proposed for execution:")
    print("--------------------")
    print(code_to_execute)
    print("--------------------")

    while True:
        approve = input("Do you approve this code for execution? (yes/no/modify): ").lower().strip()
        if approve == 'yes':
            break
        elif approve == 'no':
            print("Execution cancelled.")
            return "Code execution rejected by user."
        elif approve == 'modify':
            print("Please enter the modified code below. Press Enter twice to finish.")
            modified_code_lines = []
            while True:
                try:
                    line = input()
                    if not line: # break on empty line
                        break
                    modified_code_lines.append(line)
                except EOFError: # break on Ctrl+D
                    break
            code_to_execute = "\n".join(modified_code_lines)
            print("\n--- Modified Code ---")
            print(code_to_execute)
            print("---------------------")
            continue # Re-prompt for approval of modified code
        else:
            print("Invalid input. Please enter 'yes', 'no', or 'modify'.")

    print("\nExecuting approved code...")
    # *** SECURITY WARNING ***
    # Using exec is inherently risky, even with review, as mistakes can happen.
    # A sandboxed environment (like Docker) is strongly recommended for production.
    # Ensure the user understands the code they are approving.
    local_scope = {'df': df_global, 'pd': pd} # Provide df and pd access
    global_scope = {}
    try:
        # Capture print output and the result of the last expression
        import io
        from contextlib import redirect_stdout

        f = io.StringIO()
        with redirect_stdout(f):
            exec(code_to_execute, global_scope, local_scope)

        # Try to get the result of the last expression (if any)
        result = local_scope.get('result', None) # Convention: assign result to 'result' variable in code
        printed_output = f.getvalue()

        output = "--- Execution Result ---"
        if printed_output:
             output += f"\nOutput:\n{printed_output}"
        if result is not None:
             output += f"\nFinal Result Variable:\n{result}"
        if not printed_output and result is None:
            output += "\nCode executed successfully with no visible output or 'result' variable."

        return output

    except Exception as e:
        print(f"Error during code execution: {e}")
        return f"Execution Error:\n{traceback.format_exc()}"

# --- Autogen Agent Setup ---

# 1. Assistant Agent (Code Generator)
assistant = autogen.AssistantAgent(
    name="Data_Analyst_Agent",
    system_message="""You are a helpful AI assistant specializing in Python data analysis using the pandas library.
You are given a summary of a pandas DataFrame (column names, types, shape) and a user question.
Your goal is to determine if the question can be answered by generating Python code to be run on the DataFrame (available as a variable named 'df').

If code is needed:
1.  Analyze the question and the DataFrame summary.
2.  Write Python code using pandas (available as 'pd') to answer the question. The DataFrame is available as 'df'.
3.  To make the result easily accessible after execution, try to assign the final result to a variable named `result` (e.g., `result = df['column'].mean()`). If the result is naturally printed (like with `df.head()` or `print()`), that's okay too.
4.  Wrap the Python code you generate within triple backticks, inside a JSON block like this:
    ```json
    {
      "requires_code": true,
      "python_code": "YOUR PYTHON CODE HERE"
    }

5. Only provide the JSON block in your response. Do not add explanations outside the JSON.

If code is not needed (e.g., the question is not about the data, or is too vague):

1. Explain briefly why code generation is not appropriate or possible.
2. Output a JSON block like this:

   {
     "requires_code": false,
     "explanation": "YOUR EXPLANATION HERE"
   }

3. Only provide the JSON block in your response. Do not add explanations outside the JSON.

Make sure your generated Python code is correct, safe, and directly addresses the user's question based only on the provided DataFrame summary. Do not hallucinate columns or data. Assume the DataFrame 'df' and pandas 'pd' are already loaded and available in the execution environment. """, llm_config=llm_config, )

2. User Proxy Agent (Handles Interaction and Code Execution Request)

We use register_function to link the code execution logic to the agent.

user_proxy = autogen.UserProxyAgent( name="User_Proxy", human_input_mode="NEVER", # We handle human input explicitly outside the agent chat for review max_consecutive_auto_reply=5, # Limit agent loops is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE"), code_execution_config=False, # Important: We disable default code execution description="Represents the user. Interacts with the Data_Analyst_Agent. Can call the 'execute_code' function to run verified code." )

Register the function for execution

This allows the Assistant to "request" execution via a function call,

which our custom function then intercepts for review.

user_proxy.register_function( function_map={ "execute_reviewed_code": execute_reviewed_code } )

--- Main Application Logic ---

def main(): print("Welcome to AgenticApplication!") print("-------------------------------")

# 1. Load Data
while True:
    file_path = input("Enter the path to your CSV or Excel file: ")
    if load_data(file_path):
        break
    else:
        print("Please try again.")

# 2. Get Data Summary (only summary is sent to LLM)
data_summary = get_data_summary()
# print("\nData Summary (for LLM):") # Optional: show what's sent
# print(data_summary)

print("\n-------------------------------")
print("Data loaded. You can now ask questions about your data.")
print("Type 'exit' to quit.")

# 3. Interaction Loop
while True:
    user_query = input("\nYour question: ")
    if user_query.lower() == 'exit':
        break
    if df_global is None:
         print("Error: Dataframe not loaded correctly.")
         break

    # Construct the prompt for the assistant
    prompt = f"""Here is a summary of the data:\n{data_summary}\n\nUser Question: {user_query}\n

Analyze the question. If you can write Python code using pandas to answer it, provide the code in the specified JSON format. Otherwise, provide an explanation in the JSON format."""

    # Initiate the chat - User Proxy sends the request to Assistant
    # We expect the Assistant to either reply with JSON containing code or an explanation.
    # If it suggests code, it should ideally try to call `execute_reviewed_code`.
    # Autogen's function calling mechanism handles routing this.

    user_proxy.initiate_chat(
        assistant,
        message=prompt,
        # clear_history=True # Start fresh for each query
    )

    # Note: The `execute_reviewed_code` function handles the review and execution flow when called by the agent.
    # The result of the execution (or cancellation) is returned as a message in the chat history,
    # which could potentially be used for follow-up questions if the conversation flow is extended.

print("\nExiting AgenticApplication. Goodbye!")

if name == "main": main()