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

feat: add embedding and rerank support

Browse Source
This commit is contained in:
Kayphoon
2026-05-03 17:32:41 +08:00
parent 3e2eca4fd0
commit 5abe664d65
87 changed files with 5520 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 {