Skip to main content

split_brain_harness/backends/
mod.rs

1use crate::types::{BackendType, Config};
2use async_trait::async_trait;
3use reqwest::Client;
4
5pub mod anthropic;
6pub mod embedded;
7pub mod ollama;
8pub mod openai;
9
10#[async_trait]
11pub trait InferenceEngine: Send + Sync {
12    async fn generate(&self, system_prompt: &str, prompt_payload: &str) -> Result<String, String>;
13}
14
15pub fn init_engine(config: &Config) -> Box<dyn InferenceEngine> {
16    let client = Client::builder()
17        .pool_max_idle_per_host(10)
18        .timeout(std::time::Duration::from_secs(config.timeout_secs))
19        .build()
20        .unwrap_or_default();
21
22    match config.backend {
23        BackendType::OpenAiCompat => Box::new(openai::OpenAiEngine {
24            endpoint: config.endpoint.clone(),
25            model: config.model_name.clone(),
26            client,
27        }),
28        BackendType::OllamaNative => Box::new(ollama::OllamaNativeEngine {
29            endpoint: config.endpoint.clone(),
30            model: config.model_name.clone(),
31            client,
32        }),
33        BackendType::LocalEmbedded => Box::new(embedded::LocalEmbeddedEngine {
34            model_identifier: config.model_name.clone(),
35        }),
36        BackendType::Anthropic => Box::new(anthropic::AnthropicEngine {
37            endpoint: config.endpoint.clone(),
38            model: config.model_name.clone(),
39            api_key: config.api_key.clone().unwrap_or_default(),
40            client,
41        }),
42    }
43}