app.py

import gradio as gr
import groq
import os
import tempfile
import uuid
from dotenv import load_dotenv
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
import fitz  # PyMuPDF
import base64
from PIL import Image
import io
import requests
import json
from datetime import datetime, timedelta

# Load environment variables
load_dotenv()
client = groq.Client(api_key=os.getenv("GROQ_LEGAL_API_KEY"))
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Directory to store FAISS indexes
FAISS_INDEX_DIR = "faiss_indexes_legal"
if not os.path.exists(FAISS_INDEX_DIR):
    os.makedirs(FAISS_INDEX_DIR)

# Dictionary to store user-specific vectorstores
user_vectorstores = {}

# Custom CSS for styling
custom_css = """
:root {
    --primary-green: #10B981;
    --dark-green: #047857;
    --light-green: #D1FAE5;
    --medium-grey: #6B7280;
    --light-grey: #F3F4F6;
    --white: #FFFFFF;
    --border-grey: #E5E7EB;
}
body { background-color: var(--light-grey); font-family: 'Inter', sans-serif; }
.container { max-width: 1200px !important; margin: 0 auto !important; padding: 10px; }
.header { background-color: var(--white); border-bottom: 2px solid var(--border-grey); padding: 15px 0; margin-bottom: 20px; border-radius: 12px 12px 0 0; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
.header-title { color: var(--dark-green); font-size: 1.8rem; font-weight: 700; text-align: center; }
.header-subtitle { color: var(--medium-grey); font-size: 1rem; text-align: center; margin-top: 5px; }
.chat-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-grey) !important; min-height: 500px; }
.message-user { background-color: var(--primary-green) !important; color: var(--white) !important; border-radius: 18px 18px 4px 18px !important; padding: 12px 16px !important; margin-left: auto !important; max-width: 80% !important; }
.message-bot { background-color: var(--light-grey) !important; color: var(--medium-grey) !important; border-radius: 18px 18px 18px 4px !important; padding: 12px 16px !important; margin-right: auto !important; max-width: 80% !important; }
.input-area { background-color: var(--white) !important; border-top: 1px solid var(--border-grey) !important; padding: 12px !important; border-radius: 0 0 12px 12px !important; }
.input-box { border: 1px solid var(--border-grey) !important; border-radius: 24px !important; padding: 12px 16px !important; box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important; }
.send-btn { background-color: var(--primary-green) !important; border-radius: 24px !important; color: var(--white) !important; padding: 10px 20px !important; font-weight: 500 !important; }
.clear-btn { background-color: var(--light-grey) !important; border: 1px solid var(--border-grey) !important; border-radius: 24px !important; color: var(--medium-grey) !important; padding: 8px 16px !important; font-weight: 500 !important; }
.pdf-viewer-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-grey) !important; padding: 20px; }
.pdf-viewer-image { max-width: 100%; height: auto; border: 1px solid var(--border-grey); border-radius: 12px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
.stats-box { background-color: var(--light-green); padding: 10px; border-radius: 8px; margin-top: 10px; }
.tool-container { background-color: var(--white); border-radius: 12px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); padding: 15px; margin-bottom: 20px; }
.search-result { border-left: 3px solid var(--primary-green); padding-left: 10px; margin: 15px 0; }
.result-title { font-weight: bold; color: var(--dark-green); }
.result-meta { color: var(--medium-grey); font-size: 0.9rem; margin: 5px 0; }
.result-citation { font-style: italic; color: var(--medium-grey); }
.result-snippet { margin-top: 5px; }
"""

# Function to process PDF files
def process_pdf(pdf_file):
    if pdf_file is None:
        return None, "No file uploaded", {"page_images": [], "total_pages": 0, "total_words": 0}
    try:
        session_id = str(uuid.uuid4())
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
            temp_file.write(pdf_file)
            pdf_path = temp_file.name
        
        # Use fitz to extract text and images from PDF
        doc = fitz.open(pdf_path)
        texts = [page.get_text() for page in doc]
        page_images = []
        for page in doc:
            pix = page.get_pixmap()
            img_bytes = pix.tobytes("png")
            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
            page_images.append(img_base64)
        total_pages = len(doc)
        total_words = sum(len(text.split()) for text in texts)
        doc.close()

        # Split the extracted text into chunks
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        chunks = text_splitter.create_documents(texts)
        
        # Create and save the vector store
        vectorstore = FAISS.from_documents(chunks, embeddings)
        index_path = os.path.join(FAISS_INDEX_DIR, session_id)
        vectorstore.save_local(index_path)
        user_vectorstores[session_id] = vectorstore

        os.unlink(pdf_path)
        pdf_state = {"page_images": page_images, "total_pages": total_pages, "total_words": total_words}
        return session_id, f"✅ Successfully processed {len(chunks)} text chunks from your PDF", pdf_state
    except Exception as e:
        if "pdf_path" in locals() and os.path.exists(pdf_path):
            os.unlink(pdf_path)
        return None, f"Error processing PDF: {str(e)}", {"page_images": [], "total_pages": 0, "total_words": 0}

