snapshot from the dfatabricks documentation to build function on top of ai_query https://docs.databricks.com/aws/en/sql/language-manual/functions/ai_query

databricks_ai_query_udf.sql

 CREATE FUNCTION correct_grammar(text STRING)
  RETURNS STRING
  RETURN ai_query(
    'databricks-meta-llama-3-3-70b-instruct',
    CONCAT('Correct this to standard English:\n', text));
> GRANT EXECUTE ON correct_grammar TO ds;
- DS fixes grammar issues in a batch.
> SELECT
    * EXCEPT text,
    correct_grammar(text) AS text
  FROM articles;