updated UI and example table with enhancement in response logic

app.py

@@ -10,7 +10,11 @@ from langchain.chains import LLMChain
 from vector_db import VectorDatabase
 import re
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# 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):
@@ -18,25 +22,26 @@ class RAGSystem:
         
         self.vector_db = vector_db if vector_db else VectorDatabase()
         
-        # LLM and chain will be initialized later with user-provided API key
         self.llm = None
         self.chain = None
         
-        # Prompt template for statute-grounded answers
         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 that are explicitly grounded in the statutes provided in the context. Only use general knowledge to supplement the answer if the context lacks sufficient detail to fully answer the question, and clearly indicate when you are doing so.
+            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 information and the provided statutes as the primary source to answer the question.
-- Explicitly cite the relevant statute(s) (e.g., (AS § 34.03.220(a)(2))) in your answer to ground your response in the legal text.
-- If multiple statutes are relevant, cite all that apply.
-- If the context does not contain a relevant statute to answer the question, state that no specific statute was found and provide a general answer, clearly marking it as general knowledge.
-- Provide detailed answers with practical examples or scenarios when possible.
-- Use bullet points or numbered lists for clarity when applicable.
+- 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 found in context:
+Statutes from context:
 {statutes}
 Context information:
 {context}
@@ -44,7 +49,6 @@ Answer:"""
         )
 
     def initialize_llm(self, openai_api_key: str):
-        """Initialize the LLM and chain with the provided API key."""
         if not openai_api_key:
             raise ValueError("OpenAI API key is required.")
         
@@ -64,52 +68,60 @@ Answer:"""
             logging.error(f"Failed to initialize OpenAI LLM: {str(e)}")
             raise
 
-    def extract_statutes(self, context: str) -> str:
+    def extract_statutes(self, text: str) -> str:
         """
-        Extract statute citations from the context using a regex pattern.
-        Returns a string of statutes, one per line, or a message if none are found.
+        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-z0-9§\.\s-]+\)'
-        statutes = re.findall(statute_pattern, context)
-        if statutes:
-            return "\n".join(statutes)
+        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}")
+        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",
-                "statutes_found": "N/A"
+                "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",
-                "statutes_found": "N/A"
+                "context_used": "N/A"
             }
         
-        # Initialize LLM with the provided API key if not already initialized
         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",
-                    "statutes_found": "N/A"
+                    "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)}")
-            # Safeguard: Fallback to empty results if vector DB query fails
             results = {
                 "document_results": {"documents": [[]], "metadatas": [[]]},
                 "state_results": {"documents": [[]], "metadatas": [[]]}
@@ -118,51 +130,72 @@ Answer:"""
         
         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")
-            return {
-                "answer": "I don't have sufficient information in my database to answer this question accurately. However, I can provide some general information about tenant rights.",
-                "context_used": context,
-                "statutes_found": "N/A"
-            }
+            # 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
+                }
         
-        # Extract statutes from the context
-        statutes = self.extract_statutes(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
+                "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,
-                "statutes_found": statutes
+                "context_used": context
             }
         
         return {
             "answer": answer['text'].strip(),
-            "context_used": context,
-            "statutes_found": statutes
+            "context_used": context
         }
-    
+
     def get_states(self) -> List[str]:
         try:
             states = self.vector_db.get_states()
@@ -181,17 +214,42 @@ Answer:"""
             logging.error(f"Failed to load PDF: {str(e)}")
             return 0
 
-    def gradio_interface(self) -> gr.Interface:
+    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']}\n\n### Statutes Found:\n{result['statutes_found']}"
+            
+            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"],
@@ -200,9 +258,7 @@ Answer:"""
             ["Are there rent control laws?", "Oregon"]
         ]
 
-        # Custom CSS for a modern, readable, and responsive UI
         custom_css = """
-            /* General container styling */
             .gr-form {
                 max-width: 900px;
                 margin: 0 auto;
@@ -211,10 +267,8 @@ Answer:"""
                 border-radius: 15px;
                 box-shadow: 0 4px 20px rgba(0, 0, 0, 0.1);
             }
