Common hallucination fix

Author: KillianLucas

interpreter/core/respond.py

@@ -150,14 +150,34 @@ def respond(interpreter):
 
                 # A common hallucination
                 if code.startswith("functions.execute("):
-                    code = code.replace("functions.execute(", "").rstrip(")")
-                    code_dict = json.loads(code)
-                    language = code_dict.get("language", language)
-                    code = code_dict.get("code", code)
-                    interpreter.messages[-1]["content"] = code  # So the LLM can see it.
-                    interpreter.messages[-1][
-                        "format"
-                    ] = language  # So the LLM can see it.
+                    edited_code = code.replace("functions.execute(", "").rstrip(")")
+                    try:
+                        code_dict = json.loads(edited_code)
+                        language = code_dict.get("language", language)
+                        code = code_dict.get("code", code)
+                        interpreter.messages[-1][
+                            "content"
+                        ] = code  # So the LLM can see it.
+                        interpreter.messages[-1][
+                            "format"
+                        ] = language  # So the LLM can see it.
+                    except:
+                        pass
+
+                if code.replace("\n", "").replace(" ", "").startswith('{"language":'):
+                    try:
+                        code_dict = json.loads(code)
+                        if set(code_dict.keys()) == {"language", "code"}:
+                            language = code_dict["language"]
+                            code = code_dict["code"]
+                            interpreter.messages[-1][
+                                "content"
+                            ] = code  # So the LLM can see it.
+                            interpreter.messages[-1][
+                                "format"
+                            ] = language  # So the LLM can see it.
+                    except:
+                        pass
 
                 if language == "text" or language == "markdown":
                     # It does this sometimes just to take notes. Let it, it's useful.

tests/test_interpreter.py

@@ -22,6 +22,36 @@
 from websocket import create_connection
 
 
+def test_hallucinations():
+    # We should be resiliant to common hallucinations.
+
+    code = """{                                                                             
+    "language": "python",                                                        
+    "code": "10+12"                                                        
+  }"""
+
+    interpreter.messages = [
+        {"role": "assistant", "type": "code", "format": "python", "content": code}
+    ]
+    for chunk in interpreter._respond_and_store():
+        if chunk.get("format") == "output":
+            assert chunk.get("content") == "22"
+            break
+
+    code = """functions.execute({                                                                             
+    "language": "python",                                                        
+    "code": "10+12"                                                        
+  })"""
+
+    interpreter.messages = [
+        {"role": "assistant", "type": "code", "format": "python", "content": code}
+    ]
+    for chunk in interpreter._respond_and_store():
+        if chunk.get("format") == "output":
+            assert chunk.get("content") == "22"
+            break
+
+
 @pytest.mark.skip(reason="Requires uvicorn, which we don't require by default")
 def test_server():
     # os.system("pip install 'open-interpreter[server]'")