Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py (Gradio version for Hugging Face Spaces)
+# app.py (Updated for 3-Chatbot System)
 import gradio as gr
 import openai
 import threading
@@ -7,14 +7,21 @@ import numpy as np
 import faiss
 import os
 import pickle
-import openai
-import numpy as np
+from datetime import datetime
 
-def get_embedding(text, model="text-embedding-3-small"):
-    response = openai.embeddings.create(
-        input=[text],
-        model=model
-    )
+# Set API key
+openai.api_key = os.environ.get("OPENAI_API_KEY")
+
+# === CONFIG ===
+EMBEDDING_MODEL = "text-embedding-3-small"
+CHAT_MODEL = "gpt-4o"  # Updated to current model
+MEMORY_FILE = "memory.pkl"
+INDEX_FILE = "memory.index"
+
+# === EMBEDDING UTILS ===
+def get_embedding(text, model=EMBEDDING_MODEL):
+    text = text.replace("\n", " ")
+    response = openai.embeddings.create(input=[text], model=model)
     return response.data[0].embedding
 
 def cosine_similarity(vec1, vec2):
@@ -22,105 +29,169 @@ def cosine_similarity(vec1, vec2):
     vec2 = np.array(vec2)
     return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
 
-# === CONFIG ===
-openai.api_key = os.environ.get("OPENAI_API_KEY")
-EMBEDDING_MODEL = "text-embedding-ada-002"
-CHAT_MODEL = "gpt-4o-mini"
-
-# === MEMORY ===
+# === MEMORY INITIALIZATION ===
 memory_data = []
 try:
-    memory_index = faiss.read_index("memory.index")
-    with open("memory.pkl", "rb") as f:
+    memory_index = faiss.read_index(INDEX_FILE)
+    with open(MEMORY_FILE, "rb") as f:
         memory_data = pickle.load(f)
 except:
-    memory_index = faiss.IndexFlatL2(1536)
+    memory_index = faiss.IndexFlatL2(1536)  # 1536 dimensions for text-embedding-3-small
 
 # === SYSTEM PROMPTS ===
-AGENT_PROMPT = "You are a helpful agent in an ongoing dialogue. Respond meaningfully."
-OVERSEER_PROMPT = "You are the Overseer agent. Monitor Agent A and B, learn, and intervene when appropriate."
+AGENT_A_PROMPT = "You are Agent A, initiating conversations with thoughtful questions."
+AGENT_B_PROMPT = "You are Agent B, responding to Agent A with insightful answers."
+OVERSEER_PROMPT = """You are the Overseer (Agent C). Monitor conversations between Agent A and B. 
+Intervene when discussions become repetitive or need redirection. Ask thought-provoking questions 
+to explore new dimensions of the topic."""
 
+# === GLOBAL STATE ===
 conversation = []
+turn_count = 0
 auto_mode = False
 
-# === AGENT ===
-def chat_completion(system, messages):
+# === CHAT COMPLETION ===
+def chat_completion(system, messages, model=CHAT_MODEL):
     try:
