From 5dada822d57587a5891c138a1f49241dfb29a80e Mon Sep 17 00:00:00 2001
From: Wenzhe Xue <wenzhe@log10.io>
Date: Fri, 15 Mar 2024 11:43:55 -0700
Subject: [PATCH] call flatten_messages when retrieving the completions for
 examples

---
 log10/feedback/autofeedback.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/log10/feedback/autofeedback.py b/log10/feedback/autofeedback.py
index 833880da..fec0004e 100644
--- a/log10/feedback/autofeedback.py
+++ b/log10/feedback/autofeedback.py
@@ -50,13 +50,13 @@ def _get_examples(self):
                 print(e)
                 continue
             completion = res.json()["data"]
-            prompt = completion["request"]["messages"][1]["content"]
-            response = completion["response"]["choices"][0]["message"]["content"]
+            fm = flatten_messages(completion)
+
             few_shot_examples.append(
                 {
                     "completion_id": completion_id,
-                    "prompt": prompt,
-                    "response": response,
+                    "prompt": fm["prompt"],
+                    "response": fm["response"],
                     "feedback": json.dumps(feedback_values),
                 }
             )