parse the outermost ```s so markdown still works

thorstone137 · Jan 31, 2024 · 057b1f4 · 057b1f4
1 parent bb988ba
commit 057b1f4
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/src/rawdog/llm_client.py b/src/rawdog/llm_client.py
@@ -2,6 +2,7 @@
 import datetime
 import json
 import os
+import re
 from textwrap import dedent
 
 from litellm import completion, completion_cost
@@ -68,11 +69,12 @@ def get_response(
             timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
             script_filename = self.log_path.parent / f"script_{timestamp}.py"
 
+            script = None if not text else re.search(r'```(.*?)```', text, re.DOTALL)
             script_content = dedent(f"""\
-                    # Model: {log["model"]}
-                    # Prompt: {log["prompt"]}
-                    # Response Cost: {log.get("cost", "N/A")}
-                    """)+text.split("```")[1]
+                # Model: {log["model"]}
+                # Prompt: {log["prompt"]}
+                # Response Cost: {log.get("cost", "N/A")}
+                """) + script.group(1) if script else f"INVALID SCRIPT:\n{text}"
             with open(script_filename, "w") as script_file:
                 script_file.write(script_content)
             return text
@@ -93,7 +95,10 @@ def get_script(self, prompt: str):
         error = None
         script = None
         if response.count("```") > 1:
-            script = response.split("```")[1]
+            script = re.search(r'```(.*?)```', response, re.DOTALL)
+            if not script:
+                error = f"No script found in response: {response}"
+            script = script.group(1)
             script = dedent(script).strip()
             if script.split("\n")[0].startswith("python"):
                 script = "\n".join(script.split("\n")[1:])