gptdiff 0.1.8__tar.gz → 0.1.9__tar.gz

{gptdiff-0.1.8 → gptdiff-0.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: gptdiff
-Version: 0.1.8
+Version: 0.1.9
 Summary: A tool to generate and apply git diffs using LLMs
 Author: 255labs
 Classifier: License :: OSI Approved :: MIT License

{gptdiff-0.1.8 → gptdiff-0.1.9}/gptdiff/gptdiff.py

@@ -67,6 +67,25 @@ def create_think_toolbox():
     )
     return toolbox
 
+def color_code_diff(diff_text: str) -> str:
+    """
+    Color code lines in a diff. Lines beginning with '-' in red, and
+    lines beginning with '+' in green.
+    """
+    red = "\033[31m"
+    green = "\033[32m"
+    reset = "\033[0m"
+
+    colorized_lines = []
+    for line in diff_text.split('\n'):
+        if line.startswith('-'):
+            colorized_lines.append(f"{red}{line}{reset}")
+        elif line.startswith('+'):
+            colorized_lines.append(f"{green}{line}{reset}")
+        else:
+            colorized_lines.append(line)
+
+    return '\n'.join(colorized_lines)
 
 def load_gitignore_patterns(gitignore_path):
     with open(gitignore_path, 'r') as f:
@@ -308,8 +327,16 @@ def smartapply(diff_text, files, model=None, api_key=None, base_url=None):
             updated = call_llm_for_apply_with_think_tool_available(path, original, patch, model, api_key=api_key, base_url=base_url)
             files[path] = updated.strip()
 
+    threads = []
+
     for path, patch in parsed_diffs:
-        process_file(path, patch)
+        thread = threading.Thread(target=process_file, args=(path, patch))
+        thread.start()
+        threads.append(thread)
+
+    # Wait for all threads to complete
+    for thread in threads:
+        thread.join()
 
     return files
 
@@ -339,6 +366,7 @@ def parse_arguments():
     parser.add_argument('--temperature', type=float, default=0.7, help='Temperature parameter for model creativity (0.0 to 2.0)')
     parser.add_argument('--max_tokens', type=int, default=30000, help='Temperature parameter for model creativity (0.0 to 2.0)')
     parser.add_argument('--model', type=str, default=None, help='Model to use for the API call.')
+    parser.add_argument('--applymodel', type=str, default=None, help='Model to use for applying the diff. Defaults to the value of --model if not specified.')
 
     parser.add_argument('--nowarn', action='store_true', help='Disable large token warning')
 
@@ -404,11 +432,11 @@ def parse_diff_per_file(diff_text):
 
     return diffs
 
-def call_llm_for_apply_with_think_tool_available(file_path, original_content, file_diff, model, api_key=None, base_url=None):
+def call_llm_for_apply_with_think_tool_available(file_path, original_content, file_diff, model, api_key=None, base_url=None, extra_prompt=None, max_tokens=30000):
     parser = FlatXMLParser("think")
     formatter = FlatXMLPromptFormatter(tag="think")
     toolbox = create_think_toolbox()
-    full_response = call_llm_for_apply(file_path, original_content, file_diff, model, api_key=None, base_url=None)
+    full_response = call_llm_for_apply(file_path, original_content, file_diff, model, api_key=api_key, base_url=base_url, extra_prompt=extra_prompt, max_tokens=max_tokens)
     notool_response = ""
     events = parser.parse(full_response)
     is_in_tool = False
