don't call python function, other

4 files changed, 17 insertions, 6 deletions

diff --git a/continuedev/src/continuedev/libs/llm/openai.py b/continuedev/src/continuedev/libs/llm/openai.py
index aa12d70a..7621111f 100644
--- a/continuedev/src/continuedev/libs/llm/openai.py
+++ b/continuedev/src/continuedev/libs/llm/openai.py
@@ -70,7 +70,7 @@ class OpenAI(LLM):
                 messages=compile_chat_messages(
                     args["model"], with_history, prompt, with_functions=False),
                 **args,
-            )).choices[0].message
+            )).choices[0].message.content
         else:
             resp = (await openai.Completion.acreate(
                 prompt=prompt,
diff --git a/continuedev/src/continuedev/steps/chat.py b/continuedev/src/continuedev/steps/chat.py
index 76fead9d..a940c3ba 100644
--- a/continuedev/src/continuedev/steps/chat.py
+++ b/continuedev/src/continuedev/steps/chat.py
@@ -105,7 +105,7 @@ class ViewDirectoryTreeStep(Step):
 
 class EditFileStep(Step):
     name: str = "Edit File"
-    description: str = "Edit a file in the workspace."
+    description: str = "Edit a file in the workspace that is not currently open."
     filename: str
     instructions: str
     hide: bool = True
@@ -136,7 +136,7 @@ class ChatWithFunctions(Step):
 
         self.chat_context.append(ChatMessage(
             role="user",
-            content=self.user_input,
+            content=self.user_input + "\n**DO NOT EVER call the 'python' function.**",
             summary=self.user_input
         ))
 
@@ -183,6 +183,15 @@ class ChatWithFunctions(Step):
                 if func_name == "python" and "python" not in step_name_step_class_map:
                     # GPT must be fine-tuned to believe this exists, but it doesn't always
                     self.chat_context.append(ChatMessage(
+                        role="assistant",
+                        content=None,
+                        function_call=FunctionCall(
+                            name=func_name,
+                            arguments=func_args
+                        ),
+                        summary=f"Ran function {func_name}"
+                    ))
+                    self.chat_context.append(ChatMessage(
                         role="user",
                         content="The 'python' function does not exist. Don't call it.",
                         summary="'python' function does not exist."
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index 7ccf82cb..f146c94a 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -213,7 +213,8 @@ class DefaultModelEditCodeStep(Step):
 
             if model_to_use.name == "gpt-4":
 
-                total_tokens = model_to_use.count_tokens(full_file_contents + self._prompt)
+                total_tokens = model_to_use.count_tokens(
+                    full_file_contents + self._prompt)
                 cur_start_line, cur_end_line = cut_context(
                     model_to_use, total_tokens, cur_start_line, cur_end_line)
 
diff --git a/continuedev/src/continuedev/steps/main.py b/continuedev/src/continuedev/steps/main.py
index bb48dce9..5caac180 100644
--- a/continuedev/src/continuedev/steps/main.py
+++ b/continuedev/src/continuedev/steps/main.py
@@ -1,7 +1,7 @@
 import os
 from typing import Coroutine, List, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from ..libs.llm import LLM
 from ..models.main import Traceback, Range
@@ -246,7 +246,8 @@ class StarCoderEditHighlightedCodeStep(Step):
 
 
 class EditHighlightedCodeStep(Step):
-    user_input: str
+    user_input: str = Field(
+        ..., title="User Input", description="The natural language request describing how to edit the code")
     hide = True
     description: str = "Change the contents of the currently highlighted code or open file"