Authorize access to language model providers based on country (#15859)

This PR updates the LLM service to authorize access to language model
providers based on the requester's country.

We detect the country using Cloudflare's
[`CF-IPCountry`](https://developers.cloudflare.com/fundamentals/reference/http-request-headers/#cf-ipcountry)
header.

The country code is then checked against the list of supported countries
for the given LLM provider. Countries that are not supported will
receive an `HTTP 451: Unavailable For Legal Reasons` response.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2024-08-06 11:49:04 -04:00 committed by GitHub
parent 9c6ccaffe3
commit cf5f4dddf5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 921 additions and 1 deletions

View file

@ -1,7 +1,11 @@
mod authorization;
mod token;
use crate::api::CloudflareIpCountryHeader;
use crate::llm::authorization::authorize_access_to_language_model;
use crate::{executor::Executor, Config, Error, Result};
use anyhow::Context as _;
use axum::TypedHeader;
use axum::{
body::Body,
http::{self, HeaderName, HeaderValue, Request, StatusCode},
@ -91,9 +95,18 @@ async fn validate_api_token<B>(mut req: Request<B>, next: Next<B>) -> impl IntoR
async fn perform_completion(
Extension(state): Extension<Arc<LlmState>>,
Extension(_claims): Extension<LlmTokenClaims>,
Extension(claims): Extension<LlmTokenClaims>,
country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
Json(params): Json<PerformCompletionParams>,
) -> Result<impl IntoResponse> {
authorize_access_to_language_model(
&state.config,
&claims,
country_code_header.map(|header| header.to_string()),
params.provider,
&params.model,
)?;
match params.provider {
LanguageModelProvider::Anthropic => {
let api_key = state