Update app.py

app.py

@@ -1,83 +1,136 @@
-from flask import Flask, redirect, request, session, url_for
+# from flask import Flask, redirect, request, session, url_for
+# import os
+# from authlib.integrations.flask_client import OAuth
+# from langchain.llms.huggingface_hub import HuggingFaceHub
+# from langchain.prompts import ChatPromptTemplate
+# from langchain.schema import StrOutputParser
+# from langchain.schema.runnable import Runnable
+# from langchain.schema.runnable.config import RunnableConfig
+# import chainlit as cl
+
+# app = Flask(__name__)
+# app.secret_key = 'YourSecretKey'  # Change this to a real secret key for production
+
+# # OAuth setup with Authlib
+# oauth = OAuth(app)
+# oauth.register(
+#     name='oauth_provider',
+#     client_id=os.getenv("OAUTH_CLIENT_ID"),
+#     client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
+#     authorize_url=os.getenv("OPENID_PROVIDER_URL") + '/authorize',
+#     access_token_url=os.getenv("OPENID_PROVIDER_URL") + '/token',
+#     client_kwargs={'scope': os.getenv("OAUTH_SCOPES").split(',')},
+#     redirect_uri=f"https://{os.getenv('SPACE_HOST')}/login/callback"
+# )
+
+# print(f"REDIRECT URI: https://{os.getenv('SPACE_HOST')}/login/callback")
+
+# # Instantiate the LLM
+# llm = HuggingFaceHub(
+#     model_kwargs={"max_length": 500},
+#     repo_id="google/flan-t5-xxl",
+#     huggingfacehub_api_token=os.getenv("HUGGINGFACE_API_TOKEN"),
+# )
+
+# # Initialize ChainLit with LLM
+# def initialize_chainlit():
+#     add_llm_provider(
+#         LangchainGenericProvider(
+#             id=llm._llm_type,
+#             name="HuggingFaceHub",
+#             llm=llm,
+#             is_chat=False,
+#         )
+#     )
+
+# # Setup chainlit callbacks
+# @cl.on_chat_start
+# async def on_chat_start():
+#     prompt = ChatPromptTemplate.from_messages([("human", "{question}")])
+#     runnable = prompt | llm | StrOutputParser()
+#     cl.user_session.set("runnable", runnable)
+
+# @cl.on_message
+# async def on_message(message: cl.Message):
+#     runnable = cl.user_session.get("runnable")  # type: Runnable
+#     msg = cl.Message(content="")
+#     async for chunk in runnable.astream(
+#         {"question": message.content},
+#         config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
+#     ):
+#         await msg.stream_token(chunk)
+#     await msg.send()
+
+# @app.route('/')
+# def home():
+#     return 'Home - <a href="/login">Login with OAuth Provider</a>'
+
+# @app.route('/login')
+# def login():
+#     redirect_uri = url_for('authorize', _external=True)
+#     return oauth.oauth_provider.authorize_redirect(redirect_uri)
+
+# @app.route('/login/callback')
+# def authorize():
+#     print('Logged in and language model initialized. Proceed with operations.')
+#     token = oauth.oauth_provider.authorize_access_token()
+#     # Initialize ChainLit or perform actions based on the authenticated user
+#     initialize_chainlit()
+#     return 'Logged in and language model initialized. Proceed with operations.'
+
+# if __name__ == "__main__":
+#     app.run(debug=True)
+
+
+from flask import Flask, redirect, request, session, url_for, jsonify
 import os
-from authlib.integrations.flask_client import OAuth
-from langchain.llms.huggingface_hub import HuggingFaceHub
-from langchain.prompts import ChatPromptTemplate
-from langchain.schema import StrOutputParser
-from langchain.schema.runnable import Runnable
-from langchain.schema.runnable.config import RunnableConfig
-import chainlit as cl
+import requests
+from base64 import b64encode
 
 app = Flask(__name__)
 app.secret_key = 'YourSecretKey'  # Change this to a real secret key for production
 
