use crate::{
client::engine::Engine, client::retriever::RetrieverClient, config::Config,
events::EventEmitter, metrics::MetricsHub, storage::Workspace,
};
#[derive(Debug)]
pub struct EngineBuilder {
config: Option<Config>,
events: Option<EventEmitter>,
api_key: Option<String>,
model: Option<String>,
endpoint: Option<String>,
}
impl EngineBuilder {
#[must_use]
pub fn new() -> Self {
Self {
config: None,
events: None,
api_key: None,
model: None,
endpoint: None,
}
}
#[must_use]
pub fn with_config(mut self, config: Config) -> Self {
self.config = Some(config);
self
}
#[must_use]
pub fn with_events(mut self, events: EventEmitter) -> Self {
self.events = Some(events);
self
}
#[must_use]
pub fn with_key(mut self, key: impl Into<String>) -> Self {
self.api_key = Some(key.into());
self
}
#[must_use]
pub fn with_model(mut self, model: impl Into<String>) -> Self {
self.model = Some(model.into());
self
}
#[must_use]
pub fn with_endpoint(mut self, url: impl Into<String>) -> Self {
self.endpoint = Some(url.into());
self
}
pub async fn build(self) -> Result<Engine, BuildError> {
let mut config = self.config.unwrap_or_default();
if let Some(api_key) = self.api_key {
config.llm.api_key = Some(api_key);
}
if let Some(model) = self.model {
config.llm.model = model;
}
if let Some(endpoint) = self.endpoint {
config.llm.endpoint = Some(endpoint);
}
if config.llm.api_key.is_none() {
return Err(BuildError::MissingApiKey);
}
if config.llm.model.is_empty() {
return Err(BuildError::MissingModel);
}
if config.llm.endpoint.is_none() {
return Err(BuildError::MissingEndpoint);
}
let workspace = Workspace::new(&config.storage.workspace_dir)
.await
.map_err(|e| BuildError::Workspace(e.to_string()))?;
let metrics_hub = std::sync::Arc::new(MetricsHub::with_defaults());
let pool = crate::llm::LlmPool::from_config(&config.llm, Some(metrics_hub.clone()));
let indexer = crate::client::indexer::IndexerClient::with_llm(pool.index().clone());
let retriever = RetrieverClient::new(pool.retrieval().clone());
let events = self.events.unwrap_or_default();
Engine::with_components(config, workspace, retriever, indexer, events, metrics_hub)
.await
.map_err(|e| BuildError::Other(e.to_string()))
}
}
impl Default for EngineBuilder {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, thiserror::Error)]
pub enum BuildError {
#[error("Workspace error: {0}")]
Workspace(String),
#[error("Missing API key: call .with_key(\"sk-...\") or set api_key in config")]
MissingApiKey,
#[error("Missing model: call .with_model(\"gpt-4o\") or set model in config")]
MissingModel,
#[error(
"Missing endpoint: call .with_endpoint(\"https://api.xxx.com/v1\") or set endpoint in config"
)]
MissingEndpoint,
#[error("{0}")]
Other(String),
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_builder_with_key() {
let builder = EngineBuilder::new().with_key("sk-test-key");
assert_eq!(builder.api_key, Some("sk-test-key".to_string()));
}
#[test]
fn test_builder_with_model() {
let builder = EngineBuilder::new().with_model("gpt-4o-mini");
assert_eq!(builder.model, Some("gpt-4o-mini".to_string()));
}
#[test]
fn test_builder_with_key_and_model() {
let builder = EngineBuilder::new()
.with_model("gpt-4o-mini")
.with_key("sk-test");
assert_eq!(builder.model, Some("gpt-4o-mini".to_string()));
assert_eq!(builder.api_key, Some("sk-test".to_string()));
}
}