use crate::prompt_mfg::{PromptError, Result};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptIR {
pub sections: BTreeMap<String, Section>,
pub metadata: PromptMetadata,
pub variables: BTreeMap<String, PromptVariable>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Section {
pub section_type: SectionType,
pub blocks: Vec<ContentBlock>,
pub priority: i32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SectionType {
System,
User,
Assistant,
Custom(String),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ContentBlock {
pub block_type: BlockType,
pub content: String,
pub metadata: BTreeMap<String, String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BlockType {
Text,
Code { language: String },
Instruction,
Example,
Constraint,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptMetadata {
pub id: String,
pub version: String,
pub schema_version: String,
pub source_ontology: String,
pub construct_query: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptVariable {
pub name: String,
pub var_type: VariableType,
pub default: Option<String>,
pub description: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum VariableType {
String,
Integer,
Boolean,
List,
Object,
}
impl PromptIR {
pub fn from_construct(construct_query: &str) -> Result<Self> {
Ok(Self {
sections: BTreeMap::new(),
metadata: PromptMetadata {
id: "generated".to_string(),
version: "0.1.0".to_string(),
schema_version: "1.0.0".to_string(),
source_ontology: "unknown".to_string(),
construct_query: construct_query.to_string(),
},
variables: BTreeMap::new(),
})
}
pub fn from_store(store: &oxigraph::store::Store, construct_query: &str) -> Result<Self> {
use oxigraph::sparql::QueryResults;
#[allow(deprecated)]
let results = store
.query(construct_query)
.map_err(|e| PromptError::Sparql(e.to_string()))?;
match results {
QueryResults::Graph(triples) => {
let mut sections = BTreeMap::new();
let variables = BTreeMap::new();
for triple_result in triples {
let triple = triple_result.map_err(|e| PromptError::Sparql(e.to_string()))?;
let section_key = triple.subject.to_string();
sections.entry(section_key).or_insert_with(|| Section {
section_type: SectionType::System,
blocks: vec![],
priority: 0,
});
}
Ok(Self {
sections,
metadata: PromptMetadata {
id: "from_store".to_string(),
version: "0.1.0".to_string(),
schema_version: "1.0.0".to_string(),
source_ontology: "rdf_store".to_string(),
construct_query: construct_query.to_string(),
},
variables,
})
}
_ => Err(PromptError::Sparql(
"Expected CONSTRUCT query results".to_string(),
)),
}
}
pub fn add_section(&mut self, key: String, section: Section) {
self.sections.insert(key, section);
}
pub fn add_variable(&mut self, key: String, variable: PromptVariable) {
self.variables.insert(key, variable);
}
pub fn sections_ordered(&self) -> Vec<(&String, &Section)> {
let mut sections: Vec<_> = self.sections.iter().collect();
sections.sort_by_key(|(_, s)| s.priority);
sections
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_prompt_ir_creation() {
let ir = PromptIR::from_construct("CONSTRUCT { ?s ?p ?o } WHERE { ?s ?p ?o }");
assert!(ir.is_ok());
}
#[test]
fn test_section_ordering() {
let mut ir = PromptIR::from_construct("test").unwrap();
ir.add_section(
"section1".to_string(),
Section {
section_type: SectionType::System,
blocks: vec![],
priority: 10,
},
);
ir.add_section(
"section2".to_string(),
Section {
section_type: SectionType::User,
blocks: vec![],
priority: 5,
},
);
let ordered = ir.sections_ordered();
assert_eq!(ordered.len(), 2);
assert_eq!(ordered[0].0, "section2"); assert_eq!(ordered[1].0, "section1");
}
#[test]
fn test_btreemap_determinism() {
let mut ir1 = PromptIR::from_construct("test").unwrap();
ir1.variables.insert(
"z_var".to_string(),
PromptVariable {
name: "z_var".to_string(),
var_type: VariableType::String,
default: None,
description: "Test".to_string(),
},
);
ir1.variables.insert(
"a_var".to_string(),
PromptVariable {
name: "a_var".to_string(),
var_type: VariableType::String,
default: None,
description: "Test".to_string(),
},
);
let mut ir2 = PromptIR::from_construct("test").unwrap();
ir2.variables.insert(
"a_var".to_string(),
PromptVariable {
name: "a_var".to_string(),
var_type: VariableType::String,
default: None,
description: "Test".to_string(),
},
);
ir2.variables.insert(
"z_var".to_string(),
PromptVariable {
name: "z_var".to_string(),
var_type: VariableType::String,
default: None,
description: "Test".to_string(),
},
);
assert_eq!(ir1, ir2);
}
}