Spaces:
Build error
Build error
| 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." | |
| 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 | |
| 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)}") | |
| 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)}") | |
| async def root(): | |
| """Root endpoint that returns a welcome message.""" | |
| return HTMLResponse(content=""" | |
| <html> | |
| <head> | |
| <title>Legal Document Analysis API</title> | |
| <style> | |
| body { | |
| font-family: Arial, sans-serif; | |
| max-width: 800px; | |
| margin: 0 auto; | |
| padding: 20px; | |
| } | |
| h1 { | |
| color: #2c3e50; | |
| } | |
| .endpoint { | |
| background-color: #f8f9fa; | |
| padding: 15px; | |
| margin-bottom: 10px; | |
| border-radius: 5px; | |
| } | |
| .method { | |
| font-weight: bold; | |
| color: #e74c3c; | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <h1>Legal Document Analysis API</h1> | |
| <p>Welcome to the Legal Document Analysis API. This API provides tools for analyzing legal documents, videos, and audio.</p> | |
| <h2>Available Endpoints:</h2> | |
| <div class="endpoint"> | |
| <p><span class="method">POST</span> /analyze_legal_document - Analyze a legal document (PDF)</p> | |
| </div> | |
| <div class="endpoint"> | |
| <p><span class="method">POST</span> /analyze_legal_video - Analyze a legal video</p> | |
| </div> | |
| <div class="endpoint"> | |
| <p><span class="method">POST</span> /analyze_legal_audio - Analyze legal audio</p> | |
| </div> | |
| <div class="endpoint"> | |
| <p><span class="method">POST</span> /legal_chatbot/{task_id} - Chat with a document</p> | |
| </div> | |
| <div class="endpoint"> | |
| <p><span class="method">POST</span> /register - Register a new user</p> | |
| </div> | |
| <div class="endpoint"> | |
| <p><span class="method">POST</span> /token - Login to get an access token</p> | |
| </div> | |
| <div class="endpoint"> | |
| <p><span class="method">GET</span> /users/me - Get current user information</p> | |
| </div> | |
| <div class="endpoint"> | |
| <p><span class="method">POST</span> /subscribe/{tier} - Subscribe to a plan</p> | |
| </div> | |
| <p>For more details, visit the <a href="/docs">API documentation</a>.</p> | |
| </body> | |
| </html> | |
| """) | |
| 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)}") | |
| 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)}", | |
| ) | |
| 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)}"} | |
| 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 | |
| def health_check(): | |
| """Simple health check endpoint to verify the API is running""" | |
| return {"status": "ok", "message": "API is running"} | |
| async def read_users_me(current_user: User = Depends(get_current_active_user)): | |
| return current_user | |
| 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 | |
| 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 | |
| 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)}") | |
| 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)}") | |
| 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)}"} | |
| ) | |
| 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)}"} | |
| ) | |
| 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)}"} | |
| ) | |
| 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)}"} | |
| 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)}") | |
| 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 | |
| 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) |