Improve & Address Bugs in`test_retrieval` the Batch Test Question DAG

[davidgxue]

Bug

Describe the bug

• Traceback

  [2024-02-13, 17:31:11 EST] {taskinstance.py:2699} ERROR - Task failed with exception
  Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/airflow/models/taskinstance.py", line 433, in _execute_task
  result = execute_callable(context=context, **execute_callable_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/airflow/decorators/base.py", line 242, in execute
  return_value = super().execute(context)
                 ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/airflow/operators/python.py", line 199, in execute
  return_value = self.execute_callable()
                 ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/airflow/operators/python.py", line 216, in execute_callable
  return self.python_callable(*self.op_args, **self.op_kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/dags/monitor/test_retrieval.py", line 210, in generate_test_answers
  questions_df[["askastro_answer", "askastro_references", "langsmith_link"]] = questions_df.question.apply(
  ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/pandas/core/frame.py", line 4079, in __setitem__
  self._setitem_array(key, value)
  File "/usr/local/lib/python3.11/site-packages/pandas/core/frame.py", line 4138, in _setitem_array
  self._iset_not_inplace(key, value)
  File "/usr/local/lib/python3.11/site-packages/pandas/core/frame.py", line 4157, in _iset_not_inplace
  raise ValueError("Columns must be same length as key")
  ValueError: Columns must be same length as key

• question_number_subset
  • The questions_df after adding debug logging is empty, this only occurs if someone puts in a subset of question ids
  • The question_number_subset param isn't parsed correctly due to the incorrect code json.loads() which attempts to parse string into list of ints (but not correctly), leading to no questions being added here.

To Reproduce Steps to reproduce the behavior:

1. Have proper configuration of environment variables for the test_retrieval DAG
2. Trigger the DAG
3. Put a list of subset question ids in the parameter prompt, such as [1,2,3]
4. Errors out during DAG run

Expected behavior No errors

Improvements

1. The references saved in the csv are in random incorrect order. This is probably related to the fact that it is put into a set using {} somewhere.
2. The multi-query references and the weaviate search references are not relevant. They don't provide useful info but delays the pipeline and incurs cost.