|
1 | 1 | # (A) LOAD SETTINGS & MODULES |
2 | 2 | import settings as set |
3 | 3 | from flask import Flask, Response, request |
4 | | -import torch |
| 4 | +import torch, jwt |
5 | 5 | from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline |
6 | 6 | from langchain import PromptTemplate, HuggingFacePipeline |
7 | 7 | from langchain.vectorstores import Chroma |
|
19 | 19 | persist_directory = set.path_db, |
20 | 20 | embedding_function = HuggingFaceEmbeddings() |
21 | 21 | ) |
22 | | -pipe = "" |
23 | | -chain = "" |
24 | 22 |
|
25 | 23 | # (C) PIPE + CHAIN |
26 | 24 | pipe = pipeline( |
|
50 | 48 | verbose = set.chain_verbose |
51 | 49 | ) |
52 | 50 |
|
53 | | -# (D) FLASK |
| 51 | +""" @TODO - ENABLE THIS TO OPEN FOR REGISTERED USERS ONLY |
| 52 | +# (D) VERIFY USER |
| 53 | +def jwtVerify(cookies): |
| 54 | + try: |
| 55 | + token = jwt.decode( |
| 56 | + jwt = cookies.get("cbsess"), |
| 57 | + key = set.jwt_secret, |
| 58 | + audience = set.http_host, |
| 59 | + algorithms = [set.jwt_algo] |
| 60 | + ) |
| 61 | + # DO WHATEVER YOU WANT WITH THE DECODED USER TOKEN |
| 62 | + # print(token) |
| 63 | + return True |
| 64 | + except Exception as error: |
| 65 | + # print(error) |
| 66 | + return False |
| 67 | +""" |
| 68 | + |
| 69 | +# (E) FLASK |
54 | 70 | app = Flask(__name__) |
55 | | -@app.route("/", methods=["POST"]) |
| 71 | +@app.route("/", methods = ["POST"]) |
56 | 72 | def bot(): |
| 73 | + # (E1) CORS |
57 | 74 | if "HTTP_ORIGIN" in request.environ and request.environ["HTTP_ORIGIN"] in set.http_allow: |
| 75 | + # (E2-1) ALLOW ONLY REGISTERED USERS |
| 76 | + """ @TODO - ENABLE THIS TO OPEN FOR REGISTERED USERS ONLY |
| 77 | + if jwtVerify(request.cookies) is False: |
| 78 | + return Response("Not Allowed", status = 405) |
| 79 | + """ |
| 80 | + |
| 81 | + # (E2-2) ANSWER THE QUESTION |
58 | 82 | data = dict(request.form) |
59 | 83 | if "query" in data: |
60 | 84 | ans = chain(data["query"]) |
61 | 85 | ans = ans["result"] |
62 | 86 | else: |
63 | 87 | ans = "Where's the question, yo?" |
64 | | - response = Response(ans, status=200) |
| 88 | + response = Response(ans, status = 200) |
65 | 89 | response.headers.add("Access-Control-Allow-Origin", request.environ["HTTP_ORIGIN"] ) |
66 | 90 | response.headers.add("Access-Control-Allow-Credentials", "true") |
| 91 | + |
| 92 | + # (E2) ORIGIN NOT ALLOWED |
67 | 93 | else: |
68 | | - response = Response("Not Allowed", status=405) |
| 94 | + response = Response("Not Allowed", status = 405) |
69 | 95 | return response |
70 | 96 |
|
| 97 | +# (F) GO! |
71 | 98 | if __name__ == "__main__": |
72 | 99 | app.run( |
73 | 100 | host = set.http_host, |
|
0 commit comments