added history saving and sessions

query.py

@@ -1,11 +1,12 @@
 import gradio as gr
 
-from langchain_astradb import AstraDBVectorStore
-
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough, RunnableLambda
 from langchain_core.messages import SystemMessage, AIMessage, HumanMessage
+
+from langchain_astradb import AstraDBChatMessageHistory, AstraDBStore, AstraDBVectorStore
+
 from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 
 from elevenlabs import VoiceSettings
@@ -98,17 +99,35 @@ def just_read(pipeline_state):
 
 def new_state():
     return gr.State({
-        "user": None,
-        "system": None,
+        "user"      : None,
+        "system"    : None,
+        "history"   : None,
     })
 
-def auth(token, state):
+store = None
+def auth(token, state, request: gr.Request):
+    global store
+    
     tokens=os.environ.get("APP_TOKENS")
     if not tokens:
         state["user"] = "anonymous"
     else:
         tokens=json_loads(tokens)
         state["user"] = tokens.get(token, None)
+        
+        if store is None:
+            store = AstraDBStore(
+                collection_name=f'{os.environ.get("ASTRA_DB_COLLECTION")}_sessions',
+                token=os.environ.get("ASTRA_DB_APPLICATION_TOKEN"),
+                api_endpoint=os.environ.get("ASTRA_DB_API_ENDPOINT"),
+            )
+        user_session = f'{state["user"]}_{request.session_hash}'    
+        session_data = {
+            'user'      : state["user"],
+            'session'   : request.session_hash,
+            'timestamp' : time.asctime(time.gmtime())
+        }
+        store.mset([(user_session, session_data)])
     return "", state
 
 AUTH_JS = """function auth_js(token, state) {
@@ -126,30 +145,60 @@ def not_authenticated(state):
         gr.Warning("You need to authenticate first")
     return answer
 
+def add_history(state, request, type, message):
+    if not state["history"]:
+        session = request.session_hash
+        state["history"] = AstraDBChatMessageHistory(
+            session_id=session,
+            collection_name=f'{os.environ.get("ASTRA_DB_COLLECTION")}_chat_history',
+            token=os.environ.get("ASTRA_DB_APPLICATION_TOKEN"),
+            api_endpoint=os.environ.get("ASTRA_DB_API_ENDPOINT"),
+        )
 
-def chat(message, history, state):
+    history = state["history"]
+    if   type == "system":
+        history.add_message(message)
+    elif type == "user":
+        history.add_user_message(message)
+    elif type == "ai":
+        history.add_ai_message(message)
+        
+def chat(message, history, state, request:gr.Request):
     if not_authenticated(state):
         yield "You need to authenticate first"
-    elif AI:
-        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))
-        
-        all = ''
-        for response in llm.stream(messages):
-            all += response.content
-            yield all
     else:
-        yield f"{time.ctime()}: You said: {message}"
+        if AI:
+            if not history:
+                system_prompt = prompt_chain.invoke(message)
+                system_prompt = system_prompt.messages[0]
+                state["system"] = system_prompt
+                # add_history(state, request, "system", system_prompt)
+            else:
+                system_prompt = state["system"]
+
+            add_history(state, request, "user", message)
+
+            messages = [system_prompt]
+            for human, ai in history:
+                messages.append(HumanMessage(human))
+                messages.append(AIMessage(ai))
+            messages.append(HumanMessage(message))
+            
+            answer = ''
+            for response in llm.stream(messages):
+                answer += response.content
+                yield answer
+        else:
+            add_history(state, request, "user", message)
+
+            msg = f"{time.ctime()}: You said: {message}"
+            answer = ' '
+            for word in msg.split():
+                answer += f' {word}'
+                yield answer
+                time.sleep(0.05)
+
+        add_history(state, request, "ai", answer)
 
 def on_audio(path, state):
     if not_authenticated(state):
@@ -194,6 +243,7 @@ def gr_main():
         theme=theme
         ) as app:
             state = new_state()
+            # auto_play = gr.Checkbox(False, label="Autoplay", render=False)
             chatbot = gr.Chatbot(show_label=False, render=False, scale=1)
             iface = gr.ChatInterface(
                 chat,
@@ -225,6 +275,18 @@ def gr_main():
                     show_label=False,
                     format="mp3",
                     waveform_options=gr.WaveformOptions(sample_rate=16000))
+                player = gr.Audio(
+                    show_label=False,
+                    show_download_button=False,
+                    show_share_button=False,
+                    visible=True,
+                    autoplay=True,
+                    streaming=True,
+                    interactive=False)
+                # with gr.Column():
+                #     auto_play.render()
+                play_btn  = gr.Button("Play last ")
+
                 mic.change(
                     on_audio, [mic, state], [iface.textbox, mic]
                 ).then(
@@ -232,16 +294,9 @@ def gr_main():
                     js='function (text){if (text) document.getElementById("submit_btn").click(); return [text]}',
                     inputs=iface.textbox
                 )
-                
-                player = gr.Audio(
-                    show_label=False,
-                    show_download_button=True,
-                    visible=True,
-                    autoplay=True,
-                    streaming=True)
-                
-                play_btn = gr.Button("Play last ")
-                play_btn.click(play_last, [chatbot, state], player)
+                play_btn.click(
+                    play_last,
+                    [chatbot, state], player)
                 
             
             token = gr.Textbox(visible=False)