Update components/template_generator.py

components/template_generator.py

@@ -1,14 +1,25 @@
-# components/template_generator.py
 import streamlit as st
 from datetime import datetime
-import json
 import os
-from typing import Dict, List
+from typing import Dict
+from fpdf import FPDF
 import anthropic
 
+
 class TemplateGenerator:
     def __init__(self):
-        self.client = anthropic.Anthropic(api_key=st.secrets["Anthropic_API_KEY"])
+        # Initialize Anthropic client with API key
+        try:
+            api_key = os.getenv("ANTHROPIC_API_KEY")
+            if not api_key:
+                st.error("Please set the ANTHROPIC_API_KEY in your environment variables.")
+                st.stop()
+            self.client = anthropic.Anthropic(api_key=api_key)
+        except Exception as e:
+            st.error(f"Error initializing Anthropic client: {str(e)}")
+            st.stop()
+
+        # Define templates
         self.templates = {
             "nda": {
                 "name": "Non-Disclosure Agreement",
@@ -24,69 +35,27 @@ class TemplateGenerator:
                     "return_period": "Days to return/destroy information"
                 }
             },
-            "mou": {
-                "name": "Memorandum of Understanding",
-                "description": "Initial agreement outlining basic terms",
-                "fields": {
-                    "party1_name": "First party name",
-                    "party2_name": "Second party name",
-                    "purpose": "Purpose of the collaboration",
-                    "start_date": "Start date",
-                    "end_date": "End date",
-                    "key_responsibilities": "Key responsibilities of parties",
-                    "financial_terms": "Basic financial arrangements",
-                    "termination_terms": "Termination conditions"
-                }
-            },
             "services_agreement": {
                 "name": "Services Agreement",
-                "description": "Contract for provision of services",
+                "description": "Comprehensive contract for provision of services",
                 "fields": {
                     "service_provider": "Provider company/individual name",
                     "client_name": "Client name",
                     "services_description": "Detailed service description",
                     "start_date": "Service start date",
                     "payment_terms": "Payment amount and schedule",
-                    "service_levels": "Service level requirements",
                     "termination_notice": "Notice period for termination",
-                    "governing_law": "Jurisdiction"
-                }
-            },
-            "employment_contract": {
-                "name": "Employment Contract",
-                "description": "Standard employment agreement",
-                "fields": {
-                    "employer_name": "Company name",
-                    "employee_name": "Employee full name",
-                    "position": "Job title/position",
-                    "start_date": "Employment start date",
-                    "salary": "Annual salary amount",
-                    "benefits": "Benefits package details",
-                    "work_hours": "Working hours/schedule",
-                    "notice_period": "Notice period for termination",
-                    "probation_period": "Probation period in months"
+                    "governing_law": "Jurisdiction",
+                    "service_levels": "Service level agreements (SLAs)"
                 }
             },
-            "will": {
-                "name": "Last Will and Testament",
-                "description": "Basic will template",
-                "fields": {
-                    "testator_name": "Full name of testator",
-                    "testator_address": "Current address",
-                    "executor_name": "Executor's full name",
-                    "executor_address": "Executor's address",
-                    "beneficiaries": "List of beneficiaries",
-                    "specific_bequests": "Specific property/asset allocations",
-                    "residuary_estate": "Disposition of remaining estate",
-                    "funeral_wishes": "Funeral arrangements (optional)"
-                }
-            }
+            # Extend templates here
         }
 
     def render(self):
-        """Render template generator interface"""
+        """Render template generator interface."""
         st.header("Legal Document Generator")
