Update app.py

app.py

@@ -1,12 +1,21 @@
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, redirect, url_for
 import PyPDF2
 from transformers import pipeline
 import psycopg2
 import json
-import torch
+import requests
+from simple_salesforce import Salesforce
+import urllib.parse
 
 app = Flask(__name__)
 
+# Salesforce Connected App credentials
+SALESFORCE_CONSUMER_KEY = 'your_consumer_key'
+SALESFORCE_CONSUMER_SECRET = 'your_consumer_secret'
+SALESFORCE_REDIRECT_URI = 'http://localhost:5000/oauth/callback'
+SALESFORCE_AUTH_URL = 'https://login.salesforce.com/services/oauth2/authorize'
+SALESFORCE_TOKEN_URL = 'https://login.salesforce.com/services/oauth2/token'
+
 # Database connection
 conn = psycopg2.connect(
     dbname="contracts_db",
@@ -21,6 +30,48 @@ cursor = conn.cursor()
 ner_pipeline = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple")
 compliance_pipeline = pipeline("text-classification", model="facebook/bart-large-mnli")
 
+# Store OAuth tokens
+access_token = None
+instance_url = None
+
+@app.route('/login')
+def login():
+    params = {
+        'client_id': SALESFORCE_CONSUMER_KEY,
+        'redirect_uri': SALESFORCE_REDIRECT_URI,
+        'response_type': 'code',
+        'scope': 'api refresh_token'
+    }
+    auth_url = f"{SALESFORCE_AUTH_URL}?{urllib.parse.urlencode(params)}"
+    return redirect(auth_url)
+
+@app.route('/oauth/callback')
+def oauth_callback():
+    global access_token, instance_url
+    code = request.args.get('code')
+    if not code:
+        return jsonify({"error": "Authorization code not found"}), 400
+    
+    data = {
+        'grant_type': 'authorization_code',
+        'client_id': SALESFORCE_CONSUMER_KEY,
+        'client_secret': SALESFORCE_CONSUMER_SECRET,
+        'redirect_uri': SALESFORCE_REDIRECT_URI,
+        'code': code
+    }
+    response = requests.post(SALESFORCE_TOKEN_URL, data=data)
+    if response.status_code != 200:
+        return jsonify({"error": "Failed to obtain access token"}), 500
+    
+    token_data = response.json()
+    access_token = token_data['access_token']
+    instance_url = token_data['instance_url']
+    return redirect(url_for('index'))
+
+@app.route('/')
+def index():
+    return jsonify({"message": "Welcome to Contract Analyzer. Please log in to Salesforce."})
+
 def extract_text_from_pdf(file):
     try:
         pdf_reader = PyPDF2.PdfReader(file)
@@ -33,6 +84,10 @@ def extract_text_from_pdf(file):
 
 @app.route('/upload_contract', methods=['POST'])
 def upload_contract():
+    global access_token, instance_url
+    if not access_token or not instance_url:
+        return jsonify({"error": "Please authenticate with Salesforce first"}), 401
+    
     if 'contract' not in request.files:
         return jsonify({"error": "No file uploaded"}), 400
     
@@ -56,7 +111,7 @@ def upload_contract():
             "score": result['scores'][0]
         })
     
-    # Store in database
+    # Store in PostgreSQL
     try:
         cursor.execute(
             "INSERT INTO contracts (content, entities, compliance) VALUES (%s, %s, %s) RETURNING id",
@@ -68,6 +123,17 @@ def upload_contract():
         conn.rollback()
         return jsonify({"error": str(e)}), 500
     
+    # Push to Salesforce
+    try:
+        sf = Salesforce(instance_url=instance_url, session_id=access_token)
+        sf.Contract__c.create({
+            'Content__c': text[:1000],  # Truncate for Salesforce field limit
+            'Entities__c': json.dumps(extracted_data),
+            'Compliance__c': json.dumps(compliance_results)
+        })
+    except Exception as e:
+        return jsonify({"error": f"Failed to save to Salesforce: {str(e)}"}), 500
+    
     return jsonify({
         "contract_id": contract_id,
         "entities": extracted_data,