Databricks Integration: Integrate Databricks SDK as optional mechanism for fetching API base and token, if unspecified

[dbczumar]

Title

Integrate Databricks SDK as optional mechanism for fetching API base and token, if unspecified

Relevant issues

Type

🆕 New Feature

Changes

This PR introduces support for using the Databricks SDK to automatically fetch an API token and API base from the current environment when handling completions requests, if the user does not specify a token / base manually.

Unit tests

Manual tests

Ran the following script and confirmed that the output is sane:

import asyncio
from litellm import completion, acompletion, embedding, aembedding
import os

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

print(response)
print(list(response))

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


response = completion(
    model="databricks/databricks-dbrx-instruct", 
    messages = [{ "content": "Hello, how are you?","role": "user"}],
    stream=True
)
print(response)
print(list(response))

response = asyncio.run(acompletion(
    model="databricks/databricks-dbrx-instruct", 
    messages = [{ "content": "Hello, how are you?","role": "user"}]
))
print(response)

response = asyncio.run(acompletion(
    model="databricks/databricks-dbrx-instruct", 
    messages = [{ "content": "Hello, how are you?","role": "user"}],
    stream=True
))
print(response)
print(asyncio.run(anext(asyncio.run(response.fetch_stream()))))

response = embedding(
    model="databricks/databricks-bge-large-en",
    input=["Hello, how are you?"],
)

print(response)
print(list(response))

response = asyncio.run(aembedding(
    model="databricks/databricks-bge-large-en",
    input=["Hello, how are you?"],
))

print(response)
print(list(response))

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
litellm	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Sep 17, 2024 7:06pm

[dbczumar]

@krrishdholakia I refiled this against main, since https://github.com/BerriAI/litellm/pull/5733/files was closed when the target branch was deleted. Can you take another look?

[krrishdholakia]

Yup. Assigning myself to this ticket. Great work on this @dbczumar

[krrishdholakia]

Aiming to have this merged into tomorrow's release

[dbczumar]

@krrishdholakia Did you get a chance to take a look at this?

[krrishdholakia]

thanks for the bump @dbczumar