Upload 2 files

requirements.txt

@@ -0,0 +1,6 @@
+python-dotenv
+streamlit
+openai
+pymongo
+pinecone-client
+uuid

trainer.py

@@ -0,0 +1,628 @@
+import os
+from dotenv import load_dotenv
+import streamlit as st
+from streamlit.runtime.scriptrunner import RerunException, StopException
+from openai import OpenAI
+from pymongo import MongoClient
+from pinecone import Pinecone
+import uuid
+from datetime import datetime
+import time
+from streamlit.runtime.caching import cache_data
+from streamlit_autorefresh import st_autorefresh
+
+# Load environment variables
+load_dotenv()
+
+# Configuration
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+MONGODB_URI = os.getenv("MONGODB_URI")
+PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
+PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT")
+PINECONE_INDEX_NAME = os.getenv("PINECONE_INDEX_NAME")
+GLOBAL_MEMORY_ID = "global_common_memory_id"  # Added GLOBAL_MEMORY_ID
+
+openai_client = OpenAI(api_key=OPENAI_API_KEY)
+mongo_client = MongoClient(MONGODB_URI)
+db = mongo_client["Wall_Street"]
+conversation_history = db["conversation_history"]
+global_common_memory = db["global_common_memory"]  # New global common memory collection
+
+# Initialize GLOBAL_MEMORY_ID if it doesn't exist
+if not global_common_memory.find_one({"memory_id": GLOBAL_MEMORY_ID}):
+    global_common_memory.insert_one({
+        "memory_id": GLOBAL_MEMORY_ID,
+        "memory": []
+    })
+
+# Initialize Pinecone
+pc = Pinecone(api_key=PINECONE_API_KEY)
+pinecone_index = pc.Index(PINECONE_INDEX_NAME)
+
+# Set up Streamlit page configuration
+st.set_page_config(page_title="GPT-Driven Chat System - Tester", page_icon="🔬", layout="wide")
+
+# Custom CSS to improve the UI
+st.markdown("""
+<style>
+    /* Your custom CSS styles */
+</style>
+""", unsafe_allow_html=True)
+
+# Initialize Streamlit session state
+if 'chat_history' not in st.session_state:
+    st.session_state['chat_history'] = []
+if 'user_type' not in st.session_state:
+    st.session_state['user_type'] = None
+if 'session_id' not in st.session_state:
+    st.session_state['session_id'] = str(uuid.uuid4())
+
+# --- Common Memory Functions ---
+
+@cache_data(ttl=300)  # Cache for 5 minutes
+def get_global_common_memory():
+    """Retrieve the global common memory."""
+    memory_doc = global_common_memory.find_one({"memory_id": GLOBAL_MEMORY_ID})
+    return memory_doc.get('memory', []) if memory_doc else []
+
+def append_to_global_common_memory(content):
+    """Append content to the global common memory."""
+    try:
+        # First, ensure the document exists with an initialized memory array
+        global_common_memory.update_one(
+            {"memory_id": GLOBAL_MEMORY_ID},
+            {"$setOnInsert": {"memory": []}},
+            upsert=True
+        )
+        
+        # Then, add the new content to the memory array
+        result = global_common_memory.update_one(
+            {"memory_id": GLOBAL_MEMORY_ID},
+            {"$push": {"memory": content}}
+        )
+        
+        # Invalidate the cache after updating
+        get_global_common_memory.clear()
+        
+        st.success("Memory appended successfully!")
+        
+        # Instead of using st.rerun(), we'll set a flag in session state
+        st.session_state['memory_updated'] = True
+        
+    except Exception as e:
+        st.error(f"Failed to append to global common memory: {str(e)}")
+
+def clear_global_common_memory():
+    """Clear all items from the global common memory."""
+    try:
+        global_common_memory.update_one(
+            {"memory_id": GLOBAL_MEMORY_ID},
+            {"$set": {"memory": []}},
+            upsert=True
+        )
+        # Invalidate the cache after clearing
+        get_global_common_memory.clear()
+        st.success("Global common memory cleared successfully!")
+    except Exception as e:
+        st.error(f"Failed to clear global common memory: {str(e)}")
+
+# --- Relevant Context Retrieval ---
+
+@cache_data(ttl=60)  # Cache for 1 minute
+def get_relevant_context(query, top_k=3):
+    """
+    Retrieve relevant context from Pinecone based on the user query.
+    """
+    try:
+        query_embedding = openai_client.embeddings.create(
+            model="text-embedding-3-large",  # Updated to use the larger model
+            input=query
+        ).data[0].embedding
+
+        results = pinecone_index.query(vector=query_embedding, top_k=top_k, include_metadata=True)
+        contexts = [item['metadata']['text'] for item in results['matches']]
+        return " ".join(contexts)
+    except Exception as e:
+        print(f"Error retrieving context: {str(e)}")
+        return ""
+
+# --- GPT Response Function ---
+
+def get_gpt_response(prompt, context="", system_message=None):
+    try:
+        common_memory = get_global_common_memory()
+        system_msg = (
+            "You are a helpful assistant. Use the following context and global common memory "
+            "to inform your responses, but don't mention them explicitly unless directly relevant to the user's question."
+        )
+        
+        if system_message:
+            system_msg += f"\n\nTrainer Instructions:\n{system_message}"
+        
+        if common_memory:
+            memory_str = "\n".join(common_memory)
+            system_msg += f"\n\nGlobal Common Memory:\n{memory_str}"
+        
+        messages = [
+            {"role": "system", "content": system_msg},
+            {"role": "user", "content": f"Context: {context}\n\nUser query: {prompt}"}
+        ]
+
+        completion = openai_client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=messages
+        )
+        response = completion.choices[0].message.content.strip()
+        
+        return response
+    except Exception as e:
+        st.error(f"Error generating response: {str(e)}")
+        return None
+
+# --- Send User Message ---
+def send_message(message):
+    """
+    Sends a user message. If admin takeover is active, messages are sent to admin instead of GPT.
+    """
+    context = get_relevant_context(message)
+
+    user_message = {
+        "role": "user",
+        "content": message,
+        "timestamp": datetime.utcnow(),
+        "status": "approved"  # User messages are always approved
+    }
+
+    # Upsert the user message immediately
+    result = conversation_history.update_one(
+        {"session_id": st.session_state['session_id']},
+        {
+            "$push": {"messages": user_message},
+            "$set": {"last_updated": datetime.utcnow()},
+            "$setOnInsert": {"created_at": datetime.utcnow()}
+        },
+        upsert=True
+    )
+
+    # Update the session state with the user message
+    st.session_state['chat_history'].append(user_message)
+
+    if not st.session_state.get('admin_takeover_active'):
+        # Generate GPT response if takeover is not active
+        gpt_response = get_gpt_response(message, context)
+
+        assistant_message = {
+            "role": "assistant",
+            "content": gpt_response,
+            "timestamp": datetime.utcnow(),
+            "status": "pending"  # Set status to pending for admin approval
+        }
+
+        # Upsert the assistant message
+        result = conversation_history.update_one(
+            {"session_id": st.session_state['session_id']},
+            {
+                "$push": {"messages": assistant_message},
+                "$set": {"last_updated": datetime.utcnow()}
+            }
+        )
+
+        # Update the session state with the assistant message
+        st.session_state['chat_history'].append(assistant_message)
+
+# --- Send Admin Message ---
+
+def send_admin_message(message):
+    """
+    Sends an admin message directly to the user during a takeover.
+    """
+    admin_message = {
+        "role": "admin",
+        "content": message,
+        "timestamp": datetime.utcnow(),
+        "status": "approved"
+    }
+
+    # Upsert the admin message
+    result = conversation_history.update_one(
+        {"session_id": st.session_state['session_id']},
+        {
+            "$push": {"messages": admin_message},
+            "$set": {"last_updated": datetime.utcnow()}
+        }
+    )
+
+    # Update the session state with the admin message
+    st.session_state['chat_history'].append(admin_message)
+
+# --- Takeover Functions ---
+
+def activate_takeover(session_id):
+    """
+    Activates takeover mode for the given session.
+    """
+    try:
+        db.takeover_status.update_one(
+            {"session_id": session_id},
+            {"$set": {"active": True, "activated_at": datetime.utcnow()}},
+            upsert=True
+        )
+        st.success(f"Takeover activated for session {session_id[:8]}...")
+    except Exception as e:
+        st.error(f"Failed to activate takeover: {str(e)}")
+
+def deactivate_takeover(session_id):
+    """
+    Deactivates takeover mode for the given session.
+    """
+    try:
+        db.takeover_status.update_one(
+            {"session_id": session_id},
+            {"$set": {"active": False}},
+        )
+        st.success(f"Takeover deactivated for session {session_id[:8]}...")
+    except Exception as e:
+        st.error(f"Failed to deactivate takeover: {str(e)}")
+
+def handle_admin_takeover(session_id):
+    st.subheader("Admin Takeover")
+
+    takeover_active = db.takeover_status.find_one({"session_id": session_id})
+    is_active = takeover_active.get("active", False) if takeover_active else False
+
+    if is_active:
+        st.info("Takeover is currently active for this session.")
+        if st.button("Deactivate Takeover"):
+            deactivate_takeover(session_id)
+            st.success("Takeover deactivated.")
+            st.rerun()
+    else:
+        st.warning("Takeover is not active for this session.")
+        if st.button("Activate Takeover"):
+            activate_takeover(session_id)
+            st.success("Takeover activated.")
+            st.rerun()
+
+    if is_active:
+        admin_message = st.text_area("Send Message to User", key="admin_message")
+        if st.button("Send Admin Message"):
+            admin_message = st.session_state.get("admin_message", "")
+            if admin_message.strip():
+                send_admin_message(admin_message.strip()) 
+                st.success("Admin message sent successfully!")
+                st.session_state["admin_message"] = ""
+            else:
+                st.warning("Please enter a message to send.")
+
+# --- View Full Chat (User Perspective) ---
+
+def view_full_chat(session_id):
+    st.title(f"Full Chat View - Session: {session_id[:8]}...")
+    
+    chat = db.chat_history.find_one({"session_id": session_id})
+    
+    if not chat:
+        st.error("Chat not found.")
+        return
+    
+    col1, col2 = st.columns([2, 1])
+    with col1:
+        st.subheader(f"Session ID: {session_id}")
+    with col2:
+        st.write(f"Last Updated: {chat.get('last_updated', 'N/A')}")
+    
+    st.markdown("---")
+
+    for message in chat.get('messages', []):
+        role = message['role'].capitalize()
+        content = message['content']
+        timestamp = message.get('timestamp', 'N/A')
+        
+        if role == 'User':
+            with st.chat_message("user"):
+                st.markdown(f"**User** - {timestamp}")
+                st.markdown(content)
+        elif role == 'Assistant':
+            with st.chat_message("assistant"):
+                st.markdown(f"**Assistant** - {timestamp}")
+                st.markdown(content)
+        elif role == 'Admin':
+            with st.chat_message("human"):
+                st.markdown(f"**Admin** - {timestamp}")
+                st.markdown(content)
+    
+    st.markdown("---")
+    
+    # Add text box to append to global memory
+    st.subheader("Add to Global Memory")
+    new_memory = st.text_area("Enter new memory item", key=f"new_memory_input_{session_id}")
+    if st.button("Add Memory", key=f"add_memory_button_{session_id}"):
+        if new_memory.strip():
+            append_to_global_common_memory(new_memory.strip())
+            st.success("New memory item added to global memory!")
+            # Instead of rerunning, we'll update the session state
+            st.session_state[f'memory_added_{session_id}'] = True
+            st.rerun()
+        else:
+            st.warning("Please enter a valid memory item.")
+    
+    # Display success message if memory was added
+    if st.session_state.get(f'memory_added_{session_id}'):
+        st.success("Memory item added successfully!")
+        # Clear the flag
+        del st.session_state[f'memory_added_{session_id}']
+    
+    st.markdown("---")
+    col1, col2, col3 = st.columns([1, 1, 1])
+    with col2:
+        if st.button("Back to Chat History", use_container_width=True):
+            st.session_state.pop('full_chat_view', None)
+            st.rerun()
+
+# --- Clear Global Chat Memory---
+
+def clear_global_common_memory():
+    """Clear all items from the global common memory."""
+    try:
+        global_common_memory.update_one(
+            {"memory_id": GLOBAL_MEMORY_ID},
+            {"$set": {"memory": []}},
+            upsert=True
+        )
+        # Invalidate the cache after clearing
+        get_global_common_memory.clear()
+        st.success("Global common memory cleared successfully!")
+    except Exception as e:
+        st.error(f"Failed to clear global common memory: {str(e)}")
+
+def display_chat_history():
+    st.subheader("All Chat History")
+    
+    all_chats = list(db.chat_history.find().sort("last_updated", -1))
+    
+    if not all_chats:
+        st.info("No chat history found.")
+        return
+    
+    for idx, chat in enumerate(all_chats):
+        session_id = chat['session_id']
+        last_updated = chat.get('last_updated', 'N/A')
+        
+        with st.expander(f"Session: {session_id[:8]}... - Last Updated: {last_updated}"):
+            if chat.get('messages'):
+                last_message = chat['messages'][-1]
+                st.markdown(f"Last message ({last_message['role'].capitalize()}):")
+                st.markdown(f"> {last_message['content'][:100]}...")
+            
+            if st.button(f"Show Full Chat", key=f"show_full_chat_{idx}"):
+                st.session_state['full_chat_view'] = session_id
+                st.rerun()
+
+def trainer_intervention_tab():
+    st.subheader("Trainer Intervention")
+    
+    # Handle admin intervention
+    handle_admin_intervention()
+
+def handle_admin_intervention():
+    st.subheader("Review Pending Responses")
+    pending_responses = conversation_history.find(
+        {"messages.role": "assistant", "messages.status": "pending"}
+    )
+
+    for conversation in pending_responses:
+        st.write(f"Session ID: {conversation['session_id'][:8]}...")
+
+        for i, message in enumerate(conversation['messages']):
+            if message['role'] == 'assistant' and message.get('status') == 'pending':
+                user_message = conversation['messages'][i-1]['content'] if i > 0 else "N/A"
+                st.write(f"**User:** {user_message}")
+                st.write(f"**GPT:** {message['content']}")
+
+                col1, col2, col3, col4 = st.columns(4)
+                with col1:
+                    if st.button("Approve", key=f"approve_{conversation['session_id']}_{i}"):
+                        if approve_response(conversation['session_id'], i):
+                            st.success("Response approved")
+                            time.sleep(0.5)
+                            st.rerun()
+                with col2:
+                    if st.button("Modify", key=f"modify_{conversation['session_id']}_{i}"):
+                        st.session_state['modifying'] = (conversation['session_id'], i)
+                        st.rerun()
+                with col3:
+                    if st.button("Regenerate", key=f"regenerate_{conversation['session_id']}_{i}"):
+                        st.session_state['regenerating'] = (conversation['session_id'], i)
+                        st.rerun()
+                with col4:
+                    takeover_doc = db.takeover_status.find_one({"session_id": conversation['session_id']})
+                    takeover_active = takeover_doc.get("active", False) if takeover_doc else False
+                    if takeover_active:
+                        if st.button("Deactivate Takeover", key=f"deactivate_takeover_{conversation['session_id']}_{i}"):
+                            deactivate_takeover(conversation['session_id'])
+                            st.success("Takeover deactivated.")
+                            st.rerun()
+                    else:
+                        if st.button("Activate Takeover", key=f"activate_takeover_{conversation['session_id']}_{i}"):
+                            activate_takeover(conversation['session_id'])
+                            st.success("Takeover activated.")
+                            st.rerun()
+
+        st.divider()
+
+    if 'regenerating' in st.session_state:
+        session_id, message_index = st.session_state['regenerating']
+        with st.form(key="regenerate_form"):
+            operator_input = st.text_input("Enter additional instructions for regeneration:")
+            submit_button = st.form_submit_button("Submit")
+
+        if submit_button:
+            del st.session_state['regenerating']
+            regenerate_response(session_id, message_index, operator_input)
+            st.success("Response regenerated with operator input.")
+            st.rerun()
+
+    if 'modifying' in st.session_state:
+        session_id, message_index = st.session_state['modifying']
+        conversation = conversation_history.find_one({"session_id": session_id})
+        message = conversation['messages'][message_index]
+
+        modified_content = st.text_area("Modify the response:", value=message['content'])
+        if st.button("Save Modification"):
+            save_modified_response(session_id, message_index, modified_content)
+            st.success("Response modified and approved")
+            del st.session_state['modifying']
+            st.rerun()
+
+def approve_response(session_id, message_index):
+    try:
+        result = conversation_history.update_one(
+            {"session_id": session_id},
+            {"$set": {f"messages.{message_index}.status": "approved"}}
+        )
+        return result.modified_count > 0
+    except Exception as e:
+        st.error(f"Failed to approve response: {str(e)}")
+        return False
+
+def save_modified_response(session_id, message_index, modified_content):
+    try:
+        conversation_history.update_one(
+            {"session_id": session_id},
+            {
+                "$set": {
+                    f"messages.{message_index}.content": modified_content,
+                    f"messages.{message_index}.status": "approved"
+                }
+            }
+        )
+    except Exception as e:
+        st.error(f"Failed to save modified response: {str(e)}")
+
+def regenerate_response(session_id, message_index, operator_input):
+    try:
+        conversation = conversation_history.find_one({"session_id": session_id})
+        user_message = conversation['messages'][message_index - 1]['content'] if message_index > 0 else ""
+        new_response = get_gpt_response(user_message, system_message=operator_input)
+        
+        conversation_history.update_one(
+            {"session_id": session_id},
+            {
+                "$set": {
+                    f"messages.{message_index}.content": new_response,
+                    f"messages.{message_index}.status": "pending"
+                }
+            }
+        )
+    except Exception as e:
+        st.error(f"Failed to regenerate response: {str(e)}")
+
+def trainer_page():
+    st.title("Trainer Dashboard")
+    
+    # Add auto-refresh every 10 seconds (10000 milliseconds)
+    st_autorefresh(interval=10000, limit=None, key="trainer_autorefresh")
+    
+    tab1, tab2, tab3 = st.tabs(["Current Status", "Chat History", "Intervention"])
+    
+    with tab1:
+        # Display current global memory
+        st.subheader("Current Global Memory")
+        global_memory = get_global_common_memory()
+        if global_memory:
+            for idx, item in enumerate(global_memory, 1):
+                st.text(f"{idx}. {item}")
+        else:
+            st.info("No global memory items found.")
+        
+        # Add button to clear global memory
+        if st.button("Clear Global Memory", key="clear_global_memory"):
+            clear_global_common_memory()
+            st.success("Global memory cleared successfully!")
+            time.sleep(1)
+            st.rerun()
+        
+        # Display current chats
+        st.subheader("Active Chats")
+        chats = list(conversation_history.find().sort("last_updated", -1).limit(5))
+        for idx, chat in enumerate(chats):
+            with st.expander(f"Session: {chat['session_id'][:8]}... - Last Updated: {chat.get('last_updated', 'N/A')}"):
+                for message in chat.get('messages', [])[-5:]:
+                    role = message['role'].capitalize()
+                    content = message['content']
+                    st.markdown(f"**{role}:** {content}")
+                
+                col1, col2, col3, col4 = st.columns(4)
+                with col1:
+                    if st.button(f"View Full Chat", key=f"view_chat_{idx}"):
+                        st.session_state['selected_chat'] = chat['session_id']
+                        st.rerun()
+                with col2:
+                    takeover_doc = db.takeover_status.find_one({"session_id": chat['session_id']})
+                    takeover_active = takeover_doc.get("active", False) if takeover_doc else False
+                    if takeover_active:
+                        if st.button(f"Deactivate Takeover", key=f"deactivate_takeover_{idx}"):
+                            deactivate_takeover(chat['session_id'])
+                            st.success("Takeover deactivated.")
+                            st.rerun()
+                    else:
+                        if st.button(f"Activate Takeover", key=f"activate_takeover_{idx}"):
+                            activate_takeover(chat['session_id'])
+                            st.success("Takeover activated.")
+                            st.rerun()
+                with col3:
+                    if st.button(f"Delete Chat", key=f"delete_chat_{idx}"):
+                        delete_chat(chat['session_id'])
+                        st.success(f"Chat {chat['session_id'][:8]}... deleted.")
+                        st.rerun()
+                with col4:
+                    if takeover_active:
+                        st.text_input("Send message", key=f"takeover_message_{idx}")
+                        if st.button("Send", key=f"send_takeover_{idx}"):
+                            message = st.session_state[f"takeover_message_{idx}"]
+                            if message.strip():
+                                send_admin_message(chat['session_id'], message.strip())
+                                st.success("Message sent.")
+                                st.rerun()
+                            else:
+                                st.warning("Please enter a message to send.")
+        
+        # Manual refresh button
+        if st.button("Refresh", key="refresh_button"):
+            st.rerun()
+    
+    with tab2:
+        display_chat_history()
+    
+    with tab3:
+        trainer_intervention_tab()
+
+def delete_chat(session_id):
+    try:
+        result = conversation_history.delete_one({"session_id": session_id})
+        if result.deleted_count == 0:
+            st.error("Failed to delete chat. Please try again.")
+    except Exception as e:
+        st.error(f"Error deleting chat: {str(e)}")
+
+# --- Main Function ---
+
+def main():
+    try:
+        if 'memory_updated' in st.session_state:
+            del st.session_state['memory_updated']
+            st.rerun()
+        
+        if 'full_chat_view' in st.session_state:
+            view_full_chat(st.session_state['full_chat_view'])
+        elif 'selected_chat' in st.session_state:
+            view_full_chat(st.session_state['selected_chat'])
+        else:
+            trainer_page()
+    except (RerunException, StopException):
+        raise
+    except Exception as e:
+        st.error(f"An unexpected error occurred: {str(e)}")
+
+if __name__ == "__main__":
+    main()