As per  the doc , one can execute an ai_query in python. But I am sot sure why I can't

Works in SQL

select 1 as colA, ai_query('databricks-llama-4-maverick', 'what is 2+2') as test

But not in spark

df = spark.sql("select 1 as colA")
df = df.selectExpr("""
ai_query('databricks-llama-4-maverick', 'what is 2+2') as result
                   """)

df.show()

I am getting

AnalysisException: [UNRESOLVED_ROUTINE] Cannot resolve function `ai_query` on search path [`system`.`builtin`, `system`.`session`, `catalog`.`default`].; line 1 pos 0

I have also converted the ai_query to a sql udf but even then it does not work.