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("""
""")
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("""
Created by Calvin Allen Crawford
""")
# Launch the app
if __name__ == "__main__":
demo.launch()