Update app.py

app.py

@@ -109,7 +109,6 @@ def fallback_search(query, country_code=""):
         st.error(f"Fallback search failed: {str(e)}")
         return []
 
-# Use LLM to extract relevant information from search results
 # Use LLM to extract relevant information from search results
 def analyze_search_results(results, query, country):
     if not GROQ_API_KEY:
@@ -143,22 +142,23 @@ def analyze_search_results(results, query, country):
         
         # Call Groq API with mixed model approach (prefer cheaper model)
         completion = groq_client.chat.completions.create(
-            model="llama3-8b-8192",
+            model="llama3-8b-8192",  # Free/cheaper model
             messages=[{"role": "user", "content": prompt}],
-            temperature=0.1,
-            max_tokens=500  # Reduced token limit
+            temperature=0.0,
+            max_tokens=800
         )
         
         try:
             # Parse the response as JSON
-            return json.loads(completion.choices[0].message.content)
+            analysis = json.loads(completion.choices[0].message.content)
+            return results, analysis
         except json.JSONDecodeError:
-            # If parsing fails, return empty dict
-            return {}
+            # If parsing fails, return the original results
+            return results, None
             
     except Exception as e:
-        st.error(f"Error analyzing search results: {str(e)}")
-        return {}
+        st.error(f"LLM analysis failed: {str(e)}")
+        return results, None
 
 # Try to download PDF
 def fetch_pdf(url):
@@ -231,151 +231,19 @@ def find_pdf_in_html_page(url, html_content=None):
         return None
 
 # Display PDF safely with error handling
-# Replace your existing display_pdf function with this improved version
-
 def display_pdf(file_bytesio):
-    """Display PDF with multiple fallback methods for better browser compatibility"""
     try:
         file_bytesio.seek(0)
-        pdf_bytes = file_bytesio.read()
+        base64_pdf = base64.b64encode(file_bytesio.read()).decode('utf-8')
+        pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="800" type="application/pdf"></iframe>'
+        st.markdown(pdf_display, unsafe_allow_html=True)
+        st.success("PDF loaded successfully!")
         
-        # Method 1: Try Streamlit's native PDF display (most reliable)
-        try:
-            file_bytesio.seek(0)
-            # Use Streamlit's built-in PDF viewer (works better than iframe)
-            st.markdown("### 📄 PDF Preview")
-            
-            # Try to display using Streamlit's method
-            # This is more reliable than iframe approach
-            base64_pdf = base64.b64encode(pdf_bytes).decode('utf-8')
-            
-            # Use object tag instead of iframe (better Chrome compatibility)
-            pdf_display = f'''
-            <object data="data:application/pdf;base64,{base64_pdf}" 
-                    type="application/pdf" 
-                    width="100%" 
-                    height="800px">
-                <p>Your browser doesn't support PDF preview. 
-                   <a href="data:application/pdf;base64,{base64_pdf}" download="tax_form.pdf">
-                   Click here to download the PDF</a>
-                </p>
-            </object>
-            '''
-            
-            st.markdown(pdf_display, unsafe_allow_html=True)
-            
-        except Exception as e:
-            # Method 2: Fallback to download only
-            st.warning("⚠️ PDF preview not available in your browser. Please download to view.")
-            st.download_button(
-                label="📥 Download PDF to View",
-                data=pdf_bytes,
-                file_name="tax_form.pdf",
-                mime="application/pdf",
-                use_container_width=True
-            )
-            
-        # Method 3: Try to extract first page as image (optional)
-        try:
-            # Convert first page to image for preview
-            import fitz  # PyMuPDF
-            file_bytesio.seek(0)
-            doc = fitz.open(stream=file_bytesio, filetype="pdf")
-            
-            if len(doc) > 0:
-                page = doc[0]
-                pix = page.get_pixmap(matrix=fitz.Matrix(1.5, 1.5))  # 1.5x zoom
-                img_data = pix.tobytes("png")
-                
-                st.markdown("### 📄 First Page Preview")
-                st.image(img_data, caption="First page of the tax form", use_container_width=True)
-                
-                if len(doc) > 1:
-                    st.info(f"This form has {len(doc)} pages total. Download the PDF to see all pages.")
-                    
-            doc.close()
-            
-        except Exception as img_error:
-            st.info("PDF loaded successfully! Use the download button above to view the complete form.")
-            
-    except Exception as e:
-        st.error(f"Error loading PDF: {str(e)}")
-        st.info("Please try uploading the PDF again or use a different file.")
-
-# Alternative function for better Streamlit compatibility
-def display_pdf_streamlit_native(file_bytesio):
-    """Use Streamlit's native PDF handling"""
-    try:
+        # Add a direct download option for better user experience
         file_bytesio.seek(0)
