auth

app.py

@@ -1,18 +1,5 @@
 import subprocess
-import gradio as gr
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import Chroma
-from langchain.llms import HuggingFacePipeline
-from langchain.chains import RetrievalQA
-from transformers import AutoConfig, AutoTokenizer, pipeline, AutoModelForCausalLM
-from langchain_community.document_loaders import DirectoryLoader
-import torch
-import re
-import requests
-from urllib.parse import urlencode, urlparse, parse_qs
-import spaces
 
-# Step 1: Run the setup script
 script_path = './setup.sh'  # Adjust the path if needed
 
 # Run the script
@@ -23,6 +10,20 @@ if exit_code == 0:
 else:
     print(f"Script failed with exit code {exit_code}.")
 
+import gradio as gr
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.llms import HuggingFacePipeline
+from langchain.chains import RetrievalQA
+from transformers import AutoConfig, AutoTokenizer, pipeline, AutoModelForCausalLM
+from langchain_community.document_loaders import DirectoryLoader
+import torch
+import re
+import transformers
+import spaces
+import requests
+from urllib.parse import urlencode
+
 # Initialize embeddings and ChromaDB
 model_name = "sentence-transformers/all-mpnet-base-v2"
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -37,8 +38,9 @@ books_db_client = books_db.as_retriever()
 
 # Initialize the model and tokenizer
 model_name = "stabilityai/stablelm-zephyr-3b"
