i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #371 from Kayphoon/feature/embedding-model-support

Browse Source 
feat: add embedding and rerank support
This commit is contained in:
fawney19
2026-05-04 12:56:11 +08:00
committed by GitHub
parent 8dd4135f7c 6bdd7792eb
commit 627eb4f335
87 changed files with 5517 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
crates/aether-data/src/repository/global_models/memory.rs
View File

@@ -497,6 +497,7 @@ impl GlobalModelWriteRepository for InMemoryGlobalModelReadRepository {
Some(global_model.display_name.clone()),
global_model.default_price_per_request,
global_model.default_tiered_pricing.clone(),
global_model.supported_capabilities.clone(),
global_model.config.clone(),
)?;
self.admin_provider_model_items
@@ -542,6 +543,7 @@ impl GlobalModelWriteRepository for InMemoryGlobalModelReadRepository {
existing.global_model_display_name = Some(global_model.display_name.clone());
existing.global_model_default_price_per_request = global_model.default_price_per_request;
existing.global_model_default_tiered_pricing = global_model.default_tiered_pricing.clone();
existing.global_model_supported_capabilities = global_model.supported_capabilities.clone();
existing.global_model_config = global_model.config.clone();
Ok(Some(existing.clone()))
}
@@ -639,8 +641,9 @@ mod tests {
use super::InMemoryGlobalModelReadRepository;
use crate::repository::global_models::{
GlobalModelReadRepository, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
PublicGlobalModelQuery, StoredPublicCatalogModel, StoredPublicGlobalModel,
CreateAdminGlobalModelRecord, GlobalModelReadRepository, GlobalModelWriteRepository,
PublicCatalogModelListQuery, PublicCatalogModelSearchQuery, PublicGlobalModelQuery,
StoredPublicCatalogModel, StoredPublicGlobalModel,
};
fn sample_model(
@@ -687,11 +690,83 @@ mod tests {
Some(true),
Some(true),
Some(true),
Some(false),
true,
)
.expect("public catalog model should build")
}
#[tokio::test]
async fn embedding_model_metadata_roundtrip() {
let repository =
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new());
let record = CreateAdminGlobalModelRecord::new(
"gm-embedding".to_string(),
"text-embedding-3-small".to_string(),
"Text Embedding 3 Small".to_string(),
true,
None,
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":0.02}]})),
Some(json!(["embedding"])),
Some(json!({
"api_formats": ["openai:embedding"],
"dimensions": 1536
})),
)
.expect("embedding global model should validate");
repository
.create_admin_global_model(&record)
.await
.expect("embedding global model should persist")
.expect("embedding global model should be returned");
let stored = repository
.get_admin_global_model_by_name("text-embedding-3-small")
.await
.expect("embedding global model should read")
.expect("embedding global model should exist");
assert_eq!(stored.supported_capabilities, Some(json!(["embedding"])));
assert_eq!(
stored
.config
.as_ref()
.and_then(|value| value.get("dimensions")),
Some(&json!(1536))
);
assert_eq!(
stored
.default_tiered_pricing
.as_ref()
.and_then(|value| value.get("tiers"))
.and_then(serde_json::Value::as_array)
.and_then(|tiers| tiers.first())
.and_then(|tier| tier.get("input_price_per_1m"))
.and_then(serde_json::Value::as_f64),
Some(0.02)
);
}
#[tokio::test]
async fn embedding_missing_billing_config_rejected() {
let error = CreateAdminGlobalModelRecord::new(
"gm-embedding".to_string(),
"text-embedding-3-small".to_string(),
"Text Embedding 3 Small".to_string(),
true,
None,
None,
Some(json!(["embedding"])),
None,
)
.expect_err("embedding metadata without billing should fail closed");
assert!(error
.to_string()
.contains("embedding global model requires"));
}
#[tokio::test]
async fn defaults_to_active_models_only() {
let repository = InMemoryGlobalModelReadRepository::seed(vec![
@@ -791,6 +866,52 @@ mod tests {
assert_eq!(items[0].name, "gpt-5");
}
#[tokio::test]
async fn public_catalog_preserves_embedding_capability_without_contaminating_chat_models() {
let mut embedding_model = sample_public_catalog_model(
"model-embedding",
"provider-openai",
"openai",
"text-embedding-3-small",
"text-embedding-3-small",
"Text Embedding 3 Small",
);
embedding_model.supports_embedding = Some(true);
embedding_model.supports_streaming = Some(false);
let chat_model = sample_public_catalog_model(
"model-chat",
"provider-openai",
"openai",
"gpt-5-upstream",
"gpt-5",
"GPT 5",
);
let repository =
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new())
.with_public_catalog_models(vec![embedding_model, chat_model]);
let items = repository
.list_public_catalog_models(&PublicCatalogModelListQuery {
provider_id: Some("provider-openai".to_string()),
offset: 0,
limit: 50,
})
.await
.expect("catalog should list");
let embedding = items
.iter()
.find(|item| item.name == "text-embedding-3-small")
.expect("embedding model should be listed");
let chat = items
.iter()
.find(|item| item.name == "gpt-5")
.expect("chat model should be listed");
assert_eq!(embedding.supports_embedding, Some(true));
assert_eq!(embedding.supports_streaming, Some(false));
assert_eq!(chat.supports_embedding, Some(false));
}
#[tokio::test]
async fn searches_public_catalog_models_by_provider_and_display_name() {
let repository =

65
crates/aether-data/src/repository/global_models/sql.rs
View File

@@ -60,6 +60,27 @@ SELECT
COALESCE(m.supports_vision, CAST(gm.config->>'vision' AS BOOLEAN), FALSE) AS supports_vision,
COALESCE(m.supports_function_calling, CAST(gm.config->>'function_calling' AS BOOLEAN), FALSE) AS supports_function_calling,
COALESCE(m.supports_streaming, CAST(gm.config->>'streaming' AS BOOLEAN), TRUE) AS supports_streaming,
(
COALESCE(gm.supported_capabilities::jsonb @> '["embedding"]'::jsonb, FALSE)
OR LOWER(COALESCE(gm.config->>'embedding', 'false')) = 'true'
OR LOWER(COALESCE(gm.config->>'model_type', '')) = 'embedding'
OR LOWER(COALESCE(gm.config->>'type', '')) = 'embedding'
OR COALESCE(gm.config->'capabilities' @> '["embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'supported_capabilities' @> '["embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'api_formats' @> '["openai:embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'api_formats' @> '["jina:embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'api_formats' @> '["gemini:embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'api_formats' @> '["doubao:embedding"]'::jsonb, FALSE)
OR LOWER(COALESCE(m.config->>'embedding', 'false')) = 'true'
OR LOWER(COALESCE(m.config->>'model_type', '')) = 'embedding'
OR LOWER(COALESCE(m.config->>'type', '')) = 'embedding'
OR COALESCE(m.config::jsonb->'capabilities' @> '["embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'supported_capabilities' @> '["embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'api_formats' @> '["openai:embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'api_formats' @> '["jina:embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'api_formats' @> '["gemini:embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'api_formats' @> '["doubao:embedding"]'::jsonb, FALSE)
) AS supports_embedding,
m.is_active
FROM models m
JOIN providers p ON p.id = m.provider_id
@@ -98,6 +119,7 @@ SELECT
gm.display_name AS global_model_display_name,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
LEFT JOIN global_models gm ON gm.id = m.global_model_id
@@ -302,6 +324,7 @@ SELECT
gm.display_name AS global_model_display_name,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
LEFT JOIN global_models gm ON gm.id = m.global_model_id
@@ -348,6 +371,7 @@ SELECT
gm.display_name AS global_model_display_name,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
JOIN global_models gm ON gm.id = m.global_model_id
@@ -487,6 +511,7 @@ SELECT
gm.display_name AS global_model_display_name,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
LEFT JOIN global_models gm ON gm.id = m.global_model_id
@@ -1020,7 +1045,7 @@ fn apply_public_catalog_model_filters(
provider_id: Option<&str>,
search: Option<&str>,
) {
builder.push(" WHERE m.is_active = TRUE AND p.is_active = TRUE");
builder.push(" WHERE m.is_active = TRUE AND COALESCE(m.is_available, TRUE) = TRUE AND p.is_active = TRUE AND COALESCE(gm.is_active, TRUE) = TRUE");
if let Some(provider_id) = provider_id.map(str::trim).filter(|value| !value.is_empty()) {
builder
@@ -1060,6 +1085,7 @@ fn map_public_catalog_model_row(row: &PgRow) -> Result<StoredPublicCatalogModel,
row.try_get("supports_function_calling")
.map_postgres_err()?,
row.try_get("supports_streaming").map_postgres_err()?,
row.try_get("supports_embedding").map_postgres_err()?,
row.try_get("is_active").map_postgres_err()?,
)
}
@@ -1101,6 +1127,8 @@ fn map_admin_provider_model_row(row: &PgRow) -> Result<StoredAdminProviderModel,
.map_postgres_err()?,
row.try_get("global_model_default_tiered_pricing")
.map_postgres_err()?,
row.try_get("global_model_supported_capabilities")
.map_postgres_err()?,
row.try_get("global_model_config").map_postgres_err()?,
)
}
@@ -1177,9 +1205,42 @@ fn map_provider_active_global_model_row(
#[cfg(test)]
mod tests {
use super::SqlxGlobalModelReadRepository;
use super::{SqlxGlobalModelReadRepository, LIST_ADMIN_PROVIDER_MODELS_PREFIX};
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
const ADMIN_PROVIDER_MODEL_REQUIRED_COLUMNS: &[&str] = &[
"global_model_default_tiered_pricing",
"global_model_supported_capabilities",
"global_model_config",
];
fn assert_admin_provider_model_projection_has_required_columns(sql: &str) {
for column in ADMIN_PROVIDER_MODEL_REQUIRED_COLUMNS {
assert!(
sql.contains(column),
"admin provider model SQL projection should include {column}"
);
}
}
#[test]
fn admin_provider_model_sql_projections_include_supported_capabilities() {
assert_admin_provider_model_projection_has_required_columns(
LIST_ADMIN_PROVIDER_MODELS_PREFIX,
);
assert_admin_provider_model_projection_has_required_columns(include_str!("sql.rs"));
let supported_capabilities_projection = format!(
"{} AS {}",
"gm.supported_capabilities", "global_model_supported_capabilities"
);
assert_eq!(
include_str!("sql.rs")
.matches(&supported_capabilities_projection)
.count(),
4
);
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {