Update app.py

app.py

@@ -6,7 +6,6 @@ from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, ServiceCon
 from llama_index.llms.openai import OpenAI as LlamaOpenAI
 from llama_index.embeddings.openai import OpenAIEmbedding
 from llama_index.core.prompts import PromptTemplate
-import textwrap
 
 # ---------- AI Instructions & Prompts ----------
 SYSTEM_INSTRUCTION = """You are an elite marketing copywriter and TrustLogic expert assistant.
@@ -116,91 +115,21 @@ client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
 # ---------- Chainlit Hooks ----------
 
-
 @cl.on_chat_start
 async def start():
-    """Initialize chat session and inject CSS/JS to hide attachments and stop button"""
-    welcome_message = "How can I assist you today?"
-
-    # HTML snippet to inject CSS/JS that hides attachment and stop button selectors
-    # and provides a small helper to show an indeterminate loading bar.
-    injected_ui = textwrap.dedent("""
-    <div id="trustlogic-ui-inject" style="display:none;">
-      <style>
-        /* Hide file upload / attachment button */
-        .cl-file-input, .cl-upload, input[type="file"], button[aria-label="Upload file"] {
-          display: none !important;
-          visibility: hidden !important;
-        }
-        /* Hide stop/abort buttons commonly used in UIs */
-        button[aria-label="Stop"], .cl-stop-button, .stop-btn, button[data-testid="stop-button"] {
-          display: none !important;
-          visibility: hidden !important;
-        }
-
-        /* Simple loading bar container (hidden by default) */
-        #trustlogic-loading {
-          width: 100%;
-          max-width: 780px;
-          margin: 8px 0;
-          height: 10px;
-          background: rgba(0,0,0,0.06);
-          border-radius: 6px;
-          overflow: hidden;
-          display: none;
-        }
-        #trustlogic-loading .bar {
-          width: 0%;
-          height: 100%;
-          background: linear-gradient(90deg, rgba(0,0,0,0.12), rgba(0,0,0,0.24));
-          animation: indeterminate 1.2s infinite linear;
-        }
-        @keyframes indeterminate {
-          0% { transform: translateX(-100%); width: 40%; }
-          50% { transform: translateX(0%); width: 60%; }
-          100% { transform: translateX(100%); width: 40%; }
-        }
-      </style>
-
-      <div id="trustlogic-loading">
-        <div class="bar"></div>
-      </div>
-
-      <script>
-        // expose helpers to show/hide the loading bar from incoming messages
-        window.chainlitHelpers = window.chainlitHelpers || {};
-        window.chainlitHelpers.showLoading = function() {
-          var el = document.getElementById('trustlogic-loading');
-          if (el) el.style.display = 'block';
-        }
-        window.chainlitHelpers.hideLoading = function() {
-          var el = document.getElementById('trustlogic-loading');
-          if (el) el.style.display = 'none';
-        }
-
-        // best-effort hide for other variants: try again after a small delay (some UI renders later)
-        setTimeout(function() {
-          document.querySelectorAll('.cl-file-input, .cl-upload, button[aria-label="Upload file"], button[aria-label="Stop"], .cl-stop-button, .stop-btn').forEach(n => {
-            n.style.display = 'none';
-            n.style.visibility = 'hidden';
-          });
-        }, 600);
-      </script>
-    </div>
-    """)
-
-    # send welcome and injection; mark as HTML so it renders client-side
+    """Initialize chat session with a welcoming message"""
+    welcome_message = """How can I assist you today?"""
+    
     await cl.Message(content=welcome_message).send()
-    await cl.Message(content=injected_ui, content_type="html").send()
 
 @cl.on_message
 async def main(message: cl.Message):
     """Process incoming messages"""
     query = message.content.strip()
     
-    # Show loading message with spinner
-    loading_msg = cl.Message(content="🔄 Analyzing your query and searching for information...")
-    await loading_msg.send()
+    # Create streaming message
+    msg = cl.Message(content="")
+    await msg.send()
     
     # Check knowledge base first
     kb_query_engine = index.as_query_engine(
@@ -225,25 +154,22 @@ Decision criteria:
 Reply with only: USE_KB or USE_WEB"""
     
     decision = client.chat.completions.create(
-        model="gpt-4o",
+        model="gpt-5",
         messages=[{"role": "user", "content": decision_prompt}],
     )
     
     decision_text = decision.choices[0].message.content.strip()
     
-    # Remove loading message
-    await loading_msg.remove()
-    
-    # Create response message
-    msg = cl.Message(content="")
-    
     # --- Use Knowledge Base ---
     if "USE_KB" in decision_text:
         response_text = str(kb_response) + FOOTER_MESSAGE
         
         # Stream the response
-        for char in response_text:
+        for i, char in enumerate(response_text):
             await msg.stream_token(char)
+            # Add small delay for better UX (optional)
+            if i % 50 == 0:
+                await cl.sleep(0.01)
         
         await msg.update()
         return
@@ -298,4 +224,9 @@ Reply with only: USE_KB or USE_WEB"""
 @cl.on_chat_end
 async def end():
     """Handle chat end"""
-    print("Chat session ended")
+    print("Chat session ended")
+
+@cl.on_stop
+async def stop():
+    """Handle stop button"""
+    await cl.Message(content="⏹️ Task stopped by user.").send()