-model_config = AutoConfig.from_pretrained(model_name, max_new_tokens=1024)
-model = AutoModelForCausalLM.from_pretrained(
+
+model_config = transformers.AutoConfig.from_pretrained(model_name, max_new_tokens=1024)
+model = transformers.AutoModelForCausalLM.from_pretrained(
     model_name,
     trust_remote_code=True,
     config=model_config,
@@ -47,7 +49,7 @@ model = AutoModelForCausalLM.from_pretrained(
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-query_pipeline = pipeline(
+query_pipeline = transformers.pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
@@ -75,8 +77,8 @@ TENANT_ID = '2b093ced-2571-463f-bc3e-b4f8bcb427ee'
 CLIENT_ID = '2a7c884c-942d-49e2-9e5d-7a29d8a0d3e5'
 CLIENT_SECRET = 'EOF8Q~kKHCRgx8tnlLM-H8e93ifetxI6x7sU6bGW'
 REDIRECT_URI = 'https://sanjeevbora-chatbot.hf.space/'
-AUTH_URL = f"https://login.microsoftonline.com/{TENANT_ID}/oauth2/v2.0/authorize"
-TOKEN_URL = f"https://login.microsoftonline.com/{TENANT_ID}/oauth2/v2.0/token"
+AUTH_URL = f"https://login.microsoftonline.com/2b093ced-2571-463f-bc3e-b4f8bcb427ee/oauth2/v2.0/authorize"
+TOKEN_URL = f"https://login.microsoftonline.com/2b093ced-2571-463f-bc3e-b4f8bcb427ee/oauth2/v2.0/token"
 
 params = {
     'client_id': CLIENT_ID,
@@ -84,17 +86,30 @@ params = {
     'redirect_uri': REDIRECT_URI,
     'response_mode': 'query',
     'scope': 'User.Read',
-    'state': '12345',
-    'prompt': 'login'  # This ensures the login prompt appears even if already logged in
+    'state': '12345'  # Optional state parameter
 }
 
 # Construct the login URL
 login_url = f"{AUTH_URL}?{urlencode(params)}"
 
+# Gradio interface
 def show_login_button():
-    return f'<a href="{login_url}" class="GFG">Click here to login with Microsoft</a>'
-
-# Function to exchange auth code for token
+    return f'<a href="{login_url}" target="_blank">Click here to login with Microsoft</a>'
+
+# Dummy function to simulate token validation (you will replace this with actual validation)
+def is_logged_in(token):
+    # Check if the token exists (or check if it's valid)
+    return token is not None
+
+# Gradio interface
+def check_login(status):
+    # If logged in, show the chatbot interface, otherwise show login link
+    if status:
+        return gr.update(visible=True), gr.update(visible=True)
+    else:
+        return gr.update(visible=False), gr.update(visible=False)
+    
+# Function to exchange authorization code for access token
 def exchange_code_for_token(auth_code):
     data = {
         'grant_type': 'authorization_code',
@@ -103,6 +118,7 @@ def exchange_code_for_token(auth_code):
         'code': auth_code,
         'redirect_uri': REDIRECT_URI
     }
+    
     response = requests.post(TOKEN_URL, data=data)
     
     if response.status_code == 200:
@@ -112,17 +128,14 @@ def exchange_code_for_token(auth_code):
     else:
         return None
 
-# Function to handle redirect URL and extract the auth code
-def handle_redirect(url):
-    parsed_url = urlparse(url)
-    query_params = parse_qs(parsed_url.query)
-    auth_code = query_params.get('code')
+def login_user(auth_code):
+    # Exchange the authorization code for an access token
+    token = exchange_code_for_token(auth_code)
     
-    if auth_code:
-        token = exchange_code_for_token(auth_code[0])
-        if token:
-            return "Logged in", True  # Successfully logged in
-    return "Login failed", False
+    if token:
+        return token
+    else:
+        return None
 
 # Function to retrieve answer using the RAG system
 @spaces.GPU(duration=60)
@@ -152,43 +165,34 @@ def chat(query, history=None):
 def clear_input():
     return "",  # Return empty string to clear input field
 
-
-# Gradio Interface
 with gr.Blocks() as interface:
-    with gr.Tab("Login"):
-        gr.Markdown("## Login Page")
-        login_link = gr.HTML(show_login_button())
-        
-        # Hidden textbox for redirect URL
-        redirect_url_input = gr.Textbox(label="Redirect URL", visible=True)  # URL from the Microsoft redirect
-
-        status_label = gr.Label(value="Not logged in")  # Label to show login status
-        
-        def on_redirect(redirect_url):
-            # Extract and exchange token
-            status, logged_in = handle_redirect(redirect_url)
-            if logged_in:
-                return gr.update(visible=False), status, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
-            else:
-                return gr.update(visible=True), status, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-        
-        # Handle redirect and switch to chatbot page upon login
-        redirect_url_input.change(
-            on_redirect,
-            inputs=[redirect_url_input],
-            outputs=[redirect_url_input, status_label, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Chatbot.update(visible=True)],
-            show_progress=True
-        )
-        
-    with gr.Tab("Chatbot"):
-        gr.Markdown("## Chatbot Page")
-        
-        # Components for chat (initially hidden)
-        input_box = gr.Textbox(label="Enter your question", placeholder="Type your question here...", visible=False)
-        submit_btn = gr.Button("Submit", visible=False)
-        chat_history = gr.Chatbot(label="Chat History", visible=False)
-
-        # Chat submission
-        submit_btn.click(chat, inputs=[input_box, chat_history], outputs=[chat_history, input_box])
+    gr.Markdown("## RAG Chatbot")
+    gr.Markdown("Please log in to continue.")
+    
+    # Custom HTML to show login link
+    login_link = gr.HTML(f'<a href="{login_url}" target="_blank">Click here to login with Microsoft</a>')
+
+    # Login button to simulate the login process
+    login_button = gr.Button("Login")
+    
+    # Components for chat (initially hidden)
+    input_box = gr.Textbox(label="Enter your question", placeholder="Type your question here...", visible=False)
+    submit_btn = gr.Button("Submit", visible=False)
+    chat_history = gr.Chatbot(label="Chat History", visible=False)
+
+    # Handle login button click
+    login_button.click(
+        login_user, 
+        inputs=[], 
+        outputs=[login_button],  # You can also update the UI to show login status
+        queue=False
+    ).then(
+        lambda token: check_login(is_logged_in(token)), 
+        inputs=[], 
+        outputs=[input_box, submit_btn]
+    )
+
+    # Input submission and chat handling
+    submit_btn.click(chat, inputs=[input_box, chat_history], outputs=[chat_history, input_box])
 
 interface.launch()