lost of changes, history almost working (continuing an old chat not working, as SystemPrompt() won't be there)

query.py

@@ -14,15 +14,24 @@ from elevenlabs.client import ElevenLabs
 from openai import OpenAI
 
 from json import loads as json_loads
+import itertools
 import time
 import os
 
-prompt_template = os.environ.get("PROMPT_TEMPLATE")
-
-prompt = ChatPromptTemplate.from_messages([('system', prompt_template)])
-
 AI = True
 
+if not hasattr(itertools, "batched"):
+    def batched(iterable, n):
+        "Batch data into lists of length n. The last batch may be shorter."
+        # batched('ABCDEFG', 3) --> ABC DEF G
+        it = iter(iterable)
+        while True:
+            batch = list(itertools.islice(it, n))
+            if not batch:
+                return
+            yield batch
+    itertools.batched = batched
+    
 def ai_setup():
     global llm, prompt_chain, oai_client
     
@@ -38,17 +47,19 @@ def ai_setup():
         )
 
         retriever = vstore.as_retriever(search_kwargs={'k': 10})
+
+        prompt_template = os.environ.get("PROMPT_TEMPLATE")
+        prompt = ChatPromptTemplate.from_messages([('system', prompt_template)])
+        prompt_chain = (
+            {"context": retriever, "question": RunnablePassthrough()}
+            | RunnableLambda(format_context) 
+            | prompt
+            # | llm
+            # | StrOutputParser()
+        )
     else:
         retriever = RunnableLambda(just_read)
 
-    prompt_chain = (
-        {"context": retriever, "question": RunnablePassthrough()}
-        | RunnableLambda(format_context) 
-        | prompt
-        # | llm
-        # | StrOutputParser()
-    )
-
 def group_and_sort(documents):
     grouped = {}
     for document in documents:
@@ -107,32 +118,105 @@ def new_state():
 def session_id(state: dict, request: gr.Request) -> str:
     return f'{state["user"]}_{request.session_hash}'
 
-store = None
-def auth(token, state, request: gr.Request):
-    global store
+class History:
+    store = None
+    def __init__(self, name:str, user:str, session_id:str, id:str = None):
+        self.session_id = session_id
+        self.name = name
+        self.user = user
+        self.astra_history = None
+
+        if id:
+            self.id = id
+        else:
+            self.id = f"{user}_{session_id}"
+            self.create()
+
+    @classmethod    
+    def get_store(self):
+        if self.store is None:
+            self.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"),
+            )
+        return self.store
+
+    @classmethod
+    def from_dict(cls, id:str, data:dict):
+        name = f":{id}"
+        name = data.get("name", name)
+        answer = cls(name, user=data["user"], id = id, session_id=data["session"])
+
+        return answer
+    
+    @classmethod
+    def get_histories(cls, user:str):
+        store = cls.get_store()
+        histories = []
+        keys = [k for k in store.yield_keys(prefix=f"{user}_")]
+        for id, history in zip(keys, store.mget(keys)):
+            history = cls.from_dict(id = id, data = history)
+            histories.append(history)
+        return histories
+    
+    @classmethod
+    def load(cls, id:str):
+        data = cls.get_store().mget([id])
+        return cls.from_dict(id, data[0])
+        
+    def __str__(self):
+        return f"{self.id}:{self.name}"
+    
+    def create(self):
+        history = {
+            'session'   : self.session_id,
+            'user'      : self.user,
+            'timestamp' : time.asctime(time.gmtime()),
+            'name'      : self.name
+        }
+        self.get_store().mset([(self.id, history)])
+    
+    @staticmethod
+    def get_history_collection_name():
+        return f'{os.environ.get("ASTRA_DB_COLLECTION")}_chat_history'
+    
+    def get_astra_history(self):
+        if self.astra_history is None:
+            self.astra_history = AstraDBChatMessageHistory(
+                session_id=self.id,
+                collection_name=self.get_history_collection_name(),
+                token=os.environ.get("ASTRA_DB_APPLICATION_TOKEN"),
+                api_endpoint=os.environ.get("ASTRA_DB_API_ENDPOINT"),
+            )
+        return self.astra_history
+
+    def add(self, type:str, message):
+        if   type == "system":
+            self.get_astra_history().add_message(message)
+        elif type == "user":
+            self.get_astra_history().add_user_message(message)
+        elif type == "ai":
+            self.get_astra_history().add_ai_message(message)
+    
+    def messages(self):
+        return self.get_astra_history().messages
+
+    def clear(self):
+        self.get_astra_history().clear()
     
+    def delete(self):
+        self.clear()
+        self.get_store().mdelete([self.id])
+
+def auth(token, state, request: gr.Request):
     tokens=os.environ.get("APP_TOKENS")
     if not tokens:
         state["user"] = "anonymous"
     else:
         tokens=json_loads(tokens)
         state["user"] = tokens.get(token, None)
