diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 94375b188..83ba00a92 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -86,6 +86,7 @@ tool-parser = "1.0.0" llm-tokenizer = "1.0.0" smg-auth = "1.0.0" wfaas = "1.0.0" +data-connector = "1.0.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } rustls-pemfile = "2.2" openssl = "0.10.73" diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs index 1e1c930dc..ba3bf3f6b 100644 --- a/sgl-model-gateway/src/app_context.rs +++ b/sgl-model-gateway/src/app_context.rs @@ -11,6 +11,7 @@ use crate::{ core::{steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService}, data_connector::{ create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage, + StorageFactoryConfig, }, mcp::McpManager, middleware::TokenBucket, @@ -429,8 +430,14 @@ impl AppContextBuilder { /// Create all storage backends using the factory function fn with_storage(mut self, config: &RouterConfig) -> Result { + let storage_config = StorageFactoryConfig { + backend: &config.history_backend, + oracle: config.oracle.as_ref(), + postgres: config.postgres.as_ref(), + redis: config.redis.as_ref(), + }; let (response_storage, conversation_storage, conversation_item_storage) = - create_storage(config)?; + create_storage(storage_config)?; self.response_storage = Some(response_storage); self.conversation_storage = Some(conversation_storage); diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index 95f09bbd9..d1b521875 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -1,10 +1,11 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize}; -use url::Url; use super::ConfigResult; use crate::core::ConnectionMode; +// Re-export storage config types from data_connector +pub use crate::data_connector::{HistoryBackend, OracleConfig, PostgresConfig, RedisConfig}; /// Main router configuration #[derive(Debug, Clone, Serialize, Deserialize)] @@ -145,169 +146,6 @@ fn default_history_backend() -> HistoryBackend { HistoryBackend::Memory } -/// History backend configuration -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -#[serde(rename_all = "lowercase")] -pub enum HistoryBackend { - Memory, - None, - Oracle, - Postgres, - Redis, -} - -/// Oracle history backend configuration -#[derive(Clone, Serialize, Deserialize, PartialEq)] -pub struct OracleConfig { - /// ATP wallet or TLS config files directory - #[serde(skip_serializing_if = "Option::is_none")] - pub wallet_path: Option, - /// DSN (e.g. `tcps://host:port/service`) - pub connect_descriptor: String, - pub username: String, - pub password: String, - #[serde(default = "default_pool_min")] - pub pool_min: usize, - #[serde(default = "default_pool_max")] - pub pool_max: usize, - #[serde(default = "default_pool_timeout_secs")] - pub pool_timeout_secs: u64, -} - -impl OracleConfig { - pub fn default_pool_min() -> usize { - default_pool_min() - } - - pub fn default_pool_max() -> usize { - default_pool_max() - } - - pub fn default_pool_timeout_secs() -> u64 { - default_pool_timeout_secs() - } -} - -fn default_pool_min() -> usize { - 1 -} - -fn default_pool_max() -> usize { - 16 -} - -fn default_pool_timeout_secs() -> u64 { - 30 -} - -impl std::fmt::Debug for OracleConfig { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("OracleConfig") - .field("wallet_path", &self.wallet_path) - .field("connect_descriptor", &self.connect_descriptor) - .field("username", &self.username) - .field("pool_min", &self.pool_min) - .field("pool_max", &self.pool_max) - .field("pool_timeout_secs", &self.pool_timeout_secs) - .finish() - } -} - -#[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(()) - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct RedisConfig { - // Redis connection URL - // redis://[:password@]host[:port][/db] - pub url: String, - // Connection pool max size - #[serde(default = "default_redis_pool_max")] - pub pool_max: usize, - // Data retention in days. If None, data persists indefinitely. - #[serde(default = "default_redis_retention_days")] - pub retention_days: Option, -} - -fn default_redis_pool_max() -> usize { - 16 -} - -fn default_redis_retention_days() -> Option { - Some(30) -} - -impl RedisConfig { - pub fn validate(&self) -> Result<(), String> { - let s = self.url.trim(); - if s.is_empty() { - return Err("redis url should not be empty".to_string()); - } - - let url = Url::parse(s).map_err(|e| format!("invalid redis url: {}", e))?; - - let scheme = url.scheme(); - if scheme != "redis" && scheme != "rediss" { - return Err(format!("unsupported URL scheme: {}", scheme)); - } - - if url.host().is_none() { - return Err("redis url must have a host".to_string()); - } - - if self.pool_max == 0 { - return Err("pool_max must be greater than 0".to_string()); - } - - Ok(()) - } -} - /// Routing mode configuration #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type")] diff --git a/sgl-model-gateway/src/data_connector/common.rs b/sgl-model-gateway/src/data_connector/common.rs deleted file mode 100644 index e2fb6c8f9..000000000 --- a/sgl-model-gateway/src/data_connector/common.rs +++ /dev/null @@ -1,78 +0,0 @@ -use std::collections::HashMap; - -use serde_json::Value; - -pub(super) 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(super) 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(super) 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(super) 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-model-gateway/src/data_connector/core.rs b/sgl-model-gateway/src/data_connector/core.rs deleted file mode 100644 index e61b4163b..000000000 --- a/sgl-model-gateway/src/data_connector/core.rs +++ /dev/null @@ -1,483 +0,0 @@ -// core.rs -// -// Core types for the data connector module. -// Contains all traits, data types, error types, and IDs for all storage backends. -// -// Structure: -// 1. Conversation types + trait -// 2. ConversationItem types + trait -// 3. Response types + trait - -use std::{ - collections::HashMap, - fmt::{Display, Formatter}, -}; - -use async_trait::async_trait; -use chrono::{DateTime, Utc}; -use rand::RngCore; -use serde::{Deserialize, Serialize}; -use serde_json::{Map as JsonMap, Value}; - -// ============================================================================ -// PART 1: Conversation Storage -// ============================================================================ - -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)] -pub struct ConversationId(pub String); - -impl ConversationId { - pub fn new() -> Self { - let mut rng = rand::rng(); - let mut bytes = [0u8; 25]; - rng.fill_bytes(&mut bytes); - let hex_string: String = bytes.iter().map(|b| format!("{:02x}", b)).collect(); - Self(format!("conv_{}", hex_string)) - } -} - -impl Default for ConversationId { - fn default() -> Self { - Self::new() - } -} - -impl From for ConversationId { - fn from(value: String) -> Self { - Self(value) - } -} - -impl From<&str> for ConversationId { - fn from(value: &str) -> Self { - Self(value.to_string()) - } -} - -impl Display for ConversationId { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.write_str(&self.0) - } -} - -/// Metadata payload persisted with a conversation -pub type ConversationMetadata = JsonMap; - -/// Input payload for creating a conversation -#[derive(Debug, Clone, Serialize, Deserialize, Default)] -pub struct NewConversation { - /// Optional conversation ID (if None, a random ID will be generated) - #[serde(default, skip_serializing_if = "Option::is_none")] - pub id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -/// Stored conversation data structure -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct Conversation { - pub id: ConversationId, - pub created_at: DateTime, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -impl Conversation { - pub fn new(new_conversation: NewConversation) -> Self { - Self { - id: new_conversation.id.unwrap_or_default(), - created_at: Utc::now(), - metadata: new_conversation.metadata, - } - } - - pub fn with_parts( - id: ConversationId, - created_at: DateTime, - metadata: Option, - ) -> Self { - Self { - id, - created_at, - metadata, - } - } -} - -/// Result alias for conversation storage operations -pub type ConversationResult = Result; - -/// Error type for conversation storage operations -#[derive(Debug, thiserror::Error)] -pub enum ConversationStorageError { - #[error("Conversation not found: {0}")] - ConversationNotFound(String), - - #[error("Storage error: {0}")] - StorageError(String), - - #[error("Serialization error: {0}")] - SerializationError(#[from] serde_json::Error), -} - -/// Trait describing the CRUD interface for conversation storage backends -#[async_trait] -pub trait ConversationStorage: Send + Sync + 'static { - async fn create_conversation(&self, input: NewConversation) - -> ConversationResult; - - async fn get_conversation( - &self, - id: &ConversationId, - ) -> ConversationResult>; - - async fn update_conversation( - &self, - id: &ConversationId, - metadata: Option, - ) -> ConversationResult>; - - async fn delete_conversation(&self, id: &ConversationId) -> ConversationResult; -} - -// ============================================================================ -// PART 2: ConversationItem Storage -// ============================================================================ - -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)] -pub struct ConversationItemId(pub String); - -impl Display for ConversationItemId { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.write_str(&self.0) - } -} - -impl From for ConversationItemId { - fn from(value: String) -> Self { - Self(value) - } -} - -impl From<&str> for ConversationItemId { - fn from(value: &str) -> Self { - Self(value.to_string()) - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ConversationItem { - pub id: ConversationItemId, - pub response_id: Option, - pub item_type: String, - pub role: Option, - pub content: Value, - pub status: Option, - pub created_at: DateTime, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct NewConversationItem { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub id: Option, - pub response_id: Option, - pub item_type: String, - pub role: Option, - pub content: Value, - pub status: Option, -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -pub enum SortOrder { - Asc, - Desc, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ListParams { - pub limit: usize, - pub order: SortOrder, - pub after: Option, // item_id cursor -} - -pub type ConversationItemResult = Result; - -#[derive(Debug, thiserror::Error)] -pub enum ConversationItemStorageError { - #[error("Not found: {0}")] - NotFound(String), - - #[error("Storage error: {0}")] - StorageError(String), - - #[error("Serialization error: {0}")] - SerializationError(#[from] serde_json::Error), -} - -#[async_trait] -pub trait ConversationItemStorage: Send + Sync + 'static { - async fn create_item( - &self, - item: NewConversationItem, - ) -> ConversationItemResult; - - async fn link_item( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - added_at: DateTime, - ) -> ConversationItemResult<()>; - - async fn list_items( - &self, - conversation_id: &ConversationId, - params: ListParams, - ) -> ConversationItemResult>; - - /// Get a single item by ID - async fn get_item( - &self, - item_id: &ConversationItemId, - ) -> ConversationItemResult>; - - /// Check if an item is linked to a conversation - async fn is_item_linked( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - ) -> ConversationItemResult; - - /// Delete an item link from a conversation (does not delete the item itself) - async fn delete_item( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - ) -> ConversationItemResult<()>; -} - -/// Helper to build id prefix based on item_type -pub fn make_item_id(item_type: &str) -> ConversationItemId { - // Generate exactly 50 hex characters (25 bytes) for the part after the underscore - let mut rng = rand::rng(); - let mut bytes = [0u8; 25]; - rng.fill_bytes(&mut bytes); - let hex_string: String = bytes.iter().map(|b| format!("{:02x}", b)).collect(); - - let prefix: String = match item_type { - "message" => "msg".to_string(), - "reasoning" => "rs".to_string(), - "mcp_call" => "mcp".to_string(), - "mcp_list_tools" => "mcpl".to_string(), - "function_call" => "fc".to_string(), - other => { - // Fallback: first 3 letters of type or "itm" - let mut p = other.chars().take(3).collect::(); - if p.is_empty() { - p = "itm".to_string(); - } - p - } - }; - ConversationItemId(format!("{}_{}", prefix, hex_string)) -} - -// ============================================================================ -// PART 3: Response Storage -// ============================================================================ - -/// Response identifier -#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)] -pub struct ResponseId(pub String); - -impl ResponseId { - pub fn new() -> Self { - Self(ulid::Ulid::new().to_string()) - } -} - -impl Default for ResponseId { - fn default() -> Self { - Self::new() - } -} - -impl From for ResponseId { - fn from(value: String) -> Self { - Self(value) - } -} - -impl From<&str> for ResponseId { - fn from(value: &str) -> Self { - Self(value.to_string()) - } -} - -/// Stored response data -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct StoredResponse { - /// Unique response ID - pub id: ResponseId, - - /// ID of the previous response in the chain (if any) - pub previous_response_id: Option, - - /// Input items as JSON array - pub input: Value, - - /// System instructions used - pub instructions: Option, - - /// Output items as JSON array - pub output: Value, - - /// Tool calls made by the model (if any) - pub tool_calls: Vec, - - /// Custom metadata - pub metadata: HashMap, - - /// When this response was created - pub created_at: DateTime, - - /// Safety identifier for content moderation - pub safety_identifier: Option, - - /// Model used for generation - pub model: Option, - - /// Conversation id if associated with a conversation - #[serde(default)] - pub conversation_id: Option, - - /// Raw OpenAI response payload - #[serde(default)] - pub raw_response: Value, -} - -impl StoredResponse { - pub fn new(previous_response_id: Option) -> Self { - Self { - id: ResponseId::new(), - previous_response_id, - input: Value::Array(vec![]), - instructions: None, - output: Value::Array(vec![]), - tool_calls: Vec::new(), - metadata: HashMap::new(), - created_at: Utc::now(), - safety_identifier: None, - model: None, - conversation_id: None, - raw_response: Value::Null, - } - } -} - -/// Response chain - a sequence of related responses -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ResponseChain { - /// The responses in chronological order - pub responses: Vec, - - /// Metadata about the chain - pub metadata: HashMap, -} - -impl Default for ResponseChain { - fn default() -> Self { - Self::new() - } -} - -impl ResponseChain { - pub fn new() -> Self { - Self { - responses: Vec::new(), - metadata: HashMap::new(), - } - } - - /// Get the ID of the most recent response in the chain - pub fn latest_response_id(&self) -> Option<&ResponseId> { - self.responses.last().map(|r| &r.id) - } - - /// Add a response to the chain - pub fn add_response(&mut self, response: StoredResponse) { - self.responses.push(response); - } - - /// Build context from the chain for the next request - pub fn build_context(&self, max_responses: Option) -> Vec<(Value, Value)> { - let responses = if let Some(max) = max_responses { - let start = self.responses.len().saturating_sub(max); - &self.responses[start..] - } else { - &self.responses[..] - }; - - responses - .iter() - .map(|r| (r.input.clone(), r.output.clone())) - .collect() - } -} - -/// Error type for response storage operations -#[derive(Debug, thiserror::Error)] -pub enum ResponseStorageError { - #[error("Response not found: {0}")] - ResponseNotFound(String), - - #[error("Invalid chain: {0}")] - InvalidChain(String), - - #[error("Storage error: {0}")] - StorageError(String), - - #[error("Serialization error: {0}")] - SerializationError(#[from] serde_json::Error), -} - -pub type ResponseResult = Result; - -/// Trait for response storage -#[async_trait] -pub trait ResponseStorage: Send + Sync { - /// Store a new response - async fn store_response(&self, response: StoredResponse) -> ResponseResult; - - /// Get a response by ID - async fn get_response( - &self, - response_id: &ResponseId, - ) -> ResponseResult>; - - /// Delete a response - async fn delete_response(&self, response_id: &ResponseId) -> ResponseResult<()>; - - /// Get the chain of responses leading to a given response - /// Returns responses in chronological order (oldest first) - async fn get_response_chain( - &self, - response_id: &ResponseId, - max_depth: Option, - ) -> ResponseResult; - - /// List recent responses for a safety identifier - async fn list_identifier_responses( - &self, - identifier: &str, - limit: Option, - ) -> ResponseResult>; - - /// Delete all responses for a safety identifier - async fn delete_identifier_responses(&self, identifier: &str) -> ResponseResult; -} - -impl Default for StoredResponse { - fn default() -> Self { - Self::new(None) - } -} diff --git a/sgl-model-gateway/src/data_connector/factory.rs b/sgl-model-gateway/src/data_connector/factory.rs deleted file mode 100644 index 048a1d6bb..000000000 --- a/sgl-model-gateway/src/data_connector/factory.rs +++ /dev/null @@ -1,184 +0,0 @@ -// factory.rs -// -// Factory function to create storage backends based on configuration. - -use std::sync::Arc; - -use tracing::info; -use url::Url; - -use super::{ - core::{ConversationItemStorage, ConversationStorage, ResponseStorage}, - memory::{MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage}, - noop::{NoOpConversationItemStorage, NoOpConversationStorage, NoOpResponseStorage}, - oracle::{OracleConversationItemStorage, OracleConversationStorage, OracleResponseStorage}, -}; -use crate::{ - config::{HistoryBackend, OracleConfig, PostgresConfig, RedisConfig, RouterConfig}, - data_connector::{ - postgres::{ - PostgresConversationItemStorage, PostgresConversationStorage, PostgresResponseStorage, - PostgresStore, - }, - redis::{ - RedisConversationItemStorage, RedisConversationStorage, RedisResponseStorage, - RedisStore, - }, - }, -}; - -/// Type alias for the storage tuple returned by factory functions. -/// This avoids clippy::type_complexity warnings while keeping Arc explicit. -pub type StorageTuple = ( - Arc, - Arc, - Arc, -); - -/// Create all three storage backends based on router configuration. -/// -/// # Arguments -/// * `config` - Router configuration containing history_backend and oracle settings -/// -/// # Returns -/// Tuple of (response_storage, conversation_storage, conversation_item_storage) -/// -/// # Errors -/// Returns error string if Oracle configuration is missing or initialization fails -pub fn create_storage(config: &RouterConfig) -> Result { - match config.history_backend { - HistoryBackend::Memory => { - info!("Initializing data connector: Memory"); - Ok(( - Arc::new(MemoryResponseStorage::new()), - Arc::new(MemoryConversationStorage::new()), - Arc::new(MemoryConversationItemStorage::new()), - )) - } - HistoryBackend::None => { - info!("Initializing data connector: None (no persistence)"); - Ok(( - Arc::new(NoOpResponseStorage::new()), - Arc::new(NoOpConversationStorage::new()), - Arc::new(NoOpConversationItemStorage::new()), - )) - } - HistoryBackend::Oracle => { - let oracle_cfg = config - .oracle - .clone() - .ok_or("oracle configuration is required when history_backend=oracle")?; - - info!( - "Initializing data connector: Oracle ATP (pool: {}-{})", - oracle_cfg.pool_min, oracle_cfg.pool_max - ); - - let storages = create_oracle_storage(&oracle_cfg)?; - - 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) - } - HistoryBackend::Redis => { - let redis_cfg = config - .redis - .clone() - .ok_or("Redis configuration is required when history_backend=redis")?; - - let log_redis_url = match Url::parse(&redis_cfg.url) { - Ok(mut url) => { - if url.password().is_some() { - let _ = url.set_password(Some("****")); - } - url.to_string() - } - Err(_) => "".to_string(), - }; - - info!( - "Initializing data connector: Redis (url: {}, pool_max: {})", - log_redis_url, redis_cfg.pool_max - ); - - let storages = create_redis_storage(&redis_cfg)?; - - info!("Data connector initialized successfully: Redis"); - - Ok(storages) - } - } -} - -/// Create Oracle storage backends -fn create_oracle_storage(oracle_cfg: &OracleConfig) -> Result { - let response_storage = OracleResponseStorage::new(oracle_cfg.clone()) - .map_err(|err| format!("failed to initialize Oracle response storage: {err}"))?; - - let conversation_storage = OracleConversationStorage::new(oracle_cfg.clone()) - .map_err(|err| format!("failed to initialize Oracle conversation storage: {err}"))?; - - let conversation_item_storage = OracleConversationItemStorage::new(oracle_cfg.clone()) - .map_err(|err| format!("failed to initialize Oracle conversation item storage: {err}"))?; - - Ok(( - Arc::new(response_storage), - Arc::new(conversation_storage), - Arc::new(conversation_item_storage), - )) -} - -fn create_postgres_storage(postgres_cfg: &PostgresConfig) -> 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), - )) -} - -fn create_redis_storage(redis_cfg: &RedisConfig) -> Result { - let store = RedisStore::new(redis_cfg.clone())?; - let redis_resp = RedisResponseStorage::new(store.clone()); - let redis_conv = RedisConversationStorage::new(store.clone()); - let redis_item = RedisConversationItemStorage::new(store.clone()); - - Ok(( - Arc::new(redis_resp), - Arc::new(redis_conv), - Arc::new(redis_item), - )) -} diff --git a/sgl-model-gateway/src/data_connector/memory.rs b/sgl-model-gateway/src/data_connector/memory.rs deleted file mode 100644 index 17d17a52d..000000000 --- a/sgl-model-gateway/src/data_connector/memory.rs +++ /dev/null @@ -1,734 +0,0 @@ -//! In-memory storage implementations -//! -//! Used for development and testing - no persistence. -//! -//! Structure: -//! 1. MemoryConversationStorage -//! 2. MemoryConversationItemStorage -//! 3. MemoryResponseStorage - -use std::{ - collections::{BTreeMap, HashMap}, - sync::{Arc, RwLock}, -}; - -use async_trait::async_trait; -use chrono::{DateTime, Utc}; -use parking_lot::RwLock as ParkingLotRwLock; - -use super::core::*; - -// ============================================================================ -// PART 1: MemoryConversationStorage -// ============================================================================ - -/// In-memory conversation storage used for development and tests -#[derive(Default, Clone)] -pub struct MemoryConversationStorage { - inner: Arc>>, -} - -impl MemoryConversationStorage { - pub fn new() -> Self { - Self { - inner: Arc::new(ParkingLotRwLock::new(HashMap::new())), - } - } -} - -#[async_trait] -impl ConversationStorage for MemoryConversationStorage { - async fn create_conversation( - &self, - input: NewConversation, - ) -> ConversationResult { - let conversation = Conversation::new(input); - self.inner - .write() - .insert(conversation.id.clone(), conversation.clone()); - Ok(conversation) - } - - async fn get_conversation( - &self, - id: &ConversationId, - ) -> ConversationResult> { - Ok(self.inner.read().get(id).cloned()) - } - - async fn update_conversation( - &self, - id: &ConversationId, - metadata: Option, - ) -> ConversationResult> { - let mut store = self.inner.write(); - if let Some(entry) = store.get_mut(id) { - entry.metadata = metadata; - return Ok(Some(entry.clone())); - } - - Ok(None) - } - - async fn delete_conversation(&self, id: &ConversationId) -> ConversationResult { - let removed = self.inner.write().remove(id).is_some(); - Ok(removed) - } -} - -// ============================================================================ -// PART 2: MemoryConversationItemStorage -// ============================================================================ - -#[derive(Default)] -pub struct MemoryConversationItemStorage { - items: RwLock>, // item_id -> item - #[allow(clippy::type_complexity)] - links: RwLock>>, - // Per-conversation reverse index for fast after cursor lookup: item_id_str -> (ts, item_id_str) - #[allow(clippy::type_complexity)] - rev_index: RwLock>>, -} - -impl MemoryConversationItemStorage { - pub fn new() -> Self { - Self::default() - } -} - -#[async_trait] -impl ConversationItemStorage for MemoryConversationItemStorage { - async fn create_item( - &self, - new_item: NewConversationItem, - ) -> ConversationItemResult { - let id = new_item - .id - .clone() - .unwrap_or_else(|| make_item_id(&new_item.item_type)); - let created_at = Utc::now(); - let item = ConversationItem { - id: id.clone(), - response_id: new_item.response_id, - item_type: new_item.item_type, - role: new_item.role, - content: new_item.content, - status: new_item.status, - created_at, - }; - let mut items = self.items.write().unwrap(); - items.insert(id.clone(), item.clone()); - Ok(item) - } - - async fn link_item( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - added_at: DateTime, - ) -> ConversationItemResult<()> { - { - let mut links = self.links.write().unwrap(); - let entry = links.entry(conversation_id.clone()).or_default(); - entry.insert((added_at.timestamp(), item_id.0.clone()), item_id.clone()); - } - { - let mut rev = self.rev_index.write().unwrap(); - let entry = rev.entry(conversation_id.clone()).or_default(); - entry.insert(item_id.0.clone(), (added_at.timestamp(), item_id.0.clone())); - } - Ok(()) - } - - async fn list_items( - &self, - conversation_id: &ConversationId, - params: ListParams, - ) -> ConversationItemResult> { - let links_guard = self.links.read().unwrap(); - let map = match links_guard.get(conversation_id) { - Some(m) => m, - None => return Ok(Vec::new()), - }; - - let mut results: Vec = Vec::new(); - let after_key: Option<(i64, String)> = if let Some(after_id) = ¶ms.after { - // O(1) lookup via reverse index for this conversation - if let Some(conv_idx) = self.rev_index.read().unwrap().get(conversation_id) { - conv_idx.get(after_id).cloned() - } else { - None - } - } else { - None - }; - - let take = params.limit; - let items_guard = self.items.read().unwrap(); - - use std::ops::Bound::{Excluded, Unbounded}; - - // Helper to push item if it exists and stop when reaching the limit - let mut push_item = |key: &ConversationItemId| -> bool { - if let Some(it) = items_guard.get(key) { - results.push(it.clone()); - if results.len() == take { - return true; - } - } - false - }; - - match (params.order, after_key) { - (SortOrder::Desc, Some(k)) => { - for ((_ts, _id), item_key) in map.range(..k).rev() { - if push_item(item_key) { - break; - } - } - } - (SortOrder::Desc, None) => { - for ((_ts, _id), item_key) in map.iter().rev() { - if push_item(item_key) { - break; - } - } - } - (SortOrder::Asc, Some(k)) => { - for ((_ts, _id), item_key) in map.range((Excluded(k), Unbounded)) { - if push_item(item_key) { - break; - } - } - } - (SortOrder::Asc, None) => { - for ((_ts, _id), item_key) in map.iter() { - if push_item(item_key) { - break; - } - } - } - } - - Ok(results) - } - - async fn get_item( - &self, - item_id: &ConversationItemId, - ) -> ConversationItemResult> { - let items = self.items.read().unwrap(); - Ok(items.get(item_id).cloned()) - } - - async fn is_item_linked( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - ) -> ConversationItemResult { - let rev = self.rev_index.read().unwrap(); - if let Some(conv_idx) = rev.get(conversation_id) { - Ok(conv_idx.contains_key(&item_id.0)) - } else { - Ok(false) - } - } - - async fn delete_item( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - ) -> ConversationItemResult<()> { - // Get the key from rev_index and remove the entry at the same time - let key_to_remove = { - let mut rev = self.rev_index.write().unwrap(); - if let Some(conv_idx) = rev.get_mut(conversation_id) { - conv_idx.remove(&item_id.0) - } else { - None - } - }; - - // If the item was in rev_index, remove it from links as well - if let Some(key) = key_to_remove { - let mut links = self.links.write().unwrap(); - if let Some(conv_links) = links.get_mut(conversation_id) { - conv_links.remove(&key); - } - } - - Ok(()) - } -} - -// ============================================================================ -// PART 3: MemoryResponseStorage -// ============================================================================ - -/// Internal store structure holding both maps together -#[derive(Default)] -struct InnerStore { - /// All stored responses indexed by ID - responses: HashMap, - /// Index of response IDs by safety identifier - identifier_index: HashMap>, -} - -/// In-memory implementation of response storage -pub struct MemoryResponseStorage { - /// Single lock wrapping both maps to prevent deadlocks and ensure atomic updates - store: Arc>, -} - -impl MemoryResponseStorage { - pub fn new() -> Self { - Self { - store: Arc::new(ParkingLotRwLock::new(InnerStore::default())), - } - } - - /// Get statistics about the store - #[allow(dead_code)] - pub(super) fn stats(&self) -> MemoryStoreStats { - let store = self.store.read(); - MemoryStoreStats { - response_count: store.responses.len(), - identifier_count: store.identifier_index.len(), - } - } - - /// Clear all data (useful for testing) - pub fn clear(&self) { - let mut store = self.store.write(); - store.responses.clear(); - store.identifier_index.clear(); - } -} - -impl Default for MemoryResponseStorage { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl ResponseStorage for MemoryResponseStorage { - async fn store_response(&self, mut response: StoredResponse) -> ResponseResult { - // Generate ID if not set - if response.id.0.is_empty() { - response.id = ResponseId::new(); - } - - let response_id = response.id.clone(); - - // Single lock acquisition for atomic update - let mut store = self.store.write(); - - // Update safety identifier index if specified - if let Some(ref safety_identifier) = response.safety_identifier { - store - .identifier_index - .entry(safety_identifier.clone()) - .or_default() - .push(response_id.clone()); - } - - store.responses.insert(response_id.clone(), response); - tracing::debug!( - memory_store_size = store.responses.len(), - "Response stored in memory" - ); - - Ok(response_id) - } - - async fn get_response( - &self, - response_id: &ResponseId, - ) -> ResponseResult> { - let store = self.store.read(); - let result = store.responses.get(response_id).cloned(); - tracing::debug!(response_id = %response_id.0, found = result.is_some(), "Memory response lookup"); - Ok(result) - } - - async fn delete_response(&self, response_id: &ResponseId) -> ResponseResult<()> { - let mut store = self.store.write(); - - // Remove the response and update user index if needed - if let Some(response) = store.responses.remove(response_id) { - if let Some(ref safety_identifier) = response.safety_identifier { - if let Some(user_responses) = store.identifier_index.get_mut(safety_identifier) { - user_responses.retain(|id| id != response_id); - } - } - } - - Ok(()) - } - - async fn get_response_chain( - &self, - response_id: &ResponseId, - max_depth: Option, - ) -> ResponseResult { - let mut chain = ResponseChain::new(); - let max_depth = max_depth.unwrap_or(100); // Default max depth to prevent infinite loops - - // Collect all response IDs first - let mut response_ids = Vec::new(); - let mut current_id = Some(response_id.clone()); - let mut depth = 0; - - // Single lock acquisition to collect the chain - { - let store = self.store.read(); - while let Some(id) = current_id { - if depth >= max_depth { - break; - } - - if let Some(response) = store.responses.get(&id) { - response_ids.push(id); - current_id = response.previous_response_id.clone(); - depth += 1; - } else { - break; - } - } - } - - // Reverse to get chronological order (oldest first) - response_ids.reverse(); - - // Now collect the actual responses - let store = self.store.read(); - for id in response_ids { - if let Some(response) = store.responses.get(&id) { - chain.add_response(response.clone()); - } - } - - Ok(chain) - } - - async fn list_identifier_responses( - &self, - identifier: &str, - limit: Option, - ) -> ResponseResult> { - let store = self.store.read(); - - if let Some(user_response_ids) = store.identifier_index.get(identifier) { - // Collect responses with their timestamps for sorting - let mut responses_with_time: Vec<_> = user_response_ids - .iter() - .filter_map(|id| store.responses.get(id).map(|r| (r.created_at, id))) - .collect(); - - // Sort by creation time (newest first) - responses_with_time.sort_by(|a, b| b.0.cmp(&a.0)); - - // Apply limit and collect the actual responses - let limit = limit.unwrap_or(responses_with_time.len()); - let user_responses: Vec = responses_with_time - .into_iter() - .take(limit) - .filter_map(|(_, id)| store.responses.get(id).cloned()) - .collect(); - - Ok(user_responses) - } else { - Ok(Vec::new()) - } - } - - async fn delete_identifier_responses(&self, identifier: &str) -> ResponseResult { - let mut store = self.store.write(); - - if let Some(user_response_ids) = store.identifier_index.remove(identifier) { - let count = user_response_ids.len(); - for id in user_response_ids { - store.responses.remove(&id); - } - Ok(count) - } else { - Ok(0) - } - } -} - -/// Statistics for the memory store -#[derive(Debug, Clone)] -#[allow(dead_code)] -pub(super) struct MemoryStoreStats { - pub response_count: usize, - pub identifier_count: usize, -} - -#[cfg(test)] -mod tests { - use chrono::{TimeZone, Utc}; - use serde_json::json; - - use super::*; - - // ======================================================================== - // ConversationItem Tests - // ======================================================================== - - fn make_item( - item_type: &str, - role: Option<&str>, - content: serde_json::Value, - ) -> NewConversationItem { - NewConversationItem { - id: None, - response_id: None, - item_type: item_type.to_string(), - role: role.map(|r| r.to_string()), - content, - status: Some("completed".to_string()), - } - } - - #[tokio::test] - async fn test_list_ordering_and_cursors() { - let store = MemoryConversationItemStorage::new(); - let conv: ConversationId = "conv_test".into(); - - // Create 3 items and link them at controlled timestamps - let i1 = store - .create_item(make_item("message", Some("user"), json!([]))) - .await - .unwrap(); - let i2 = store - .create_item(make_item("message", Some("assistant"), json!([]))) - .await - .unwrap(); - let i3 = store - .create_item(make_item("reasoning", None, json!([]))) - .await - .unwrap(); - - let t1 = Utc.timestamp_opt(1_700_000_001, 0).single().unwrap(); - let t2 = Utc.timestamp_opt(1_700_000_002, 0).single().unwrap(); - let t3 = Utc.timestamp_opt(1_700_000_003, 0).single().unwrap(); - - store.link_item(&conv, &i1.id, t1).await.unwrap(); - store.link_item(&conv, &i2.id, t2).await.unwrap(); - store.link_item(&conv, &i3.id, t3).await.unwrap(); - - // Desc order, no cursor - let desc = store - .list_items( - &conv, - ListParams { - limit: 2, - order: SortOrder::Desc, - after: None, - }, - ) - .await - .unwrap(); - assert!(desc.len() >= 2); - assert_eq!(desc[0].id, i3.id); - assert_eq!(desc[1].id, i2.id); - - // Desc with cursor = i2 -> expect i1 next - let desc_after = store - .list_items( - &conv, - ListParams { - limit: 2, - order: SortOrder::Desc, - after: Some(i2.id.0.clone()), - }, - ) - .await - .unwrap(); - assert!(!desc_after.is_empty()); - assert_eq!(desc_after[0].id, i1.id); - - // Asc order, no cursor - let asc = store - .list_items( - &conv, - ListParams { - limit: 2, - order: SortOrder::Asc, - after: None, - }, - ) - .await - .unwrap(); - assert!(asc.len() >= 2); - assert_eq!(asc[0].id, i1.id); - assert_eq!(asc[1].id, i2.id); - - // Asc with cursor = i2 -> expect i3 next - let asc_after = store - .list_items( - &conv, - ListParams { - limit: 2, - order: SortOrder::Asc, - after: Some(i2.id.0.clone()), - }, - ) - .await - .unwrap(); - assert!(!asc_after.is_empty()); - assert_eq!(asc_after[0].id, i3.id); - } - - // ======================================================================== - // Response Tests - // ======================================================================== - - #[tokio::test] - async fn test_store_with_custom_id() { - let store = MemoryResponseStorage::new(); - let mut response = StoredResponse::new(None); - response.id = ResponseId::from("resp_custom"); - response.input = json!("Input"); - response.output = json!("Output"); - store.store_response(response.clone()).await.unwrap(); - let retrieved = store - .get_response(&ResponseId::from("resp_custom")) - .await - .unwrap(); - assert!(retrieved.is_some()); - assert_eq!(retrieved.unwrap().output, json!("Output")); - } - - #[tokio::test] - async fn test_memory_store_basic() { - let store = MemoryResponseStorage::new(); - - // Store a response - let mut response = StoredResponse::new(None); - response.input = json!("Hello"); - response.output = json!("Hi there!"); - let response_id = store.store_response(response).await.unwrap(); - - // Retrieve it - let retrieved = store.get_response(&response_id).await.unwrap(); - assert!(retrieved.is_some()); - assert_eq!(retrieved.unwrap().input, json!("Hello")); - - // Delete it - store.delete_response(&response_id).await.unwrap(); - let deleted = store.get_response(&response_id).await.unwrap(); - assert!(deleted.is_none()); - } - - #[tokio::test] - async fn test_response_chain() { - let store = MemoryResponseStorage::new(); - - // Create a chain of responses - let mut response1 = StoredResponse::new(None); - response1.input = json!("First"); - response1.output = json!("First response"); - let id1 = store.store_response(response1).await.unwrap(); - - let mut response2 = StoredResponse::new(Some(id1.clone())); - response2.input = json!("Second"); - response2.output = json!("Second response"); - let id2 = store.store_response(response2).await.unwrap(); - - let mut response3 = StoredResponse::new(Some(id2.clone())); - response3.input = json!("Third"); - response3.output = json!("Third response"); - let id3 = store.store_response(response3).await.unwrap(); - - // Get the chain - let chain = store.get_response_chain(&id3, None).await.unwrap(); - assert_eq!(chain.responses.len(), 3); - assert_eq!(chain.responses[0].input, json!("First")); - assert_eq!(chain.responses[1].input, json!("Second")); - assert_eq!(chain.responses[2].input, json!("Third")); - - let limited_chain = store.get_response_chain(&id3, Some(2)).await.unwrap(); - assert_eq!(limited_chain.responses.len(), 2); - assert_eq!(limited_chain.responses[0].input, json!("Second")); - assert_eq!(limited_chain.responses[1].input, json!("Third")); - } - - #[tokio::test] - async fn test_user_responses() { - let store = MemoryResponseStorage::new(); - - // Store responses for different users - let mut response1 = StoredResponse::new(None); - response1.input = json!("User1 message"); - response1.output = json!("Response to user1"); - response1.safety_identifier = Some("user1".to_string()); - store.store_response(response1).await.unwrap(); - - let mut response2 = StoredResponse::new(None); - response2.input = json!("Another user1 message"); - response2.output = json!("Another response to user1"); - response2.safety_identifier = Some("user1".to_string()); - store.store_response(response2).await.unwrap(); - - let mut response3 = StoredResponse::new(None); - response3.input = json!("User2 message"); - response3.output = json!("Response to user2"); - response3.safety_identifier = Some("user2".to_string()); - store.store_response(response3).await.unwrap(); - - // List user1's responses - let user1_responses = store - .list_identifier_responses("user1", None) - .await - .unwrap(); - assert_eq!(user1_responses.len(), 2); - - // List user2's responses - let user2_responses = store - .list_identifier_responses("user2", None) - .await - .unwrap(); - assert_eq!(user2_responses.len(), 1); - - // Delete user1's responses - let deleted_count = store.delete_identifier_responses("user1").await.unwrap(); - assert_eq!(deleted_count, 2); - - let user1_responses_after = store - .list_identifier_responses("user1", None) - .await - .unwrap(); - assert_eq!(user1_responses_after.len(), 0); - - // User2's responses should still be there - let user2_responses_after = store - .list_identifier_responses("user2", None) - .await - .unwrap(); - assert_eq!(user2_responses_after.len(), 1); - } - - #[tokio::test] - async fn test_memory_store_stats() { - let store = MemoryResponseStorage::new(); - - let mut response1 = StoredResponse::new(None); - response1.input = json!("Test1"); - response1.output = json!("Reply1"); - response1.safety_identifier = Some("user1".to_string()); - store.store_response(response1).await.unwrap(); - - let mut response2 = StoredResponse::new(None); - response2.input = json!("Test2"); - response2.output = json!("Reply2"); - response2.safety_identifier = Some("user2".to_string()); - store.store_response(response2).await.unwrap(); - - let stats = store.stats(); - assert_eq!(stats.response_count, 2); - assert_eq!(stats.identifier_count, 2); - } -} diff --git a/sgl-model-gateway/src/data_connector/mod.rs b/sgl-model-gateway/src/data_connector/mod.rs deleted file mode 100644 index e73c0828f..000000000 --- a/sgl-model-gateway/src/data_connector/mod.rs +++ /dev/null @@ -1,26 +0,0 @@ -// Data connector module for response storage and conversation storage -// -// Simplified module structure: -// - core.rs: All traits, data types, and errors -// - memory.rs: All in-memory storage implementations -// - noop.rs: All no-op storage implementations -// - 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; -mod redis; - -pub use core::{ - Conversation, ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage, - ConversationStorage, ListParams, NewConversation, NewConversationItem, ResponseId, - ResponseStorage, SortOrder, StoredResponse, -}; - -pub use factory::create_storage; -pub use memory::{MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage}; diff --git a/sgl-model-gateway/src/data_connector/noop.rs b/sgl-model-gateway/src/data_connector/noop.rs deleted file mode 100644 index ed041bb86..000000000 --- a/sgl-model-gateway/src/data_connector/noop.rs +++ /dev/null @@ -1,189 +0,0 @@ -//! NoOp storage implementations -//! -//! These implementations do nothing - useful for when persistence is disabled. -//! -//! Structure: -//! 1. NoOpConversationStorage -//! 2. NoOpConversationItemStorage -//! 3. NoOpResponseStorage - -use async_trait::async_trait; -use chrono::{DateTime, Utc}; - -use super::core::*; - -// ============================================================================ -// PART 1: NoOpConversationStorage -// ============================================================================ - -/// No-op implementation that synthesizes conversation responses without persistence -#[derive(Default, Debug, Clone)] -pub(super) struct NoOpConversationStorage; - -impl NoOpConversationStorage { - pub fn new() -> Self { - Self - } -} - -#[async_trait] -impl ConversationStorage for NoOpConversationStorage { - async fn create_conversation( - &self, - input: NewConversation, - ) -> ConversationResult { - Ok(Conversation::new(input)) - } - - async fn get_conversation( - &self, - _id: &ConversationId, - ) -> ConversationResult> { - Ok(None) - } - - async fn update_conversation( - &self, - _id: &ConversationId, - _metadata: Option, - ) -> ConversationResult> { - Ok(None) - } - - async fn delete_conversation(&self, _id: &ConversationId) -> ConversationResult { - Ok(false) - } -} - -// ============================================================================ -// PART 2: NoOpConversationItemStorage -// ============================================================================ - -/// No-op conversation item storage (does nothing) -#[derive(Clone, Copy, Default)] -pub(super) struct NoOpConversationItemStorage; - -impl NoOpConversationItemStorage { - pub fn new() -> Self { - Self - } -} - -#[async_trait] -impl ConversationItemStorage for NoOpConversationItemStorage { - async fn create_item( - &self, - item: NewConversationItem, - ) -> ConversationItemResult { - let id = item - .id - .clone() - .unwrap_or_else(|| make_item_id(&item.item_type)); - Ok(ConversationItem { - id, - response_id: item.response_id, - item_type: item.item_type, - role: item.role, - content: item.content, - status: item.status, - created_at: Utc::now(), - }) - } - - async fn link_item( - &self, - _conversation_id: &ConversationId, - _item_id: &ConversationItemId, - _added_at: DateTime, - ) -> ConversationItemResult<()> { - Ok(()) - } - - async fn list_items( - &self, - _conversation_id: &ConversationId, - _params: ListParams, - ) -> ConversationItemResult> { - Ok(Vec::new()) - } - - async fn get_item( - &self, - _item_id: &ConversationItemId, - ) -> ConversationItemResult> { - Ok(None) - } - - async fn is_item_linked( - &self, - _conversation_id: &ConversationId, - _item_id: &ConversationItemId, - ) -> ConversationItemResult { - Ok(false) - } - - async fn delete_item( - &self, - _conversation_id: &ConversationId, - _item_id: &ConversationItemId, - ) -> ConversationItemResult<()> { - Ok(()) - } -} - -// ============================================================================ -// PART 3: NoOpResponseStorage -// ============================================================================ - -/// No-op implementation of response storage (does nothing) -pub(super) struct NoOpResponseStorage; - -impl NoOpResponseStorage { - pub fn new() -> Self { - Self - } -} - -impl Default for NoOpResponseStorage { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl ResponseStorage for NoOpResponseStorage { - async fn store_response(&self, response: StoredResponse) -> ResponseResult { - Ok(response.id) - } - - async fn get_response( - &self, - _response_id: &ResponseId, - ) -> ResponseResult> { - Ok(None) - } - - async fn delete_response(&self, _response_id: &ResponseId) -> ResponseResult<()> { - Ok(()) - } - - async fn get_response_chain( - &self, - _response_id: &ResponseId, - _max_depth: Option, - ) -> ResponseResult { - Ok(ResponseChain::new()) - } - - async fn list_identifier_responses( - &self, - _identifier: &str, - _limit: Option, - ) -> ResponseResult> { - Ok(Vec::new()) - } - - async fn delete_identifier_responses(&self, _identifier: &str) -> ResponseResult { - Ok(0) - } -} diff --git a/sgl-model-gateway/src/data_connector/oracle.rs b/sgl-model-gateway/src/data_connector/oracle.rs deleted file mode 100644 index b98ea3949..000000000 --- a/sgl-model-gateway/src/data_connector/oracle.rs +++ /dev/null @@ -1,1130 +0,0 @@ -//! Oracle storage implementation using OracleStore helper. -//! -//! Structure: -//! 1. OracleStore helper and common utilities -//! 2. OracleConversationStorage -//! 3. OracleConversationItemStorage -//! 4. OracleResponseStorage - -use std::{path::Path, sync::Arc, time::Duration}; - -use async_trait::async_trait; -use chrono::{DateTime, Utc}; -use deadpool::managed::{Manager, Metrics, Pool, RecycleError, RecycleResult}; -use oracle::{ - sql_type::{OracleType, ToSql}, - Connection, Row, -}; -use serde_json::Value; - -use super::core::{ - make_item_id, Conversation, ConversationId, ConversationItem, ConversationItemId, - ConversationItemStorage, ConversationItemStorageError, ConversationMetadata, - ConversationStorage, ConversationStorageError, ListParams, NewConversation, - NewConversationItem, ResponseChain, ResponseId, ResponseStorage, ResponseStorageError, - SortOrder, StoredResponse, -}; -use crate::{ - config::OracleConfig, - data_connector::common::{ - parse_json_value, parse_metadata, parse_raw_response, parse_tool_calls, - }, -}; -// ============================================================================ -// PART 1: OracleStore Helper + Common Utilities -// ============================================================================ - -/// Shared Oracle connection pool infrastructure -/// -/// This helper eliminates ~540 LOC of duplication across storage implementations. -/// It handles connection pooling, error mapping, and client configuration. -pub(crate) struct OracleStore { - pool: Pool, -} - -impl OracleStore { - /// Create pool with custom schema initialization - /// - /// The `init_schema` function receives a connection and should: - /// - Check if tables/indexes exist - /// - Create them if needed - /// - Return Ok(()) on success or Err(message) on failure - pub fn new( - config: &OracleConfig, - init_schema: impl FnOnce(&Connection) -> Result<(), String>, - ) -> Result { - // Configure Oracle client (wallet, etc.) - configure_oracle_client(config)?; - - // Initialize schema using the provided function - let conn = Connection::connect( - &config.username, - &config.password, - &config.connect_descriptor, - ) - .map_err(map_oracle_error)?; - - init_schema(&conn)?; - drop(conn); - - // Create connection pool - let config_arc = Arc::new(config.clone()); - let manager = OracleConnectionManager { - params: Arc::new(OracleConnectParams::from_config(&config_arc)), - }; - - let mut builder = Pool::builder(manager) - .max_size(config.pool_max) - .runtime(deadpool::Runtime::Tokio1); - - if config.pool_timeout_secs > 0 { - builder = builder.wait_timeout(Some(Duration::from_secs(config.pool_timeout_secs))); - } - - let pool = builder - .build() - .map_err(|e| format!("Failed to build Oracle pool: {e}"))?; - - Ok(Self { pool }) - } - - /// Execute function with a connection from the pool - pub async fn execute(&self, func: F) -> Result - where - F: FnOnce(&Connection) -> Result + Send + 'static, - T: Send + 'static, - { - let connection = self - .pool - .get() - .await - .map_err(|e| format!("Failed to get Oracle connection: {e}"))?; - - tokio::task::spawn_blocking(move || { - let result = func(&connection); - drop(connection); - result - }) - .await - .map_err(|e| format!("Task execution failed: {e}"))? - } -} - -impl Clone for OracleStore { - fn clone(&self) -> Self { - Self { - pool: self.pool.clone(), - } - } -} - -// Error mapping helper -pub(crate) fn map_oracle_error(err: oracle::Error) -> String { - if let Some(db_err) = err.db_error() { - format!( - "Oracle error (code {}): {}", - db_err.code(), - db_err.message() - ) - } else { - err.to_string() - } -} - -// Client configuration helper -fn configure_oracle_client(config: &OracleConfig) -> Result<(), String> { - if let Some(wallet_path) = &config.wallet_path { - let path = Path::new(wallet_path); - - if !path.is_dir() { - return Err(format!( - "Oracle wallet path '{}' is not a directory", - wallet_path - )); - } - - if !path.join("tnsnames.ora").exists() && !path.join("sqlnet.ora").exists() { - return Err(format!( - "Oracle wallet path '{}' is missing tnsnames.ora or sqlnet.ora", - wallet_path - )); - } - - std::env::set_var("TNS_ADMIN", wallet_path); - } - Ok(()) -} - -// Connection parameters -#[derive(Clone)] -pub(crate) struct OracleConnectParams { - pub username: String, - pub password: String, - pub connect_descriptor: String, -} - -impl OracleConnectParams { - pub fn from_config(config: &OracleConfig) -> Self { - Self { - username: config.username.clone(), - password: config.password.clone(), - connect_descriptor: config.connect_descriptor.clone(), - } - } -} - -impl std::fmt::Debug for OracleConnectParams { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("OracleConnectParams") - .field("username", &self.username) - .field("connect_descriptor", &self.connect_descriptor) - .finish() - } -} - -// Connection manager (same for all stores) -#[derive(Clone)] -struct OracleConnectionManager { - params: Arc, -} - -impl std::fmt::Debug for OracleConnectionManager { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("OracleConnectionManager") - .field("username", &self.params.username) - .field("connect_descriptor", &self.params.connect_descriptor) - .finish() - } -} - -#[async_trait] -impl Manager for OracleConnectionManager { - type Type = Connection; - type Error = oracle::Error; - - fn create( - &self, - ) -> impl std::future::Future> + Send { - let params = self.params.clone(); - async move { - let mut conn = Connection::connect( - ¶ms.username, - ¶ms.password, - ¶ms.connect_descriptor, - )?; - conn.set_autocommit(true); - Ok(conn) - } - } - - #[allow(clippy::manual_async_fn)] - fn recycle( - &self, - conn: &mut Connection, - _: &Metrics, - ) -> impl std::future::Future> + Send { - async move { conn.ping().map_err(RecycleError::Backend) } - } -} - -// ============================================================================ -// PART 2: OracleConversationStorage -// ============================================================================ - -#[derive(Clone)] -pub(super) struct OracleConversationStorage { - store: OracleStore, -} - -impl OracleConversationStorage { - pub fn new(config: OracleConfig) -> Result { - let store = OracleStore::new(&config, |conn| { - // Check if table exists - let exists: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_tables WHERE table_name = 'CONVERSATIONS'", - &[], - ) - .map_err(map_oracle_error)?; - - // Create table if missing - if exists == 0 { - conn.execute( - "CREATE TABLE conversations ( - id VARCHAR2(64) PRIMARY KEY, - created_at TIMESTAMP WITH TIME ZONE, - metadata CLOB - )", - &[], - ) - .map_err(map_oracle_error)?; - } - - Ok(()) - }) - .map_err(ConversationStorageError::StorageError)?; - - Ok(Self { store }) - } - - fn parse_metadata( - raw: Option, - ) -> Result, ConversationStorageError> { - match raw { - Some(json) if !json.is_empty() => { - let value: Value = serde_json::from_str(&json)?; - match value { - Value::Object(map) => Ok(Some(map)), - Value::Null => Ok(None), - other => Err(ConversationStorageError::StorageError(format!( - "conversation metadata expected object, got {other}" - ))), - } - } - _ => Ok(None), - } - } -} - -#[async_trait] -impl ConversationStorage for OracleConversationStorage { - async fn create_conversation( - &self, - input: NewConversation, - ) -> Result { - let conversation = Conversation::new(input); - let id_str = conversation.id.0.clone(); - let created_at = conversation.created_at; - let metadata_json = conversation - .metadata - .as_ref() - .map(serde_json::to_string) - .transpose()?; - - self.store - .execute(move |conn| { - conn.execute( - "INSERT INTO conversations (id, created_at, metadata) VALUES (:1, :2, :3)", - &[&id_str, &created_at, &metadata_json], - ) - .map(|_| ()) - .map_err(map_oracle_error) - }) - .await - .map_err(ConversationStorageError::StorageError)?; - - Ok(conversation) - } - - async fn get_conversation( - &self, - id: &ConversationId, - ) -> Result, ConversationStorageError> { - let lookup = id.0.clone(); - self.store - .execute(move |conn| { - let mut stmt = conn - .statement("SELECT id, created_at, metadata FROM conversations WHERE id = :1") - .build() - .map_err(map_oracle_error)?; - let mut rows = stmt.query(&[&lookup]).map_err(map_oracle_error)?; - - if let Some(row_res) = rows.next() { - let row = row_res.map_err(map_oracle_error)?; - let id: String = row.get(0).map_err(map_oracle_error)?; - let created_at: DateTime = row.get(1).map_err(map_oracle_error)?; - let metadata_raw: Option = row.get(2).map_err(map_oracle_error)?; - let metadata = Self::parse_metadata(metadata_raw).map_err(|e| e.to_string())?; - Ok(Some(Conversation::with_parts( - ConversationId(id), - created_at, - metadata, - ))) - } else { - Ok(None) - } - }) - .await - .map_err(ConversationStorageError::StorageError) - } - - async fn update_conversation( - &self, - id: &ConversationId, - metadata: Option, - ) -> Result, ConversationStorageError> { - let id_str = id.0.clone(); - let metadata_json = metadata.as_ref().map(serde_json::to_string).transpose()?; - let conversation_id = id.clone(); - - self.store - .execute(move |conn| { - let mut stmt = conn - .statement( - "UPDATE conversations \ - SET metadata = :1 \ - WHERE id = :2 \ - RETURNING created_at INTO :3", - ) - .build() - .map_err(map_oracle_error)?; - - stmt.bind(3, &OracleType::TimestampTZ(6)) - .map_err(map_oracle_error)?; - stmt.execute(&[&metadata_json, &id_str]) - .map_err(map_oracle_error)?; - - if stmt.row_count().map_err(map_oracle_error)? == 0 { - return Ok(None); - } - - let mut created_at: Vec> = - stmt.returned_values(3).map_err(map_oracle_error)?; - let created_at = created_at - .pop() - .ok_or_else(|| "Oracle update did not return created_at".to_string())?; - - Ok(Some(Conversation::with_parts( - conversation_id, - created_at, - metadata, - ))) - }) - .await - .map_err(ConversationStorageError::StorageError) - } - - async fn delete_conversation( - &self, - id: &ConversationId, - ) -> Result { - let id_str = id.0.clone(); - let res = self - .store - .execute(move |conn| { - conn.execute("DELETE FROM conversations WHERE id = :1", &[&id_str]) - .map_err(map_oracle_error) - }) - .await - .map_err(ConversationStorageError::StorageError)?; - - Ok(res - .row_count() - .map_err(|e| ConversationStorageError::StorageError(map_oracle_error(e)))? - > 0) - } -} - -// ============================================================================ -// PART 3: OracleConversationItemStorage -// ============================================================================ - -#[derive(Clone)] -pub(super) struct OracleConversationItemStorage { - store: OracleStore, -} - -impl OracleConversationItemStorage { - pub fn new(config: OracleConfig) -> Result { - let store = OracleStore::new(&config, |conn| { - // Create conversation_items table - let exists_items: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_tables WHERE table_name = 'CONVERSATION_ITEMS'", - &[], - ) - .map_err(map_oracle_error)?; - - if exists_items == 0 { - conn.execute( - "CREATE TABLE conversation_items ( - id VARCHAR2(64) PRIMARY KEY, - response_id VARCHAR2(64), - item_type VARCHAR2(32) NOT NULL, - role VARCHAR2(32), - content CLOB, - status VARCHAR2(32), - created_at TIMESTAMP WITH TIME ZONE - )", - &[], - ) - .map_err(map_oracle_error)?; - } - - // Create conversation_item_links table - let exists_links: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_tables WHERE table_name = 'CONVERSATION_ITEM_LINKS'", - &[], - ) - .map_err(map_oracle_error)?; - - if exists_links == 0 { - conn.execute( - "CREATE TABLE conversation_item_links ( - conversation_id VARCHAR2(64) NOT NULL, - item_id VARCHAR2(64) NOT NULL, - added_at TIMESTAMP WITH TIME ZONE, - CONSTRAINT pk_conv_item_link PRIMARY KEY (conversation_id, item_id) - )", - &[], - ) - .map_err(map_oracle_error)?; - - conn.execute( - "CREATE INDEX conv_item_links_conv_idx ON conversation_item_links (conversation_id, added_at)", - &[], - ) - .map_err(map_oracle_error)?; - } - - Ok(()) - }) - .map_err(ConversationItemStorageError::StorageError)?; - - Ok(Self { store }) - } -} - -#[async_trait] -impl ConversationItemStorage for OracleConversationItemStorage { - 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, - 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(); - - self.store - .execute(move |conn| { - conn.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], - ) - .map_err(map_oracle_error)?; - Ok(()) - }) - .await - .map_err(ConversationItemStorageError::StorageError)?; - - Ok(conversation_item) - } - - async fn link_item( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - added_at: DateTime, - ) -> Result<(), ConversationItemStorageError> { - let cid = conversation_id.0.clone(); - let iid = item_id.0.clone(); - self.store - .execute(move |conn| { - conn.execute( - "INSERT INTO conversation_item_links (conversation_id, item_id, added_at) VALUES (:1, :2, :3)", - &[&cid, &iid, &added_at], - ) - .map_err(map_oracle_error)?; - Ok(()) - }) - .await - .map_err(ConversationItemStorageError::StorageError) - } - - async fn list_items( - &self, - conversation_id: &ConversationId, - params: ListParams, - ) -> Result, ConversationItemStorageError> { - 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(); - - // Resolve the added_at of the after cursor if provided - let after_key: Option<(DateTime, String)> = if let Some(ref aid) = after_id { - self.store - .execute({ - let cid = cid.clone(); - let aid = aid.clone(); - move |conn| { - let mut stmt = conn - .statement( - "SELECT added_at FROM conversation_item_links WHERE conversation_id = :1 AND item_id = :2", - ) - .build() - .map_err(map_oracle_error)?; - let mut rows = stmt.query(&[&cid, &aid]).map_err(map_oracle_error)?; - if let Some(row_res) = rows.next() { - let row = row_res.map_err(map_oracle_error)?; - let ts: DateTime = row.get(0).map_err(map_oracle_error)?; - Ok(Some((ts, aid))) - } else { - Ok(None) - } - } - }) - .await - .map_err(ConversationItemStorageError::StorageError)? - } else { - None - }; - - // Build the main list query - let rows: Vec<(String, Option, String, Option, Option, Option, DateTime)> = - self.store - .execute({ - let cid = cid.clone(); - move |conn| { - 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 = :cid", - ); - - // Cursor predicate - if let Some((_ts, _iid)) = &after_key { - if order_desc { - sql.push_str(" AND (l.added_at < :ats OR (l.added_at = :ats AND l.item_id < :iid))"); - } else { - sql.push_str(" AND (l.added_at > :ats OR (l.added_at = :ats AND l.item_id > :iid))"); - } - } - - // 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"); - } - sql.push_str(" FETCH NEXT :limit ROWS ONLY"); - - // Build params and perform a named SELECT query - let mut params_vec: Vec<(&str, &dyn ToSql)> = vec![("cid", &cid)]; - if let Some((ts, iid)) = &after_key { - params_vec.push(("ats", ts)); - params_vec.push(("iid", iid)); - } - params_vec.push(("limit", &limit)); - - let rows_iter = conn.query_named(&sql, ¶ms_vec).map_err(map_oracle_error)?; - - let mut out = Vec::new(); - for row_res in rows_iter { - let row = row_res.map_err(map_oracle_error)?; - let id: String = row.get(0).map_err(map_oracle_error)?; - let resp_id: Option = row.get(1).map_err(map_oracle_error)?; - let item_type: String = row.get(2).map_err(map_oracle_error)?; - let role: Option = row.get(3).map_err(map_oracle_error)?; - let content_raw: Option = row.get(4).map_err(map_oracle_error)?; - let status: Option = row.get(5).map_err(map_oracle_error)?; - let created_at: DateTime = row.get(6).map_err(map_oracle_error)?; - out.push((id, resp_id, item_type, role, content_raw, status, created_at)); - } - Ok(out) - } - }) - .await - .map_err(ConversationItemStorageError::StorageError)?; - - // Map rows to ConversationItem - rows.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(); - - self.store - .execute(move |conn| { - let mut stmt = conn - .statement( - "SELECT id, response_id, item_type, role, content, status, created_at \ - FROM conversation_items WHERE id = :1", - ) - .build() - .map_err(map_oracle_error)?; - - let mut rows = stmt.query(&[&iid]).map_err(map_oracle_error)?; - - if let Some(row_res) = rows.next() { - let row = row_res.map_err(map_oracle_error)?; - let id: String = row.get(0).map_err(map_oracle_error)?; - let response_id: Option = row.get(1).map_err(map_oracle_error)?; - let item_type: String = row.get(2).map_err(map_oracle_error)?; - let role: Option = row.get(3).map_err(map_oracle_error)?; - let content_raw: Option = row.get(4).map_err(map_oracle_error)?; - let status: Option = row.get(5).map_err(map_oracle_error)?; - let created_at: DateTime = row.get(6).map_err(map_oracle_error)?; - - let content = match content_raw { - Some(s) => serde_json::from_str(&s).map_err(|e| e.to_string())?, - None => Value::Null, - }; - - Ok(Some(ConversationItem { - id: ConversationItemId(id), - response_id, - item_type, - role, - content, - status, - created_at, - })) - } else { - Ok(None) - } - }) - .await - .map_err(ConversationItemStorageError::StorageError) - } - - async fn is_item_linked( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - ) -> Result { - let cid = conversation_id.0.clone(); - let iid = item_id.0.clone(); - - self.store - .execute(move |conn| { - let count: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM conversation_item_links WHERE conversation_id = :1 AND item_id = :2", - &[&cid, &iid], - ) - .map_err(map_oracle_error)?; - Ok(count > 0) - }) - .await - .map_err(ConversationItemStorageError::StorageError) - } - - async fn delete_item( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - ) -> Result<(), ConversationItemStorageError> { - let cid = conversation_id.0.clone(); - let iid = item_id.0.clone(); - - self.store - .execute(move |conn| { - conn.execute( - "DELETE FROM conversation_item_links WHERE conversation_id = :1 AND item_id = :2", - &[&cid, &iid], - ) - .map_err(map_oracle_error)?; - Ok(()) - }) - .await - .map_err(ConversationItemStorageError::StorageError) - } -} - -// ============================================================================ -// PART 4: OracleResponseStorage -// ============================================================================ - -const SELECT_BASE: &str = "SELECT id, previous_response_id, input, instructions, output, \ - tool_calls, metadata, created_at, safety_identifier, model, conversation_id, raw_response FROM responses"; - -#[derive(Clone)] -pub(super) struct OracleResponseStorage { - store: OracleStore, -} - -impl OracleResponseStorage { - pub fn new(config: OracleConfig) -> Result { - let store = OracleStore::new(&config, |conn| { - // Create responses table - let exists: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_tables WHERE table_name = 'RESPONSES'", - &[], - ) - .map_err(map_oracle_error)?; - - if exists == 0 { - conn.execute( - "CREATE TABLE responses ( - id VARCHAR2(64) PRIMARY KEY, - conversation_id VARCHAR2(64), - previous_response_id VARCHAR2(64), - input CLOB, - instructions CLOB, - output CLOB, - tool_calls CLOB, - metadata CLOB, - created_at TIMESTAMP WITH TIME ZONE, - safety_identifier VARCHAR2(128), - model VARCHAR2(128), - raw_response CLOB - )", - &[], - ) - .map_err(map_oracle_error)?; - } else { - Self::alter_safety_identifier_column(conn)?; - Self::remove_user_id_column_if_exists(conn)?; - } - - // Create indexes - create_index_if_missing( - conn, - "RESPONSES_PREV_IDX", - "CREATE INDEX responses_prev_idx ON responses(previous_response_id)", - )?; - create_index_if_missing( - conn, - "RESPONSES_USER_IDX", - "CREATE INDEX responses_user_idx ON responses(safety_identifier)", - )?; - - Ok(()) - }) - .map_err(ResponseStorageError::StorageError)?; - - Ok(Self { store }) - } - - // Alter safety_identifier column if missing - fn alter_safety_identifier_column(conn: &Connection) -> Result<(), String> { - let present: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_tab_columns WHERE table_name = 'RESPONSES' AND column_name = 'SAFETY_IDENTIFIER'", - &[], - ) - .map_err(map_oracle_error)?; - - if present == 0 { - if let Err(err) = conn.execute( - "ALTER TABLE responses ADD (safety_identifier VARCHAR2(128))", - &[], - ) { - let present_after: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_tab_columns WHERE table_name = 'RESPONSES' AND column_name = 'SAFETY_IDENTIFIER'", - &[], - ) - .map_err(map_oracle_error)?; - if present_after == 0 { - return Err(map_oracle_error(err)); - } - } - } - - Ok(()) - } - - // Remove user_id column if exists - fn remove_user_id_column_if_exists(conn: &Connection) -> Result<(), String> { - let present: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_tab_columns WHERE table_name = 'RESPONSES' AND column_name = 'USER_ID'", - &[], - ) - .map_err(map_oracle_error)?; - - if present > 0 { - if let Err(err) = conn.execute("ALTER TABLE responses DROP COLUMN USER_ID", &[]) { - let present_after: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_tab_columns WHERE table_name = 'RESPONSES' AND column_name = 'USER_ID'", - &[], - ) - .map_err(map_oracle_error)?; - if present_after > 0 { - return Err(map_oracle_error(err)); - } - } - } - - Ok(()) - } - - fn build_response_from_row(row: &Row) -> Result { - let id: String = row.get(0).map_err(map_oracle_error)?; - let previous: Option = row.get(1).map_err(map_oracle_error)?; - let input_json: Option = row.get(2).map_err(map_oracle_error)?; - let instructions: Option = row.get(3).map_err(map_oracle_error)?; - let output_json: Option = row.get(4).map_err(map_oracle_error)?; - let tool_calls_json: Option = row.get(5).map_err(map_oracle_error)?; - let metadata_json: Option = row.get(6).map_err(map_oracle_error)?; - let created_at: DateTime = row.get(7).map_err(map_oracle_error)?; - let safety_identifier: Option = row.get(8).map_err(map_oracle_error)?; - let model: Option = row.get(9).map_err(map_oracle_error)?; - let conversation_id: Option = row.get(10).map_err(map_oracle_error)?; - let raw_response_json: Option = row.get(11).map_err(map_oracle_error)?; - - 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 OracleResponseStorage { - 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 = serde_json::to_string(&response.input)?; - let json_output = serde_json::to_string(&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 = serde_json::to_string(&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(); - - self.store - .execute(move |conn| { - conn.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, - &json_tool_calls, - &json_metadata, - &created_at, - &safety_identifier, - &model, - &conversation_id, - &json_raw_response, - ], - ) - .map(|_| ()) - .map_err(map_oracle_error) - }) - .await - .map_err(ResponseStorageError::StorageError)?; - - Ok(response_id) - } - - async fn get_response( - &self, - response_id: &ResponseId, - ) -> Result, ResponseStorageError> { - let id = response_id.0.clone(); - self.store - .execute(move |conn| { - let mut stmt = conn - .statement(&format!("{} WHERE id = :1", SELECT_BASE)) - .build() - .map_err(map_oracle_error)?; - let mut rows = stmt.query(&[&id]).map_err(map_oracle_error)?; - match rows.next() { - Some(row) => { - let row = row.map_err(map_oracle_error)?; - Self::build_response_from_row(&row).map(Some) - } - None => Ok(None), - } - }) - .await - .map_err(ResponseStorageError::StorageError) - } - - async fn delete_response(&self, response_id: &ResponseId) -> Result<(), ResponseStorageError> { - let id = response_id.0.clone(); - self.store - .execute(move |conn| { - conn.execute("DELETE FROM responses WHERE id = :1", &[&id]) - .map(|_| ()) - .map_err(map_oracle_error) - }) - .await - .map_err(ResponseStorageError::StorageError) - } - - async fn get_response_chain( - &self, - response_id: &ResponseId, - max_depth: Option, - ) -> Result { - 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, - ) -> Result, ResponseStorageError> { - let identifier = identifier.to_string(); - - self.store - .execute(move |conn| { - let sql = if let Some(limit) = limit { - format!( - "SELECT * FROM ({} WHERE safety_identifier = :1 ORDER BY created_at DESC) WHERE ROWNUM <= {}", - SELECT_BASE, limit - ) - } else { - format!("{} WHERE safety_identifier = :1 ORDER BY created_at DESC", SELECT_BASE) - }; - - let mut stmt = conn.statement(&sql).build().map_err(map_oracle_error)?; - let mut rows = stmt.query(&[&identifier]).map_err(map_oracle_error)?; - let mut results = Vec::new(); - - for row in &mut rows { - let row = row.map_err(map_oracle_error)?; - results.push(Self::build_response_from_row(&row)?); - } - - Ok(results) - }) - .await - .map_err(ResponseStorageError::StorageError) - } - - async fn delete_identifier_responses( - &self, - identifier: &str, - ) -> Result { - let identifier = identifier.to_string(); - let affected = self - .store - .execute(move |conn| { - conn.execute( - "DELETE FROM responses WHERE safety_identifier = :1", - &[&identifier], - ) - .map_err(map_oracle_error) - }) - .await - .map_err(ResponseStorageError::StorageError)?; - - let deleted = affected - .row_count() - .map_err(|e| ResponseStorageError::StorageError(map_oracle_error(e)))? - as usize; - Ok(deleted) - } -} - -// Helper functions for response parsing - -fn create_index_if_missing(conn: &Connection, index_name: &str, ddl: &str) -> Result<(), String> { - let count: i64 = conn - .query_row_as( - "SELECT COUNT(*) FROM user_indexes WHERE table_name = 'RESPONSES' AND index_name = :1", - &[&index_name], - ) - .map_err(map_oracle_error)?; - - if count == 0 { - if let Err(err) = conn.execute(ddl, &[]) { - if let Some(db_err) = err.db_error() { - // ORA-00955: name is already used by an existing object - // ORA-01408: such column list already indexed - // Both errors indicate the index already exists (race condition) - if db_err.code() != 955 && db_err.code() != 1408 { - return Err(map_oracle_error(err)); - } - } else { - return Err(map_oracle_error(err)); - } - } - } - - Ok(()) -} diff --git a/sgl-model-gateway/src/data_connector/postgres.rs b/sgl-model-gateway/src/data_connector/postgres.rs deleted file mode 100644 index f0891250b..000000000 --- a/sgl-model-gateway/src/data_connector/postgres.rs +++ /dev/null @@ -1,700 +0,0 @@ -//! 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 tracing; - -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(super) 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(super) 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(super) 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(); - tracing::debug!(rows_affected = insert_count, "Response stored in Postgres"); - 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-model-gateway/src/data_connector/redis.rs b/sgl-model-gateway/src/data_connector/redis.rs deleted file mode 100644 index 9e2dd0b95..000000000 --- a/sgl-model-gateway/src/data_connector/redis.rs +++ /dev/null @@ -1,913 +0,0 @@ -//! Redis storage implementation using RedisStore helper -//! -//! Structure: -//! 1. RedisStore helper and common utilities -//! 2. RedisConversationStorage -//! 3. RedisConversationItemStorage -//! 4. RedisResponseStorage - -use async_trait::async_trait; -use chrono::{DateTime, Utc}; -use deadpool_redis::{Config, Pool, Runtime}; -use redis::AsyncCommands; -use serde_json::Value; - -use crate::{ - config::RedisConfig, - 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 RedisStore { - pool: Pool, - retention_days: Option, -} - -impl RedisStore { - pub fn new(config: RedisConfig) -> Result { - let mut cfg = Config::from_url(config.url); - cfg.pool = Some(deadpool_redis::PoolConfig::new(config.pool_max)); - let pool = cfg - .create_pool(Some(Runtime::Tokio1)) - .map_err(|e| e.to_string())?; - Ok(Self { - pool, - retention_days: config.retention_days, - }) - } -} - -impl Clone for RedisStore { - fn clone(&self) -> Self { - Self { - pool: self.pool.clone(), - retention_days: self.retention_days, - } - } -} - -pub struct RedisConversationStorage { - store: RedisStore, -} - -impl RedisConversationStorage { - pub fn new(store: RedisStore) -> Self { - Self { store } - } - - fn conversation_key(id: &str) -> String { - format!("conversation:{}", id) - } - - 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 RedisConversationStorage { - 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 mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - let key = Self::conversation_key(id_str); - - let mut pipe = redis::pipe(); - pipe.hset(&key, "id", id_str); - pipe.hset(&key, "created_at", created_at.to_rfc3339()); - if let Some(meta) = metadata_json { - pipe.hset(&key, "metadata", meta); - } - - // Expire after configured retention days (optional) - if let Some(days) = self.store.retention_days { - pipe.expire(&key, (days * 24 * 60 * 60) as i64); - } - - pipe.query_async::<()>(&mut conn) - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - - Ok(conversation) - } - - async fn get_conversation( - &self, - id: &ConversationId, - ) -> Result, ConversationStorageError> { - let id_str = id.0.as_str(); - let key = Self::conversation_key(id_str); - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - - let exists: bool = conn - .exists(&key) - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - if !exists { - return Ok(None); - } - - let (created_at_str, metadata_json): (String, Option) = redis::pipe() - .hget(&key, "created_at") - .hget(&key, "metadata") - .query_async(&mut conn) - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - - let created_at = DateTime::parse_from_rfc3339(&created_at_str) - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))? - .with_timezone(&Utc); - - let metadata = Self::parse_metadata(metadata_json)?; - - Ok(Some(Conversation::with_parts( - id.clone(), - created_at, - metadata, - ))) - } - - async fn update_conversation( - &self, - id: &ConversationId, - metadata: Option, - ) -> Result, ConversationStorageError> { - let id_str = id.0.as_str(); - let key = Self::conversation_key(id_str); - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - - let exists: bool = conn - .exists(&key) - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - if !exists { - return Ok(None); - } - - let metadata_json = metadata.as_ref().map(serde_json::to_string).transpose()?; - - if let Some(meta) = metadata_json { - conn.hset::<_, _, _, ()>(&key, "metadata", meta) - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - } else { - conn.hdel::<_, _, ()>(&key, "metadata") - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - } - - // We need to fetch created_at to return the full object - let created_at_str: String = conn - .hget(&key, "created_at") - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - let created_at = DateTime::parse_from_rfc3339(&created_at_str) - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))? - .with_timezone(&Utc); - - Ok(Some(Conversation::with_parts( - id.clone(), - created_at, - metadata, - ))) - } - - async fn delete_conversation(&self, id: &ConversationId) -> ConversationResult { - let id_str = id.0.as_str(); - let key = Self::conversation_key(id_str); - // Also delete the items list for this conversation - let items_key = format!("{}:items", key); - - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - - let count: usize = redis::pipe() - .del(&key) - .del(&items_key) - .query_async(&mut conn) - .await - .map_err(|e| ConversationStorageError::StorageError(e.to_string()))?; - - Ok(count > 0) - } -} - -pub struct RedisConversationItemStorage { - store: RedisStore, -} - -impl RedisConversationItemStorage { - pub fn new(store: RedisStore) -> Self { - Self { store } - } - - fn item_key(id: &str) -> String { - format!("item:{}", id) - } - - fn conv_items_key(conv_id: &str) -> String { - format!("conversation:{}:items", conv_id) - } -} - -#[async_trait] -impl ConversationItemStorage for RedisConversationItemStorage { - 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.as_str(); - let key = Self::item_key(id_str); - - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - let mut pipe = redis::pipe(); - - pipe.hset(&key, "id", id_str); - if let Some(rid) = &conversation_item.response_id { - pipe.hset(&key, "response_id", rid); - } - pipe.hset(&key, "item_type", &conversation_item.item_type); - if let Some(r) = &conversation_item.role { - pipe.hset(&key, "role", r); - } - pipe.hset(&key, "content", content_json); - if let Some(s) = &conversation_item.status { - pipe.hset(&key, "status", s); - } - pipe.hset(&key, "created_at", created_at.to_rfc3339()); - - // Expire after configured retention days - if let Some(days) = self.store.retention_days { - pipe.expire(&key, (days * 24 * 60 * 60) as i64); - } - - pipe.query_async::<()>(&mut conn) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - Ok(conversation_item) - } - - async fn link_item( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - added_at: DateTime, - ) -> ConversationItemResult<()> { - let cid = conversation_id.0.as_str(); - let iid = item_id.0.as_str(); - let key = Self::conv_items_key(cid); - - let score = added_at.timestamp_millis() as f64; - - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - conn.zadd::<_, _, _, ()>(&key, iid, score) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - Ok(()) - } - - async fn list_items( - &self, - conversation_id: &ConversationId, - params: ListParams, - ) -> ConversationItemResult> { - let cid = conversation_id.0.as_str(); - let key = Self::conv_items_key(cid); - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - let mut min = "-inf".to_string(); - let mut max = "+inf".to_string(); - - if let Some(after_id) = ¶ms.after { - let score: Option = conn - .zscore(&key, after_id) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - if let Some(s) = score { - match params.order { - SortOrder::Asc => min = format!("({}", s), - SortOrder::Desc => max = format!("({}", s), - } - } - } - - let item_ids: Vec = match params.order { - SortOrder::Asc => { - // ZRANGEBYSCORE key min max LIMIT offset count - conn.zrangebyscore_limit(&key, min, max, 0, params.limit as isize) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))? - } - SortOrder::Desc => { - // ZREVRANGEBYSCORE key max min LIMIT offset count - conn.zrevrangebyscore_limit(&key, max, min, 0, params.limit as isize) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))? - } - }; - - if item_ids.is_empty() { - return Ok(Vec::::new()); - } - - // Fetch all items in pipeline - let mut pipe = redis::pipe(); - for iid in &item_ids { - pipe.hgetall(Self::item_key(iid)); - } - - let results: Vec> = pipe - .query_async(&mut conn) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - let mut items: Vec = Vec::new(); - for (i, map) in results.into_iter().enumerate() { - if map.is_empty() { - // Item might have been deleted or expired, skip - continue; - } - - let id = ConversationItemId( - map.get("id") - .cloned() - .unwrap_or_else(|| item_ids[i].clone()), - ); - let response_id = map.get("response_id").cloned(); - let item_type = map.get("item_type").cloned().unwrap_or_default(); - let role = map.get("role").cloned(); - let status = map.get("status").cloned(); - - let content_raw = map.get("content"); - let content = match content_raw { - Some(s) => serde_json::from_str(s).unwrap_or(Value::Null), - None => Value::Null, - }; - - let created_at_str = map - .get("created_at") - .cloned() - .unwrap_or_else(|| Utc::now().to_rfc3339()); - let created_at = DateTime::parse_from_rfc3339(&created_at_str) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or_else(|_| Utc::now()); - - items.push(ConversationItem { - id, - response_id, - item_type, - role, - content, - status, - created_at, - }); - } - - Ok(items) - } - - async fn get_item( - &self, - item_id: &ConversationItemId, - ) -> ConversationItemResult> { - let iid = item_id.0.as_str(); - let key = Self::item_key(iid); - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - let map: std::collections::HashMap = conn - .hgetall(&key) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - if map.is_empty() { - return Ok(None); - } - - let id = ConversationItemId(map.get("id").cloned().unwrap_or_else(|| iid.to_string())); - let response_id = map.get("response_id").cloned(); - let item_type = map.get("item_type").cloned().unwrap_or_default(); - let role = map.get("role").cloned(); - let status = map.get("status").cloned(); - - let content_raw = map.get("content"); - let content = match content_raw { - Some(s) => serde_json::from_str(s).unwrap_or(Value::Null), - None => Value::Null, - }; - - let created_at_str = map - .get("created_at") - .cloned() - .unwrap_or_else(|| Utc::now().to_rfc3339()); - let created_at = DateTime::parse_from_rfc3339(&created_at_str) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or_else(|_| Utc::now()); - - Ok(Some(ConversationItem { - 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.as_str(); - let iid = item_id.0.as_str(); - let key = Self::conv_items_key(cid); - - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - let score: Option = conn - .zscore(&key, iid) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - Ok(score.is_some()) - } - - async fn delete_item( - &self, - conversation_id: &ConversationId, - item_id: &ConversationItemId, - ) -> ConversationItemResult<()> { - let cid = conversation_id.0.as_str(); - let iid = item_id.0.as_str(); - let key = Self::conv_items_key(cid); - - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - conn.zrem::<_, _, ()>(&key, iid) - .await - .map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?; - - Ok(()) - } -} - -pub struct RedisResponseStorage { - store: RedisStore, -} - -impl RedisResponseStorage { - pub fn new(store: RedisStore) -> Self { - Self { store } - } - - fn response_key(id: &str) -> String { - format!("response:{}", id) - } - - fn safety_key(identifier: &str) -> String { - format!("safety:{}:responses", identifier) - } -} - -#[async_trait] -impl ResponseStorage for RedisResponseStorage { - async fn store_response( - &self, - response: StoredResponse, - ) -> Result { - let response_id = response.id.clone(); - let response_id_str = response_id.0.as_str(); - let key = Self::response_key(response_id_str); - - let json_input = serde_json::to_string(&response.input)?; - let json_output = serde_json::to_string(&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 = serde_json::to_string(&response.raw_response)?; - - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - let mut pipe = redis::pipe(); - - pipe.hset(&key, "id", response_id_str); - if let Some(prev) = &response.previous_response_id { - pipe.hset(&key, "previous_response_id", &prev.0); - } - pipe.hset(&key, "input", json_input); - if let Some(inst) = &response.instructions { - pipe.hset(&key, "instructions", inst); - } - pipe.hset(&key, "output", json_output); - pipe.hset(&key, "tool_calls", json_tool_calls); - pipe.hset(&key, "metadata", json_metadata); - pipe.hset(&key, "created_at", response.created_at.to_rfc3339()); - if let Some(safety) = &response.safety_identifier { - pipe.hset(&key, "safety_identifier", safety); - } - if let Some(model) = &response.model { - pipe.hset(&key, "model", model); - } - if let Some(cid) = &response.conversation_id { - pipe.hset(&key, "conversation_id", cid); - } - pipe.hset(&key, "raw_response", json_raw_response); - - // Expire after configured retention days - if let Some(days) = self.store.retention_days { - pipe.expire(&key, (days * 24 * 60 * 60) as i64); - } - - pipe.query_async::<()>(&mut conn) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - // Index by safety identifier if present - if let Some(safety) = &response.safety_identifier { - let safety_key = Self::safety_key(safety); - let score = response.created_at.timestamp_millis() as f64; - conn.zadd::<_, _, _, ()>(safety_key, response_id_str, score) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - } - - Ok(response_id) - } - - async fn get_response( - &self, - response_id: &ResponseId, - ) -> Result, ResponseStorageError> { - let id = response_id.0.as_str(); - let key = Self::response_key(id); - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - let map: std::collections::HashMap = conn - .hgetall(&key) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - if map.is_empty() { - return Ok(None); - } - - let id = ResponseId(map.get("id").cloned().unwrap_or_else(|| id.to_string())); - let previous_response_id = map - .get("previous_response_id") - .map(|s| ResponseId(s.clone())); - let conversation_id = map.get("conversation_id").cloned(); - - let input = match parse_json_value(map.get("input").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - let instructions = map.get("instructions").cloned(); - let output = match parse_json_value(map.get("output").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - let tool_calls = match parse_tool_calls(map.get("tool_calls").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - let metadata = match parse_metadata(map.get("metadata").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - - let created_at_str = map - .get("created_at") - .cloned() - .unwrap_or_else(|| Utc::now().to_rfc3339()); - let created_at = DateTime::parse_from_rfc3339(&created_at_str) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or_else(|_| Utc::now()); - - let safety_identifier = map.get("safety_identifier").cloned(); - let model = map.get("model").cloned(); - let raw_response = match parse_raw_response(map.get("raw_response").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - - Ok(Some(StoredResponse { - id, - previous_response_id, - input, - instructions, - output, - tool_calls, - metadata, - created_at, - safety_identifier, - model, - conversation_id, - raw_response, - })) - } - - async fn delete_response(&self, response_id: &ResponseId) -> ResponseResult<()> { - let id = response_id.0.as_str(); - let key = Self::response_key(id); - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - // First check if it has a safety identifier to remove from index - let safety: Option = conn.hget(&key, "safety_identifier").await.ok(); - - conn.del::<_, ()>(&key) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - if let Some(s) = safety { - conn.zrem::<_, _, ()>(Self::safety_key(&s), id) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - } - - 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 key = Self::safety_key(identifier); - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - // ZREVRANGE key 0 limit-1 - let stop = match limit { - Some(l) => (l as isize) - 1, - None => -1, - }; - - let response_ids: Vec = conn - .zrevrange(&key, 0, stop) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - if response_ids.is_empty() { - return Ok(Vec::::new()); - } - - let mut pipe = redis::pipe(); - for id in &response_ids { - pipe.hgetall(Self::response_key(id)); - } - - let results: Vec> = pipe - .query_async(&mut conn) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - let mut out: Vec = Vec::with_capacity(results.len()); - for (i, map) in results.into_iter().enumerate() { - if map.is_empty() { - continue; - } - - let id = ResponseId( - map.get("id") - .cloned() - .unwrap_or_else(|| response_ids[i].clone()), - ); - let previous_response_id = map - .get("previous_response_id") - .map(|s| ResponseId(s.clone())); - let conversation_id = map.get("conversation_id").cloned(); - - let input = match parse_json_value(map.get("input").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - let instructions = map.get("instructions").cloned(); - let output = match parse_json_value(map.get("output").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - let tool_calls = match parse_tool_calls(map.get("tool_calls").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - let metadata = match parse_metadata(map.get("metadata").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - - let created_at_str = map - .get("created_at") - .cloned() - .unwrap_or_else(|| Utc::now().to_rfc3339()); - let created_at = DateTime::parse_from_rfc3339(&created_at_str) - .map(|dt| dt.with_timezone(&Utc)) - .unwrap_or_else(|_| Utc::now()); - - let safety_identifier = map.get("safety_identifier").cloned(); - let model = map.get("model").cloned(); - let raw_response = match parse_raw_response(map.get("raw_response").cloned()) { - Ok(v) => v, - Err(e) => return Err(ResponseStorageError::StorageError(e)), - }; - - out.push(StoredResponse { - id, - previous_response_id, - input, - instructions, - output, - tool_calls, - metadata, - created_at, - safety_identifier, - model, - conversation_id, - raw_response, - }); - } - - Ok(out) - } - - async fn delete_identifier_responses(&self, identifier: &str) -> ResponseResult { - let key = Self::safety_key(identifier); - let mut conn = self - .store - .pool - .get() - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - // Get all IDs - let response_ids: Vec = conn - .zrange(&key, 0, -1) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - let count = response_ids.len(); - - if count == 0 { - return Ok(0); - } - - let mut pipe = redis::pipe(); - for id in response_ids { - pipe.del(Self::response_key(&id)); - } - pipe.del(&key); - - pipe.query_async::<()>(&mut conn) - .await - .map_err(|e| ResponseStorageError::StorageError(e.to_string()))?; - - Ok(count) - } -} diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs index a75f0e5a4..3ce0f5c8c 100644 --- a/sgl-model-gateway/src/lib.rs +++ b/sgl-model-gateway/src/lib.rs @@ -2,7 +2,7 @@ pub mod app_context; pub use smg_auth as auth; pub mod config; pub mod core; -pub mod data_connector; +pub use data_connector; pub mod grpc_client; pub mod mcp; pub mod mesh; diff --git a/sgl-model-gateway/src/service_discovery.rs b/sgl-model-gateway/src/service_discovery.rs index 986ab1f19..a743cb771 100644 --- a/sgl-model-gateway/src/service_discovery.rs +++ b/sgl-model-gateway/src/service_discovery.rs @@ -841,10 +841,10 @@ mod tests { reasoning_parser_factory: None, tool_parser_factory: None, router_manager: None, - response_storage: Arc::new(crate::data_connector::MemoryResponseStorage::new()), - conversation_storage: Arc::new(crate::data_connector::MemoryConversationStorage::new()), + response_storage: Arc::new(data_connector::MemoryResponseStorage::new()), + conversation_storage: Arc::new(data_connector::MemoryConversationStorage::new()), conversation_item_storage: Arc::new( - crate::data_connector::MemoryConversationItemStorage::new(), + data_connector::MemoryConversationItemStorage::new(), ), load_monitor: None, configured_reasoning_parser: None,