Update app.py

app.py

@@ -1,297 +1,366 @@
 import os
 import subprocess
-#import openai
 from huggingface_hub import InferenceClient
 import gradio as gr
-import prompts
-client = InferenceClient(
-    "mistralai/Mixtral-8x7B-Instruct-v0.1"
-)
+import random
+import time
+from typing import List, Dict
+
+# Simulated agent and tool libraries
+AGENT_TYPES = ["Task Executor", "Information Retriever", "Decision Maker", "Data Analyzer"]
+TOOL_TYPES = ["Web Scraper", "Database Connector", "API Caller", "File Handler", "Text Processor"]
+
+# Initialize Hugging Face client
+client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+
+VERBOSE = False
+MAX_HISTORY = 100
+MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
+# Import necessary prompts and functions from the existing code
 from agent.prompts import (
-    ACTION_PROMPT,
-    ADD_PROMPT,
-    COMPRESS_HISTORY_PROMPT,
-    LOG_PROMPT,
-    LOG_RESPONSE,
-    MODIFY_PROMPT,
-    PREFIX,
-    READ_PROMPT,
-    TASK_PROMPT,
+    ACTION_PROMPT, ADD_PROMPT, COMPRESS_HISTORY_PROMPT, LOG_PROMPT,
+    LOG_RESPONSE, MODIFY_PROMPT, PREFIX, READ_PROMPT, TASK_PROMPT,
     UNDERSTAND_TEST_RESULTS_PROMPT,
 )
 from agent.utils import parse_action, parse_file_content, read_python_module_structure
 
