app.py

import os
import json
from typing import Dict, List, Optional
import logging
from functools import lru_cache
import gradio as gr
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from vector_db import VectorDatabase
import re

# Enhanced logging for better debugging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s'
)

class RAGSystem:
    def __init__(self, vector_db: Optional[VectorDatabase] = None):
        logging.info("Initializing RAGSystem")
        
        self.vector_db = vector_db if vector_db else VectorDatabase()
        
        self.llm = None
        self.chain = None
        
        self.prompt_template = PromptTemplate(
            input_variables=["query", "context", "state", "statutes"],
            template="""You are a legal assistant specializing in tenant rights and landlord-tenant laws. Your goal is to provide accurate, detailed, and helpful answers grounded in legal authority. Use the provided statutes as the primary source when available. If no relevant statutes are found in the context, rely on your general knowledge to provide a pertinent and practical response, clearly indicating when you are doing so and prioritizing state-specific information over federal laws for state-specific queries.

Instructions:
- Use the context and statutes as the primary basis for your answer when available.
- For state-specific queries, prioritize statutes or legal principles from the specified state over federal laws.
- Cite relevant statutes (e.g., (AS § 34.03.220(a)(2))) explicitly in your answer when applicable.
- If multiple statutes apply, list all relevant ones.
- If no specific statute is found in the context, state this clearly (e.g., 'No specific statute was found in the provided context'), then provide a general answer based on common legal principles or practices, marked as such.
- Include practical examples or scenarios to enhance clarity and usefulness.
- Use bullet points or numbered lists for readability when appropriate.
- Maintain a professional and neutral tone.

Question: {query}
State: {state}
Statutes from context:
{statutes}
Context information:
{context}
Answer:"""
        )

    def initialize_llm(self, openai_api_key: str):
        if not openai_api_key:
            raise ValueError("OpenAI API key is required.")
        
        try:
            self.llm = ChatOpenAI(
                temperature=0.2,
                openai_api_key=openai_api_key,
                model_name="gpt-3.5-turbo",
                max_tokens=1500,
                request_timeout=30
            )
            logging.info("OpenAI LLM initialized successfully")
            
            self.chain = LLMChain(llm=self.llm, prompt=self.prompt_template)
            logging.info("LLMChain created successfully")
        except Exception as e:
            logging.error(f"Failed to initialize OpenAI LLM: {str(e)}")
            raise

    def extract_statutes(self, text: str) -> str:
        """
        Extract statute citations from the given text using a refined regex pattern.
        Returns a string of valid statutes, one per line, or a message if none are found.
        """
        statute_pattern = r'\((?:[A-Za-z\s]+\s*(?:Code|Laws|Statutes|CCP)\s*§\s*[0-9-]+(?:\([a-z0-9]+\))?|[A-Za-z0-9\s]+\s*§\s*[0-9-]+(?:\([a-z0-9]+\))?|[A-Z]{2,3}\s*§\s*[0-9-]+(?:\([a-z0-9]+\))?|[0-9]+\s*ILCS\s*[0-9]+/[0-9-]+(?:\([a-z0-9]+\))?|Title\s*[0-9]+\s*USC\s*§\s*[0-9]+-[0-9]+|[A-Za-z\s]+\s*Laws\s*[0-9]+\s*§\s*[0-9-]+(?:\([a-z0-9]+\))?|[A-Za-z\s]+\s*CCP\s*§\s*[0-9-]+(?:\([a-z0-9]+\))?)\)'
        statutes = re.findall(statute_pattern, text)
        
        valid_statutes = []
        for statute in statutes:
            if '§' in statute and any(char.isdigit() for char in statute) and not re.match(r'\([a-z]\)', statute) and 'found here' not in statute:
                valid_statutes.append(statute)
        
        if valid_statutes:
            seen = set()
            unique_statutes = [statute for statute in valid_statutes if not (statute in seen or seen.add(statute))]
            return "\n".join(unique_statutes)
        return "No statutes found in the context."

    @lru_cache(maxsize=100)
    def process_query(self, query: str, state: str, openai_api_key: str, n_results: int = 5) -> Dict[str, any]:
        logging.info(f"Processing query: '{query}' for state: '{state}' with n_results={n_results}")
        
        if not state:
            logging.warning("No state provided for query")
            return {
                "answer": "Please select a state to proceed with your query.",
                "context_used": "N/A"
            }
        
        if not openai_api_key:
            logging.warning("No OpenAI API key provided")
            return {
                "answer": "Please provide an OpenAI API key to proceed.",
                "context_used": "N/A"
            }
        
        if not self.llm or not self.chain:
            try:
                self.initialize_llm(openai_api_key)
            except Exception as e:
                logging.error(f"Failed to initialize LLM: {str(e)}")
                return {
                    "answer": f"Failed to initialize LLM with the provided API key: {str(e)}",
                    "context_used": "N/A"
                }
        
        # Query the vector database
        try:
            results = self.vector_db.query(query, state=state, n_results=n_results)
            logging.info("Vector database query successful")
            logging.debug(f"Query results: {json.dumps(results, indent=2)}")
        except Exception as e:
            logging.error(f"Vector database query failed: {str(e)}")
            results = {
                "document_results": {"documents": [[]], "metadatas": [[]]},
                "state_results": {"documents": [[]], "metadatas": [[]]}
            }
            logging.info("Applied safeguard: Using empty results due to vector DB failure")
        
        context_parts = []
        
        # Process document results
        if results["document_results"]["documents"] and results["document_results"]["documents"][0]:
            for i, doc in enumerate(results["document_results"]["documents"][0]):
                metadata = results["document_results"]["metadatas"][0][i]
                context_parts.append(f"[{metadata['state']} - Chunk {metadata.get('chunk_id', 'N/A')}] {doc}")
        else:
            logging.warning("No document results found in query response")
        
        # Process state summary results
        if results["state_results"]["documents"] and results["state_results"]["documents"][0]:
            for i, doc in enumerate(results["state_results"]["documents"][0]):
                metadata = results["state_results"]["metadatas"][0][i]
                context_parts.append(f"[{metadata['state']} - Summary] {doc}")
        else:
            logging.warning("No state summary results found in query response")
        
        context = "\n\n---\n\n".join(context_parts) if context_parts else "No relevant context found."
        
        logging.info(f"Raw context for query: {context}")
        
        if not context_parts:
            logging.info("No relevant context found for query")
            # Fallback to general knowledge
            statutes_from_context = "No statutes found in the context."
            try:
                answer = self.chain.invoke({
                    "query": query,
                    "context": "No specific legal documents available.",
                    "state": state,
                    "statutes": statutes_from_context
                })
                return {
                    "answer": answer['text'].strip(),
                    "context_used": context
                }
            except Exception as e:
                logging.error(f"LLM fallback processing failed: {str(e)}")
                return {
                    "answer": "I don’t have sufficient information to answer this accurately, and an error occurred while generating a general response. Please try again.",
                    "context_used": context
                }
        
        statutes_from_context = self.extract_statutes(context)
        logging.info(f"Statutes extracted from context: {statutes_from_context}")
        
        try:
            answer = self.chain.invoke({
                "query": query,
                "context": context,
                "state": state,
                "statutes": statutes_from_context
            })
            logging.info("LLM generated answer successfully")
            logging.debug(f"Raw answer text: {answer['text']}")
        except Exception as e:
            logging.error(f"LLM processing failed: {str(e)}")
            return {
                "answer": "An error occurred while generating the answer. Please try again.",
                "context_used": context
            }
        
        return {
            "answer": answer['text'].strip(),
            "context_used": context
        }

    def get_states(self) -> List[str]:
        try:
            states = self.vector_db.get_states()
            logging.info(f"Retrieved {len(states)} states from database")
            return states
        except Exception as e:
            logging.error(f"Failed to get states: {str(e)}")
            return []

    def load_pdf(self, pdf_path: str) -> int:
        try:
            num_states = self.vector_db.process_and_load_pdf(pdf_path)
            logging.info(f"Loaded PDF with {num_states} states")
            return num_states
        except Exception as e:
            logging.error(f"Failed to load PDF: {str(e)}")
            return 0

    def gradio_interface(self):
        def query_interface(api_key: str, query: str, state: str) -> str:
            if not api_key:
                logging.warning("No OpenAI API key provided in interface")
                return "⚠️ **Error:** Please provide an OpenAI API key to proceed."
            if not state:
                logging.warning("No state selected in interface")
                return "⚠️ **Error:** Please select a state to proceed with your query."
            result = self.process_query(query, state=state, openai_api_key=api_key)
            
            return f"### Answer:\n{result['answer']}"

        states = self.get_states()
        
        # Define the inputs
        api_key_input = gr.Textbox(
            label="Open AI API Key",
            type="password",
            placeholder="e.g., sk-abc123",
            elem_classes="input-field"
        )
        query_input = gr.Textbox(
            label="Query",
            placeholder="e.g., What are the eviction rules?",
            lines=3,
            elem_classes="input-field"
        )
        state_input = gr.Dropdown(
            label="Select a state (required)",
            choices=states,
            value=None,
            allow_custom_value=False,
            elem_classes="input-field"
        )

        # Define the example queries (only for query and state)
        example_queries = [
            ["What is the rent due date law?", "California"],
            ["What are the rules for security deposit returns?", "New York"],
            ["Can a landlord enter without notice?", "Texas"],
            ["What are the eviction notice requirements?", "Florida"],
            ["Are there rent control laws?", "Oregon"]
        ]

        custom_css = """
            .gr-form {
                max-width: 900px;
                margin: 0 auto;
                padding: 20px;
                background-color: #ffffff;
                border-radius: 15px;
                box-shadow: 0 4px 20px rgba(0, 0, 0, 0.1);
            }
            .gr-title {
                font-size: 3em;
                font-weight: bold;
                color: #2c3e50;
                text-align: center;
                margin-bottom: 10px;
            }
            .gr-description {
                font-size: 1.1em;
                color: #7f8c8d;
                text-align: center;
                margin-bottom: 30px;
            }
            .footnote {
                font-size: 1.1em;
                color: #34495e;
                text-align: center;
                margin-top: 20px;
                padding-top: 10px;
                border-top: 1px solid #e0e0e0;
            }
            .footnote a {
                color: #3498db;
                text-decoration: none;
                transition: color 0.3s ease;
            }
            .footnote a:hover {
                color: #2980b9;
                text-decoration: underline;
            }
            .gr-textbox, .gr-dropdown {
                border: 1px solid #dcdcdc !important;
                border-radius: 8px !important;
                padding: 12px !important;
                font-size: 1em !important;
                transition: border-color 0.3s ease;
            }
            .gr-textbox:focus, .gr-dropdown:focus {
                border-color: #3498db !important;
                box-shadow: 0 0 5px rgba(52, 152, 219, 0.3) !important;
            }
            .gr-textbox label, .gr-dropdown label {
                font-weight: 600;
                color: #34495e;
                margin-bottom: 8px;
            }
            .gr-button-primary {
                background-color: #3498db !important;
                border: none !important;
                padding: 12px 30px !important;
                font-weight: bold !important;
                font-size: 1em !important;
                border-radius: 8px !important;
                transition: background-color 0.3s ease, transform 0.1s ease;
            }
            .gr-button-primary:hover {
                background-color: #2980b9 !important;
                transform: translateY(-2px);
            }
            .gr-button-secondary {
                background-color: #95a5a6 !important;
                border: none !important;
                padding: 12px 30px !important;
                font-weight: bold !important;
                font-size: 1em !important;
                border-radius: 8px !important;
                transition: background-color 0.3s ease;
            }
            .gr-button-secondary:hover {
                background-color: #7f8c8d !important;
            }
            .output-markdown {
                background-color: #f9f9f9 !important;
                color: #2c3e50 !important;
                padding: 25px !important;
                border-radius: 10px !important;
                border: 1px solid #e0e0e0 !important;
                font-size: 1.1em !important;
                line-height: 1.8 !important;
                box-shadow: 0 2px 10px rgba(0, 0, 0, 0.05);
            }
            .gr-examples {
                background-color: #ecf0f1;
                padding: 15px;
                border-radius: 10px;
                margin-top: 20px;
            }
            .gr-examples table {
                background-color: transparent !important;
            }
            @media (prefers-color-scheme: dark) {
                .gr-form {
                    background-color: #2c3e50;
                    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.3);
                }
                .gr-title {
                    color: #ecf0f1;
                }
                .gr-description {
                    color: #bdc3c7;
                }
                .footnote {
                    color: #bdc3c7;
                    border-top: 1px solid #7f8c8d;
                }
                .footnote a {
                    color: #66b0ff;
                }
                .footnote a:hover {
                    color: #4a90e2;
                }
                .gr-textbox, .gr-dropdown {
                    background-color: #34495e !important;
                    color: #ecf0f1 !important;
                    border-color: #7f8c8d !important;
                }
                .gr-textbox label, .gr-dropdown label {
                    color: #ecf0f1;
                }
                .output-markdown {
                    background-color: #34495e !important;
                    color: #ecf0f1 !important;
                    border-color: #7f8c8d !important;
                }
                .gr-examples {
                    background-color: #3e5367;
                }
            }
            @media (max-width: 600px) {
                .gr-form {
                    padding: 15px;
                }
                .gr-title {
                    font-size: 2.2em;
                }
                .gr-description {
                    font-size: 1em;
                }
                .footnote {
                    font-size: 1em;
                }
                .gr-textbox, .gr-dropdown {
                    font-size: 0.9em !important;
                }
                .gr-button-primary, .gr-button-secondary {
                    padding: 10px 20px !important;
                    font-size: 0.9em !important;
                }
                .output-markdown {
                    font-size: 1em !important;
                    padding: 15px !important;
                }
            }
        """

        with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
            gr.Markdown(
                """
                # 🏠 Landlord-Tenant Rights Bot
                Ask questions about tenant rights and landlord-tenant laws based on state-specific legal documents. Provide your OpenAI API key, select a state, and enter your question below. You can get an API key from [OpenAI](https://platform.openai.com/api-keys).

                # <div class='footnote'>Developed by Nischal Subedi. Follow me on <a href='https://www.linkedin.com/in/nischal1/' target='_blank'>LinkedIn</a> or read my insights on <a href='https://datascientistinsights.substack.com/' target='_blank'>Substack</a>.</div>
                """
            )

            with gr.Column(elem_classes="gr-form"):
                api_key_input = gr.Textbox(
                    label="Open AI API Key",
                    type="password",
                    placeholder="e.g., sk-abc123",
                    elem_classes="input-field"
                )
                query_input = gr.Textbox(
                    label="Query",
                    placeholder="e.g., What are the eviction rules?",
                    lines=3,
                    elem_classes="input-field"
                )
                state_input = gr.Dropdown(
                    label="Select a state (required)",
                    choices=states,
                    value=None,
                    allow_custom_value=False,
                    elem_classes="input-field"
                )

                with gr.Row():
                    clear_button = gr.Button("Clear", variant="secondary")
                    submit_button = gr.Button("Submit", variant="primary")

                output = gr.Markdown(
                    label="Response",
                    elem_classes="output-markdown"
                )

                gr.Examples(
                    examples=example_queries,
                    inputs=[query_input, state_input],
                    outputs=output,
                    fn=query_interface,
                    examples_per_page=5
                )

            submit_button.click(
                fn=query_interface,
                inputs=[api_key_input, query_input, state_input],
                outputs=output
            )
            clear_button.click(
                fn=lambda: ("", "", None, ""),
                inputs=[],
                outputs=[api_key_input, query_input, state_input, output]
            )

        return demo

if __name__ == "__main__":
    try:
        rag = RAGSystem()
        
        pdf_path = "data/tenant-landlord.pdf"
        rag.load_pdf(pdf_path)
        
        demo = rag.gradio_interface()
        demo.launch(share=True)
        
    except Exception as e:
        logging.error(f"Main execution failed: {str(e)}")
        raise