@@ -424,7 +452,7 @@ def call_llm_for_apply_with_think_tool_available(file_path, original_content, fi
 
     return notool_response
 
-def call_llm_for_apply(file_path, original_content, file_diff, model, api_key=None, base_url=None):
+def call_llm_for_apply(file_path, original_content, file_diff, model, api_key=None, base_url=None, extra_prompt=None, max_tokens=30000):
     """AI-powered diff application with conflict resolution.
     
     Internal workhorse for smartapply that handles individual file patches.
@@ -474,6 +502,8 @@ Diff to apply:
 {file_diff}
 </diff>"""
 
+    if extra_prompt:
+        user_prompt += f"\n\n{extra_prompt}"
     if model == "gemini-2.0-flash-thinking-exp-01-21":
         user_prompt = system_prompt+"\n"+user_prompt
     messages = [
@@ -490,7 +520,7 @@ Diff to apply:
     response = client.chat.completions.create(model=model,
         messages=messages,
         temperature=0.0,
-        max_tokens=30000)
+        max_tokens=max_tokens)
     full_response = response.choices[0].message.content
 
     elapsed = time.time() - start_time
@@ -590,12 +620,21 @@ def main():
             if confirmation != 'y':
                 print("Request canceled")
                 sys.exit(0)
-        full_text, diff_text, prompt_tokens, completion_tokens, total_tokens, cost = call_llm_for_diff(system_prompt, user_prompt, files_content, args.model, 
+        try:
+            full_text, diff_text, prompt_tokens, completion_tokens, total_tokens, cost = call_llm_for_diff(system_prompt, user_prompt, files_content, args.model, 
                                                                                                     temperature=args.temperature,
                                                                                                     api_key=os.getenv('GPTDIFF_LLM_API_KEY'),
                                                                                                     base_url=os.getenv('GPTDIFF_LLM_BASE_URL', "https://nano-gpt.com/api/v1/"),
                                                                                                     max_tokens=args.max_tokens
                                                                                                     ) 
+        except Exception as e:
+            full_text = f"{e}"
+            diff_text = ""
+            prompt_tokens = 0
+            completion_tokens = 0
+            total_tokens = 0
+            cost = 0
+            print(f"Error in LLM response {e}")
 
     if(diff_text.strip() == ""):
         print(f"\033[1;33mThere was no data in this diff. The LLM may have returned something invalid.\033[0m")
@@ -608,7 +647,7 @@ def main():
     elif args.apply:
         print("\nAttempting apply with the following diff:")
         print("\n<diff>")
-        print(diff_text)
+        print(color_code_diff(diff_text))
         print("\n</diff>")
         print("Saved to patch.diff")
         if apply_diff(project_dir, diff_text):
@@ -651,8 +690,11 @@ def main():
                 print("SMARTAPPLY")
                 print(file_diff)
                 print("-" * 40)
+                if args.applymodel is None:
+                    args.applymodel = args.model
+
                 try:
-                    updated_content = call_llm_for_apply_with_think_tool_available(file_path, original_content, file_diff, args.model)
+                    updated_content = call_llm_for_apply_with_think_tool_available(file_path, original_content, file_diff, args.applymodel, extra_prompt=f"This changeset is from the following instructions:\n{user_prompt}", max_tokens=args.max_tokens)
 
                     if updated_content.strip() == "":
                         print("Cowardly refusing to write empty file to", file_path, "merge failed")

{gptdiff-0.1.8 → gptdiff-0.1.9}/gptdiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: gptdiff
-Version: 0.1.8
+Version: 0.1.9
 Summary: A tool to generate and apply git diffs using LLMs
 Author: 255labs
 Classifier: License :: OSI Approved :: MIT License

{gptdiff-0.1.8 → gptdiff-0.1.9}/gptdiff.egg-info/SOURCES.txt

@@ -9,4 +9,5 @@ gptdiff.egg-info/dependency_links.txt
 gptdiff.egg-info/entry_points.txt
 gptdiff.egg-info/requires.txt
 gptdiff.egg-info/top_level.txt
+tests/test_diff_parse.py
 tests/test_smartapply.py

{gptdiff-0.1.8 → gptdiff-0.1.9}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='gptdiff',
-    version='0.1.8',
+    version='0.1.9',
     description='A tool to generate and apply git diffs using LLMs',
     author='255labs',
     packages=find_packages(),  # Use find_packages() to automatically discover packages

gptdiff-0.1.9/tests/test_diff_parse.py

@@ -0,0 +1,35 @@
+import pytest
+from gptdiff.gptdiff import color_code_diff
+
+def test_color_code_diff():
+    """
+    Test that lines starting with '-' are wrapped in red color codes
+    and lines starting with '+' are wrapped in green color codes.
+    """
+    diff_text = """-this line is removed
++this line is added
+ unchanged line
+-removed again
++and added again
+ some other neutral line"""
+
+    colorized = color_code_diff(diff_text)
+
+    # We expect lines beginning with '-' to be in red
+    assert "\033[31m-this line is removed\033[0m" in colorized
+    assert "\033[31m-removed again\033[0m" in colorized
+
+    # We expect lines beginning with '+' to be in green
+    assert "\033[32m+this line is added\033[0m" in colorized
+    assert "\033[32m+and added again\033[0m" in colorized
+
+    # Lines unchanged should remain uncolored
+    assert "unchanged line" in colorized
+    assert "some other neutral line" in colorized
+
+    # Ensure no erroneous color codes are added
+    # by counting them in the final output
+    color_code_count = colorized.count('\033[')
+    # We have four lines that should be color-coded, so we expect 4 * 2 = 8 color code inserts
+    # (start code + reset code per line)
+    assert color_code_count == 8

{gptdiff-0.1.8 → gptdiff-0.1.9}/tests/test_smartapply.py

@@ -160,7 +160,7 @@ diff --git a/file2.py b/file2.py
     }
 
     # Mock LLM to return modified content based on file path
-    def mock_call_llm(file_path, original_content, file_diff, model, api_key, base_url):
+    def mock_call_llm(file_path, original_content, file_diff, model, api_key, base_url, extra_prompt=None, max_tokens=None):
         if file_path == "file1.py":
             return "def func1():\n    print('New func1')"
         elif file_path == "file2.py":