add output condition

app.py

@@ -114,7 +114,10 @@ real_docstring = r"""
 def docgen_func(function_code):
     req_data = {"inputs": function_code}
     output = query(req_data)
-    return f'"""\n{output[0]["generated_text"]}\n"""'
+    if type(output) is list:
+        return f'"""\n{output[0]["generated_text"]}\n"""'
+    else:
+        return str(output)
 
 def pygen_func(nl_code_intent):
     pass # TODO: generate code PL from intent NL + search in corpus