Spaces:
Running
Running
| """ | |
| Prompt Templates Module | |
| Manages prompt templates for different AI tools and providers | |
| Includes system for storing, loading, and managing user-created prompts | |
| """ | |
| import os | |
| import json | |
| import uuid | |
| import logging | |
| from typing import Dict, Any, List, Optional | |
| from datetime import datetime | |
| # Setup logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger("prompts") | |
| class PromptTemplate: | |
| """Represents a prompt template that can be used with AI providers""" | |
| def __init__(self, | |
| id: str = None, | |
| name: str = "", | |
| description: str = "", | |
| template: str = "", | |
| system_message: str = "", | |
| category: str = "", | |
| is_public: bool = False, | |
| created_by: str = "system", | |
| created_at: str = None, | |
| price: float = 0.0, | |
| parameters: Dict[str, Any] = None, | |
| provider_defaults: Dict[str, Any] = None): | |
| """Initialize a prompt template""" | |
| self.id = id or str(uuid.uuid4()) | |
| self.name = name | |
| self.description = description | |
| self.template = template | |
| self.system_message = system_message | |
| self.category = category | |
| self.is_public = is_public | |
| self.created_by = created_by | |
| self.created_at = created_at or datetime.now().isoformat() | |
| self.price = price | |
| self.parameters = parameters or {} | |
| self.provider_defaults = provider_defaults or {} | |
| def to_dict(self) -> Dict[str, Any]: | |
| """Convert template to dictionary""" | |
| return { | |
| "id": self.id, | |
| "name": self.name, | |
| "description": self.description, | |
| "template": self.template, | |
| "system_message": self.system_message, | |
| "category": self.category, | |
| "is_public": self.is_public, | |
| "created_by": self.created_by, | |
| "created_at": self.created_at, | |
| "price": self.price, | |
| "parameters": self.parameters, | |
| "provider_defaults": self.provider_defaults | |
| } | |
| def from_dict(cls, data: Dict[str, Any]) -> 'PromptTemplate': | |
| """Create template from dictionary""" | |
| return cls( | |
| id=data.get("id"), | |
| name=data.get("name", ""), | |
| description=data.get("description", ""), | |
| template=data.get("template", ""), | |
| system_message=data.get("system_message", ""), | |
| category=data.get("category", ""), | |
| is_public=data.get("is_public", False), | |
| created_by=data.get("created_by", "system"), | |
| created_at=data.get("created_at"), | |
| price=data.get("price", 0.0), | |
| parameters=data.get("parameters", {}), | |
| provider_defaults=data.get("provider_defaults", {}) | |
| ) | |
| def render(self, variables: Dict[str, Any] = None) -> Dict[str, str]: | |
| """ | |
| Render the prompt template with the provided variables | |
| Returns both the rendered prompt and system message | |
| """ | |
| variables = variables or {} | |
| prompt = self.template | |
| system = self.system_message | |
| # Replace variables in the template | |
| for key, value in variables.items(): | |
| placeholder = "{" + key + "}" | |
| prompt = prompt.replace(placeholder, str(value)) | |
| system = system.replace(placeholder, str(value)) | |
| return { | |
| "prompt": prompt, | |
| "system_message": system | |
| } | |
| class PromptTemplateManager: | |
| """Manages prompt templates, including default and user-created ones""" | |
| def __init__(self, templates_dir: str = None): | |
| """Initialize the prompt template manager""" | |
| self.templates_dir = templates_dir or os.path.join(os.path.dirname(__file__), "templates") | |
| self.user_templates_dir = os.path.join(self.templates_dir, "user") | |
| # Create directories if they don't exist | |
| os.makedirs(self.templates_dir, exist_ok=True) | |
| os.makedirs(self.user_templates_dir, exist_ok=True) | |
| # Load default templates | |
| self.default_templates = self._load_default_templates() | |
| # Load user templates | |
| self.user_templates = self._load_user_templates() | |
| def _load_default_templates(self) -> Dict[str, PromptTemplate]: | |
| """Load default templates from files""" | |
| templates = {} | |
| # Define default templates if no files exist yet | |
| default_templates = { | |
| "general_chat": PromptTemplate( | |
| id="general_chat", | |
| name="General Chat", | |
| description="A general-purpose chat prompt", | |
| template="Please answer the following question or respond to the message: {input}", | |
| system_message="You are a helpful assistant that provides accurate and concise responses.", | |
| category="general", | |
| created_by="system" | |
| ), | |
| "creative_writing": PromptTemplate( | |
| id="creative_writing", | |
| name="Creative Writing", | |
| description="Generate creative writing based on a premise", | |
| template="Write a {genre} {format} about {topic}.", | |
| system_message="You are a creative writer with expertise in different genres and formats.", | |
| category="creative", | |
| created_by="system", | |
| parameters={ | |
| "genre": {"type": "string", "description": "Genre of the writing", "default": "science fiction"}, | |
| "format": {"type": "string", "description": "Format of the writing", "default": "short story"}, | |
| "topic": {"type": "string", "description": "Topic or premise", "required": True} | |
| } | |
| ), | |
| "code_assistant": PromptTemplate( | |
| id="code_assistant", | |
| name="Code Assistant", | |
| description="Generate or debug code in various languages", | |
| template="I need help with the following code task in {language}:\n\n{task}\n\n{code}", | |
| system_message="You are an expert programmer. Provide well-commented, efficient, and correct code solutions.", | |
| category="development", | |
| created_by="system", | |
| parameters={ | |
| "language": {"type": "string", "description": "Programming language", "required": True}, | |
| "task": {"type": "string", "description": "Description of the coding task", "required": True}, | |
| "code": {"type": "string", "description": "Existing code (if any)", "default": ""} | |
| }, | |
| provider_defaults={ | |
| "openai": {"model": "gpt-4-turbo"}, | |
| "deepseek": {"model": "deepseek-coder"} | |
| } | |
| ), | |
| "image_prompt": PromptTemplate( | |
| id="image_prompt", | |
| name="Image Generation", | |
| description="Detailed prompt for image generation", | |
| template="{subject} {style}, {details}, {quality}", | |
| system_message="", | |
| category="images", | |
| created_by="system", | |
| parameters={ | |
| "subject": {"type": "string", "description": "Main subject of the image", "required": True}, | |
| "style": {"type": "string", "description": "Art style", "default": "digital art"}, | |
| "details": {"type": "string", "description": "Additional details", "default": "detailed, vibrant colors"}, | |
| "quality": {"type": "string", "description": "Quality descriptors", "default": "high quality, 4k, trending on artstation"} | |
| }, | |
| provider_defaults={ | |
| "openai": {"model": "dall-e-3"} | |
| } | |
| ) | |
| } | |
| # Save default templates if they don't exist | |
| for template_id, template in default_templates.items(): | |
| template_path = os.path.join(self.templates_dir, f"{template_id}.json") | |
| if not os.path.exists(template_path): | |
| with open(template_path, "w") as f: | |
| json.dump(template.to_dict(), f, indent=2) | |
| templates[template_id] = template | |
| return templates | |
| def _load_user_templates(self) -> Dict[str, PromptTemplate]: | |
| """Load user-created templates""" | |
| templates = {} | |
| try: | |
| # List all JSON files in user_templates_dir | |
| for filename in os.listdir(self.user_templates_dir): | |
| if filename.endswith(".json"): | |
| template_path = os.path.join(self.user_templates_dir, filename) | |
| with open(template_path, "r") as f: | |
| template_data = json.load(f) | |
| template = PromptTemplate.from_dict(template_data) | |
| templates[template.id] = template | |
| except Exception as e: | |
| logger.error(f"Error loading user templates: {e}") | |
| return templates | |
| def get_all_templates(self) -> List[PromptTemplate]: | |
| """Get all available templates (default + user)""" | |
| all_templates = list(self.default_templates.values()) | |
| # Add public user templates and user's own templates | |
| for template in self.user_templates.values(): | |
| if template.is_public: | |
| all_templates.append(template) | |
| return all_templates | |
| def get_user_templates(self, user_id: str) -> List[PromptTemplate]: | |
| """Get templates created by a specific user""" | |
| return [t for t in self.user_templates.values() if t.created_by == user_id] | |
| def get_template(self, template_id: str) -> Optional[PromptTemplate]: | |
| """Get a specific template by ID""" | |
| if template_id in self.default_templates: | |
| return self.default_templates[template_id] | |
| if template_id in self.user_templates: | |
| return self.user_templates[template_id] | |
| return None | |
| def create_template(self, template: PromptTemplate) -> PromptTemplate: | |
| """Create a new user template""" | |
| # Ensure unique ID | |
| if template.id in self.default_templates or template.id in self.user_templates: | |
| template.id = str(uuid.uuid4()) | |
| # Save template to file | |
| template_path = os.path.join(self.user_templates_dir, f"{template.id}.json") | |
| with open(template_path, "w") as f: | |
| json.dump(template.to_dict(), f, indent=2) | |
| # Add to user templates | |
| self.user_templates[template.id] = template | |
| return template | |
| def update_template(self, template: PromptTemplate) -> Optional[PromptTemplate]: | |
| """Update an existing user template""" | |
| if template.id not in self.user_templates: | |
| logger.error(f"Template {template.id} not found or not a user template") | |
| return None | |
| # Save updated template | |
| template_path = os.path.join(self.user_templates_dir, f"{template.id}.json") | |
| with open(template_path, "w") as f: | |
| json.dump(template.to_dict(), f, indent=2) | |
| # Update in memory | |
| self.user_templates[template.id] = template | |
| return template | |
| def delete_template(self, template_id: str, user_id: str) -> bool: | |
| """Delete a user template""" | |
| if template_id not in self.user_templates: | |
| logger.error(f"Template {template_id} not found or not a user template") | |
| return False | |
| # Check ownership | |
| template = self.user_templates[template_id] | |
| if template.created_by != user_id and user_id != "admin": | |
| logger.error(f"User {user_id} does not own template {template_id}") | |
| return False | |
| # Delete template file | |
| template_path = os.path.join(self.user_templates_dir, f"{template_id}.json") | |
| try: | |
| os.remove(template_path) | |
| del self.user_templates[template_id] | |
| return True | |
| except Exception as e: | |
| logger.error(f"Error deleting template {template_id}: {e}") | |
| return False | |
| def get_templates_by_category(self, category: str) -> List[PromptTemplate]: | |
| """Get templates by category""" | |
| all_templates = self.get_all_templates() | |
| return [t for t in all_templates if t.category.lower() == category.lower()] | |
| def get_public_templates(self) -> List[PromptTemplate]: | |
| """Get all public templates created by users""" | |
| return [t for t in self.user_templates.values() if t.is_public] | |
| # Example usage | |
| if __name__ == "__main__": | |
| # Initialize manager | |
| manager = PromptTemplateManager() | |
| # Get all templates | |
| all_templates = manager.get_all_templates() | |
| print(f"Loaded {len(all_templates)} templates") | |
| # Create a new template | |
| new_template = PromptTemplate( | |
| name="SEO Content", | |
| description="Generate SEO-optimized content for websites", | |
| template="Write SEO-optimized content about {topic} targeting the keyword {keyword}. {tone} tone, {length} words.", | |
| system_message="You are an expert SEO content writer who creates engaging, well-researched content that ranks well in search engines.", | |
| category="marketing", | |
| created_by="user123", | |
| is_public=True, | |
| parameters={ | |
| "topic": {"type": "string", "description": "Main topic", "required": True}, | |
| "keyword": {"type": "string", "description": "Target keyword", "required": True}, | |
| "tone": {"type": "string", "description": "Content tone", "default": "professional"}, | |
| "length": {"type": "number", "description": "Content length in words", "default": 500} | |
| } | |
| ) | |
| created = manager.create_template(new_template) | |
| print(f"Created new template: {created.id} - {created.name}") | |
| # Test rendering a template | |
| code_template = manager.get_template("code_assistant") | |
| if code_template: | |
| rendered = code_template.render({ | |
| "language": "Python", | |
| "task": "Create a function to calculate Fibonacci numbers", | |
| "code": "def fibonacci(n):\n # TODO: Implement" | |
| }) | |
| print("\nRendered prompt:") | |
| print(rendered["prompt"]) | |
| print("\nSystem message:") | |