relay-knowledge 1.1.16

Graph-database-based knowledge graph project.
Documentation
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use tokio::io::AsyncWriteExt;

use super::profile::{default_connect_timeout_seconds, default_temperature, default_top_p};
use super::*;
use crate::{
    net::http::{HttpBindAddress, HttpConfig, HttpProxyConfig},
    retrieval::{
        DEFAULT_EMBEDDING_BATCH_SIZE, DEFAULT_EMBEDDING_MAX_CONCURRENCY, DEFAULT_EMBEDDING_TIMEOUT,
        EmbeddingProviderKind, ReadModelBackendConfig, ReadModelBackendMode, RemoteEmbeddingConfig,
    },
};

pub(super) fn test_service(label: &str) -> ModelProviderConfigService {
    ModelProviderConfigService::new(test_paths(label))
}

fn test_paths(label: &str) -> RuntimePaths {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("clock should be after epoch")
        .as_nanos();
    let root = std::env::temp_dir().join(format!("relay-model-provider-{label}-{now}"));
    RuntimePaths {
        config_dir: root.join("config"),
        data_dir: root.join("data"),
        state_dir: root.join("state"),
        cache_dir: root.join("cache"),
        log_dir: root.join("logs"),
        temp_dir: root.join("tmp"),
        runtime_dir: root.join("run"),
        service_dir: root.join("service"),
    }
}

pub(super) fn test_http_config() -> HttpConfig {
    HttpConfig::new(
        HttpBindAddress::parse("127.0.0.1:8791").expect("bind address"),
        Duration::from_millis(50),
        Duration::from_millis(50),
        crate::net::http::DEFAULT_MAX_BODY_BYTES,
        HttpProxyConfig::new(None, Vec::new(), true).expect("proxy config"),
    )
    .expect("http config")
}

pub(super) fn remote_retrieval() -> ReadModelBackendConfig {
    let mut config = ReadModelBackendConfig::local();
    config.semantic_mode = ReadModelBackendMode::External;
    config.vector_mode = ReadModelBackendMode::External;
    config.vector_model.name = "text-embedding-3-small".to_owned();
    config.remote_embedding = Some(RemoteEmbeddingConfig {
        provider: EmbeddingProviderKind::OpenAiCompatible,
        base_url: "https://api.openai.example/v1".to_owned(),
        api_key: "env-secret".to_owned(),
        batch_size: DEFAULT_EMBEDDING_BATCH_SIZE,
        timeout: DEFAULT_EMBEDDING_TIMEOUT,
        max_concurrency: DEFAULT_EMBEDDING_MAX_CONCURRENCY,
    });
    config
}

pub(super) async fn failing_catalog_url() -> String {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
        .await
        .expect("catalog fixture should bind");
    let address = listener.local_addr().expect("catalog fixture address");
    tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.expect("catalog request");
        stream
            .write_all(
                b"HTTP/1.1 503 Service Unavailable\r\nContent-Type: application/json\r\nContent-Length: 2\r\n\r\n{}",
            )
            .await
            .expect("catalog response");
    });
    format!("http://{address}/api.json")
}

pub(super) fn openai_request(model: &str, api_key: Option<&str>) -> ModelProfileSaveRequest {
    ModelProfileSaveRequest {
        provider: ModelProviderKind::OpenAiCompatible,
        model: model.to_owned(),
        base_url: Some("https://user:pass@api.example.com/v1".to_owned()),
        api_key: api_key.map(ToOwned::to_owned),
        clear_api_key: false,
        headers: vec![ModelRequestHeader {
            name: "x-extra-secret".to_owned(),
            value: Some("header-secret".to_owned()),
            secret: true,
            configured: false,
        }],
        ssl_verify: Some(true),
        context_window: Some(128_000),
        max_tokens: Some(4096),
        temperature: default_temperature(),
        top_p: default_top_p(),
        connect_timeout_seconds: default_connect_timeout_seconds(),
        capabilities: Some(ModelCapabilities {
            input: ModelModalityMatrix {
                text: Some(true),
                image: Some(true),
                audio: None,
                video: None,
                pdf: None,
            },
            output: ModelModalityMatrix {
                text: Some(true),
                image: None,
                audio: None,
                video: None,
                pdf: None,
            },
        }),
        fallback_policy_id: Some(" same_provider_then_other_provider ".to_owned()),
        fallback_priority: 1,
        catalog_provider_id: Some(" openai ".to_owned()),
        catalog_provider_name: Some(" OpenAI ".to_owned()),
        catalog_model_name: Some(model.to_owned()),
        is_default: true,
    }
}

pub(super) fn echo_request(model: &str, is_default: bool) -> ModelProfileSaveRequest {
    ModelProfileSaveRequest {
        provider: ModelProviderKind::Echo,
        model: model.to_owned(),
        base_url: None,
        api_key: None,
        clear_api_key: false,
        headers: Vec::new(),
        ssl_verify: None,
        context_window: None,
        max_tokens: None,
        temperature: 0.2,
        top_p: default_top_p(),
        connect_timeout_seconds: 5.0,
        capabilities: None,
        fallback_policy_id: None,
        fallback_priority: 0,
        catalog_provider_id: None,
        catalog_provider_name: None,
        catalog_model_name: None,
        is_default,
    }
}