From 39b8cbdb62a3ab9b6ff506ae1b69d8aeca9fd01e Mon Sep 17 00:00:00 2001 From: nisha2003 <56770365+nisha2003@users.noreply.github.com> Date: Sun, 14 Jun 2026 22:44:13 +0000 Subject: [PATCH 1/5] genie: return final summary text from _parse_attachments --- src/databricks_ai_bridge/genie.py | 46 +++++++++++++++--------- tests/databricks_ai_bridge/test_genie.py | 15 ++++++++ 2 files changed, 44 insertions(+), 17 deletions(-) diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index 3243d64d..faedb833 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/src/databricks_ai_bridge/genie.py @@ -195,29 +195,41 @@ def _parse_attachments(resp: Dict[str, Any]) -> Dict[str, Any]: "suggested_questions_attachment": None, } - attachments = resp.get("attachments") or [] - if not isinstance(attachments, list): + attachments = [a for a in (resp.get("attachments") or []) if isinstance(a, dict)] + if not attachments: 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 - + # Genie may self-correct, producing multiple query+text pairs. Text emitted + # strictly between the first and last query is a superseded attempt and is + # dropped; leading text (before the first query) and trailing text (after the + # last query) are part of the answer. Among the kept text attachments, the + # final summary is the one *without* an "attachment_id" -- internally-created + # messages also emit a follow-up/clarifying-question text attachment (which + # *does* have an attachment_id), and we don't want to surface that instead of + # the answer. Prefer the id-less summary; otherwise fall back to the first + # kept text attachment. + query_indices = [i for i, a in enumerate(attachments) if "query" in a] + first_query, last_query = ( + (query_indices[0], query_indices[-1]) if query_indices else (None, None) + ) + + text_candidates = [] + for i, a in enumerate(attachments): 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 - + result["query_attachment"] = a # last query wins + elif "text" in a: + if first_query is not None and first_query < i < last_query: + continue + text_candidates.append(a) elif "suggested_questions" in a: result["suggested_questions_attachment"] = a + id_less_text = [a for a in text_candidates if a.get("attachment_id") is None] + if id_less_text: + result["text_attachment"] = id_less_text[-1] + elif text_candidates: + result["text_attachment"] = text_candidates[0] + return result diff --git a/tests/databricks_ai_bridge/test_genie.py b/tests/databricks_ai_bridge/test_genie.py index acd1ec1e..9405f2a9 100644 --- a/tests/databricks_ai_bridge/test_genie.py +++ b/tests/databricks_ai_bridge/test_genie.py @@ -936,6 +936,21 @@ def test_poll_for_result_continues_on_mlflow_tracing_exceptions(genie, mock_work {"attachment_id": "5", "text": {"content": "explains correct"}}, {"attachment_id": "7", "suggested_questions": {"questions": ["Q2?"]}}, ), + # Follow-up question (has attachment_id) emitted BEFORE the final summary + # (no attachment_id) - the id-less summary is the answer and must win over + # the preceding follow-up text. + ( + { + "attachments": [ + {"attachment_id": "1", "query": {"query": "SELECT *"}}, + {"attachment_id": "2", "text": {"content": "Which proteins?"}}, + {"text": {"content": "Summary:\n- point a\n- point b"}}, + ] + }, + {"attachment_id": "1", "query": {"query": "SELECT *"}}, + {"text": {"content": "Summary:\n- point a\n- point b"}}, + None, + ), ], ) def test_parse_attachments(resp, exp_query, exp_text, exp_questions): From 4080152a9737ad022d66d05157e6cfcfb300bfee Mon Sep 17 00:00:00 2001 From: nisha2003 <56770365+nisha2003@users.noreply.github.com> Date: Sun, 14 Jun 2026 22:48:04 +0000 Subject: [PATCH 2/5] genie: fix typecheck by computing superseded index set --- src/databricks_ai_bridge/genie.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index faedb833..456dd48b 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/src/databricks_ai_bridge/genie.py @@ -209,17 +209,14 @@ def _parse_attachments(resp: Dict[str, Any]) -> Dict[str, Any]: # the answer. Prefer the id-less summary; otherwise fall back to the first # kept text attachment. query_indices = [i for i, a in enumerate(attachments) if "query" in a] - first_query, last_query = ( - (query_indices[0], query_indices[-1]) if query_indices else (None, None) - ) + # Indices of text strictly between the first and last query: superseded attempts. + superseded = set(range(query_indices[0] + 1, query_indices[-1])) if query_indices else set() text_candidates = [] for i, a in enumerate(attachments): if "query" in a: result["query_attachment"] = a # last query wins - elif "text" in a: - if first_query is not None and first_query < i < last_query: - continue + elif "text" in a and i not in superseded: text_candidates.append(a) elif "suggested_questions" in a: result["suggested_questions_attachment"] = a From 87e026b2a221c886d3a86af5b9768fa9bdb694d0 Mon Sep 17 00:00:00 2001 From: nisha2003 <56770365+nisha2003@users.noreply.github.com> Date: Sun, 14 Jun 2026 22:54:45 +0000 Subject: [PATCH 3/5] genie: simplify _parse_attachments self-correction filter --- src/databricks_ai_bridge/genie.py | 28 +++++++++++++--------------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index 456dd48b..24ab892a 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/src/databricks_ai_bridge/genie.py @@ -199,31 +199,29 @@ def _parse_attachments(resp: Dict[str, Any]) -> Dict[str, Any]: if not attachments: return result - # Genie may self-correct, producing multiple query+text pairs. Text emitted - # strictly between the first and last query is a superseded attempt and is - # dropped; leading text (before the first query) and trailing text (after the - # last query) are part of the answer. Among the kept text attachments, the - # final summary is the one *without* an "attachment_id" -- internally-created - # messages also emit a follow-up/clarifying-question text attachment (which - # *does* have an attachment_id), and we don't want to surface that instead of - # the answer. Prefer the id-less summary; otherwise fall back to the first - # kept text attachment. + # Genie may self-correct, producing multiple query+text pairs. Text strictly + # between the first and last query is a superseded attempt; leading and trailing + # text is part of the answer. (With 0-1 queries the bounds collapse and nothing + # is dropped.) Among the kept text, the final summary is the one *without* an + # "attachment_id" -- internally-created messages also emit a follow-up/clarifying + # question text attachment (which *does* have an attachment_id) that we must not + # surface instead of the answer. Prefer the id-less summary, else the first text. query_indices = [i for i, a in enumerate(attachments) if "query" in a] - # Indices of text strictly between the first and last query: superseded attempts. - superseded = set(range(query_indices[0] + 1, query_indices[-1])) if query_indices else set() + first_query = query_indices[0] if query_indices else 0 + last_query = query_indices[-1] if query_indices else 0 text_candidates = [] for i, a in enumerate(attachments): if "query" in a: result["query_attachment"] = a # last query wins - elif "text" in a and i not in superseded: + elif "text" in a and not (first_query < i < last_query): text_candidates.append(a) elif "suggested_questions" in a: result["suggested_questions_attachment"] = a - id_less_text = [a for a in text_candidates if a.get("attachment_id") is None] - if id_less_text: - result["text_attachment"] = id_less_text[-1] + summaries = [a for a in text_candidates if a.get("attachment_id") is None] + if summaries: + result["text_attachment"] = summaries[-1] # final summary has no attachment_id elif text_candidates: result["text_attachment"] = text_candidates[0] From 5ee76832f69b1b064ad5946e011015088d92c2ec Mon Sep 17 00:00:00 2001 From: nisha2003 <56770365+nisha2003@users.noreply.github.com> Date: Sun, 14 Jun 2026 22:58:33 +0000 Subject: [PATCH 4/5] genie: take last text attachment without an attachment_id --- src/databricks_ai_bridge/genie.py | 31 ++++++++---------------- tests/databricks_ai_bridge/test_genie.py | 12 ++++----- 2 files changed, 16 insertions(+), 27 deletions(-) diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index 24ab892a..64ad760c 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/src/databricks_ai_bridge/genie.py @@ -196,34 +196,23 @@ def _parse_attachments(resp: Dict[str, Any]) -> Dict[str, Any]: } attachments = [a for a in (resp.get("attachments") or []) if isinstance(a, dict)] - if not attachments: - return result - # Genie may self-correct, producing multiple query+text pairs. Text strictly - # between the first and last query is a superseded attempt; leading and trailing - # text is part of the answer. (With 0-1 queries the bounds collapse and nothing - # is dropped.) Among the kept text, the final summary is the one *without* an - # "attachment_id" -- internally-created messages also emit a follow-up/clarifying - # question text attachment (which *does* have an attachment_id) that we must not - # surface instead of the answer. Prefer the id-less summary, else the first text. - query_indices = [i for i, a in enumerate(attachments) if "query" in a] - first_query = query_indices[0] if query_indices else 0 - last_query = query_indices[-1] if query_indices else 0 - - text_candidates = [] - for i, a in enumerate(attachments): + for a in attachments: if "query" in a: result["query_attachment"] = a # last query wins - elif "text" in a and not (first_query < i < last_query): - text_candidates.append(a) elif "suggested_questions" in a: result["suggested_questions_attachment"] = a - summaries = [a for a in text_candidates if a.get("attachment_id") is None] + # The answer is the final summary, which is the text attachment without an + # "attachment_id" -- internally-created messages also emit a follow-up/clarifying + # question text attachment (which *does* have one) that we must not surface + # instead. Take the last id-less text; fall back to the last text otherwise. + texts = [a for a in attachments if "text" in a] + summaries = [a for a in texts if a.get("attachment_id") is None] if summaries: - result["text_attachment"] = summaries[-1] # final summary has no attachment_id - elif text_candidates: - result["text_attachment"] = text_candidates[0] + result["text_attachment"] = summaries[-1] + elif texts: + result["text_attachment"] = texts[-1] return result diff --git a/tests/databricks_ai_bridge/test_genie.py b/tests/databricks_ai_bridge/test_genie.py index 9405f2a9..57ffb031 100644 --- a/tests/databricks_ai_bridge/test_genie.py +++ b/tests/databricks_ai_bridge/test_genie.py @@ -911,9 +911,9 @@ 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-attempt text and a follow-up question both carry + # an attachment_id; only the final summary has none. The id-less summary wins + # over both. ( { "attachments": [ @@ -924,8 +924,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 +933,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?"]}}, ), # Follow-up question (has attachment_id) emitted BEFORE the final summary From 9c77c9d662f2266c87eafa98c9818d1ac8cab4a6 Mon Sep 17 00:00:00 2001 From: nisha2003 <56770365+nisha2003@users.noreply.github.com> Date: Mon, 15 Jun 2026 16:30:39 +0000 Subject: [PATCH 5/5] genie: minimal diff for _parse_attachments cleanup --- src/databricks_ai_bridge/genie.py | 28 +++++++++++++----------- tests/databricks_ai_bridge/test_genie.py | 20 ++--------------- 2 files changed, 17 insertions(+), 31 deletions(-) diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index 64ad760c..4f84ed33 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/src/databricks_ai_bridge/genie.py @@ -195,25 +195,27 @@ def _parse_attachments(resp: Dict[str, Any]) -> Dict[str, Any]: "suggested_questions_attachment": None, } - attachments = [a for a in (resp.get("attachments") or []) if isinstance(a, dict)] + attachments = resp.get("attachments") or [] + if not isinstance(attachments, list): + return result for a in attachments: + if not isinstance(a, dict): + continue + if "query" in a: - result["query_attachment"] = a # last query wins + result["query_attachment"] = a + + 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 - # The answer is the final summary, which is the text attachment without an - # "attachment_id" -- internally-created messages also emit a follow-up/clarifying - # question text attachment (which *does* have one) that we must not surface - # instead. Take the last id-less text; fall back to the last text otherwise. - texts = [a for a in attachments if "text" in a] - summaries = [a for a in texts if a.get("attachment_id") is None] - if summaries: - result["text_attachment"] = summaries[-1] - elif texts: - result["text_attachment"] = texts[-1] - return result diff --git a/tests/databricks_ai_bridge/test_genie.py b/tests/databricks_ai_bridge/test_genie.py index 57ffb031..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: earlier-attempt text and a follow-up question both carry - # an attachment_id; only the final summary has none. The id-less summary wins - # over both. + # Self-correction: earlier attempts and the follow-up question carry an + # attachment_id; only the final summary has none, and it wins. ( { "attachments": [ @@ -936,21 +935,6 @@ def test_poll_for_result_continues_on_mlflow_tracing_exceptions(genie, mock_work {"text": {"content": "final summary"}}, {"attachment_id": "7", "suggested_questions": {"questions": ["Q2?"]}}, ), - # Follow-up question (has attachment_id) emitted BEFORE the final summary - # (no attachment_id) - the id-less summary is the answer and must win over - # the preceding follow-up text. - ( - { - "attachments": [ - {"attachment_id": "1", "query": {"query": "SELECT *"}}, - {"attachment_id": "2", "text": {"content": "Which proteins?"}}, - {"text": {"content": "Summary:\n- point a\n- point b"}}, - ] - }, - {"attachment_id": "1", "query": {"query": "SELECT *"}}, - {"text": {"content": "Summary:\n- point a\n- point b"}}, - None, - ), ], ) def test_parse_attachments(resp, exp_query, exp_text, exp_questions):