-VERBOSE = False
-MAX_HISTORY = 100
-MODEL = "gpt-3.5-turbo"  # "gpt-4"
-
-
-def run_gpt(
-    prompt_template,
-    stop_tokens,
-    max_tokens,
-    module_summary,
-    purpose,
-    **prompt_kwargs,
-):
-    content = PREFIX.format(
-        module_summary=module_summary,
-        purpose=purpose,
-    ) + prompt_template.format(**prompt_kwargs)
-    if VERBOSE:
-        print(LOG_PROMPT.format(content))
-    
-    
-    #formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
-    stream = client.text_generation({"role": "system", "content": content}, **generate_kwargs, stream=True, details=True, return_full_text=False)["choices"][0]["message"]["content"]
-    output = ""
-    
-    '''
-    resp = openai.ChatCompletion.create(
-        model=MODEL,
-        messages=[
-            {"role": "system", "content": content},
-        ],
-        temperature=0.0,
-        max_tokens=max_tokens,
-        stop=stop_tokens if stop_tokens else None,
-    )["choices"][0]["message"]["content"]
-    '''
-    if VERBOSE:
-        print(LOG_RESPONSE.format(resp))
-    return resp
-
-
-def compress_history(purpose, task, history, directory):
-    module_summary, _, _ = read_python_module_structure(directory)
-    resp = run_gpt(
-        COMPRESS_HISTORY_PROMPT,
-        stop_tokens=["observation:", "task:", "action:", "thought:"],
-        max_tokens=512,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-    )
-    history = "observation: {}\n".format(resp)
-    return history
-
-
-def call_main(purpose, task, history, directory, action_input):
-    module_summary, _, _ = read_python_module_structure(directory)
-    resp = run_gpt(
-        ACTION_PROMPT,
-        stop_tokens=["observation:", "task:"],
-        max_tokens=256,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-    )
-    lines = resp.strip().strip("\n").split("\n")
-    for line in lines:
-        if line == "":
-            continue
-        if line.startswith("thought: "):
-            history += "{}\n".format(line)
-        elif line.startswith("action: "):
-            action_name, action_input = parse_action(line)
-            history += "{}\n".format(line)
-            return action_name, action_input, history, task
-        else:
-            assert False, "unknown action: {}".format(line)
-    return "MAIN", None, history, task
-
-
-def call_test(purpose, task, history, directory, action_input):
-    result = subprocess.run(
-        ["python", "-m", "pytest", "--collect-only", directory],
-        capture_output=True,
-        text=True,
-    )
-    if result.returncode != 0:
-        history += "observation: there are no tests! Test should be written in a test folder under {}\n".format(
-            directory
+class Agent:
+    def __init__(self, name: str, agent_type: str, complexity: int):
+        self.name = name
+        self.type = agent_type
+        self.complexity = complexity
+        self.tools = []
+
+    def add_tool(self, tool):
+        self.tools.append(tool)
+
+    def __str__(self):
+        return f"{self.name} ({self.type}) - Complexity: {self.complexity}"
+
+class Tool:
+    def __init__(self, name: str, tool_type: str):
+        self.name = name
+        self.type = tool_type
+
+    def __str__(self):
+        return f"{self.name} ({self.type})"
+
+class Pypelyne:
+    def __init__(self):
+        self.agents: List[Agent] = []
+        self.tools: List[Tool] = []
+        self.history = ""
+        self.task = None
+        self.purpose = None
+        self.directory = None
+
+    def add_agent(self, agent: Agent):
+        self.agents.append(agent)
+
+    def add_tool(self, tool: Tool):
+        self.tools.append(tool)
+
+    def generate_chat_app(self):
+        time.sleep(2)  # Simulate processing time
+        return f"Chat app generated with {len(self.agents)} agents and {len(self.tools)} tools."
+
+    def run_gpt(self, prompt_template, stop_tokens, max_tokens, **prompt_kwargs):
+        content = PREFIX.format(
+            module_summary=read_python_module_structure(self.directory)[0],
+            purpose=self.purpose,
+        ) + prompt_template.format(**prompt_kwargs)
+        
+        if VERBOSE:
+            print(LOG_PROMPT.format(content))
+        
+        stream = client.text_generation(
+            prompt=content,
+            max_new_tokens=max_tokens,
+            stop_sequences=stop_tokens if stop_tokens else None,
+            do_sample=True,
+            temperature=0.7,
         )
-        return "MAIN", None, history, task
-    result = subprocess.run(
-        ["python", "-m", "pytest", directory], capture_output=True, text=True
-    )
-    if result.returncode == 0:
-        history += "observation: tests pass\n"
-        return "MAIN", None, history, task
-    module_summary, content, _ = read_python_module_structure(directory)
-    resp = run_gpt(
-        UNDERSTAND_TEST_RESULTS_PROMPT,
-        stop_tokens=[],
-        max_tokens=256,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-        stdout=result.stdout[:5000],  # limit amount of text
-        stderr=result.stderr[:5000],  # limit amount of text
-    )
-    history += "observation: tests failed: {}\n".format(resp)
-    return "MAIN", None, history, task
-
-
-def call_set_task(purpose, task, history, directory, action_input):
-    module_summary, content, _ = read_python_module_structure(directory)
-    task = run_gpt(
-        TASK_PROMPT,
-        stop_tokens=[],
-        max_tokens=64,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-    ).strip("\n")
-    history += "observation: task has been updated to: {}\n".format(task)
-    return "MAIN", None, history, task
-
-
-def call_read(purpose, task, history, directory, action_input):
-    if not os.path.exists(action_input):
-        history += "observation: file does not exist\n"
-        return "MAIN", None, history, task
-    module_summary, content, _ = read_python_module_structure(directory)
-    f_content = (
-        content[action_input] if content[action_input] else "< document is empty >"
-    )
-    resp = run_gpt(
-        READ_PROMPT,
-        stop_tokens=[],
-        max_tokens=256,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-        file_path=action_input,
-        file_contents=f_content,
-    ).strip("\n")
-    history += "observation: {}\n".format(resp)
-    return "MAIN", None, history, task
-
-
-def call_modify(purpose, task, history, directory, action_input):
-    if not os.path.exists(action_input):
-        history += "observation: file does not exist\n"
-        return "MAIN", None, history, task
-    (
-        module_summary,
-        content,
-        _,
-    ) = read_python_module_structure(directory)
-    f_content = (
-        content[action_input] if content[action_input] else "< document is empty >"
-    )
-    resp = run_gpt(
-        MODIFY_PROMPT,
-        stop_tokens=["action:", "thought:", "observation:"],
-        max_tokens=2048,
-        module_summary=module_summary,
-        purpose=purpose,
-        task=task,
-        history=history,
-        file_path=action_input,
-        file_contents=f_content,
-    )
-    new_contents, description = parse_file_content(resp)
-    if new_contents is None:
-        history += "observation: failed to modify file\n"
-        return "MAIN", None, history, task
-
-    with open(action_input, "w") as f:
-        f.write(new_contents)
-
-    history += "observation: file successfully modified\n"
-    history += "observation: {}\n".format(description)
-    return "MAIN", None, history, task
-
-
-def call_add(purpose, task, history, directory, action_input):
-    d = os.path.dirname(action_input)
-    if not d.startswith(directory):
-        history += "observation: files must be under directory {}\n".format(directory)
-    elif not action_input.endswith(".py"):
-        history += "observation: can only write .py files\n"
-    else:
-        if d and not os.path.exists(d):
-            os.makedirs(d)
+        
+        resp = "".join(token for token in stream)
+        
+        if VERBOSE:
+            print(LOG_RESPONSE.format(resp))
+        return resp
+
+    def compress_history(self):
+        resp = self.run_gpt(
+            COMPRESS_HISTORY_PROMPT,
+            stop_tokens=["observation:", "task:", "action:", "thought:"],
+            max_tokens=512,
+            task=self.task,
+            history=self.history,
+        )
+        self.history = f"observation: {resp}\n"
+
+    def run_action(self, action_name, action_input):
+        if action_name == "COMPLETE":
+            return "Task completed."
+
+        if len(self.history.split("\n")) > MAX_HISTORY:
+            if VERBOSE:
+                print("COMPRESSING HISTORY")
+            self.compress_history()
+
+        action_funcs = {
+            "MAIN": self.call_main,
+            "UPDATE-TASK": self.call_set_task,
+            "MODIFY-FILE": self.call_modify,
+            "READ-FILE": self.call_read,
+            "ADD-FILE": self.call_add,
+            "TEST": self.call_test,
+        }
+
+        if action_name not in action_funcs:
+            return f"Unknown action: {action_name}"
+
+        print(f"RUN: {action_name} {action_input}")
+        return action_funcs[action_name](action_input)
+
+    def call_main(self, action_input):
+        resp = self.run_gpt(
+            ACTION_PROMPT,
+            stop_tokens=["observation:", "task:"],
+            max_tokens=256,
+            task=self.task,
+            history=self.history,
+        )
+        lines = resp.strip().strip("\n").split("\n")
+        for line in lines:
+            if line == "":
+                continue
+            if line.startswith("thought: "):
+                self.history += f"{line}\n"
+            elif line.startswith("action: "):
+                action_name, action_input = parse_action(line)
+                self.history += f"{line}\n"
+                return self.run_action(action_name, action_input)
+        return "No valid action found."
+
+    def call_set_task(self, action_input):
+        self.task = self.run_gpt(
+            TASK_PROMPT,
+            stop_tokens=[],
+            max_tokens=64,
+            task=self.task,
+            history=self.history,
+        ).strip("\n")
+        self.history += f"observation: task has been updated to: {self.task}\n"
+        return f"Task updated: {self.task}"
+
+    def call_modify(self, action_input):
         if not os.path.exists(action_input):
-            module_summary, _, _ = read_python_module_structure(directory)
-            resp = run_gpt(
-                ADD_PROMPT,
-                stop_tokens=["action:", "thought:", "observation:"],
-                max_tokens=2048,
-                module_summary=module_summary,
-                purpose=purpose,
-                task=task,
-                history=history,
-                file_path=action_input,
-            )
-            new_contents, description = parse_file_content(resp)
-            if new_contents is None:
-                history += "observation: failed to write file\n"
-                return "MAIN", None, history, task
-
-            with open(action_input, "w") as f:
-                f.write(new_contents)
-
-            history += "observation: file successfully written\n"
-            history += "obsertation: {}\n".format(description)
-        else:
-            history += "observation: file already exists\n"
-    return "MAIN", None, history, task
+            self.history += "observation: file does not exist\n"
+            return "File does not exist."
+        
+        content = read_python_module_structure(self.directory)[1]
+        f_content = content[action_input] if content[action_input] else "< document is empty >"
+        
+        resp = self.run_gpt(
+            MODIFY_PROMPT,
+            stop_tokens=["action:", "thought:", "observation:"],
+            max_tokens=2048,
+            task=self.task,
+            history=self.history,
+            file_path=action_input,
+            file_contents=f_content,
+        )
+        new_contents, description = parse_file_content(resp)
+        if new_contents is None:
+            self.history += "observation: failed to modify file\n"
+            return "Failed to modify file."
 
+        with open(action_input, "w") as f:
+            f.write(new_contents)
 
-NAME_TO_FUNC = {
-    "MAIN": call_main,
-    "UPDATE-TASK": call_set_task,
-    "MODIFY-FILE": call_modify,
-    "READ-FILE": call_read,
-    "ADD-FILE": call_add,
-    "TEST": call_test,
-}
+        self.history += f"observation: file successfully modified\n"
+        self.history += f"observation: {description}\n"
+        return f"File modified: {action_input}"
 
+    def call_read(self, action_input):
+        if not os.path.exists(action_input):
+            self.history += "observation: file does not exist\n"
+            return "File does not exist."
+        
+        content = read_python_module_structure(self.directory)[1]
+        f_content = content[action_input] if content[action_input] else "< document is empty >"
+        
+        resp = self.run_gpt(
+            READ_PROMPT,
+            stop_tokens=[],
+            max_tokens=256,
+            task=self.task,
+            history=self.history,
+            file_path=action_input,
+            file_contents=f_content,
+        ).strip("\n")
+        self.history += f"observation: {resp}\n"
+        return f"File read: {action_input}"
+
+    def call_add(self, action_input):
+        d = os.path.dirname(action_input)
+        if not d.startswith(self.directory):
+            self.history += f"observation: files must be under directory {self.directory}\n"
+            return f"Invalid directory: {d}"
+        elif not action_input.endswith(".py"):
+            self.history += "observation: can only write .py files\n"
+            return "Only .py files are allowed."
+        else:
+            if d and not os.path.exists(d):
+                os.makedirs(d)
+            if not os.path.exists(action_input):
+                resp = self.run_gpt(
+                    ADD_PROMPT,
+                    stop_tokens=["action:", "thought:", "observation:"],
+                    max_tokens=2048,
+                    task=self.task,
+                    history=self.history,
+                    file_path=action_input,
+                )
+                new_contents, description = parse_file_content(resp)
+                if new_contents is None:
+                    self.history += "observation: failed to write file\n"
+                    return "Failed to write file."
+
+                with open(action_input, "w") as f:
+                    f.write(new_contents)
+
+                self.history += "observation: file successfully written\n"
+                self.history += f"observation: {description}\n"
+                return f"File added: {action_input}"
+            else:
+                self.history += "observation: file already exists\n"
+                return "File already exists."
+
+    def call_test(self, action_input):
+        result = subprocess.run(
+            ["python", "-m", "pytest", "--collect-only", self.directory],
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            self.history += f"observation: there are no tests! Test should be written in a test folder under {self.directory}\n"
+            return "No tests found."
+        result = subprocess.run(
+            ["python", "-m", "pytest", self.directory], capture_output=True, text=True
+        )
+        if result.returncode == 0:
+            self.history += "observation: tests pass\n"
+            return "All tests passed."
+        
+        resp = self.run_gpt(
+            UNDERSTAND_TEST_RESULTS_PROMPT,
+            stop_tokens=[],
+            max_tokens=256,
+            task=self.task,
+            history=self.history,
+            stdout=result.stdout[:5000],
+            stderr=result.stderr[:5000],
+        )
+        self.history += f"observation: tests failed: {resp}\n"
+        return f"Tests failed: {resp}"
 
-def run_action(purpose, task, history, directory, action_name, action_input):
-    if action_name == "COMPLETE":
-        exit(0)
+pypelyne = Pypelyne()
 
-    # compress the history when it is long
-    if len(history.split("\n")) > MAX_HISTORY:
-        if VERBOSE:
-            print("COMPRESSING HISTORY")
-        history = compress_history(purpose, task, history, directory)
-
-    assert action_name in NAME_TO_FUNC
-
-    print("RUN: ", action_name, action_input)
-    return NAME_TO_FUNC[action_name](purpose, task, history, directory, action_input)
-
-
-def run(purpose, directory, task=None):
-    history = ""
-    action_name = "UPDATE-TASK" if task is None else "MAIN"
-    action_input = None
-    while True:
-        print("")
-        print("")
-        print("---")
-        print("purpose:", purpose)
-        print("task:", task)
-        print("---")
-        print(history)
-        print("---")
-
-        action_name, action_input, history, task = run_action(
-            purpose,
-            task,
-            history,
-            directory,
-            action_name,
-            action_input,
-        )
+def create_agent(name: str, agent_type: str, complexity: int) -> str:
+    agent = Agent(name, agent_type, complexity)
+    pypelyne.add_agent(agent)
+    return f"Agent created: {agent}"
+
+def create_tool(name: str, tool_type: str) -> str:
+    tool = Tool(name, tool_type)
+    pypelyne.add_tool(tool)
+    return f"Tool created: {tool}"
+
+def assign_tool(agent_name: str, tool_name: str) -> str:
+    agent = next((a for a in pypelyne.agents if a.name == agent_name), None)
+    tool = next((t for t in pypelyne.tools if t.name == tool_name), None)
+    
+    if agent and tool:
+        agent.add_tool(tool)
+        return f"Tool '{tool.name}' assigned to agent '{agent.name}'"
+    else:
+        return "Agent or tool not found."
+
+def generate_chat_app() -> str:
+    return pypelyne.generate_chat_app()
+
+def list_agents() -> str:
+    return "\n".join(str(agent) for agent in pypelyne.agents) or "No agents created yet."
+
+def list_tools() -> str:
+    return "\n".join(str(tool) for tool in pypelyne.tools) or "No tools created yet."
+
+def chat_with_pypelyne(message: str) -> str:
+    return pypelyne.run_action("MAIN", message)
+
+def set_purpose_and_directory(purpose: str, directory: str) -> str:
+    pypelyne.purpose = purpose
+    pypelyne.directory = directory
+    return f"Purpose set to: {purpose}\nWorking directory set to: {directory}"
+
+with gr.Blocks() as app:
+    gr.Markdown("# Welcome to Pypelyne")
+    gr.Markdown("Create your custom pipeline with agents and tools, then chat with it!")
+
+    with gr.Tab("Setup"):
+        purpose_input = gr.Textbox(label="Set Purpose")
+        directory_input = gr.Textbox(label="Set Working Directory")
+        setup_btn = gr.Button("Set Purpose and Directory")
+        setup_output = gr.Textbox(label="Setup Output")
+        setup_btn.click(set_purpose_and_directory, inputs=[purpose_input, directory_input], outputs=setup_output)
+
+    with gr.Tab("Create Agents"):
+        agent_name = gr.Textbox(label="Agent Name")
+        agent_type = gr.Dropdown(choices=AGENT_TYPES, label="Agent Type")
+        agent_complexity = gr.Slider(minimum=1, maximum=10, step=1, label="Agent Complexity")
+        create_agent_btn = gr.Button("Create Agent")
+        agent_output = gr.Textbox(label="Output")
+        create_agent_btn.click(create_agent, inputs=[agent_name, agent_type, agent_complexity], outputs=agent_output)
+
+    with gr.Tab("Create Tools"):
+        tool_name = gr.Textbox(label="Tool Name")
+        tool_type = gr.Dropdown(choices=TOOL_TYPES, label="Tool Type")
+        create_tool_btn = gr.Button("Create Tool")
+        tool_output = gr.Textbox(label="Output")
+        create_tool_btn.click(create_tool, inputs=[tool_name, tool_type], outputs=tool_output)
+
+    with gr.Tab("Assign Tools"):
+        agent_select = gr.Dropdown(choices=[], label="Select Agent")
+        tool_select = gr.Dropdown(choices=[], label="Select Tool")
+        assign_tool_btn = gr.Button("Assign Tool")
+        assign_output = gr.Textbox(label="Output")
+        assign_tool_btn.click(assign_tool, inputs=[agent_select, tool_select], outputs=assign_output)
+
+    with gr.Tab("Generate Chat App"):
+        generate_btn = gr.Button("Generate Chat App")
+        generate_output = gr.Textbox(label="Output")
+        generate_btn.click(generate_chat_app, outputs=generate_output)
+
+    with gr.Tab("Chat with Pypelyne"):
+        chat_input = gr.Textbox(label="Your Message")
+        chat_output = gr.Textbox(label="Pypelyne's Response")
+        chat_btn = gr.Button("Send")
+        chat_btn.click(chat_with_pypelyne, inputs=chat_input, outputs=chat_output)
+
+    with gr.Tab("View Pypelyne"):
+        view_agents_btn = gr.Button("View Agents")
+        view_tools_btn = gr.Button("View Tools")
+        view_output = gr.Textbox(label="Pypelyne Components")
+        view_agents_btn.click(list_agents, outputs=view_output)
+        view_tools_btn.click(list_tools, outputs=view_output)
+
+    def update_dropdowns():
+        return gr.Dropdown.update(choices=[agent.name for agent in pypelyne.agents]), \
+               gr.Dropdown.update(choices=[tool.name for tool in pypelyne.tools])
+
+    create_agent_btn.click(update_dropdowns, outputs=[agent_select, tool_select])
+    create_tool_btn.click(update_dropdowns, outputs=[agent_select, tool_select])
+
+if __name__ == "__main__":
+    app.launch()