fix: improved transform_examples_to_chat_directives to be more robust

ErikBjare · Aug 8, 2024 · b86b3ce · b86b3ce
1 parent f0a581d
commit b86b3ce
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 9 deletions.
diff --git a/gptme/tools/browser.py b/gptme/tools/browser.py
@@ -62,9 +62,9 @@
 
 ### Searching
 User: who is the founder of ActivityWatch?
-Let's search for that.
+Assistant: Let's search for that.
 ```python
-search("who is the founder of ActivityWatch?")
+search("ActivityWatch founder")
 ```
 System:
 ```Results:

diff --git a/gptme/tools/python.py b/gptme/tools/python.py
@@ -62,7 +62,7 @@ def get_installed_python_libraries() -> set[str]:
 """
 
 examples = """
-#### Results of the last expression will be displayed, IPython-style.
+#### Results of the last expression will be displayed, IPython-style:
 User: What is 2 + 2?
 Assistant:
 ```python
@@ -73,7 +73,7 @@ def get_installed_python_libraries() -> set[str]:
 4
 ```
 
-The user can also run Python code with the /python command:
+#### The user can also run Python code with the /python command:
 
 User: /python 2 + 2
 System: Executed code block.

diff --git a/gptme/tools/terminal.py b/gptme/tools/terminal.py
@@ -224,18 +224,19 @@ def execute_terminal(
 ? Project name: › vue-project
 ```
 Assistant: vue-project is a placeholder we can fill in. What would you like to name your project?
-User: Lets go with 'test-project'
+User: fancy-project
 Assistant:
 ```terminal
-send_keys 0 test-project Enter
+send_keys 0 fancy-project Enter
 ```
-System: Sent 'test-project Enter' to pane 0
+System: Sent 'fancy-project Enter' to pane 0
 Assistant: Lets check that the name was set and move on
 ```terminal
 inspect_pane 0
 ```
 System:
 ```output
+? Project name: › fancy-project
 ? Add TypeScript? › No / Yes
 ```
 Assistant: The project name has been set, now we choose TypeScript

diff --git a/gptme/util.py b/gptme/util.py
@@ -159,12 +159,21 @@ def ask_execute(question="Execute code?", default=True) -> bool:  # pragma: no c
 def transform_examples_to_chat_directives(s: str, strict=False) -> str:
     # transforms an example with "> Role:" dividers into ".. chat::" directive
     orig = s
-    s = re.sub(r"(^|\n)([>] )?(.+):", r"\1\3:", s, re.DOTALL)
+    s = re.sub(
+        r"(^|\n)([>] )?(.+):",
+        r"\1\3:",
+        s,
+        re.DOTALL,
+    )
     if strict:
         assert s != orig, "Couldn't find a message"
     s = textwrap.indent(s, "   ")
     orig = s
-    s = re.sub(r"(^|\n)(   [# ]+(.+)\n)?   User:", r"\1\3\n\n.. chat::\n\n   User:", s)
+    s = re.sub(
+        r"(^|\n)(   [# ]+(.+)(\n\s*)+)?   User:",
+        r"\1\3\n\n.. chat::\n\n   User:",
+        s,
+    )
     if strict:
         assert s != orig, "Couldn't find place to put start of directive"
     return s