Update app.py

app.py

@@ -116,33 +116,25 @@ def create_vector_store(_documents, _embedding_model_name: str):
         return None
 
 @st.cache_resource(show_spinner="Initializing LLM...")
-def get_llm(api_key: str, model_name: str = "mixtral-8x7b-32768"): # "llama3-70b-8192" is another option
+def get_llm(api_key: str, model_name: str = "llama3-8b-8192"): # UPDATED MODEL
     """Initializes the Groq LLM."""
     if not api_key:
         st.error("GROQ_API_KEY not found! Please set it in your environment variables or a .env file.")
         return None
     try:
+        # Available models (check Groq documentation for the latest):
+        # "llama3-8b-8192" (good balance of speed and capability)
+        # "llama3-70b-8192" (more powerful, potentially slower)
+        # "gemma-7b-it"
         llm = ChatGroq(temperature=0, groq_api_key=api_key, model_name=model_name)
+        st.sidebar.info(f"LLM Initialized: {model_name}") # Add info about which model is used
         return llm
     except Exception as e:
         st.error(f"Error initializing Groq LLM: {e}")
         return None
 
 # --- RAG Chain Setup ---
-def get_rag_chain(llm, retriever, prompt_template_str):
-    """Creates a RAG chain with the given LLM, retriever, and prompt template."""
-    prompt = PromptTemplate(
-        template=prompt_template_str,
-        input_variables=["context", "question"]
-    )
-    
-    rag_chain = (
-        {"context": retriever, "question": RunnablePassthrough()}
-        | prompt
-        | llm
-        | StrOutputParser()
-    )
-    return rag_chain
+# ... (get_rag_chain function remains the same) ...
 
 # --- Main Application Logic ---
 def main():
@@ -152,86 +144,35 @@ def main():
     # --- UI Setup ---
     st.set_page_config(page_title="Internal Knowledge Base AI", layout="wide", initial_sidebar_state="expanded")
 
-    # Custom CSS for a "catchy and elegant" design
+    # Custom CSS (remains the same)
     st.markdown("""
     <style>
-        /* General body style */
-        body {
-            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-            background-color: #f0f2f6; /* Light gray background */
-        }
-        /* Main content area */
-        .main .block-container {
-            padding-top: 2rem;
-            padding-bottom: 2rem;
-            padding-left: 3rem;
-            padding-right: 3rem;
-            background-color: #ffffff; /* White content background */
-            border-radius: 10px;
-            box-shadow: 0 4px 12px rgba(0,0,0,0.1); /* Subtle shadow */
-        }
-        /* Title style */
-        h1 {
-            color: #1E88E5; /* Catchy blue */
-            text-align: center;
-            font-weight: 600;
-        }
-        /* Sidebar style */
-        .stSidebar {
-            background-color: #E3F2FD; /* Light blue sidebar */
-            padding: 10px;
-        }
-        .stSidebar .sidebar-content {
-             background-color: #E3F2FD;
-        }
-        /* Input box style */
-        .stTextInput > div > div > input {
-            background-color: #f8f9fa;
-            border-radius: 5px;
-            border: 1px solid #ced4da;
-        }
-        /* Button style */
-        .stButton > button {
-            background-color: #1E88E5; /* Catchy blue */
-            color: white;
-            border-radius: 5px;
-            padding: 0.5rem 1rem;
-            font-weight: 500;
-            border: none;
-            transition: background-color 0.3s ease;
-        }
-        .stButton > button:hover {
-            background-color: #1565C0; /* Darker blue on hover */
-        }
-        /* Status messages */
-        .stAlert { /* For st.info, st.success, st.warning, st.error */
-            border-radius: 5px;
-        }
-        /* Response area */
-        .response-area {
-            background-color: #f8f9fa;
-            padding: 1rem;
-            border-radius: 5px;
-            border: 1px solid #e0e0e0;
-            margin-top: 1rem;
-            min-height: 100px;
-        }
+        # ... (CSS content remains the same) ...
     </style>
     """, unsafe_allow_html=True)
 
     st.title("📚 Internal Knowledge Base AI 💡")
     
-    # Sidebar for status and information
+    st.sidebar.header("System Settings") # Changed from System Status for clarity
+    
+    # Model selection in sidebar (New Feature)
+    available_models = ["llama3-8b-8192", "llama3-70b-8192", "gemma-7b-it"] # Add more as Groq supports them
+    selected_model = st.sidebar.selectbox(
+        "Select LLM Model:",
+        available_models,
+        index=available_models.index("llama3-8b-8192") # Default selection
+    )
+    st.sidebar.markdown("---")
     st.sidebar.header("System Status")
     status_placeholder = st.sidebar.empty()
     status_placeholder.info("Initializing...")
 
