Databricks Provider

# Configure using the CLI
goose configure

# Or set environment variables
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="dapi..."

export DATABRICKS_TOKEN="dapi1234567890abcdef"

# Only set the host
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"

# When you run goose, you'll be prompted to authenticate via browser
goose session start

use goose::providers::create;
use goose::model::ModelConfig;

// Create with default model
let model_config = ModelConfig::new("databricks-claude-sonnet-4")?;
let provider = create("databricks", model_config, vec![]).await?;

// Stream a response
let messages = vec![Message::user().with_text("Hello!")];
let stream = provider.stream(
    &provider.get_model_config(),
    "session-123",
    "You are a helpful assistant.",
    &messages,
    &[],
).await?;

let model_config = ModelConfig::new("databricks-meta-llama-3-3-70b-instruct")?
    .with_temperature(0.7)
    .with_max_tokens(2048);

let provider = create("databricks", model_config, vec![]).await?;

// Automatically tries databricks-claude-haiku-4-5 with 0 retries
let (response, usage) = provider.complete_fast(
    "session-123",
    "You are a helpful assistant.",
    &messages,
    &[],
).await?;

if provider.supports_embeddings() {
    let texts = vec![
        "Hello world".to_string(),
        "Databricks embeddings".to_string(),
    ];
    
    let embeddings = provider.create_embeddings(
        "session-123",
        texts,
    ).await?;
    
    println!("Generated {} embeddings", embeddings.len());
}

export DATABRICKS_MAX_RETRIES="5"
export DATABRICKS_INITIAL_RETRY_INTERVAL_MS="2000"
export DATABRICKS_BACKOFF_MULTIPLIER="2.5"
export DATABRICKS_MAX_RETRY_INTERVAL_MS="120000"

// Regular models
GET /serving-endpoints/{model_name}/invocations

// Codex models (if model name contains "codex")
GET /serving-endpoints/responses

// Embeddings
GET /serving-endpoints/text-embedding-3-small/invocations

impl ProviderDef for DatabricksProvider {
    fn metadata() -> ProviderMetadata {
        ProviderMetadata::new(
            "databricks",
            "Databricks",
            "Models on Databricks AI Gateway",
            "databricks-claude-sonnet-4",
            DATABRICKS_KNOWN_MODELS.to_vec(),
            "https://docs.databricks.com/en/generative-ai/external-models/",
            vec![
                ConfigKey::new("DATABRICKS_HOST", true, false, None, true),
                ConfigKey::new("DATABRICKS_TOKEN", false, true, None, true),
            ],
        )
    }
}

pub enum DatabricksAuth {
    Token(String),
    OAuth {
        host: String,
        client_id: String,
        redirect_url: String,
        scopes: Vec<String>,
    },
}

impl AuthProvider for DatabricksAuthProvider {
    async fn get_auth_header(&self) -> Result<(String, String)> {
        let token = match &self.auth {
            DatabricksAuth::Token(token) => token.clone(),
            DatabricksAuth::OAuth { host, client_id, redirect_url, scopes } => {
                oauth::get_oauth_token_async(host, client_id, redirect_url, scopes).await?
            }
        };
        Ok(("Authorization".to_string(), format!("Bearer {}", token)))
    }
}

{
  "messages": [
    {
      "role": "system",
      "content": "You are a helpful assistant."
    },
    {
      "role": "user",
      "content": "Hello!"
    }
  ],
  "temperature": 0.7,
  "max_tokens": 2048,
  "stream": true
}

pub struct RetryConfig {
    pub max_retries: usize,
    pub initial_interval_ms: u64,
    pub backoff_multiplier: f64,
    pub max_interval_ms: u64,
}

// Default config
RetryConfig {
    max_retries: 3,
    initial_interval_ms: 1000,
    backoff_multiplier: 2.0,
    max_interval_ms: 60000,
}

// Fast model config (fail fast)
RetryConfig {
    max_retries: 0,
    initial_interval_ms: 0,
    backoff_multiplier: 1.0,
    max_interval_ms: 0,
}

// Get all serving endpoints
let models = provider.fetch_supported_models().await?;

// Queries: GET /api/2.0/serving-endpoints
// Returns endpoint names that can be used as model names

{
  "endpoints": [
    {
      "name": "databricks-claude-sonnet-4-5",
      "creator": "[email protected]",
      "creation_timestamp": 1234567890,
      "config": { ... }
    }
  ]
}

match provider.stream(...).await {
    Ok(stream) => { /* handle stream */ },
    Err(ProviderError::Authentication(msg)) => {
        eprintln!("Auth failed: {}", msg);
        eprintln!("Try running: goose configure");
    },
    Err(ProviderError::RateLimited { retry_after }) => {
        eprintln!("Rate limited");
    },
    Err(e) => eprintln!("Error: {}", e),
}

use goose::providers::databricks::DatabricksProvider;

let provider = DatabricksProvider::from_params(
    "https://your-workspace.cloud.databricks.com".to_string(),
    "dapi1234567890abcdef".to_string(),
    model_config,
)?;