databricks_ai_bridge_model_serving_and_unity_catalog_functions.py

1from databricks_ai_bridge import AIBridge
2
3# Initialize the AI Bridge
4# Note: Ensure DATABRICKS_HOST and DATABRICKS_TOKEN environment variables are set
5bridge = AIBridge()
6
7# Example 1: Call a model serving endpoint
8response = bridge.predict(
9    endpoint_name="databricks-llama-2-70b-chat",
10    inputs={
11        "messages": [
12            {"role": "user", "content": "What is the capital of France?"}
13        ]
14    }
15)
16
17print(f"Model response: {response}")
18
19# Example 2: Call a Unity Catalog registered function
20# Assumes a function 'main.default.hello_world' exists
21result = bridge.call_function(
22    function_name="main.default.hello_world",
23    parameters={"name": "Databricks User"}
24)
25
26print(f"Function result: {result}")