genie: return final summary text from _parse_attachments

src/databricks_ai_bridge/genie.py

@@ -199,21 +199,19 @@ def _parse_attachments(resp: Dict[str, Any]) -> Dict[str, Any]:
     if not isinstance(attachments, list):
         return result
 
-    # Genie may self-correct, producing multiple query+text pairs. We want
-    # the final query and its paired text (the first text attachment following
-    # the final query).
-    want_new_text = True
     for a in attachments:
         if not isinstance(a, dict):
             continue
 
         if "query" in a:
             result["query_attachment"] = a
-            want_new_text = True
 
-        elif "text" in a and want_new_text:
-            result["text_attachment"] = a
-            want_new_text = False
+        elif "text" in a:
+            # Genie's final summary is the text attachment with no "attachment_id";
+            # one that has an id is a follow-up/clarifying question. Prefer the summary
+            # (last wins), but keep any text as a fallback so we never drop answer text.
+            if a.get("attachment_id") is None or result["text_attachment"] is None:
+                result["text_attachment"] = a
 
         elif "suggested_questions" in a:
             result["suggested_questions_attachment"] = a

tests/databricks_ai_bridge/test_genie.py

@@ -911,9 +911,8 @@ def test_poll_for_result_continues_on_mlflow_tracing_exceptions(genie, mock_work
             None,
             None,
         ),
-        # Self-correction with paired text attachments - text should be paired
-        # with the final query (i.e., the first text AFTER the last query), not
-        # the first text overall or the final text (which may be a follow-up).
+        # Self-correction: earlier attempts and the follow-up question carry an
+        # attachment_id; only the final summary has none, and it wins.
         (
             {
                 "attachments": [
@@ -924,16 +923,16 @@ def test_poll_for_result_continues_on_mlflow_tracing_exceptions(genie, mock_work
                         "suggested_questions": {"questions": ["Q1?"]},
                     },
                     {"attachment_id": "4", "query": {"query": "SELECT correct"}},
-                    {"attachment_id": "5", "text": {"content": "explains correct"}},
-                    {"attachment_id": "6", "text": {"content": "follow-up prompt"}},
+                    {"attachment_id": "5", "text": {"content": "follow-up prompt"}},
+                    {"text": {"content": "final summary"}},
                     {
                         "attachment_id": "7",
                         "suggested_questions": {"questions": ["Q2?"]},
                     },
                 ]
             },
             {"attachment_id": "4", "query": {"query": "SELECT correct"}},
-            {"attachment_id": "5", "text": {"content": "explains correct"}},
+            {"text": {"content": "final summary"}},
             {"attachment_id": "7", "suggested_questions": {"questions": ["Q2?"]}},
         ),
     ],