Refactor/file store (#18)

* feat: change port variable name

* feat: chdir to playground

* refactor: etc

* feat: add verify decorator

README.md

@@ -77,7 +77,7 @@ Manatory envs are required in order to serve EVAL.
 
 Each optional env has default value, so you don't need to set unless you want to change it.
 
-- `PORT` - port (default: 8000)
+- `EVAL_PORT` - port (default: 8000)
 - `SERVER` - server address (default: http://localhost:8000)
 - `LOG_LEVEL` - INFO | DEBUG (default: INFO)
 - `BOT_NAME` - give it a name! (default: Orca)

api/main.py

@@ -1,3 +1,4 @@
+import os
 import re
 from typing import Dict, List, TypedDict
 
@@ -23,7 +24,7 @@ app = FastAPI()
 
 app.mount("/static", StaticFiles(directory=StaticUploader.STATIC_DIR), name="static")
 uploader = StaticUploader.from_settings(settings)
-
+os.chdir(settings["PLAYGROUND_DIR"])
 
 toolsets: List[BaseToolSet] = [
     Terminal(),
@@ -91,7 +92,7 @@ async def command(request: Request) -> Response:
     except Exception as e:
         return {"response": str(e), "files": []}
 
-    files = re.findall("(image/\S*png)|(dataframe/\S*csv)", res["output"])
+    files = re.findall("image/\S*png|dataframe/\S*csv", res["output"])
 
     return {
         "response": res["output"],
@@ -100,4 +101,4 @@ async def command(request: Request) -> Response:
 
 
 def serve():
-    uvicorn.run("api.main:app", host="0.0.0.0", port=settings["PORT"])
+    uvicorn.run("api.main:app", host="0.0.0.0", port=settings["EVAL_PORT"])

core/handlers/base.py

@@ -1,10 +1,13 @@
 import os
 import uuid
 from enum import Enum
+from pathlib import Path
 from typing import Dict
 
 import requests
 
+from env import settings
+
 
 class FileType(Enum):
     IMAGE = "image"
@@ -64,6 +67,7 @@ class FileHandler:
         local_filename = os.path.join(
             filetype.value, str(uuid.uuid4())[0:8] + filetype.to_extension()
         )
+        os.makedirs(os.path.dirname(local_filename), exist_ok=True)
         with open(local_filename, "wb") as f:
             size = f.write(data)
         print(f"Inputs: {url} ({size//1000}MB)  => {local_filename}")

core/handlers/dataframe.py

@@ -9,7 +9,7 @@ class CsvToDataframe(BaseHandler):
     def handle(self, filename: str):
         df = pd.read_csv(filename)
         description = (
-            f"Dataframe with {len(df)} rows and {len(df.columns)} columns."
+            f"Dataframe with {len(df)} rows and {len(df.columns)} columns. "
             "Columns are: "
             f"{', '.join(df.columns)}"
         )

core/prompts/file.py

@@ -2,6 +2,9 @@ IMAGE_PROMPT = """
 provide a figure named {filename}. The description is: {description}.
 
 Please understand and answer the image based on this information. The image understanding is complete, so don't try to understand the image again.
+
+USER INPUT
+============
 """
 
 
@@ -9,12 +12,18 @@ AUDIO_PROMPT = """
 provide a audio named {filename}. The description is: {description}.
 
 Please understand and answer the audio based on this information. The audio understanding is complete, so don't try to understand the audio again.
+
+USER INPUT
+============
 """
 
 VIDEO_PROMPT = """
 provide a video named {filename}. The description is: {description}.
 
 Please understand and answer the video based on this information. The video understanding is complete, so don't try to understand the video again.
+
+USER INPUT
+============
 """
 
 DATAFRAME_PROMPT = """
@@ -22,4 +31,7 @@ provide a dataframe named {filename}. The description is: {description}.
 
 You are able to use the dataframe to answer the question.
 You have to act like an data analyst who can do an effective analysis through dataframe.
+
+USER INPUT
+============
 """

core/prompts/input.py

@@ -1,7 +1,7 @@
 EVAL_PREFIX = """{bot_name} can execute any user's request.
 
-{bot_name} has permission to handle one instance and can handle the environment in it at will. File creation is only possible in playground folder, and other folders and files can never be modified under any circumstances. 
-You can code, run, debug, and test yourself. You can correct the code appropriately by looking at the error message. Please do TDD.
+{bot_name} has permission to handle one instance and can handle the environment in it at will.
+You can code, run, debug, and test yourself. You can correct the code appropriately by looking at the error message.
 
 I can understand, process, and create various types of files. Images must be stored in the ./image/, audio in the ./audio/, video in the ./video/, and dataframes must be stored in the ./dataframe/.
 

core/tools/editor/__init__.py

@@ -55,16 +55,12 @@ class CodeEditor(BaseToolSet):
         "Input should be filename and code to append. "
         "Input code must be the code that should be appended, NOT whole code. "
         "ex. test.py\nprint('hello world')\n "
-        "and the output will be last 3 line.",
+        "and the output will be last 3 lines.",
     )
     def append(self, inputs: str) -> str:
         try:
             code = CodeWriter.append(inputs)
-            output = (
-                "Last 3 line was:\n"
-                + "\n".join(code.split("\n")[-3:])
-                + "\nYou can use CodeEditor.APPEND tool to append the code if this file is not completed."
-            )
+            output = "Last 3 line was:\n" + "\n".join(code.split("\n")[-3:])
         except Exception as e:
             output = str(e)
 
@@ -80,16 +76,12 @@ class CodeEditor(BaseToolSet):
         "If the code is completed, use the Terminal tool to execute it, if not, append the code through the CodeEditor.APPEND tool. "
         "Input should be filename and code. This file must be in playground folder. "
         "ex. test.py\nprint('hello world')\n "
-        "and the output will be last 3 line.",
+        "and the output will be last 3 lines.",
     )
     def write(self, inputs: str) -> str:
         try:
             code = CodeWriter.write(inputs)
-            output = (
-                "Last 3 line was:\n"
-                + "\n".join(code.split("\n")[-3:])
-                + "\nYou can use CodeEditor.APPEND tool to append the code if this file is not completed."
-            )
+            output = "Last 3 line was:\n" + "\n".join(code.split("\n")[-3:])
         except Exception as e:
             output = str(e)
 
@@ -136,7 +128,6 @@ class CodeEditor(BaseToolSet):
         "Output will be success or error message.",
     )
     def delete(self, inputs: str) -> str:
-        filepath: str = str(Path(settings["PLAYGROUND_DIR"]) / Path(inputs))
         try:
             with open(filepath, "w") as f:
                 f.write("")

core/tools/editor/patch.py

@@ -63,6 +63,8 @@ from typing import Tuple
 
 from env import settings
 
+from .verify import verify
+
 
 class Position:
     separator = ","
@@ -84,7 +86,7 @@ class PatchCommand:
     separator = "|"
 
     def __init__(self, filepath: str, start: Position, end: Position, content: str):
-        self.filepath: str = str(Path(settings["PLAYGROUND_DIR"]) / Path(filepath))
+        self.filepath: str = filepath
         self.start: Position = start
         self.end: Position = end
         self.content: str = content
@@ -99,13 +101,8 @@ class PatchCommand:
             f.writelines(lines)
         return sum([len(line) for line in lines])
 
+    @verify
     def execute(self) -> Tuple[int, int]:
-        # make sure the directory exists
-        if not str(Path(self.filepath).resolve()).startswith(
-            str(Path(settings["PLAYGROUND_DIR"]).resolve())
-        ):
-            return "You can't write file outside of current directory."
-
         os.makedirs(os.path.dirname(self.filepath), exist_ok=True)
         lines = self.read_lines()
         before = sum([len(line) for line in lines])

core/tools/editor/read.py

@@ -8,6 +8,8 @@ from typing import List, Optional, Tuple
 
 from env import settings
 
+from .verify import verify
+
 
 class Line:
     def __init__(self, content: str, line_number: int, depth: int):
@@ -104,16 +106,12 @@ class ReadCommand:
     separator = "|"
 
     def __init__(self, filepath: str, start: int, end: int):
-        self.filepath: str = str(Path(settings["PLAYGROUND_DIR"]) / Path(filepath))
+        self.filepath: str = filepath
         self.start: int = start
         self.end: int = end
 
+    @verify
     def execute(self) -> str:
-        if not str(Path(self.filepath).resolve()).startswith(
-            str(Path(settings["PLAYGROUND_DIR"]).resolve())
-        ):
-            return "You can't write file outside of current directory."
-
         with open(self.filepath, "r") as f:
             code = f.readlines()
 
@@ -134,16 +132,12 @@ class SummaryCommand:
     separator = "|"
 
     def __init__(self, filepath: str, depth: int, parent_content: Optional[str] = None):
-        self.filepath: str = str(Path(settings["PLAYGROUND_DIR"]) / Path(filepath))
+        self.filepath: str = filepath
         self.depth: int = depth
         self.parent_content: Optional[str] = parent_content
 
+    @verify
     def execute(self) -> str:
-        if not str(Path(self.filepath).resolve()).startswith(
-            str(Path(settings["PLAYGROUND_DIR"]).resolve())
-        ):
-            return "You can't write file outside of current directory."
-
         with open(self.filepath, "r") as f:
             code = f.readlines()
 
@@ -152,8 +146,6 @@ class SummaryCommand:
             if line.strip() != "":
                 code_tree.append(line, i + 1)
 
-        # code_tree.print()
-
         if self.parent_content is None:
             lines = code_tree.find_from_root(self.depth)
         else:

core/tools/editor/verify.py

@@ -0,0 +1,18 @@
+from pathlib import Path
+
+from env import settings
+
+
+def verify(func):
+    def wrapper(*args, **kwargs):
+        try:
+            filepath = args[0].filepath
+        except:
+            raise Exception("This tool doesn't have filepath. Please check your code.")
+        if not str(Path(filepath).resolve()).startswith(
+            str(Path(settings["PLAYGROUND_DIR"]).resolve())
+        ):
+            return "You can't access file outside of playground."
+        return func(*args, **kwargs)
+
+    return wrapper

core/tools/editor/write.py

@@ -9,12 +9,14 @@ from pathlib import Path
 
 from env import settings
 
+from .verify import verify
+
 
 class WriteCommand:
     separator = "\n"
 
     def __init__(self, filepath: str, content: int):
-        self.filepath: str = str(Path(settings["PLAYGROUND_DIR"]) / Path(filepath))
+        self.filepath: str = filepath
         self.content: str = content
         self.mode: str = "w"
 
@@ -22,14 +24,8 @@ class WriteCommand:
         self.mode = mode
         return self
 
+    @verify
     def execute(self) -> str:
-        # make sure the directory exists
-        if not str(Path(self.filepath).resolve()).startswith(
-            str(Path(settings["PLAYGROUND_DIR"]).resolve())
-        ):
-            return "You can't write file outside of current directory."
-
-        os.makedirs(os.path.dirname(self.filepath), exist_ok=True)
         with open(self.filepath, self.mode) as f:
             f.write(self.content)
         return self.content

core/tools/terminal/__init__.py

@@ -31,7 +31,6 @@ class Terminal(BaseToolSet):
             process = subprocess.Popen(
                 commands,
                 shell=True,
-                cwd=settings["PLAYGROUND_DIR"],
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
             )

env.py

@@ -9,7 +9,7 @@ load_dotenv()
 class DotEnv(TypedDict):
     OPENAI_API_KEY: str
 
-    PORT: int
+    EVAL_PORT: int
     SERVER: str
 
     USE_GPU: bool  # optional
@@ -27,10 +27,10 @@ class DotEnv(TypedDict):
     SERPAPI_API_KEY: str  # optional
 
 
-PORT = int(os.getenv("PORT", 8000))
+EVAL_PORT = int(os.getenv("EVAL_PORT", 8000))
 settings: DotEnv = {
-    "PORT": PORT,
-    "SERVER": os.getenv("SERVER", f"http://localhost:{PORT}"),
+    "EVAL_PORT": EVAL_PORT,
+    "SERVER": os.getenv("SERVER", f"http://localhost:{EVAL_PORT}"),
     "USE_GPU": os.getenv("USE_GPU", "False").lower() == "true",
     "PLAYGROUND_DIR": os.getenv("PLAYGROUND_DIR", "playground"),
     "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY"),

utils.py

@@ -4,12 +4,6 @@ import uuid
 
 import numpy as np
 
-os.makedirs("image", exist_ok=True)
-os.makedirs("audio", exist_ok=True)
-os.makedirs("video", exist_ok=True)
-os.makedirs("dataframe", exist_ok=True)
-os.makedirs("playground", exist_ok=True)
-
 
 def seed_everything(seed):
     random.seed(seed)
@@ -24,15 +18,6 @@ def seed_everything(seed):
     return seed
 
 
-def prompts(name, description):
-    def decorator(func):
-        func.name = name
-        func.description = description
-        return func
-
-    return decorator
-
-
 def cut_dialogue_history(history_memory, keep_last_n_words=500):
     tokens = history_memory.split()
     n_tokens = len(tokens)