Update app.py

app.py

@@ -1,77 +1,85 @@
+import os
 import gradio as gr
-import logging
-import json
 from together import Together
 
 # ----------------------------------------------------------------------------
 # Configuration & Constants
 # ----------------------------------------------------------------------------
 MODEL_NAME = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
-SYSTEM_PROMPT = (
+SYSTEM_PROMPT_BASE = (
     "You are CyberGuard, a senior-level cybersecurity expert assistant. "
     "You autonomously enforce security best practices, making informed decisions when rule-based policies fail."
 )
-HISTORY_FILE = "conversation_history.json"
 
 # ----------------------------------------------------------------------------
-# Setup Logging
+# Together Client Initialization with API Key
 # ----------------------------------------------------------------------------
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
+# 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")
+together_client = Together(api_key=api_key)
 
 # ----------------------------------------------------------------------------
-# Conversation Persistence Utilities
+# In-Memory Context Tracking
 # ----------------------------------------------------------------------------
-def load_history(filepath: str) -> list:
-    """Load conversation history from a JSON file."""
-    try:
-        with open(filepath, 'r') as f:
-            history = json.load(f)
-        logger.info("Loaded existing conversation history.")
-        return history
-    except FileNotFoundError:
-        logger.info("No existing history found, starting fresh.")
-        return []
-
-
-def save_history(history: list, filepath: str) -> None:
-    """Persist conversation history to disk."""
-    with open(filepath, 'w') as f:
-        json.dump(history, f, indent=2)
-    logger.info("Conversation history saved.")
+conversation_history = []        # List of {'role': str, 'content': str}
+context_summary = ""           # Running summary of the conversation
 
 # ----------------------------------------------------------------------------
-# Together Client Initialization
+# Helper: Update Context Summary
 # ----------------------------------------------------------------------------
-together_client = Together()
+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])
+        )}
+    ]
+
+    response = 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()
 
 # ----------------------------------------------------------------------------
 # Core Chat Functionality
 # ----------------------------------------------------------------------------
-def append_and_stream(user_input: str, history: list):
-    """
-    Append user input to history, call the LLM streaming API, and yield token-by-token.
-    """
+def append_and_stream(user_input: str):
+    global context_summary
+
     # Frame as cybersecurity expert if not already
     if not user_input.lower().startswith("as a cybersecurity expert"):
         user_input = f"(As a cybersecurity expert) {user_input}"
 
-    # Append system prompt at start if missing
-    if not history or history[0]['role'] != 'system':
-        history.insert(0, {'role': 'system', 'content': SYSTEM_PROMPT})
+    # Append user message
+    conversation_history.append({'role': 'user', 'content': user_input})
+
+    # Build system prompt including updated context summary
+    system_prompt = SYSTEM_PROMPT_BASE
+    if context_summary:
+        system_prompt += "\n\nPrevious context summary:\n" + context_summary
 
-    # Append user message and prepare assistant placeholder
-    history.append({'role': 'user', 'content': user_input})
-    history.append({'role': 'assistant', 'content': ''})
-    save_history(history, HISTORY_FILE)
+    # 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
     stream = together_client.chat.completions.create(
         model=MODEL_NAME,
-        messages=history,
+        messages=messages,
         stream=True
     )
 
@@ -79,38 +87,32 @@ def append_and_stream(user_input: str, history: list):
     for token in stream:
         if hasattr(token, 'choices'):
             delta = token.choices[0].delta.content
-            history[-1]['content'] += delta
-            save_history(history, HISTORY_FILE)
-            yield history
+            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)
 
 # ----------------------------------------------------------------------------
 # Gradio Interface Definition
 # ----------------------------------------------------------------------------
 def launch_interface():
-    # Load previous history or start a new one
-    history = load_history(HISTORY_FILE)
-
     with gr.Blocks() as demo:
         gr.Markdown("## CyberGuard – Autonomous Cybersecurity Chat")
-        chatbot = gr.Chatbot(value=[(msg['role'], msg['content']) for msg in history if msg['role'] != 'system'])
-        state = gr.State(history)
+        chatbot = gr.Chatbot()
+        state = gr.State([])  # placeholder to trigger streaming
 
         txt = gr.Textbox(show_label=False, placeholder="Enter your security query...")
 
-        # Handle user submission
-        def on_submit(user_msg, hist):
-            return hist, hist + []  # trigger state change
+        # 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)  # Clear input box
+        txt.submit(lambda *_: None, None, txt)
         txt.submit(on_submit, [txt, state], [state, chatbot], queue=False)
-
-        # Stream assistant response when history updates
-        state.change(fn=append_and_stream, inputs=state, outputs=chatbot)
+        state.change(fn=lambda s: append_and_stream(txt.value), inputs=state, outputs=chatbot)
 
         demo.launch(share=True, server_name='0.0.0.0', server_port=7860)
 
 if __name__ == "__main__":
-    try:
-        launch_interface()
-    except Exception as e:
-        logger.exception("Failed to launch CyberGuard chat interface.")
+    launch_interface()