use serde::{Deserialize, Serialize};
use crate::error::{ContragError, Result};
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ContragConfig {
pub entities: Vec<EntityConfig>,
pub embedder: EmbedderConfigDef,
pub chunking: ChunkingConfig,
pub vector_store: VectorStoreConfig,
pub system_prompt: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EntityConfig {
pub name: String,
pub canister_id: String,
pub fetch_method: String,
pub fetch_many_method: Option<String>,
pub relationships: Vec<RelationshipConfig>,
pub auto_include: bool,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RelationshipConfig {
pub field_name: String,
pub target_entity: String,
pub relationship_type: String, }
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EmbedderConfigDef {
pub provider: String,
pub model: String,
pub dimensions: usize,
pub api_endpoint: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ChunkingConfig {
pub chunk_size: usize,
pub overlap: usize,
pub include_field_names: bool,
}
impl Default for ChunkingConfig {
fn default() -> Self {
Self {
chunk_size: 1000,
overlap: 100,
include_field_names: true,
}
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct VectorStoreConfig {
pub storage_type: String,
pub max_hot_vectors: Option<usize>,
pub enable_cache: bool,
}
impl Default for VectorStoreConfig {
fn default() -> Self {
Self {
storage_type: "stable_memory".to_string(),
max_hot_vectors: Some(10000),
enable_cache: true,
}
}
}
#[derive(Clone, Debug)]
pub struct EnvVars {
pub openai_api_key: Option<String>,
pub gemini_api_key: Option<String>,
}
pub fn load_config_from_json(json_str: &str) -> Result<ContragConfig> {
serde_json::from_str(json_str)
.map_err(|e| ContragError::ConfigError(format!("Failed to parse config: {}", e)))
}
pub fn validate_config(config: &ContragConfig) -> Result<()> {
if config.entities.is_empty() {
return Err(ContragError::InvalidConfig(
"At least one entity must be configured".to_string(),
));
}
if config.embedder.dimensions == 0 {
return Err(ContragError::InvalidConfig(
"Embedder dimensions must be greater than 0".to_string(),
));
}
if config.chunking.chunk_size == 0 {
return Err(ContragError::InvalidConfig(
"Chunk size must be greater than 0".to_string(),
));
}
if config.chunking.overlap >= config.chunking.chunk_size {
return Err(ContragError::InvalidConfig(
"Overlap must be less than chunk size".to_string(),
));
}
for entity in &config.entities {
if entity.name.is_empty() {
return Err(ContragError::InvalidConfig(
"Entity name cannot be empty".to_string(),
));
}
if entity.canister_id.is_empty() {
return Err(ContragError::InvalidConfig(
format!("Canister ID for entity '{}' cannot be empty", entity.name),
));
}
}
Ok(())
}
pub fn create_default_config() -> ContragConfig {
ContragConfig {
entities: vec![],
embedder: EmbedderConfigDef {
provider: "openai".to_string(),
model: "text-embedding-3-small".to_string(),
dimensions: 1536,
api_endpoint: None,
},
chunking: ChunkingConfig::default(),
vector_store: VectorStoreConfig::default(),
system_prompt: None,
}
}
pub fn load_config(config_json: &str, env_vars: EnvVars) -> Result<ContragConfig> {
let config = load_config_from_json(config_json)?;
validate_config(&config)?;
Ok(config)
}