Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,90 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from langchain_groq import ChatGroq
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+from langchain.schema import SystemMessage
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# 1. Setup Page
+st.set_page_config(page_title="Serenity AI", page_icon="🌿")
+st.title("🌿 Serenity: Your CBT Companion")
+
+# 2. Sidebar - Privacy & Resources
+with st.sidebar:
+    st.header("About")
+    st.info("This is a support tool, NOT a doctor. Data is processed locally for privacy.")
+    groq_api_key = st.text_input("Groq API Key", type="password")
+
+# 3. Load & Process Knowledge Base (Only runs once)
+@st.cache_resource
+def setup_rag():
+    # Load your manual text file
+    with open("cbt_resources.txt", "r") as f:
+        raw_text = f.read()
+    
+    # Split into chunks
+    text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    texts = text_splitter.split_text(raw_text)
+    
+    # Create Embeddings (Local & Free)
+    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+    
+    # Store in Vector DB
+    db = Chroma.from_texts(texts, embeddings)
+    return db.as_retriever()
+
+# 4. Initialize Chat Logic
+if groq_api_key:
+    # Setup LLM
+    llm = ChatGroq(
+        temperature=0.6, 
+        model_name="llama3-70b-8192", 
+        groq_api_key=groq_api_key
+    )
+    
+    # Setup Memory
+    if "memory" not in st.session_state:
+        st.session_state.memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True
+        )
+
+    # Setup RAG Chain
+    retriever = setup_rag()
+    chain = ConversationalRetrievalChain.from_llm(
+        llm=llm,
+        retriever=retriever,
+        memory=st.session_state.memory,
+        verbose=False
+    )
+    
+    # Inject System Prompt logic (Simplified for Streamlit)
+    # Ideally, you wrap the LLM with the system prompt in the chain creation.
+
+    # 5. Chat Interface
+    if "messages" not in st.session_state:
+        st.session_state.messages = [{"role": "assistant", "content": "Hello. I'm Serenity. How are you feeling today?"}]
+
+    for msg in st.session_state.messages:
+        st.chat_message(msg["role"]).write(msg["content"])
+
+    if prompt := st.chat_input():
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.chat_message("user").write(prompt)
+
+        # Safety Check (Simple Keyword Filter)
+        dangerous_keywords = ["suicide", "kill myself", "end it all"]
+        if any(word in prompt.lower() for word in dangerous_keywords):
+            response = "I'm really concerned about you, but I am an AI. Please call your local emergency number immediately. You are not alone."
+        else:
+            # Generate Response using RAG
+            response_dict = chain.invoke({"question": prompt})
+            response = response_dict['answer']
+
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        st.chat_message("assistant").write(response)
+
+else:
+    st.warning("Please enter your Groq API Key to start.")