cognee_llm/prompts/
mod.rs

1//! Compiled-in LLM prompt templates ported from Python.
2//!
3//! Mirrors [`cognee/infrastructure/llm/prompts/`](https://github.com/topoteretes/cognee/tree/main/cognee/infrastructure/llm/prompts).
4//! Filenames are kept identical for cross-SDK diffing.
5//!
6//! Templates use `{{KEY}}` placeholders. Only flat substitution is supported —
7//! the Python prompts use a tiny subset of Jinja that this is compatible with.
8
9use std::collections::HashMap;
10
11use thiserror::Error;
12
13// ─── Compiled-in prompt files ─────────────────────────────────────────────────
14
15pub const CUSTOM_PROMPT_GENERATION_USER: &str = include_str!("custom_prompt_generation_user.txt");
16pub const CUSTOM_PROMPT_GENERATION_SYSTEM: &str =
17    include_str!("custom_prompt_generation_system.txt");
18pub const INFER_SCHEMA_USER: &str = include_str!("infer_schema_user.txt");
19pub const INFER_SCHEMA_SYSTEM: &str = include_str!("infer_schema_system.txt");
20
21// ─── PromptError ──────────────────────────────────────────────────────────────
22
23#[derive(Debug, Error)]
24pub enum PromptError {
25    #[error("unknown prompt template: {0}")]
26    UnknownTemplate(String),
27}
28
29// ─── render_prompt ────────────────────────────────────────────────────────────
30
31/// Render the named prompt template with the given context map.
32///
33/// Returns the prompt string with every `{{KEY}}` placeholder replaced by the
34/// corresponding `ctx[key]` value. Keys not in `ctx` are left as-is so a
35/// missing context entry surfaces visibly in the rendered prompt rather than
36/// silently swallowed.
37///
38/// Recognised template names: `custom_prompt_generation_user`,
39/// `custom_prompt_generation_system`, `infer_schema_user`, `infer_schema_system`.
40pub fn render_prompt(name: &str, ctx: &HashMap<&str, &str>) -> Result<String, PromptError> {
41    let raw = match name {
42        "custom_prompt_generation_user" => CUSTOM_PROMPT_GENERATION_USER,
43        "custom_prompt_generation_system" => CUSTOM_PROMPT_GENERATION_SYSTEM,
44        "infer_schema_user" => INFER_SCHEMA_USER,
45        "infer_schema_system" => INFER_SCHEMA_SYSTEM,
46        other => return Err(PromptError::UnknownTemplate(other.to_string())),
47    };
48
49    let mut out = raw.to_string();
50    for (key, value) in ctx {
51        let needle = format!("{{{{{key}}}}}");
52        out = out.replace(&needle, value);
53    }
54    Ok(out)
55}
56
57#[cfg(test)]
58mod tests {
59    #![allow(
60        clippy::unwrap_used,
61        clippy::expect_used,
62        reason = "test code — panics are acceptable"
63    )]
64    use super::*;
65
66    #[test]
67    fn test_custom_prompt_user_renders_graph_schema() {
68        let mut ctx = HashMap::new();
69        ctx.insert("GRAPH_SCHEMA_JSON", "{\"entity_types\": []}");
70        let out = render_prompt("custom_prompt_generation_user", &ctx).expect("render");
71        assert!(out.contains("{\"entity_types\": []}"));
72        assert!(!out.contains("{{GRAPH_SCHEMA_JSON}}"));
73    }
74
75    #[test]
76    fn test_custom_prompt_system_loads() {
77        let ctx = HashMap::new();
78        let out = render_prompt("custom_prompt_generation_system", &ctx).expect("render");
79        assert!(!out.is_empty());
80    }
81
82    #[test]
83    fn test_infer_schema_user_renders_sample_text() {
84        let mut ctx = HashMap::new();
85        ctx.insert("SAMPLE_TEXT", "Alice met Bob.");
86        let out = render_prompt("infer_schema_user", &ctx).expect("render");
87        assert!(out.contains("Alice met Bob."));
88        assert!(!out.contains("{{SAMPLE_TEXT}}"));
89    }
90
91    #[test]
92    fn test_infer_schema_system_loads() {
93        let ctx = HashMap::new();
94        let out = render_prompt("infer_schema_system", &ctx).expect("render");
95        assert!(!out.is_empty());
96    }
97
98    #[test]
99    fn test_unknown_template_errors() {
100        let ctx = HashMap::new();
101        let err = render_prompt("does_not_exist", &ctx).unwrap_err();
102        match err {
103            PromptError::UnknownTemplate(name) => assert_eq!(name, "does_not_exist"),
104        }
105    }
106}
cognee_llm/prompts/mod.rs

cognee_llm/prompts/
mod.rs