-        
-        # Simply use st.download_button which always works
-        st.download_button(
-            label="📥 Download PDF to View",
-            data=file_bytesio,
-            file_name="tax_form.pdf",
-            mime="application/pdf",
-            use_container_width=True
-        )
-        
-        # Try to show first page as image
-        try:
-            import fitz
-            file_bytesio.seek(0)
-            doc = fitz.open(stream=file_bytesio, filetype="pdf")
-            
-            if len(doc) > 0:
-                # Show first page as image
-                page = doc[0]
-                pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x zoom for better quality
-                img_data = pix.tobytes("png")
-                
-                st.image(img_data, caption=f"Preview: Page 1 of {len(doc)}", use_column_width=True)
-                
-                if len(doc) > 1:
-                    st.info(f"📄 This form has {len(doc)} pages. Download the PDF to see all pages.")
-                    
-            doc.close()
-            
-        except Exception:
-            st.info("PDF ready for download. Click the button above to view the complete form.")
-            
-    except Exception as e:
-        st.error(f"Error processing PDF: {str(e)}")
-
-# Even simpler fallback function
-def display_pdf_simple(file_bytesio):
-    """Simple PDF display with download only"""
-    try:
-        file_bytesio.seek(0)
-        pdf_bytes = file_bytesio.read()
-        
-        st.markdown("### 📄 Tax Form Ready")
-        st.success("✅ PDF loaded successfully!")
-        
-        # Provide download button
-        st.download_button(
-            label="📥 Download PDF to View",
-            data=pdf_bytes,
-            file_name="pakistan_tax_form.pdf",
-            mime="application/pdf",
-            use_container_width=True
-        )
-        
-        # Show PDF info
-        try:
-            import fitz
-            file_bytesio.seek(0)
-            doc = fitz.open(stream=file_bytesio, filetype="pdf")
-            page_count = len(doc)
-            doc.close()
-            
-            st.info(f"📄 Form contains {page_count} page(s)")
-            
-        except Exception:
-            st.info("📄 PDF is ready for download")
-            
     except Exception as e:
-        st.error(f"Error: {str(e)}")
+        st.error(f"Error displaying PDF: {str(e)}")
+        st.info("If the PDF isn't displaying, you can try using the direct link.")
 
 
 # Extract interactive fields from PDF
@@ -410,32 +278,38 @@ def extract_form_fields(file_bytesio):
     except Exception as e:
         st.error(f"Error extracting form fields: {str(e)}")
         return []
+
 # Use LLM to explain form fields
 def explain_form_fields(fields, country, form_name):
     if not GROQ_API_KEY or not fields:
         return {}
         
     try:
