Update app.py

app.py

@@ -1,11 +1,359 @@
+import os
+import subprocess
+import random
 from huggingface_hub import InferenceClient
 import gradio as gr
-import random
-import prompts
+from i_search import google
+from i_search import i_search as i_s
+from prompts import (
+    ACTION_PROMPT,
+    ADD_PROMPT,
+    COMPRESS_HISTORY_PROMPT,
+    LOG_PROMPT,
+    LOG_RESPONSE,
+    MODIFY_PROMPT,
+    PREFIX,
+    SEARCH_QUERY,
+    READ_PROMPT,
+    TASK_PROMPT,
+    UNDERSTAND_TEST_RESULTS_PROMPT,
+)
+from utils import parse_action, parse_file_content, read_python_module_structure
 client = InferenceClient(
     "mistralai/Mixtral-8x7B-Instruct-v0.1"
 )
 
+############################################
+
+
+VERBOSE = True
+MAX_HISTORY = 100
+#MODEL = "gpt-3.5-turbo"  # "gpt-4"
+
+
+def format_prompt(message, history):
+  prompt = "<s>"
+  for user_prompt, bot_response in history:
+    prompt += f"[INST] {user_prompt} [/INST]"
+    prompt += f" {bot_response}</s> "
+  prompt += f"[INST] {message} [/INST]"
+  return prompt
+
+
+
+def run_gpt(
+    prompt_template,
+    stop_tokens,
+    max_tokens,
+    module_summary,
+    purpose,
+    **prompt_kwargs,
+):
+    seed = random.randint(1,1111111111111111)
+
+    generate_kwargs = dict(
+        temperature=0.9,
+        max_new_tokens=10480,
+        top_p=0.95,
+        repetition_penalty=1.0,
+        do_sample=True,
+        seed=seed,
+    )
+
+    
+    content = PREFIX.format(
+        purpose=purpose,
+    ) + prompt_template.format(**prompt_kwargs)
+    if VERBOSE:
+        print(LOG_PROMPT.format(content))
+    
+    
+    #formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+    #formatted_prompt = format_prompt(f'{content}', history)
+
+    stream = client.text_generation(content, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    resp = ""
+    for response in stream:
+        resp += response.token.text
+        #yield resp
+    '''
+    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_search(purpose, task, history, directory, action_input):
+    print("CALLING SEARCH")
+    if "http" in action_input:
+        if "<" in action_input:
+            action_input = action_input.strip("<")
+        if ">" in action_input:
+            action_input = action_input.strip(">")
+        response = i_s(action_input)
+        #response = google(search_return)
+        print(response)
+        history += "observation: search result is: {}\n".format(response)
+    else:
+        history += "observation: I need to provide a valid URL to 'action: SEARCH action_input=URL'"
+    return "MAIN", None, history, task
+
+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)
+            print (f'ACTION_NAME :: {action_name}')
+            print (f'ACTION_INPUT :: {action_input}')
+            
+            history += "{}\n".format(line)
+            if action_name=="COMPLETE" or action_input=="COMPLETE":
+                task = "END"
+                return action_name, action_input, history, task
+            else:
+                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
+        )
+        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)
+        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
+def end_fn(purpose, task, history, directory, action_input):
+    task = "END"
+    return "COMPLETE", None, history, task
+NAME_TO_FUNC = {
+    "MAIN": call_main,
+    "UPDATE-TASK": call_set_task,
+    "SEARCH": call_search,
+    "COMPLETE": end_fn,
+
+}
+
+
+def run_action(purpose, task, history, directory, action_name, action_input):
+    if action_name == "COMPLETE":
+        task="END"
+        return action_name, action_input, history, task
+
+    # 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,
+        )
+        if task == "END":
+            return history
+
+
+
+################################################
+
 def format_prompt(message, history):
   prompt = "<s>"
   for user_prompt, bot_response in history: