Skip to content

Commit

Permalink
Add support for kluster.ai provider
Browse files Browse the repository at this point in the history
  • Loading branch information
thxsh committed Feb 2, 2025
1 parent 0ca4cf0 commit 1efb045
Show file tree
Hide file tree
Showing 6 changed files with 165 additions and 0 deletions.
6 changes: 6 additions & 0 deletions crates/goose-server/src/routes/providers_and_keys.json
Original file line number Diff line number Diff line change
Expand Up @@ -46,5 +46,11 @@
"description": "Connect to Azure OpenAI Service",
"models": ["gpt-4o", "gpt-4o-mini"],
"required_keys": ["AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_DEPLOYMENT_NAME"]
},
"kluster": {
"name": "Kluster",
"description": "Connect to Kluster",
"models": ["deepseek-ai/DeepSeek-R1", "klusterai/Meta-Llama-3.1-405B-Instruct-Turbo"],
"required_keys": ["KLUSTER_API_KEY"]
}
}
3 changes: 3 additions & 0 deletions crates/goose/src/providers/factory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use super::{
ollama::OllamaProvider,
openai::OpenAiProvider,
openrouter::OpenRouterProvider,
kluster::KlusterProvider,
};
use crate::model::ModelConfig;
use anyhow::Result;
Expand All @@ -22,6 +23,7 @@ pub fn providers() -> Vec<ProviderMetadata> {
OllamaProvider::metadata(),
OpenAiProvider::metadata(),
OpenRouterProvider::metadata(),
KlusterProvider::metadata(),
]
}

Expand All @@ -35,6 +37,7 @@ pub fn create(name: &str, model: ModelConfig) -> Result<Box<dyn Provider + Send
"ollama" => Ok(Box::new(OllamaProvider::from_env(model)?)),
"openrouter" => Ok(Box::new(OpenRouterProvider::from_env(model)?)),
"google" => Ok(Box::new(GoogleProvider::from_env(model)?)),
"kluster" => Ok(Box::new(KlusterProvider::from_env(model)?)),
_ => Err(anyhow::anyhow!("Unknown provider: {}", name)),
}
}
135 changes: 135 additions & 0 deletions crates/goose/src/providers/kluster.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,135 @@
use anyhow::Result;
use async_trait::async_trait;
use reqwest::Client;
use serde_json::Value;
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::utils::{emit_debug_trace, get_model, handle_response_openai_compat, ImageFormat};
use crate::message::Message;
use crate::model::ModelConfig;
use mcp_core::tool::Tool;

pub const KLUSTER_DEFAULT_MODEL: &str = "deepseek-ai/DeepSeek-R1";
pub const KLUSTER_KNOWN_MODELS: &[&str] = &[
"deepseek-ai/DeepSeek-R1",
"klusterai/Meta-Llama-3.1-405B-Instruct-Turbo",
];

pub const KLUSTER_DOC_URL: &str = "https://docs.kluster.ai/";

#[derive(Debug, serde::Serialize)]
pub struct KlusterProvider {
#[serde(skip)]
client: Client,
host: String,
api_key: String,
model: ModelConfig,
}

impl Default for KlusterProvider {
fn default() -> Self {
let model = ModelConfig::new(KlusterProvider::metadata().default_model);
KlusterProvider::from_env(model).expect("Failed to initialize Kluster provider")
}
}

impl KlusterProvider {
pub fn from_env(model: ModelConfig) -> Result<Self> {
let config = crate::config::Config::global();
let api_key: String = config.get_secret("KLUSTER_API_KEY")?;
let host: String = config
.get("KLUSTER_HOST")
.unwrap_or_else(|_| "https://api.kluster.ai/v1".to_string());
let client = Client::builder()
.timeout(Duration::from_secs(600))
.build()?;

Ok(Self {
client,
host,
api_key,
model,
})
}

async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
let base_url = url::Url::parse(&self.host)
.map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
let url = base_url.join("v1/chat/completions").map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;

let response = self
.client
.post(url)
.header("Authorization", format!("Bearer {}", self.api_key))
.json(&payload)
.send()
.await?;

handle_response_openai_compat(response).await
}
}

#[async_trait]
impl Provider for KlusterProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::new(
"kluster",
"Kluster",
"Kluster models",
KLUSTER_DEFAULT_MODEL,
KLUSTER_KNOWN_MODELS
.iter()
.map(|&s| s.to_string())
.collect(),
KLUSTER_DOC_URL,
vec![
ConfigKey::new("KLUSTER_API_KEY", true, true, None),
ConfigKey::new(
"KLUSTER_HOST",
false,
false,
Some("https://api.kluster.ai/v1"),
),
],
)
}

fn get_model_config(&self) -> ModelConfig {
self.model.clone()
}