-        
-    if state["user"]:
-        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 = session_id(state, request)
-        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) {
@@ -150,24 +234,40 @@ def not_authenticated(state):
         gr.Warning("You need to authenticate first")
     return answer
 
-def add_history(state, request, type, message):
+def list_histories(state):
+    if not_authenticated(state):
+        return gr.update()
+    
+    histories = History.get_histories(state["user"])
+    answer = [(h.name, h.id) for h in histories]
+    return gr.update(choices=answer, value=None)
+
+def add_history(state, request, type, message, name:str = None):
     if not state["history"]:
-        session = session_id(state, request)
-        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"),
+        name = name or message[:60]
+        state["history"] = History(
+            name = name,
+            user = state["user"],
+            session_id = request.session_hash
         )
-
-    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)
         
+    state["history"].add(type, message)
+
+def load_history(state, history_id):
+    state["history"] = History.load(history_id)
+
+    history = [m.content for m in state["history"].messages()]
+    history = itertools.batched(history, 2)
+    history = [m for m in history]
+    
+    if len(history) and len(history[-1]) == 1:
+        user_input = history[-1][0]
+        history = history[:-1]
+    else:
+        user_input = ''
+
+    return state, history, history, user_input   # state, Chatbot, ChatInterface.state, ChatInterface.textbox
+ 
 def chat(message, history, state, request:gr.Request):
     if not_authenticated(state):
         yield "You need to authenticate first"
@@ -177,7 +277,9 @@ def chat(message, history, state, request:gr.Request):
                 system_prompt = prompt_chain.invoke(message)
                 system_prompt = system_prompt.messages[0]
                 state["system"] = system_prompt
-                # add_history(state, request, "system", system_prompt)
+                
+                # Next is commented out because astra has a limit on document size
+                # add_history(state, request, "system", system_prompt, name=message)
             else:
                 system_prompt = state["system"]
 
@@ -235,14 +337,13 @@ def play_last(history, state):
             response = lab11.generate(text=text, voice=whatson, stream=True)
             yield from response
 
-
-def chat_chage(history):
+def chat_change(history):
     if history:
         if not history[-1][1]:
             return gr.update(interactive=False)
         elif history[-1][1][-1] != '…':
             return gr.update(interactive=True)
-    return gr.update()
+    return gr.update()      # play_last_btn
 
 TEXT_TALK = "🎤 Talk"
 TEXT_STOP = "⏹ Stop"
@@ -250,8 +351,9 @@ TEXT_STOP = "⏹ Stop"
 def gr_main():
     theme = gr.Theme.from_hub("freddyaboulton/[email protected]")
     theme.set(
-        color_accent_soft="#818eb6",            # ChatBot.svelte / .message-row.panel.user-row
-        background_fill_secondary="#6272a4",    # ChatBot.svelte / .message-row.panel.bot-row
+        color_accent_soft="#818eb6",            # ChatBot.svelte / .user / .message-row.panel.user-row . neutral_500 -> neutral_200
+        background_fill_secondary="#6272a4",    # ChatBot.svelte / .bot / .message-row.panel.bot-row . neutral_500 -> neutral_400
+        background_fill_primary="#818eb6",      # DropdownOptions.svelte / item 
         button_primary_text_color="*button_secondary_text_color",
         button_primary_background_fill="*button_secondary_background_fill")
 
@@ -262,12 +364,20 @@ def gr_main():
         css="footer {visibility: hidden}"
         ) as app:
             state = new_state()
-            # auto_play = gr.Checkbox(False, label="Autoplay", render=False)
             chatbot = gr.Chatbot(show_label=False, render=False, scale=1)
+            gr.HTML('<h1 style="text-align: center">Sherlock Holmes stories</h1>')
+            history_choice = gr.Dropdown(
+                choices=[("History", "History")],
+                value="History",
+                show_label=False,
+                container=False,
+                interactive=True,
+                filterable=True)
+                
             iface = gr.ChatInterface(
                 chat,
                 chatbot=chatbot,
-                title="Sherlock Holmes stories",
+                title=None,
                 submit_btn=gr.Button(
                                 "Send",
                                 variant="primary",
@@ -313,7 +423,7 @@ def gr_main():
                     play_last,
                     [chatbot, state], player)
 
-            chatbot.change(chat_chage, inputs=chatbot, outputs=play_last_btn)
+            chatbot.change(chat_change, inputs=chatbot, outputs=play_last_btn)
             start_stop_rec.click(
                 lambda x:x,
                 inputs=start_stop_rec,
@@ -335,6 +445,16 @@ def gr_main():
                 inputs=iface.textbox,
                 js='function (text){if (text) document.getElementById("submit_btn").click(); return [text]}'
             )
+            
+            history_choice.focus(
+                list_histories,
+                inputs=state,
+                outputs=history_choice
+            )
+            history_choice.input(
+                load_history,
+                inputs=[state, history_choice],
+                outputs=[state, chatbot, iface.chatbot_state, iface.textbox])
 
             token = gr.Textbox(visible=False)
 
@@ -344,9 +464,9 @@ def gr_main():
                 js=AUTH_JS)
         
     app.queue(default_concurrency_limit=None, api_open=False)
-    app.launch(show_api=False)
+    return app
     
 if __name__ == "__main__":
     ai_setup()
-    gr_main()
-    
+    app = gr_main()
+    app.launch(show_api=False)