feat: add databricks format that can use claude 3.7 with thinking (#1834)

* databricks added a new "reasoning" type field for reasoner models
This commit is contained in:
Salman Mohammed
2025-03-26 13:43:55 -04:00
committed by GitHub
parent fb35d875b4
commit 4f2e193554
5 changed files with 1145 additions and 8 deletions

View File

@@ -161,7 +161,7 @@ impl Provider for AnthropicProvider {
headers.insert("x-api-key", self.api_key.parse().unwrap());
headers.insert("anthropic-version", "2023-06-01".parse().unwrap());
let is_thinking_enabled = std::env::var("ANTHROPIC_THINKING_ENABLED").is_ok();
let is_thinking_enabled = std::env::var("CLAUDE_THINKING_ENABLED").is_ok();
if self.model.model_name.starts_with("claude-3-7-sonnet-") && is_thinking_enabled {
// https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#extended-output-capabilities-beta
headers.insert("anthropic-beta", "output-128k-2025-02-19".parse().unwrap());

View File

@@ -7,7 +7,7 @@ use std::time::Duration;
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::formats::openai::{create_request, get_usage, response_to_message};
use super::formats::databricks::{create_request, get_usage, response_to_message};
use super::oauth;
use super::utils::{get_model, ImageFormat};
use crate::config::ConfigError;

View File

@@ -311,10 +311,10 @@ pub fn create_request(
}
// Add thinking parameters for claude-3-7-sonnet model
let is_thinking_enabled = std::env::var("ANTHROPIC_THINKING_ENABLED").is_ok();
let is_thinking_enabled = std::env::var("CLAUDE_THINKING_ENABLED").is_ok();
if model_config.model_name.starts_with("claude-3-7-sonnet-") && is_thinking_enabled {
// Minimum budget_tokens is 1024
let budget_tokens = std::env::var("ANTHROPIC_THINKING_BUDGET")
let budget_tokens = std::env::var("CLAUDE_THINKING_BUDGET")
.unwrap_or_else(|_| "16000".to_string())
.parse()
.unwrap_or(16000);
@@ -574,10 +574,10 @@ mod tests {
#[test]
fn test_create_request_with_thinking() -> Result<()> {
// Save the original env var value if it exists
let original_value = std::env::var("ANTHROPIC_THINKING_ENABLED").ok();
let original_value = std::env::var("CLAUDE_THINKING_ENABLED").ok();
// Set the env var for this test
std::env::set_var("ANTHROPIC_THINKING_ENABLED", "true");
std::env::set_var("CLAUDE_THINKING_ENABLED", "true");
// Execute the test
let result = (|| {
@@ -606,8 +606,8 @@ mod tests {
// Restore the original env var state
match original_value {
Some(val) => std::env::set_var("ANTHROPIC_THINKING_ENABLED", val),
None => std::env::remove_var("ANTHROPIC_THINKING_ENABLED"),
Some(val) => std::env::set_var("CLAUDE_THINKING_ENABLED", val),
None => std::env::remove_var("CLAUDE_THINKING_ENABLED"),
}
// Return the test result

File diff suppressed because it is too large Load Diff

View File

@@ -1,5 +1,6 @@
pub mod anthropic;
pub mod bedrock;
pub mod databricks;
pub mod gcpvertexai;
pub mod google;
pub mod openai;