try-except

Author: nisha2003

src/databricks_ai_bridge/genie.py

@@ -404,60 +404,73 @@ def ask_question(self, question, conversation_id: Optional[str] = None):
         Returns:
             GenieResponse with result, query, description, and conversation_id
         """
-        args = {"query": question}
-        if conversation_id:
-            args["conversation_id"] = conversation_id
+        try:
+            # Try using the new MCP approach
+            args = {"query": question}
+            if conversation_id:
+                args["conversation_id"] = conversation_id
 
-        mcp_result = self._mcp_client.call_tool(self._tool_name, args)
+            mcp_result = self._mcp_client.call_tool(self._tool_name, args)
 
-        if not mcp_result.content or len(mcp_result.content) == 0:
-            return GenieResponse(
-                result="No content returned from Genie",
-                conversation_id=conversation_id,
-            )
+            if not mcp_result.content or len(mcp_result.content) == 0:
+                return GenieResponse(
+                    result="No content returned from Genie",
+                    conversation_id=conversation_id,
+                )
 
-        # Genie backend always returns exactly 1 content block with JSON
-        content_block = mcp_result.content[0]
-        content_text = content_block.text if hasattr(content_block, "text") else "{}"
+            # Genie backend always returns exactly 1 content block with JSON
+            content_block = mcp_result.content[0]
+            content_text = content_block.text if hasattr(content_block, "text") else "{}"
 
-        try:
-            genie_response = json.loads(content_text)
-        except json.JSONDecodeError:
-            return GenieResponse(
-                result=f"Failed to parse response: {content_text}",
-                conversation_id=conversation_id,
-            )
+            try:
+                genie_response = json.loads(content_text)
+            except json.JSONDecodeError:
+                return GenieResponse(
+                    result=f"Failed to parse response: {content_text}",
+                    conversation_id=conversation_id,
+                )
 
-        content = genie_response.get("content", "")
-        conv_id = genie_response.get("conversationId", conversation_id)
-        status = genie_response.get("status", "")
+            content = genie_response.get("content", "")
+            conv_id = genie_response.get("conversationId", conversation_id)
+            status = genie_response.get("status", "")
+
+            try:
+                content_data = json.loads(content)
+                query_str = content_data.get("query", "")
+                description = content_data.get("description", "")
+                statement_response = content_data.get("statement_response")
+
+                if (
+                    statement_response
+                    and statement_response.get("status", {}).get("state") == "SUCCEEDED"
+                ):
+                    result = _parse_query_result(
+                        statement_response, self.truncate_results, self.return_pandas
+                    )
+                else:
+                    result = content
+                    query_str = ""
+                    description = ""
 
-        try:
-            content_data = json.loads(content)
-            query_str = content_data.get("query", "")
-            description = content_data.get("description", "")
-            statement_response = content_data.get("statement_response")
-
-            if (
-                statement_response
-                and statement_response.get("status", {}).get("state") == "SUCCEEDED"
-            ):
-                result = _parse_query_result(
-                    statement_response, self.truncate_results, self.return_pandas
-                )
-            else:
+            except (json.JSONDecodeError, KeyError, TypeError, AttributeError):
                 result = content
                 query_str = ""
                 description = ""
 
-        except (json.JSONDecodeError, KeyError, TypeError, AttributeError):
-            result = content
-            query_str = ""
-            description = ""
+            return GenieResponse(
+                result=result,
+                query=query_str,
+                description=description,
+                conversation_id=conv_id,
+            )
 
-        return GenieResponse(
-            result=result,
-            query=query_str,
-            description=description,
-            conversation_id=conv_id,
-        )
+        except Exception as e:
+            # Fall back to deprecated REST methods
+            logging.warning(f"MCP call failed, falling back to REST API: {e}")
+
+            if conversation_id:
+                resp = self.create_message(conversation_id, question)
+                return self.poll_for_result(conversation_id, resp["id"])
+            else:
+                resp = self.start_conversation(question)
+                return self.poll_for_result(resp["conversation_id"], resp["id"])