[model-gateway] Add Redis support as a history backend (#16300)
This commit is contained in:
@@ -858,6 +858,7 @@ Prefill pods can expose bootstrap ports via the `sglang.ai/bootstrap-port` annot
|
|||||||
| `none` | No persistence | `--history-backend none` |
|
| `none` | No persistence | `--history-backend none` |
|
||||||
| `oracle` | Oracle Autonomous Database | `--history-backend oracle` |
|
| `oracle` | Oracle Autonomous Database | `--history-backend oracle` |
|
||||||
| `postgres` | PostgreSQL Database | `--history-backend postgres` |
|
| `postgres` | PostgreSQL Database | `--history-backend postgres` |
|
||||||
|
| `redis` | Redis | `--history-backend redis` |
|
||||||
|
|
||||||
### Oracle Configuration
|
### Oracle Configuration
|
||||||
|
|
||||||
@@ -892,6 +893,22 @@ python -m sglang_router.launch_router \
|
|||||||
--history-backend postgres
|
--history-backend postgres
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### Redis Configuration
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export REDIS_URL="redis://localhost:6379"
|
||||||
|
export REDIS_POOL_MAX=16
|
||||||
|
export REDIS_RETENTION_DAYS=30
|
||||||
|
|
||||||
|
python -m sglang_router.launch_router \
|
||||||
|
--backend openai \
|
||||||
|
--worker-urls https://api.openai.com \
|
||||||
|
--history-backend redis \
|
||||||
|
--redis-retention-days 30
|
||||||
|
```
|
||||||
|
|
||||||
|
Use `--redis-retention-days -1` for persistent storage (default is 30 days).
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## WASM Middleware
|
## WASM Middleware
|
||||||
|
|||||||
@@ -121,6 +121,8 @@ bitflags = "2.10.0"
|
|||||||
once_cell = "1.21.3"
|
once_cell = "1.21.3"
|
||||||
tokio-postgres = { version = "0.7.15", features = ["runtime","with-chrono-0_4","with-serde_json-1","array-impls"] }
|
tokio-postgres = { version = "0.7.15", features = ["runtime","with-chrono-0_4","with-serde_json-1","array-impls"] }
|
||||||
deadpool-postgres = "0.14.1"
|
deadpool-postgres = "0.14.1"
|
||||||
|
redis = { version = "0.27.6", features = ["tokio-comp", "json", "connection-manager"] }
|
||||||
|
deadpool-redis = "0.18.0"
|
||||||
|
|
||||||
|
|
||||||
# wasm dependencies
|
# wasm dependencies
|
||||||
|
|||||||
@@ -639,6 +639,7 @@ Supported tool parsers: `json`, `python`, `xml`.
|
|||||||
- `--history-backend none` disables persistence while keeping APIs.
|
- `--history-backend none` disables persistence while keeping APIs.
|
||||||
- `--history-backend oracle` uses Oracle Autonomous Database; provide credentials via flags or environment variables.
|
- `--history-backend oracle` uses Oracle Autonomous Database; provide credentials via flags or environment variables.
|
||||||
- `--history-backend postgres` uses PostgreSQL Database.
|
- `--history-backend postgres` uses PostgreSQL Database.
|
||||||
|
- `--history-backend redis` uses Redis.
|
||||||
- Conversation item storage mirrors the history backend (Oracle or memory). The same storage powers OpenAI `/responses` and conversation APIs.
|
- Conversation item storage mirrors the history backend (Oracle or memory). The same storage powers OpenAI `/responses` and conversation APIs.
|
||||||
|
|
||||||
### History Backend (OpenAI Router Mode)
|
### History Backend (OpenAI Router Mode)
|
||||||
@@ -651,6 +652,7 @@ Store conversation and response data for tracking, debugging, or analytics.
|
|||||||
- **None**: No storage, minimal overhead.
|
- **None**: No storage, minimal overhead.
|
||||||
- **Oracle**: Persistent storage backed by Oracle Autonomous Database.
|
- **Oracle**: Persistent storage backed by Oracle Autonomous Database.
|
||||||
- **Postgres**: Persistent storage backed by PostgreSQL Database.
|
- **Postgres**: Persistent storage backed by PostgreSQL Database.
|
||||||
|
- **Redis**: Persistent storage backed by Redis.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# Memory backend (default)
|
# Memory backend (default)
|
||||||
@@ -676,6 +678,12 @@ python3 -m sglang_router.launch_router \
|
|||||||
--backend openai \
|
--backend openai \
|
||||||
--worker-urls https://api.openai.com \
|
--worker-urls https://api.openai.com \
|
||||||
--history-backend postgres
|
--history-backend postgres
|
||||||
|
|
||||||
|
# Redis backend
|
||||||
|
python3 -m sglang_router.launch_router \
|
||||||
|
--backend openai \
|
||||||
|
--worker-urls https://api.openai.com \
|
||||||
|
--history-backend redis
|
||||||
```
|
```
|
||||||
|
|
||||||
#### Oracle configuration
|
#### Oracle configuration
|
||||||
@@ -704,6 +712,19 @@ Router flags map to these values:
|
|||||||
|
|
||||||
Only one of `--oracle-dsn` or `--oracle-tns-alias` should be supplied.
|
Only one of `--oracle-dsn` or `--oracle-tns-alias` should be supplied.
|
||||||
|
|
||||||
|
#### Redis configuration
|
||||||
|
Provide Redis connection URL and optional pool sizing:
|
||||||
|
```bash
|
||||||
|
export REDIS_URL="redis://localhost:6379"
|
||||||
|
export REDIS_POOL_MAX=16
|
||||||
|
export REDIS_RETENTION_DAYS=30
|
||||||
|
```
|
||||||
|
|
||||||
|
Router flags map to these values:
|
||||||
|
- `--redis-url` (env: `REDIS_URL`)
|
||||||
|
- `--redis-pool-max` (env: `REDIS_POOL_MAX`)
|
||||||
|
- `--redis-retention-days` (env: `REDIS_RETENTION_DAYS`). Set to `-1` for persistent storage (default: 30 days).
|
||||||
|
|
||||||
## Reliability & Flow Control
|
## Reliability & Flow Control
|
||||||
- **Retries**: Default max retries = 5 with exponential backoff (`--retry-max-retries`, `--retry-initial-backoff-ms`, `--retry-max-backoff-ms`, `--retry-backoff-multiplier`, `--retry-jitter-factor`). Retries trigger on 408/429/500/502/503/504.
|
- **Retries**: Default max retries = 5 with exponential backoff (`--retry-max-retries`, `--retry-initial-backoff-ms`, `--retry-max-backoff-ms`, `--retry-backoff-multiplier`, `--retry-jitter-factor`). Retries trigger on 408/429/500/502/503/504.
|
||||||
- **Circuit Breakers**: Per worker thresholds (`--cb-failure-threshold`, `--cb-success-threshold`, `--cb-timeout-duration-secs`, `--cb-window-duration-secs`). Disable via `--disable-circuit-breaker`.
|
- **Circuit Breakers**: Per worker thresholds (`--cb-failure-threshold`, `--cb-success-threshold`, `--cb-timeout-duration-secs`, `--cb-window-duration-secs`). Disable via `--disable-circuit-breaker`.
|
||||||
|
|||||||
@@ -31,6 +31,7 @@ pub enum HistoryBackendType {
|
|||||||
None,
|
None,
|
||||||
Oracle,
|
Oracle,
|
||||||
Postgres,
|
Postgres,
|
||||||
|
Redis,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[pyclass(eq)]
|
#[pyclass(eq)]
|
||||||
@@ -270,6 +271,40 @@ impl PyOracleConfig {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[pyclass]
|
||||||
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
pub struct PyRedisConfig {
|
||||||
|
#[pyo3(get, set)]
|
||||||
|
pub url: String,
|
||||||
|
#[pyo3(get, set)]
|
||||||
|
pub pool_max: usize,
|
||||||
|
#[pyo3(get, set)]
|
||||||
|
pub retention_days: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[pymethods]
|
||||||
|
impl PyRedisConfig {
|
||||||
|
#[new]
|
||||||
|
#[pyo3(signature = (url, pool_max = 16, retention_days = Some(30)))]
|
||||||
|
fn new(url: String, pool_max: usize, retention_days: Option<u64>) -> PyResult<Self> {
|
||||||
|
Ok(PyRedisConfig {
|
||||||
|
url,
|
||||||
|
pool_max,
|
||||||
|
retention_days,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PyRedisConfig {
|
||||||
|
pub fn to_config_redis(&self) -> config::RedisConfig {
|
||||||
|
config::RedisConfig {
|
||||||
|
url: self.url.clone(),
|
||||||
|
pool_max: self.pool_max,
|
||||||
|
retention_days: self.retention_days,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[pyclass]
|
#[pyclass]
|
||||||
#[derive(Debug, Clone, PartialEq)]
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
pub struct PyPostgresConfig {
|
pub struct PyPostgresConfig {
|
||||||
@@ -374,6 +409,7 @@ struct Router {
|
|||||||
history_backend: HistoryBackendType,
|
history_backend: HistoryBackendType,
|
||||||
oracle_config: Option<PyOracleConfig>,
|
oracle_config: Option<PyOracleConfig>,
|
||||||
postgres_config: Option<PyPostgresConfig>,
|
postgres_config: Option<PyPostgresConfig>,
|
||||||
|
redis_config: Option<PyRedisConfig>,
|
||||||
client_cert_path: Option<String>,
|
client_cert_path: Option<String>,
|
||||||
client_key_path: Option<String>,
|
client_key_path: Option<String>,
|
||||||
ca_cert_paths: Vec<String>,
|
ca_cert_paths: Vec<String>,
|
||||||
@@ -486,6 +522,7 @@ impl Router {
|
|||||||
HistoryBackendType::None => config::HistoryBackend::None,
|
HistoryBackendType::None => config::HistoryBackend::None,
|
||||||
HistoryBackendType::Oracle => config::HistoryBackend::Oracle,
|
HistoryBackendType::Oracle => config::HistoryBackend::Oracle,
|
||||||
HistoryBackendType::Postgres => config::HistoryBackend::Postgres,
|
HistoryBackendType::Postgres => config::HistoryBackend::Postgres,
|
||||||
|
HistoryBackendType::Redis => config::HistoryBackend::Redis,
|
||||||
};
|
};
|
||||||
|
|
||||||
let oracle = if matches!(self.history_backend, HistoryBackendType::Oracle) {
|
let oracle = if matches!(self.history_backend, HistoryBackendType::Oracle) {
|
||||||
@@ -504,6 +541,14 @@ impl Router {
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
|
let redis_config = if matches!(self.history_backend, HistoryBackendType::Redis) {
|
||||||
|
self.redis_config
|
||||||
|
.as_ref()
|
||||||
|
.map(|cfg| cfg.to_config_redis())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
config::RouterConfig::builder()
|
config::RouterConfig::builder()
|
||||||
.mode(mode)
|
.mode(mode)
|
||||||
.policy(policy)
|
.policy(policy)
|
||||||
@@ -558,6 +603,7 @@ impl Router {
|
|||||||
.maybe_chat_template(self.chat_template.as_ref())
|
.maybe_chat_template(self.chat_template.as_ref())
|
||||||
.maybe_oracle(oracle)
|
.maybe_oracle(oracle)
|
||||||
.maybe_postgres(postgres_config)
|
.maybe_postgres(postgres_config)
|
||||||
|
.maybe_redis(redis_config)
|
||||||
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
|
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
|
||||||
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
|
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
|
||||||
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
|
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
|
||||||
@@ -654,6 +700,7 @@ impl Router {
|
|||||||
history_backend = HistoryBackendType::Memory,
|
history_backend = HistoryBackendType::Memory,
|
||||||
oracle_config = None,
|
oracle_config = None,
|
||||||
postgres_config = None,
|
postgres_config = None,
|
||||||
|
redis_config = None,
|
||||||
client_cert_path = None,
|
client_cert_path = None,
|
||||||
client_key_path = None,
|
client_key_path = None,
|
||||||
ca_cert_paths = vec![],
|
ca_cert_paths = vec![],
|
||||||
@@ -737,6 +784,7 @@ impl Router {
|
|||||||
history_backend: HistoryBackendType,
|
history_backend: HistoryBackendType,
|
||||||
oracle_config: Option<PyOracleConfig>,
|
oracle_config: Option<PyOracleConfig>,
|
||||||
postgres_config: Option<PyPostgresConfig>,
|
postgres_config: Option<PyPostgresConfig>,
|
||||||
|
redis_config: Option<PyRedisConfig>,
|
||||||
client_cert_path: Option<String>,
|
client_cert_path: Option<String>,
|
||||||
client_key_path: Option<String>,
|
client_key_path: Option<String>,
|
||||||
ca_cert_paths: Vec<String>,
|
ca_cert_paths: Vec<String>,
|
||||||
@@ -834,6 +882,7 @@ impl Router {
|
|||||||
history_backend,
|
history_backend,
|
||||||
oracle_config,
|
oracle_config,
|
||||||
postgres_config,
|
postgres_config,
|
||||||
|
redis_config,
|
||||||
client_cert_path,
|
client_cert_path,
|
||||||
client_key_path,
|
client_key_path,
|
||||||
ca_cert_paths,
|
ca_cert_paths,
|
||||||
@@ -946,6 +995,7 @@ fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|||||||
m.add_class::<PyControlPlaneAuthConfig>()?;
|
m.add_class::<PyControlPlaneAuthConfig>()?;
|
||||||
m.add_class::<PyOracleConfig>()?;
|
m.add_class::<PyOracleConfig>()?;
|
||||||
m.add_class::<PyPostgresConfig>()?;
|
m.add_class::<PyPostgresConfig>()?;
|
||||||
|
m.add_class::<PyRedisConfig>()?;
|
||||||
m.add_class::<Router>()?;
|
m.add_class::<Router>()?;
|
||||||
m.add_function(wrap_pyfunction!(get_version_string, m)?)?;
|
m.add_function(wrap_pyfunction!(get_version_string, m)?)?;
|
||||||
m.add_function(wrap_pyfunction!(get_verbose_version_string, m)?)?;
|
m.add_function(wrap_pyfunction!(get_verbose_version_string, m)?)?;
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ from sglang_router.sglang_router_rs import (
|
|||||||
PyJwtConfig,
|
PyJwtConfig,
|
||||||
PyOracleConfig,
|
PyOracleConfig,
|
||||||
PyPostgresConfig,
|
PyPostgresConfig,
|
||||||
|
PyRedisConfig,
|
||||||
PyRole,
|
PyRole,
|
||||||
)
|
)
|
||||||
from sglang_router.sglang_router_rs import Router as _Router
|
from sglang_router.sglang_router_rs import Router as _Router
|
||||||
@@ -63,6 +64,8 @@ def history_backend_from_str(backend_str: Optional[str]) -> HistoryBackendType:
|
|||||||
return HistoryBackendType.Oracle
|
return HistoryBackendType.Oracle
|
||||||
elif backend_lower == "postgres":
|
elif backend_lower == "postgres":
|
||||||
return HistoryBackendType.Postgres
|
return HistoryBackendType.Postgres
|
||||||
|
elif backend_lower == "redis":
|
||||||
|
return HistoryBackendType.Redis
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Unknown history backend: {backend_str}")
|
raise ValueError(f"Unknown history backend: {backend_str}")
|
||||||
|
|
||||||
@@ -262,6 +265,20 @@ class Router:
|
|||||||
)
|
)
|
||||||
args_dict["postgres_config"] = postgres_config
|
args_dict["postgres_config"] = postgres_config
|
||||||
|
|
||||||
|
# Convert Redis config if needed
|
||||||
|
redis_config = None
|
||||||
|
if history_backend == HistoryBackendType.Redis:
|
||||||
|
retention_days = args_dict.get("redis_retention_days", 30)
|
||||||
|
# If retention_days is negative, it means persistent storage (None in Rust)
|
||||||
|
retention_arg = None if retention_days < 0 else retention_days
|
||||||
|
|
||||||
|
redis_config = PyRedisConfig(
|
||||||
|
url=args_dict.get("redis_url"),
|
||||||
|
pool_max=args_dict.get("redis_pool_max", 16),
|
||||||
|
retention_days=retention_arg,
|
||||||
|
)
|
||||||
|
args_dict["redis_config"] = redis_config
|
||||||
|
|
||||||
# Build control plane auth config
|
# Build control plane auth config
|
||||||
args_dict["control_plane_auth"] = build_control_plane_auth_config(args_dict)
|
args_dict["control_plane_auth"] = build_control_plane_auth_config(args_dict)
|
||||||
|
|
||||||
@@ -278,6 +295,9 @@ class Router:
|
|||||||
"oracle_pool_timeout_secs",
|
"oracle_pool_timeout_secs",
|
||||||
"postgres_db_url",
|
"postgres_db_url",
|
||||||
"postgres_pool_max",
|
"postgres_pool_max",
|
||||||
|
"redis_url",
|
||||||
|
"redis_pool_max",
|
||||||
|
"redis_retention_days",
|
||||||
# Control plane auth fields (converted to control_plane_auth)
|
# Control plane auth fields (converted to control_plane_auth)
|
||||||
"control_plane_api_keys",
|
"control_plane_api_keys",
|
||||||
"control_plane_audit_enabled",
|
"control_plane_audit_enabled",
|
||||||
|
|||||||
@@ -117,6 +117,9 @@ class RouterArgs:
|
|||||||
oracle_pool_timeout_secs: int = 30
|
oracle_pool_timeout_secs: int = 30
|
||||||
postgres_db_url: Optional[str] = None
|
postgres_db_url: Optional[str] = None
|
||||||
postgres_pool_max: int = 16
|
postgres_pool_max: int = 16
|
||||||
|
redis_url: Optional[str] = None
|
||||||
|
redis_pool_max: int = 16
|
||||||
|
redis_retention_days: int = 30
|
||||||
# mTLS configuration for worker communication
|
# mTLS configuration for worker communication
|
||||||
client_cert_path: Optional[str] = None
|
client_cert_path: Optional[str] = None
|
||||||
client_key_path: Optional[str] = None
|
client_key_path: Optional[str] = None
|
||||||
@@ -200,6 +203,9 @@ class RouterArgs:
|
|||||||
postgres_group = parser.add_argument_group(
|
postgres_group = parser.add_argument_group(
|
||||||
"PostgreSQL Database", "PostgreSQL database backend configuration"
|
"PostgreSQL Database", "PostgreSQL database backend configuration"
|
||||||
)
|
)
|
||||||
|
redis_group = parser.add_argument_group(
|
||||||
|
"Redis Database", "Redis database backend configuration"
|
||||||
|
)
|
||||||
tls_group = parser.add_argument_group(
|
tls_group = parser.add_argument_group(
|
||||||
"TLS/mTLS Security", "TLS certificates for server and worker communication"
|
"TLS/mTLS Security", "TLS certificates for server and worker communication"
|
||||||
)
|
)
|
||||||
@@ -663,7 +669,7 @@ class RouterArgs:
|
|||||||
f"--{prefix}history-backend",
|
f"--{prefix}history-backend",
|
||||||
type=str,
|
type=str,
|
||||||
default=RouterArgs.history_backend,
|
default=RouterArgs.history_backend,
|
||||||
choices=["memory", "none", "oracle", "postgres"],
|
choices=["memory", "none", "oracle", "postgres", "redis"],
|
||||||
help="History storage backend for conversations and responses (default: memory)",
|
help="History storage backend for conversations and responses (default: memory)",
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -733,6 +739,28 @@ class RouterArgs:
|
|||||||
help="Maximum PostgreSQL connection pool size (default: 16, env: POSTGRES_POOL_MAX)",
|
help="Maximum PostgreSQL connection pool size (default: 16, env: POSTGRES_POOL_MAX)",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Redis configuration
|
||||||
|
redis_group.add_argument(
|
||||||
|
f"--{prefix}redis-url",
|
||||||
|
type=str,
|
||||||
|
default=os.getenv("REDIS_URL"),
|
||||||
|
help="Redis connection URL (env: REDIS_URL)",
|
||||||
|
)
|
||||||
|
redis_group.add_argument(
|
||||||
|
f"--{prefix}redis-pool-max",
|
||||||
|
type=int,
|
||||||
|
default=int(os.getenv("REDIS_POOL_MAX", RouterArgs.redis_pool_max)),
|
||||||
|
help="Maximum Redis connection pool size (default: 16, env: REDIS_POOL_MAX)",
|
||||||
|
)
|
||||||
|
redis_group.add_argument(
|
||||||
|
f"--{prefix}redis-retention-days",
|
||||||
|
type=int,
|
||||||
|
default=int(
|
||||||
|
os.getenv("REDIS_RETENTION_DAYS", RouterArgs.redis_retention_days)
|
||||||
|
),
|
||||||
|
help="Redis data retention in days (-1 for persistent, default: 30, env: REDIS_RETENTION_DAYS)",
|
||||||
|
)
|
||||||
|
|
||||||
# TLS/mTLS configuration
|
# TLS/mTLS configuration
|
||||||
tls_group.add_argument(
|
tls_group.add_argument(
|
||||||
f"--{prefix}client-cert-path",
|
f"--{prefix}client-cert-path",
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
use super::{
|
use super::{
|
||||||
CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig,
|
CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig,
|
||||||
HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RetryConfig,
|
HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RedisConfig,
|
||||||
RouterConfig, RoutingMode, TokenizerCacheConfig, TraceConfig,
|
RetryConfig, RouterConfig, RoutingMode, TokenizerCacheConfig, TraceConfig,
|
||||||
};
|
};
|
||||||
use crate::{core::ConnectionMode, mcp::McpConfig};
|
use crate::{core::ConnectionMode, mcp::McpConfig};
|
||||||
|
|
||||||
@@ -395,6 +395,12 @@ impl RouterConfigBuilder {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn redis_history(mut self, redis_config: RedisConfig) -> Self {
|
||||||
|
self.config.history_backend = HistoryBackend::Redis;
|
||||||
|
self.config.redis = Some(redis_config);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
// ==================== Parsers ====================
|
// ==================== Parsers ====================
|
||||||
|
|
||||||
pub fn reasoning_parser<S: Into<String>>(mut self, parser: S) -> Self {
|
pub fn reasoning_parser<S: Into<String>>(mut self, parser: S) -> Self {
|
||||||
@@ -539,6 +545,14 @@ impl RouterConfigBuilder {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn maybe_redis(mut self, redis: Option<RedisConfig>) -> Self {
|
||||||
|
if let Some(cfg) = redis {
|
||||||
|
self.config.history_backend = HistoryBackend::Redis;
|
||||||
|
self.config.redis = Some(cfg);
|
||||||
|
}
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
pub fn maybe_reasoning_parser(mut self, parser: Option<impl Into<String>>) -> Self {
|
pub fn maybe_reasoning_parser(mut self, parser: Option<impl Into<String>>) -> Self {
|
||||||
self.config.reasoning_parser = parser.map(|p| p.into());
|
self.config.reasoning_parser = parser.map(|p| p.into());
|
||||||
self
|
self
|
||||||
|
|||||||
@@ -58,6 +58,9 @@ pub struct RouterConfig {
|
|||||||
/// Required when history_backend = "postgres"
|
/// Required when history_backend = "postgres"
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub postgres: Option<PostgresConfig>,
|
pub postgres: Option<PostgresConfig>,
|
||||||
|
/// Required when history_backend = "redis"
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub redis: Option<RedisConfig>,
|
||||||
/// For reasoning models (e.g., deepseek-r1, qwen3)
|
/// For reasoning models (e.g., deepseek-r1, qwen3)
|
||||||
pub reasoning_parser: Option<String>,
|
pub reasoning_parser: Option<String>,
|
||||||
/// For tool-call interactions
|
/// For tool-call interactions
|
||||||
@@ -138,6 +141,7 @@ pub enum HistoryBackend {
|
|||||||
None,
|
None,
|
||||||
Oracle,
|
Oracle,
|
||||||
Postgres,
|
Postgres,
|
||||||
|
Redis,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Oracle history backend configuration
|
/// Oracle history backend configuration
|
||||||
@@ -245,6 +249,53 @@ impl PostgresConfig {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[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<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_redis_pool_max() -> usize {
|
||||||
|
16
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_redis_retention_days() -> Option<u64> {
|
||||||
|
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
|
/// Routing mode configuration
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
#[serde(tag = "type")]
|
#[serde(tag = "type")]
|
||||||
@@ -579,6 +630,7 @@ impl Default for RouterConfig {
|
|||||||
history_backend: default_history_backend(),
|
history_backend: default_history_backend(),
|
||||||
oracle: None,
|
oracle: None,
|
||||||
postgres: None,
|
postgres: None,
|
||||||
|
redis: None,
|
||||||
reasoning_parser: None,
|
reasoning_parser: None,
|
||||||
tool_call_parser: None,
|
tool_call_parser: None,
|
||||||
tokenizer_cache: TokenizerCacheConfig::default(),
|
tokenizer_cache: TokenizerCacheConfig::default(),
|
||||||
|
|||||||
@@ -14,11 +14,17 @@ use super::{
|
|||||||
oracle::{OracleConversationItemStorage, OracleConversationStorage, OracleResponseStorage},
|
oracle::{OracleConversationItemStorage, OracleConversationStorage, OracleResponseStorage},
|
||||||
};
|
};
|
||||||
use crate::{
|
use crate::{
|
||||||
config::{HistoryBackend, OracleConfig, PostgresConfig, RouterConfig},
|
config::{HistoryBackend, OracleConfig, PostgresConfig, RedisConfig, RouterConfig},
|
||||||
data_connector::postgres::{
|
data_connector::{
|
||||||
|
postgres::{
|
||||||
PostgresConversationItemStorage, PostgresConversationStorage, PostgresResponseStorage,
|
PostgresConversationItemStorage, PostgresConversationStorage, PostgresResponseStorage,
|
||||||
PostgresStore,
|
PostgresStore,
|
||||||
},
|
},
|
||||||
|
redis::{
|
||||||
|
RedisConversationItemStorage, RedisConversationStorage, RedisResponseStorage,
|
||||||
|
RedisStore,
|
||||||
|
},
|
||||||
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Type alias for the storage tuple returned by factory functions.
|
/// Type alias for the storage tuple returned by factory functions.
|
||||||
@@ -98,6 +104,33 @@ pub fn create_storage(config: &RouterConfig) -> Result<StorageTuple, String> {
|
|||||||
|
|
||||||
info!("Data connector initialized successfully: Postgres");
|
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(_) => "<redacted>".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)
|
Ok(storages)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -136,3 +169,16 @@ fn create_postgres_storage(postgres_cfg: &PostgresConfig) -> Result<StorageTuple
|
|||||||
Arc::new(postgres_item),
|
Arc::new(postgres_item),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn create_redis_storage(redis_cfg: &RedisConfig) -> Result<StorageTuple, String> {
|
||||||
|
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),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ mod memory;
|
|||||||
mod noop;
|
mod noop;
|
||||||
mod oracle;
|
mod oracle;
|
||||||
mod postgres;
|
mod postgres;
|
||||||
|
mod redis;
|
||||||
|
|
||||||
pub use core::{
|
pub use core::{
|
||||||
Conversation, ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage,
|
Conversation, ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage,
|
||||||
|
|||||||
@@ -0,0 +1,913 @@
|
|||||||
|
//! 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<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RedisStore {
|
||||||
|
pub fn new(config: RedisConfig) -> Result<Self, String> {
|
||||||
|
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<String>,
|
||||||
|
) -> Result<Option<ConversationMetadata>, 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::<ConversationMetadata>(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<Conversation, ConversationStorageError> {
|
||||||
|
let conversation = Conversation::new(input);
|
||||||
|
let id_str = conversation.id.0.as_str();
|
||||||
|
let created_at: DateTime<Utc> = 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<Option<Conversation>, 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<String>) = 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<ConversationMetadata>,
|
||||||
|
) -> Result<Option<Conversation>, 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<bool> {
|
||||||
|
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<ConversationItem, ConversationItemStorageError> {
|
||||||
|
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<Utc>,
|
||||||
|
) -> 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<Vec<ConversationItem>> {
|
||||||
|
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<f64> = 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<String> = 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::<ConversationItem>::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<std::collections::HashMap<String, String>> = pipe
|
||||||
|
.query_async(&mut conn)
|
||||||
|
.await
|
||||||
|
.map_err(|e| ConversationItemStorageError::StorageError(e.to_string()))?;
|
||||||
|
|
||||||
|
let mut items: Vec<ConversationItem> = 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<Option<ConversationItem>> {
|
||||||
|
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<String, String> = 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<bool> {
|
||||||
|
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<f64> = 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<ResponseId, ResponseStorageError> {
|
||||||
|
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<Option<StoredResponse>, 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<String, String> = 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<String> = 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<usize>,
|
||||||
|
) -> ResponseResult<ResponseChain> {
|
||||||
|
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<usize>,
|
||||||
|
) -> ResponseResult<Vec<StoredResponse>> {
|
||||||
|
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<String> = conn
|
||||||
|
.zrevrange(&key, 0, stop)
|
||||||
|
.await
|
||||||
|
.map_err(|e| ResponseStorageError::StorageError(e.to_string()))?;
|
||||||
|
|
||||||
|
if response_ids.is_empty() {
|
||||||
|
return Ok(Vec::<StoredResponse>::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut pipe = redis::pipe();
|
||||||
|
for id in &response_ids {
|
||||||
|
pipe.hgetall(Self::response_key(id));
|
||||||
|
}
|
||||||
|
|
||||||
|
let results: Vec<std::collections::HashMap<String, String>> = pipe
|
||||||
|
.query_async(&mut conn)
|
||||||
|
.await
|
||||||
|
.map_err(|e| ResponseStorageError::StorageError(e.to_string()))?;
|
||||||
|
|
||||||
|
let mut out: Vec<StoredResponse> = 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<usize> {
|
||||||
|
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<String> = 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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -5,8 +5,8 @@ use smg::{
|
|||||||
auth::{ApiKeyEntry, ControlPlaneAuthConfig, JwtConfig, Role},
|
auth::{ApiKeyEntry, ControlPlaneAuthConfig, JwtConfig, Role},
|
||||||
config::{
|
config::{
|
||||||
CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig,
|
CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig,
|
||||||
HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RetryConfig,
|
HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RedisConfig,
|
||||||
RouterConfig, RoutingMode, TokenizerCacheConfig, TraceConfig,
|
RetryConfig, RouterConfig, RoutingMode, TokenizerCacheConfig, TraceConfig,
|
||||||
},
|
},
|
||||||
core::ConnectionMode,
|
core::ConnectionMode,
|
||||||
observability::{
|
observability::{
|
||||||
@@ -419,7 +419,7 @@ struct CliArgs {
|
|||||||
backend: Backend,
|
backend: Backend,
|
||||||
|
|
||||||
/// History storage backend
|
/// History storage backend
|
||||||
#[arg(long, default_value = "memory", value_parser = ["memory", "none", "oracle","postgres"], help_heading = "Backend")]
|
#[arg(long, default_value = "memory", value_parser = ["memory", "none", "oracle", "postgres", "redis"], help_heading = "Backend")]
|
||||||
history_backend: String,
|
history_backend: String,
|
||||||
|
|
||||||
/// Enable WebAssembly support
|
/// Enable WebAssembly support
|
||||||
@@ -468,6 +468,19 @@ struct CliArgs {
|
|||||||
#[arg(long, help_heading = "PostgreSQL Database")]
|
#[arg(long, help_heading = "PostgreSQL Database")]
|
||||||
postgres_pool_max_size: Option<usize>,
|
postgres_pool_max_size: Option<usize>,
|
||||||
|
|
||||||
|
// ==================== Redis Database ====================
|
||||||
|
/// Redis connection URL
|
||||||
|
#[arg(long, help_heading = "Redis Database")]
|
||||||
|
redis_url: Option<String>,
|
||||||
|
|
||||||
|
/// Maximum Redis connection pool size
|
||||||
|
#[arg(long, help_heading = "Redis Database")]
|
||||||
|
redis_pool_max_size: Option<usize>,
|
||||||
|
|
||||||
|
/// Redis data retention in days (-1 for persistent, default 30)
|
||||||
|
#[arg(long, help_heading = "Redis Database")]
|
||||||
|
redis_retention_days: Option<i64>,
|
||||||
|
|
||||||
// ==================== TLS/mTLS Security ====================
|
// ==================== TLS/mTLS Security ====================
|
||||||
/// Path to server TLS certificate (PEM format)
|
/// Path to server TLS certificate (PEM format)
|
||||||
#[arg(long, help_heading = "TLS/mTLS Security")]
|
#[arg(long, help_heading = "TLS/mTLS Security")]
|
||||||
@@ -793,6 +806,27 @@ impl CliArgs {
|
|||||||
Ok(pcf)
|
Ok(pcf)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn build_redis_config(&self) -> ConfigResult<RedisConfig> {
|
||||||
|
let url = self.redis_url.clone().unwrap_or_default();
|
||||||
|
let pool_max = self.redis_pool_max_size.unwrap_or(16);
|
||||||
|
|
||||||
|
let retention_days = match self.redis_retention_days {
|
||||||
|
Some(d) if d < 0 => None, // Persistent
|
||||||
|
Some(d) => Some(d as u64),
|
||||||
|
None => Some(30), // Default 30 days
|
||||||
|
};
|
||||||
|
|
||||||
|
let rcf = RedisConfig {
|
||||||
|
url,
|
||||||
|
pool_max,
|
||||||
|
retention_days,
|
||||||
|
};
|
||||||
|
rcf.validate().map_err(|e| ConfigError::ValidationFailed {
|
||||||
|
reason: e.to_string(),
|
||||||
|
})?;
|
||||||
|
Ok(rcf)
|
||||||
|
}
|
||||||
|
|
||||||
fn to_router_config(
|
fn to_router_config(
|
||||||
&self,
|
&self,
|
||||||
prefill_urls: Vec<(String, Option<u16>)>,
|
prefill_urls: Vec<(String, Option<u16>)>,
|
||||||
@@ -869,6 +903,7 @@ impl CliArgs {
|
|||||||
"none" => HistoryBackend::None,
|
"none" => HistoryBackend::None,
|
||||||
"oracle" => HistoryBackend::Oracle,
|
"oracle" => HistoryBackend::Oracle,
|
||||||
"postgres" => HistoryBackend::Postgres,
|
"postgres" => HistoryBackend::Postgres,
|
||||||
|
"redis" => HistoryBackend::Redis,
|
||||||
_ => HistoryBackend::Memory,
|
_ => HistoryBackend::Memory,
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -882,6 +917,11 @@ impl CliArgs {
|
|||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
let redis = if history_backend == HistoryBackend::Redis {
|
||||||
|
Some(self.build_redis_config()?)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
let builder = RouterConfig::builder()
|
let builder = RouterConfig::builder()
|
||||||
.mode(mode)
|
.mode(mode)
|
||||||
@@ -939,6 +979,7 @@ impl CliArgs {
|
|||||||
.maybe_chat_template(self.chat_template.as_ref())
|
.maybe_chat_template(self.chat_template.as_ref())
|
||||||
.maybe_oracle(oracle)
|
.maybe_oracle(oracle)
|
||||||
.maybe_postgres(postgres)
|
.maybe_postgres(postgres)
|
||||||
|
.maybe_redis(redis)
|
||||||
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
|
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
|
||||||
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
|
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
|
||||||
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
|
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
|
||||||
|
|||||||
Reference in New Issue
Block a user