diff --git a/crates/aisix-core/src/models/model.rs b/crates/aisix-core/src/models/model.rs index ab01e541..77e54085 100644 --- a/crates/aisix-core/src/models/model.rs +++ b/crates/aisix-core/src/models/model.rs @@ -25,9 +25,17 @@ pub enum Provider { Anthropic, Google, Deepseek, - /// Cohere — currently exposed for `/v1/rerank` only (#213 Phase 1). - /// Cohere's chat / generate APIs are not OpenAI-compatible; a - /// future bridge implementation can extend coverage. + /// Cohere — `/v1/rerank` (native, via `aisix-proxy::rerank`) and + /// chat/completions / embeddings (via `OpenAiBridge::with_name("cohere")` + /// against `https://api.cohere.com/compatibility/v1`, per + /// ). Chat-compat is the + /// OpenAI-shape namespace Cohere ships alongside its native + /// `/v1/chat`. Chat-compat coverage today: the `command-r` / + /// `command-a` family per + /// ; legacy + /// `command` / `command-light` / `command-nightly` go to native + /// (not yet bridged) and will return upstream 400 if configured + /// against the chat-compat path. Cohere, /// Jina AI — currently exposed for `/v1/rerank` only (#213 Phase 2). /// Jina's rerank wire shape is identity-mapped to the OpenAI-compat diff --git a/crates/aisix-provider-openai/src/bridge.rs b/crates/aisix-provider-openai/src/bridge.rs index 0086c31c..d8497bd3 100644 --- a/crates/aisix-provider-openai/src/bridge.rs +++ b/crates/aisix-provider-openai/src/bridge.rs @@ -61,6 +61,18 @@ const DEEPSEEK_DEFAULT_BASE: &str = "https://api.deepseek.com"; /// dispatches to Google's OpenAI-compatible Gemini endpoint. const GOOGLE_DEFAULT_BASE: &str = "https://generativelanguage.googleapis.com/v1beta/openai"; +/// Fallback host for the `cohere`-named variant of this bridge, so a +/// `with_name("cohere")` instance without an explicit `api_base` +/// dispatches to Cohere's [OpenAI-compatible chat +/// endpoint](https://docs.cohere.com/reference/chat) at +/// `/compatibility/v1`. Cohere's native chat surface (`/v1/chat`) has a +/// different wire shape; the `/compatibility/v1` namespace mirrors the +/// OpenAI `/chat/completions` shape verbatim, so `OpenAiBridge` can +/// serve it directly. `Provider::Cohere.default_base_url()` returns +/// the bare host because the rerank path (`/v1/rerank`) builds its own +/// URL — that constant stays as-is. +const COHERE_DEFAULT_BASE: &str = "https://api.cohere.com/compatibility/v1"; + /// Path suffixes the bridge appends to `api_base` when building upstream /// URLs. If an operator accidentally pastes the full upstream URL into /// `api_base` (e.g. `https://api.openai.com/v1/chat/completions`), @@ -108,6 +120,7 @@ impl OpenAiBridge { match self.name { "deepseek" => DEEPSEEK_DEFAULT_BASE, "google" => GOOGLE_DEFAULT_BASE, + "cohere" => COHERE_DEFAULT_BASE, _ => OPENAI_DEFAULT_BASE, } } @@ -184,6 +197,7 @@ fn normalize_api_base(base: &str, provider: &str) -> String { match provider { "openai" => normalize_canonical_openai(stripped), "deepseek" => normalize_canonical_deepseek(stripped), + "cohere" => normalize_canonical_cohere(stripped), _ => stripped.to_string(), } } @@ -220,6 +234,28 @@ fn normalize_canonical_deepseek(base: &str) -> String { base.to_string() } +/// Canonical Cohere hosts. +const COHERE_CANONICAL_HOSTS: &[&str] = &["https://api.cohere.com", "http://api.cohere.com"]; + +/// Add the `/compatibility/v1` segment if and only if the operator +/// pasted the bare canonical Cohere host. The Cohere rerank endpoint +/// at `/v1/rerank` builds its own URL outside this bridge — that path +/// continues to use the bare host. For chat completions Cohere serves +/// an OpenAI-shape envelope at `/compatibility/v1/chat/completions` +/// (per ), so the bridge +/// synthesizes the right prefix when the operator left it off. +/// +/// Anything past the host root is left as-is — corporate proxies and +/// alternative deployments win. +fn normalize_canonical_cohere(base: &str) -> String { + for host in COHERE_CANONICAL_HOSTS { + if base == *host { + return format!("{host}/compatibility/v1"); + } + } + base.to_string() +} + fn api_key(ctx: &BridgeContext) -> Result<&str, BridgeError> { let k = &ctx.provider_key.secret; if k.is_empty() { @@ -1144,6 +1180,97 @@ data: [DONE]\n\n"; ); } + /// `with_name("cohere")` default must target Cohere's + /// OpenAI-compatible `/compatibility/v1` namespace rather than + /// falling through to OpenAI's host. The dashboard placeholder + /// (`https://api.cohere.com`) is the rerank path; for chat the + /// bridge synthesizes the right suffix (closes #332). + #[test] + fn cohere_default_base_targets_compatibility_v1() { + let bridge = OpenAiBridge::new().with_name("cohere"); + let pk: ProviderKey = serde_json::from_str(r#"{"display_name":"x","secret":"k"}"#).unwrap(); + let ctx = BridgeContext::new("rid", sample_model(), Arc::new(pk)); + assert_eq!( + bridge.resolve_base(&ctx), + "https://api.cohere.com/compatibility/v1", + ); + } + + /// Operators copy-paste the bare canonical Cohere host + /// (`https://api.cohere.com`) from rerank docs / the dashboard + /// placeholder. For chat the bridge synthesizes + /// `/compatibility/v1` so a misconfigured-but-recoverable + /// `api_base` still routes to Cohere's chat-compat endpoint. + /// Non-canonical hosts (corporate proxies) pass through verbatim + /// — operator-intent on a custom host wins. + #[test] + fn cohere_api_base_tolerance_bare_host_synthesizes_compatibility_prefix() { + let bridge = OpenAiBridge::new().with_name("cohere"); + let canonical = "https://api.cohere.com/compatibility/v1"; + + for form in [ + "https://api.cohere.com", + "https://api.cohere.com/", + "https://api.cohere.com/compatibility/v1", + "https://api.cohere.com/compatibility/v1/", + "https://api.cohere.com/compatibility/v1/chat/completions", + ] { + let ctx = BridgeContext::new("rid", sample_model(), Arc::new(pk_with_base(form))); + assert_eq!( + bridge.resolve_base(&ctx), + canonical, + "form {form:?} should normalize to {canonical}", + ); + } + + // Non-canonical host passes through after suffix stripping. + let custom = "https://proxy.acme.internal/cohere-chat"; + let ctx = BridgeContext::new("rid", sample_model(), Arc::new(pk_with_base(custom))); + assert_eq!( + bridge.resolve_base(&ctx), + custom, + "non-canonical host must NOT be rewritten", + ); + } + + /// End-to-end chat flow through the `cohere`-named bridge: + /// outbound URL matches the chat-compat namespace and the + /// OpenAI envelope round-trips without translation. Pins the + /// contract Hub.register relies on. + #[tokio::test] + async fn cohere_chat_compat_round_trips_openai_envelope() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .and(header("authorization", "Bearer cohere-key")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-cohere", + "model": "command-r", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello from cohere"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7} + }))) + .mount(&server) + .await; + + let bridge = OpenAiBridge::new().with_name("cohere"); + let pk_json = format!( + r#"{{"display_name":"cohere-prod","secret":"cohere-key","api_base":"{}"}}"#, + server.uri() + ); + let pk: Arc = Arc::new(serde_json::from_str(&pk_json).unwrap()); + let ctx = BridgeContext::new("rid", sample_model(), pk); + let resp = bridge.chat(&req(), &ctx).await.unwrap(); + + assert_eq!(resp.id, "cmpl-cohere"); + assert_eq!(resp.message.role, Role::Assistant); + assert_eq!(resp.message.content, "hello from cohere"); + assert_eq!(resp.usage.total_tokens, 7); + } + /// For Gemini and any future `with_name` variant, the bridge does not /// synthesize a `/v1beta/openai` prefix — the path is non-trivial. It /// still strips an accidentally-pasted endpoint suffix. diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 68762afa..49bea5c4 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -766,12 +766,20 @@ fn load_heartbeat_config_from_disk( /// Hub. The Hub is created once at startup; future dynamic reload /// lands behind the same `register()` call. /// -/// `Provider::Cohere` and `Provider::Jina` are intentionally NOT -/// registered: per #213 Phases 1–2 they are exposed only via -/// `/v1/rerank`, which is a verbatim HTTP forward (`aisix-proxy:: -/// rerank`) and bypasses the Bridge trait entirely. A bridge for -/// either provider would be needed only when chat completions / -/// embeddings on those providers are added. +/// `Provider::Jina` is intentionally NOT registered: per #213 Phase 2 +/// Jina is exposed only via `/v1/rerank`, which is a verbatim HTTP +/// forward (`aisix-proxy::rerank`) and bypasses the Bridge trait +/// entirely. +/// +/// `Provider::Cohere` is registered against the OpenAI-compatible +/// chat endpoint at `https://api.cohere.com/compatibility/v1` (per +/// ). Cohere's rerank surface +/// at `/v1/rerank` continues to bypass the Bridge via +/// `aisix-proxy::rerank` — the bridge here only serves `chat/completions`, +/// `embeddings`, and the other OpenAI-shape endpoints the bridge +/// supports. The chat-compat namespace gives an exact OpenAI envelope +/// shape so `OpenAiBridge::with_name("cohere")` can serve it directly +/// (closes #332). fn build_hub() -> Hub { let hub = Hub::new(); hub.register(Provider::Openai, Arc::new(OpenAiBridge::new())); @@ -784,6 +792,10 @@ fn build_hub() -> Hub { Provider::Deepseek, Arc::new(OpenAiBridge::new().with_name("deepseek")), ); + hub.register( + Provider::Cohere, + Arc::new(OpenAiBridge::new().with_name("cohere")), + ); // Family bridges (issue #302 Phase A/D two-tier dispatch). The // legacy `Provider`-keyed register() above stays the live @@ -1075,4 +1087,39 @@ mod tests { let err = derive_cp_etcd_url(&m).unwrap_err(); assert!(err.to_string().contains("cp_base_url"), "unexpected: {err}"); } + + /// `build_hub()` must register `Provider::Cohere` against the + /// `with_name("cohere")` variant of [`OpenAiBridge`] — the only + /// thing that ties the Provider enum to the chat-compat URL + /// (closes #332). A regression that registered `OpenAiBridge::new()` + /// (default name = `"openai"`) or omitted the registration would + /// flip the bridge label on metrics and (more importantly) the + /// `default_base()` fallback, silently routing Cohere chat to + /// OpenAI's host. + #[test] + fn build_hub_registers_cohere_chat_compat_variant() { + let hub = build_hub(); + let bridge = hub + .get(aisix_core::Provider::Cohere) + .expect("Provider::Cohere must have a Hub bridge registered for chat-compat"); + assert_eq!( + bridge.name(), + "cohere", + "Hub.register(Provider::Cohere, …) MUST use OpenAiBridge::with_name(\"cohere\") — \ + a `with_name(\"openai\")` fallback would route Cohere chat to OpenAI's host", + ); + } + + /// Companion to the cohere check above: Jina deliberately stays + /// rerank-only per #213 Phase 2. A future PR that flips Jina to + /// chat-compat must update this assertion deliberately. + #[test] + fn build_hub_does_not_register_jina_for_chat() { + let hub = build_hub(); + assert!( + hub.get(aisix_core::Provider::Jina).is_none(), + "Provider::Jina is rerank-only (#213 Phase 2); a Hub registration here would \ + silently route /v1/chat/completions on Jina to whichever bridge name was picked", + ); + } } diff --git a/schemas/resources/model.schema.json b/schemas/resources/model.schema.json index e5f3eae5..93f24058 100644 --- a/schemas/resources/model.schema.json +++ b/schemas/resources/model.schema.json @@ -267,7 +267,7 @@ ] }, { - "description": "Cohere — currently exposed for `/v1/rerank` only (#213 Phase 1). Cohere's chat / generate APIs are not OpenAI-compatible; a future bridge implementation can extend coverage.", + "description": "Cohere — `/v1/rerank` (native, via `aisix-proxy::rerank`) and chat/completions / embeddings (via `OpenAiBridge::with_name(\"cohere\")` against `https://api.cohere.com/compatibility/v1`, per ). Chat-compat is the OpenAI-shape namespace Cohere ships alongside its native `/v1/chat`. Chat-compat coverage today: the `command-r` / `command-a` family per ; legacy `command` / `command-light` / `command-nightly` go to native (not yet bridged) and will return upstream 400 if configured against the chat-compat path.", "type": "string", "enum": [ "cohere"