Genie tool uses MCP Adapters under the hood

[nisha2003]

Migrates the Genie class to use MCP (Model Context Protocol) adapters, replacing direct REST API calls in poll_for_result and ask_question.

Updated unit tests.

Manual test here: https://e2-spog.staging.cloud.databricks.com/editor/notebooks/2260046326096166?o=6051921418418893

[bbqiu]

can we do something like this instead since we know it's only 1 tool?

# Use next() with a generator:
  self._query_tool_name = next(
      (tool.name for tool in tools if tool.name.startswith("query_space_")),
      None
  )
  self._poll_tool_name = next(
      (tool.name for tool in tools if tool.name.startswith("poll_response_")),
      None
  )

[bbqiu]

what is this for? i think for dev, we can manually do local installs of databricks-mcp instead of this

does this always try to install from local instead of pypi? looks like it's only for local work, and on pypi, this won't apply -- if that's right, we should do this for all of our subpackage definitions to rely on the local installs when possible for dev

[nisha2003]

added for the other subpackages!

[bbqiu]

overall looks great! left a few comments

[aravind-segu]

just curious here, do we always return only one query attachment?

[nisha2003]

Addressed in comment below.

[aravind-segu]

we are adding the query_str here even if the statement response failed. Is that ok? Is that what we do currently?

[nisha2003]

Yep, I think the old code did this (snippet below of the old poll_query_results). I imagine this is so that even if a SQL query fails there's some use in sending it back for debugging?

if state == "SUCCEEDED":
    result = _parse_query_result(resp, self.truncate_results, self.return_pandas)
    return GenieResponse(result, query_str, description, returned_conversation_id)
elif state in ["RUNNING", "PENDING"]:
    logging.debug("Waiting for query result...")
    time.sleep(ITERATION_FREQUENCY)
else:
     return GenieResponse(
        f"No query result: {resp['state']}",
        query_str,
        description,
        returned_conversation_id,
     )

[aravind-segu]

If we are keeping start_conversation/create_message and just adding a deprecated, should we keep original poll_for_result, mark as deprecated and then create a new poll for ask_question

[nisha2003]

I think the old poll_for_result is used by ask_question. I guess it's internally used but the pattern of using start_conversation --> create_message --> poll_for_result is deprecated. Since on the user side it shouldn't be used I can mark poll_for_result to deprecated as well.

I feel like we shouldn't need a new poll for ask_question right? Let me know if I'm missing smth.

[nisha2003]

Added an internal poller that is called by ask_question and poll_for_result.

[aravind-segu]

There are some cases where we return both query_attachments and text_attachments. Previously we were returning all attachments to the user. Now we will only be returnign query_attachements. Is this a conscious decision?

If not maybe its better to also return text_attachments here if query_attachments exist. Text attachments are usually what genie summarizes and basically is an answer to the user.

[nisha2003]

Hm okay -- I did this because the old code had the following logic. It's also I took only one query attachment. If we want to change this to send all the attachments we probably need to update GenieResponse. Is this okay? I imagine it is since we'll just feed it to the agent + will have more information but let me know what you think.

if current_status == "COMPLETED": attachment = next( (r for r in resp["attachments"] if "query" in r), None ) if attachment: query_obj = attachment["query"] description = query_obj.get("description", "") query_str = query_obj.get("query", "") attachment_id = attachment["attachment_id"] return poll_query_results( attachment_id, query_str, description, returned_conversation_id, ) if current_status == "COMPLETED": text_content = next(r for r in resp["attachments"] if "text" in r)[ "text" ]["content"] return GenieResponse( result=text_content, conversation_id=returned_conversation_id, )