cognee_http_server/dto/
llm.rs

1//! DTOs for `/api/v1/llm`.
2//!
3//! Mirrors Python's
4//! [`get_llm_router.py`](https://github.com/topoteretes/cognee/blob/main/cognee/api/v1/llm/routers/get_llm_router.py)
5//! — `_ALLOWED_LLM_PARAMS` is the wire filter; anything else is silently dropped.
6//!
7//! See `docs/http-server/routers/llm.md` §4 for the per-router spec.
8
9use serde::{Deserialize, Serialize};
10use serde_json::{Map, Value};
11use utoipa::ToSchema;
12
13/// LLM kwargs the wire is allowed to forward into the underlying adapter.
14/// Mirrors Python's `_ALLOWED_LLM_PARAMS` constant in
15/// [`get_llm_router.py:20`](https://github.com/topoteretes/cognee/blob/main/cognee/api/v1/llm/routers/get_llm_router.py#L20).
16pub const ALLOWED_LLM_PARAMS: &[&str] = &["temperature", "max_tokens", "top_p", "seed"];
17
18/// Filter `parameters` against the `ALLOWED_LLM_PARAMS` whitelist.
19///
20/// Drops any key not in the allow-list silently (no error). Non-object inputs
21/// produce an empty object. Matches Python's `_safe_params()` semantics.
22pub fn safe_params(input: &Value) -> Value {
23    let mut out = Map::new();
24    if let Some(obj) = input.as_object() {
25        for (k, v) in obj {
26            if ALLOWED_LLM_PARAMS.contains(&k.as_str()) {
27                out.insert(k.clone(), v.clone());
28            }
29        }
30    }
31    Value::Object(out)
32}
33
34// ─── /custom-prompt ───────────────────────────────────────────────────────────
35
36/// Mirrors Python `CustomPromptGenerationPayloadDTO`
37/// ([`get_llm_router.py:27-32`](https://github.com/topoteretes/cognee/blob/main/cognee/api/v1/llm/routers/get_llm_router.py#L27-L32)).
38///
39/// Inherits `InDTO` in Python, so the wire is camelCase per Decision 10
40/// (`graphModel` is the canonical key; `graph_model` is accepted as an alias).
41#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)]
42#[serde(rename_all = "camelCase")]
43pub struct CustomPromptGenerationPayloadDTO {
44    /// Free-form JSON object describing the desired graph model.
45    #[serde(alias = "graph_model")]
46    pub graph_model: Value,
47
48    /// Kwargs forwarded to the LLM adapter (filtered via `safe_params`).
49    #[serde(default)]
50    pub parameters: Value,
51}
52
53/// Inherits `OutDTO` in Python — wire is camelCase (`customPrompt`).
54#[derive(Debug, Clone, Serialize, ToSchema)]
55#[serde(rename_all = "camelCase")]
56pub struct CustomPromptGenerationResponseDTO {
57    pub custom_prompt: String,
58}
59
60// ─── /infer-schema ────────────────────────────────────────────────────────────
61
62/// JSON-body adapter for `POST /api/v1/llm/infer-schema`.
63///
64/// Note: Python's endpoint takes multipart `Form(...)` fields rather than a
65/// JSON body — the Rust port uses a JSON body (acknowledged divergence). All
66/// fields here are single-word, so camelCase has no wire effect, but the
67/// attribute is added for forward consistency.
68#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)]
69#[serde(rename_all = "camelCase")]
70pub struct InferSchemaPayloadDTO {
71    /// Sample text to analyze for entity types and relationships.
72    pub text: String,
73
74    /// Same `safe_params` filter rules as `CustomPromptGenerationPayloadDTO`.
75    #[serde(default)]
76    pub parameters: Value,
77}
78
79/// Inherits `OutDTO` in Python — wire is camelCase (`graphSchema`).
80#[derive(Debug, Clone, Serialize, ToSchema)]
81#[serde(rename_all = "camelCase")]
82pub struct InferSchemaResponseDTO {
83    /// Parsed-and-validated JSON object the LLM produced.
84    pub graph_schema: Value,
85}
86
87#[cfg(test)]
88#[allow(
89    clippy::unwrap_used,
90    clippy::expect_used,
91    reason = "test code — panics are acceptable failures"
92)]
93mod tests {
94    use super::*;
95    use serde_json::json;
96
97    #[test]
98    fn test_safe_params_keeps_all_allowed_keys() {
99        let input = json!({
100            "temperature": 0.7,
101            "max_tokens": 256,
102            "top_p": 0.9,
103            "seed": 42,
104        });
105        let out = safe_params(&input);
106        assert_eq!(out["temperature"], json!(0.7));
107        assert_eq!(out["max_tokens"], json!(256));
108        assert_eq!(out["top_p"], json!(0.9));
109        assert_eq!(out["seed"], json!(42));
110    }
111
112    #[test]
113    fn test_safe_params_drops_unknown_keys() {
114        let input = json!({
115            "temperature": 0.5,
116            "junk_key": "x",
117            "model": "gpt-4o",
118            "stream": true,
119        });
120        let out = safe_params(&input);
121        assert_eq!(out["temperature"], json!(0.5));
122        assert!(out.get("junk_key").is_none());
123        assert!(out.get("model").is_none());
124        assert!(out.get("stream").is_none());
125    }
126
127    #[test]
128    fn test_safe_params_non_object_returns_empty_object() {
129        assert_eq!(safe_params(&json!(null)), json!({}));
130        assert_eq!(safe_params(&json!([1, 2, 3])), json!({}));
131        assert_eq!(safe_params(&json!("string")), json!({}));
132    }
133
134    #[test]
135    fn test_safe_params_empty_object_round_trips() {
136        assert_eq!(safe_params(&json!({})), json!({}));
137    }
138
139    #[test]
140    fn test_custom_prompt_dto_deserializes_snake_case_via_alias() {
141        let json = r#"{
142            "graph_model": {"entity_types": []},
143            "parameters": {"temperature": 0.0}
144        }"#;
145        let payload: CustomPromptGenerationPayloadDTO = serde_json::from_str(json).unwrap();
146        assert!(payload.graph_model.is_object());
147        assert_eq!(payload.parameters["temperature"], json!(0.0));
148    }
149
150    #[test]
151    fn test_custom_prompt_dto_deserializes_camelcase() {
152        let json = r#"{
153            "graphModel": {"entity_types": []},
154            "parameters": {"temperature": 0.0}
155        }"#;
156        let payload: CustomPromptGenerationPayloadDTO = serde_json::from_str(json).unwrap();
157        assert!(payload.graph_model.is_object());
158    }
159
160    #[test]
161    fn custom_prompt_response_dto_serializes_camelcase_only() {
162        let dto = CustomPromptGenerationResponseDTO {
163            custom_prompt: "p".into(),
164        };
165        let s = serde_json::to_string(&dto).expect("serialize");
166        assert!(s.contains("\"customPrompt\""), "missing customPrompt: {s}");
167        assert!(
168            !s.contains("\"custom_prompt\""),
169            "snake_case custom_prompt leaked: {s}"
170        );
171    }
172
173    #[test]
174    fn infer_schema_response_dto_serializes_camelcase_only() {
175        let dto = InferSchemaResponseDTO {
176            graph_schema: serde_json::json!({"x": 1}),
177        };
178        let s = serde_json::to_string(&dto).expect("serialize");
179        assert!(s.contains("\"graphSchema\""), "missing graphSchema: {s}");
180        assert!(
181            !s.contains("\"graph_schema\""),
182            "snake_case graph_schema leaked: {s}"
183        );
184    }
185
186    #[test]
187    fn test_infer_schema_dto_deserializes_with_default_parameters() {
188        let json = r#"{"text": "Alice met Bob."}"#;
189        let payload: InferSchemaPayloadDTO = serde_json::from_str(json).unwrap();
190        assert_eq!(payload.text, "Alice met Bob.");
191        assert!(payload.parameters.is_null());
192    }
193}
cognee_http_server/dto/llm.rs

cognee_http_server/dto/
llm.rs
