implement message history management with trimming functionality

app/app.py

@@ -11,7 +11,7 @@ from ulid import ulid
 import asyncio
 import time
 
-from utils import MODELS_PROVIDERS_MAP, PROVIDERS_API_KEYS, REASONER_PROVIDERS_MAP, check_openai_api_key
+from utils import MODELS_PROVIDERS_MAP, PROVIDERS_API_KEYS, REASONER_PROVIDERS_MAP, check_openai_api_key, trim_messages
 from settings import PROFILES_SETTINGS
 
 DEFAULT_REASONER_CONFIG = LlmConfig(
@@ -90,6 +90,7 @@ async def setup_agent(settings):
 @cl.on_chat_start
 async def start_chat():
     user_profile = cl.user_session.get("chat_profile")
+    cl.user_session.set("history", [])
     llm_config = DEFAULT_LLM_CONFIG.get(user_profile)
     llm = Llm.from_config(llm_config)
     llm.session_id = ulid()
@@ -100,6 +101,7 @@ async def start_chat():
     settings = await cl.ChatSettings(
         PROFILES_SETTINGS.get(user_profile)
     ).send()
+    
 
 async def run_concurrent_tasks(llm, message):
     asyncio.create_task(llm.acomplete(message))
@@ -128,7 +130,11 @@ async def main(message: cl.Message):
     
     start = time.time()
     thinking=False
-
+    
+    history = cl.user_session.get("history")
+    history.append(message.content)
+    history = trim_messages(history, llm.tokenizer)
+    model_id = None
     try:
         if llm.reasoner is not None or llm.config.model in REASONER_PROVIDERS_MAP:
             # Streaming the thinking
@@ -136,7 +142,7 @@ async def main(message: cl.Message):
                 msg = cl.Message(content="")
                 async for chunk in run_concurrent_tasks(
                         llm,
-                        message=message.content
+                        message=history
                     ):
                     if chunk == STREAM_START_TOKEN:
                         continue
@@ -150,6 +156,7 @@ async def main(message: cl.Message):
                         thinking_step.name = f"{llm.reasoner.config.model} to think for {thought_for}s"
                         await thinking_step.update()
                         chunk = f"```{llm.config.model}```\n"
+                        model_id = f"```{llm.config.model}```\n"
 
                     if chunk == STREAM_END_TOKEN:
                         break
@@ -162,7 +169,7 @@ async def main(message: cl.Message):
             msg = cl.Message(content="")
             async for chunk in run_concurrent_tasks(
                     llm,
-                    message=message.content
+                    message=history
                 ):
                 if chunk == STREAM_START_TOKEN:
                     continue
@@ -171,14 +178,12 @@ async def main(message: cl.Message):
                         break
                 
                 await msg.stream_token(chunk)
-
-        # message_history.append({"role": "assistant", "content": msg.content})
+            
+        hst_msg = msg.content.replace(model_id, "") if model_id else msg.content
+        history.append(hst_msg)
         await msg.send()
     
     except Exception as e:
         await cl.ErrorMessage("Internal Server Error").send()
 
-### TODO add support for history capping via tokenier fn
-### Add suppor to deactivate history in settings
-### TODO add future todos, include support for images and pdf upload for conversation
-### TODO openrouter missing'
+### TODO add future todos, include support for images and pdf upload for conversation

app/utils.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from typing import Optional
 import openai
 import os
 
@@ -35,4 +36,10 @@ def check_openai_api_key(api_key, base_url=None):
     except openai.AuthenticationError:
         return False
     else:
-        return True
+        return True
+
+def trim_messages(messages, tokenizer_fn, max_tokens :Optional[int]=None):
+    max_tokens = max_tokens or int(os.environ.get("MAX_HISTORY_TOKENS", 1028))
+    while messages and sum(len(tokenizer_fn(msg)) for msg in messages) > max_tokens:
+        messages.pop(0)  # Remove from the beginning
+    return messages