Update app.py

app.py

@@ -1,73 +1,79 @@
 import os
 import json
+import time
 import tempfile
+from datetime import datetime
+
 import gradio as gr
 from huggingface_hub import InferenceClient
 from gtts import gTTS
 
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
+from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
+from reportlab.lib import colors
+from reportlab.lib.units import inch
+
 try:
     import PyPDF2
 except ImportError:
     PyPDF2 = None
 
 
-# =========================================
-# Configuration
-# =========================================
+# =========================================================
+# CONFIGURATION
+# =========================================================
 
 MODEL_OPTIONS = {
-    "Zephyr 7B (Serverless)": "HuggingFaceH4/zephyr-7b-beta",
-    "Mixtral 8x7B (Serverless)": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "Zephyr 7B (Stable)": "HuggingFaceH4/zephyr-7b-beta",
+    "Mixtral 8x7B": "mistralai/Mixtral-8x7B-Instruct-v0.1",
 }
 
 FALLBACK_MODEL = "HuggingFaceH4/zephyr-7b-beta"
 
+MAX_CONTEXT_CHARS = 10000
+MAX_HISTORY_PAIRS = 6
+MAX_INPUT_CHARS = 3000
+CHAT_LOG_FILE = "chat_master_log.json"
+
+
 PERSONA_PRESETS = {
     "Balanced Assistant": "You are a helpful, intelligent AI assistant.",
-
     "Forensic Psychologist": """
 You are acting as a forensic psychologist.
 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.
 """
 }
 
 TASK_MODES = {
     "General Chat": "{user_input}",
-
     "Behavioral Analysis": """
-Conduct a behavioral analysis.
+Conduct a behavioral analysis:
 1. Behavioral Indicators
 2. Emotional Tone
 3. Cognitive Patterns
-4. Stress/Deception Markers
-5. Risk-Relevant Observations
-6. Alternative Explanations
-7. Limitations
+4. Risk Observations
+5. Alternative Explanations
+6. Limitations
+
 Material:
 {user_input}
 """,
-
-    "Explain Code (Natural Voice)": """
+    "Explain Code": """
 Explain the following code clearly and conversationally,
-as if teaching a developer step-by-step in a calm, natural voice.
+as if teaching a developer step-by-step.
 
 Code:
 {user_input}
 """
 }
 
-MAX_CONTEXT_CHARS = 10000
-MAX_HISTORY_PAIRS = 6
-MAX_INPUT_CHARS = 3000
 
-
-# =========================================
-# Helpers
-# =========================================
+# =========================================================
+# INFERENCE LAYER
+# =========================================================
 
 def get_client(model_id):
     token = os.environ.get("HF_TOKEN")
@@ -76,6 +82,39 @@ def get_client(model_id):
     return InferenceClient(model=model_id, token=token)
 
 
+def run_with_retry(client, messages, temperature, retries=3, delay=2):
+    for attempt in range(retries):
+        try:
+            response = client.chat_completion(
+                messages=messages,
+                max_tokens=700,
+                temperature=temperature,
+                top_p=0.95,
+                stream=False,
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            if attempt < retries - 1:
+                time.sleep(delay)
+            else:
+                raise e
+
+
+def generate_ai_response(model_id, messages, temperature):
+    try:
+        client = get_client(model_id)
+        return run_with_retry(client, messages, temperature)
+    except Exception:
+        fallback_client = get_client(FALLBACK_MODEL)
+        return "⚠️ Fallback model used.\n\n" + run_with_retry(
+            fallback_client, messages, temperature
+        )
+
+
+# =========================================================
+# UTILITIES
+# =========================================================
+
 def extract_text(file):
     if not file:
         return ""
@@ -104,87 +143,112 @@ def text_to_speech(text):
     return temp_file.name
 
 
-def run_model(client, messages, temperature):
-    response = client.chat_completion(
-        messages=messages,
-        max_tokens=700,
-        temperature=temperature,
-        top_p=0.95,
-        stream=False,
+def append_to_master_log(user_msg, assistant_msg):
+    entry = {
+        "timestamp": datetime.utcnow().isoformat(),
+        "user": user_msg,
+        "assistant": assistant_msg
+    }
+
+    if os.path.exists(CHAT_LOG_FILE):
+        with open(CHAT_LOG_FILE, "r") as f:
+            data = json.load(f)
+    else:
+        data = []
+
+    data.append(entry)
+
+    with open(CHAT_LOG_FILE, "w") as f:
+        json.dump(data, f, indent=2)
+
+
+def export_pdf(history):
+    file_path = f"transcript_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.pdf"
+    doc = SimpleDocTemplate(file_path)
+    elements = []
+
+    styles = getSampleStyleSheet()
+    user_style = ParagraphStyle(
+        "UserStyle",
+        parent=styles["Normal"],
+        textColor=colors.blue,
+        spaceAfter=6
+    )
+
+    assistant_style = ParagraphStyle(
+        "AssistantStyle",
+        parent=styles["Normal"],
+        textColor=colors.darkgreen,
+        spaceAfter=12
     )
-    return response.choices[0].message.content
 
+    elements.append(Paragraph("Omniscient IRIS Transcript", styles["Heading1"]))
+    elements.append(Spacer(1, 0.3 * inch))
 
-# =========================================
-# Core Logic
-# =========================================
+    for msg in history:
+        role = msg["role"]
+        content = msg["content"].replace("\n", "<br/>")
 
-def generate_response(message, history, model_label, persona_label, task_mode, uploaded_file, voice_enabled):
+        if role == "user":
+            elements.append(Paragraph(f"<b>User:</b> {content}", user_style))
+        elif role == "assistant":
+            elements.append(Paragraph(f"<b>Assistant:</b> {content}", assistant_style))
+
+    doc.build(elements)
+    return file_path
+
+
+# =========================================================
+# CORE CHAT LOGIC
+# =========================================================
+
+def generate_response(message, history, model_label, persona_label,
+                      task_mode, uploaded_file, voice_enabled):
 
     if not message:
         return history, None
 
-    try:
-        message = message[:MAX_INPUT_CHARS]
-        history = history[-MAX_HISTORY_PAIRS * 2:]
-
-        model_id = MODEL_OPTIONS[model_label]
-        system_prompt = PERSONA_PRESETS.get(persona_label, "")
-        temperature = 0.4 if "Forensic" in persona_label else 0.7
+    message = message[:MAX_INPUT_CHARS]
+    history = history[-MAX_HISTORY_PAIRS * 2:]
 
-        file_text = extract_text(uploaded_file)
-        if file_text:
-            file_text = file_text[:MAX_CONTEXT_CHARS]
-            system_prompt += f"\n\nUse this document context when relevant:\n\n{file_text}"
+    model_id = MODEL_OPTIONS[model_label]
+    system_prompt = PERSONA_PRESETS[persona_label]
+    temperature = 0.4 if "Forensic" in persona_label else 0.7
 
-        messages = [{"role": "system", "content": system_prompt}]
-        messages.extend(history)
+    file_text = extract_text(uploaded_file)
+    if file_text:
+        file_text = file_text[:MAX_CONTEXT_CHARS]
+        system_prompt += f"\n\nDocument Context:\n{file_text}"
 
-        formatted_input = TASK_MODES[task_mode].format(user_input=message)
-        messages.append({"role": "user", "content": formatted_input})
+    messages = [{"role": "system", "content": system_prompt}]
+    messages.extend(history)
 
-        try:
-            client = get_client(model_id)
-            answer = run_model(client, messages, temperature)
-        except Exception:
-            fallback_client = get_client(FALLBACK_MODEL)
-            answer = "⚠️ Fallback model used.\n\n" + run_model(fallback_client, messages, temperature)
+    formatted_input = TASK_MODES[task_mode].format(user_input=message)
+    messages.append({"role": "user", "content": formatted_input})
 
-        history.append({"role": "user", "content": message})
-        history.append({"role": "assistant", "content": answer})
+    answer = generate_ai_response(model_id, messages, temperature)
 
-        audio_file = None
-        if voice_enabled:
-            audio_file = text_to_speech(answer)
+    history.append({"role": "user", "content": message})
+    history.append({"role": "assistant", "content": answer})
 
-        return history, audio_file
+    append_to_master_log(message, answer)
 
-    except Exception as e:
-        history.append({"role": "assistant", "content": f"⚠️ ERROR:\n\n{str(e)}"})
-        return history, None
+    audio_file = text_to_speech(answer) if voice_enabled else None
 
+    return history, audio_file
 
-# =========================================
-# Utilities
-# =========================================
 
 def clear_chat():
     return [], None
 
-def export_chat(history):
-    file_path = "conversation.json"
-    with open(file_path, "w") as f:
-        json.dump(history, f, indent=2)
-    return file_path
 
-
-# =========================================
+# =========================================================
 # UI
-# =========================================
+# =========================================================
 
 with gr.Blocks(theme=gr.themes.Soft(), title="Omniscient IRIS") as demo:
 
-    gr.Markdown("## Omniscient IRIS — Voice Enabled")
+    gr.Markdown("## Omniscient IRIS — Resilient Voice Enabled Assistant")
 
     with gr.Row():
 
@@ -192,7 +256,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Omniscient IRIS") as demo:
 
             chatbot = gr.Chatbot(height=500, type="messages")
             msg = gr.Textbox(
-                placeholder="Enter material for analysis or ask a question...",
+                placeholder="Enter material for analysis...",
                 show_label=False
             )
 
@@ -221,26 +285,33 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Omniscient IRIS") as demo:
                 label="Mode"
             )
 
-            voice_toggle = gr.Checkbox(label="Enable Voice Output", value=True)
+            voice_toggle = gr.Checkbox(
+                label="Enable Voice Output",
+                value=True
+            )
 
             file_upload = gr.File(
-                label="Upload Document (txt, pdf, md)",
+                label="Upload Document",
                 file_types=[".txt", ".pdf", ".md"]
             )
 
-            export_btn = gr.Button("Export Conversation")
+            pdf_btn = gr.Button("Download PDF Transcript")
 
     state = gr.State([])
 
     send_btn.click(
         generate_response,
-        inputs=[msg, state, model_selector, persona_selector, task_selector, file_upload, voice_toggle],
+        inputs=[msg, state, model_selector,
+                persona_selector, task_selector,
+                file_upload, voice_toggle],
         outputs=[chatbot, audio_output],
     )
 
     msg.submit(
         generate_response,
-        inputs=[msg, state, model_selector, persona_selector, task_selector, file_upload, voice_toggle],
+        inputs=[msg, state, model_selector,
+                persona_selector, task_selector,
+                file_upload, voice_toggle],
         outputs=[chatbot, audio_output],
     )
 
@@ -249,10 +320,10 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Omniscient IRIS") as demo:
         outputs=[chatbot, audio_output],
     ).then(lambda: [], outputs=state)
 
-    export_btn.click(
-        export_chat,
+    pdf_btn.click(
+        export_pdf,
         inputs=chatbot,
-        outputs=gr.File(label="Download JSON")
+        outputs=gr.File(label="Download PDF")
     )
 
 if __name__ == "__main__":