Update app.py

app.py

@@ -43,6 +43,28 @@ def get_ai_response(prompt, chat_history):
     chat_history.insert(0, {"role": "system", "content": system_prompt})
     chat_history.append({"role": "user", "content": prompt})
 
+    system_tokens = num_tokens_from_string(system_message["content"])
+    user_tokens = num_tokens_from_string(user_message["content"])
+
+    # Maximum allowed tokens
+    context_length = 4096 - max_tokens
+    available_tokens = context_length - system_tokens - user_tokens
+
+    # Trim chat history if necessary
+    trimmed_history = []
+    total_tokens = 0
+
+    for message in reversed(chat_history):
+        message_tokens = num_tokens_from_string(message["content"])
+        if total_tokens + message_tokens <= available_tokens:
+            trimmed_history.insert(0, message)
+            total_tokens += message_tokens
+        else:
+            break
+
+    # Construct final message list
+    final_messages = [system_message] + trimmed_history + [user_message]
+
     try:
         with st.chat_message("assistant", avatar=st.session_state.assistant_avatar):
             stream = client.chat.completions.create(
@@ -50,7 +72,7 @@ def get_ai_response(prompt, chat_history):
                 messages=chat_history,
                 temperature=temperature,
                 top_p=top_p,
-                max_tokens=1024,
+                max_tokens=max_tokens,
                 stream=True,
             )
     
@@ -215,7 +237,7 @@ with st.sidebar:
         system_prompt = st.text_area("System Prompt", height=200)
         temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.3, step=0.1)
         top_p = st.slider("Top-P", min_value=0.01, max_value=1.00, value=1.00, step=0.01)
-        
+        max_tokens = st.slider("Max Tokens (Output)", min_value=1, max_value=1024, value=1024, step=1)
         st.header("Restore History")
         history_input = st.text_area("Paste conversation history:", height=200)
         if st.button("Restore History"):