Fix blocking errors in notebook

[huyouare]

Ran into issues while running the colab, all now resolved:

First error (cell 7):

* postgresql://readonly:***@db.jccajjvblmajkbwqsmaz.supabase.co:5432/postgres
(psycopg2.errors.NumericValueOutOfRange) integer out of range

[SQL: select
  task.name,
  agent_name,
  model_name,
  avg(
    (
      n_input_tokens * cents_per_million_input_tokens + n_output_tokens * cents_per_million_output_tokens
    ) [/](https://file+.vscode-resource.vscode-cdn.net/) 1e6
  ) as avg_cost_cents,
  AVG(
    EXTRACT(
      EPOCH
      from
        (ended_at - started_at)
    )
  ) as avg_completion_seconds
from
  trial
  inner join trial_model on trial.id = trial_model.trial_id
  inner join model on model.name = trial_model.model_name
  inner join task on task.checksum = trial.task_checksum
  and model.provider = trial_model.model_provider
...
  agent_version
order by
  avg_cost_cents desc;]
(Background on this error at: https://sqlalche.me/e/20/9h9h)

Second error (cell 7):

KeyError                                  Traceback (most recent call last)
Cell In[20], line 1
----> 1 get_ipython().run_cell_magic('sql', 'result <<', 'select\n  task.name,\n  agent_name,\n  model_name,\n  avg(\n    (\n      n_input_tokens::numeric * cents_per_million_input_tokens + n_output_tokens::numeric * cents_per_million_output_tokens\n    ) / 1e6\n  ) as avg_cost_cents,\n  AVG(\n    EXTRACT(\n      EPOCH\n      from\n        (ended_at - started_at)\n    )\n  ) as avg_completion_seconds\nfrom\n  trial\n  inner join trial_model on trial.id = trial_model.trial_id\n  inner join model on model.name = trial_model.model_name\n  inner join task on task.checksum = trial.task_checksum\n  and model.provider = trial_model.model_provider\nwhere\n  n_input_tokens is not null\n  and n_output_tokens is not null\ngroup by\n  task.name,\n  agent_name,\n  model_name,\n  agent_version\norder by\n  avg_cost_cents desc;\n')

File ~/code/tbench-colab/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py:2565, in InteractiveShell.run_cell_magic(self, magic_name, line, cell)
   2563 with self.builtin_trap:
   2564     args = (magic_arg_s, cell)
-> 2565     result = fn(*args, **kwargs)
   2567 # The code below prevents the output from being displayed
   2568 # when using magics with decorator @output_can_be_silenced
   2569 # when the last Python token in the expression is a ';'.
   2570 if getattr(fn, magic.MAGIC_OUTPUT_CAN_BE_SILENCED, False):

File ~/code/tbench-colab/.venv/lib/python3.11/site-packages/sql/magic.py:219, in SqlMagic.execute(self, line, cell, local_ns)
    216     return
    218 try:
--> 219     result = sql.run.run(conn, parsed["sql"], self, user_ns)
    221     if (
    222         result is not None
    223         and not isinstance(result, str)
   (...)    226         # Instead of returning values, set variables directly in the
    227         # user's namespace. Variable names given by column names
    229         if self.autopandas:
...
--> 116     self.pretty = PrettyTable(self.field_names, style=prettytable.__dict__[config.style.upper()])
    117 else:
    118     list.__init__(self, [])

KeyError: 'DEFAULT'

Third issue:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[24], line 1
----> 1 costs = [float(row[3]) for row in result]
      2 times = [float(row[4]) for row in result]
      4 plt.figure(figsize=(8, 6))

Cell In[24], line 1, in <listcomp>(.0)
----> 1 costs = [float(row[3]) for row in result]
      2 times = [float(row[4]) for row in result]
      4 plt.figure(figsize=(8, 6))

TypeError: float() argument must be a string or a real number, not 'NoneType'

[huyouare]

@alexgshaw could you PTAL for review? Would like to add add'l notebook branched from this one