Merge pull request #96 from DL4DS/cleanup

removing debug print statements and minor fix

Author: trgardos

apps/ai_tutor/chainlit_app.py

@@ -13,7 +13,7 @@
     get_sources,
     get_history_chat_resume,
     get_history_setup_llm,
-    get_last_config,
+    # get_last_config,
 )
 from modules.chat_processor.helpers import (
     update_user_info,
@@ -27,7 +27,6 @@
 import copy
 from typing import Optional
 from chainlit.types import ThreadDict
-import time
 import base64
 from langchain_community.callbacks import get_openai_callback
 from datetime import datetime, timezone
@@ -90,7 +89,6 @@ async def setup_llm(self):
 
         #TODO: Clean this up.
         """
-        start_time = time.time()
 
         llm_settings = cl.user_session.get("llm_settings", {})
         (
@@ -138,8 +136,6 @@ async def setup_llm(self):
         cl.user_session.set("chain", self.chain)
         cl.user_session.set("llm_tutor", self.llm_tutor)
 
-        print("Time taken to setup LLM: ", time.time() - start_time)
-
     @no_type_check
     async def update_llm(self, new_settings: Dict[str, Any]):
         """
@@ -222,32 +218,9 @@ async def inform_llm_settings(self):
         """
         Inform the user about the updated LLM settings and display them as a message.
         """
-        llm_settings: Dict[str, Any] = cl.user_session.get("llm_settings", {})
-        llm_tutor = cl.user_session.get("llm_tutor")
-        settings_dict = {
-            "model": llm_settings.get("chat_model"),
-            "retriever": llm_settings.get("retriever_method"),
-            "memory_window": llm_settings.get("memory_window"),
-            "num_docs_in_db": (
-                len(llm_tutor.vector_db)
-                if llm_tutor and hasattr(llm_tutor, "vector_db")
-                else 0
-            ),
-            "view_sources": llm_settings.get("view_sources"),
-            "follow_up_questions": llm_settings.get("follow_up_questions"),
-        }
-        print("Settings Dict: ", settings_dict)
         await cl.Message(
             author=SYSTEM,
             content="LLM settings have been updated. You can continue with your Query!",
-            # elements=[
-            #     cl.Text(
-            #         name="settings",
-            #         display="side",
-            #         content=json.dumps(settings_dict, indent=4),
-            #         language="json",
-            #     ),
-            # ],
         ).send()
 
     async def set_starters(self):
@@ -306,8 +279,6 @@ async def start(self):
         and display and load previous conversation if chat logging is enabled.
         """
 
-        start_time = time.time()
-
         await self.make_llm_settings_widgets(self.config)  # Reload the settings widgets
 
         user = cl.user_session.get("user")
@@ -335,8 +306,6 @@ async def start(self):
         cl.user_session.set("llm_tutor", self.llm_tutor)
         cl.user_session.set("chain", self.chain)
 
-        print("Time taken to start LLM: ", time.time() - start_time)
-
     async def stream_response(self, response):
         """
         Stream the response from the LLM.
@@ -367,8 +336,6 @@ async def main(self, message):
             message: The incoming chat message.
         """
 
-        start_time = time.time()
-
         chain = cl.user_session.get("chain")
         token_count = 0  # initialize token count
         if not chain:
@@ -386,8 +353,6 @@ async def main(self, message):
         user.metadata = updated_user.metadata
         cl.user_session.set("user", user)
 
-        print("\n\n User Tokens Left: ", user.metadata["tokens_left"])
-
         # see if user has token credits left
         # if not, return message saying they have run out of tokens
         if user.metadata["tokens_left"] <= 0 and "admin" not in user.metadata["role"]:
@@ -478,12 +443,9 @@ async def main(self, message):
         )
         answer_with_sources = answer_with_sources.replace("$$", "$")
 
-        print("Time taken to process the message: ", time.time() - start_time)
-
         actions = []
 
         if self.config["llm_params"]["generate_follow_up"]:
-            start_time = time.time()
             cb_follow_up = cl.AsyncLangchainCallbackHandler()
             config = {
                 "callbacks": (
@@ -513,8 +475,6 @@ async def main(self, message):
                     )
                 )
 
-            print("Time taken to generate questions: ", time.time() - start_time)
-
         # # update user info with token count
         tokens_left = await update_user_from_chainlit(user, token_count)
 
@@ -532,21 +492,20 @@ async def main(self, message):
         ).send()
 
     async def on_chat_resume(self, thread: ThreadDict):
-        thread_config = None
+        # thread_config = None
         steps = thread["steps"]
         k = self.config["llm_params"][
             "memory_window"
         ]  # on resume, alwyas use the default memory window
         conversation_list = get_history_chat_resume(steps, k, SYSTEM, LLM)
-        thread_config = get_last_config(
-            steps
-        )  # TODO: Returns None for now - which causes config to be reloaded with default values
+        # thread_config = get_last_config(
+        #     steps
+        # )  # TODO: Returns None for now - which causes config to be reloaded with default values
         cl.user_session.set("memory", conversation_list)
-        await self.start(config=thread_config)
+        await self.start()
 
     @cl.header_auth_callback
     def header_auth_callback(headers: dict) -> Optional[cl.User]:
-        print("\n\n\nI am here\n\n\n")
         # try: # TODO: Add try-except block after testing
         # TODO: Implement to get the user information from the headers (not the cookie)
         cookie = headers.get("cookie")  # gets back a str
@@ -562,10 +521,6 @@ def header_auth_callback(headers: dict) -> Optional[cl.User]:
         ).decode()
         decoded_user_info = json.loads(decoded_user_info)
 
-        print(
-            f"\n\n USER ROLE: {decoded_user_info['literalai_info']['metadata']['role']} \n\n"
-        )
-
         return cl.User(
             id=decoded_user_info["literalai_info"]["id"],
             identifier=decoded_user_info["literalai_info"]["identifier"],

apps/ai_tutor/helpers.py

@@ -32,9 +32,6 @@ async def check_user_cooldown(
     cooldown_end_time = last_message_time + timedelta(seconds=COOLDOWN_TIME)
     cooldown_end_time_iso = cooldown_end_time.isoformat()
 
-    # Debug: Print the cooldown end time
-    print(f"Cooldown end time (ISO): {cooldown_end_time_iso}")
-
     # Check if the user is still in cooldown
     if elapsed_time_in_seconds < COOLDOWN_TIME:
         return True, cooldown_end_time_iso  # Return in ISO 8601 format
@@ -81,10 +78,6 @@ async def reset_tokens_for_user(user_info, TOKENS_LEFT, REGEN_TIME):
         # Ensure the new token count does not exceed max_tokens
         new_token_count = min(current_tokens + tokens_to_regenerate, max_tokens)
 
-        print(
-            f"\n\n Adding {tokens_to_regenerate} tokens to the user, Time elapsed: {elapsed_time_in_seconds} seconds, Tokens after regeneration: {new_token_count}, Tokens before: {current_tokens} \n\n"
-        )
-
         # Update the user's token count
         user_info["metadata"]["tokens_left"] = new_token_count
 

apps/chainlit_base/chainlit_base.py

@@ -8,7 +8,6 @@
     get_history_setup_llm,
 )
 import copy
-import time
 from langchain_community.callbacks import get_openai_callback
 from config.config_manager import config_manager
 
@@ -36,7 +35,6 @@ async def setup_llm(self):
 
         #TODO: Clean this up.
         """
-        start_time = time.time()
 
         llm_settings = cl.user_session.get("llm_settings", {})
         (
@@ -84,8 +82,6 @@ async def setup_llm(self):
         cl.user_session.set("chain", self.chain)
         cl.user_session.set("llm_tutor", self.llm_tutor)
 
-        print("Time taken to setup LLM: ", time.time() - start_time)
-
     @no_type_check
     async def update_llm(self, new_settings: Dict[str, Any]):
         """
@@ -168,32 +164,9 @@ async def inform_llm_settings(self):
         """
         Inform the user about the updated LLM settings and display them as a message.
         """
-        llm_settings: Dict[str, Any] = cl.user_session.get("llm_settings", {})
-        llm_tutor = cl.user_session.get("llm_tutor")
-        settings_dict = {
-            "model": llm_settings.get("chat_model"),
-            "retriever": llm_settings.get("retriever_method"),
-            "memory_window": llm_settings.get("memory_window"),
-            "num_docs_in_db": (
-                len(llm_tutor.vector_db)
-                if llm_tutor and hasattr(llm_tutor, "vector_db")
-                else 0
-            ),
-            "view_sources": llm_settings.get("view_sources"),
-            "follow_up_questions": llm_settings.get("follow_up_questions"),
-        }
-        print("Settings Dict: ", settings_dict)
         await cl.Message(
             author=SYSTEM,
             content="LLM settings have been updated. You can continue with your Query!",
-            # elements=[
-            #     cl.Text(
-            #         name="settings",
-            #         display="side",
-            #         content=json.dumps(settings_dict, indent=4),
-            #         language="json",
-            #     ),
-            # ],
         ).send()
 
     async def set_starters(self):
@@ -243,8 +216,6 @@ async def start(self):
         and display and load previous conversation if chat logging is enabled.
         """
 
-        start_time = time.time()
-
         await self.make_llm_settings_widgets(self.config)  # Reload the settings widgets
 
         # TODO: remove self.user with cl.user_session.get("user")
@@ -263,8 +234,6 @@ async def start(self):
         cl.user_session.set("llm_tutor", self.llm_tutor)
         cl.user_session.set("chain", self.chain)
 
-        print("Time taken to start LLM: ", time.time() - start_time)
-
     async def stream_response(self, response):
         """
         Stream the response from the LLM.
@@ -295,8 +264,6 @@ async def main(self, message):
             message: The incoming chat message.
         """
 
-        start_time = time.time()
-
         chain = cl.user_session.get("chain")
         token_count = 0  # initialize token count
         if not chain:
@@ -342,12 +309,9 @@ async def main(self, message):
         )
         answer_with_sources = answer_with_sources.replace("$$", "$")
 
-        print("Time taken to process the message: ", time.time() - start_time)
-
         actions = []
 
         if self.config["llm_params"]["generate_follow_up"]:
-            start_time = time.time()
             cb_follow_up = cl.AsyncLangchainCallbackHandler()
             config = {
                 "callbacks": (
@@ -377,9 +341,6 @@ async def main(self, message):
                     )
                 )
 
-            print("Time taken to generate questions: ", time.time() - start_time)
-            print("Total Tokens Used: ", token_count)
-
         await cl.Message(
             content=answer_with_sources,
             elements=source_elements,

modules/chat/helpers.py

@@ -137,31 +137,59 @@ def get_history_chat_resume(steps, k, SYSTEM, LLM):
 
 def get_history_setup_llm(memory_list):
     conversation_list = []
-    for message in memory_list:
-        message_dict = message.to_dict() if hasattr(message, "to_dict") else message
-
-        # Check if the type attribute is present as a key or attribute
-        message_type = (
-            message_dict.get("type", None)
-            if isinstance(message_dict, dict)
-            else getattr(message, "type", None)
+    i = 0
+    while i < len(memory_list) - 1:
+        # Process the current and next message
+        current_message = memory_list[i]
+        next_message = memory_list[i + 1]
+
+        # Convert messages to dictionary if necessary
+        current_message_dict = (
+            current_message.to_dict()
+            if hasattr(current_message, "to_dict")
+            else current_message
+        )
+        next_message_dict = (
+            next_message.to_dict() if hasattr(next_message, "to_dict") else next_message
         )
 
-        # Check if content is present as a key or attribute
-        message_content = (
-            message_dict.get("content", None)
-            if isinstance(message_dict, dict)
-            else getattr(message, "content", None)
+        # Check message type and content for current and next message
+        current_message_type = (
+            current_message_dict.get("type", None)
+            if isinstance(current_message_dict, dict)
+            else getattr(current_message, "type", None)
+        )
+        current_message_content = (
+            current_message_dict.get("content", None)
+            if isinstance(current_message_dict, dict)
+            else getattr(current_message, "content", None)
         )
 
-        if message_type in ["ai", "ai_message"]:
-            conversation_list.append({"type": "ai_message", "content": message_content})
-        elif message_type in ["human", "user_message"]:
+        next_message_type = (
+            next_message_dict.get("type", None)
+            if isinstance(next_message_dict, dict)
+            else getattr(next_message, "type", None)
+        )
+        next_message_content = (
+            next_message_dict.get("content", None)
+            if isinstance(next_message_dict, dict)
+            else getattr(next_message, "content", None)
+        )
+
+        # Check if the current message is user message and the next one is AI message
+        if current_message_type in ["human", "user_message"] and next_message_type in [
+            "ai",
+            "ai_message",
+        ]:
+            conversation_list.append(
+                {"type": "user_message", "content": current_message_content}
+            )
             conversation_list.append(
-                {"type": "user_message", "content": message_content}
+                {"type": "ai_message", "content": next_message_content}
             )
+            i += 2  # Skip the next message since it has been paired
         else:
-            raise ValueError("Invalid message type")
+            i += 1  # Move to the next message if not a valid pair (example user message, followed by the cooldown system message)
 
     return conversation_list