#[tracing::instrument(
skip(self, system, messages, tools),
fields(model_config, input, output, input_tokens, output_tokens, total_tokens)
)]
async fn complete(
&self,
system: &str,
messages: &[Message],
tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
let payload = create_request(&self.model, system, messages, tools, &ImageFormat::OpenAi)?;

// Make request
let response = self.post(payload.clone()).await?;

// Parse response
let message = response_to_message(response.clone())?;
let usage = match get_usage(&response) {
Ok(usage) => usage,
Err(ProviderError::UsageError(e)) => {
tracing::warn!("Failed to get usage data: {}", e);
Usage::default()
}
Err(e) => return Err(e),
};
let model = get_model(&response);
emit_debug_trace(self, &payload, &response, &usage);
Ok((message, ProviderUsage::new(model, usage)))
}
}
1 change: 1 addition & 0 deletions crates/goose/src/providers/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ pub mod groq;
pub mod oauth;
pub mod ollama;
pub mod openai;
pub mod kluster;
pub mod openrouter;
pub mod utils;

Expand Down
9 changes: 9 additions & 0 deletions ui/desktop/src/components/ApiKeyWarning.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,11 @@ export GOOSE_PROVIDER__HOST=https://openrouter.ai
export GOOSE_PROVIDER__MODEL=anthropic/claude-3.5-sonnet
export GOOSE_PROVIDER__API_KEY=your_api_key_here`;

const KLUSTER_CONFIG = `export GOOSE_PROVIDER__TYPE=kluster
export GOOSE_PROVIDER__HOST=https://api.kluster.ai
export GOOSE_PROVIDER__MODEL=klusterai/Meta-Llama-3.1-405B-Instruct-Turbo
export GOOSE_PROVIDER__API_KEY=your_api_key_here`;

export function ApiKeyWarning({ className }: ApiKeyWarningProps) {
return (
<Card
Expand Down Expand Up @@ -81,6 +86,10 @@ export function ApiKeyWarning({ className }: ApiKeyWarningProps) {
<Collapsible title="OpenRouter Configuration">
<pre className="bg-gray-50 p-4 rounded-md text-sm">{OPENROUTER_CONFIG}</pre>
</Collapsible>

<Collapsible title="Kluster Configuration">
<pre className="bg-gray-50 p-4 rounded-md text-sm">{KLUSTER_CONFIG}</pre>
</Collapsible>
</div>
<p className="text-gray-600 mt-4">
After setting these variables, restart Goose for the changes to take effect.
Expand Down
11 changes: 11 additions & 0 deletions ui/desktop/src/components/settings/models/hardcoded_stuff.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ export const goose_models: Model[] = [
{ id: 16, name: 'qwen2.5', provider: 'Ollama' },
{ id: 17, name: 'anthropic/claude-3.5-sonnet', provider: 'OpenRouter' },
{ id: 18, name: 'gpt-4o', provider: 'Azure OpenAI' },
{ id: 19, name: 'klusterai/Meta-Llama-3.1-405B-Instruct-Turbo', provider: 'Kluster' },
];

export const openai_models = ['gpt-4o-mini', 'gpt-4o', 'gpt-4-turbo', 'o1'];
Expand Down Expand Up @@ -45,6 +46,11 @@ export const openrouter_models = ['anthropic/claude-3.5-sonnet'];

export const azure_openai_models = ['gpt-4o'];

export const kluster_models = [
'deepseek-ai/DeepSeek-R1',
'klusterai/Meta-Llama-3.1-405B-Instruct-Turbo',
];

export const default_models = {
openai: 'gpt-4o',
anthropic: 'claude-3-5-sonnet-latest',
Expand All @@ -54,6 +60,7 @@ export const default_models = {
openrouter: 'anthropic/claude-3.5-sonnet',
ollama: 'qwen2.5',
azure_openai: 'gpt-4o',
kluster: 'deepseek-ai/DeepSeek-R1',
};

export function getDefaultModel(key: string): string | undefined {
Expand All @@ -71,6 +78,7 @@ export const required_keys = {
Google: ['GOOGLE_API_KEY'],
OpenRouter: ['OPENROUTER_API_KEY'],
'Azure OpenAI': ['AZURE_OPENAI_API_KEY', 'AZURE_OPENAI_ENDPOINT', 'AZURE_OPENAI_DEPLOYMENT_NAME'],
Kluster: ['KLUSTER_API_KEY'],
};

export const supported_providers = [
Expand All @@ -82,6 +90,7 @@ export const supported_providers = [
'Ollama',
'OpenRouter',
'Azure OpenAI',
'Kluster',
];

export const model_docs_link = [
Expand All @@ -95,6 +104,7 @@ export const model_docs_link = [
},
{ name: 'OpenRouter', href: 'https://openrouter.ai/models' },
{ name: 'Ollama', href: 'https://ollama.com/library' },
{ name: 'Kluster', href: 'https://docs.kluster.ai' },
];

export const provider_aliases = [
Expand All @@ -106,4 +116,5 @@ export const provider_aliases = [
{ provider: 'OpenRouter', alias: 'openrouter' },
{ provider: 'Google', alias: 'google' },
{ provider: 'Azure OpenAI', alias: 'azure_openai' },
{ provider: 'Kluster', alias: 'kluster' },
];

0 comments on commit 1efb045

Please sign in to comment.