Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions src/config/entities/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,7 @@ impl<T: Clone + 'static> ResourceStore<T> {
deleted
}

#[cfg(test)]
/// Get an entry by its primary key
fn get(&self, key: &str) -> Option<ResourceEntry<T>> {
self.data.load().primary.get(key).cloned()
}
Expand Down Expand Up @@ -295,16 +295,16 @@ impl<T: DeserializeOwned + Clone + Send + Sync + 'static> EntityStore<T> {
Self { store }
}

/// Get a snapshot of all entries
pub fn list(&self) -> Arc<HashMap<String, ResourceEntry<T>>> {
self.store.primary_snapshot()
}

/// Get the value of the specified key
#[cfg(test)]
pub fn get(&self, key: &str) -> Option<ResourceEntry<T>> {
self.store.get(key)
}

pub fn list(&self) -> Arc<HashMap<String, ResourceEntry<T>>> {
self.store.primary_snapshot()
}

/// Get an entry via a secondary index
fn get_by_secondary(&self, index: &str, sec_key: &str) -> Option<ResourceEntry<T>> {
self.store.get_by_secondary(index, sec_key)
Expand Down
12 changes: 11 additions & 1 deletion src/config/entities/models.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,10 @@ use utoipa::ToSchema;

use super::{ConfigProvider, EntityStore, IndexFn, ResourceEntry};
use crate::{
config::entities::types::{HasRateLimit, RateLimit, RateLimitMetric},
config::entities::{
Provider, ResourceRegistry,
types::{HasRateLimit, RateLimit, RateLimitMetric},
},
utils::jsonschema::format_evaluation_error,
};

Expand All @@ -32,6 +35,13 @@ pub struct Model {
pub rate_limit: Option<RateLimit>,
}

impl Model {
/// Get provider of current model
pub fn provider(&self, resources: &ResourceRegistry) -> Option<ResourceEntry<Provider>> {
resources.providers.get_by_id(&self.provider_id)
}
}
Comment thread
bzp2010 marked this conversation as resolved.

impl HasRateLimit for ResourceEntry<Model> {
fn rate_limit(&self) -> Option<RateLimit> {
self.rate_limit.clone()
Expand Down
4 changes: 4 additions & 0 deletions src/config/entities/providers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,10 @@ impl ProvidersStore {
pub fn list(&self) -> Arc<HashMap<String, ResourceEntry<Provider>>> {
self.store.list()
}

pub fn get_by_id(&self, id: &str) -> Option<ResourceEntry<Provider>> {
self.store.get(id)
}
}

#[cfg(test)]
Expand Down
23 changes: 21 additions & 2 deletions src/gateway/providers/azure.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use crate::gateway::{
provider_instance::ProviderAuth,
traits::{
ChatTransform, CompatQuirks, EmbedTransform, ProviderCapabilities, ProviderMeta,
provider::encode_path_segment,
ProviderSemanticConventions, provider::encode_path_segment,
},
types::{
embed::{EmbedRequestBody, EmbeddingRequest},
Expand Down Expand Up @@ -51,6 +51,14 @@ impl ProviderMeta for AzureDef {
DEFAULT_BASE_URL
}

fn semantic_conventions(&self) -> ProviderSemanticConventions {
ProviderSemanticConventions {
gen_ai_provider_name: "azure.ai.openai",
llm_system: "openai",
llm_provider: Some("azure"),
}
}

fn chat_endpoint_path(&self, model: &str) -> Cow<'static, str> {
Cow::Owned(format!(
"/openai/deployments/{}/chat/completions",
Expand Down Expand Up @@ -124,7 +132,10 @@ mod tests {
use super::{AzureDef, DEFAULT_API_VERSION};
use crate::gateway::{
provider_instance::ProviderAuth,
traits::{ChatTransform, EmbedTransform, ProviderCapabilities, ProviderMeta},
traits::{
ChatTransform, EmbedTransform, ProviderCapabilities, ProviderMeta,
ProviderSemanticConventions,
},
types::{embed::EmbedRequestBody, openai::ChatCompletionRequest},
};

Expand Down Expand Up @@ -171,6 +182,14 @@ mod tests {
);
assert_eq!(embed_url.query(), Some("api-version=v1"));
assert!(provider.as_embed_transform().is_some());
assert_eq!(
provider.semantic_conventions(),
ProviderSemanticConventions {
gen_ai_provider_name: "azure.ai.openai",
llm_system: "openai",
llm_provider: Some("azure"),
}
);
}

#[test]
Expand Down
21 changes: 19 additions & 2 deletions src/gateway/providers/bedrock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use crate::gateway::{
provider_instance::ProviderAuth,
traits::{
ChatStreamState, ChatTransform, PreparedRequest, ProviderCapabilities, ProviderMeta,
StreamReaderKind, provider::encode_path_segment,
ProviderSemanticConventions, StreamReaderKind, provider::encode_path_segment,
},
types::openai::{ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse},
};
Expand Down Expand Up @@ -77,6 +77,14 @@ impl ProviderMeta for BedrockDef {
DEFAULT_BASE_URL
}

fn semantic_conventions(&self) -> ProviderSemanticConventions {
ProviderSemanticConventions {
gen_ai_provider_name: "aws.bedrock",
llm_system: "amazon",
llm_provider: Some("aws"),
}
}

fn chat_endpoint_path(&self, model: &str) -> Cow<'static, str> {
Cow::Owned(format!("/model/{}/converse", encode_path_segment(model)))
}
Expand Down Expand Up @@ -217,7 +225,7 @@ mod tests {
use super::{BedrockDef, BedrockProviderConfig};
use crate::gateway::{
provider_instance::ProviderAuth,
traits::{PreparedRequest, ProviderMeta},
traits::{PreparedRequest, ProviderMeta, ProviderSemanticConventions},
};

#[test]
Expand Down Expand Up @@ -262,6 +270,15 @@ mod tests {
fn build_url_uses_overlap_handling_and_encodes_model_ids_with_slashes() {
let provider = BedrockDef;

assert_eq!(
provider.semantic_conventions(),
ProviderSemanticConventions {
gen_ai_provider_name: "aws.bedrock",
llm_system: "amazon",
llm_provider: Some("aws"),
}
);

let url = provider.build_url(
"https://bedrock-runtime.us-east-1.amazonaws.com/model",
"inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
Expand Down
23 changes: 21 additions & 2 deletions src/gateway/providers/gemini.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,10 @@ use serde::{Deserialize, Serialize};
use crate::gateway::{
error::{GatewayError, Result},
provider_instance::ProviderAuth,
traits::{ChatTransform, EmbedTransform, ProviderCapabilities, ProviderMeta},
traits::{
ChatTransform, EmbedTransform, ProviderCapabilities, ProviderMeta,
ProviderSemanticConventions,
},
};

pub const IDENTIFIER: &str = "gemini";
Expand All @@ -30,6 +33,14 @@ impl ProviderMeta for GoogleDef {
"https://generativelanguage.googleapis.com/v1beta/openai"
}

fn semantic_conventions(&self) -> ProviderSemanticConventions {
ProviderSemanticConventions {
gen_ai_provider_name: "gcp.gemini",
llm_system: "gemini",
llm_provider: Some("google"),
}
}

fn chat_endpoint_path(&self, _model: &str) -> Cow<'static, str> {
Cow::Borrowed("/chat/completions")
}
Expand Down Expand Up @@ -67,7 +78,7 @@ mod tests {
use super::GoogleDef;
use crate::gateway::{
provider_instance::ProviderAuth,
traits::{EmbedTransform, ProviderCapabilities, ProviderMeta},
traits::{EmbedTransform, ProviderCapabilities, ProviderMeta, ProviderSemanticConventions},
};

#[test]
Expand All @@ -92,5 +103,13 @@ mod tests {
);
assert_eq!(headers["x-goog-api-key"], "gemini-key");
assert!(provider.as_embed_transform().is_some());
assert_eq!(
provider.semantic_conventions(),
ProviderSemanticConventions {
gen_ai_provider_name: "gcp.gemini",
llm_system: "gemini",
llm_provider: Some("google"),
}
);
}
}
1 change: 1 addition & 0 deletions src/gateway/traits/chat_format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ pub trait ChatFormat: Send + Sync + 'static {
type NativeStreamState: Default + Send + Unpin;

/// Stable format name used for logs and diagnostics.
#[allow(unused)]
fn name() -> &'static str;

/// Whether the request expects a streaming response.
Expand Down
2 changes: 1 addition & 1 deletion src/gateway/traits/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ pub use native::{
pub use native::{NativeOpenAIResponsesSupport, OpenAIResponsesNativeStreamState};
pub use provider::{
ChatTransform, CompatQuirks, EmbedTransform, PreparedRequest, ProviderCapabilities,
ProviderMeta, StreamReaderKind,
ProviderMeta, ProviderSemanticConventions, StreamReaderKind,
};
#[allow(unused_imports)]
pub use provider::{ImageGenTransform, SttTransform, TtsTransform};
37 changes: 36 additions & 1 deletion src/gateway/traits/provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,11 +45,29 @@ pub(crate) fn encode_path_segment(segment: &str) -> String {
utf8_percent_encode(segment, PATH_SEGMENT_ENCODE_SET).to_string()
}

/// OpenTelemetry and OpenInference semantic conventions for the provider.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ProviderSemanticConventions {
pub gen_ai_provider_name: &'static str,
pub llm_system: &'static str,
pub llm_provider: Option<&'static str>,
}
Comment thread
bzp2010 marked this conversation as resolved.

/// Provider metadata with no data transformation logic.
pub trait ProviderMeta: Send + Sync + 'static {
fn name(&self) -> &'static str;
fn default_base_url(&self) -> &'static str;

/// Get the provider's semantic conventions.
/// Used for OpenTelemetry and OpenInference semantic conventions.
fn semantic_conventions(&self) -> ProviderSemanticConventions {
ProviderSemanticConventions {
gen_ai_provider_name: self.name(),
llm_system: self.name(),
llm_provider: None,
}
}

/// Chat endpoint path for the provider. Implementations may use `model`
/// for providers whose route shape depends on the model name.
fn chat_endpoint_path(&self, _model: &str) -> Cow<'static, str> {
Expand Down Expand Up @@ -306,7 +324,10 @@ mod tests {
use http::HeaderMap;
use serde_json::json;

use super::{ChatTransform, CompatQuirks, EmbedTransform, ProviderMeta, StreamReaderKind};
use super::{
ChatTransform, CompatQuirks, EmbedTransform, ProviderMeta, ProviderSemanticConventions,
StreamReaderKind,
};
Comment thread
bzp2010 marked this conversation as resolved.
use crate::gateway::{
provider_instance::ProviderAuth,
traits::chat_format::ChatStreamState,
Expand Down Expand Up @@ -421,6 +442,20 @@ mod tests {
assert_eq!(body["stream_options"]["include_usage"], true);
}

#[test]
fn provider_meta_uses_name_based_default_semantic_conventions() {
let provider = DummyProvider;

assert_eq!(
provider.semantic_conventions(),
ProviderSemanticConventions {
gen_ai_provider_name: "dummy",
llm_system: "dummy",
llm_provider: None,
}
);
}

#[test]
fn apply_to_request_skips_stream_usage_for_non_streaming_requests() {
let quirks = CompatQuirks {
Expand Down
Loading
Loading