from huggingface_hub import snapshot_download
import os
import io
import time
import uuid
import tempfile
import numpy as np
import matplotlib.pyplot as plt
import pdfplumber
import spacy
import torch
import sqlite3
import uvicorn
import moviepy.editor as mp
from threading import Thread
from datetime import datetime, timedelta
from typing import List, Dict, Optional
from fastapi import FastAPI, File, UploadFile, Form, Depends, HTTPException, status, Header
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
import logging
from pydantic import BaseModel
from transformers import (
AutoTokenizer,
AutoModelForQuestionAnswering,
pipeline,
TrainingArguments,
Trainer
)
from sentence_transformers import SentenceTransformer
from passlib.context import CryptContext
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
import jwt
from dotenv import load_dotenv
# Import get_db_connection from auth
from auth import (
User, UserCreate, Token, get_current_active_user, authenticate_user,
create_access_token, hash_password, register_user, check_subscription_access,
SUBSCRIPTION_TIERS, JWT_EXPIRATION_DELTA, get_db_connection, update_auth_db_schema
)
# Add this import near the top with your other imports
from paypal_integration import (
create_user_subscription, verify_subscription_payment,
update_user_subscription, handle_subscription_webhook, initialize_database
)
from fastapi import Request # Add this if not already imported
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("app")
# Initialize the database
# Initialize FastAPI app
app = FastAPI(
title="Legal Document Analysis API",
description="API for analyzing legal documents, videos, and audio",
version="1.0.0"
)
# Set up CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["https://testing-78wtxfqt0-hardikkandpals-projects.vercel.app", "http://localhost:3000"], # Frontend URL
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
initialize_database()
try:
update_auth_db_schema()
logger.info("Database schema updated successfully")
except Exception as e:
logger.error(f"Database schema update error: {e}")
# Create static directory for file storage
os.makedirs("static", exist_ok=True)
os.makedirs("uploads", exist_ok=True)
os.makedirs("temp", exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")
# Set device for model inference
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
# Initialize chat history
chat_history = []
# Document context storage
document_contexts = {}
def store_document_context(task_id, text):
"""Store document text for later retrieval."""
document_contexts[task_id] = text
def load_document_context(task_id):
"""Load document text for a given task ID."""
return document_contexts.get(task_id, "")
def get_db_connection():
"""Get a connection to the SQLite database."""
db_path = os.path.join(os.path.dirname(__file__), "legal_analysis.db")
conn = sqlite3.connect(db_path)
conn.row_factory = sqlite3.Row
return conn
load_dotenv()
DB_PATH = os.getenv("DB_PATH", os.path.join(os.path.dirname(__file__), "data/user_data.db"))
os.makedirs(os.path.join(os.path.dirname(__file__), "data"), exist_ok=True)
def fine_tune_qa_model():
"""Fine-tunes a QA model on the CUAD dataset."""
print("Loading base model for fine-tuning...")
tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
# Load and preprocess CUAD dataset
print("Loading CUAD dataset...")
from datasets import load_dataset
try:
dataset = load_dataset("cuad")
except Exception as e:
print(f"Error loading CUAD dataset: {str(e)}")
print("Downloading CUAD dataset from alternative source...")
# Implement alternative dataset loading here
return tokenizer, model
print(f"Dataset loaded with {len(dataset['train'])} training examples")
# Preprocess the dataset
def preprocess_function(examples):
questions = [q.strip() for q in examples["question"]]
contexts = [c.strip() for c in examples["context"]]
inputs = tokenizer(
questions,
contexts,
max_length=384,
truncation="only_second",
stride=128,
return_overflowing_tokens=True,
return_offsets_mapping=True,
padding="max_length",
)
offset_mapping = inputs.pop("offset_mapping")
sample_map = inputs.pop("overflow_to_sample_mapping")
answers = examples["answers"]
start_positions = []
end_positions = []
for i, offset in enumerate(offset_mapping):
sample_idx = sample_map[i]
answer = answers[sample_idx]
start_char = answer["answer_start"][0] if len(answer["answer_start"]) > 0 else 0
end_char = start_char + len(answer["text"][0]) if len(answer["text"]) > 0 else 0
sequence_ids = inputs.sequence_ids(i)
# Find the start and end of the context
idx = 0
while sequence_ids[idx] != 1:
idx += 1
context_start = idx
while idx < len(sequence_ids) and sequence_ids[idx] == 1:
idx += 1
context_end = idx - 1
# If the answer is not fully inside the context, label is (0, 0)
if offset[context_start][0] > start_char or offset[context_end][1] < end_char:
start_positions.append(0)
end_positions.append(0)
else:
# Otherwise it's the start and end token positions
idx = context_start
while idx <= context_end and offset[idx][0] <= start_char:
idx += 1
start_positions.append(idx - 1)
idx = context_end
while idx >= context_start and offset[idx][1] >= end_char:
idx -= 1
end_positions.append(idx + 1)
inputs["start_positions"] = start_positions
inputs["end_positions"] = end_positions
return inputs
print("Preprocessing dataset...")
processed_dataset = dataset.map(
preprocess_function,
batched=True,
remove_columns=dataset["train"].column_names,
)
print("Splitting dataset...")
train_dataset = processed_dataset["train"]
val_dataset = processed_dataset["validation"]
train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "start_positions", "end_positions"])
val_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "start_positions", "end_positions"])
training_args = TrainingArguments(
output_dir="./fine_tuned_legal_qa",
evaluation_strategy="steps",
eval_steps=100,
learning_rate=2e-5,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
num_train_epochs=1,
weight_decay=0.01,
logging_steps=50,
save_steps=100,
load_best_model_at_end=True,
report_to=[]
)
print("✅ Starting fine tuning on CUAD QA dataset...")
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=val_dataset,
tokenizer=tokenizer,
)
trainer.train()
print("✅ Fine tuning completed. Saving model...")
model.save_pretrained("./fine_tuned_legal_qa")
tokenizer.save_pretrained("./fine_tuned_legal_qa")
return tokenizer, model
#############################
# Load NLP Models #
#############################
# Initialize model variables
nlp = None
summarizer = None
embedding_model = None
ner_model = None
speech_to_text = None
cuad_model = None
cuad_tokenizer = None
qa_model = None
# Add model caching functionality
import pickle
import os.path
#MODELS_CACHE_DIR = "c:\\Users\\hardi\\OneDrive\\Desktop\\New folder (7)\\doc-vid-analyze-main\\models_cache"
MODELS_CACHE_DIR = os.getenv("MODELS_CACHE_DIR", "models_cache")
os.makedirs(MODELS_CACHE_DIR, exist_ok=True)
def download_model_from_hub(model_id, subfolder=None):
"""Download a model from Hugging Face Hub"""
try:
local_dir = snapshot_download(
repo_id=model_id,
subfolder=subfolder,
local_dir=os.path.join(MODELS_CACHE_DIR, model_id.replace("/", "_"))
)
print(f"✅ Downloaded model {model_id} to {local_dir}")
return local_dir
except Exception as e:
print(f"⚠️ Error downloading model {model_id}: {str(e)}")
return None
def save_model_to_cache(model, model_name):
"""Save a model to the cache directory"""
try:
cache_path = os.path.join(MODELS_CACHE_DIR, f"{model_name}.pkl")
with open(cache_path, 'wb') as f:
pickle.dump(model, f)
print(f"✅ Saved {model_name} to cache")
return True
except Exception as e:
print(f"⚠️ Failed to save {model_name} to cache: {str(e)}")
return False
def load_model_from_cache(model_name):
"""Load a model from the cache directory"""
try:
cache_path = os.path.join(MODELS_CACHE_DIR, f"{model_name}.pkl")
if os.path.exists(cache_path):
with open(cache_path, 'rb') as f:
model = pickle.load(f)
print(f"✅ Loaded {model_name} from cache")
return model
return None
except Exception as e:
print(f"⚠️ Failed to load {model_name} from cache: {str(e)}")
return None
# Add a flag to control model loading
LOAD_MODELS = os.getenv("LOAD_MODELS", "True").lower() in ("true", "1", "t")
try:
if LOAD_MODELS:
# Try to load SpaCy from cache first
nlp = load_model_from_cache("spacy_model")
if nlp is None:
try:
nlp = spacy.load("en_core_web_sm")
save_model_to_cache(nlp, "spacy_model")
except:
print("⚠️ SpaCy model not found, downloading...")
spacy.cli.download("en_core_web_sm")
nlp = spacy.load("en_core_web_sm")
save_model_to_cache(nlp, "spacy_model")
print("✅ Loading NLP models...")
# Load the summarizer with caching
print("Loading summarizer model...")
summarizer = load_model_from_cache("summarizer_model")
if summarizer is None:
try:
summarizer = pipeline("summarization", model="facebook/bart-large-cnn",
device=0 if torch.cuda.is_available() else -1)
save_model_to_cache(summarizer, "summarizer_model")
print("✅ Summarizer loaded successfully")
except Exception as e:
print(f"⚠️ Error loading summarizer: {str(e)}")
try:
print("Trying alternative summarizer model...")
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6",
device=0 if torch.cuda.is_available() else -1)
save_model_to_cache(summarizer, "summarizer_model")
print("✅ Alternative summarizer loaded successfully")
except Exception as e2:
print(f"⚠️ Error loading alternative summarizer: {str(e2)}")
summarizer = None
# Load the embedding model with caching
print("Loading embedding model...")
embedding_model = load_model_from_cache("embedding_model")
if embedding_model is None:
try:
embedding_model = SentenceTransformer("all-mpnet-base-v2", device=device)
save_model_to_cache(embedding_model, "embedding_model")
print("✅ Embedding model loaded successfully")
except Exception as e:
print(f"⚠️ Error loading embedding model: {str(e)}")
embedding_model = None
# Load the NER model with caching
print("Loading NER model...")
ner_model = load_model_from_cache("ner_model")
if ner_model is None:
try:
ner_model = pipeline("ner", model="dslim/bert-base-NER",
device=0 if torch.cuda.is_available() else -1)
save_model_to_cache(ner_model, "ner_model")
print("✅ NER model loaded successfully")
except Exception as e:
print(f"⚠️ Error loading NER model: {str(e)}")
ner_model = None
# Speech to text model with caching
print("Loading speech to text model...")
speech_to_text = load_model_from_cache("speech_to_text_model")
if speech_to_text is None:
try:
speech_to_text = pipeline("automatic-speech-recognition",
model="openai/whisper-medium",
chunk_length_s=30,
device_map="auto" if torch.cuda.is_available() else "cpu")
save_model_to_cache(speech_to_text, "speech_to_text_model")
print("✅ Speech to text model loaded successfully")
except Exception as e:
print(f"⚠️ Error loading speech to text model: {str(e)}")
speech_to_text = None
# Load the fine-tuned model with caching
print("Loading fine-tuned CUAD QA model...")
cuad_model = load_model_from_cache("cuad_model")
cuad_tokenizer = load_model_from_cache("cuad_tokenizer")
if cuad_model is None or cuad_tokenizer is None:
try:
cuad_tokenizer = AutoTokenizer.from_pretrained("hardik8588/fine-tuned-legal-qa")
from transformers import AutoModelForQuestionAnswering
cuad_model = AutoModelForQuestionAnswering.from_pretrained("hardik8588/fine-tuned-legal-qa")
cuad_model.to(device)
save_model_to_cache(cuad_tokenizer, "cuad_tokenizer")
save_model_to_cache(cuad_model, "cuad_model")
print("✅ Successfully loaded fine-tuned model")
except Exception as e:
print(f"⚠️ Error loading fine-tuned model: {str(e)}")
print("⚠️ Falling back to pre-trained model...")
try:
cuad_tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
from transformers import AutoModelForQuestionAnswering
cuad_model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
cuad_model.to(device)
save_model_to_cache(cuad_tokenizer, "cuad_tokenizer")
save_model_to_cache(cuad_model, "cuad_model")
print("✅ Pre-trained model loaded successfully")
except Exception as e2:
print(f"⚠️ Error loading pre-trained model: {str(e2)}")
cuad_model = None
cuad_tokenizer = None
# Load a general QA model with caching
print("Loading general QA model...")
qa_model = load_model_from_cache("qa_model")
if qa_model is None:
try:
qa_model = pipeline("question-answering", model="deepset/roberta-base-squad2")
save_model_to_cache(qa_model, "qa_model")
print("✅ QA model loaded successfully")
except Exception as e:
print(f"⚠️ Error loading QA model: {str(e)}")
qa_model = None
print("✅ All models loaded successfully")
else:
print("⚠️ Model loading skipped (LOAD_MODELS=False)")
except Exception as e:
print(f"⚠️ Error loading models: {str(e)}")
# Instead of raising an error, set fallback behavior
nlp = None
summarizer = None
embedding_model = None
ner_model = None
speech_to_text = None
cuad_model = None
cuad_tokenizer = None
qa_model = None
print("⚠️ Running with limited functionality due to model loading errors")
def legal_chatbot(user_input, context):
"""Uses a real NLP model for legal Q&A."""
global chat_history
chat_history.append({"role": "user", "content": user_input})
response = qa_model(question=user_input, context=context)["answer"]
chat_history.append({"role": "assistant", "content": response})
return response
def extract_text_from_pdf(pdf_file):
"""Extracts text from a PDF file using pdfplumber."""
try:
# Suppress pdfplumber warnings about CropBox
import logging
logging.getLogger("pdfminer").setLevel(logging.ERROR)
with pdfplumber.open(pdf_file) as pdf:
print(f"Processing PDF with {len(pdf.pages)} pages")
text = ""
for i, page in enumerate(pdf.pages):
page_text = page.extract_text() or ""
text += page_text + "\n"
if (i + 1) % 10 == 0: # Log progress every 10 pages
print(f"Processed {i + 1} pages...")
print(f"✅ PDF text extraction complete: {len(text)} characters extracted")
return text.strip() if text else None
except Exception as e:
print(f"❌ PDF extraction error: {str(e)}")
raise HTTPException(status_code=400, detail=f"PDF extraction failed: {str(e)}")
def process_video_to_text(video_file_path):
"""Extract audio from video and convert to text."""
try:
print(f"Processing video file at {video_file_path}")
temp_audio_path = os.path.join("temp", "extracted_audio.wav")
video = mp.VideoFileClip(video_file_path)
video.audio.write_audiofile(temp_audio_path, codec='pcm_s16le')
print(f"Audio extracted to {temp_audio_path}")
result = speech_to_text(temp_audio_path)
transcript = result["text"]
print(f"Transcription completed: {len(transcript)} characters")
if os.path.exists(temp_audio_path):
os.remove(temp_audio_path)
return transcript
except Exception as e:
print(f"Error in video processing: {str(e)}")
raise HTTPException(status_code=400, detail=f"Video processing failed: {str(e)}")
def process_audio_to_text(audio_file_path):
"""Process audio file and convert to text."""
try:
print(f"Processing audio file at {audio_file_path}")
result = speech_to_text(audio_file_path)
transcript = result["text"]
print(f"Transcription completed: {len(transcript)} characters")
return transcript
except Exception as e:
print(f"Error in audio processing: {str(e)}")
raise HTTPException(status_code=400, detail=f"Audio processing failed: {str(e)}")
def extract_named_entities(text):
"""Extracts named entities from legal text."""
max_length = 10000
entities = []
for i in range(0, len(text), max_length):
chunk = text[i:i+max_length]
doc = nlp(chunk)
entities.extend([{"entity": ent.text, "label": ent.label_} for ent in doc.ents])
return entities
def analyze_risk(text):
"""Analyzes legal risk in the document using keyword-based analysis."""
risk_keywords = {
"Liability": ["liability", "responsible", "responsibility", "legal obligation"],
"Termination": ["termination", "breach", "contract end", "default"],
"Indemnification": ["indemnification", "indemnify", "hold harmless", "compensate", "compensation"],
"Payment Risk": ["payment", "terms", "reimbursement", "fee", "schedule", "invoice", "money"],
"Insurance": ["insurance", "coverage", "policy", "claims"],
}
risk_scores = {category: 0 for category in risk_keywords}
lower_text = text.lower()
for category, keywords in risk_keywords.items():
for keyword in keywords:
risk_scores[category] += lower_text.count(keyword.lower())
return risk_scores
def extract_context_for_risk_terms(text, risk_keywords, window=1):
"""
Extracts and summarizes the context around risk terms.
"""
doc = nlp(text)
sentences = list(doc.sents)
risk_contexts = {category: [] for category in risk_keywords}
for i, sent in enumerate(sentences):
sent_text_lower = sent.text.lower()
for category, details in risk_keywords.items():
for keyword in details["keywords"]:
if keyword.lower() in sent_text_lower:
start_idx = max(0, i - window)
end_idx = min(len(sentences), i + window + 1)
context_chunk = " ".join([s.text for s in sentences[start_idx:end_idx]])
risk_contexts[category].append(context_chunk)
summarized_contexts = {}
for category, contexts in risk_contexts.items():
if contexts:
combined_context = " ".join(contexts)
try:
summary_result = summarizer(combined_context, max_length=100, min_length=30, do_sample=False)
summary = summary_result[0]['summary_text']
except Exception as e:
summary = "Context summarization failed."
summarized_contexts[category] = summary
else:
summarized_contexts[category] = "No contextual details found."
return summarized_contexts
def get_detailed_risk_info(text):
"""
Returns detailed risk information by merging risk scores with descriptive details
and contextual summaries from the document.
"""
risk_details = {
"Liability": {
"description": "Liability refers to the legal responsibility for losses or damages.",
"common_concerns": "Broad liability clauses may expose parties to unforeseen risks.",
"recommendations": "Review and negotiate clear limits on liability.",
"example": "E.g., 'The party shall be liable for direct damages due to negligence.'"
},
"Termination": {
"description": "Termination involves conditions under which a contract can be ended.",
"common_concerns": "Unilateral termination rights or ambiguous conditions can be risky.",
"recommendations": "Ensure termination clauses are balanced and include notice periods.",
"example": "E.g., 'Either party may terminate the agreement with 30 days notice.'"
},
"Indemnification": {
"description": "Indemnification requires one party to compensate for losses incurred by the other.",
"common_concerns": "Overly broad indemnification can shift significant risk.",
"recommendations": "Negotiate clear limits and carve-outs where necessary.",
"example": "E.g., 'The seller shall indemnify the buyer against claims from product defects.'"
},
"Payment Risk": {
"description": "Payment risk pertains to terms regarding fees, schedules, and reimbursements.",
"common_concerns": "Vague payment terms or hidden charges increase risk.",
"recommendations": "Clarify payment conditions and include penalties for delays.",
"example": "E.g., 'Payments must be made within 30 days, with a 2% late fee thereafter.'"
},
"Insurance": {
"description": "Insurance risk covers the adequacy and scope of required coverage.",
"common_concerns": "Insufficient insurance can leave parties exposed in unexpected events.",
"recommendations": "Review insurance requirements to ensure they meet the risk profile.",
"example": "E.g., 'The contractor must maintain liability insurance with at least $1M coverage.'"
}
}
risk_scores = analyze_risk(text)
risk_keywords_context = {
"Liability": {"keywords": ["liability", "responsible", "responsibility", "legal obligation"]},
"Termination": {"keywords": ["termination", "breach", "contract end", "default"]},
"Indemnification": {"keywords": ["indemnification", "indemnify", "hold harmless", "compensate", "compensation"]},
"Payment Risk": {"keywords": ["payment", "terms", "reimbursement", "fee", "schedule", "invoice", "money"]},
"Insurance": {"keywords": ["insurance", "coverage", "policy", "claims"]}
}
risk_contexts = extract_context_for_risk_terms(text, risk_keywords_context, window=1)
detailed_info = {}
for risk_term, score in risk_scores.items():
if score > 0:
info = risk_details.get(risk_term, {"description": "No details available."})
detailed_info[risk_term] = {
"score": score,
"description": info.get("description", ""),
"common_concerns": info.get("common_concerns", ""),
"recommendations": info.get("recommendations", ""),
"example": info.get("example", ""),
"context_summary": risk_contexts.get(risk_term, "No context available.")
}
return detailed_info
def analyze_contract_clauses(text):
"""Analyzes contract clauses using the fine-tuned CUAD QA model."""
max_length = 512
step = 256
clauses_detected = []
try:
clause_types = list(cuad_model.config.id2label.values())
except Exception as e:
clause_types = [
"Obligations of Seller", "Governing Law", "Termination", "Indemnification",
"Confidentiality", "Insurance", "Non-Compete", "Change of Control",
"Assignment", "Warranty", "Limitation of Liability", "Arbitration",
"IP Rights", "Force Majeure", "Revenue/Profit Sharing", "Audit Rights"
]
chunks = [text[i:i+max_length] for i in range(0, len(text), step) if i+step < len(text)]
for chunk in chunks:
inputs = cuad_tokenizer(chunk, return_tensors="pt", truncation=True, max_length=512).to(device)
with torch.no_grad():
outputs = cuad_model(**inputs)
predictions = torch.sigmoid(outputs.start_logits).cpu().numpy()[0]
for idx, confidence in enumerate(predictions):
if confidence > 0.5 and idx < len(clause_types):
clauses_detected.append({"type": clause_types[idx], "confidence": float(confidence)})
aggregated_clauses = {}
for clause in clauses_detected:
clause_type = clause["type"]
if clause_type not in aggregated_clauses or clause["confidence"] > aggregated_clauses[clause_type]["confidence"]:
aggregated_clauses[clause_type] = clause
return list(aggregated_clauses.values())
def summarize_text(text):
"""Summarizes legal text using the summarizer model."""
try:
if summarizer is None:
return "Basic analysis (NLP models not available)"
# Split text into chunks if it's too long
max_chunk_size = 1024
if len(text) > max_chunk_size:
chunks = [text[i:i+max_chunk_size] for i in range(0, len(text), max_chunk_size)]
summaries = []
for chunk in chunks:
summary = summarizer(chunk, max_length=100, min_length=30, do_sample=False)
summaries.append(summary[0]['summary_text'])
return " ".join(summaries)
else:
summary = summarizer(text, max_length=100, min_length=30, do_sample=False)
return summary[0]['summary_text']
except Exception as e:
print(f"Error in summarization: {str(e)}")
return "Summarization failed. Please try again later."
@app.post("/analyze_legal_document")
async def analyze_legal_document(
file: UploadFile = File(...),
current_user: User = Depends(get_current_active_user)
):
"""Analyzes a legal document (PDF) and returns insights based on subscription tier."""
try:
# Calculate file size in MB
file_content = await file.read()
file_size_mb = len(file_content) / (1024 * 1024)
# Check subscription access for document analysis
check_subscription_access(current_user, "document_analysis", file_size_mb)
print(f"Processing file: {file.filename}")
# Create a temporary file to store the uploaded PDF
with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp:
tmp.write(file_content)
tmp_path = tmp.name
# Extract text from PDF
text = extract_text_from_pdf(tmp_path)
# Clean up the temporary file
os.unlink(tmp_path)
if not text:
raise HTTPException(status_code=400, detail="Could not extract text from PDF")
# Generate a task ID
task_id = str(uuid.uuid4())
# Store document context for later retrieval
store_document_context(task_id, text)
# Basic analysis available to all tiers
summary = summarize_text(text)
entities = extract_named_entities(text)
risk_scores = analyze_risk(text)
# Prepare response based on subscription tier
response = {
"task_id": task_id,
"summary": summary,
"entities": entities,
"risk_assessment": risk_scores,
"subscription_tier": current_user.subscription_tier
}
# Add premium features if user has access
if current_user.subscription_tier == "premium_tier":
# Add detailed risk assessment
if "detailed_risk_assessment" in SUBSCRIPTION_TIERS[current_user.subscription_tier]["features"]:
detailed_risk = get_detailed_risk_info(text)
response["detailed_risk_assessment"] = detailed_risk
# Add contract clause analysis
if "contract_clause_analysis" in SUBSCRIPTION_TIERS[current_user.subscription_tier]["features"]:
clauses = analyze_contract_clauses(text)
response["contract_clauses"] = clauses
return response
except Exception as e:
print(f"Error analyzing document: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error analyzing document: {str(e)}")
# Add this function to check resource limits based on subscription tier
def check_resource_limits(user: User, resource_type: str, size_mb: float = None, count: int = 1):
"""
Check if the user has exceeded their subscription limits for a specific resource
Args:
user: The user making the request
resource_type: Type of resource (document, video, audio)
size_mb: Size of the resource in MB
count: Number of resources being used (default 1)
Returns:
bool: True if within limits, raises HTTPException otherwise
"""
# Get the user's subscription tier limits
tier = user.subscription_tier
tier_limits = SUBSCRIPTION_TIERS.get(tier, SUBSCRIPTION_TIERS["free_tier"])["limits"]
# Check size limits
if size_mb is not None:
if resource_type == "document" and size_mb > tier_limits["document_size_mb"]:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"Document size exceeds the {tier_limits['document_size_mb']}MB limit for your {tier} subscription"
)
elif resource_type == "video" and size_mb > tier_limits["video_size_mb"]:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"Video size exceeds the {tier_limits['video_size_mb']}MB limit for your {tier} subscription"
)
elif resource_type == "audio" and size_mb > tier_limits["audio_size_mb"]:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"Audio size exceeds the {tier_limits['audio_size_mb']}MB limit for your {tier} subscription"
)
# Check monthly document count
if resource_type == "document":
# Get current month and year
now = datetime.now()
month, year = now.month, now.year
# Check usage stats for current month
conn = get_db_connection()
cursor = conn.cursor()
cursor.execute(
"SELECT analyses_used FROM usage_stats WHERE user_id = ? AND month = ? AND year = ?",
(user.id, month, year)
)
result = cursor.fetchone()
current_usage = result[0] if result else 0
# Check if adding this usage would exceed the limit
if current_usage + count > tier_limits["documents_per_month"]:
conn.close()
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"You have reached your monthly limit of {tier_limits['documents_per_month']} document analyses for your {tier} subscription"
)
# Update usage stats
if result:
cursor.execute(
"UPDATE usage_stats SET analyses_used = ? WHERE user_id = ? AND month = ? AND year = ?",
(current_usage + count, user.id, month, year)
)
else:
usage_id = str(uuid.uuid4())
cursor.execute(
"INSERT INTO usage_stats (id, user_id, month, year, analyses_used) VALUES (?, ?, ?, ?, ?)",
(usage_id, user.id, month, year, count)
)
conn.commit()
conn.close()
# Check if feature is available in the tier
if resource_type == "video" and tier_limits["video_size_mb"] == 0:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"Video analysis is not available in your {tier} subscription"
)
if resource_type == "audio" and tier_limits["audio_size_mb"] == 0:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"Audio analysis is not available in your {tier} subscription"
)
return True
@app.post("/analyze_legal_video")
async def analyze_legal_video(
file: UploadFile = File(...),
current_user: User = Depends(get_current_active_user)
):
"""Analyzes legal video by transcribing and analyzing the transcript."""
try:
# Calculate file size in MB
file_content = await file.read()
file_size_mb = len(file_content) / (1024 * 1024)
# Check subscription access for video analysis
check_subscription_access(current_user, "video_analysis", file_size_mb)
print(f"Processing video file: {file.filename}")
# Create a temporary file to store the uploaded video
with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp:
tmp.write(file_content)
tmp_path = tmp.name
# Process video to extract transcript
transcript = process_video_to_text(tmp_path)
# Clean up the temporary file
os.unlink(tmp_path)
if not transcript:
raise HTTPException(status_code=400, detail="Could not extract transcript from video")
# Generate a task ID
task_id = str(uuid.uuid4())
# Store document context for later retrieval
store_document_context(task_id, transcript)
# Basic analysis
summary = summarize_text(transcript)
entities = extract_named_entities(transcript)
risk_scores = analyze_risk(transcript)
# Prepare response
response = {
"task_id": task_id,
"transcript": transcript,
"summary": summary,
"entities": entities,
"risk_assessment": risk_scores,
"subscription_tier": current_user.subscription_tier
}
# Add premium features if user has access
if current_user.subscription_tier == "premium_tier":
# Add detailed risk assessment
if "detailed_risk_assessment" in SUBSCRIPTION_TIERS[current_user.subscription_tier]["features"]:
detailed_risk = get_detailed_risk_info(transcript)
response["detailed_risk_assessment"] = detailed_risk
return response
except Exception as e:
print(f"Error analyzing video: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error analyzing video: {str(e)}")
@app.post("/legal_chatbot/{task_id}")
async def chat_with_document(
task_id: str,
question: str = Form(...),
current_user: User = Depends(get_current_active_user)
):
"""Chat with a document using the legal chatbot."""
try:
# Check if user has access to chatbot feature
if "chatbot" not in SUBSCRIPTION_TIERS[current_user.subscription_tier]["features"]:
raise HTTPException(
status_code=403,
detail=f"The chatbot feature is not available in your {current_user.subscription_tier} subscription. Please upgrade to access this feature."
)
# Check if document context exists
context = load_document_context(task_id)
if not context:
raise HTTPException(status_code=404, detail="Document context not found. Please analyze a document first.")
# Use the chatbot to answer the question
answer = legal_chatbot(question, context)
return {"answer": answer, "chat_history": chat_history}
except Exception as e:
print(f"Error in chatbot: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error in chatbot: {str(e)}")
@app.get("/")
async def root():
"""Root endpoint that returns a welcome message."""
return HTMLResponse(content="""
Legal Document Analysis API
Legal Document Analysis API
Welcome to the Legal Document Analysis API. This API provides tools for analyzing legal documents, videos, and audio.
Available Endpoints:
POST /analyze_legal_document - Analyze a legal document (PDF)
POST /analyze_legal_video - Analyze a legal video
POST /analyze_legal_audio - Analyze legal audio
POST /legal_chatbot/{task_id} - Chat with a document
POST /register - Register a new user
POST /token - Login to get an access token
GET /users/me - Get current user information
POST /subscribe/{tier} - Subscribe to a plan
For more details, visit the API documentation.
""")
@app.post("/register", response_model=Token)
async def register_new_user(user_data: UserCreate):
"""Register a new user with a free subscription"""
try:
success, result = register_user(user_data.email, user_data.password)
if not success:
raise HTTPException(status_code=400, detail=result)
return {"access_token": result["access_token"], "token_type": "bearer"}
except HTTPException:
# Re-raise HTTP exceptions
raise
except Exception as e:
print(f"Registration error: {str(e)}")
raise HTTPException(status_code=500, detail=f"Registration failed: {str(e)}")
@app.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
"""Endpoint for OAuth2 token generation"""
try:
# Add debug logging
logger.info(f"Token request for username: {form_data.username}")
user = authenticate_user(form_data.username, form_data.password)
if not user:
logger.warning(f"Authentication failed for: {form_data.username}")
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
headers={"WWW-Authenticate": "Bearer"},
)
access_token = create_access_token(user.id)
if not access_token:
logger.error(f"Failed to create access token for user: {user.id}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Could not create access token",
)
logger.info(f"Login successful for: {form_data.username}")
return {"access_token": access_token, "token_type": "bearer"}
except Exception as e:
logger.error(f"Token endpoint error: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Login error: {str(e)}",
)
@app.get("/debug/token")
async def debug_token(authorization: str = Header(None)):
"""Debug endpoint to check token validity"""
try:
if not authorization:
return {"valid": False, "error": "No authorization header provided"}
# Extract token from Authorization header
scheme, token = authorization.split()
if scheme.lower() != 'bearer':
return {"valid": False, "error": "Not a bearer token"}
# Log the token for debugging
logger.info(f"Debugging token: {token[:10]}...")
# Try to validate the token
try:
user = await get_current_active_user(token)
return {"valid": True, "user_id": user.id, "email": user.email}
except Exception as e:
return {"valid": False, "error": str(e)}
except Exception as e:
return {"valid": False, "error": f"Token debug error: {str(e)}"}
@app.post("/login")
async def api_login(email: str, password: str):
success, result = login_user(email, password)
if not success:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=result
)
return result
@app.get("/health")
def health_check():
"""Simple health check endpoint to verify the API is running"""
return {"status": "ok", "message": "API is running"}
@app.get("/users/me", response_model=User)
async def read_users_me(current_user: User = Depends(get_current_active_user)):
return current_user
@app.post("/analyze_legal_audio")
async def analyze_legal_audio(
file: UploadFile = File(...),
current_user: User = Depends(get_current_active_user)
):
"""Analyzes legal audio by transcribing and analyzing the transcript."""
try:
# Calculate file size in MB
file_content = await file.read()
file_size_mb = len(file_content) / (1024 * 1024)
# Check subscription access for audio analysis
check_subscription_access(current_user, "audio_analysis", file_size_mb)
print(f"Processing audio file: {file.filename}")
# Create a temporary file to store the uploaded audio
with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp:
tmp.write(file_content)
tmp_path = tmp.name
# Process audio to extract transcript
transcript = process_audio_to_text(tmp_path)
# Clean up the temporary file
os.unlink(tmp_path)
if not transcript:
raise HTTPException(status_code=400, detail="Could not extract transcript from audio")
# Generate a task ID
task_id = str(uuid.uuid4())
# Store document context for later retrieval
store_document_context(task_id, transcript)
# Basic analysis
summary = summarize_text(transcript)
entities = extract_named_entities(transcript)
risk_scores = analyze_risk(transcript)
# Prepare response
response = {
"task_id": task_id,
"transcript": transcript,
"summary": summary,
"entities": entities,
"risk_assessment": risk_scores,
"subscription_tier": current_user.subscription_tier
}
# Add premium features if user has access
if current_user.subscription_tier == "premium_tier": # Change from premium_tier to premium
# Add detailed risk assessment
if "detailed_risk_assessment" in SUBSCRIPTION_TIERS[current_user.subscription_tier]["features"]:
detailed_risk = get_detailed_risk_info(transcript)
response["detailed_risk_assessment"] = detailed_risk
return response
except Exception as e:
print(f"Error analyzing audio: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error analyzing audio: {str(e)}")
# Add these new endpoints before the if __name__ == "__main__" line
@app.get("/users/me/subscription")
async def get_user_subscription(current_user: User = Depends(get_current_active_user)):
"""Get the current user's subscription details"""
try:
# Get subscription details from database
conn = get_db_connection()
cursor = conn.cursor()
# Get the most recent active subscription
try:
cursor.execute(
"SELECT id, tier, status, created_at, expires_at, paypal_subscription_id FROM subscriptions "
"WHERE user_id = ? AND status = 'active' ORDER BY created_at DESC LIMIT 1",
(current_user.id,)
)
subscription = cursor.fetchone()
except sqlite3.OperationalError as e:
# Handle missing tier column
if "no such column: tier" in str(e):
logger.warning("Subscriptions table missing 'tier' column. Returning default subscription.")
subscription = None
else:
raise
# Get subscription tiers with pricing directly from SUBSCRIPTION_TIERS
subscription_tiers = {
"free_tier": {
"price": SUBSCRIPTION_TIERS["free_tier"]["price"],
"currency": SUBSCRIPTION_TIERS["free_tier"]["currency"],
"features": SUBSCRIPTION_TIERS["free_tier"]["features"]
},
"standard_tier": {
"price": SUBSCRIPTION_TIERS["standard_tier"]["price"],
"currency": SUBSCRIPTION_TIERS["standard_tier"]["currency"],
"features": SUBSCRIPTION_TIERS["standard_tier"]["features"]
},
"premium_tier": {
"price": SUBSCRIPTION_TIERS["premium_tier"]["price"],
"currency": SUBSCRIPTION_TIERS["premium_tier"]["currency"],
"features": SUBSCRIPTION_TIERS["premium_tier"]["features"]
}
}
if subscription:
sub_id, tier, status, created_at, expires_at, paypal_id = subscription
result = {
"id": sub_id,
"tier": tier,
"status": status,
"created_at": created_at,
"expires_at": expires_at,
"paypal_subscription_id": paypal_id,
"current_tier": current_user.subscription_tier,
"subscription_tiers": subscription_tiers
}
else:
result = {
"tier": "free_tier",
"status": "active",
"current_tier": current_user.subscription_tier,
"subscription_tiers": subscription_tiers
}
conn.close()
return result
except Exception as e:
logger.error(f"Error getting subscription: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error getting subscription: {str(e)}")
# Add this model definition before your endpoints
class SubscriptionCreate(BaseModel):
tier: str
@app.post("/create_subscription")
async def create_subscription(
subscription: SubscriptionCreate,
current_user: User = Depends(get_current_active_user)
):
"""Create a subscription for the current user"""
try:
# Log the request for debugging
logger.info(f"Creating subscription for user {current_user.email} with tier {subscription.tier}")
logger.info(f"Available tiers: {list(SUBSCRIPTION_TIERS.keys())}")
# Validate tier
valid_tiers = ["standard_tier", "premium_tier"]
if subscription.tier not in valid_tiers:
logger.warning(f"Invalid tier requested: {subscription.tier}")
raise HTTPException(status_code=400, detail=f"Invalid tier: {subscription.tier}. Must be one of {valid_tiers}")
# Create subscription
logger.info(f"Calling create_user_subscription with email: {current_user.email}, tier: {subscription.tier}")
success, result = create_user_subscription(current_user.email, subscription.tier)
if not success:
logger.error(f"Failed to create subscription: {result}")
raise HTTPException(status_code=400, detail=result)
logger.info(f"Subscription created successfully: {result}")
return result
except Exception as e:
logger.error(f"Error creating subscription: {str(e)}")
# Include the full traceback for better debugging
import traceback
logger.error(f"Traceback: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=f"Error creating subscription: {str(e)}")
@app.post("/subscribe/{tier}")
async def subscribe_to_tier(
tier: str,
current_user: User = Depends(get_current_active_user)
):
"""Subscribe to a specific tier"""
try:
# Validate tier
valid_tiers = ["standard_tier", "premium_tier"]
if tier not in valid_tiers:
raise HTTPException(status_code=400, detail=f"Invalid tier: {tier}. Must be one of {valid_tiers}")
# Create subscription
success, result = create_user_subscription(current_user.email, tier)
if not success:
raise HTTPException(status_code=400, detail=result)
return result
except Exception as e:
logger.error(f"Error creating subscription: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error creating subscription: {str(e)}")
@app.post("/subscription/create")
async def create_subscription(request: Request, current_user: User = Depends(get_current_active_user)):
"""Create a subscription for the current user"""
try:
data = await request.json()
tier = data.get("tier")
if not tier:
return JSONResponse(
status_code=400,
content={"detail": "Tier is required"}
)
# Log the request for debugging
logger.info(f"Creating subscription for user {current_user.email} with tier {tier}")
# Create the subscription using the imported function directly
success, result = create_user_subscription(current_user.email, tier)
if success:
# Make sure we're returning the approval_url in the response
logger.info(f"Subscription created successfully: {result}")
logger.info(f"Approval URL: {result.get('approval_url')}")
return {
"success": True,
"data": {
"approval_url": result["approval_url"],
"subscription_id": result["subscription_id"],
"tier": result["tier"]
}
}
else:
logger.error(f"Failed to create subscription: {result}")
return JSONResponse(
status_code=400,
content={"success": False, "detail": result}
)
except Exception as e:
logger.error(f"Error creating subscription: {str(e)}")
import traceback
logger.error(f"Traceback: {traceback.format_exc()}")
return JSONResponse(
status_code=500,
content={"success": False, "detail": f"Error creating subscription: {str(e)}"}
)
@app.post("/admin/initialize-paypal-plans")
async def initialize_paypal_plans(request: Request):
"""Initialize PayPal subscription plans"""
try:
# This should be protected with admin authentication in production
plans = initialize_subscription_plans()
if plans:
return JSONResponse(
status_code=200,
content={"success": True, "plans": plans}
)
else:
return JSONResponse(
status_code=500,
content={"success": False, "detail": "Failed to initialize plans"}
)
except Exception as e:
logger.error(f"Error initializing PayPal plans: {str(e)}")
return JSONResponse(
status_code=500,
content={"success": False, "detail": f"Error initializing plans: {str(e)}"}
)
@app.post("/subscription/verify")
async def verify_subscription(request: Request, current_user: User = Depends(get_current_active_user)):
"""Verify a subscription after payment"""
try:
data = await request.json()
subscription_id = data.get("subscription_id")
if not subscription_id:
return JSONResponse(
status_code=400,
content={"success": False, "detail": "Subscription ID is required"}
)
logger.info(f"Verifying subscription: {subscription_id}")
# Verify the subscription with PayPal
success, result = verify_paypal_subscription(subscription_id)
if not success:
logger.error(f"Subscription verification failed: {result}")
return JSONResponse(
status_code=400,
content={"success": False, "detail": str(result)}
)
# Update the user's subscription in the database
conn = get_db_connection()
cursor = conn.cursor()
# Get the subscription details
cursor.execute(
"SELECT tier FROM subscriptions WHERE paypal_subscription_id = ?",
(subscription_id,)
)
subscription = cursor.fetchone()
if not subscription:
# This is a new subscription, get the tier from the PayPal response
tier = "standard_tier" # Default to standard tier
# You could extract the tier from the PayPal plan ID if needed
# Create a new subscription record
sub_id = str(uuid.uuid4())
start_date = datetime.now()
expires_at = start_date + timedelta(days=30)
cursor.execute(
"INSERT INTO subscriptions (id, user_id, tier, status, created_at, expires_at, paypal_subscription_id) VALUES (?, ?, ?, ?, ?, ?, ?)",
(sub_id, current_user.id, tier, "active", start_date, expires_at, subscription_id)
)
else:
# Update existing subscription
tier = subscription[0]
cursor.execute(
"UPDATE subscriptions SET status = 'active' WHERE paypal_subscription_id = ?",
(subscription_id,)
)
# Update user's subscription tier
cursor.execute(
"UPDATE users SET subscription_tier = ? WHERE id = ?",
(tier, current_user.id)
)
conn.commit()
conn.close()
return JSONResponse(
status_code=200,
content={"success": True, "detail": "Subscription verified successfully"}
)
except Exception as e:
logger.error(f"Error verifying subscription: {str(e)}")
return JSONResponse(
status_code=500,
content={"success": False, "detail": f"Error verifying subscription: {str(e)}"}
)
@app.post("/subscription/webhook")
async def subscription_webhook(request: Request):
"""Handle PayPal subscription webhooks"""
try:
payload = await request.json()
success, result = handle_subscription_webhook(payload)
if not success:
logger.error(f"Webhook processing failed: {result}")
return {"status": "error", "message": result}
return {"status": "success", "message": result}
except Exception as e:
logger.error(f"Error processing webhook: {str(e)}")
return {"status": "error", "message": f"Error processing webhook: {str(e)}"}
@app.get("/subscription/verify/{subscription_id}")
async def verify_subscription(
subscription_id: str,
current_user: User = Depends(get_current_active_user)
):
"""Verify a subscription payment and update user tier"""
try:
# Verify the subscription
success, result = verify_subscription_payment(subscription_id)
if not success:
raise HTTPException(status_code=400, detail=f"Subscription verification failed: {result}")
# Get the plan ID from the subscription to determine tier
plan_id = result.get("plan_id", "")
# Connect to DB to get the tier for this plan
conn = get_db_connection()
cursor = conn.cursor()
cursor.execute("SELECT tier FROM paypal_plans WHERE plan_id = ?", (plan_id,))
tier_result = cursor.fetchone()
conn.close()
if not tier_result:
raise HTTPException(status_code=400, detail="Could not determine subscription tier")
tier = tier_result[0]
# Update the user's subscription
success, update_result = update_user_subscription(current_user.email, subscription_id, tier)
if not success:
raise HTTPException(status_code=500, detail=f"Failed to update subscription: {update_result}")
return {
"message": f"Successfully subscribed to {tier} tier",
"subscription_id": subscription_id,
"status": result.get("status", ""),
"next_billing_time": result.get("billing_info", {}).get("next_billing_time", "")
}
except HTTPException:
raise
except Exception as e:
print(f"Subscription verification error: {str(e)}")
raise HTTPException(status_code=500, detail=f"Subscription verification failed: {str(e)}")
@app.post("/webhook/paypal")
async def paypal_webhook(request: Request):
"""Handle PayPal subscription webhooks"""
try:
payload = await request.json()
logger.info(f"Received PayPal webhook: {payload.get('event_type', 'unknown event')}")
# Process the webhook
result = handle_subscription_webhook(payload)
return {"status": "success", "message": "Webhook processed"}
except Exception as e:
logger.error(f"Webhook processing error: {str(e)}")
# Return 200 even on error to acknowledge receipt to PayPal
return {"status": "error", "message": str(e)}
# Add this to your startup code
@app.on_event("startup")
async def startup_event():
"""Initialize subscription plans on startup"""
try:
# Initialize PayPal subscription plans if needed
# If you have an initialize_subscription_plans function in your paypal_integration.py,
# you can call it here
print("Application started successfully")
except Exception as e:
print(f"Error during startup: {str(e)}")
if __name__ == "__main__":
import uvicorn
port = int(os.environ.get("PORT", 7860))
host = os.environ.get("HOST", "0.0.0.0")
uvicorn.run("app:app", host=host, port=port, reload=True)