Skip to main content

cognee_http_server/dto/
settings.rs

1//! DTOs for `/api/v1/settings/*` per `routers/settings.md §4`.
2
3use serde::{Deserialize, Serialize};
4use std::collections::BTreeMap;
5use utoipa::ToSchema;
6
7// ── Selectable provider/model lists ────────────────────────────────────────
8
9#[derive(Debug, Clone, Serialize, ToSchema)]
10#[serde(rename_all = "camelCase")]
11pub struct ConfigChoice {
12    pub value: String,
13    pub label: String,
14}
15
16// ── GET response ────────────────────────────────────────────────────────────
17
18#[derive(Debug, Clone, Serialize, ToSchema)]
19#[serde(rename_all = "camelCase")]
20pub struct LLMConfigOutputDTO {
21    pub provider: String,
22    pub model: String,
23    pub endpoint: Option<String>,
24    pub api_version: Option<String>,
25    pub api_key: Option<String>,
26    pub providers: Vec<ConfigChoice>,
27    pub models: BTreeMap<String, Vec<ConfigChoice>>,
28}
29
30#[derive(Debug, Clone, Serialize, ToSchema)]
31#[serde(rename_all = "camelCase")]
32pub struct VectorDBConfigOutputDTO {
33    pub provider: String,
34    pub url: String,
35    pub api_key: String,
36    pub providers: Vec<ConfigChoice>,
37}
38
39#[derive(Debug, Clone, Serialize, ToSchema)]
40#[serde(rename_all = "camelCase")]
41pub struct SettingsDTO {
42    pub llm: LLMConfigOutputDTO,
43    pub vector_db: VectorDBConfigOutputDTO,
44}
45
46// ── POST request body ───────────────────────────────────────────────────────
47
48#[derive(Debug, Clone, Deserialize, ToSchema)]
49#[serde(rename_all = "camelCase")]
50pub struct LLMConfigInputDTO {
51    pub provider: LlmProvider,
52    pub model: String,
53    #[serde(alias = "api_key")]
54    pub api_key: String,
55}
56
57/// Provider enum for `LLMConfigInputDTO::provider`. Note that `bedrock`
58/// is **not** in this list — Python's GET advertises it but the save
59/// `Literal` rejects it (`routers/settings.md §6.4`).
60#[derive(Debug, Clone, Copy, Deserialize, Serialize, ToSchema, PartialEq, Eq)]
61#[serde(rename_all = "lowercase")]
62pub enum LlmProvider {
63    Openai,
64    Ollama,
65    Anthropic,
66    Gemini,
67    Mistral,
68}
69
70#[derive(Debug, Clone, Deserialize, ToSchema)]
71#[serde(rename_all = "camelCase")]
72pub struct VectorDBConfigInputDTO {
73    pub provider: VectorDbProvider,
74    pub url: String,
75    #[serde(alias = "api_key")]
76    pub api_key: String,
77}
78
79#[derive(Debug, Clone, Copy, Deserialize, Serialize, ToSchema, PartialEq, Eq)]
80#[serde(rename_all = "lowercase")]
81pub enum VectorDbProvider {
82    Lancedb,
83    Chromadb,
84    Pgvector,
85    #[serde(rename = "brute-force", alias = "brute_force", alias = "bruteforce")]
86    BruteForce,
87}
88
89#[derive(Debug, Clone, Deserialize, ToSchema)]
90#[serde(rename_all = "camelCase")]
91pub struct SettingsPayloadDTO {
92    #[serde(default)]
93    pub llm: Option<LLMConfigInputDTO>,
94    #[serde(default, alias = "vector_db")]
95    pub vector_db: Option<VectorDBConfigInputDTO>,
96}
97
98// ── Helpers ────────────────────────────────────────────────────────────────
99
100/// Mirrors Python's `(key[0:10] + "*" * (len(key) - 10)) if key else None`.
101///
102/// - `None` / empty → `None` (Python returns the empty-key short-circuit).
103/// - Up to 10 chars → return the key as-is (no stars).
104/// - Longer → first 10 chars + `(len - 10)` stars.
105pub fn redact_api_key(key: Option<&str>) -> Option<String> {
106    let key = key.filter(|k| !k.is_empty())?;
107    let len = key.len();
108    if len <= 10 {
109        // No stars; return as-is.
110        return Some(key.to_string());
111    }
112    let mut head = String::with_capacity(len);
113    head.push_str(&key[..10]);
114    head.push_str(&"*".repeat(len - 10));
115    Some(head)
116}
117
118/// Mirrors Python's `'*****' not in key and len(key.strip()) > 0` guard.
119pub fn should_persist_api_key(submitted: &str) -> bool {
120    !submitted.contains("*****") && !submitted.trim().is_empty()
121}
122
123#[cfg(test)]
124#[allow(
125    clippy::unwrap_used,
126    clippy::expect_used,
127    reason = "test code — panics are acceptable failures"
128)]
129mod tests {
130    use super::*;
131
132    #[test]
133    fn redact_empty_returns_none() {
134        assert_eq!(redact_api_key(None), None);
135        assert_eq!(redact_api_key(Some("")), None);
136    }
137
138    #[test]
139    fn redact_short_key_returns_as_is() {
140        assert_eq!(redact_api_key(Some("short")), Some("short".into()));
141    }
142
143    #[test]
144    fn redact_long_key_masks_tail() {
145        let r = redact_api_key(Some("sk-1234567890ABC")).expect("some");
146        // 10 chars + 6 stars
147        assert_eq!(r, "sk-1234567******");
148    }
149
150    #[test]
151    fn should_persist_rejects_empty() {
152        assert!(!should_persist_api_key(""));
153        assert!(!should_persist_api_key("   "));
154    }
155
156    #[test]
157    fn should_persist_rejects_redacted() {
158        assert!(!should_persist_api_key("sk-prefix*****abc"));
159        assert!(!should_persist_api_key("AAAAAAAAAA*****"));
160    }
161
162    #[test]
163    fn should_persist_accepts_real_key() {
164        assert!(should_persist_api_key("sk-real-key"));
165    }
166}