feat: add inline external generation mode
This commit is contained in:
@@ -22,6 +22,7 @@ pub struct AppConfig {
|
||||
pub listen_backlog: i32,
|
||||
pub worker_threads: Option<usize>,
|
||||
pub process_role: ProcessRole,
|
||||
pub external_generation_mode: ExternalGenerationMode,
|
||||
pub external_generation_worker_id: String,
|
||||
pub external_generation_worker_concurrency: usize,
|
||||
pub external_generation_worker_poll_interval: Duration,
|
||||
@@ -171,6 +172,25 @@ pub enum ProcessRole {
|
||||
All,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum ExternalGenerationMode {
|
||||
Inline,
|
||||
Queue,
|
||||
}
|
||||
|
||||
impl ExternalGenerationMode {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Inline => "inline",
|
||||
Self::Queue => "queue",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_inline(self) -> bool {
|
||||
matches!(self, Self::Inline)
|
||||
}
|
||||
}
|
||||
|
||||
impl ProcessRole {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
@@ -197,6 +217,7 @@ impl Default for AppConfig {
|
||||
listen_backlog: 1024,
|
||||
worker_threads: None,
|
||||
process_role: ProcessRole::Api,
|
||||
external_generation_mode: ExternalGenerationMode::Queue,
|
||||
external_generation_worker_id: default_external_generation_worker_id(),
|
||||
external_generation_worker_concurrency: 2,
|
||||
external_generation_worker_poll_interval: Duration::from_millis(2_000),
|
||||
@@ -385,6 +406,11 @@ impl AppConfig {
|
||||
if let Some(process_role) = read_first_process_role_env(&["GENARRATIVE_PROCESS_ROLE"]) {
|
||||
config.process_role = process_role;
|
||||
}
|
||||
if let Some(external_generation_mode) =
|
||||
read_first_external_generation_mode_env(&["GENARRATIVE_EXTERNAL_GENERATION_MODE"])
|
||||
{
|
||||
config.external_generation_mode = external_generation_mode;
|
||||
}
|
||||
if let Some(worker_id) =
|
||||
read_first_non_empty_env(&["GENARRATIVE_EXTERNAL_GENERATION_WORKER_ID"])
|
||||
{
|
||||
@@ -1046,6 +1072,14 @@ fn read_first_process_role_env(keys: &[&str]) -> Option<ProcessRole> {
|
||||
})
|
||||
}
|
||||
|
||||
fn read_first_external_generation_mode_env(keys: &[&str]) -> Option<ExternalGenerationMode> {
|
||||
keys.iter().find_map(|key| {
|
||||
env::var(key)
|
||||
.ok()
|
||||
.and_then(|value| parse_external_generation_mode(&value))
|
||||
})
|
||||
}
|
||||
|
||||
fn read_first_positive_u32_env(keys: &[&str]) -> Option<u32> {
|
||||
keys.iter().find_map(|key| {
|
||||
env::var(key)
|
||||
@@ -1111,6 +1145,16 @@ fn parse_process_role(value: &str) -> Option<ProcessRole> {
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_external_generation_mode(value: &str) -> Option<ExternalGenerationMode> {
|
||||
match trim_quoted_env_value(value).to_ascii_lowercase().as_str() {
|
||||
"inline" | "sync" | "synchronous" => Some(ExternalGenerationMode::Inline),
|
||||
"queue" | "queued" | "worker" | "async" | "asynchronous" => {
|
||||
Some(ExternalGenerationMode::Queue)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn trim_quoted_env_value(raw: &str) -> &str {
|
||||
let raw = raw.trim();
|
||||
raw.strip_prefix('"')
|
||||
@@ -1243,8 +1287,8 @@ fn parse_positive_u16(raw: &str) -> Option<u16> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
AppConfig, DEFAULT_VECTOR_ENGINE_IMAGE_REQUEST_TIMEOUT_MS, LlmProvider, ProcessRole,
|
||||
parse_bool, parse_process_role,
|
||||
AppConfig, DEFAULT_VECTOR_ENGINE_IMAGE_REQUEST_TIMEOUT_MS, ExternalGenerationMode,
|
||||
LlmProvider, ProcessRole, parse_bool, parse_external_generation_mode, parse_process_role,
|
||||
};
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
|
||||
@@ -1312,6 +1356,51 @@ mod tests {
|
||||
assert!(ProcessRole::All.runs_external_generation_worker());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn external_generation_mode_parses_inline_and_queue_aliases() {
|
||||
assert_eq!(
|
||||
parse_external_generation_mode("inline"),
|
||||
Some(ExternalGenerationMode::Inline)
|
||||
);
|
||||
assert_eq!(
|
||||
parse_external_generation_mode("'sync'"),
|
||||
Some(ExternalGenerationMode::Inline)
|
||||
);
|
||||
assert_eq!(
|
||||
parse_external_generation_mode("\"queue\""),
|
||||
Some(ExternalGenerationMode::Queue)
|
||||
);
|
||||
assert_eq!(
|
||||
parse_external_generation_mode("worker"),
|
||||
Some(ExternalGenerationMode::Queue)
|
||||
);
|
||||
assert_eq!(parse_external_generation_mode("unknown"), None);
|
||||
|
||||
assert!(ExternalGenerationMode::Inline.is_inline());
|
||||
assert!(!ExternalGenerationMode::Queue.is_inline());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_env_reads_external_generation_mode() {
|
||||
let _guard = ENV_LOCK
|
||||
.get_or_init(|| Mutex::new(()))
|
||||
.lock()
|
||||
.expect("env lock");
|
||||
unsafe {
|
||||
std::env::set_var("GENARRATIVE_EXTERNAL_GENERATION_MODE", "inline");
|
||||
}
|
||||
|
||||
let config = AppConfig::from_env();
|
||||
|
||||
assert_eq!(
|
||||
config.external_generation_mode,
|
||||
ExternalGenerationMode::Inline
|
||||
);
|
||||
unsafe {
|
||||
std::env::remove_var("GENARRATIVE_EXTERNAL_GENERATION_MODE");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_env_reads_sms_enabled_when_shell_value_keeps_quotes() {
|
||||
let _guard = ENV_LOCK
|
||||
|
||||
Reference in New Issue
Block a user