-# OAuth setup with Authlib
-oauth = OAuth(app)
-oauth.register(
-    name='oauth_provider',
-    client_id=os.getenv("OAUTH_CLIENT_ID"),
-    client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
-    authorize_url=os.getenv("OPENID_PROVIDER_URL") + '/authorize',
-    access_token_url=os.getenv("OPENID_PROVIDER_URL") + '/token',
-    client_kwargs={'scope': os.getenv("OAUTH_SCOPES").split(',')},
-    redirect_uri=f"https://{os.getenv('SPACE_HOST')}/login/callback"
-)
-
-print(f"REDIRECT URI: https://{os.getenv('SPACE_HOST')}/login/callback")
-
-# Instantiate the LLM
-llm = HuggingFaceHub(
-    model_kwargs={"max_length": 500},
-    repo_id="google/flan-t5-xxl",
-    huggingfacehub_api_token=os.getenv("HUGGINGFACE_API_TOKEN"),
-)
-
-# Initialize ChainLit with LLM
-def initialize_chainlit():
-    add_llm_provider(
-        LangchainGenericProvider(
-            id=llm._llm_type,
-            name="HuggingFaceHub",
-            llm=llm,
-            is_chat=False,
-        )
-    )
-
-# Setup chainlit callbacks
-@cl.on_chat_start
-async def on_chat_start():
-    prompt = ChatPromptTemplate.from_messages([("human", "{question}")])
-    runnable = prompt | llm | StrOutputParser()
-    cl.user_session.set("runnable", runnable)
-
-@cl.on_message
-async def on_message(message: cl.Message):
-    runnable = cl.user_session.get("runnable")  # type: Runnable
-    msg = cl.Message(content="")
-    async for chunk in runnable.astream(
-        {"question": message.content},
-        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
-    ):
-        await msg.stream_token(chunk)
-    await msg.send()
+# OAuth Configuration
+CLIENT_ID = os.getenv("OAUTH_CLIENT_ID")
+CLIENT_SECRET = os.getenv("OAUTH_CLIENT_SECRET")
+REDIRECT_URI = f"https://{os.getenv('SPACE_HOST')}/login/callback"
+AUTHORIZE_URL = "https://huggingface.co/oauth/authorize"
+TOKEN_URL = "https://huggingface.co/oauth/token"
 
 @app.route('/')
 def home():
-    return 'Home - <a href="/login">Login with OAuth Provider</a>'
-
-@app.route('/login')
-def login():
-    redirect_uri = url_for('authorize', _external=True)
-    return oauth.oauth_provider.authorize_redirect(redirect_uri)
+    # Generate a random state for CSRF protection
+    state = os.urandom(16).hex()
+    session['state'] = state
+    # Redirect URL for "Sign-in with HF"
+    return redirect(f"{AUTHORIZE_URL}?redirect_uri={REDIRECT_URI}&scope=openid%20profile&client_id={CLIENT_ID}&state={state}")
 
 @app.route('/login/callback')
-def authorize():
-    print('Logged in and language model initialized. Proceed with operations.')
-    token = oauth.oauth_provider.authorize_access_token()
-    # Initialize ChainLit or perform actions based on the authenticated user
-    initialize_chainlit()
-    return 'Logged in and language model initialized. Proceed with operations.'
+def login_callback():
+    # Verify state matches
+    state = request.args.get('state')
+    if state != session.pop('state', None):
+        return 'State mismatch', 400
+
+    # Exchange code for token
+    code = request.args.get('code')
+    headers = {'Authorization': 'Basic ' + b64encode(f"{CLIENT_ID}:{CLIENT_SECRET}".encode()).decode()}
+    data = {
+        'grant_type': 'authorization_code',
+        'code': code,
+        'redirect_uri': REDIRECT_URI,
+        'client_id': CLIENT_ID
+    }
+    response = requests.post(TOKEN_URL, headers=headers, data=data)
+    response_data = response.json()
+
+    if response.status_code != 200:
+        return jsonify(response_data), response.status_code
+
+    # At this point, you have access_token and id_token in response_data
+    # You can use these tokens to authenticate against the Hugging Face API or your application's backend
+
+    return 'Logged in successfully.', 200
 
 if __name__ == "__main__":
     app.run(debug=True)
-