-        response = openai.ChatCompletion.create(
-            model=CHAT_MODEL,
+        response = openai.chat.completions.create(
+            model=model,
             messages=[{"role": "system", "content": system}] + messages,
-            temperature=0.7
+            temperature=0.7,
+            max_tokens=150
         )
         return response.choices[0].message.content.strip()
     except Exception as e:
-        return f"[Error: {e}]"
+        return f"[API Error: {str(e)}]"
 
-# === FAISS EMBED ===
+# === MEMORY MANAGEMENT ===
 def embed_and_store(text):
     try:
-        vec = get_embedding(text, engine=EMBEDDING_MODEL)
+        vec = get_embedding(text)
         memory_index.add(np.array([vec], dtype='float32'))
-        memory_data.append(text)
-        with open("memory.pkl", "wb") as f:
-            pickle.dump(memory_data, f)
-        faiss.write_index(memory_index, "memory.index")
+        memory_data.append({
+            "text": text,
+            "timestamp": datetime.now().isoformat()
+        })
+        # Periodic save to avoid constant I/O
+        if len(memory_data) % 5 == 0:
+            with open(MEMORY_FILE, "wb") as f:
+                pickle.dump(memory_data, f)
+            faiss.write_index(memory_index, INDEX_FILE)
     except Exception as e:
-        print(f"Embed error: {e}")
-
-# === CONVERSATION ===
-def step():
-    global conversation
-    turn = len(conversation)
-    agent = "Agent A" if turn % 2 == 0 else "Agent B"
-    msgs = [{"role": "assistant", "content": m['text']} for m in conversation]
-    reply = chat_completion(AGENT_PROMPT, msgs)
-    conversation.append({"agent": agent, "text": reply})
-    embed_and_store(reply)
-    return format_convo(), ""
+        print(f"Memory Error: {str(e)}")
 
+# === CONVERSATION MANAGEMENT ===
 def format_convo():
     return "\n".join([f"**{m['agent']}**: {m['text']}" for m in conversation])
 
-# === OVERSEER ===
+def detect_repetition():
+    """Check if recent messages are similar using embeddings"""
+    if len(conversation) < 4:
+        return False
+        
+    # Get embeddings of last 2 pairs
+    recent = [m['text'] for m in conversation[-4:]]
+    embeddings = [get_embedding(text) for text in recent]
+    
+    # Compare current with 2 messages back
+    similarity = cosine_similarity(embeddings[-1], embeddings[-3])
+    print(f"Similarity: {similarity:.4f}")
+    return similarity > 0.85
+
+# === CORE CONVERSATION FLOW ===
+def step():
+    global conversation, turn_count
+    
+    if not conversation:  # Initial message
+        msg = chat_completion(AGENT_A_PROMPT, [])
+        conversation.append({"agent": "Agent A", "text": msg})
+        embed_and_store(msg)
+        turn_count = 0
+        return format_convo(), ""
+    
+    # Agent B responds to last message
+    last_msg = conversation[-1]['text']
+    b_msg = chat_completion(
+        AGENT_B_PROMPT, 
+        [{"role": "user", "content": last_msg}]
+    )
+    conversation.append({"agent": "Agent B", "text": b_msg})
+    embed_and_store(b_msg)
+    
+    # Agent A responds to Agent B
+    a_msg = chat_completion(
+        AGENT_A_PROMPT, 
+        [{"role": "user", "content": b_msg}]
+    )
+    conversation.append({"agent": "Agent A", "text": a_msg})
+    embed_and_store(a_msg)
+    
+    # Overseer intervention logic
+    intervention = None
+    if turn_count % 3 == 0 or detect_repetition():
+        context = "\n".join([m['text'] for m in conversation[-4:]])
+        prompt = f"Conversation Context:\n{context}\n\nIntervene to redirect or deepen the discussion:"
+        intervention = chat_completion(OVERSEER_PROMPT, [{"role": "user", "content": prompt}])
+        conversation.append({"agent": "Overseer", "text": intervention})
+        embed_and_store(intervention)
+    
+    turn_count += 1
+    return format_convo(), intervention or ""
+
+# === OVERSEER QUERY HANDLER ===
 def overseer_respond(query):
     try:
-        qvec = get_embedding(query, engine=EMBEDDING_MODEL)
-        sims = cosine_similarity(qvec, [get_embedding(m, engine=EMBEDDING_MODEL) for m in memory_data])
-        top_idxs = np.argsort(sims)[-3:][::-1]
-        context = "\n".join([memory_data[i] for i in top_idxs])
-        msgs = [{"role": "user", "content": f"Context:\n{context}\nQuestion:{query}"}]
-        return chat_completion(OVERSEER_PROMPT, msgs)
+        # Add context from recent conversation
+        context = "\n".join([m['text'] for m in conversation[-3:]]) if conversation else "No context"
+        messages = [
+            {"role": "user", "content": f"Recent conversation:\n{context}\n\nQuery: {query}"}
+        ]
+        return chat_completion(OVERSEER_PROMPT, messages)
     except Exception as e:
-        return f"[Overseer Error: {e}]"
+        return f"[Overseer Error: {str(e)}]"
 
-# === AUTO LOOP ===
+# === AUTO MODE HANDLER ===
 def auto_loop():
     global auto_mode
     while auto_mode:
         step()
         time.sleep(5)
 
+def toggle_auto():
+    global auto_mode
+    auto_mode = not auto_mode
+    if auto_mode:
+        threading.Thread(target=auto_loop, daemon=True).start()
+    return "🔴 Auto: OFF" if not auto_mode else "🟢 Auto: ON"
+
 # === GRADIO UI ===
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Tri-Agent Conversational AI (Hugging Face Edition)")
+    gr.Markdown("# 🤖 Tri-Agent Conversational System")
+    gr.Markdown("**Agents**: A (Initiator) → B (Responder) → C (Overseer)")
+    
     with gr.Row():
-        convo_display = gr.Markdown(value="")
-        step_btn = gr.Button("Manual Step")
-        auto_btn = gr.Button("Toggle Auto Mode")
-
-    with gr.Accordion("🧠 Overseer Panel", open=False):
-        qbox = gr.Textbox(label="Ask the Overseer")
-        overseer_out = gr.Textbox(label="Overseer's Response")
-
-    def toggle_auto():
-        global auto_mode
-        auto_mode = not auto_mode
-        if auto_mode:
-            threading.Thread(target=auto_loop, daemon=True).start()
-        return "Auto Mode: ON" if auto_mode else "Auto Mode: OFF"
-
+        convo_display = gr.Markdown(value="**Conversation will appear here**")
+        
+    with gr.Row():
+        step_btn = gr.Button("▶️ Next Conversation Step")
+        auto_btn = gr.Button("🔴 Auto: OFF", variant="secondary")
+        clear_btn = gr.Button("🔄 Reset Conversation")
+    
+    with gr.Accordion("🧠 Overseer Query Panel", open=False):
+        gr.Markdown("Ask the Overseer (Agent C) for insights:")
+        qbox = gr.Textbox(label="Your Question", placeholder="What should we discuss next?")
+        overseer_out = gr.Textbox(label="Overseer's Response", interactive=False)
+    
+    # Event handlers
+    def clear_convo():
+        global conversation, turn_count
+        conversation = []
+        turn_count = 0
+        return "**Conversation reset**", "🔴 Auto: OFF"
+    
     step_btn.click(step, outputs=[convo_display, overseer_out])
     qbox.submit(overseer_respond, inputs=qbox, outputs=overseer_out)
     auto_btn.click(toggle_auto, outputs=auto_btn)
+    clear_btn.click(clear_convo, outputs=[convo_display, auto_btn])
 
 demo.launch()