Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions crates/buzz-agent/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ Everything is environment variables. No flags, no config files. (We are a subpro

| Variable | Default | Notes |
|---|---|---|
| `BUZZ_AGENT_PROVIDER` | — | `anthropic`, `openai`, or `databricks`. If unset, or if `anthropic`/`openai` is selected but its API key is missing, Databricks is auto-selected when `DATABRICKS_HOST` + `DATABRICKS_MODEL` are set. |
| `BUZZ_AGENT_PROVIDER` | — | `anthropic`, `openai`, `databricks`, or `databricks_v2`. If unset, or if `anthropic`/`openai` is selected but its API key is missing, Databricks is auto-selected when `DATABRICKS_HOST` + `DATABRICKS_MODEL` are set. |
| `ANTHROPIC_API_KEY` | — | Required when provider=anthropic unless Databricks fallback is configured. |
| `ANTHROPIC_MODEL` | — | Required when provider=anthropic. |
| `ANTHROPIC_BASE_URL` | `https://api.anthropic.com` | |
Expand Down Expand Up @@ -158,7 +158,7 @@ Everything is environment variables. No flags, no config files. (We are a subpro

## Providers

`buzz-agent` speaks two HTTP dialects. Pick with `BUZZ_AGENT_PROVIDER`.
`buzz-agent` speaks a few HTTP dialects. Pick with `BUZZ_AGENT_PROVIDER`.

| Provider | `BUZZ_AGENT_PROVIDER` | Endpoint (auto) | Tested with |
|---|---|---|---|
Expand All @@ -170,14 +170,15 @@ Everything is environment variables. No flags, no config files. (We are a subpro
| OpenRouter | `openai` | `POST {base}/chat/completions` | anything they route |
| Block Gateway | `openai` | `POST {base}/chat/completions` | gpt-5, claude |
| Databricks | `databricks` | `POST {host}/serving-endpoints/{model}/invocations` | goose-claude-4-6-sonnet |
| Databricks AI Gateway v2 | `databricks_v2` | `POST {host}/ai-gateway/{provider}/v1/...` | databricks-gpt-5-5, databricks-claude-opus-4-7 |

If `BUZZ_AGENT_PROVIDER=anthropic` is selected without `ANTHROPIC_API_KEY`, or `BUZZ_AGENT_PROVIDER=openai` is selected without `OPENAI_COMPAT_API_KEY`, the agent automatically falls back to Databricks OAuth when `DATABRICKS_HOST` and `DATABRICKS_MODEL` are set. The same Databricks fallback applies when `BUZZ_AGENT_PROVIDER` is unset. Explicit Anthropic/OpenAI API keys always win.

`provider=openai` speaks two HTTP dialects: the [Responses API](https://platform.openai.com/docs/api-reference/responses) (`/v1/responses`, required for GPT-5 / o-series tool-calling on OpenAI's own service) and the [Chat Completions API](https://platform.openai.com/docs/api-reference/chat) (`/chat/completions`, the broadly-supported OpenAI-compatible wire format).

By default (`OPENAI_COMPAT_API=auto`) the agent picks **Responses** when `OPENAI_COMPAT_BASE_URL` points at an `*.openai.com` host and **Chat Completions** everywhere else. Pin the choice explicitly with `OPENAI_COMPAT_API=chat` or `OPENAI_COMPAT_API=responses` for providers that diverge from the default (e.g. a Responses-compatible self-hosted gateway).

`Provider` is a Rust `enum` with one `match` in `Llm::complete`. There is no trait, no `Box<dyn>`, no async-trait. Adding a third provider is a `match` arm and one `body`/`parse` pair in `llm.rs`.
`Provider` is a Rust `enum` with one `match` in `Llm::complete`. There is no trait, no `Box<dyn>`, no async-trait. Adding a provider is a `match` arm and one `body`/`parse` pair in `llm.rs`.

## MCP Servers

Expand Down
12 changes: 8 additions & 4 deletions crates/buzz-agent/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,9 @@ pub enum Provider {
/// with a dynamically-acquired bearer (OAuth 2.0 PKCE, or static `DATABRICKS_TOKEN`).
/// Wire format is OpenAI-chat-compatible — reuses the same body builder and parser.
Databricks,
/// Databricks AI Gateway v2. Routes by model family through the gateway's
/// OpenAI Responses, Anthropic Messages, or MLflow Chat Completions paths.
DatabricksV2,
}

/// Which OpenAI-family HTTP API to call. Set via `OPENAI_COMPAT_API`
Expand Down Expand Up @@ -114,8 +117,8 @@ impl Config {
// bad value can't break an Anthropic-only deployment.
//
// Databricks borrows api_key as the *optional* `DATABRICKS_TOKEN` escape
// hatch — empty means "use OAuth PKCE." The model lives in the URL path,
// not the request body (see `EndpointStrategy::DatabricksServing`).
// hatch — empty means "use OAuth PKCE." Legacy Databricks encodes the
// model in the URL path; Databricks v2 keeps it in the request body.
let (api_key, model, base_url, openai_api) = match provider {
Provider::Anthropic => (
req("ANTHROPIC_API_KEY")?,
Expand All @@ -137,12 +140,12 @@ impl Config {
env_or("OPENAI_COMPAT_BASE_URL", "https://api.openai.com/v1"),
parse_openai_api(env("OPENAI_COMPAT_API").as_deref())?,
),
Provider::Databricks => (
Provider::Databricks | Provider::DatabricksV2 => (
env("DATABRICKS_TOKEN").unwrap_or_default(),
resolve_model(databricks_model.as_deref(), buzz_agent_model.as_deref())
.ok_or_else(|| "config: DATABRICKS_MODEL required".to_string())?,
databricks_host.ok_or_else(|| "config: DATABRICKS_HOST required".to_string())?,
OpenAiApi::Chat, // Databricks invocations is chat-shaped
OpenAiApi::Chat, // unused by Databricks v2; legacy Databricks is chat-shaped
),
};
let system_prompt = match (env("BUZZ_AGENT_SYSTEM_PROMPT"), env("BUZZ_AGENT_SYSTEM_PROMPT_FILE")) {
Expand Down Expand Up @@ -318,6 +321,7 @@ fn resolve_provider(
"config: OPENAI_COMPAT_API_KEY required (or set DATABRICKS_HOST and DATABRICKS_MODEL for Databricks OAuth fallback)".into(),
),
"databricks" => Ok(Provider::Databricks),
"databricks_v2" | "databricks-v2" => Ok(Provider::DatabricksV2),
_ => Err(format!(
"config: BUZZ_AGENT_PROVIDER={raw} not supported"
)),
Expand Down
8 changes: 4 additions & 4 deletions crates/buzz-agent/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -78,13 +78,13 @@ pub fn run() -> Result<(), Box<dyn std::error::Error>> {
}

/// `buzz-agent auth <provider>` — run the interactive auth flow for a
/// provider and persist the result, then exit. Today the only provider is
/// `databricks` (OAuth 2.0 PKCE). Reads `DATABRICKS_HOST` from env; needs
/// a browser on the machine.
/// provider and persist the result, then exit. Today this supports Databricks
/// OAuth 2.0 PKCE. Reads `DATABRICKS_HOST` from env; needs a browser on the
/// machine.
async fn auth_subcommand(args: &[String]) -> Result<(), Box<dyn std::error::Error>> {
let provider = args.first().map(String::as_str);
match provider {
Some("databricks") => {
Some("databricks" | "databricks_v2" | "databricks-v2") => {
let host = std::env::var("DATABRICKS_HOST")
.map_err(|_| "auth databricks: DATABRICKS_HOST required")?;
let pkce = auth::PkceOAuthConfig {
Expand Down
126 changes: 125 additions & 1 deletion crates/buzz-agent/src/llm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,13 @@ const MAX_LLM_ERROR_BODY_BYTES: usize = 4 * 1024;
/// alongside its `_body` serializer.
type OpenAiParse = fn(Value) -> Result<LlmResponse, AgentError>;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DatabricksV2Route {
OpenAiResponses,
AnthropicMessages,
MlflowChatCompletions,
}

pub struct Llm {
http: Client,
/// One-shot sticky flag: set when a Chat Completions request comes
Expand Down Expand Up @@ -83,6 +90,23 @@ impl Llm {
})
.await
}
Provider::DatabricksV2 => {
self.databricks_v2_request(cfg, |route| match route {
DatabricksV2Route::OpenAiResponses => (
responses_body(cfg, system_prompt, history, tools),
parse_responses as OpenAiParse,
),
DatabricksV2Route::AnthropicMessages => (
anthropic_body(cfg, system_prompt, history, tools),
parse_anthropic as OpenAiParse,
),
DatabricksV2Route::MlflowChatCompletions => (
openai_body(cfg, system_prompt, history, tools),
parse_openai as OpenAiParse,
),
})
.await
}
}
}

Expand Down Expand Up @@ -137,6 +161,46 @@ impl Llm {
.await?;
Ok(r.text)
}
Provider::DatabricksV2 => {
let r = self
.databricks_v2_request(cfg, |route| match route {
DatabricksV2Route::OpenAiResponses => (
json!({
"model": cfg.model,
"max_output_tokens": max_output_tokens,
"instructions": system_prompt,
"input": user_prompt,
}),
parse_responses as OpenAiParse,
),
DatabricksV2Route::AnthropicMessages => (
json!({
"model": cfg.model,
"max_tokens": max_output_tokens,
"system": system_prompt,
"messages": [{
"role": "user",
"content": [{ "type": "text", "text": user_prompt }],
}],
}),
parse_anthropic as OpenAiParse,
),
DatabricksV2Route::MlflowChatCompletions => (
json!({
"model": cfg.model,
"stream": false,
"max_completion_tokens": max_output_tokens,
"messages": [
{ "role": "system", "content": system_prompt },
{ "role": "user", "content": user_prompt },
],
}),
parse_openai as OpenAiParse,
),
})
.await?;
Ok(r.text)
}
}
}

Expand Down Expand Up @@ -176,6 +240,22 @@ impl Llm {
}
}

async fn databricks_v2_request<F>(
&self,
cfg: &Config,
mut build: F,
) -> Result<LlmResponse, AgentError>
where
F: FnMut(DatabricksV2Route) -> (Value, OpenAiParse) + Send,
{
let route = databricks_v2_route_for_model(&cfg.model);
let (body, parse) = build(route);
parse(
self.post_openai(cfg, databricks_v2_path(route), &body)
.await?,
)
}

/// POST to an OpenAI-family endpoint. For OpenAI-compat this is just
/// `{base_url}{path}` with the body untouched. For Databricks the URL
/// becomes `{base_url}/serving-endpoints/{model}/invocations` and the
Expand Down Expand Up @@ -517,6 +597,25 @@ fn is_responses_required_error(body: &str) -> bool {
|| b.contains("use the responses api")
}

fn databricks_v2_route_for_model(model: &str) -> DatabricksV2Route {
let lower = model.to_ascii_lowercase();
if lower.contains("gpt-5") || lower.contains("gpt5") {
DatabricksV2Route::OpenAiResponses
} else if lower.contains("claude") {
DatabricksV2Route::AnthropicMessages
} else {
DatabricksV2Route::MlflowChatCompletions
}
}

fn databricks_v2_path(route: DatabricksV2Route) -> &'static str {
match route {
DatabricksV2Route::OpenAiResponses => "/ai-gateway/openai/v1/responses",
DatabricksV2Route::AnthropicMessages => "/ai-gateway/anthropic/v1/messages",
DatabricksV2Route::MlflowChatCompletions => "/ai-gateway/mlflow/v1/chat/completions",
}
}

fn parse_responses(v: Value) -> Result<LlmResponse, AgentError> {
let mut text = String::new();
let mut tool_calls = Vec::new();
Expand Down Expand Up @@ -886,7 +985,7 @@ fn build_token_source(cfg: &Config) -> Result<Arc<dyn TokenSource>, AgentError>
Provider::Anthropic | Provider::OpenAi => {
Ok(Arc::new(StaticTokenSource::new(cfg.api_key.clone())))
}
Provider::Databricks => {
Provider::Databricks | Provider::DatabricksV2 => {
if !cfg.api_key.is_empty() {
return Ok(Arc::new(StaticTokenSource::new(cfg.api_key.clone())));
}
Expand Down Expand Up @@ -1199,6 +1298,31 @@ mod tests {
}
}

#[test]
fn databricks_v2_routes_by_model_family() {
for (model, route, path) in [
(
"databricks-gpt-5-5",
DatabricksV2Route::OpenAiResponses,
"/ai-gateway/openai/v1/responses",
),
(
"databricks-claude-opus-4-7",
DatabricksV2Route::AnthropicMessages,
"/ai-gateway/anthropic/v1/messages",
),
(
"custom-tool-model",
DatabricksV2Route::MlflowChatCompletions,
"/ai-gateway/mlflow/v1/chat/completions",
),
] {
let got = databricks_v2_route_for_model(model);
assert_eq!(got, route, "model={model}");
assert_eq!(databricks_v2_path(got), path, "model={model}");
}
}

#[test]
fn parse_responses_rejects_malformed_function_arguments() {
let v = serde_json::json!({
Expand Down
Loading
Loading