diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index 3243d64d..4f84ed33 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/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 diff --git a/tests/databricks_ai_bridge/test_genie.py b/tests/databricks_ai_bridge/test_genie.py index acd1ec1e..c5991ad9 100644 --- a/tests/databricks_ai_bridge/test_genie.py +++ b/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,8 +923,8 @@ 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?"]}, @@ -933,7 +932,7 @@ def test_poll_for_result_continues_on_mlflow_tracing_exceptions(genie, mock_work ] }, {"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?"]}}, ), ],