-        st.write("Generate legal documents using AI-assisted templates")
+        st.write("Generate comprehensive legal documents using customizable templates.")
 
         # Template selection
         template_type = st.selectbox(
@@ -101,73 +70,36 @@ class TemplateGenerator:
         # Input form
         with st.form(f"template_form_{template_type}"):
             inputs = {}
-            
-            # Add fields based on template type
             for field, description in self.templates[template_type]["fields"].items():
                 if "date" in field:
                     inputs[field] = st.date_input(f"{description}")
-                elif any(word in field for word in ["description", "terms", "responsibilities", "bequests"]):
+                elif "description" in field or "terms" in field:
                     inputs[field] = st.text_area(f"{description}", height=100)
-                elif "amount" in field or "salary" in field:
+                elif "amount" in field or "months" in field:
                     inputs[field] = st.number_input(f"{description}", min_value=0)
-                elif "period" in field:
-                    inputs[field] = st.number_input(f"{description}", min_value=1)
                 else:
                     inputs[field] = st.text_input(f"{description}")
 
-            # Jurisdiction selection where applicable
-            if template_type in ["nda", "services_agreement"]:
-                jurisdiction = st.selectbox(
-                    "Select Jurisdiction",
-                    ["UK", "India", "UAE"],
-                    help="Choose the governing law for this document"
-                )
-                inputs["jurisdiction"] = jurisdiction
-
+            # Submit form
             generate = st.form_submit_button("Generate Document")
 
         if generate:
             if all(inputs.values()):
                 with st.spinner("Generating document..."):
                     document = self._generate_document(template_type, inputs)
-                    
-                    # Preview
-                    st.subheader("Generated Document")
-                    with st.expander("Preview", expanded=True):
-                        st.markdown(document)
-                    
-                    # Download options
-                    col1, col2 = st.columns(2)
-                    with col1:
-                        st.download_button(
-                            "Download as TXT",
-                            document,
-                            file_name=f"{template_type}_{datetime.now().strftime('%Y%m%d')}.txt",
-                            mime="text/plain"
-                        )
-                    with col2:
-                        st.download_button(
-                            "Download as PDF",
-                            document,
-                            file_name=f"{template_type}_{datetime.now().strftime('%Y%m%d')}.pdf",
-                            mime="application/pdf"
-                        )
+                    self._display_document(document, template_type)
             else:
-                st.error("Please fill in all required fields")
+                st.error("Please fill in all required fields.")
 
     def _generate_document(self, template_type: str, inputs: Dict) -> str:
-        """Generate document using Claude"""
+        """Generate document using Claude."""
         try:
             prompt = self._create_prompt(template_type, inputs)
-            
             message = self.client.messages.create(
                 model="claude-3-sonnet-20240229",
                 max_tokens=2000,
                 temperature=0.7,
-                messages=[{
-                    "role": "user",
-                    "content": prompt
-                }]
+                messages=[{"role": "user", "content": prompt}]
             )
             return message.content[0].text
         except Exception as e:
@@ -175,34 +107,63 @@ class TemplateGenerator:
             return ""
 
     def _create_prompt(self, template_type: str, inputs: Dict) -> str:
-        """Create prompt for document generation"""
+        """Create prompt for document generation."""
         template_name = self.templates[template_type]["name"]
-        
+
         prompt = f"""Generate a detailed {template_name} using the following information:
 
 Template Type: {template_name}
 """
-        
-        # Add inputs to prompt
         for field, value in inputs.items():
             prompt += f"{field}: {value}\n"
 
-        prompt += f"""
+        prompt += """
+
+Ensure the document:
+1. Follows professional legal formatting.
+2. Includes all necessary clauses and sections.
+3. Is tailored for {inputs.get('jurisdiction', 'general')} jurisdiction.
+4. Uses formal legal language.
+5. Includes signature blocks and clear numbering.
+"""
+        return prompt
 
-Please generate a complete and professionally formatted legal document that:
-1. Follows standard legal document structure
-2. Includes all necessary clauses and sections
-3. Uses formal legal language
-4. Incorporates all provided information
-5. Follows the legal requirements for {inputs.get('jurisdiction', 'general')} jurisdiction
-6. Includes proper signature blocks and dates
-7. Has clear section numbering and formatting
+    def _display_document(self, document: str, template_type: str):
+        """Display and provide download options for the generated document."""
+        st.subheader("Generated Document")
+        with st.expander("Preview", expanded=True):
+            st.markdown(document)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.download_button(
+                "Download as TXT",
+                document,
+                file_name=f"{template_type}_{datetime.now().strftime('%Y%m%d')}.txt",
+                mime="text/plain"
+            )
+        with col2:
+            pdf_document = self._generate_pdf(document)
+            st.download_button(
+                "Download as PDF",
+                pdf_document,
+                file_name=f"{template_type}_{datetime.now().strftime('%Y%m%d')}.pdf",
+                mime="application/pdf"
+            )
 
-Return the complete document in a format ready for use."""
+    def _generate_pdf(self, document: str) -> bytes:
+        """Generate a PDF file for the document."""
+        pdf = FPDF()
+        pdf.add_page()
+        pdf.set_font("Arial", size=12)
+        for line in document.split("\n"):
+            pdf.cell(200, 10, txt=line, ln=True)
+        output = bytes()
+        pdf.output(name="document.pdf", dest="S").encode("latin1")
+        return output
 
-        return prompt
 
 def render_template_generator():
-    """Main function to render template generator"""
+    """Main function to render the template generator."""
     generator = TemplateGenerator()
-    generator.render()
+    generator.render()