Accessing the MLflow AI Gateway¶

MLflow AI Gateway provides a centralized, auditable way to connect to various LLM APIs.

The following figure shows the basic system setup for accessing an MLflow AI Gateway endpoint from within a UDF:

In order to use such a setup, you need to

Have access to an external Model Inference service or start it locally (e.g. Ollama).
Have a model inside the inference service.
Create an AI Gateway endpoint in MLflow, see MLflow documentation.

Accessing an AI Gateway Endpoint From Python¶

The MLflow documentation on MLflow Invocations API contains examples for cURL and Python, while the Python example actually only uses the REST API via the Python library requests.

import requests

def send_request_to_ai_gateway(
    endpoint: str,
    mlflow_tracking_uri: str = "http://localhost:5000",
    auth: tuple[str, str] | None = None,
    question: str = "",
) -> dict:
    url = f"{mlflow_tracking_uri}/gateway/{endpoint}/mlflow/invocations"
    jreq = {
        "messages": [{ "role": "user", "content": question}],
        "max_tokens": 400,
        "temperature": 0.7,
    }
    response = requests.post(url, json=jreq, auth=auth)
    return response.json()

Accessing AI Gateway Endpoints From Within a UDF¶

Creating an Exasol Connection¶

To access an MLflow AI Gateway endpoint from within a UDF, store the MLflow tracking URI in an Exasol Connection. The Connection can also hold the authentication credentials for the MLflow server which is usually required and recommended for production setups.

Here is an example for creating such a Connection for basic authentication with username and password. See the MLflow documentation for other authentication variants.

CREATE OR REPLACE CONNECTION "MLFLOW"
    TO '<mlflow_tracking_uri>'
    USER '<user_name>'
    IDENTIFIED BY '<password>';

UDF Content¶

A UDF can securely read the MLflow tracking URI and the authentication secrets from the connection object and call the function send_request_to_ai_gateway() defined above:

--/
CREATE OR REPLACE PYTHON3 SCALAR SCRIPT ASK_AI (
    endpoint VARCHAR(2000000),
    question VARCHAR(2000000)
)
EMITS (
    answer VARCHAR(2000000)
) AS

# Add function send_request_to_ai_gateway() here!

def run(ctx):
    conn = exa.get_connection("MLFLOW")
    resp = send_request_to_ai_gateway(
        ctx.endpoint,
        conn.address,
        auth=(conn.user, conn.password),
        question=ctx.question,
    )
    try:
        answer = resp["choices"][0]["message"]["content"]
    except:
        answer = ""
    ctx.emit(answer)
/
;

SELECT ASK_AI('What is the capital of Germany? (Keep it short)');

For more details, please see Exasol’s documentation on Python UDFs and UDFs in general.