# Function to generate chatbot responses
def generate_response(message, session_id, model_name, history):
    if not message:
        return history
    try:
        context = ""
        if session_id and session_id in user_vectorstores:
            vectorstore = user_vectorstores[session_id]
            docs = vectorstore.similarity_search(message, k=3)
            if docs:
                context = "\n\nRelevant information from uploaded PDF:\n" + "\n".join(f"- {doc.page_content}" for doc in docs)
        
        # Check if it's a special command for case law search
        if message.lower().startswith("/case ") or message.lower().startswith("/search "):
            query = message.split(" ", 1)[1]
            case_results = search_case_law(query)
            if case_results:
                response = "**Case Law Search Results:**\n\n"
                for case in case_results[:5]:  # Limit to top 5 results
                    response += f"**{case['title']}**\n"
                    response += f"Link: {case.get('link', 'N/A')}\n"
                    response += f"Source: {case.get('source', 'N/A')}\n"
                    if case.get('snippet'):
                        response += f"Snippet: \"{case.get('snippet')}\"\n"
                    response += "\n"
                history.append((message, response))
                return history
            else:
                history.append((message, "No case law results found for your query."))
                return history
                
        system_prompt = "You are a legal assistant specializing in contract analysis and case law."
        system_prompt += " You can help with legal terminology, precedent cases, and statutory interpretation."
        if context:
            system_prompt += " Use the following context to answer the question if relevant: " + context
        completion = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": message}
            ],
            temperature=0.7,
            max_tokens=1024
        )
        response = completion.choices[0].message.content
        history.append((message, response))
        return history
    except Exception as e:
        history.append((message, f"Error generating response: {str(e)}"))
        return history

# Function to update the PDF viewer with the first page
def update_pdf_viewer(pdf_state):
    if not pdf_state["total_pages"]:
        return 0, None, "No PDF uploaded yet"
    try:
        # Decode the base64 image data for the first page
        img_data = base64.b64decode(pdf_state["page_images"][0])
        # Convert to a PIL image
        img = Image.open(io.BytesIO(img_data))
        return pdf_state["total_pages"], img, f"**Total Pages:** {pdf_state['total_pages']}\n**Total Words:** {pdf_state['total_words']}"
    except Exception as e:
        print(f"Error decoding image: {e}")
        return 0, None, "Error displaying PDF"

# Function to update the displayed PDF page based on the slider value
def update_image(page_num, pdf_state):
    if not pdf_state["total_pages"] or page_num < 1 or page_num > pdf_state["total_pages"]:
        return None
    try:
        # Decode the base64 image data
        img_data = base64.b64decode(pdf_state["page_images"][page_num - 1])
        # Convert to a PIL image
        img = Image.open(io.BytesIO(img_data))
        return img
    except Exception as e:
        print(f"Error decoding image: {e}")
        return None

# Legal-specific tools using SerpApi
def search_case_law(query, jurisdiction="", court="", date_range=""):
    """Search for relevant case law using SerpApi Google Scholar API"""
    serp_api_key = os.getenv("SERP_API_KEY", "")
    if not serp_api_key:
        print("SerpApi API key not configured")
        return []
    
    try:
        # Build the search query with legal focus
        search_query = query
        if jurisdiction:
            search_query += f" {jurisdiction} jurisdiction"
        if court:
            search_query += f" {court} court"
        if date_range:
            search_query += f" {date_range}"
            
        # Add legal terms to focus the search on case law
        search_query += " case law legal opinion precedent"
        
        # Call SerpApi Google Scholar API
        url = "https://serpapi.com/search"
        params = {
            "engine": "google_scholar",
            "q": search_query,
            "api_key": serp_api_key,
            "num": 10,  # Number of results
            "as_sdt": "6",  # Limit to case law
            "hl": "en"  # Language
        }
        
        response = requests.get(url, params=params)
        
        if response.status_code != 200:
            print(f"API Error: {response.status_code} - {response.text}")
            return []
            
        data = response.json()
        results = []
        
        # Process the organic results
        for case in data.get("organic_results", []):
            case_data = {
                "title": case.get("title", "Unknown Case"),
                "link": case.get("link", ""),
                "snippet": case.get("snippet", ""),
                "source": case.get("publication_info", {}).get("summary", "Unknown Source")
            }
            results.append(case_data)
            
        return results
    except Exception as e:
        print(f"Error in case law search: {e}")
        return []

