Implement generic pool member scoring and probing

This commit is contained in:
fawney19
2026-05-12 01:46:22 +08:00
parent cb0ccb9cdb
commit b9e62d1667
76 changed files with 6276 additions and 68 deletions

View File

@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
MysqlOAuthProviderRepository, OAuthProviderReadRepository, OAuthProviderWriteRepository,
};
use crate::repository::pool_scores::{
MysqlPoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
};
use crate::repository::provider_catalog::{
MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository,
@@ -184,6 +187,14 @@ impl MysqlBackend {
Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
}

View File

@@ -39,6 +39,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
};
use crate::repository::pool_scores::{
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PostgresPoolMemberScoreRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
SqlxProviderCatalogReadRepository,
@@ -195,6 +198,14 @@ impl PostgresBackend {
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
}

View File

@@ -14,6 +14,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingReadRepository;
use crate::repository::global_models::GlobalModelReadRepository;
use crate::repository::management_tokens::ManagementTokenReadRepository;
use crate::repository::oauth_providers::OAuthProviderReadRepository;
use crate::repository::pool_scores::PoolScoreReadRepository;
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
use crate::repository::quota::ProviderQuotaReadRepository;
@@ -34,6 +35,7 @@ pub struct DataReadRepositories {
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
pool_scores: Option<Arc<dyn PoolScoreReadRepository>>,
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
@@ -61,6 +63,7 @@ impl fmt::Debug for DataReadRepositories {
.field("has_global_models", &self.global_models.is_some())
.field("has_management_tokens", &self.management_tokens.is_some())
.field("has_oauth_providers", &self.oauth_providers.is_some())
.field("has_pool_scores", &self.pool_scores.is_some())
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field(
"has_minimal_candidate_selection",
@@ -124,6 +127,10 @@ impl DataReadRepositories {
.map(PostgresBackend::oauth_provider_read_repository)
.or_else(|| mysql.map(MysqlBackend::oauth_provider_read_repository))
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_read_repository)),
pool_scores: postgres
.map(PostgresBackend::pool_score_read_repository)
.or_else(|| mysql.map(MysqlBackend::pool_score_read_repository))
.or_else(|| sqlite.map(SqliteBackend::pool_score_read_repository)),
proxy_nodes: postgres
.map(PostgresBackend::proxy_node_read_repository)
.or_else(|| mysql.map(MysqlBackend::proxy_node_read_repository))
@@ -208,6 +215,10 @@ impl DataReadRepositories {
self.oauth_providers.clone()
}
pub fn pool_scores(&self) -> Option<Arc<dyn PoolScoreReadRepository>> {
self.pool_scores.clone()
}
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeReadRepository>> {
self.proxy_nodes.clone()
}
@@ -257,6 +268,7 @@ impl DataReadRepositories {
|| self.global_models.is_some()
|| self.management_tokens.is_some()
|| self.oauth_providers.is_some()
|| self.pool_scores.is_some()
|| self.proxy_nodes.is_some()
|| self.minimal_candidate_selection.is_some()
|| self.request_candidates.is_some()

View File

@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository,
};
use crate::repository::pool_scores::{
PoolMemberScoreWriteRepository, PoolScoreReadRepository, SqlitePoolMemberScoreRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
SqliteProviderCatalogReadRepository,
@@ -197,6 +200,14 @@ impl SqliteBackend {
Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
}

View File

@@ -11,6 +11,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository;
use crate::repository::global_models::GlobalModelWriteRepository;
use crate::repository::management_tokens::ManagementTokenWriteRepository;
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
use crate::repository::pool_scores::PoolMemberScoreWriteRepository;
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
use crate::repository::quota::ProviderQuotaWriteRepository;
@@ -30,6 +31,7 @@ pub struct DataWriteRepositories {
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
pool_scores: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
@@ -54,6 +56,7 @@ impl fmt::Debug for DataWriteRepositories {
.field("has_global_models", &self.global_models.is_some())
.field("has_management_tokens", &self.management_tokens.is_some())
.field("has_oauth_providers", &self.oauth_providers.is_some())
.field("has_pool_scores", &self.pool_scores.is_some())
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_provider_quotas", &self.provider_quotas.is_some())
@@ -108,6 +111,10 @@ impl DataWriteRepositories {
.map(PostgresBackend::oauth_provider_write_repository)
.or_else(|| mysql.map(MysqlBackend::oauth_provider_write_repository))
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_write_repository)),
pool_scores: postgres
.map(PostgresBackend::pool_score_write_repository)
.or_else(|| mysql.map(MysqlBackend::pool_score_write_repository))
.or_else(|| sqlite.map(SqliteBackend::pool_score_write_repository)),
proxy_nodes: postgres
.map(PostgresBackend::proxy_node_write_repository)
.or_else(|| mysql.map(MysqlBackend::proxy_node_write_repository))
@@ -184,6 +191,10 @@ impl DataWriteRepositories {
self.oauth_providers.clone()
}
pub fn pool_scores(&self) -> Option<Arc<dyn PoolMemberScoreWriteRepository>> {
self.pool_scores.clone()
}
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeWriteRepository>> {
self.proxy_nodes.clone()
}
@@ -218,6 +229,7 @@ impl DataWriteRepositories {
|| self.global_models.is_some()
|| self.management_tokens.is_some()
|| self.oauth_providers.is_some()
|| self.pool_scores.is_some()
|| self.proxy_nodes.is_some()
|| self.provider_catalog.is_some()
|| self.provider_quotas.is_some()