Spaces:
Sleeping
Sleeping
| """ | |
| Output Generation Agent | |
| Consolidates all agent outputs into a user-friendly format. | |
| """ | |
| from typing import Dict, List, Any | |
| import json | |
| from datetime import datetime | |
| class OutputGenerationAgent: | |
| def __init__(self): | |
| self.version = "1.0" | |
| def generate_output(self, | |
| analysis_result: Dict[str, Any], | |
| technique_selection: Dict[str, Any], | |
| prompt_generation: Dict[str, Any], | |
| risk_analysis: Dict[str, Any], | |
| cost_estimation: Dict[str, Any]) -> Dict[str, Any]: | |
| """ | |
| Generate comprehensive output combining all agent results. | |
| Args: | |
| analysis_result: From InputAnalysisAgent | |
| technique_selection: From TechniqueSelectionAgent | |
| prompt_generation: From PromptGenerationAgent | |
| risk_analysis: From RiskAnalysisAgent | |
| cost_estimation: From CostEstimationAgent | |
| Returns: | |
| Comprehensive formatted output | |
| """ | |
| # Generate the main output structure | |
| output = { | |
| "timestamp": datetime.now().isoformat(), | |
| "version": self.version, | |
| "original_prompt": analysis_result.get("original_prompt", ""), | |
| "analysis": self._format_analysis(analysis_result), | |
| "technique": self._format_technique_selection(technique_selection), | |
| "optimized_prompt": self._format_optimized_prompt(prompt_generation), | |
| "safety": self._format_safety_analysis(risk_analysis), | |
| "cost": self._format_cost_analysis(cost_estimation), | |
| "recommendations": self._generate_comprehensive_recommendations( | |
| analysis_result, technique_selection, risk_analysis, cost_estimation | |
| ), | |
| "metadata": self._generate_metadata( | |
| analysis_result, technique_selection, prompt_generation, | |
| risk_analysis, cost_estimation | |
| ) | |
| } | |
| return output | |
| def _format_analysis(self, analysis_result: Dict[str, Any]) -> Dict[str, Any]: | |
| """Format the input analysis results.""" | |
| return { | |
| "task_type": analysis_result.get("task_type", "unknown"), | |
| "complexity": analysis_result.get("complexity", "unknown"), | |
| "domain": analysis_result.get("domain", "general"), | |
| "intent": analysis_result.get("intent", ""), | |
| "key_entities": analysis_result.get("entities", []), | |
| "metrics": { | |
| "word_count": analysis_result.get("word_count", 0), | |
| "character_count": analysis_result.get("char_count", 0) | |
| } | |
| } | |
| def _format_technique_selection(self, technique_info: Dict[str, Any]) -> Dict[str, Any]: | |
| """Format the technique selection results.""" | |
| return { | |
| "selected_technique": technique_info.get("name", "Unknown"), | |
| "description": technique_info.get("description", ""), | |
| "reasoning": technique_info.get("reasoning", ""), | |
| "confidence": technique_info.get("confidence", 0.0), | |
| "use_cases": technique_info.get("use_cases", []), | |
| "complexity_handled": technique_info.get("complexity", []) | |
| } | |
| def _format_optimized_prompt(self, prompt_generation: Dict[str, Any]) -> Dict[str, Any]: | |
| """Format the optimized prompt results.""" | |
| return { | |
| "prompt": prompt_generation.get("optimized_prompt", ""), | |
| "technique_used": prompt_generation.get("technique_used", ""), | |
| "enhancements_applied": prompt_generation.get("enhancements_applied", []), | |
| "estimated_tokens": prompt_generation.get("estimated_tokens", 0), | |
| "structure": prompt_generation.get("prompt_structure", {}), | |
| "improvements": self._identify_improvements(prompt_generation) | |
| } | |
| def _format_safety_analysis(self, risk_analysis: Dict[str, Any]) -> Dict[str, Any]: | |
| """Format the safety analysis results.""" | |
| return { | |
| "overall_risk": risk_analysis.get("overall_risk", "unknown"), | |
| "risk_score": risk_analysis.get("risk_score", 0), | |
| "safe_to_proceed": risk_analysis.get("safe_to_proceed", True), | |
| "should_block": risk_analysis.get("should_block", False), | |
| "risks_found": len(risk_analysis.get("risks_found", [])), | |
| "risk_categories": list(set( | |
| risk["category"] for risk in risk_analysis.get("risks_found", []) | |
| )), | |
| "recommendations": risk_analysis.get("recommendations", []), | |
| "details": risk_analysis.get("risks_found", []) | |
| } | |
| def _format_cost_analysis(self, cost_estimation: Dict[str, Any]) -> Dict[str, Any]: | |
| """Format the cost analysis results.""" | |
| summary = cost_estimation.get("summary", {}) | |
| return { | |
| "input_tokens": cost_estimation.get("input_tokens", 0), | |
| "expected_output_tokens": cost_estimation.get("expected_output_tokens", 0), | |
| "total_tokens": cost_estimation.get("total_expected_tokens", 0), | |
| "cheapest_option": summary.get("cheapest_model", {}), | |
| "most_expensive_option": summary.get("most_expensive_model", {}), | |
| "cost_range": summary.get("cost_range", {}), | |
| "average_cost": summary.get("average_cost", 0), | |
| "potential_savings": summary.get("cost_savings", 0), | |
| "models_analyzed": summary.get("total_models", 0), | |
| "recommendations": cost_estimation.get("recommendations", []), | |
| "detailed_estimates": self._format_detailed_costs(cost_estimation.get("model_estimates", {})) | |
| } | |
| def _format_detailed_costs(self, model_estimates: Dict[str, Any]) -> List[Dict[str, Any]]: | |
| """Format detailed cost estimates for display.""" | |
| formatted_costs = [] | |
| for model_name, estimate in model_estimates.items(): | |
| formatted_costs.append({ | |
| "model": model_name, | |
| "provider": estimate.get("provider", "Unknown"), | |
| "total_cost": estimate.get("total_cost", 0), | |
| "input_cost": estimate.get("input_cost", 0), | |
| "output_cost": estimate.get("output_cost", 0), | |
| "exceeds_limit": estimate.get("exceeds_limit", False), | |
| "max_tokens": estimate.get("max_tokens", 0), | |
| "cost_per_token": estimate.get("cost_per_token", 0) | |
| }) | |
| # Sort by total cost | |
| formatted_costs.sort(key=lambda x: x["total_cost"]) | |
| return formatted_costs | |
| def _identify_improvements(self, prompt_generation: Dict[str, Any]) -> List[str]: | |
| """Identify improvements made to the original prompt.""" | |
| improvements = [] | |
| enhancements = prompt_generation.get("enhancements_applied", []) | |
| structure = prompt_generation.get("prompt_structure", {}) | |
| if "Step-by-step reasoning" in enhancements: | |
| improvements.append("Added structured reasoning approach") | |
| if "Example-based learning" in enhancements: | |
| improvements.append("Included relevant examples for better guidance") | |
| if structure.get("has_instructions"): | |
| improvements.append("Added clear instructions") | |
| if structure.get("has_step_markers"): | |
| improvements.append("Structured with step-by-step markers") | |
| if structure.get("structure_complexity") == "complex": | |
| improvements.append("Enhanced with comprehensive structure") | |
| return improvements | |
| def _generate_comprehensive_recommendations(self, | |
| analysis_result: Dict[str, Any], | |
| technique_selection: Dict[str, Any], | |
| risk_analysis: Dict[str, Any], | |
| cost_estimation: Dict[str, Any]) -> List[str]: | |
| """Generate comprehensive recommendations based on all analyses.""" | |
| recommendations = [] | |
| # Safety recommendations (highest priority) | |
| if not risk_analysis.get("safe_to_proceed", True): | |
| recommendations.append("⚠️ SAFETY: Review and modify prompt to address safety concerns before use") | |
| # Technique recommendations | |
| confidence = technique_selection.get("confidence", 0) | |
| if confidence < 0.7: | |
| recommendations.append("💡 TECHNIQUE: Consider manual review of selected technique") | |
| # Cost recommendations | |
| cost_summary = cost_estimation.get("summary", {}) | |
| if cost_summary.get("cost_savings", 0) > 0.01: # More than 1 cent savings | |
| cheapest = cost_summary.get("cheapest_model", {}) | |
| if cheapest: | |
| recommendations.append(f"💰 COST: Consider using {cheapest.get('name')} for cost optimization") | |
| # Complexity recommendations | |
| complexity = analysis_result.get("complexity", "simple") | |
| if complexity == "complex": | |
| recommendations.append("🔧 COMPLEXITY: Consider breaking down into smaller, focused prompts") | |
| # Token recommendations | |
| total_tokens = cost_estimation.get("total_expected_tokens", 0) | |
| if total_tokens > 2000: | |
| recommendations.append("📏 LENGTH: Consider shortening prompt to reduce token usage") | |
| # General best practices | |
| recommendations.append("✅ BEST PRACTICE: Test the optimized prompt with your specific use case") | |
| recommendations.append("📊 MONITORING: Track actual costs and adjust estimates accordingly") | |
| return recommendations | |
| def _generate_metadata(self, *agent_results) -> Dict[str, Any]: | |
| """Generate metadata about the analysis process.""" | |
| return { | |
| "agents_used": [ | |
| "InputAnalysisAgent", | |
| "TechniqueSelectionAgent", | |
| "PromptGenerationAgent", | |
| "RiskAnalysisAgent", | |
| "CostEstimationAgent" | |
| ], | |
| "processing_timestamp": datetime.now().isoformat(), | |
| "system_version": self.version, | |
| "analysis_complete": True, | |
| "total_processing_steps": 5 | |
| } | |
| def generate_simple_output(self, optimized_prompt: str, technique_name: str, | |
| cost_estimate: float, is_safe: bool) -> Dict[str, Any]: | |
| """Generate a simplified output for basic use cases.""" | |
| return { | |
| "optimized_prompt": optimized_prompt, | |
| "technique_used": technique_name, | |
| "estimated_cost": cost_estimate, | |
| "is_safe": is_safe, | |
| "timestamp": datetime.now().isoformat(), | |
| "version": "simple_" + self.version | |
| } | |
| def format_for_display(self, output: Dict[str, Any], format_type: str = "detailed") -> str: | |
| """Format output for different display types.""" | |
| if format_type == "json": | |
| return json.dumps(output, indent=2) | |
| elif format_type == "summary": | |
| return self._format_summary_display(output) | |
| elif format_type == "detailed": | |
| return self._format_detailed_display(output) | |
| else: | |
| return str(output) | |
| def _format_summary_display(self, output: Dict[str, Any]) -> str: | |
| """Format a summary display of the output.""" | |
| technique = output.get("technique", {}) | |
| safety = output.get("safety", {}) | |
| cost = output.get("cost", {}) | |
| summary = f""" | |
| PRONTO ANALYSIS SUMMARY | |
| ====================== | |
| ✨ Technique Selected: {technique.get('selected_technique', 'Unknown')} | |
| 🛡️ Safety Status: {'✅ Safe' if safety.get('safe_to_proceed') else '⚠️ Needs Review'} | |
| 💰 Estimated Cost: ${cost.get('cheapest_option', {}).get('cost', 0):.4f} (cheapest option) | |
| 📊 Total Tokens: {cost.get('total_tokens', 0)} | |
| 🎯 Key Recommendations: | |
| {chr(10).join('• ' + rec for rec in output.get('recommendations', [])[:3])} | |
| """ | |
| return summary.strip() | |
| def _format_detailed_display(self, output: Dict[str, Any]) -> str: | |
| """Format a detailed display of the output.""" | |
| # This would be a comprehensive formatted output | |
| # For brevity, returning JSON format | |
| return json.dumps(output, indent=2) | |