-        # CHANGE: Limit fields and reduce content
-        limited_fields = fields[:10]  # Only take first 10 fields
-        fields_summary = []
-        for field in limited_fields:
-            fields_summary.append({
-                "name": field.get("name", "Unknown"),
-                "type": field.get("type", "Unknown")
-            })
-        
-        fields_json = json.dumps(fields_summary, indent=1)  # Less indentation
+        fields_json = json.dumps(fields, indent=2)
         
         prompt = f"""
-        Tax form ({form_name}) from {country} has these fields:
+        These are form fields from a tax form ({form_name}) from {country}:
         {fields_json}
         
-        Provide brief explanations for these field types in JSON format:
+        Please analyze these fields and:
+        1. Group them into logical sections (personal info, income, deductions, etc.)
+        2. Explain any technical tax terms in simple language
+        3. Identify which fields are mandatory vs. optional if possible
+        
+        Format your response as JSON with the following structure:
         {{
-            "field_explanations": {{
-                "field_name": "brief explanation"
-            }}
+            "sections": [
+                {{
+                    "name": "section name",
+                    "fields": ["field1", "field2"],
+                    "explanation": "explanation of this section"
+                }}
+            ],
+            "key_terms": {{
+                "term1": "simple explanation",
+                "term2": "simple explanation"
+            }},
+            "mandatory_fields": ["field1", "field2"]
         }}
         """
         
@@ -443,19 +317,19 @@ def explain_form_fields(fields, country, form_name):
             model="llama3-8b-8192",
             messages=[{"role": "user", "content": prompt}],
             temperature=0.1,
-            max_tokens=500  # Reduced token limit
+            max_tokens=1000
         )
         
         try:
-            # Parse the response as JSON
-            return json.loads(completion.choices[0].message.content)
+            explanation = json.loads(completion.choices[0].message.content)
+            return explanation
         except json.JSONDecodeError:
-            # If parsing fails, return empty dict
             return {}
             
     except Exception as e:
         st.error(f"Error explaining form fields: {str(e)}")
         return {}
+
 # Fill PDF form with user data
 def fill_pdf_form(file_bytesio, field_values):
     try:
@@ -627,97 +501,6 @@ def recommend_tax_form_type(user_query):
             
     except Exception as e:
         return ["Income Tax Return", "Sales Tax Return", "Withholding Tax Statement"]
