use super::template::{
PromptComposition, PromptError, PromptResult, PromptTemplate, PromptVariable,
};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use uuid::Uuid;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptEntity {
pub id: Uuid,
pub template_id: String,
pub name: Option<String>,
pub description: Option<String>,
pub content: String,
pub variables: serde_json::Value,
pub tags: Vec<String>,
pub version: Option<String>,
pub metadata: serde_json::Value,
pub enabled: bool,
pub created_at: chrono::DateTime<chrono::Utc>,
pub updated_at: chrono::DateTime<chrono::Utc>,
pub created_by: Option<Uuid>,
pub tenant_id: Option<Uuid>,
}
impl PromptEntity {
pub fn from_template(template: &PromptTemplate) -> Self {
let now = chrono::Utc::now();
let variables = serde_json::to_value(&template.variables).unwrap_or_default();
let metadata = serde_json::to_value(&template.metadata).unwrap_or_default();
Self {
id: Uuid::now_v7(),
template_id: template.id.clone(),
name: template.name.clone(),
description: template.description.clone(),
content: template.content.clone(),
variables,
tags: template.tags.clone(),
version: template.version.clone(),
metadata,
enabled: true,
created_at: now,
updated_at: now,
created_by: None,
tenant_id: None,
}
}
pub fn to_template(&self) -> PromptResult<PromptTemplate> {
let variables: Vec<PromptVariable> = serde_json::from_value(self.variables.clone())
.map_err(|e| PromptError::ParseError(e.to_string()))?;
let metadata: HashMap<String, String> =
serde_json::from_value(self.metadata.clone()).unwrap_or_default();
Ok(PromptTemplate {
id: self.template_id.clone(),
name: self.name.clone(),
description: self.description.clone(),
content: self.content.clone(),
variables,
tags: self.tags.clone(),
version: self.version.clone(),
metadata,
})
}
pub fn with_creator(mut self, creator_id: Uuid) -> Self {
self.created_by = Some(creator_id);
self
}
pub fn with_tenant(mut self, tenant_id: Uuid) -> Self {
self.tenant_id = Some(tenant_id);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptCompositionEntity {
pub id: Uuid,
pub composition_id: String,
pub description: Option<String>,
pub template_ids: Vec<String>,
pub separator: String,
pub enabled: bool,
pub created_at: chrono::DateTime<chrono::Utc>,
pub updated_at: chrono::DateTime<chrono::Utc>,
pub tenant_id: Option<Uuid>,
}
impl PromptCompositionEntity {
pub fn from_composition(composition: &PromptComposition) -> Self {
let now = chrono::Utc::now();
Self {
id: Uuid::now_v7(),
composition_id: composition.id.clone(),
description: composition.description.clone(),
template_ids: composition.template_ids.clone(),
separator: composition.separator.clone(),
enabled: true,
created_at: now,
updated_at: now,
tenant_id: None,
}
}
pub fn to_composition(&self) -> PromptComposition {
PromptComposition {
id: self.composition_id.clone(),
description: self.description.clone(),
template_ids: self.template_ids.clone(),
separator: self.separator.clone(),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct PromptFilter {
pub template_id: Option<String>,
pub tags: Option<Vec<String>>,
pub search: Option<String>,
pub enabled_only: bool,
pub tenant_id: Option<Uuid>,
pub offset: Option<i64>,
pub limit: Option<i64>,
}
impl PromptFilter {
pub fn new() -> Self {
Self {
enabled_only: true,
..Default::default()
}
}
pub fn template_id(mut self, id: impl Into<String>) -> Self {
self.template_id = Some(id.into());
self
}
pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
self.tags.get_or_insert_with(Vec::new).push(tag.into());
self
}
pub fn with_tags(mut self, tags: Vec<String>) -> Self {
self.tags = Some(tags);
self
}
pub fn search(mut self, keyword: impl Into<String>) -> Self {
self.search = Some(keyword.into());
self
}
pub fn include_disabled(mut self) -> Self {
self.enabled_only = false;
self
}
pub fn tenant(mut self, tenant_id: Uuid) -> Self {
self.tenant_id = Some(tenant_id);
self
}
pub fn paginate(mut self, offset: i64, limit: i64) -> Self {
self.offset = Some(offset);
self.limit = Some(limit);
self
}
}
#[async_trait]
pub trait PromptStore: Send + Sync {
async fn save_template(&self, entity: &PromptEntity) -> PromptResult<()>;
async fn save_templates(&self, entities: &[PromptEntity]) -> PromptResult<()> {
for entity in entities {
self.save_template(entity).await?;
}
Ok(())
}
async fn get_template_by_id(&self, id: Uuid) -> PromptResult<Option<PromptEntity>>;
async fn get_template(&self, template_id: &str) -> PromptResult<Option<PromptEntity>>;
async fn query_templates(&self, filter: &PromptFilter) -> PromptResult<Vec<PromptEntity>>;
async fn find_by_tag(&self, tag: &str) -> PromptResult<Vec<PromptEntity>>;
async fn search_templates(&self, keyword: &str) -> PromptResult<Vec<PromptEntity>>;
async fn update_template(&self, entity: &PromptEntity) -> PromptResult<()>;
async fn delete_template_by_id(&self, id: Uuid) -> PromptResult<bool>;
async fn delete_template(&self, template_id: &str) -> PromptResult<bool>;
async fn set_template_enabled(&self, template_id: &str, enabled: bool) -> PromptResult<()>;
async fn exists(&self, template_id: &str) -> PromptResult<bool>;
async fn count(&self, filter: &PromptFilter) -> PromptResult<i64>;
async fn get_all_tags(&self) -> PromptResult<Vec<String>>;
async fn save_composition(&self, entity: &PromptCompositionEntity) -> PromptResult<()>;
async fn get_composition(
&self,
composition_id: &str,
) -> PromptResult<Option<PromptCompositionEntity>>;
async fn query_compositions(&self) -> PromptResult<Vec<PromptCompositionEntity>>;
async fn delete_composition(&self, composition_id: &str) -> PromptResult<bool>;
}
pub type DynPromptStore = std::sync::Arc<dyn PromptStore>;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_prompt_entity_from_template() {
let template = PromptTemplate::new("test")
.with_name("Test Template")
.with_content("Hello, {name}!")
.with_tag("greeting");
let entity = PromptEntity::from_template(&template);
assert_eq!(entity.template_id, "test");
assert_eq!(entity.name, Some("Test Template".to_string()));
assert!(entity.enabled);
}
#[test]
fn test_prompt_entity_to_template() {
let template = PromptTemplate::new("test")
.with_name("Test Template")
.with_content("Hello, {name}!")
.with_tag("greeting");
let entity = PromptEntity::from_template(&template);
let converted = entity.to_template().unwrap();
assert_eq!(converted.id, template.id);
assert_eq!(converted.name, template.name);
assert_eq!(converted.content, template.content);
}
#[test]
fn test_prompt_filter_builder() {
let filter = PromptFilter::new()
.with_tag("code")
.search("review")
.paginate(0, 10);
assert_eq!(filter.tags, Some(vec!["code".to_string()]));
assert_eq!(filter.search, Some("review".to_string()));
assert_eq!(filter.limit, Some(10));
}
}