Update app.py

app.py

@@ -14,7 +14,6 @@ SYSTEM_PROMPT_BASE = (
 # ----------------------------------------------------------------------------
 # Together Client Initialization with API Key
 # ----------------------------------------------------------------------------
-# Expecting TOGETHER_API_KEY set in environment (e.g., Gradio Secrets)
 api_key = os.environ.get("TOGETHER_API_KEY")
 if not api_key:
     raise ValueError("Missing TOGETHER_API_KEY environment variable")
@@ -23,94 +22,82 @@ together_client = Together(api_key=api_key)
 # ----------------------------------------------------------------------------
 # In-Memory Context Tracking
 # ----------------------------------------------------------------------------
-conversation_history = []        # List of {'role': str, 'content': str}
-context_summary = ""           # Running summary of the conversation
+conversation_history = []  # List of (user_msg, assistant_msg)
+context_summary = ""     # Running summary of conversation
 
 # ----------------------------------------------------------------------------
 # Helper: Update Context Summary
 # ----------------------------------------------------------------------------
 def update_summary(history):
-    """
-    Call the model to generate a concise summary of the conversation history.
-    """
-    # Build summarization prompt
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT_BASE},
         {"role": "user", "content": (
             "Summarize the following cybersecurity conversation in 3-5 bullet points, "
             "focusing on key decisions and context:\n" +
-            "\n".join([f"{msg['role']}: {msg['content']}" for msg in history])
+            "\n".join([f"user: {u}\nassistant: {a}" for u, a in history])
         )}
     ]
-
-    response = together_client.chat.completions.create(
+    resp = together_client.chat.completions.create(
         model=MODEL_NAME,
         messages=messages,
         stream=False
     )
-    # Extract full summary text
-    summary_text = response.choices[0].message.content
-    return summary_text.strip()
+    return resp.choices[0].message.content.strip()
 
 # ----------------------------------------------------------------------------
-# Core Chat Functionality
+# Core Chat Functions for Gradio
 # ----------------------------------------------------------------------------
-def append_and_stream(user_input: str):
-    global context_summary
-
-    # Frame as cybersecurity expert if not already
+def user_submit(user_input, history):
+    history = history or []
+    if not user_input:
+        return "", history
+    # Frame as cybersecurity expert
     if not user_input.lower().startswith("as a cybersecurity expert"):
         user_input = f"(As a cybersecurity expert) {user_input}"
+    # Append placeholder for assistant
+    history.append((user_input, ""))
+    return "", history
 
-    # Append user message
-    conversation_history.append({'role': 'user', 'content': user_input})
-
-    # Build system prompt including updated context summary
-    system_prompt = SYSTEM_PROMPT_BASE
+def assistant_stream(history):
+    global context_summary
+    if not history:
+        return history
+    # Build messages for model
+    model_msgs = [{"role": "system", "content": SYSTEM_PROMPT_BASE}]
     if context_summary:
-        system_prompt += "\n\nPrevious context summary:\n" + context_summary
-
-    # Assemble messages for streaming
-    messages = ([{'role': 'system', 'content': system_prompt}] +
-                conversation_history)
-    # Add placeholder for assistant
-    conversation_history.append({'role': 'assistant', 'content': ''})
-
-    # Stream tokens from the model
+        model_msgs[0]["content"] += f"\n\nPrevious summary:\n{context_summary}"
+    for user_msg, assistant_msg in history[:-1]:
+        model_msgs.append({"role": "user", "content": user_msg})
+        model_msgs.append({"role": "assistant", "content": assistant_msg})
+    # Current user turn
+    user_msg, _ = history[-1]
+    model_msgs.append({"role": "user", "content": user_msg})
+    # Stream from together
     stream = together_client.chat.completions.create(
         model=MODEL_NAME,
-        messages=messages,
+        messages=model_msgs,
         stream=True
     )
-
-    # Incrementally build assistant reply
     for token in stream:
         if hasattr(token, 'choices'):
             delta = token.choices[0].delta.content
-            conversation_history[-1]['content'] += delta
-            yield [(msg['role'], msg['content']) for msg in conversation_history if msg['role'] != 'system']
-
-    # After full reply, update the context summary
-    context_summary = update_summary(conversation_history)
+            history[-1] = (history[-1][0], history[-1][1] + delta)
+            yield history
+    # After full response, update summary
+    context_summary = update_summary(history)
 
 # ----------------------------------------------------------------------------
-# Gradio Interface Definition
+# Launch Gradio Interface
 # ----------------------------------------------------------------------------
 def launch_interface():
     with gr.Blocks() as demo:
         gr.Markdown("## CyberGuard – Autonomous Cybersecurity Chat")
         chatbot = gr.Chatbot()
-        state = gr.State([])  # placeholder to trigger streaming
-
         txt = gr.Textbox(show_label=False, placeholder="Enter your security query...")
 
-        # On submit: trigger streaming
-        def on_submit(user_msg, _state):
-            return [[('user', user_msg)]], None  # just triggers state change
-
-        txt.submit(lambda *_: None, None, txt)
-        txt.submit(on_submit, [txt, state], [state, chatbot], queue=False)
-        state.change(fn=lambda s: append_and_stream(txt.value), inputs=state, outputs=chatbot)
+        txt.submit(user_submit, [txt, chatbot], [txt, chatbot], queue=False)
+        txt.submit(lambda: None, None, txt)  # clear input
+        chatbot.stream(assistant_stream, chatbot)
 
         demo.launch(share=True, server_name='0.0.0.0', server_port=7860)