Spaces:
Sleeping
Sleeping
| """ | |
| Prompt Integration Module | |
| This module provides utilities for integrating shared components into existing prompts | |
| while maintaining backward compatibility with the current prompt system. | |
| """ | |
| from typing import Dict, List, Optional, Any | |
| from .prompt_controller import PromptController | |
| from .data_models import Indicator, Rule, Template, IndicatorCategory | |
| class PromptIntegrator: | |
| """Integrates shared components with existing prompt system.""" | |
| def __init__(self): | |
| self.controller = PromptController() | |
| def generate_indicators_section(self, category_filter: Optional[IndicatorCategory] = None) -> str: | |
| """ | |
| Generate indicators section for prompt files. | |
| Args: | |
| category_filter: Optional filter to include only specific category indicators | |
| Returns: | |
| Formatted indicators section for inclusion in prompts | |
| """ | |
| if category_filter: | |
| indicators = self.controller.indicator_catalog.get_indicators_by_category(category_filter) | |
| else: | |
| indicators = self.controller.indicator_catalog.get_all_indicators() | |
| if not indicators: | |
| return "" | |
| # Group indicators by category | |
| by_category = {} | |
| for indicator in indicators: | |
| cat_name = indicator.category.value | |
| if cat_name not in by_category: | |
| by_category[cat_name] = [] | |
| by_category[cat_name].append(indicator) | |
| # Generate formatted section | |
| sections = [] | |
| for category, cat_indicators in by_category.items(): | |
| section_lines = [f"<{category}_indicators>"] | |
| for indicator in cat_indicators: | |
| section_lines.append(f"- {indicator.definition}") | |
| if indicator.examples: | |
| example_text = ", ".join(f'"{ex}"' for ex in indicator.examples[:3]) | |
| section_lines.append(f" Examples: {example_text}") | |
| section_lines.append(f"</{category}_indicators>") | |
| sections.append("\n".join(section_lines)) | |
| return "\n\n".join(sections) | |
| def generate_rules_section(self, action_filter: Optional[str] = None) -> str: | |
| """ | |
| Generate rules section for prompt files. | |
| Args: | |
| action_filter: Optional filter to include only rules with specific actions | |
| Returns: | |
| Formatted rules section for inclusion in prompts | |
| """ | |
| if action_filter: | |
| rules = self.controller.rules_catalog.get_rules_by_action(action_filter) | |
| else: | |
| rules = self.controller.rules_catalog.get_rules_by_priority() | |
| if not rules: | |
| return "" | |
| section_lines = ["<critical_rules>"] | |
| for i, rule in enumerate(rules, 1): | |
| section_lines.append(f"{i}. {rule.description}") | |
| if rule.examples: | |
| example_text = ", ".join(f'"{ex}"' for ex in rule.examples[:2]) | |
| section_lines.append(f" Examples: {example_text}") | |
| section_lines.append("</critical_rules>") | |
| return "\n".join(section_lines) | |
| def generate_categories_section(self) -> str: | |
| """ | |
| Generate categories section for prompt files. | |
| Returns: | |
| Formatted categories section for inclusion in prompts | |
| """ | |
| categories = self.controller.category_definitions.get_all_categories() | |
| if not categories: | |
| return "" | |
| section_lines = ["<classification_categories>"] | |
| section_lines.append("You must classify this message into exactly ONE of the following three categories:") | |
| section_lines.append("") | |
| for cat_name, cat_data in categories.items(): | |
| section_lines.append(f'<category name="{cat_name}" severity="{cat_data["severity"]}">') | |
| section_lines.append(cat_data["description"]) | |
| section_lines.append("") | |
| if "criteria" in cat_data: | |
| section_lines.append("Key criteria:") | |
| for criterion in cat_data["criteria"]: | |
| section_lines.append(f"- {criterion}") | |
| section_lines.append("") | |
| section_lines.append("</category>") | |
| section_lines.append("") | |
| section_lines.append("</classification_categories>") | |
| return "\n".join(section_lines) | |
| def get_enhanced_prompt(self, agent_type: str, session_id: Optional[str] = None) -> str: | |
| """ | |
| Get enhanced prompt with integrated shared components. | |
| Args: | |
| agent_type: Type of AI agent | |
| session_id: Optional session ID for session-level overrides | |
| Returns: | |
| Enhanced prompt content with shared components integrated | |
| """ | |
| config = self.controller.get_prompt(agent_type, session_id=session_id) | |
| # Start with base prompt | |
| enhanced_prompt = config.base_prompt | |
| # Add shared components sections if not already present | |
| if "<shared_indicators>" not in enhanced_prompt: | |
| indicators_section = self.generate_indicators_section() | |
| if indicators_section: | |
| # Insert after system_role if present, otherwise at the beginning | |
| if "<system_role>" in enhanced_prompt and "</system_role>" in enhanced_prompt: | |
| role_end = enhanced_prompt.find("</system_role>") + len("</system_role>") | |
| enhanced_prompt = (enhanced_prompt[:role_end] + | |
| f"\n\n<shared_indicators>\n{indicators_section}\n</shared_indicators>" + | |
| enhanced_prompt[role_end:]) | |
| else: | |
| enhanced_prompt = f"<shared_indicators>\n{indicators_section}\n</shared_indicators>\n\n{enhanced_prompt}" | |
| if "<shared_rules>" not in enhanced_prompt: | |
| rules_section = self.generate_rules_section() | |
| if rules_section: | |
| # Insert before output_format if present, otherwise at the end | |
| if "<output_format>" in enhanced_prompt: | |
| format_start = enhanced_prompt.find("<output_format>") | |
| enhanced_prompt = (enhanced_prompt[:format_start] + | |
| f"<shared_rules>\n{rules_section}\n</shared_rules>\n\n" + | |
| enhanced_prompt[format_start:]) | |
| else: | |
| enhanced_prompt += f"\n\n<shared_rules>\n{rules_section}\n</shared_rules>" | |
| return enhanced_prompt | |
| def update_prompt_file(self, agent_type: str, backup: bool = True) -> bool: | |
| """ | |
| Update a prompt file to use shared components. | |
| Args: | |
| agent_type: Type of AI agent | |
| backup: Whether to create a backup of the original file | |
| Returns: | |
| True if update was successful | |
| """ | |
| try: | |
| from ..prompt_loader import PROMPTS_DIR | |
| from datetime import datetime | |
| filename = f"{agent_type}.txt" | |
| filepath = PROMPTS_DIR / filename | |
| if not filepath.exists(): | |
| print(f"Prompt file not found: {filepath}") | |
| return False | |
| # Create backup if requested | |
| if backup: | |
| backup_path = filepath.with_suffix(f".backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt") | |
| with open(filepath, 'r', encoding='utf-8') as f: | |
| original_content = f.read() | |
| with open(backup_path, 'w', encoding='utf-8') as f: | |
| f.write(original_content) | |
| print(f"Backup created: {backup_path}") | |
| # Generate enhanced prompt | |
| enhanced_prompt = self.get_enhanced_prompt(agent_type) | |
| # Write updated prompt | |
| with open(filepath, 'w', encoding='utf-8') as f: | |
| f.write(enhanced_prompt) | |
| print(f"Updated prompt file: {filepath}") | |
| return True | |
| except Exception as e: | |
| print(f"Error updating prompt file: {e}") | |
| return False | |
| def validate_prompt_integration(self, agent_type: str) -> Dict[str, Any]: | |
| """ | |
| Validate that a prompt properly integrates shared components. | |
| Args: | |
| agent_type: Type of AI agent | |
| Returns: | |
| Dictionary with validation results | |
| """ | |
| config = self.controller.get_prompt(agent_type) | |
| result = { | |
| "agent_type": agent_type, | |
| "has_shared_indicators": len(config.shared_indicators) > 0, | |
| "has_shared_rules": len(config.shared_rules) > 0, | |
| "has_templates": len(config.templates) > 0, | |
| "indicator_count": len(config.shared_indicators), | |
| "rule_count": len(config.shared_rules), | |
| "template_count": len(config.templates), | |
| "validation_errors": [], | |
| "recommendations": [] | |
| } | |
| # Check for common integration issues | |
| prompt_content = config.base_prompt.lower() | |
| if "indicator" in prompt_content and not config.shared_indicators: | |
| result["validation_errors"].append("Prompt mentions indicators but has no shared indicators") | |
| if "rule" in prompt_content and not config.shared_rules: | |
| result["validation_errors"].append("Prompt mentions rules but has no shared rules") | |
| if len(config.shared_indicators) == 0: | |
| result["recommendations"].append("Consider adding shared indicators for consistency") | |
| if len(config.shared_rules) == 0: | |
| result["recommendations"].append("Consider adding shared rules for consistency") | |
| return result | |
| def create_integrator() -> PromptIntegrator: | |
| """Create a new PromptIntegrator instance.""" | |
| return PromptIntegrator() |