Logs to Langsmith

app.py

@@ -93,11 +93,10 @@ class SQLExecutorTool(Tool):
 tool = SQLExecutorTool()
 
 def process_outputs(output) :
-    if 'data' in output:
-        output['data'] = "<DataFrame is hidden>"
-    if 'fig' in output:
-        output['fig'] = "<Figure is hidden>"
-    return output
+    return {
+        'sql': output.get('sql', None),
+        'code': output.get('code', None)
+    }
 
 @traceable(process_outputs=process_outputs)
 def get_visualization(question, schema, table_name):
@@ -106,7 +105,7 @@ def get_visualization(question, schema, table_name):
                                                  'pandas', 'plotly.express',
                                                  'seaborn'], max_iterations=10)
     results = agent.run(
-    task= prompt.format(question=question, schema=schema, table_name=table_name)
+        task= prompt.format(question=question, schema=schema, table_name=table_name)
     )
 
     return results