File size: 3,136 Bytes
cdf2dcf
 
 
 
 
 
 
 
 
 
55972ed
cdf2dcf
 
 
 
 
 
 
 
 
 
 
 
55972ed
cdf2dcf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55972ed
cdf2dcf
 
55972ed
 
 
 
 
cdf2dcf
 
 
 
 
55972ed
cdf2dcf
55972ed
 
 
 
 
cdf2dcf
55972ed
 
 
 
 
 
 
 
 
 
cdf2dcf
 
 
 
 
225b614
cdf2dcf
 
55972ed
 
 
 
 
aa52af1
55972ed
cdf2dcf
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
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()