From faf0035323ce04afe287324b3df54e36566b3488 Mon Sep 17 00:00:00 2001
From: Aman Rusia <aman@kombai.io>
Date: Mon, 18 Nov 2024 17:07:27 +0530
Subject: [PATCH] Added CreateFileNew for claude

---
 src/wcgw/client/anthropic_client.py | 12 +++++++-----
 src/wcgw/client/tools.py            | 26 ++++++++++++++++++++++----
 src/wcgw/types_.py                  |  5 +++++
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/src/wcgw/client/anthropic_client.py b/src/wcgw/client/anthropic_client.py
index 266d62c..097d6f7 100644
--- a/src/wcgw/client/anthropic_client.py
+++ b/src/wcgw/client/anthropic_client.py
@@ -24,6 +24,7 @@
 from ..types_ import (
     BashCommand,
     BashInteraction,
+    CreateFileNew,
     FileEditFindReplace,
     ReadImage,
     Writefile,
@@ -178,12 +179,14 @@ def loop(
 """,
         ),
         ToolParam(
-            input_schema=Writefile.model_json_schema(),
-            name="WriteFile",
+            input_schema=CreateFileNew.model_json_schema(),
+            name="CreateFileNew",
             description="""
-- Write content to a file. Provide file path and content. Use this instead of BashCommand for writing files.
+- Write content to a new file. Provide file path and content. Use this instead of BashCommand for writing new files.
 - This doesn't create any directories, please create directories using `mkdir -p` BashCommand.
-- Provide absolute file path only.""",
+- Provide absolute file path only.
+- For editing existing files, use FileEditFindReplace.
+""",
         ),
         ToolParam(
             input_schema=ReadImage.model_json_schema(),
@@ -202,7 +205,6 @@ def loop(
 - Find and replace multiple lines in a file.
 - Use absolute file path only.
 - Replaces complete lines.
-- Prefer this over WriteFile if edits are to be made on large files.
 """,
         ),
     ]
diff --git a/src/wcgw/client/tools.py b/src/wcgw/client/tools.py
index 93b6886..4654dce 100644
--- a/src/wcgw/client/tools.py
+++ b/src/wcgw/client/tools.py
@@ -43,7 +43,7 @@
     ParsedChatCompletionMessage,
 )
 from nltk.metrics.distance import edit_distance
-from ..types_ import FileEditFindReplace, ResetShell, Writefile
+from ..types_ import CreateFileNew, FileEditFindReplace, ResetShell, Writefile
 
 from ..types_ import BashCommand
 
@@ -415,12 +415,17 @@ def read_image_from_shell(file_path: str) -> ImageData:
         return ImageData(media_type=image_type, data=image_b64)
 
 
-def write_file(writefile: Writefile) -> str:
+def write_file(writefile: Writefile | CreateFileNew, error_on_exist: bool) -> str:
     if not os.path.isabs(writefile.file_path):
         return "Failure: file_path should be absolute path"
     else:
         path_ = writefile.file_path
 
+    if error_on_exist and os.path.exists(path_):
+        file_data = Path(path_).read_text()
+        if file_data:
+            return f"Error: can't write to existing file {path_}, use other functions to edit the file"
+
     path = Path(path_)
     path.parent.mkdir(parents=True, exist_ok=True)
 
@@ -555,6 +560,7 @@ def get_tool_output(
     | BashInteraction
     | ResetShell
     | Writefile
+    | CreateFileNew
     | FileEditFindReplace
     | AIAssistant
     | DoneFlag
@@ -571,6 +577,7 @@ def get_tool_output(
             | BashInteraction
             | ResetShell
             | Writefile
+            | CreateFileNew
             | FileEditFindReplace
             | AIAssistant
             | DoneFlag
@@ -581,6 +588,7 @@ def get_tool_output(
             | BashInteraction
             | ResetShell
             | Writefile
+            | CreateFileNew
             | FileEditFindReplace
             | AIAssistant
             | DoneFlag
@@ -598,7 +606,10 @@ def get_tool_output(
         output = execute_bash(enc, arg, max_tokens)
     elif isinstance(arg, Writefile):
         console.print("Calling write file tool")
-        output = write_file(arg), 0
+        output = write_file(arg, False), 0
+    elif isinstance(arg, CreateFileNew):
+        console.print("Calling write file tool")
+        output = write_file(arg, True), 0
     elif isinstance(arg, FileEditFindReplace):
         console.print("Calling file edit tool")
         output = file_edit(arg), 0.0
@@ -630,7 +641,14 @@ def get_tool_output(
 
 
 class Mdata(BaseModel):
-    data: BashCommand | BashInteraction | Writefile | ResetShell | FileEditFindReplace
+    data: (
+        BashCommand
+        | BashInteraction
+        | Writefile
+        | CreateFileNew
+        | ResetShell
+        | FileEditFindReplace
+    )
 
 
 execution_lock = threading.Lock()
diff --git a/src/wcgw/types_.py b/src/wcgw/types_.py
index aa24e77..adebb19 100644
--- a/src/wcgw/types_.py
+++ b/src/wcgw/types_.py
@@ -43,6 +43,11 @@ class Writefile(BaseModel):
     file_content: str
 
 
+class CreateFileNew(BaseModel):
+    file_path: str
+    file_content: str
+
+
 class FileEditFindReplace(BaseModel):
     file_path: str
     find_lines: str