+
     if not groq_api_key:
         status_placeholder.error("GROQ API Key not configured. Application cannot start.")
         st.stop()
 
     # --- Knowledge Base Loading ---
-    # This will be cached after the first run
     with st.spinner("Knowledge Base is loading... Please wait."):
         start_time = time.time()
         processed_documents = load_and_process_documents(DOCS_DIR)
@@ -244,17 +185,22 @@ def main():
             status_placeholder.error("Failed to create vector store. Application cannot proceed.")
             st.stop()
         
-        llm = get_llm(groq_api_key)
+        # Pass the selected model to get_llm
+        llm = get_llm(groq_api_key, model_name=selected_model) 
         if not llm:
+            # Error is already shown by get_llm, but update status_placeholder too
             status_placeholder.error("Failed to initialize LLM. Application cannot proceed.")
             st.stop()
             
         end_time = time.time()
+        # status_placeholder is updated by get_llm or on success below
         status_placeholder.success(f"Application Ready! (Loaded in {end_time - start_time:.2f}s)")
     
-    retriever = vector_store.as_retriever(search_kwargs={"k": 5}) # Retrieve top 5 relevant chunks
+    retriever = vector_store.as_retriever(search_kwargs={"k": 5})
 
     # --- Query Input and Response ---
+    # ... (rest of the main function remains the same, including prompt templates, query input, button, and response display logic) ...
+
     st.markdown("---")
     st.subheader("Ask a question about our documents:")
 
@@ -293,7 +239,6 @@ def main():
     Answer:
     """
 
-    # Use session state to store conversation history if desired, or just last query/response
     if "messages" not in st.session_state:
         st.session_state.messages = []
     
@@ -303,14 +248,14 @@ def main():
         if query:
             st.session_state.messages.append({"role": "user", "content": query})
             
-            # Determine prompt based on query type (simple keyword check)
-            # A more sophisticated intent detection could be used here (e.g., another LLM call, classifier)
-            if "order" in query.lower() and ("status" in query.lower() or "track" in query.lower() or "update" in query.lower() or any(name_part.lower() in query.lower() for name_part in ["customer", "client", "name"])): # Basic check for order status
+            current_model_info = st.sidebar.empty() # Placeholder for current mode info
+
+            if "order" in query.lower() and ("status" in query.lower() or "track" in query.lower() or "update" in query.lower() or any(name_part.lower() in query.lower() for name_part in ["customer", "client", "name"])):
                 active_prompt_template = ORDER_STATUS_PROMPT
-                st.sidebar.info("Using: Order Status Query Mode")
+                current_model_info.info("Mode: Order Status Query")
             else:
                 active_prompt_template = GENERAL_QA_PROMPT
-                st.sidebar.info("Using: General Query Mode")
+                current_model_info.info("Mode: General Query")
 
             rag_chain = get_rag_chain(llm, retriever, active_prompt_template)
 
@@ -325,22 +270,16 @@ def main():
         else:
             st.warning("Please enter a question.")
 
-    # Display chat messages
     st.markdown("---")
     st.subheader("Response:")
     response_area = st.container()
-    response_area.add_rows([ # Create a container with fixed height and scroll
-            st.markdown(f"<div class='response-area'>{st.session_state.messages[-1]['content'] if st.session_state.messages and st.session_state.messages[-1]['role'] == 'assistant' else 'Ask a question to see the answer here.'}</div>", unsafe_allow_html=True)
-    ])
-
-
-    # Optional: Display retrieved context for debugging or transparency
-    # if st.sidebar.checkbox("Show Retrieved Context (for debugging)"):
-    #     if query and vector_store: # Check if query and vector_store exist
-    #         docs = retriever.get_relevant_documents(query)
-    #         st.sidebar.subheader("Retrieved Context:")
-    #         for i, doc in enumerate(docs):
-    #             st.sidebar.text_area(f"Chunk {i+1} (Source: {doc.metadata.get('source', 'N/A')})", doc.page_content, height=150)
+    # Ensure response_area is robust against empty messages or incorrect last role
+    last_assistant_message = "Ask a question to see the answer here."
+    if st.session_state.messages and st.session_state.messages[-1]['role'] == 'assistant':
+        last_assistant_message = st.session_state.messages[-1]['content']
+    
+    response_area.markdown(f"<div class='response-area'>{last_assistant_message}</div>", unsafe_allow_html=True)
+
 
     st.sidebar.markdown("---")
     st.sidebar.markdown("Built with ❤️ using Streamlit & Langchain & Groq")