-
-            /* Title and description */
             .gr-title {
-                font-size: 3em; /* Increased font size for the title */
+                font-size: 3em;
                 font-weight: bold;
                 color: #2c3e50;
                 text-align: center;
@@ -226,8 +280,6 @@ Answer:"""
                 text-align: center;
                 margin-bottom: 30px;
             }
-
-            /* Footnote styling */
             .footnote {
                 font-size: 1.1em;
                 color: #34495e;
@@ -245,8 +297,6 @@ Answer:"""
                 color: #2980b9;
                 text-decoration: underline;
             }
-
-            /* Input fields */
             .gr-textbox, .gr-dropdown {
                 border: 1px solid #dcdcdc !important;
                 border-radius: 8px !important;
@@ -263,8 +313,6 @@ Answer:"""
                 color: #34495e;
                 margin-bottom: 8px;
             }
-
-            /* Buttons */
             .gr-button-primary {
                 background-color: #3498db !important;
                 border: none !important;
@@ -290,11 +338,9 @@ Answer:"""
             .gr-button-secondary:hover {
                 background-color: #7f8c8d !important;
             }
-
-            /* Output area */
             .output-markdown {
                 background-color: #f9f9f9 !important;
-                color: #2c3e50 !important; /* Dark text for readability */
+                color: #2c3e50 !important;
                 padding: 25px !important;
                 border-radius: 10px !important;
                 border: 1px solid #e0e0e0 !important;
@@ -302,8 +348,6 @@ Answer:"""
                 line-height: 1.8 !important;
                 box-shadow: 0 2px 10px rgba(0, 0, 0, 0.05);
             }
-
-            /* Examples section */
             .gr-examples {
                 background-color: #ecf0f1;
                 padding: 15px;
@@ -313,8 +357,6 @@ Answer:"""
             .gr-examples table {
                 background-color: transparent !important;
             }
-
-            /* Dark mode */
             @media (prefers-color-scheme: dark) {
                 .gr-form {
                     background-color: #2c3e50;
@@ -353,20 +395,18 @@ Answer:"""
                     background-color: #3e5367;
                 }
             }
-
-            /* Responsive design */
             @media (max-width: 600px) {
                 .gr-form {
                     padding: 15px;
                 }
                 .gr-title {
-                    font-size: 2.2em; /* Adjusted for smaller screens */
+                    font-size: 2.2em;
                 }
                 .gr-description {
                     font-size: 1em;
                 }
                 .footnote {
-                    font-size: 1em; /* Slightly smaller but still larger than before */
+                    font-size: 1em;
                 }
                 .gr-textbox, .gr-dropdown {
                     font-size: 0.9em !important;
@@ -382,40 +422,66 @@ Answer:"""
             }
         """
 
-        interface = gr.Interface(
-            fn=query_interface,
-            inputs=[
-                gr.Textbox(
-                    label="Enter your OpenAI API Key",
+        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"
-                ),
-                gr.Textbox(
-                    label="Enter your question about Landlord-Tenant laws",
+                )
+                query_input = gr.Textbox(
+                    label="Query",
                     placeholder="e.g., What are the eviction rules?",
                     lines=3,
                     elem_classes="input-field"
-                ),
-                gr.Dropdown(
+                )
+                state_input = gr.Dropdown(
                     label="Select a state (required)",
                     choices=states,
                     value=None,
                     allow_custom_value=False,
                     elem_classes="input-field"
                 )
-            ],
-            outputs=gr.Markdown(
-                label="Response",
-                elem_classes="output-markdown"
-            ),
-            title="🏠 Landlord-Tenant Rights Bot",
-            description="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).\n\n<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>",
-            examples=example_queries,
-            theme=gr.themes.Default(),
-            css=custom_css
-        )
-        return interface
+
+                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:
@@ -424,8 +490,8 @@ if __name__ == "__main__":
         pdf_path = "data/tenant-landlord.pdf"
         rag.load_pdf(pdf_path)
         
-        interface = rag.gradio_interface()
-        interface.launch(share=True)
+        demo = rag.gradio_interface()
+        demo.launch(share=True)
         
     except Exception as e:
         logging.error(f"Main execution failed: {str(e)}")