Properly merge original Discord bot with JurySystem DBT integration

bot/bot.py

@@ -1,14 +1,53 @@
+"""
+bot.py - Discord bot client with session management and command routing
+
+Features:
+- Login flow with username/password
+- Session management with JWT tokens
+- AI-powered command parsing via registry
+- Background task loop for polling
+- JurySystem DBT integration for mental health support
+"""
+
+import discord
+from discord.ext import tasks
+import os
+import sys
 import json
-import time
+import base64
+import requests
+import bcrypt
+import pickle
 import numpy as np
 from openai import OpenAI
 
-# --- Configuration ---
+from bot.command_registry import get_handler, list_registered, register_module
-CONFIG_PATH = 'config.json'
+import ai.parser as ai_parser
-KNOWLEDGE_BASE_PATH = 'dbt_knowledge.json'
+import bot.commands.routines  # noqa: F401 - registers handler
+import bot.commands.medications  # noqa: F401 - registers handler
+import bot.commands.knowledge  # noqa: F401 - registers handler
+
+DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
+API_URL = os.getenv("API_URL", "http://app:5000")
+
+user_sessions = {}
+login_state = {}
+message_history = {}
+user_cache = {}
+CACHE_FILE = "/app/user_cache.pkl"
+
+intents = discord.Intents.default()
+intents.message_content = True
+
+client = discord.Client(intents=intents)
+
+
+# ==================== JurySystem Integration ====================
+
 
 class SimpleVectorStore:
     """A simple in-memory vector store using NumPy."""
+
     def __init__(self):
         self.vectors = []
         self.metadata = []
@@ -21,175 +60,522 @@ class SimpleVectorStore:
         if not self.vectors:
             return []
 
-        # Convert to numpy arrays for efficient math
         query_vec = np.array(query_vector)
         doc_vecs = np.array(self.vectors)
-
-        # Cosine Similarity: (A . B) / (||A|| * ||B||)
-        # Note: Both vectors must have the same dimension (e.g., 4096)
         norms = np.linalg.norm(doc_vecs, axis=1)
-        
-        # Avoid division by zero
         valid_indices = norms > 0
         scores = np.zeros(len(doc_vecs))
-        
-        # Calculate dot product
         dot_products = np.dot(doc_vecs, query_vec)
-        
+        scores[valid_indices] = dot_products[valid_indices] / (
-        # Calculate cosine similarity only for valid norms
+            norms[valid_indices] * np.linalg.norm(query_vec)
-        scores[valid_indices] = dot_products[valid_indices] / (norms[valid_indices] * np.linalg.norm(query_vec))
+        )
-        
-        # Get top_k indices
         top_indices = np.argsort(scores)[-top_k:][::-1]
 
         results = []
         for idx in top_indices:
-            results.append({
+            results.append({"metadata": self.metadata[idx], "score": scores[idx]})
-                "metadata": self.metadata[idx],
-                "score": scores[idx]
-            })
         return results
 
-class JurySystem:
-    def __init__(self):
-        self.config = self.load_config()
 
-        # Initialize OpenRouter Client
+class JurySystem:
+    """DBT Knowledge Base Query System"""
+
+    def __init__(self):
+        config_path = os.getenv("CONFIG_PATH", "config.json")
+        kb_path = os.getenv(
+            "KNOWLEDGE_BASE_PATH", "bot/data/dbt_knowledge.embeddings.json"
+        )
+
+        with open(config_path, "r") as f:
+            self.config = json.load(f)
+
         self.client = OpenAI(
             base_url="https://openrouter.ai/api/v1",
-            api_key=self.config['openrouter_api_key']
+            api_key=self.config["openrouter_api_key"],
         )
-        
         self.vector_store = SimpleVectorStore()
-        self.load_knowledge_base()
+        self._load_knowledge_base(kb_path)
 
-    def load_config(self):
+    def _load_knowledge_base(self, kb_path):
-        with open(CONFIG_PATH, 'r') as f:
+        print(f"Loading DBT knowledge base from {kb_path}...")
-            return json.load(f)
-
-    def load_knowledge_base(self):
-        """Loads the pre-computed embeddings from the JSON file."""
-        print(f"Loading knowledge base from {KNOWLEDGE_BASE_PATH}...")
         try:
-            with open(KNOWLEDGE_BASE_PATH, 'r', encoding='utf-8') as f:
+            with open(kb_path, "r", encoding="utf-8") as f:
                 data = json.load(f)
