enable auth, enable AI

query.py

@@ -8,13 +8,14 @@ from langchain_core.runnables import RunnablePassthrough, RunnableLambda
 from langchain_core.messages import SystemMessage, AIMessage, HumanMessage
 from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 
+from json import loads as json_loads
 import os
 
 prompt_template = os.environ.get("PROMPT_TEMPLATE")
 
 prompt = ChatPromptTemplate.from_messages([('system', prompt_template)])
 
-AI = False
+AI = True
 
 def ai_setup():
     global llm, prompt_chain
@@ -91,27 +92,50 @@ def just_read(pipeline_state):
 
 def new_state():
     return gr.State({
+        "user": None,
         "system": None,
     })
 
-def chat(message, history, state):    
-    if not history:
-        system_prompt = prompt_chain.invoke(message)
-        system_prompt = system_prompt.messages[0]
-        state["system"] = system_prompt
+def auth(token, state):
+    tokens=os.environ.get("APP_TOKENS", None)
+    if tokens is None:
+        state["user"] = "anonymous"
     else:
-        system_prompt = state["system"]
+        tokens=json_loads(tokens)
+        state["user"] = tokens.get(token, None)
+    return "", state
+
+AUTH_JS = """function auth_js(token, state) {
+    if (!!document.location.hash) {
+        token = document.location.hash
+        document.location.hash=""
+    }        
+    return [token, state]
+}
+"""
+
+def chat(message, history, state):
+    if (state is None) or (not state['user']):
+        gr.Warning("You need to authenticate first")
+        yield "You need to authenticate first"
+    else:
+        if not history:
+            system_prompt = prompt_chain.invoke(message)
+            system_prompt = system_prompt.messages[0]
+            state["system"] = system_prompt
+        else:
+            system_prompt = state["system"]
+            
+        messages = [system_prompt]
+        for human, ai in history:
+            messages.append(HumanMessage(human))
+            messages.append(AIMessage(ai))
+        messages.append(HumanMessage(message))
         
-    messages = [system_prompt]
-    for human, ai in history:
-        messages.append(HumanMessage(human))
-        messages.append(AIMessage(ai))
-    messages.append(HumanMessage(message))
-    
-    all = ''
-    for response in llm.stream(messages):
-        all += response.content
-        yield all
+        all = ''
+        for response in llm.stream(messages):
+            all += response.content
+            yield all
 
 def gr_main():
     theme = gr.Theme.from_hub("freddyaboulton/[email protected]")
@@ -137,8 +161,15 @@ def gr_main():
                     "and I just received a letter threatening me to make public some things"
                     "of my past I'd rather keep quiet, unless I don't marry"],
                 ],
-            additional_inputs=[state])
+                additional_inputs=[state])
+            token = gr.Textbox(visible=False)
+            app.load(auth,
+                [token,state],
+                [token,state],
+                js=AUTH_JS)
+            
     app.launch(show_api=False)
+    
 if __name__ == "__main__":
     ai_setup()
     gr_main()