update gradio sdk and add lfu caching

README.md

@@ -4,7 +4,7 @@ emoji: 💬
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.36.1
+sdk_version: 5.17.1
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -13,24 +13,30 @@ config.set_envs()  # Set environment variables using the config module
 store = stores.chroma.ChromaDB(config.EMBEDDINGS)
 service = services.scraper.Service(store)
 
-
 # Scrape data and get the store vector retriever
 service.scrape_and_get_store_vector_retriever(config.URLS)
 
 # Initialize the LLMService with logger, prompt, and store vector retriever
 llm_svc = LLMService(logger, config.SYSTEM_PROMPT, store.get_chroma_instance().as_retriever())
 
-def respond(user_input, history):
+def respond(user_input,session_hash):
     if user_input == "clear_chat_history_aisdb_override":
         llm_svc.store={}
         return "Memory Cache cleared"
     response = llm_svc.conversational_rag_chain().invoke(
         {"input": user_input},
-        config={"configurable": {"session_id": "abc"}},
+        config={"configurable": {"session_id": session_hash}},
     )["answer"]
 
     return response
 
+def echo(text, chat_history, request: gr.Request):
+    if request:
+        session_hash = request.session_hash
+        return respond(text, session_hash)
+    else:
+        return "No request object received."
+
 
 def on_reset_button_click():
     llm_svc.store={}
@@ -40,18 +46,7 @@ if __name__ == '__main__':
     logging.info("Starting AIVIz Bot")
 
     with gr.Blocks() as demo:
-        gr.Markdown("# 🚢 AIVIz Bot - Vessel Trajectory Prediction")
-        gr.Markdown("Welcome! Ask me anything about vessel tracking, AI models.")
-
-        with gr.Row():
-            chat_interface = gr.ChatInterface(fn=respond)
-        
-        with gr.Row():
-            reset_button = gr.Button("🔄 Reset Chat Memory Cache")
-            reset_status = gr.Textbox(label="Status", interactive=False)
-
-        # Bind reset button to function
-        reset_button.click(fn=on_reset_button_click, outputs=reset_status)
+        gr.ChatInterface(fn=echo, type="messages")
 
     # Launch the interface
-    demo.launch(share=True)
+    demo.launch()

caching/lfu.py

@@ -0,0 +1,43 @@
+from collections import defaultdict, OrderedDict
+
+class LFUCache:
+    def __init__(self, capacity: int):
+        self.capacity = capacity
+        self.data = {}  # session_id -> (value, freq)
+        self.freq_map = defaultdict(OrderedDict)  # freq -> {session_id: None}
+        self.min_freq = 0
+
+    def _update_freq(self, session_id):
+        value, freq = self.data[session_id]
+        del self.freq_map[freq][session_id]
+        if not self.freq_map[freq]:
+            del self.freq_map[freq]
+            if self.min_freq == freq:
+                self.min_freq += 1
+        
+        new_freq = freq + 1
+        self.data[session_id] = (value, new_freq)
+        self.freq_map[new_freq][session_id] = None
+
+    def get(self, session_id):
+        if session_id not in self.data:
+            return None
+        self._update_freq(session_id)
+        return self.data[session_id][0]
+
+    def put(self, session_id, value):
+        if self.capacity == 0:
+            return
+        
+        if session_id in self.data:
+            self.data[session_id] = (value, self.data[session_id][1])
+            self._update_freq(session_id)
+        else:
+            if len(self.data) >= self.capacity:
+                # Evict the least frequently used item
+                lfu_session_id, _ = self.freq_map[self.min_freq].popitem(last=False)
+                del self.data[lfu_session_id]
+                
+            self.data[session_id] = (value, 1)
+            self.freq_map[1][session_id] = None
+            self.min_freq = 1

llm_setup/llm_setup.py

@@ -12,7 +12,7 @@ from langchain_core.chat_history import BaseChatMessageHistory
 from langchain_community.chat_message_histories import ChatMessageHistory
 from langchain_core.runnables.history import RunnableWithMessageHistory
 from processing.documents import format_documents
-
+from caching.lfu import LFUCache
 
 def _initialize_llm() -> ChatGoogleGenerativeAI:
     """
@@ -23,7 +23,7 @@ def _initialize_llm() -> ChatGoogleGenerativeAI:
 
 
 class LLMService:
-    def __init__(self, logger, system_prompt: str, web_retriever: VectorStoreRetriever):
+    def __init__(self, logger, system_prompt: str, web_retriever: VectorStoreRetriever,cache_capacity: int = 50):
         self._conversational_rag_chain = None
         self._logger = logger
         self.system_prompt = system_prompt
@@ -34,7 +34,7 @@ class LLMService:
         self._initialize_conversational_rag_chain()
 
         ### Statefully manage chat history ###
-        self.store = {}
+        self.store = LFUCache(capacity=cache_capacity)
 
     def _initialize_conversational_rag_chain(self):
         """
@@ -55,7 +55,6 @@ class LLMService:
         )
 
 
-
         history_aware_retriever = create_history_aware_retriever(
         self.llm, self._web_retriever, contextualize_q_prompt)
 
@@ -79,9 +78,11 @@ class LLMService:
         )
 
     def _get_session_history(self, session_id: str) -> BaseChatMessageHistory:
-            if session_id not in self.store:
-                self.store[session_id] = ChatMessageHistory()
-            return self.store[session_id]
+        history = self.store.get(session_id)
+        if history is None:
+            history = ChatMessageHistory()
+            self.store.put(session_id, history)
+        return history
 
     def conversational_rag_chain(self):
         """