-                
             vectors = []
             metadata = []
-            
             for item in data:
-                vectors.append(item['embedding'])
+                vectors.append(item["embedding"])
-                metadata.append({
+                metadata.append(
-                    "id": item['id'],
+                    {"id": item["id"], "source": item["source"], "text": item["text"]}
-                    "source": item['source'],
-                    "text": item['text']
-                })
-            
-            self.vector_store.add(vectors, metadata)
-            print(f"Loaded {len(vectors)} chunks into vector store.")
-            
-        except FileNotFoundError:
-            print(f"Error: {KNOWLEDGE_BASE_PATH} not found. Did you run the embedder script?")
-            exit(1)
-        except Exception as e:
-            print(f"Error loading knowledge base: {e}")
-            exit(1)
-
-    def retrieve_context(self, query, top_k=5):
-        print("[1. Retrieving Context...]")
-        
-        try:
-            # --- CRITICAL FIX: Use the EXACT same model as the embedder ---
-            # Embedder used: "qwen/qwen3-embedding-8b" -> Dimension 4096
-            # We must use the same here to avoid shape mismatch.
-            response = self.client.embeddings.create(
-                model="qwen/qwen3-embedding-8b", 
-                input=query
                 )
+            self.vector_store.add(vectors, metadata)
+            print(f"Loaded {len(vectors)} chunks into DBT vector store.")
+        except Exception as e:
+            print(f"Error loading DBT knowledge base: {e}")
+            raise
 
+    async def query(self, query_text):
+        """Query the DBT knowledge base"""
+        try:
+            # Get embedding
+            response = self.client.embeddings.create(
+                model="qwen/qwen3-embedding-8b", input=query_text
+            )
             query_emb = response.data[0].embedding
 
-            # Search the vector store
+            # Search
-            context_chunks = self.vector_store.search(query_emb, top_k=top_k)
+            context_chunks = self.vector_store.search(query_emb, top_k=5)
+            if not context_chunks:
+                return "I couldn't find relevant DBT information for that query."
 
-            return context_chunks
+            # Generate answer
+            context_text = "\n\n---\n\n".join(
+                [chunk["metadata"]["text"] for chunk in context_chunks]
+            )
 
-        except Exception as e:
+            system_prompt = """You are a helpful DBT (Dialectical Behavior Therapy) assistant.
-            print(f"Error retrieving context: {e}")
+Use the provided context from the DBT Skills Training Handouts to answer the user's question.
-            return []
-
-    def generate_answer(self, query, context_chunks):
-        print("[2. Generating Answer...]")
-        
-        # Build the context string
-        context_text = "\n\n---\n\n".join([chunk['metadata']['text'] for chunk in context_chunks])
-        
-        system_prompt = """You are a helpful AI assistant specializing in DBT (Dialectical Behavior Therapy). 
-Use the provided context to answer the user's question. 
 If the answer is not in the context, say you don't know based on the provided text.
-Be concise and compassionate."""
+Be concise, compassionate, and practical."""
 
-        user_prompt = f"""Context:
+            user_prompt = f"Context:\n{context_text}\n\nQuestion: {query_text}"
-{context_text}
 
-Question: {query}"""
-
-        try:
-            # Using a strong model for the final generation
             response = self.client.chat.completions.create(
-                model="openai/gpt-4o-mini", # You can change this to "qwen/qwen-3-8b" or similar if desired
+                model=self.config.get("models", {}).get(
+                    "generator", "openai/gpt-4o-mini"
+                ),
                 messages=[
                     {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt}
+                    {"role": "user", "content": user_prompt},
                 ],
-                temperature=0.7
+                temperature=0.7,
             )
 
             return response.choices[0].message.content
-            
         except Exception as e:
-            return f"Error generating answer: {e}"
+            return f"Error querying DBT knowledge base: {e}"
 
-    def process_query(self, query):
-        # 1. Retrieve
-        context = self.retrieve_context(query)
 
-        if not context:
+# Initialize JurySystem
-            return "I couldn't find any relevant information in the knowledge base."
+jury_system = None
-        
-        # Optional: Print sources for debugging
-        print(f"   Found {len(context)} relevant chunks (Top score: {context[0]['score']:.4f})")
-        
-        # 2. Generate
-        answer = self.generate_answer(query, context)
-        
-        return answer
-
-def main():
-    print("Initializing AI Jury System...")
-    system = JurySystem()
-    
-    print("\nSystem Ready. Ask a question (or type 'exit').")
-    
-    while True:
 try:
-            user_query = input("\nYou: ").strip()
+    jury_system = JurySystem()
-            
+    print("✓ JurySystem (DBT) initialized successfully")
-            if user_query.lower() in ['exit', 'quit']:
-                print("Goodbye!")
-                break
-            
-            if not user_query:
-                continue
-                
-            response = system.process_query(user_query)
-            print(f"\nAI: {response}")
-            
-        except KeyboardInterrupt:
-            print("\nGoodbye!")
-            break
 except Exception as e:
-            print(f"\nAn error occurred: {e}")
+    print(f"⚠ JurySystem initialization failed: {e}")
+
+
+# ==================== Original Bot Functions ====================
+
+
+def decodeJwtPayload(token):
+    payload = token.split(".")[1]
+    payload += "=" * (4 - len(payload) % 4)
+    return json.loads(base64.urlsafe_b64decode(payload))
+
+
+def apiRequest(method, endpoint, token=None, data=None):
+    url = f"{API_URL}{endpoint}"
+    headers = {"Content-Type": "application/json"}
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+    try:
+        resp = getattr(requests, method)(url, headers=headers, json=data, timeout=10)
+        try:
+            return resp.json(), resp.status_code
+        except ValueError:
+            return {}, resp.status_code
+    except requests.RequestException:
+        return {"error": "API unavailable"}, 503
+
+
+def loadCache():
+    try:
+        if os.path.exists(CACHE_FILE):
+            with open(CACHE_FILE, "rb") as f:
+                global user_cache
+                user_cache = pickle.load(f)
+                print(f"Loaded cache for {len(user_cache)} users")
+    except Exception as e:
+        print(f"Error loading cache: {e}")
+
+
+def saveCache():
+    try:
+        with open(CACHE_FILE, "wb") as f:
+            pickle.dump(user_cache, f)
+    except Exception as e:
+        print(f"Error saving cache: {e}")
+
+
+def hashPassword(password):
+    return bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8")
+
+
+def verifyPassword(password, hashed):
+    return bcrypt.checkpw(password.encode("utf-8"), hashed.encode("utf-8"))
+
+
+def getCachedUser(discord_id):
+    return user_cache.get(discord_id)
+
+
+def setCachedUser(discord_id, user_data):
+    user_cache[discord_id] = user_data
+    saveCache()
+
+
+def negotiateToken(discord_id, username, password):
+    cached = getCachedUser(discord_id)
+    if (
+        cached
+        and cached.get("username") == username
+        and verifyPassword(password, cached.get("hashed_password"))
+    ):
+        result, status = apiRequest(
+            "post", "/api/login", data={"username": username, "password": password}
+        )
+        if status == 200 and "token" in result:
+            token = result["token"]
+            payload = decodeJwtPayload(token)
+            user_uuid = payload["sub"]
+            setCachedUser(
+                discord_id,
+                {
+                    "hashed_password": cached["hashed_password"],
+                    "user_uuid": user_uuid,
+                    "username": username,
+                },
+            )
+            return token, user_uuid
+        return None, None
+
+    result, status = apiRequest(
+        "post", "/api/login", data={"username": username, "password": password}
+    )
+    if status == 200 and "token" in result:
+        token = result["token"]
+        payload = decodeJwtPayload(token)
+        user_uuid = payload["sub"]
+        setCachedUser(
+            discord_id,
+            {
+                "hashed_password": hashPassword(password),
+                "user_uuid": user_uuid,
+                "username": username,
+            },
+        )
+        return token, user_uuid
+    return None, None
+
+
+async def handleAuthFailure(message):
+    discord_id = message.author.id
+    user_sessions.pop(discord_id, None)
+    await message.channel.send(
+        "Your session has expired. Send any message to log in again."
+    )
+
+
+async def handleLoginStep(message):
+    discord_id = message.author.id
+    state = login_state[discord_id]
+
+    if state["step"] == "username":
+        state["username"] = message.content.strip()
+        state["step"] = "password"
+        await message.channel.send("Password?")
+
+    elif state["step"] == "password":
+        username = state["username"]
+        password = message.content.strip()
+        del login_state[discord_id]
+
+        token, user_uuid = negotiateToken(discord_id, username, password)
+
+        if token and user_uuid:
+            user_sessions[discord_id] = {
+                "token": token,
+                "user_uuid": user_uuid,
+                "username": username,
+            }
+            registered = ", ".join(list_registered()) or "none"
+            await message.channel.send(
+                f"Welcome back **{username}**!\n\n"
+                f"Registered modules: {registered}\n\n"
+                f"Send 'help' for available commands."
+            )
+        else:
+            await message.channel.send(
+                "Invalid credentials. Send any message to try again."
+            )
+
+
+async def sendHelpMessage(message):
+    help_msg = """**🤖 Synculous Bot - Natural Language Commands**
+
+Just talk to me naturally! Here are some examples:
+
+**💊 Medications:**
+• "add lsd 50 mcg every tuesday at 4:20pm"
+• "take my wellbutrin"
+• "what meds do i have today?"
+• "show my refills"
+• "snooze my reminder for 30 minutes"
+• "check adherence"
+
+**📋 Routines:**
+• "create morning routine with brush teeth, shower, eat"
+• "start my morning routine"
+• "done" (complete current step)
+• "skip" (skip current step)
+• "pause/resume" (pause or continue)
+• "what steps are in my routine?"
+• "schedule workout for monday wednesday friday at 7am"
+• "show my stats"
+
+**🧠 DBT Support:**
+• "how do I use distress tolerance?"
+• "explain radical acceptance"
+• "give me a DBT skill for anger"
+• "what are the TIPP skills?"
+
+**💡 Tips:**
+• I understand natural language, typos, and slang
+• If I'm unsure, I'll ask for clarification
+• For important actions, I'll ask you to confirm with "yes" or "no"
+• When you're in a routine, shortcuts like "done", "skip", "pause" work automatically"""
+    await message.channel.send(help_msg)
+
+
+async def checkActiveSession(session):
+    """Check if user has an active routine session and return details."""
+    token = session.get("token")
+    if not token:
+        return None
+
+    resp, status = apiRequest("get", "/api/sessions/active", token)
+    if status == 200 and "session" in resp:
+        return resp
+    return None
+
+
+async def handleConfirmation(message, session):
+    """Handle yes/no confirmation responses. Returns True if handled."""
+    discord_id = message.author.id
+    user_input = message.content.lower().strip()
+
+    if "pending_confirmations" not in session:
+        return False
+
+    pending = session["pending_confirmations"]
+    if not pending:
+        return False
+
+    confirmation_id = list(pending.keys())[-1]
+    confirmation_data = pending[confirmation_id]
+
+    if user_input in ("yes", "y", "yeah", "sure", "ok", "confirm"):
+        del pending[confirmation_id]
+        interaction_type = confirmation_data.get("interaction_type")
+        handler = get_handler(interaction_type)
+
+        if handler:
+            fake_parsed = confirmation_data.copy()
+            fake_parsed["needs_confirmation"] = False
+            await handler(message, session, fake_parsed)
+        return True
+
+    elif user_input in ("no", "n", "nah", "cancel", "abort"):
+        del pending[confirmation_id]
+        await message.channel.send("❌ Cancelled.")
+        return True
+
+    return False
+
+
+async def handleActiveSessionShortcuts(message, session, active_session):
+    """Handle shortcuts like 'done', 'skip', 'next' when in active session."""
+    user_input = message.content.lower().strip()
+
+    shortcuts = {
+        "done": ("routine", "complete"),
+        "finished": ("routine", "complete"),
+        "complete": ("routine", "complete"),
+        "next": ("routine", "complete"),
+        "skip": ("routine", "skip"),
+        "pass": ("routine", "skip"),
+        "pause": ("routine", "pause"),
+        "hold": ("routine", "pause"),
+        "resume": ("routine", "resume"),
+        "continue": ("routine", "resume"),
+        "stop": ("routine", "cancel"),
+        "quit": ("routine", "cancel"),
+        "abort": ("routine", "abort"),
+    }
+
+    if user_input in shortcuts:
+        interaction_type, action = shortcuts[user_input]
+        handler = get_handler(interaction_type)
+        if handler:
+            fake_parsed = {"action": action}
+            await handler(message, session, fake_parsed)
+            return True
+
+    return False
+
+
+async def handleDBTQuery(message):
+    """Handle DBT-related queries using JurySystem"""
+    if not jury_system:
+        return False
+
+    # Keywords that indicate a DBT query
+    dbt_keywords = [
+        "dbt",
+        "distress tolerance",
+        "emotion regulation",
+        "interpersonal effectiveness",
+        "mindfulness",
+        "radical acceptance",
+        "wise mind",
+        "tipp",
+        "dearman",
+        "check the facts",
+        "opposite action",
+        "cope ahead",
+        "abc please",
+        "stop skill",
+        "pros and cons",
+        "half smile",
+        "willing hands",
+    ]
+
+    user_input_lower = message.content.lower()
+    is_dbt_query = any(keyword in user_input_lower for keyword in dbt_keywords)
+
+    if is_dbt_query:
+        async with message.channel.typing():
+            response = await jury_system.query(message.content)
+        await message.channel.send(f"🧠 **DBT Support:**\n{response}")
+        return True
+
+    return False
+
+
+async def routeCommand(message):
+    discord_id = message.author.id
+    session = user_sessions[discord_id]
+    user_input = message.content.lower()
+
+    if "help" in user_input or "what can i say" in user_input:
+        await sendHelpMessage(message)
+        return
+
+    # Check for active session first
+    active_session = await checkActiveSession(session)
+
+    # Handle confirmation responses
+    confirmation_handled = await handleConfirmation(message, session)
+    if confirmation_handled:
+        return
+
+    # Handle shortcuts when in active session
+    if active_session:
+        shortcut_handled = await handleActiveSessionShortcuts(
+            message, session, active_session
+        )
+        if shortcut_handled:
+            return
+
+    # Check for DBT queries
+    dbt_handled = await handleDBTQuery(message)
+    if dbt_handled:
+        return
+
+    async with message.channel.typing():
+        history = message_history.get(discord_id, [])
+
+        # Add context about active session to help AI understand
+        context = ""
+        if active_session:
+            session_data = active_session.get("session", {})
+            routine_name = session_data.get("routine_name", "a routine")
+            current_step = session_data.get("current_step_index", 0) + 1
+            total_steps = active_session.get("total_steps", 0)
+            context = f"\n[Context: User is currently in active session for '{routine_name}', on step {current_step} of {total_steps}. They can say 'done', 'skip', 'pause', 'resume', or 'stop'.]"
+
+        parsed = await ai_parser.parse(
+            message.content + context, "command_parser", history=history
+        )
+
+        if discord_id not in message_history:
+            message_history[discord_id] = []
+        message_history[discord_id].append((message.content, parsed))
+        message_history[discord_id] = message_history[discord_id][-5:]
+
+    if "needs_clarification" in parsed:
+        await message.channel.send(
+            f"I'm not quite sure what you mean. {parsed['needs_clarification']}"
+        )
+        return
+
+    if "error" in parsed:
+        await message.channel.send(
+            f"I had trouble understanding that: {parsed['error']}"
+        )
+        return
+
+    interaction_type = parsed.get("interaction_type")
+    handler = get_handler(interaction_type)
+
+    if handler:
+        await handler(message, session, parsed)
+    else:
+        registered = ", ".join(list_registered()) or "none"
+        await message.channel.send(
+            f"Unknown command type '{interaction_type}'. Registered modules: {registered}"
+        )
+
+
+@client.event
+async def on_ready():
+    print(f"Bot logged in as {client.user}")
+    loadCache()
+    backgroundLoop.start()
+
+
+@client.event
+async def on_message(message):
+    if message.author == client.user:
+        return
+    if not isinstance(message.channel, discord.DMChannel):
+        return
+
+    discord_id = message.author.id
+
+    if discord_id in login_state:
+        await handleLoginStep(message)
+        return
+
+    if discord_id not in user_sessions:
+        login_state[discord_id] = {"step": "username"}
+        await message.channel.send("Welcome! Send your username to log in.")
+        return
+
+    await routeCommand(message)
+
+
+@tasks.loop(seconds=60)
+async def backgroundLoop():
+    """Override this in your domain module or extend as needed."""
+    pass
+
+
+@backgroundLoop.before_loop
+async def beforeBackgroundLoop():
+    await client.wait_until_ready()
+
 
 if __name__ == "__main__":
-    main()
+    client.run(DISCORD_BOT_TOKEN)

config.json

@@ -1,6 +1,6 @@
 {
   "openrouter_api_key": "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",
-  "embedding_file": "dbt_knowledge.json",
+  "embedding_file": "bot/data/dbt_knowledge.embeddings.json",
   "models": {
     "generator": "moonshotai/kimi-k2.5",
     "jury_clinical": "z-ai/glm-5",