(core) Improve parsing formula from completion

Summary:
The previous code for extracting a Python formula from the LLM completion involved some shaky string manipulation which this improves on.
Overall the 'test results' from `runCompletion` went from 37/47 to 45/47 for `gpt-3.5-turbo-0613`.

The biggest problem that motivated these changes was that it assumed that code was always inside a markdown code block
(i.e. triple backticks) and so if there was no block there was no code. But the completion often consists of *only* code
with no accompanying explanation or markdown. By parsing the completion in Python instead of JS,
we can easily check if the entire completion is valid Python syntax and accept it if it is.

I also noticed one failure resulting from the completion containing the full function (instead of just the body)
and necessary imports before that function instead of inside. The new parsing moves import inside.

Test Plan: Added a Python unit test

Reviewers: paulfitz

Reviewed By: paulfitz

Subscribers: paulfitz

Differential Revision: https://phab.getgrist.com/D3922

sandbox/grist/formula_prompt.py

@@ -1,6 +1,9 @@
+import ast
 import json
+import re
 import textwrap
 
+import asttokens
 import six
 
 from column import is_visible_column, BaseReferenceColumn
@@ -180,6 +183,57 @@ def indent(text, prefix, predicate=None):
       yield (prefix + line if predicate(line) else line)
   return ''.join(prefixed_lines())
 
+
 def convert_completion(completion):
+  # Extract code from a markdown code block if needed.
+  match = re.search(r"```\w*\n(.*)```", completion, re.DOTALL)
+  if match:
+    completion = match.group(1)
+
   result = textwrap.dedent(completion)
-  return result
+
+  try:
+    atok = asttokens.ASTTokens(result, parse=True)
+  except SyntaxError:
+    # If we don't have valid Python code, don't suggest a formula at all
+    return ""
+
+  stmts = atok.tree.body
+
+  # If the code starts with imports, save them for later.
+  # In particular, the model may return something like:
+  #  from datetime import date
+  #  def my_column():
+  #     ...
+  # We want to return just the function body, but we need to keep the import,
+  # i.e. move it 'inside the function'.
+  imports = ""
+  while stmts and isinstance(stmts[0], (ast.Import, ast.ImportFrom)):
+    imports += atok.get_text(stmts.pop(0)) + "\n"
+
+  # If the non-import code consists only of a function definition, extract the body.
+  if len(stmts) == 1 and isinstance(stmts[0], ast.FunctionDef):
+    func_body_stmts = stmts[0].body
+    if (
+      len(func_body_stmts) > 1 and
+      isinstance(func_body_stmts[0], ast.Expr) and
+      isinstance(func_body_stmts[0].value, ast.Str)
+    ):
+      # Skip the docstring.
+      first_stmt = func_body_stmts[1]
+    else:
+      first_stmt = func_body_stmts[0]
+    result_lines = result.splitlines()[first_stmt.lineno - 1:]
+    result = "\n".join(result_lines)
+    result = textwrap.dedent(result)
+
+    if imports:
+      result = imports + "\n" + result
+
+    # Check that we still have valid code.
+    try:
+      ast.parse(result)
+    except SyntaxError:
+      return ""
+
+  return result.strip()

sandbox/grist/test_formula_prompt.py

@@ -5,7 +5,7 @@ import test_engine
 import testutil
 
 from formula_prompt import (
-  values_type, column_type, referenced_tables, get_formula_prompt,
+  values_type, column_type, referenced_tables, get_formula_prompt, convert_completion,
 )
 from objtypes import RaisedException
 from records import Record as BaseRecord, RecordSet as BaseRecordSet
@@ -223,3 +223,33 @@ class Table3:
         description here
         """
 ''')
+
+  def test_convert_completion(self):
+    completion = """
+Here's some code:
+
+```python
+import os
+from x import (
+  y,
+  z,
+)
+
+@property
+def foo():
+    '''This is a docstring'''
+    x = 5
+    return 1
+```
+
+Hope you like it!
+"""
+    self.assertEqual(convert_completion(completion), """\
+import os
+from x import (
+  y,
+  z,
+)
+
+x = 5
+return 1""")