app.py

# app.py (Compatible with OpenAI v0.x and v1.x)
import gradio as gr
import openai
import threading
import time
import numpy as np
import faiss
import os
import pickle
from datetime import datetime

# === CONFIG ===
EMBEDDING_MODEL = "text-embedding-3-small"
CHAT_MODEL = "gpt-4o"  # Updated to current model
MEMORY_FILE = "memory.pkl"
INDEX_FILE = "memory.index"

# Initialize OpenAI API
openai.api_key = os.environ.get("OPENAI_API_KEY")

# === EMBEDDING UTILS ===
def get_embedding(text, model=EMBEDDING_MODEL):
    text = text.replace("\n", " ")
    
    # Compatible API call for both v0.x and v1.x
    try:
        # Try v1.x API first
        response = openai.embeddings.create(input=[text], model=model)
        return response.data[0].embedding
    except AttributeError:
        # Fallback to v0.x API
        response = openai.Embedding.create(input=[text], model=model)
        return response['data'][0]['embedding']

def cosine_similarity(vec1, vec2):
    vec1 = np.array(vec1)
    vec2 = np.array(vec2)
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

# === MEMORY INITIALIZATION ===
memory_data = []
try:
    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)  # 1536 dimensions for text-embedding-3-small

# === SYSTEM PROMPTS ===
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

# === CHAT COMPLETION (Compatible with both APIs) ===
def chat_completion(system, messages, model=CHAT_MODEL):
    try:
        # Build message list with system prompt
        full_messages = [{"role": "system", "content": system}] 
        full_messages.extend(messages)
        
        # Try v1.x API first
        try:
            response = openai.chat.completions.create(
                model=model,
                messages=full_messages,
                temperature=0.7,
                max_tokens=150
            )
            return response.choices[0].message.content.strip()
        except AttributeError:
            # Fallback to v0.x API
            response = openai.ChatCompletion.create(
                model=model,
                messages=full_messages,
                temperature=0.7,
                max_tokens=150
            )
            return response['choices'][0]['message']['content'].strip()
            
    except Exception as e:
        return f"[API Error: {str(e)}]"

# === MEMORY MANAGEMENT ===
def embed_and_store(text):
    try:
        vec = get_embedding(text)
        memory_index.add(np.array([vec], dtype='float32'))
        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"Memory Error: {str(e)}")

# === CONVERSATION MANAGEMENT ===
def format_convo():
    return "\n".join([f"**{m['agent']}**: {m['text']}" for m in conversation])

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:
        # 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: {str(e)}]"

# === 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 System")
    gr.Markdown("**Agents**: A (Initiator) → B (Responder) → C (Overseer)")
    
    with gr.Row():
        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()