fix: make sure that context logging works in development mode

Author: gustavocidornelas

src/openlayer/lib/core/base_model.py

@@ -42,9 +42,7 @@ class OpenlayerModel(abc.ABC):
     def run_from_cli(self) -> None:
         """Run the model from the command line."""
         parser = argparse.ArgumentParser(description="Run data through a model.")
-        parser.add_argument(
-            "--dataset-path", type=str, required=True, help="Path to the dataset"
-        )
+        parser.add_argument("--dataset-path", type=str, required=True, help="Path to the dataset")
         parser.add_argument(
             "--output-dir",
             type=str,
@@ -87,9 +85,7 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
             # Filter row_dict to only include keys that are valid parameters
             # for the 'run' method
             row_dict = row.to_dict()
-            filtered_kwargs = {
-                k: v for k, v in row_dict.items() if k in run_signature.parameters
-            }
+            filtered_kwargs = {k: v for k, v in row_dict.items() if k in run_signature.parameters}
 
             # Call the run method with filtered kwargs
             output = self.run(**filtered_kwargs)
@@ -111,6 +107,8 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
                     df.at[index, "cost"] = processed_trace["cost"]
                 if "tokens" in processed_trace:
                     df.at[index, "tokens"] = processed_trace["tokens"]
+                if "context" in processed_trace:
+                    df.at[index, "context"] = processed_trace["context"]
 
         config = {
             "outputColumnName": "output",
@@ -126,6 +124,8 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
             config["costColumnName"] = "cost"
         if "tokens" in df.columns:
             config["numOfTokenColumnName"] = "tokens"
+        if "context" in df.columns:
+            config["contextColumnName"] = "context"
 
         return df, config