How to use LLM pipeline library with models served by databricks?

[4x4notfound]

In my notebook I'm currently setting up an agent llm the following way:

from txtai import Agent, LLM

DATABRICKS_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

llm = LLM('Name-of-databricks-model-served', token=DATABRICKS_TOKEN, api_base='https://<my-db-instance>.cloud.databricks.com/serving-endpoints')

agent = Agent(
    model=llm,
    tools=[{
        "name": "test agent",
        "description": "what is the capital of Portugal",
        "path": "test",
    }],
    max_steps=10,
)

Currently it keeps trying to use the hugging face api.

401 Client Error: Unauthorized for url: https://huggingface.co/Name-of-databricks-model-served/resolve/main/config.json 

[davidmezzetti]

Hello. This documentation page from LiteLLM explains more on this: https://docs.litellm.ai/docs/providers/databricks

from txtai import Agent, LLM

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

llm = LLM('databricks/Name-of-databricks-model-served')

agent = Agent(
    model=llm,
    tools=[{
        "name": "test agent",
        "description": "what is the capital of Portugal",
        "path": "test",
    }],
    max_steps=10,
)