File size: 6,009 Bytes
e13d87a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
from typing import Dict, List, Any, Optional, Union
import pandas as pd
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.image import MIMEImage
import base64
import io
from pathlib import Path
import json
import datetime
from llama_index.tools import FunctionTool
class ExportTools:
"""Tools for exporting data, generating reports, and sending emails."""
def __init__(self, output_directory: str = "./exports"):
"""Initialize with directory for saved exports."""
self.output_directory = Path(output_directory)
self.output_directory.mkdir(exist_ok=True, parents=True)
self.tools = self._create_tools()
def _create_tools(self) -> List[FunctionTool]:
"""Create LlamaIndex function tools for export operations."""
tools = [
FunctionTool.from_defaults(
name="generate_report",
description="Generate a report from conversation and results",
fn=self.generate_report
),
FunctionTool.from_defaults(
name="save_results_to_csv",
description="Save query results to a CSV file",
fn=self.save_results_to_csv
),
FunctionTool.from_defaults(
name="send_email",
description="Send results via email",
fn=self.send_email
)
]
return tools
def get_tools(self) -> List[FunctionTool]:
"""Get all available export tools."""
return self.tools
def generate_report(self, title: str, content: str,
images: List[str] = None) -> Dict[str, Any]:
"""Generate HTML report from content and images."""
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"{title.replace(' ', '_')}_{timestamp}.html"
file_path = self.output_directory / filename
# Basic HTML template
html = f"""
<!DOCTYPE html>
<html>
<head>
<title>{title}</title>
<style>
body {{ font-family: Arial, sans-serif; margin: 40px; line-height: 1.6; }}
h1 {{ color: #333366; }}
.report-container {{ max-width: 800px; margin: 0 auto; }}
.timestamp {{ color: #666; font-size: 0.8em; }}
img {{ max-width: 100%; height: auto; margin: 20px 0; }}
</style>
</head>
<body>
<div class="report-container">
<h1>{title}</h1>
<div class="timestamp">Generated on: {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}</div>
<div class="content">
{content.replace('\n', '<br>')}
</div>
"""
# Add images if provided
if images and len(images) > 0:
html += "<div class='images'>"
for i, img_base64 in enumerate(images):
html += f"<img src='data:image/png;base64,{img_base64}' alt='Figure {i+1}'>"
html += "</div>"
html += """
</div>
</body>
</html>
"""
# Write to file
with open(file_path, "w", encoding="utf-8") as f:
f.write(html)
return {
"success": True,
"report_path": str(file_path),
"title": title,
"timestamp": timestamp
}
def save_results_to_csv(self, data: List[Dict[str, Any]],
filename: str = None) -> Dict[str, Any]:
"""Save query results to a CSV file."""
if not data or len(data) == 0:
return {"success": False, "error": "No data provided"}
# Create DataFrame from data
df = pd.DataFrame(data)
# Generate filename if not provided
if not filename:
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"query_results_{timestamp}.csv"
# Ensure filename has .csv extension
if not filename.lower().endswith('.csv'):
filename += '.csv'
file_path = self.output_directory / filename
# Save to CSV
df.to_csv(file_path, index=False)
return {
"success": True,
"file_path": str(file_path),
"row_count": len(df),
"column_count": len(df.columns)
}
def send_email(self, to_email: str, subject: str, body: str,
from_email: str = None, smtp_server: str = None,
smtp_port: int = 587, username: str = None,
password: str = None, images: List[str] = None) -> Dict[str, Any]:
"""
Send email with results.
Note: In production, credentials should be securely managed.
For demo purposes, this will log the email content instead.
"""
# For safety in a demo app, don't actually send emails
# Just log what would be sent and return success
email_content = {
"to": to_email,
"subject": subject,
"body": body[:100] + "..." if len(body) > 100 else body,
"images": f"{len(images) if images else 0} images would be attached",
"note": "Email sending is simulated for demo purposes"
}
# Log the email content
print(f"SIMULATED EMAIL: {json.dumps(email_content, indent=2)}")
return {
"success": True,
"to": to_email,
"subject": subject,
"simulated": True,
"timestamp": datetime.datetime.now().isoformat()
}
|