def search_legislation(query, congress="current"):
    """Search for legislation using SerpApi"""
    serp_api_key = os.getenv("SERP_API_KEY", "")
    if not serp_api_key:
        print("SerpApi API key not configured")
        return []
    
    try:
        # Build search query for legislation
        search_query = f"{query} legislation law statute {congress} congress"
        
        url = "https://serpapi.com/search"
        params = {
            "engine": "google",
            "q": search_query,
            "api_key": serp_api_key,
            "num": 10,
            "hl": "en"
        }
        
        response = requests.get(url, params=params)
        if response.status_code != 200:
            print(f"API Error: {response.status_code} - {response.text}")
            return []
            
        data = response.json()
        
        results = []
        for result in data.get("organic_results", []):
            if any(term in result.get("link", "").lower() for term in [".gov", "congress", "legislation", "statute"]):
                results.append({
                    "title": result.get("title", "Unknown Legislation"),
                    "link": result.get("link", ""),
                    "snippet": result.get("snippet", "No description available"),
                    "source": result.get("source", "Unknown Source")
                })
            
        return results
    except Exception as e:
        print(f"Error in legislation search: {e}")
        return []

def analyze_legal_terms(text):
    """Extract and define legal terms from text"""
    try:
        # This would ideally use a legal term dictionary or API
        # For now, we'll use a simplified approach with Groq API
        system_prompt = """
        You are a legal assistant tasked with identifying legal terms in a text and providing their definitions.
        Extract up to 5 key legal terms and provide a brief definition for each.
        Format your response as a JSON array with "term" and "definition" keys.
        """
        
        completion = client.chat.completions.create(
            model="llama3-70b-8192",  # Using a capable model for this task
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": text}
            ],
            temperature=0.3,
            max_tokens=800
        )
        
        response = completion.choices[0].message.content
        
        # Try to parse as JSON (assuming the model outputs well-formed JSON)
        try:
            terms = json.loads(response)
            return terms
        except json.JSONDecodeError:
            # If not JSON, extract terms manually with simple parsing
            terms = []
            lines = response.split("\n")
            current_term = None
            current_def = ""
            
            for line in lines:
                if line.strip().startswith('"') or line.strip().startswith("'") or line.strip().startswith("{"):
                    continue
                    
                if ":" in line and not current_term:
                    parts = line.split(":", 1)
                    current_term = parts[0].strip().strip('"\'').replace("term", "").strip()
                    current_def = parts[1].strip().strip('"\'')
                elif current_term and line.strip():
                    current_def += " " + line.strip()
                elif current_term and not line.strip():
                    terms.append({"term": current_term, "definition": current_def})
                    current_term = None
                    current_def = ""
            
            if current_term:  # Don't forget the last one
                terms.append({"term": current_term, "definition": current_def})
                
            return terms if terms else [{"term": "Error", "definition": "Could not parse legal terms"}]
    except Exception as e:
        print(f"Error analyzing legal terms: {e}")
        return [{"term": "Error", "definition": f"An error occurred: {str(e)}"}]

def search_cases_with_form(query, jurisdiction, court_type, date_min, date_max):
    """Search case law with form inputs"""
    date_range = ""
    if date_min or date_max:
        date_range = f"{date_min or ''} to {date_max or ''}"
    
    results = search_case_law(query, jurisdiction, court_type, date_range)
    
    if not results:
        return "No results found. Try different search terms or criteria."
    
    # Format results as markdown
    markdown_results = "## Case Law Search Results\n\n"
    for i, case in enumerate(results, 1):
        markdown_results += f"### {i}. {case['title']}\n"
        markdown_results += f"**Source:** {case.get('source', 'Unknown Source')}\n"
        if case.get('snippet'):
            markdown_results += f"**Excerpt:** \"{case['snippet']}\"\n"
        if case.get('link'):
            markdown_results += f"[View Case]({case['link']})\n"
        markdown_results += "\n---\n\n"
    
    return markdown_results

# Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    current_session_id = gr.State(None)
    pdf_state = gr.State({"page_images": [], "total_pages": 0, "total_words": 0})
    gr.HTML("""
    <div class="header">
        <div class="header-title">Legal-Vision</div>
        <div class="header-subtitle">Analyze legal documents with Groq's LLM API.</div>
    </div>
    """)
    with gr.Row(elem_classes="container"):
        with gr.Column(scale=1, min_width=300):
            pdf_file = gr.File(label="Upload PDF Document", file_types=[".pdf"], type="binary")
            upload_button = gr.Button("Process PDF", variant="primary")
            pdf_status = gr.Markdown("No PDF uploaded yet")
            model_dropdown = gr.Dropdown(
                choices=["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"],
                value="llama3-70b-8192",
                label="Select Groq Model"
            )
            
            # Legal Tools Section
            gr.Markdown("### Legal Tools", elem_classes="tool-title")
            with gr.Group(elem_classes="tool-container"):
                with gr.Tabs():
                    with gr.TabItem("Case Law Search"):
                        case_search = gr.Textbox(label="Search Query", placeholder="Enter search terms")
                        with gr.Row():
                            jurisdiction = gr.Dropdown(
                                choices=["", "US Federal", "US State", "International", "UK", "EU", "Canadian"], 
                                value="", 
                                label="Jurisdiction"
                            )
                            court_type = gr.Dropdown(
                                choices=["", "Supreme Court", "Appellate Court", "Trial Court", "Administrative"], 
                                value="", 
                                label="Court Type"
                            )
                        with gr.Row():
                            date_min = gr.Textbox(label="From Date (YYYY-MM-DD)", placeholder="e.g., 2000-01-01")
                            date_max = gr.Textbox(label="To Date (YYYY-MM-DD)", placeholder="e.g., 2023-12-31")
                        case_search_btn = gr.Button("Search Cases")
                    
                    with gr.TabItem("Legal Term Analysis"):
                        legal_text = gr.Textbox(label="Text to Analyze", lines=5, placeholder="Enter legal text for term extraction and analysis")
                        term_analyze_btn = gr.Button("Analyze Terms")
        
        with gr.Column(scale=2, min_width=600):
            with gr.Tabs():
                with gr.TabItem("PDF Viewer"):
                    with gr.Column(elem_classes="pdf-viewer-container"):
                        page_slider = gr.Slider(minimum=1, maximum=1, step=1, label="Page Number", value=1)
                        pdf_image = gr.Image(label="PDF Page", type="pil", elem_classes="pdf-viewer-image")
                        stats_display = gr.Markdown("No PDF uploaded yet", elem_classes="stats-box")
                
                with gr.TabItem("Case Law Results"):
                    case_results = gr.Markdown("Search for case law to see results here")
                
                with gr.TabItem("Legal Terms"):
                    terms_results = gr.JSON(label="Extracted Legal Terms")
    
    # Chatbot at the bottom
    with gr.Row(elem_classes="container"):
        with gr.Column(scale=2, min_width=600):
            chatbot = gr.Chatbot(height=500, bubble_full_width=False, show_copy_button=True, elem_classes="chat-container")
            with gr.Row():
                msg = gr.Textbox(show_label=False, placeholder="Ask about your legal document or type /case to search case law...", scale=5)
                send_btn = gr.Button("Send", scale=1)
            clear_btn = gr.Button("Clear Conversation")
    
    # Event Handlers
    upload_button.click(
        process_pdf,
        inputs=[pdf_file],
        outputs=[current_session_id, pdf_status, pdf_state]
    ).then(
        update_pdf_viewer,
        inputs=[pdf_state],
        outputs=[page_slider, pdf_image, stats_display]
    )
    
    msg.submit(
        generate_response,
        inputs=[msg, current_session_id, model_dropdown, chatbot],
        outputs=[chatbot]
    ).then(lambda: "", None, [msg])
    
    send_btn.click(
        generate_response,
        inputs=[msg, current_session_id, model_dropdown, chatbot],
        outputs=[chatbot]
    ).then(lambda: "", None, [msg])
    
    clear_btn.click(
        lambda: ([], None, "No PDF uploaded yet", {"page_images": [], "total_pages": 0, "total_words": 0}, 0, None, "No PDF uploaded yet"),
        None,
        [chatbot, current_session_id, pdf_status, pdf_state, page_slider, pdf_image, stats_display]
    )
    
    page_slider.change(
        update_image,
        inputs=[page_slider, pdf_state],
        outputs=[pdf_image]
    )
    
    # Legal tool handlers
    case_search_btn.click(
        search_cases_with_form,
        inputs=[case_search, jurisdiction, court_type, date_min, date_max],
        outputs=[case_results]
    )
    
    term_analyze_btn.click(
        analyze_legal_terms,
        inputs=[legal_text],
        outputs=[terms_results]
    )

# Add footer with attribution
gr.HTML("""
<div style="text-align: center; margin-top: 20px; padding: 10px; color: #666; font-size: 0.8rem; border-top: 1px solid #eee;">
    Created by Calvin Allen Crawford
</div>
""")

# Launch the app
if __name__ == "__main__":
    demo.launch()