litellm-mirror/docs/my-website/docs/providers/databricks.md
2024-08-01 16:52:52 -07:00

6.2 KiB

import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem';

🆕 Databricks

LiteLLM supports all models on Databricks

:::tip

We support ALL Databricks models, just set model=databricks/<any-model-on-databricks> as a prefix when sending litellm requests

:::

Usage

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"}]
)
  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
  1. Start the proxy
$ litellm --config /path/to/config.yaml --debug
  1. Send Request to LiteLLM Proxy Server
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 --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

:::tip

We support ALL Databricks models, just set model=databricks/<any-model-on-databricks> as a prefix when sending litellm requests

:::

Model Name Command
databricks-meta-llama-3-1-70b-instruct completion(model='databricks/databricks-meta-llama-3-1-70b-instruct', messages=messages)
databricks-meta-llama-3-1-405b-instruct completion(model='databricks/databricks-meta-llama-3-1-405b-instruct', messages=messages)
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

:::tip

We support ALL Databricks models, just set model=databricks/<any-model-on-databricks> as a prefix when sending litellm requests

:::

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)