Update app.py

app.py

@@ -19,32 +19,23 @@ MODEL_OPTIONS = {
 
 PERSONA_PRESETS = {
     "Balanced Assistant": "You are a helpful, intelligent AI assistant.",
-
     "Creative": "You are highly creative and imaginative.",
-
     "Precise": "You are concise, factual, and analytical.",
-
-    "Code Expert": "You are a senior software engineer who writes clean, well-documented code.",
+    "Code Expert": "You are a senior software engineer who writes clean code.",
 
     "Forensic Psychologist": """
 You are acting as a forensic psychologist.
-
-You provide structured, objective, and evidence-based psychological analysis 
-within legal and investigative contexts.
-
-You do NOT diagnose real individuals.
-You avoid speculation beyond available evidence.
-You clearly separate facts, interpretations, and limitations.
+Provide structured, objective, evidence-based analysis.
+Do NOT diagnose real individuals.
+Separate facts, interpretations, and limitations.
 Maintain professional neutrality.
-
 This analysis is educational and not a clinical diagnosis.
 """,
 
     "Forensic Psychologist (Structured Report)": """
-You are acting as a forensic psychologist preparing a structured analytical summary.
-
-When responding, structure your output using:
+You are acting as a forensic psychologist preparing a structured report.
 
+Structure:
 1. Case Summary
 2. Observed Behaviors
 3. Psychological Considerations (Hypothetical)
@@ -53,8 +44,7 @@ When responding, structure your output using:
 6. Evidentiary Limitations
 7. Professional Disclaimer
 
-Remain objective and avoid definitive diagnoses.
-This analysis is educational and not a clinical diagnosis.
+Remain objective and non-diagnostic.
 """
 }
 
@@ -65,9 +55,9 @@ TASK_MODES = {
 Conduct a behavioral analysis.
 
 1. Behavioral Indicators
-2. Emotional Tone & Affect
+2. Emotional Tone
 3. Cognitive Patterns
-4. Stress or Deception Markers
+4. Stress/Deception Markers
 5. Risk-Relevant Observations
 6. Alternative Explanations
 7. Limitations
@@ -108,7 +98,7 @@ Material:
 }
 
 MAX_CONTEXT_CHARS = 12000
-MAX_HISTORY = 8  # memory cap for stability
+MAX_HISTORY_PAIRS = 8  # 8 exchanges
 
 # =========================================
 # File Processing
@@ -144,49 +134,43 @@ def stream_response(message, history, model_label, persona_label, task_mode, upl
         yield history
         return
 
-    # Guard message length
+    # Cap input length
     message = message[:4000]
 
-    # Limit memory
-    history = history[-MAX_HISTORY:]
+    # Cap memory (each pair = 2 messages)
+    history = history[-MAX_HISTORY_PAIRS * 2:]
 
     model_name = MODEL_OPTIONS[model_label]
     base_system_prompt = PERSONA_PRESETS[persona_label]
 
-    # Temperature adjustment
-    if "Forensic Psychologist" in persona_label:
-        temperature = 0.4
-    else:
-        temperature = 0.7
+    temperature = 0.4 if "Forensic Psychologist" in persona_label else 0.7
 
     hf_token = os.environ.get("HF_TOKEN")
     client = InferenceClient(model=model_name, token=hf_token)
 
     # File context
     file_text = extract_text(uploaded_file)
-
     if file_text:
         file_text = file_text[:MAX_CONTEXT_CHARS]
         system_prompt = (
             f"{base_system_prompt}\n\n"
-            f"Use the following document as context when relevant:\n\n"
-            f"{file_text}"
+            f"Use this document context when relevant:\n\n{file_text}"
         )
     else:
         system_prompt = base_system_prompt
 
+    # Build messages
     messages = [{"role": "system", "content": system_prompt}]
+    messages.extend(history)
 
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
+    formatted_input = TASK_MODES[task_mode].format(user_input=message)
+    messages.append({"role": "user", "content": formatted_input})
 
-    formatted_message = TASK_MODES[task_mode].format(user_input=message)
-    messages.append({"role": "user", "content": formatted_message})
+    # Update UI history (Gradio 5 format)
+    history.append({"role": "user", "content": message})
+    history.append({"role": "assistant", "content": ""})
 
-    history.append((message, ""))
-
-    partial_response = ""
+    partial = ""
 
     for chunk in client.chat_completion(
         messages=messages,
@@ -197,8 +181,8 @@ def stream_response(message, history, model_label, persona_label, task_mode, upl
     ):
         if chunk.choices and chunk.choices[0].delta.content:
             token = chunk.choices[0].delta.content
-            partial_response += token
-            history[-1] = (message, partial_response)
+            partial += token
+            history[-1]["content"] = partial
             yield history
 
 # =========================================
@@ -209,10 +193,7 @@ def clear_chat():
     return []
 
 def export_chat(history):
-    return json.dumps(
-        [{"user": u, "assistant": a} for u, a in history],
-        indent=2
-    )
+    return json.dumps(history, indent=2)
 
 # =========================================
 # UI
@@ -223,16 +204,17 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Omniscient IRIS") as demo:
     gr.Markdown("## Omniscient IRIS — Intelligent Analysis Assistant")
 
     with gr.Row():
+
         with gr.Column(scale=3):
 
-            chatbot = gr.Chatbot(height=500)
+            chatbot = gr.Chatbot(height=500, type="messages")
             msg = gr.Textbox(
                 placeholder="Enter material for analysis or ask a question...",
                 show_label=False
             )
 
             with gr.Row():
-                submit_btn = gr.Button("Send", variant="primary")
+                send_btn = gr.Button("Send", variant="primary")
                 stop_btn = gr.Button("Stop")
                 clear_btn = gr.Button("Clear")
 
@@ -265,7 +247,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Omniscient IRIS") as demo:
 
     state = gr.State([])
 
-    submit_event = submit_btn.click(
+    submit_event = send_btn.click(
         stream_response,
         inputs=[msg, state, model_selector, persona_selector, task_selector, file_upload],
         outputs=chatbot,