🆕 Databricks

LiteLLM supports all models on Databricks

Usage

SDK

ENV VAR

import os 
os.environ["DATABRICKS_API_KEY"] = ""
os.environ["DATABRICKS_API_BASE"] = ""

Example Call

from litellm import completion
import os
## set ENV variables
os.environ["DATABRICKS_API_KEY"] = "databricks key"
os.environ["DATABRICKS_API_BASE"] = "databricks base url" # e.g.: https://adb-3064715882934586.6.azuredatabricks.net/serving-endpoints

# Databricks dbrx-instruct call
response = completion(
    model="databricks/databricks-dbrx-instruct", 
    messages = [{ "content": "Hello, how are you?","role": "user"}]
)

PROXY

1. Add models to your config.yaml

   model_list:
     - model_name: dbrx-instruct
       litellm_params:
         model: databricks/databricks-dbrx-instruct
         api_key: os.environ/DATABRICKS_API_KEY
         api_base: os.environ/DATABRICKS_API_BASE

2. Start the proxy

   $ litellm --config /path/to/config.yaml --debug

3. Send Request to LiteLLM Proxy Server

   OpenAI Python v1.0.0+

   import openai
   client = openai.OpenAI(
       api_key="sk-1234",             # pass litellm proxy key, if you're using virtual keys
       base_url="http://0.0.0.0:4000" # litellm-proxy-base url
   )
   
   response = client.chat.completions.create(
       model="dbrx-instruct",
       messages = [
         {
             "role": "system",
             "content": "Be a good human!"
         },
         {
             "role": "user",
             "content": "What do you know about earth?"
         }
     ]
   )
   
   print(response)

   curl

   curl --location 'http://0.0.0.0:4000/chat/completions' \
       --header 'Authorization: Bearer sk-1234' \
       --header 'Content-Type: application/json' \
       --data '{
       "model": "dbrx-instruct",
       "messages": [
         {
             "role": "system",
             "content": "Be a good human!"
         },
         {
             "role": "user",
             "content": "What do you know about earth?"
         }
         ],
   }'

Passing additional params - max_tokens, temperature

See all litellm.completion supported params here

# !pip install litellm
from litellm import completion
import os
## set ENV variables
os.environ["DATABRICKS_API_KEY"] = "databricks key"
os.environ["DATABRICKS_API_BASE"] = "databricks api base"

# databricks dbrx call
response = completion(
    model="databricks/databricks-dbrx-instruct", 
    messages = [{ "content": "Hello, how are you?","role": "user"}],
    max_tokens=20,
    temperature=0.5
)

proxy

  model_list:
    - model_name: llama-3
      litellm_params:
        model: databricks/databricks-meta-llama-3-70b-instruct
        api_key: os.environ/DATABRICKS_API_KEY
        max_tokens: 20
        temperature: 0.5

Passings Databricks specific params - 'instruction'

For embedding models, databricks lets you pass in an additional param 'instruction'. Full Spec

# !pip install litellm
from litellm import embedding
import os
## set ENV variables
os.environ["DATABRICKS_API_KEY"] = "databricks key"
os.environ["DATABRICKS_API_BASE"] = "databricks url"

# Databricks bge-large-en call
response = litellm.embedding(
      model="databricks/databricks-bge-large-en",
      input=["good morning from litellm"],
      instruction="Represent this sentence for searching relevant passages:",
  )

proxy

  model_list:
    - model_name: bge-large
      litellm_params:
        model: databricks/databricks-bge-large-en
        api_key: os.environ/DATABRICKS_API_KEY
        api_base: os.environ/DATABRICKS_API_BASE
        instruction: "Represent this sentence for searching relevant passages:"

Supported Databricks Chat Completion Models

Model Name	Command
databricks-dbrx-instruct	completion(model='databricks/databricks-dbrx-instruct', messages=messages)
databricks-meta-llama-3-70b-instruct	completion(model='databricks/databricks-meta-llama-3-70b-instruct', messages=messages)
databricks-llama-2-70b-chat	completion(model='databricks/databricks-llama-2-70b-chat', messages=messages)
databricks-mixtral-8x7b-instruct	completion(model='databricks/databricks-mixtral-8x7b-instruct', messages=messages)
databricks-mpt-30b-instruct	completion(model='databricks/databricks-mpt-30b-instruct', messages=messages)
databricks-mpt-7b-instruct	completion(model='databricks/databricks-mpt-7b-instruct', messages=messages)

Supported Databricks Embedding Models

Model Name	Command
databricks-bge-large-en	embedding(model='databricks/databricks-bge-large-en', messages=messages)
databricks-gte-large-en	embedding(model='databricks/databricks-gte-large-en', messages=messages)