multi_agent / app.py
ucKaizen's picture
Update app.py
55972ed verified
import os
import shutil
import gradio as gr
from autogen import ConversableAgent, AssistantAgent
from autogen.coding import LocalCommandLineCodeExecutor
# Setup
WORK_DIR = "coding"
CSV_PATH = os.path.join(WORK_DIR, "input.csv")
PLOT_PATH = os.path.join(WORK_DIR, "plot.png")
CODE_PATH = os.path.join(WORK_DIR, "generated_code.py")
os.makedirs(WORK_DIR, exist_ok=True)
# Agents
llm_config = {"model": "gpt-4-turbo"}
code_writer_agent = AssistantAgent(
name="code_writer_agent",
llm_config=llm_config,
human_input_mode="NEVER",
system_message=(
"You are a coding assistant. Write Python code that reads 'input.csv' in the current directory, "
"performs analysis based on the user prompt, and saves the resulting plot to 'plot.png'. "
"If plot is not possible, output text or tabular results to 'results.txt'. "
"Only generate the code, no explanations or comments."
),
)
executor = LocalCommandLineCodeExecutor(timeout=60, work_dir=WORK_DIR)
code_executor_agent = ConversableAgent(
name="code_executor_agent",
llm_config=False,
code_execution_config={"executor": executor},
human_input_mode="NEVER",
default_auto_reply="TERMINATE",
)
# Main function
def process_csv_and_prompt(prompt, file):
# Save uploaded file
if file is not None:
shutil.copyfile(file.name, CSV_PATH)
# Clean up old outputs
if os.path.exists(PLOT_PATH):
os.remove(PLOT_PATH)
if os.path.exists(CODE_PATH):
os.remove(CODE_PATH)
results_path = os.path.join(WORK_DIR, "results.txt")
if os.path.exists(results_path):
os.remove(results_path)
# Run agents
code_executor_agent.reset()
code_writer_agent.reset()
full_prompt = f"{prompt}\n(Remember: the file is 'input.csv' in current directory)"
chat = code_executor_agent.initiate_chat(code_writer_agent, message=full_prompt)
# Save generated code
if chat and "content" in chat.messages[-1]:
generated_code = chat.messages[-1]["content"]
with open(CODE_PATH, "w") as f:
f.write(generated_code)
else:
generated_code = "No code captured."
# Decide outputs
plot = PLOT_PATH if os.path.exists(PLOT_PATH) else None
table_text = None
if os.path.exists(results_path):
with open(results_path, "r") as f:
table_text = f.read()
return plot, generated_code, table_text or "No tabular/text results produced."
# Gradio Interface
iface = gr.Interface(
fn=process_csv_and_prompt,
inputs=[
gr.Textbox(label="Enter Analysis Prompt", placeholder="e.g. Generate 1000 Random numbers and plot histogram"),
gr.File(label="Upload CSV", file_types=[".csv"]),
],
outputs=[
gr.Image(label="Generated Plot"),
gr.Code(label="Generated Code"),
gr.Textbox(label="Tabular/Text Results")
],
title="Multi agent coding example",
description="Upload CSV + analysis prompt. System generates Python code, executes it, and shows results as plot, code, and/or text.",
allow_flagging="never",
)
if __name__ == "__main__":
iface.launch()