-# Generate summary of uploaded form
-def generate_form_summary(form_fields, file_name):
-    """Generate a summary of the uploaded tax form using LLM"""
-    if not GROQ_API_KEY or not form_fields:
-        return None
-        
-    try:
-        # Prepare form fields information
-        fields_info = []
-        for field in form_fields[:20]:  # Limit to first 20 fields to avoid token limits
-            fields_info.append({
-                "name": field.get("name", ""),
-                "type": field.get("type", ""),
-                "page": field.get("page", 1)
-            })
-        
-        fields_json = json.dumps(fields_info, indent=2)
-        
-        prompt = f"""
-        Analyze this Pakistani tax form: "{file_name}"
-        
-        Form fields found:
-        {fields_json}
-        
-        Please provide a summary that includes:
-        1. What type of tax form this appears to be
-        2. Who should use this form (individuals, businesses, etc.)
-        3. Main sections/categories of information required
-        4. Key requirements or important notes
-        
-        Keep the summary concise but informative, focused on Pakistani tax context.
-        """
-        
-        completion = groq_client.chat.completions.create(
-            model="llama3-8b-8192",
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0.2,
-            max_tokens=600
-        )
-        
-        return completion.choices[0].message.content
-            
-    except Exception as e:
-        st.error(f"Error generating summary: {str(e)}")
-        return None
-
-# Generate response for uploaded form questions
-def generate_upload_form_response(question, file_name, form_fields):
-    """Generate responses to questions about uploaded forms"""
-    if not GROQ_API_KEY:
-        return "I need an LLM API key to analyze and answer questions about your uploaded form."
-    
-    try:
-        # Prepare context from form fields
-        context = f"User uploaded a tax form: '{file_name}'\n"
-        
-        if form_fields:
-            # Get field names and types for context
-            field_info = []
-            for field in form_fields[:15]:  # Limit for token efficiency
-                field_info.append(f"{field.get('name', 'Unknown')} ({field.get('type', 'Unknown')})")
-            
-            context += f"Form contains fields: {', '.join(field_info)}\n"
-        
-        context += f"User question: {question}\n"
-        
-        prompt = f"""
-        {context}
-        
-        As a Pakistani tax expert, answer the user's question about their uploaded form.
-        
-        Provide helpful guidance that:
-        1. Addresses their specific question
-        2. Relates to Pakistani tax requirements
-        3. Explains any technical terms simply
-        4. Suggests next steps if relevant
-        
-        Be conversational and helpful.
-        """
-        
-        completion = groq_client.chat.completions.create(
-            model="llama3-8b-8192",
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0.3,
-            max_tokens=700
-        )
-        
-        return completion.choices[0].message.content
-            
-    except Exception as e:
-        return f"I encountered an error while analyzing your form: {str(e)}"
 # Main Streamlit app
 def main():
     st.set_page_config(
@@ -994,148 +777,63 @@ def main():
             
             # If form has fields, show them
             if 'form_fields' in st.session_state and st.session_state.form_fields:
-                st.markdown("### 📝 Form Fields Found")
+                st.markdown("### 📝 Form Fields")
                 
-                # Simple field listing
-                field_count = len(st.session_state.form_fields)
-                st.info(f"Found {field_count} fillable fields in this form")
+                # Get field explanations if LLM is available
+                field_explanations = {}
+                if GROQ_API_KEY:
+                    with st.spinner("Analyzing form fields..."):
+                        explanations = explain_form_fields(
+                            st.session_state.form_fields, 
+                            "Pakistan", 
+                            f"{form_type} form"
+                        )
+                        if explanations:
+                            field_explanations = explanations
                 
-                # Group fields by type
-                field_types = {}
-                for field in st.session_state.form_fields:
-                    field_type = field.get('type', 'Unknown')
-                    if field_type not in field_types:
-                        field_types[field_type] = []
-                    field_types[field_type].append(field.get('name', 'Unnamed'))
+                # Display fields with explanations if available
+                if field_explanations and "sections" in field_explanations:
+                    for section in field_explanations["sections"]:
+                        with st.expander(f"📑 {section['name']}"):
+                            st.write(section["explanation"])
+                            for field_name in section["fields"]:
+                                matching_fields = [f for f in st.session_state.form_fields 
+                                                if f["name"] == field_name]
+                                if matching_fields:
+                                    field = matching_fields[0]
+                                    st.write(f"**{field['name']}** ({field['type']})")
+                else:
+                    # Simple field display without explanations
+                    for field in st.session_state.form_fields:
+                        st.write(f"**{field['name']}** ({field['type']})")
+            
+                # Option to fill form
+                st.markdown("### ✏️ Fill This Form")
+                st.info("This feature will help you fill in the form fields.")
                 
-                # Display field types and count
-                for field_type, field_names in field_types.items():
-                    with st.expander(f"{field_type} Fields ({len(field_names)})"):
-                        for name in field_names:
-                            st.write(f"• {name}")
+                if st.button("Start Filling Form"):
+                    st.session_state.is_filling = True
             else:
                 st.warning("⚠️ This form doesn't have fillable fields. It may be a scanned document or not an interactive form.")
                 st.info("You can still view and download the form, but automatic filling isn't available.")
-            
-            # Download form
-            st.download_button(
-                label="📥 Download Form",
-                data=st.session_state.pdf_bytes,
-                file_name="pakistan_tax_form.pdf",
-                mime="application/pdf"
-            )
-            
-            # Add a question and answer section for this form
-            st.markdown("### ❓ Questions about this form?")
-            form_question = st.text_input("Ask a question about this form:", key="form_question_input")
-            
-            if form_question:
-                with st.spinner("Getting answer..."):
-                    form_response = tax_agent_response(form_question, tax_form_type=form_type)
-                    st.markdown(form_response)
-    # Tab 3: Upload Form
-    with tab3:
-        st.header("Upload Your Tax Form")
-        st.markdown("Upload a PDF tax form to analyze, get summaries, and ask questions about it.")
-        
-        # File uploader
-        uploaded_file = st.file_uploader(
-            "Choose a PDF file", 
-            type=['pdf'], 
-            help="Upload any Pakistani tax form PDF for analysis"
-        )
-        
-        if uploaded_file is not None:
-            # Store uploaded file in session state
-            st.session_state.pdf_bytes = BytesIO(uploaded_file.read())
-            st.session_state.uploaded_file_name = uploaded_file.name
-            
-            # Display success message
-            st.success(f"✅ Successfully uploaded: {uploaded_file.name}")
-            
-            # Extract form fields from uploaded PDF
-            with st.spinner("Analyzing uploaded form..."):
-                st.session_state.form_fields = extract_form_fields(st.session_state.pdf_bytes)
-            
-            # Display PDF preview
-            st.markdown("### 📄 Form Preview")
-            display_pdf(st.session_state.pdf_bytes)
-            
-            # Generate summary using LLM
-            if GROQ_API_KEY:
-                with st.spinner("Generating form summary..."):
-                    summary = generate_form_summary(st.session_state.form_fields, uploaded_file.name)
-                    if summary:
-                        st.markdown("### 📋 Form Summary")
-                        st.markdown(summary)
-            # Show form fields if available
-            if st.session_state.form_fields:
-                st.markdown("### 📝 Form Fields Found")
-                
-                # Simple field listing
-                field_count = len(st.session_state.form_fields)
-                st.info(f"Found {field_count} fillable fields in this form")
                 
-                # Group fields by type
-                field_types = {}
-                for field in st.session_state.form_fields:
-                    field_type = field.get('type', 'Unknown')
-                    if field_type not in field_types:
-                        field_types[field_type] = []
-                    field_types[field_type].append(field.get('name', 'Unnamed'))
-                
-                # Display field types and count
-                for field_type, field_names in field_types.items():
-                    with st.expander(f"{field_type} Fields ({len(field_names)})"):
-                        for name in field_names:
-                            st.write(f"• {name}")
-            else:
-                st.warning("⚠️ This form doesn't have fillable fields or couldn't be analyzed.")
-            
-            # Q&A Section for uploaded form
-            st.markdown("### ❓ Ask Questions About This Form")
-            st.markdown("Get help understanding your uploaded tax form.")
-            
-            upload_question = st.text_input(
-                "Ask a question about this form:", 
-                placeholder="What is this form used for? How do I fill section X?",
-                key="upload_question_input"
-            )
-            
-            if upload_question:
-                with st.spinner("Analyzing your question..."):
-                    # Generate response using the uploaded form context
-                    upload_response = generate_upload_form_response(
-                        upload_question, 
-                        uploaded_file.name,
-                        st.session_state.form_fields
-                    )
-                    
-                    st.markdown("### 💬 Response:")
-                    st.markdown(upload_response)
-            
-            # Download options
-            st.markdown("### 📥 Download Options")
-            col1, col2 = st.columns(2)
-            
-            with col1:
-                # Download original form
+                # Option to download the non-fillable form
                 st.download_button(
-                    label="📥 Download Original Form",
+                    label="📥 Download Form",
                     data=st.session_state.pdf_bytes,
-                    file_name=uploaded_file.name,
+                    file_name="pakistan_tax_form.pdf",
                     mime="application/pdf"
                 )
-        
-        else:
-            st.info("👆 Please upload a PDF tax form to get started.")
-            st.markdown("""
-            **What you can do with uploaded forms:**
-            - 📄 View and preview the form
-            - 📋 Get AI-generated summaries
-            - 🔍 Analyze form fields and sections
-            - ❓ Ask questions about the form
-            """)
+                
+                # Add a question and answer section for this form
+                st.markdown("### ❓ Questions about this form?")
+                form_question = st.text_input("Ask a question about this form:", key="form_question_input")
+                
+                if form_question:
+                    with st.spinner("Getting answer..."):
+                        form_response = tax_agent_response(form_question, tax_form_type=form_type)
+                        st.markdown(form_response)
 
 if __name__ == "__main__":
     main()
+