diff --git a/sgl-router/Cargo.toml b/sgl-router/Cargo.toml index a8a1b4bde..8d5f44dd0 100644 --- a/sgl-router/Cargo.toml +++ b/sgl-router/Cargo.toml @@ -95,6 +95,9 @@ deadpool = { version = "0.12", features = ["managed", "rt_tokio_1"] } backoff = { version = "0.4", features = ["tokio"] } strum = { version = "0.26", features = ["derive"] } once_cell = "1.21.3" +tokio-postgres = { version = "0.7.15", features = ["runtime","with-chrono-0_4","with-serde_json-1","array-impls"] } +deadpool-postgres = "0.14.1" + [build-dependencies] tonic-prost-build = "0.14.2" diff --git a/sgl-router/py_src/sglang_router/router.py b/sgl-router/py_src/sglang_router/router.py index a71febfab..0ee130749 100644 --- a/sgl-router/py_src/sglang_router/router.py +++ b/sgl-router/py_src/sglang_router/router.py @@ -6,6 +6,7 @@ from sglang_router.sglang_router_rs import ( HistoryBackendType, PolicyType, PyOracleConfig, + PyPostgresConfig, ) from sglang_router.sglang_router_rs import Router as _Router @@ -53,6 +54,8 @@ def history_backend_from_str(backend_str: Optional[str]) -> HistoryBackendType: return getattr(HistoryBackendType, "None") elif backend_lower == "oracle": return HistoryBackendType.Oracle + elif backend_lower == "postgres": + return HistoryBackendType.Postgres else: raise ValueError(f"Unknown history backend: {backend_str}") @@ -187,6 +190,15 @@ class Router: args_dict["oracle_config"] = oracle_config args_dict["history_backend"] = history_backend + # Convert Postgres config if needed + postgres_config = None + if history_backend == HistoryBackendType.Postgres: + postgres_config = PyPostgresConfig( + db_url=args_dict.get("postgres_db_url"), + pool_max=args_dict.get("postgres_pool_max", 16), + ) + args_dict["postgres_config"] = postgres_config + # Remove fields that shouldn't be passed to Rust Router constructor fields_to_remove = [ "mini_lb", diff --git a/sgl-router/py_src/sglang_router/router_args.py b/sgl-router/py_src/sglang_router/router_args.py index bec64c8d9..d2d44542c 100644 --- a/sgl-router/py_src/sglang_router/router_args.py +++ b/sgl-router/py_src/sglang_router/router_args.py @@ -544,7 +544,7 @@ class RouterArgs: f"--{prefix}history-backend", type=str, default=RouterArgs.history_backend, - choices=["memory", "none", "oracle"], + choices=["memory", "none", "oracle", "postgres"], help="History storage backend for conversations and responses (default: memory)", ) # Oracle configuration diff --git a/sgl-router/src/config/builder.rs b/sgl-router/src/config/builder.rs index f84c3df33..f7e15aae5 100644 --- a/sgl-router/src/config/builder.rs +++ b/sgl-router/src/config/builder.rs @@ -1,7 +1,7 @@ use super::{ CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig, - HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, RetryConfig, RouterConfig, - RoutingMode, TokenizerCacheConfig, + HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RetryConfig, + RouterConfig, RoutingMode, TokenizerCacheConfig, }; use crate::{core::ConnectionMode, mcp::McpConfig}; @@ -497,6 +497,14 @@ impl RouterConfigBuilder { self } + pub fn maybe_postgres(mut self, postgres: Option) -> Self { + if let Some(cfg) = postgres { + self.config.history_backend = HistoryBackend::Postgres; + self.config.postgres = Some(cfg); + } + self + } + pub fn maybe_reasoning_parser(mut self, parser: Option>) -> Self { self.config.reasoning_parser = parser.map(|p| p.into()); self diff --git a/sgl-router/src/config/types.rs b/sgl-router/src/config/types.rs index ae267d069..90174f070 100644 --- a/sgl-router/src/config/types.rs +++ b/sgl-router/src/config/types.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize}; +use url::Url; use super::ConfigResult; use crate::core::ConnectionMode; @@ -53,6 +54,9 @@ pub struct RouterConfig { /// Required when history_backend = "oracle" #[serde(skip_serializing_if = "Option::is_none")] pub oracle: Option, + /// Required when history_backend = "postgres" + #[serde(skip_serializing_if = "Option::is_none")] + pub postgres: Option, /// For reasoning models (e.g., deepseek-r1, qwen3) pub reasoning_parser: Option, /// For tool-call interactions @@ -123,6 +127,7 @@ pub enum HistoryBackend { Memory, None, Oracle, + Postgres, } /// Oracle history backend configuration @@ -182,6 +187,54 @@ impl std::fmt::Debug for OracleConfig { } } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct PostgresConfig { + // Database connection URL, + // postgres://[user[:password]@][netloc][:port][/dbname][?param1=value1&...] + pub db_url: String, + // Database pool max size + pub pool_max: usize, +} + +impl PostgresConfig { + pub fn default_pool_max() -> usize { + 16 + } + + pub fn validate(&self) -> Result<(), String> { + let s = self.db_url.trim(); + if s.is_empty() { + return Err("is it db-url should be not empty".to_string()); + } + + let url = Url::parse(s).map_err(|e| format!("invalid db_url: {}", e))?; + + let scheme = url.scheme(); + if scheme != "postgres" && scheme != "postgresql" { + return Err(format!("don't support URL scheme: {}", scheme)); + } + + if url.host().is_none() { + return Err("db_url must need host".to_string()); + } + + let path = url.path(); + let dbname = path + .strip_prefix('/') + .filter(|p| !p.is_empty()) + .map(|s| s.to_string()); + if dbname.is_none() { + return Err("db_url must need database name".to_string()); + } + + if self.pool_max == 0 { + return Err("pool_max must be greater 1, default is 16".to_string()); + } + + Ok(()) + } +} + /// Routing mode configuration #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type")] @@ -442,6 +495,7 @@ impl Default for RouterConfig { chat_template: None, history_backend: default_history_backend(), oracle: None, + postgres: None, reasoning_parser: None, tool_call_parser: None, tokenizer_cache: TokenizerCacheConfig::default(), diff --git a/sgl-router/src/data_connector/common.rs b/sgl-router/src/data_connector/common.rs new file mode 100644 index 000000000..79b7b5e42 --- /dev/null +++ b/sgl-router/src/data_connector/common.rs @@ -0,0 +1,78 @@ +use std::collections::HashMap; + +use serde_json::Value; + +pub fn parse_tool_calls(raw: Option) -> Result, String> { + match raw { + Some(s) if !s.is_empty() => serde_json::from_str(&s).map_err(|e| e.to_string()), + _ => Ok(Vec::new()), + } +} + +pub fn parse_metadata(raw: Option) -> Result, String> { + match raw { + Some(s) if !s.is_empty() => serde_json::from_str(&s).map_err(|e| e.to_string()), + _ => Ok(HashMap::new()), + } +} + +pub fn parse_raw_response(raw: Option) -> Result { + match raw { + Some(s) if !s.is_empty() => serde_json::from_str(&s).map_err(|e| e.to_string()), + _ => Ok(Value::Null), + } +} + +pub fn parse_json_value(raw: Option) -> Result { + match raw { + Some(s) if !s.is_empty() => serde_json::from_str(&s).map_err(|e| e.to_string()), + _ => Ok(Value::Array(vec![])), + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn parse_tool_calls_handles_empty_input() { + assert!(parse_tool_calls(None).unwrap().is_empty()); + assert!(parse_tool_calls(Some(String::new())).unwrap().is_empty()); + } + + #[test] + fn parse_tool_calls_round_trips() { + let payload = json!([{ "type": "test", "value": 1 }]).to_string(); + let parsed = parse_tool_calls(Some(payload)).unwrap(); + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0]["type"], "test"); + assert_eq!(parsed[0]["value"], 1); + } + + #[test] + fn parse_metadata_defaults_to_empty_map() { + assert!(parse_metadata(None).unwrap().is_empty()); + } + + #[test] + fn parse_metadata_round_trips() { + let payload = json!({"key": "value", "nested": {"bool": true}}).to_string(); + let parsed = parse_metadata(Some(payload)).unwrap(); + assert_eq!(parsed.get("key").unwrap(), "value"); + assert_eq!(parsed["nested"]["bool"], true); + } + + #[test] + fn parse_raw_response_handles_null() { + assert_eq!(parse_raw_response(None).unwrap(), Value::Null); + } + + #[test] + fn parse_raw_response_round_trips() { + let payload = json!({"id": "abc"}).to_string(); + let parsed = parse_raw_response(Some(payload)).unwrap(); + assert_eq!(parsed["id"], "abc"); + } +} diff --git a/sgl-router/src/data_connector/factory.rs b/sgl-router/src/data_connector/factory.rs index e351d5340..5c62934df 100644 --- a/sgl-router/src/data_connector/factory.rs +++ b/sgl-router/src/data_connector/factory.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use tracing::info; +use url::Url; use super::{ core::{ConversationItemStorage, ConversationStorage, ResponseStorage}, @@ -14,7 +15,13 @@ use super::{ noop::{NoOpConversationItemStorage, NoOpConversationStorage, NoOpResponseStorage}, oracle::{OracleConversationItemStorage, OracleConversationStorage, OracleResponseStorage}, }; -use crate::config::{HistoryBackend, OracleConfig, RouterConfig}; +use crate::{ + config::{HistoryBackend, OracleConfig, PostgresConfig, RouterConfig}, + data_connector::postgres::{ + PostgresConversationItemStorage, PostgresConversationStorage, PostgresResponseStorage, + PostgresStore, + }, +}; /// Type alias for the storage tuple returned by factory functions. /// This avoids clippy::type_complexity warnings while keeping Arc explicit. @@ -68,6 +75,33 @@ pub fn create_storage(config: &RouterConfig) -> Result { info!("Data connector initialized successfully: Oracle ATP"); Ok(storages) } + HistoryBackend::Postgres => { + let postgres_cfg = config + .postgres + .clone() + .ok_or("Postgres configuration is required when history_backend=postgres")?; + + let log_db_url = match Url::parse(&postgres_cfg.db_url) { + Ok(mut url) => { + if url.password().is_some() { + let _ = url.set_password(Some("****")); + } + url.to_string() + } + Err(_) => "".to_string(), + }; + + info!( + "Initializing data connector: Postgres (db_url: {}, pool_max: {})", + log_db_url, postgres_cfg.pool_max + ); + + let storages = create_postgres_storage(&postgres_cfg)?; + + info!("Data connector initialized successfully: Postgres"); + + Ok(storages) + } } } @@ -88,3 +122,19 @@ fn create_oracle_storage(oracle_cfg: &OracleConfig) -> Result Result { + let store = PostgresStore::new(postgres_cfg.clone())?; + let postgres_resp = PostgresResponseStorage::new(store.clone()) + .map_err(|err| format!("failed to initialize Postgres response storage: {err}"))?; + let postgres_conv = PostgresConversationStorage::new(store.clone()) + .map_err(|err| format!("failed to initialize Postgres conversation storage: {err}"))?; + let postgres_item = PostgresConversationItemStorage::new(store.clone()) + .map_err(|err| format!("failed to initialize Postgres conversation item storage: {err}"))?; + + Ok(( + Arc::new(postgres_resp), + Arc::new(postgres_conv), + Arc::new(postgres_item), + )) +} diff --git a/sgl-router/src/data_connector/mod.rs b/sgl-router/src/data_connector/mod.rs index 9fdf75fd8..dc1e2deca 100644 --- a/sgl-router/src/data_connector/mod.rs +++ b/sgl-router/src/data_connector/mod.rs @@ -7,11 +7,13 @@ // - oracle.rs: All Oracle ATP storage implementations // - factory.rs: Storage creation function +mod common; mod core; mod factory; mod memory; mod noop; mod oracle; +mod postgres; pub use core::{ Conversation, ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage, diff --git a/sgl-router/src/data_connector/oracle.rs b/sgl-router/src/data_connector/oracle.rs index fe42cfefb..2ab4cb77c 100644 --- a/sgl-router/src/data_connector/oracle.rs +++ b/sgl-router/src/data_connector/oracle.rs @@ -6,7 +6,7 @@ //! 3. OracleConversationItemStorage //! 4. OracleResponseStorage -use std::{collections::HashMap, path::Path, sync::Arc, time::Duration}; +use std::{path::Path, sync::Arc, time::Duration}; use async_trait::async_trait; use chrono::{DateTime, Utc}; @@ -24,8 +24,12 @@ use super::core::{ NewConversationItem, ResponseChain, ResponseId, ResponseStorage, ResponseStorageError, SortOrder, StoredResponse, }; -use crate::config::OracleConfig; - +use crate::{ + config::OracleConfig, + data_connector::common::{ + parse_json_value, parse_metadata, parse_raw_response, parse_tool_calls, + }, +}; // ============================================================================ // PART 1: OracleStore Helper + Common Utilities // ============================================================================ @@ -1124,78 +1128,3 @@ fn create_index_if_missing(conn: &Connection, index_name: &str, ddl: &str) -> Re Ok(()) } - -fn parse_tool_calls(raw: Option) -> Result, String> { - match raw { - Some(s) if !s.is_empty() => serde_json::from_str(&s).map_err(|e| e.to_string()), - _ => Ok(Vec::new()), - } -} - -fn parse_metadata(raw: Option) -> Result, String> { - match raw { - Some(s) if !s.is_empty() => serde_json::from_str(&s).map_err(|e| e.to_string()), - _ => Ok(HashMap::new()), - } -} - -fn parse_raw_response(raw: Option) -> Result { - match raw { - Some(s) if !s.is_empty() => serde_json::from_str(&s).map_err(|e| e.to_string()), - _ => Ok(Value::Null), - } -} - -fn parse_json_value(raw: Option) -> Result { - match raw { - Some(s) if !s.is_empty() => serde_json::from_str(&s).map_err(|e| e.to_string()), - _ => Ok(Value::Array(vec![])), - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - - #[test] - fn parse_tool_calls_handles_empty_input() { - assert!(parse_tool_calls(None).unwrap().is_empty()); - assert!(parse_tool_calls(Some(String::new())).unwrap().is_empty()); - } - - #[test] - fn parse_tool_calls_round_trips() { - let payload = json!([{ "type": "test", "value": 1 }]).to_string(); - let parsed = parse_tool_calls(Some(payload)).unwrap(); - assert_eq!(parsed.len(), 1); - assert_eq!(parsed[0]["type"], "test"); - assert_eq!(parsed[0]["value"], 1); - } - - #[test] - fn parse_metadata_defaults_to_empty_map() { - assert!(parse_metadata(None).unwrap().is_empty()); - } - - #[test] - fn parse_metadata_round_trips() { - let payload = json!({"key": "value", "nested": {"bool": true}}).to_string(); - let parsed = parse_metadata(Some(payload)).unwrap(); - assert_eq!(parsed.get("key").unwrap(), "value"); - assert_eq!(parsed["nested"]["bool"], true); - } - - #[test] - fn parse_raw_response_handles_null() { - assert_eq!(parse_raw_response(None).unwrap(), Value::Null); - } - - #[test] - fn parse_raw_response_round_trips() { - let payload = json!({"id": "abc"}).to_string(); - let parsed = parse_raw_response(Some(payload)).unwrap(); - assert_eq!(parsed["id"], "abc"); - } -} diff --git a/sgl-router/src/data_connector/postgres.rs b/sgl-router/src/data_connector/postgres.rs new file mode 100644 index 000000000..96aabde1d --- /dev/null +++ b/sgl-router/src/data_connector/postgres.rs @@ -0,0 +1,699 @@ +//! Postgres storage implementation using PostgresStore helper +//! +//! Structure: +//! 1. PostgresStore helper and common utilities +//! 2. PostgresConversationStorage +//! 3. PostgresConversationItemStorage +//! 4. PostgresResponseStorage + +use std::str::FromStr; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod}; +use serde_json::Value; +use tokio_postgres::{NoTls, Row}; + +use crate::{ + config::PostgresConfig, + data_connector::{ + common::{parse_json_value, parse_metadata, parse_raw_response, parse_tool_calls}, + core::{ + make_item_id, ConversationItemResult, ConversationItemStorageError, + ConversationMetadata, ConversationResult, ConversationStorageError, ResponseChain, + ResponseResult, ResponseStorageError, + }, + Conversation, ConversationId, ConversationItem, ConversationItemId, + ConversationItemStorage, ConversationStorage, ListParams, NewConversation, + NewConversationItem, ResponseId, ResponseStorage, SortOrder, StoredResponse, + }, +}; + +pub(crate) struct PostgresStore { + pool: Pool, +} + +impl PostgresStore { + pub fn new(config: PostgresConfig) -> Result { + let pg_config = tokio_postgres::Config::from_str(config.db_url.as_str()).unwrap(); + let mgr_config = ManagerConfig { + recycling_method: RecyclingMethod::Fast, + }; + let mgr = Manager::from_config(pg_config, NoTls, mgr_config); + let pool = Pool::builder(mgr) + .max_size(config.pool_max) + .build() + .unwrap(); + + Ok(Self { pool }) + } +} + +impl Clone for PostgresStore { + fn clone(&self) -> Self { + Self { + pool: self.pool.clone(), + } + } +} + +pub struct PostgresConversationStorage { + store: PostgresStore, +} + +impl PostgresConversationStorage { + pub fn new(store: PostgresStore) -> Result { + futures::executor::block_on(Self::initialize_schema(store.clone())) + .expect("Failed to initialize conversations schema"); + Ok(Self { store }) + } + + async fn initialize_schema(store: PostgresStore) -> Result<(), ConversationStorageError> { + let client = store.pool.get().await.unwrap(); + client + .batch_execute( + " + CREATE TABLE IF NOT EXISTS conversations ( + id VARCHAR(64) PRIMARY KEY, + created_at TIMESTAMPTZ, + metadata JSON + );", + ) + .await + .unwrap(); + + Ok(()) + } + + fn parse_metadata( + metadata: Option, + ) -> Result, ConversationStorageError> { + match metadata { + None => Ok(None), + Some(s) => { + let s = s.trim(); + if s.is_empty() || s.eq_ignore_ascii_case("null") { + return Ok(None); + } + serde_json::from_str::(s) + .map(Some) + .map_err(|e| ConversationStorageError::StorageError(e.to_string())) + } + } + } +} + +#[async_trait] +impl ConversationStorage for PostgresConversationStorage { + async fn create_conversation( + &self, + input: NewConversation, + ) -> Result { + let conversation = Conversation::new(input); + let id_str = conversation.id.0.as_str(); + let created_at: DateTime = conversation.created_at; + let metadata_json = conversation + .metadata + .as_ref() + .map(serde_json::to_string) + .transpose()?; + let client = self.store.pool.get().await.unwrap(); + client + .execute( + "INSERT INTO conversations (id, created_at, metadata) VALUES ($1, $2, $3)", + &[&id_str, &created_at, &metadata_json], + ) + .await + .unwrap(); + Ok(conversation) + } + + async fn get_conversation( + &self, + id: &ConversationId, + ) -> Result, ConversationStorageError> { + let conversation_id = id.0.clone(); + let client = self.store.pool.get().await.unwrap(); + let rows = client + .query( + "SELECT id, created_at, metadata FROM conversations WHERE id = $1", + &[&conversation_id], + ) + .await + .unwrap(); + if rows.is_empty() { + return Ok(None); + } + let row = &rows[0]; + let id_str: String = row.get(0); + let created_at: DateTime = row.get(1); + let metadata_json: Option = row.get(2); + let metadata = Self::parse_metadata(metadata_json)?; + Ok(Some(Conversation::with_parts( + ConversationId(id_str), + created_at, + metadata, + ))) + } + + async fn update_conversation( + &self, + id: &ConversationId, + metadata: Option, + ) -> Result, ConversationStorageError> { + let conversation_id = id.0.clone(); + let metadata_json = metadata.as_ref().map(serde_json::to_string).transpose()?; + let client = self.store.pool.get().await.unwrap(); + let rows = client + .query( + "UPDATE conversations SET metadata = $1 WHERE id = $2 RETURNING created_at", + &[&metadata_json, &conversation_id], + ) + .await + .unwrap(); + if rows.is_empty() { + return Ok(None); + } + let row = &rows[0]; + let created_at: DateTime = row.get(0); + Ok(Some(Conversation::with_parts( + ConversationId(conversation_id), + created_at, + metadata, + ))) + } + + async fn delete_conversation(&self, id: &ConversationId) -> ConversationResult { + let conversation_id = id.0.clone(); + let client = self.store.pool.get().await.unwrap(); + let rows_deleted = client + .execute( + "DELETE FROM conversations WHERE id = $1", + &[&conversation_id], + ) + .await + .unwrap(); + Ok(rows_deleted > 0) + } +} + +pub struct PostgresConversationItemStorage { + store: PostgresStore, +} + +impl PostgresConversationItemStorage { + pub fn new(store: PostgresStore) -> Result { + futures::executor::block_on(Self::initialize_schema(store.clone())) + .expect("Failed to initialize conversation_items or conversation_item_links schema"); + Ok(Self { store }) + } + + async fn initialize_schema(store: PostgresStore) -> Result<(), ConversationItemStorageError> { + let client = store.pool.get().await.unwrap(); + client + .batch_execute( + " + CREATE TABLE IF NOT EXISTS conversation_items ( + id SERIAL PRIMARY KEY, + response_id VARCHAR(64), + item_type VARCHAR(32) NOT NULL, + role VARCHAR(32), + content JSON, + status VARCHAR(32), + created_at TIMESTAMPTZ + );", + ) + .await + .unwrap(); + + // Create conversation_item_links table + client + .batch_execute( + " + CREATE TABLE IF NOT EXISTS conversation_item_links ( + conversation_id VARCHAR(64), + item_id VARCHAR(64) NOT NULL, + added_at TIMESTAMPTZ, + CONSTRAINT pk_conv_item_link PRIMARY KEY (conversation_id, item_id) + ); + CREATE INDEX IF NOT EXISTS conv_item_links_conv_idx ON conversation_item_links (conversation_id, added_at);", + ) + .await + .unwrap(); + Ok(()) + } +} + +#[async_trait] +impl ConversationItemStorage for PostgresConversationItemStorage { + async fn create_item( + &self, + item: NewConversationItem, + ) -> Result { + let id = item + .id + .clone() + .unwrap_or_else(|| make_item_id(&item.item_type)); + let created_at = Utc::now(); + let content_json = serde_json::to_string(&item.content)?; + let conversation_item = ConversationItem { + id: id.clone(), + response_id: item.response_id.clone(), + item_type: item.item_type.clone(), + role: item.role.clone(), + content: item.content.clone(), + status: item.status.clone(), + created_at, + }; + let id_str = conversation_item.id.0.clone(); + let response_id = conversation_item.response_id.clone(); + let item_type = conversation_item.item_type.clone(); + let role = conversation_item.role.clone(); + let status = conversation_item.status.clone(); + + let client = self.store.pool.get().await.unwrap(); + client.execute("INSERT INTO conversation_items (id, response_id, item_type, role, content, status, created_at) VALUES ($1, $2, $3, $4, $5, $6, $7)", + &[&id_str, &response_id, &item_type, &role, &content_json, &status, &created_at]) + .await + .unwrap(); + Ok(conversation_item) + } + + async fn link_item( + &self, + conversation_id: &ConversationId, + item_id: &ConversationItemId, + added_at: DateTime, + ) -> ConversationItemResult<()> { + let cid = conversation_id.0.clone(); + let iid = item_id.0.clone(); + let client = self.store.pool.get().await.unwrap(); + client.execute("INSERT INTO conversation_item_links (conversation_id, item_id, added_at) VALUES ($1, $2, $3)", + &[&cid, &iid, &added_at]) + .await + .unwrap(); + Ok(()) + } + + async fn list_items( + &self, + conversation_id: &ConversationId, + params: ListParams, + ) -> ConversationItemResult> { + let cid = conversation_id.0.clone(); + let limit: i64 = params.limit as i64; + let order_desc = matches!(params.order, SortOrder::Desc); + let after_id = params.after.clone(); + + let after_key: Option<(DateTime, String)> = if let Some(ref aid) = after_id { + let cid = cid.clone(); + let aid = aid.clone(); + let client = self.store.pool.get().await.unwrap(); + let rows = client + .query( + "SELECT added_at FROM conversation_item_links WHERE conversation_id = $1 AND item_id = $2", + &[&cid, &aid], + ) + .await + .unwrap(); + if !rows.is_empty() { + let row = &rows[0]; + let ts: DateTime = row.get(0); + Some((ts, aid)) + } else { + None + } + } else { + None + }; + + let cid = cid.clone(); + let mut sql = String::from( + "SELECT i.id, i.response_id, i.item_type, i.role, i.content, i.status, i.created_at \ + FROM conversation_item_links l \ + JOIN conversation_items i ON i.id = l.item_id \ + WHERE l.conversation_id = $1", + ); + // If cursor provided, append predicate using $2/$3 + if let Some((_ts, _iid)) = &after_key { + if order_desc { + sql.push_str(" AND (l.added_at < $2 OR (l.added_at = $2 AND l.item_id < $3))"); + } else { + sql.push_str(" AND (l.added_at > $2 OR (l.added_at = $2 AND l.item_id > $3))"); + } + } + // Order and limit + if order_desc { + sql.push_str(" ORDER BY l.added_at DESC, l.item_id DESC"); + } else { + sql.push_str(" ORDER BY l.added_at ASC, l.item_id ASC"); + } + // PostgreSQL LIMIT + if after_key.is_some() { + sql.push_str(" LIMIT $4"); + } else { + sql.push_str(" LIMIT $2"); + } + let client = self.store.pool.get().await.unwrap(); + let rows = if let Some((ts, iid)) = &after_key { + client.query(&sql, &[&cid, ts, iid, &limit]).await.unwrap() + } else { + client.query(&sql, &[&cid, &limit]).await.unwrap() + }; + let mut out = Vec::new(); + for row in rows { + let id = row.get(0); + let resp_id: Option = row.get(1); + let item_type: String = row.get(2); + let role: Option = row.get(3); + let content_raw: Option = row.get(4); + let status: Option = row.get(5); + let created_at: DateTime = row.get(6); + out.push(( + id, + resp_id, + item_type, + role, + content_raw, + status, + created_at, + )); + } + out.into_iter() + .map( + |(id, resp_id, item_type, role, content_raw, status, created_at)| { + let content = match content_raw { + Some(s) => { + serde_json::from_str(&s).map_err(ConversationItemStorageError::from)? + } + None => Value::Null, + }; + Ok(ConversationItem { + id: ConversationItemId(id), + response_id: resp_id, + item_type, + role, + content, + status, + created_at, + }) + }, + ) + .collect() + } + + async fn get_item( + &self, + item_id: &ConversationItemId, + ) -> Result, ConversationItemStorageError> { + let iid = item_id.0.clone(); + let client = self.store.pool.get().await.unwrap(); + let row = client.query_one("SELECT id, response_id, item_type, role, content, status, created_at FROM converstation_items WHERE id = $1", &[&iid]).await.unwrap(); + if row.is_empty() { + Ok(None) + } else { + let id: String = row.get(0); + let response_id: Option = row.get(1); + let item_type: String = row.get(2); + let role: Option = row.get(3); + let content_raw: Option = row.get(4); + let status: Option = row.get(5); + let created_at: DateTime = row.get(6); + + let content = match content_raw { + Some(s) => serde_json::from_str(&s) + .map_err(ConversationItemStorageError::SerializationError)?, + None => Value::Null, + }; + + Ok(Some(ConversationItem { + id: ConversationItemId(id), + response_id, + item_type, + role, + content, + status, + created_at, + })) + } + } + + async fn is_item_linked( + &self, + conversation_id: &ConversationId, + item_id: &ConversationItemId, + ) -> ConversationItemResult { + let cid = conversation_id.0.clone(); + let iid = item_id.0.clone(); + let client = self.store.pool.get().await.unwrap(); + let row = client + .query_one( + "SELECT COUNT(*) FROM conversation_item_links WHERE conversation_id = $1 AND item_id = $2", + &[&cid, &iid], + ) + .await + .unwrap(); + let count: i64 = row.get(0); + Ok(count > 0) + } + + async fn delete_item( + &self, + conversation_id: &ConversationId, + item_id: &ConversationItemId, + ) -> ConversationItemResult<()> { + let cid = conversation_id.0.clone(); + let iid = item_id.0.clone(); + let client = self.store.pool.get().await.unwrap(); + client + .execute( + "DELETE FROM conversation_item_links WHERE conversation_id = $1 AND item_id = $2", + &[&cid, &iid], + ) + .await + .unwrap(); + Ok(()) + } +} + +pub struct PostgresResponseStorage { + store: PostgresStore, +} + +impl PostgresResponseStorage { + pub fn new(store: PostgresStore) -> Result { + futures::executor::block_on(Self::initialize_schema(store.clone())) + .expect("Failed to initialize responses schema"); + Ok(Self { store }) + } + + async fn initialize_schema(store: PostgresStore) -> Result<(), ResponseStorageError> { + let client = store.pool.get().await.unwrap(); + client + .batch_execute( + " + CREATE TABLE IF NOT EXISTS responses ( + id VARCHAR(64) PRIMARY KEY, + conversation_id VARCHAR(64), + previous_response_id VARCHAR(64), + input JSON, + instructions TEXT, + output JSON, + tool_calls JSON, + metadata JSON, + created_at TIMESTAMPTZ, + safety_identifier VARCHAR(128), + model VARCHAR(128), + raw_response JSON + );", + ) + .await + .unwrap(); + Ok(()) + } + + pub fn build_response_from_now(row: &Row) -> Result { + let id: String = row.get("id"); + let conversation_id: Option = row.get("conversation_id"); + let previous: Option = row.get("previous_response_id"); + let input_json: Option = row.get("input"); + let instructions: Option = row.get("instructions"); + let output_json: Option = row.get("output"); + let tool_calls_json: Option = row.get("tool_calls"); + let metadata_json: Option = row.get("metadata"); + let created_at: DateTime = row.get("created_at"); + let safety_identifier: Option = row.get("safety_identifier"); + let model: Option = row.get("model"); + let raw_response_json: Option = row.get("raw_response"); + + let previous_response_id = previous.map(ResponseId); + let tool_calls = parse_tool_calls(tool_calls_json)?; + let metadata = parse_metadata(metadata_json)?; + let raw_response = parse_raw_response(raw_response_json)?; + let input = parse_json_value(input_json)?; + let output = parse_json_value(output_json)?; + + Ok(StoredResponse { + id: ResponseId(id), + previous_response_id, + input, + instructions, + output, + tool_calls, + metadata, + created_at, + safety_identifier, + model, + conversation_id, + raw_response, + }) + } +} + +#[async_trait] +impl ResponseStorage for PostgresResponseStorage { + async fn store_response( + &self, + response: StoredResponse, + ) -> Result { + let response_id = response.id.clone(); + let response_id_str = response_id.0.clone(); + let previous_id = response.previous_response_id.map(|r| r.0); + let json_input = &response.input; + let json_output = &response.output; + let json_tool_calls = serde_json::to_string(&response.tool_calls)?; + let json_metadata = serde_json::to_string(&response.metadata)?; + let json_raw_response = &response.raw_response; + let instructions = response.instructions.clone(); + let created_at = response.created_at; + let safety_identifier = response.safety_identifier.clone(); + let model = response.model.clone(); + let conversation_id = response.conversation_id.clone(); + let client = self.store.pool.get().await.unwrap(); + let insert_count = client.execute( + "INSERT INTO responses (id, previous_response_id, input, instructions, output, \ + tool_calls, metadata, created_at, safety_identifier, model, conversation_id, raw_response) \ + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)", + &[ + &response_id_str, + &previous_id, + &json_input, + &instructions, + &json_output, + &serde_json::json!(&json_tool_calls), + &serde_json::json!(&json_metadata), + &created_at, + &safety_identifier, + &model, + &conversation_id, + &json_raw_response, + ]).await.unwrap(); + println!("INSERT INTO responses VALUES {insert_count}"); + Ok(response_id) + } + + async fn get_response( + &self, + response_id: &ResponseId, + ) -> Result, ResponseStorageError> { + let id = response_id.0.clone(); + let client = self.store.pool.get().await.unwrap(); + let row = client + .query_one("SELECT * FROM responses WHERE id = $1", &[&id]) + .await + .unwrap(); + Self::build_response_from_now(&row) + .map(Some) + .map_err(|err| ResponseStorageError::StorageError(err.to_string())) + } + + async fn delete_response(&self, response_id: &ResponseId) -> ResponseResult<()> { + let id = response_id.0.clone(); + let client = self.store.pool.get().await.unwrap(); + client + .execute("DELETE FROM responses WHERE id = $1", &[&id]) + .await + .unwrap(); + Ok(()) + } + + async fn get_response_chain( + &self, + response_id: &ResponseId, + max_depth: Option, + ) -> ResponseResult { + let mut chain = ResponseChain::new(); + let mut current_id = Some(response_id.clone()); + let mut visited = 0usize; + + while let Some(ref lookup_id) = current_id { + if let Some(limit) = max_depth { + if visited >= limit { + break; + } + } + + let fetched = self.get_response(lookup_id).await?; + match fetched { + Some(response) => { + current_id = response.previous_response_id.clone(); + chain.responses.push(response); + visited += 1; + } + None => break, + } + } + + chain.responses.reverse(); + Ok(chain) + } + + async fn list_identifier_responses( + &self, + identifier: &str, + limit: Option, + ) -> ResponseResult> { + let identifier = identifier.to_string(); + let client = self.store.pool.get().await.unwrap(); + let rows = if let Some(l) = limit { + let l_i64: i64 = l as i64; + client + .query( + "SELECT * FROM responses WHERE safety_identifier = $1 ORDER BY created_at DESC LIMIT $2", + &[&identifier, &l_i64], + ) + .await + .unwrap() + } else { + client + .query( + "SELECT * FROM responses WHERE safety_identifier = $1 ORDER BY created_at DESC", + &[&identifier], + ) + .await + .unwrap() + }; + + let mut out = Vec::with_capacity(rows.len()); + for row in rows { + let resp = + Self::build_response_from_now(&row).map_err(ResponseStorageError::StorageError)?; + out.push(resp); + } + + Ok(out) + } + + async fn delete_identifier_responses(&self, identifier: &str) -> ResponseResult { + let identifier = identifier.to_string(); + let client = self.store.pool.get().await.unwrap(); + let rows_deleted = client + .execute( + "DELETE FROM responses WHERE safety_identifier = $1", + &[&identifier], + ) + .await + .unwrap(); + Ok(rows_deleted as usize) + } +} diff --git a/sgl-router/src/lib.rs b/sgl-router/src/lib.rs index b18e76749..2af3561e3 100644 --- a/sgl-router/src/lib.rs +++ b/sgl-router/src/lib.rs @@ -45,6 +45,7 @@ pub enum HistoryBackendType { Memory, None, Oracle, + Postgres, } #[pyclass] @@ -139,6 +140,34 @@ impl PyOracleConfig { } } +#[pyclass] +#[derive(Debug, Clone, PartialEq)] +pub struct PyPostgresConfig { + #[pyo3(get, set)] + pub db_url: Option, + + #[pyo3(get, set)] + pub pool_max: usize, +} + +#[pymethods] +impl PyPostgresConfig { + #[new] + #[pyo3(signature = (db_url = None,pool_max = 16,))] + fn new(db_url: Option, pool_max: usize) -> PyResult { + Ok(PyPostgresConfig { db_url, pool_max }) + } +} + +impl PyPostgresConfig { + fn to_config_postgres(&self) -> config::PostgresConfig { + config::PostgresConfig { + db_url: self.db_url.clone().unwrap_or_default(), + pool_max: self.pool_max, + } + } +} + #[pyclass] #[derive(Debug, Clone, PartialEq)] struct Router { @@ -212,6 +241,7 @@ struct Router { backend: BackendType, history_backend: HistoryBackendType, oracle_config: Option, + postgres_config: Option, client_cert_path: Option, client_key_path: Option, ca_cert_paths: Vec, @@ -305,6 +335,7 @@ impl Router { HistoryBackendType::Memory => config::HistoryBackend::Memory, HistoryBackendType::None => config::HistoryBackend::None, HistoryBackendType::Oracle => config::HistoryBackend::Oracle, + HistoryBackendType::Postgres => config::HistoryBackend::Postgres, }; let oracle = if matches!(self.history_backend, HistoryBackendType::Oracle) { @@ -315,6 +346,14 @@ impl Router { None }; + let postgres_config = if matches!(self.history_backend, HistoryBackendType::Postgres) { + self.postgres_config + .as_ref() + .map(|cfg| cfg.to_config_postgres()) + } else { + None + }; + config::RouterConfig::builder() .mode(mode) .policy(policy) @@ -367,6 +406,7 @@ impl Router { .maybe_tokenizer_path(self.tokenizer_path.as_ref()) .maybe_chat_template(self.chat_template.as_ref()) .maybe_oracle(oracle) + .maybe_postgres(postgres_config) .maybe_reasoning_parser(self.reasoning_parser.as_ref()) .maybe_tool_call_parser(self.tool_call_parser.as_ref()) .maybe_mcp_config_path(self.mcp_config_path.as_ref()) @@ -455,6 +495,7 @@ impl Router { backend = BackendType::Sglang, history_backend = HistoryBackendType::Memory, oracle_config = None, + postgres_config = None, client_cert_path = None, client_key_path = None, ca_cert_paths = vec![], @@ -529,6 +570,7 @@ impl Router { backend: BackendType, history_backend: HistoryBackendType, oracle_config: Option, + postgres_config: Option, client_cert_path: Option, client_key_path: Option, ca_cert_paths: Vec, @@ -617,6 +659,7 @@ impl Router { backend, history_backend, oracle_config, + postgres_config, client_cert_path, client_key_path, ca_cert_paths, @@ -687,6 +730,7 @@ fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/sgl-router/src/main.rs b/sgl-router/src/main.rs index 12987cb41..e87b57139 100644 --- a/sgl-router/src/main.rs +++ b/sgl-router/src/main.rs @@ -4,15 +4,14 @@ use clap::{ArgAction, Parser, ValueEnum}; use sglang_router_rs::{ config::{ CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig, - HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, RetryConfig, RouterConfig, - RoutingMode, TokenizerCacheConfig, + HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RetryConfig, + RouterConfig, RoutingMode, TokenizerCacheConfig, }, core::ConnectionMode, metrics::PrometheusConfig, server::{self, ServerConfig}, service_discovery::ServiceDiscoveryConfig, }; - fn parse_prefill_args() -> Vec<(String, Option)> { let args: Vec = std::env::args().collect(); let mut prefill_entries = Vec::new(); @@ -283,7 +282,7 @@ struct CliArgs { #[arg(long, default_value_t = 52428800)] tokenizer_cache_l1_max_memory: usize, - #[arg(long, default_value = "memory", value_parser = ["memory", "none", "oracle"])] + #[arg(long, default_value = "memory", value_parser = ["memory", "none", "oracle","postgres"])] history_backend: String, #[arg(long, env = "ATP_WALLET_PATH")] @@ -310,6 +309,12 @@ struct CliArgs { #[arg(long, env = "ATP_POOL_TIMEOUT_SECS")] oracle_pool_timeout_secs: Option, + #[arg(long)] + postgres_db_url: Option, + + #[arg(long)] + postgres_pool_max_size: Option, + #[arg(long)] reasoning_parser: Option, @@ -446,6 +451,18 @@ impl CliArgs { }) } + fn build_postgres_config(&self) -> ConfigResult { + let db_url = self.postgres_db_url.clone().unwrap_or_default(); + let pool_max = self + .postgres_pool_max_size + .unwrap_or_else(PostgresConfig::default_pool_max); + let pcf = PostgresConfig { db_url, pool_max }; + pcf.validate().map_err(|e| ConfigError::ValidationFailed { + reason: e.to_string(), + })?; + Ok(pcf) + } + fn to_router_config( &self, prefill_urls: Vec<(String, Option)>, @@ -532,6 +549,7 @@ impl CliArgs { let history_backend = match self.history_backend.as_str() { "none" => HistoryBackend::None, "oracle" => HistoryBackend::Oracle, + "postgres" => HistoryBackend::Postgres, _ => HistoryBackend::Memory, }; @@ -540,6 +558,11 @@ impl CliArgs { } else { None }; + let postgres = if history_backend == HistoryBackend::Postgres { + Some(self.build_postgres_config()?) + } else { + None + }; let builder = RouterConfig::builder() .mode(mode) @@ -595,6 +618,7 @@ impl CliArgs { .maybe_tokenizer_path(self.tokenizer_path.as_ref()) .maybe_chat_template(self.chat_template.as_ref()) .maybe_oracle(oracle) + .maybe_postgres(postgres) .maybe_reasoning_parser(self.reasoning_parser.as_ref()) .maybe_tool_call_parser(self.tool_call_parser.as_ref()) .maybe_mcp_config_path(self.mcp_config_path.as_ref())