127 lines
3.8 KiB
Rust
127 lines
3.8 KiB
Rust
use crate::http::context::{AppContext, OpenApiResult, bad_request};
|
|
use crate::llm::chat;
|
|
use crate::store::StoreOps;
|
|
use poem_openapi::{Object, OpenApi, Tags, payload::Json};
|
|
use reqwest::header::{HeaderMap, HeaderValue};
|
|
use serde::Deserialize;
|
|
use std::sync::Arc;
|
|
|
|
const ANTHROPIC_MODELS_URL: &str = "https://api.anthropic.com/v1/models";
|
|
const ANTHROPIC_VERSION: &str = "2023-06-01";
|
|
const KEY_ANTHROPIC_API_KEY: &str = "anthropic_api_key";
|
|
|
|
#[derive(Deserialize)]
|
|
struct AnthropicModelsResponse {
|
|
data: Vec<AnthropicModelInfo>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct AnthropicModelInfo {
|
|
id: String,
|
|
}
|
|
|
|
fn get_anthropic_api_key(ctx: &AppContext) -> Result<String, String> {
|
|
match ctx.store.get(KEY_ANTHROPIC_API_KEY) {
|
|
Some(value) => {
|
|
if let Some(key) = value.as_str() {
|
|
if key.is_empty() {
|
|
Err("Anthropic API key is empty. Please set your API key.".to_string())
|
|
} else {
|
|
Ok(key.to_string())
|
|
}
|
|
} else {
|
|
Err("Stored API key is not a string".to_string())
|
|
}
|
|
}
|
|
None => Err("Anthropic API key not found. Please set your API key.".to_string()),
|
|
}
|
|
}
|
|
|
|
#[derive(Deserialize, Object)]
|
|
struct ApiKeyPayload {
|
|
api_key: String,
|
|
}
|
|
|
|
#[derive(Tags)]
|
|
enum AnthropicTags {
|
|
Anthropic,
|
|
}
|
|
|
|
pub struct AnthropicApi {
|
|
ctx: Arc<AppContext>,
|
|
}
|
|
|
|
impl AnthropicApi {
|
|
pub fn new(ctx: Arc<AppContext>) -> Self {
|
|
Self { ctx }
|
|
}
|
|
}
|
|
|
|
#[OpenApi(tag = "AnthropicTags::Anthropic")]
|
|
impl AnthropicApi {
|
|
/// Check whether an Anthropic API key is stored.
|
|
///
|
|
/// Returns `true` if a non-empty key is present, otherwise `false`.
|
|
#[oai(path = "/anthropic/key/exists", method = "get")]
|
|
async fn get_anthropic_api_key_exists(&self) -> OpenApiResult<Json<bool>> {
|
|
let exists =
|
|
chat::get_anthropic_api_key_exists(self.ctx.store.as_ref()).map_err(bad_request)?;
|
|
Ok(Json(exists))
|
|
}
|
|
|
|
/// Store or update the Anthropic API key used for requests.
|
|
///
|
|
/// Returns `true` when the key is saved successfully.
|
|
#[oai(path = "/anthropic/key", method = "post")]
|
|
async fn set_anthropic_api_key(
|
|
&self,
|
|
payload: Json<ApiKeyPayload>,
|
|
) -> OpenApiResult<Json<bool>> {
|
|
chat::set_anthropic_api_key(self.ctx.store.as_ref(), payload.0.api_key)
|
|
.map_err(bad_request)?;
|
|
Ok(Json(true))
|
|
}
|
|
|
|
/// List available Anthropic models.
|
|
#[oai(path = "/anthropic/models", method = "get")]
|
|
async fn list_anthropic_models(&self) -> OpenApiResult<Json<Vec<String>>> {
|
|
let api_key = get_anthropic_api_key(self.ctx.as_ref()).map_err(bad_request)?;
|
|
let client = reqwest::Client::new();
|
|
let mut headers = HeaderMap::new();
|
|
headers.insert(
|
|
"x-api-key",
|
|
HeaderValue::from_str(&api_key).map_err(|e| bad_request(e.to_string()))?,
|
|
);
|
|
headers.insert(
|
|
"anthropic-version",
|
|
HeaderValue::from_static(ANTHROPIC_VERSION),
|
|
);
|
|
|
|
let response = client
|
|
.get(ANTHROPIC_MODELS_URL)
|
|
.headers(headers)
|
|
.send()
|
|
.await
|
|
.map_err(|e| bad_request(e.to_string()))?;
|
|
|
|
if !response.status().is_success() {
|
|
let status = response.status();
|
|
let error_text = response
|
|
.text()
|
|
.await
|
|
.unwrap_or_else(|_| "Unknown error".to_string());
|
|
return Err(bad_request(format!(
|
|
"Anthropic API error {status}: {error_text}"
|
|
)));
|
|
}
|
|
|
|
let body = response
|
|
.json::<AnthropicModelsResponse>()
|
|
.await
|
|
.map_err(|e| bad_request(e.to_string()))?;
|
|
let models = body.data.into_iter().map(|m| m.id).collect();
|
|
|
|
Ok(Json(models))
|
|
}
|
|
}
|