Genie tool uses MCP Adapters under the hood

.github/workflows/main.yml

@@ -150,7 +150,6 @@ jobs:
         run: |
           # Only testing initialization since functionality can change
           pytest integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py::test_init
-          pytest integrations/langchain/tests/unit_tests/test_genie.py
           pytest integrations/langchain/tests/unit_tests/test_embeddings.py
           pytest integrations/langchain/tests/unit_tests/test_chat_models.py
 

integrations/langchain/src/databricks_langchain/genie.py

@@ -6,6 +6,13 @@
 from databricks_ai_bridge.genie import Genie
 
 
+def _extract_query_attachment_fields(query_attachments):
+    descriptions = [qa.description for qa in query_attachments if qa.description]
+    queries = [qa.query for qa in query_attachments if qa.query]
+    results = [qa.result for qa in query_attachments if qa.result is not None]
+    return descriptions, queries, results
+
+
 @mlflow.trace()
 def _concat_messages_array(messages):
     concatenated_message = "\n".join(
@@ -25,6 +32,7 @@ def _query_genie_as_agent(
     genie: Genie,
     genie_agent_name,
     include_context: bool = False,
+    include_suggested_questions: bool = False,
     message_processor: Optional[Callable] = None,
 ):
     """
@@ -42,7 +50,7 @@ def _query_genie_as_agent(
     Returns:
         A dictionary containing the messages and conversation_id.
         If include_context is True, the dictionary will also contain the query_reasoning and query_sql fields.
-        If Genie returned a dataframe because it was told to do returns in Pandas format, the dictionary will also contain the dataframe field.
+        If Genie returned DataFrames (return_pandas=True), the dictionary will also contain a "dataframes" field with the list of DataFrames.
     """
     from langchain_core.messages import AIMessage
 
@@ -61,36 +69,51 @@ def _query_genie_as_agent(
     # Send the message and wait for a response, passing conversation_id if available
     genie_response = genie.ask_question(query, conversation_id=conversation_id)
 
-    query_reasoning = genie_response.description or ""
-    query_sql = genie_response.query or ""
-    query_result = genie_response.result if genie_response.result is not None else ""
     query_conversation_id = genie_response.conversation_id or ""
+    query_message_id = genie_response.message_id or ""
+    query_attachments = genie_response.query_attachments
+    text_attachments = genie_response.text_attachments
+    suggested_questions = genie_response.suggested_questions
+    error_msg = genie_response.error_msg
+
+    descriptions, queries, results = _extract_query_attachment_fields(query_attachments)
 
     # Create a list of AIMessage to return
     messages = []
 
     if include_context:
-        messages.append(AIMessage(content=query_reasoning, name="query_reasoning"))
-        messages.append(AIMessage(content=query_sql, name="query_sql"))
-
-    # Handle DataFrame vs string results
-    if isinstance(query_result, pd.DataFrame):  # if we asked for Pandas return
-        # Convert to markdown for message display
-        query_result_content = query_result.to_markdown(index=False)
-        messages.append(AIMessage(content=query_result_content, name="query_result"))
-
-        # Return with DataFrame included
-        return {
-            "messages": messages,
-            "conversation_id": query_conversation_id,
-            "dataframe": query_result,  # Include raw DataFrame if Genie returned dataframe
-        }
-    else:
-        # String result - just add to messages
-        messages.append(AIMessage(content=query_result, name="query_result"))
-
-        # Return without DataFrame field
-        return {"messages": messages, "conversation_id": query_conversation_id}
+        if descriptions:
+            messages.append(AIMessage(content="\n\n".join(descriptions), name="query_reasoning"))
+        if queries:
+            messages.append(AIMessage(content="\n\n".join(queries), name="query_sql"))
+
+    if include_suggested_questions and suggested_questions:
+        messages.append(
+            AIMessage(content="\n\n".join(suggested_questions), name="suggested_questions")
+        )
+
+    query_result_parts = []
+    if results:
+        query_result_parts.extend(
+            r.to_markdown(index=False) if isinstance(r, pd.DataFrame) else r for r in results
+        )
+    if text_attachments:
+        query_result_parts.extend(text_attachments)
+    if error_msg:
+        query_result_parts.append(error_msg)
+    if query_result_parts:
+        messages.append(AIMessage(content="\n\n".join(query_result_parts), name="query_result"))
+
+    result = {
+        "messages": messages,
+        "conversation_id": query_conversation_id,
+        "message_id": query_message_id,
+    }
+
+    if results and isinstance(results[0], pd.DataFrame):
+        result["dataframes"] = results
+
+    return result
 
 
 @mlflow.trace(span_type="AGENT")
@@ -99,6 +122,7 @@ def GenieAgent(
     genie_agent_name: str = "Genie",
     description: str = "",
     include_context: bool = False,
+    include_suggested_questions: bool = False,
     message_processor: Optional[Callable] = None,
     client: Optional["WorkspaceClient"] = None,
     return_pandas: bool = False,
@@ -110,6 +134,7 @@ def GenieAgent(
         genie_agent_name: Name for the agent (default: "Genie")
         description: Custom description for the agent
         include_context: Whether to include query reasoning and SQL in the response
+        include_suggested_questions: Whether to include suggested follow-up questions in the response
         message_processor: Optional function to process messages before querying. It should accept a list of either dict
                             or LangChain Message objects and return a query string. If not provided, the agent will
                             use the chat history to form the query.
@@ -165,6 +190,7 @@ def custom_processor(messages):
         genie=genie,
         genie_agent_name=genie_agent_name,
         include_context=include_context,
+        include_suggested_questions=include_suggested_questions,
         message_processor=message_processor,
     )