Update app.py

app.py

@@ -1,137 +1,84 @@
-# 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
+from flask import Flask, redirect, request, session, url_for
 import os
-import requests
-from base64 import b64encode
+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 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"
+# 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():
-    # 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 login_callback():
-    # Verify state matches
-    state = request.args.get('state')
-    if state != session.pop('state', None):
-        return 'State mismatch', 400
+    return 'Home - <a href="/login">Login with OAuth Provider</a>'
 
-    # 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()
-    print(response_data)
+@app.route('/login')
+def login():
+    redirect_uri = url_for('authorize', _external=True)
+    return oauth.oauth_provider.authorize_redirect(redirect_uri)
 
-    if response.status_code != 200:
-        return jsonify(response_data), response.status_code
+@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.'
 
-    # 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
+if __name__ == "__main__":
+    app.run(debug=True)
 
-    return 'Logged in successfully.', 200
 
-# if __name__ == "__main__":
-app.run(debug=True, port=7860)