Upload app.py with huggingface_hub

app.py

@@ -1,210 +1,320 @@
 
-# Import necessary libraries
-import os  # Interacting with the operating system (reading/writing files)
-import chromadb  # High-performance vector database for storing/querying dense vectors
-from dotenv import load_dotenv  # Loading environment variables from a .env file
-import json  # Parsing and handling JSON data
-
-# LangChain imports
-from langchain_core.documents import Document  # Document data structures
-from langchain_core.runnables import RunnablePassthrough  # LangChain core library for running pipelines
-from langchain_core.output_parsers import StrOutputParser  # String output parser
-from langchain.prompts import ChatPromptTemplate  # Template for chat prompts
-from langchain.chains.query_constructor.base import AttributeInfo  # Base classes for query construction
-from langchain.retrievers.self_query.base import SelfQueryRetriever  # Base classes for self-querying retrievers
-from langchain.retrievers.document_compressors import LLMChainExtractor, CrossEncoderReranker  # Document compressors
-from langchain.retrievers import ContextualCompressionRetriever  # Contextual compression retrievers
-
-# LangChain community & experimental imports
-from langchain_community.vectorstores import Chroma  # Implementations of vector stores like Chroma
-from langchain_community.document_loaders import PyPDFDirectoryLoader, PyPDFLoader  # Document loaders for PDFs
-from langchain_community.cross_encoders import HuggingFaceCrossEncoder  # Cross-encoders from HuggingFace
-from langchain_experimental.text_splitter import SemanticChunker  # Experimental text splitting methods
-from langchain.text_splitter import (
-    CharacterTextSplitter,  # Splitting text by characters
-    RecursiveCharacterTextSplitter  # Recursive splitting of text by characters
-)
-from langchain_core.tools import tool
-from langchain.agents import create_tool_calling_agent, AgentExecutor
+# --- 0. Library Imports ---
+import os
+import json
+import shutil
+import time
+import numpy as np
+from datetime import datetime
+from typing import Dict, List, Any, TypedDict, Tuple
+
+# LangChain Core & Community
+from langchain_core.documents import Document
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.tools import tool
+from langchain_community.vectorstores import Chroma
+from langchain_community.document_loaders import PyPDFDirectoryLoader
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter
 
-# LangChain OpenAI imports
-# Commented cout - not used
-#from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI  # OpenAI embeddings and models
-#from langchain.embeddings.openai import OpenAIEmbeddings  # OpenAI embeddings for text vectors
-# Added and used below
+# LangChain OpenAI
 from langchain_openai import OpenAIEmbeddings, ChatOpenAI
-#from langchain_openai import ChatOpenAI
 
-# LlamaParse & LlamaIndex imports
-from llama_parse import LlamaParse  # Document parsing library
-from llama_index.core import Settings, SimpleDirectoryReader  # Core functionalities of the LlamaIndex
+# LangChain Experimental
+from langchain_experimental.text_splitter import SemanticChunker
 
-# LangGraph import
-from langgraph.graph import StateGraph, END, START  # State graph for managing states in LangChain
+# LangChain Agents & Graph
+from langchain.agents import create_tool_calling_agent, AgentExecutor
+from langgraph.graph import StateGraph, END, START
 
-# Pydantic import
-from pydantic import BaseModel  # Pydantic for data validation
+# External Libraries
+import chromadb
+from llama_parse import LlamaParse # For PDF parsing
+from groq import Groq # For Llama Guard
+from mem0 import MemoryClient # For memory
+import streamlit as st # For Web UI
 
-# Typing imports
-from typing import Dict, List, Tuple, Any, TypedDict  # Python typing for function annotations
+# Fix for numpy depreciation warning if necessary
+np.float_ = np.float64
 
-# Other utilities
-import numpy as np  # Numpy for numerical operations
-from groq import Groq
-from mem0 import MemoryClient
-import streamlit as st
-from datetime import datetime
+#import nest_asyncio
 
-#====================================SETUP=====================================#
-# Fetch secrets from Hugging Face Spaces
-api_key = os.environ["API_KEY"]
-endpoint = os.environ["OPENAI_API_BASE"]
-groq_api_key = os.environ['LLAMA_API_KEY'] # llama_api_key = os.environ['GROQ_API_KEY']
-MEM0_api_key = os.environ['MEM0_API_KEY'] # MEM0_api_key = os.environ['mem0']
-#my_api_key = os.environ["MY_API_KEY"]
-
-# Initialize the OpenAI embedding function for Chroma
-embedding_function = chromadb.utils.embedding_functions.OpenAIEmbeddingFunction(
-    api_base=endpoint, # Complete the code to define the API base endpoint
-    api_key=api_key # Complete the code to define the API key
-)
-#model_name='text-embedding-ada-002' # This is a fixed value and does not need modification
-
-# This initializes the OpenAI embedding function for the Chroma vectorstore, using the provided endpoint and API key.
-
-# Initialize the OpenAI Embeddings
-embedding_model = OpenAIEmbeddings(
-    openai_api_base=endpoint,
-    openai_api_key=api_key
-)
-#model='text-embedding-ada-002'
-
-# Initialize the Chat OpenAI model
-llm = ChatOpenAI(
-    openai_api_base=endpoint,
-    openai_api_key=api_key,
-    model="gpt-4o-mini", # used gpt4o instead of gpt-4o-mini to get improved results
-    streaming=False
-)
-
-# This initializes the Chat OpenAI model with the provided endpoint, API key, deployment name, and a temperature setting of 0 (to control response variability).
-
-# set the LLM and embedding model in the LlamaIndex settings.
-Settings.llm = llm # Complete the code to define the LLM model
-Settings.embedding = embedding_model # Complete the code to define the embedding model
-
-#================================Creating Langgraph agent======================#
+# --- 1. Configuration and Setup Utilities ---
+
+from typing import Dict, List, Any, TypedDict, Tuple
+
+def load_config_from_env() -> Dict:
+    """Loads API keys and endpoints from environment variables."""
+    # Prioritize environment variables for deployment
+    config = {
+        "AZURE_OPENAI_API_KEY": os.getenv("AZURE_OPENAI_API_KEY"),
+        "AZURE_OPENAI_API_BASE": os.getenv("AZURE_OPENAI_API_BASE"),
+        "LLAMA_KEY": os.getenv("LLAMA_KEY"), # LlamaParse API Key
+        "MEM0_API_KEY": os.getenv("MEM0_API_KEY"),
+        "GROQ_API_KEY": os.getenv("GROQ_API_KEY"), # For Llama Guard via Groq
+    }
+    # Basic validation
+    for key, value in config.items():
+        if not value:
+            st.warning(f"Warning: Environment variable '{key}' is not set.")
+    return config
+
+def initialize_llms_and_embeddings(config: Dict) -> Tuple[OpenAIEmbeddings, ChatOpenAI, chromadb.utils.embedding_functions.OpenAIEmbeddingFunction, Groq]:
+    """Initializes LLM, Embedding models, and API clients."""
+    api_key = config["AZURE_OPENAI_API_KEY"]
+    endpoint = config["AZURE_OPENAI_API_BASE"]
+    groq_api_key = config["GROQ_API_KEY"]
+
+    # Initialize ChromaDB embedding function (used for collection creation)
+    embedding_function = chromadb.utils.embedding_functions.OpenAIEmbeddingFunction(
+        api_base=endpoint,
+        api_key=api_key,
+        model_name='text-embedding-ada-002' # Specify model explicitly
+    )
+
+    # Initialize LangChain OpenAI Embeddings (used for `SemanticChunker` and `Chroma` vectorstore)
+    embedding_model = OpenAIEmbeddings(
+        openai_api_base=endpoint,
+        openai_api_key=api_key,
+        model='text-embedding-ada-002' # Specify model explicitly
+    )
+
+    # Initialize LangChain Chat OpenAI model
+    llm = ChatOpenAI(
+        openai_api_base=endpoint,
+        openai_api_key=api_key,
+        model="gpt-4o-mini",
+        streaming=False,
+        temperature=0.0 # Ensure deterministic behavior for evaluations
+    )
+
+    # Initialize Groq client for Llama Guard
+    llama_guard_client = Groq(api_key=groq_api_key)
+
+    return embedding_model, llm, embedding_function, llama_guard_client
+
+def filter_input_with_llama_guard(user_input: str, llama_guard_client: Groq, model: str = "meta-llama/llama-guard-4-12b") -> str:
+    """
+    Filters user input using Llama Guard to ensure it is safe.
+    Returns "safe", "UNSAFE" (with category), or None on error.
+    """
+    try:
+        response = llama_guard_client.chat.completions.create(
+            messages=[{"role": "user", "content": user_input}],
+            model=model,
+        )
+        return response.choices[0].message.content.strip()
+    except Exception as e:
+        st.error(f"Error with Llama Guard: {e}")
+        return None
+
+# --- 2. Data Preparation (Parsing & Chunking) ---
+
+# Note: In a deployed app, PDF parsing and vector DB creation would typically be
+# a separate, offline process, and the pre-built vector DB would be loaded.
+# For this template, we'll assume the nutritional_db is pre-existing and loaded.
+
+def load_and_split_documents(folder_path: str, embedding_model: OpenAIEmbeddings) -> List[Document]:
+    """Loads PDF documents from a folder and semantically chunks them."""
+    semantic_text_splitter = SemanticChunker(
+        embedding_model,
+        breakpoint_threshold_type='percentile',
+        breakpoint_threshold_amount=80
+    )
+    pdf_loader = PyPDFDirectoryLoader(folder_path)
+    chunks = pdf_loader.load_and_split(semantic_text_splitter)
+    return chunks
+
+def parse_pdf_tables_with_llamaparse(pdf_path: str, llamaparse_api_key: str) -> Tuple[Dict, Dict]:
+    """Parses a PDF file using LlamaParse and extracts page texts and tables."""
+    # This requires `nest_asyncio.apply()` to be called once at the start of the app if running async.
+    # In a Streamlit app, ensure it's at the very top level if needed.
+    # import nest_asyncio; nest_asyncio.apply() # Uncomment if needed for async parser
+
+    parser = LlamaParse(
+        result_type="markdown",
+        skip_diagonal_text=True,
+        fast_mode=False,
+        num_workers=1, # Adjust as per environment capabilities
+        check_interval=10,
+        api_key=llamaparse_api_key
+    )
+    json_objs = parser.get_json_result(pdf_path)
+    page_texts, tables = {}, {}
+    for obj in json_objs:
+        json_list = obj['pages']
+        name = obj["file_path"].split("/")[-1]
+        page_texts[name] = {}
+        tables[name] = {}
+        for json_item in json_list:
+            for component in json_item['items']:
+                if component['type'] == 'table':
+                    tables[name][json_item['page']] = component['rows']
+    return page_texts, tables
+
+def generate_hypothetical_questions(llm: ChatOpenAI, docs: List[Document], is_table: bool = False) -> List[Document]:
+    """Generates hypothetical questions for text chunks or tables."""
+    prompt_template = """
+    Generate a list of exactly three (3) hypothetical questions that the below nutritional disorder {content_type} could be used to answer:
+    {content}
+    Ensure that the questions are specific in the context of nutrition, dietary deficiencies, metabolic disorders, vitamin and mineral imbalances, obesity, and related health conditions.
+    Generate only a list of questions.
+    Do not mention anything before or after the list.
+    If the content cannot answer any questions, return an empty list.
+    """
+    hyp_docs = []
+    content_type = "table" if is_table else "document"
+
+    for i, doc in enumerate(docs):
+        content_to_use = str(doc) if is_table else doc.page_content # Tables are often raw data, stringify
+        try:
+            response = llm.invoke(prompt_template.format(content_type=content_type, content=content_to_use))
+            questions = response.content
+        except Exception as e:
+            st.error(f"Error generating hypothetical questions for {'table' if is_table else 'text'} chunk ID {doc.id}: {e}")
+            questions = "[]" # Return empty list string on error
+
+        questions_metadata = {
+            'original_content': content_to_use,
+            'source': doc.metadata.get('source', 'unknown'),
+            'page': doc.metadata.get('page', -1),
+            'type': content_type
+        }
+        hyp_docs.append(
+            Document(
+                id=f"{'table_' if is_table else 'text_chunk_'}{doc.id or i}", # Ensure unique IDs
+                page_content=questions,
+                metadata=questions_metadata
+            )
+        )
+        time.sleep(0.1) # Small delay to avoid rate limits
+    return hyp_docs
+
+
+def create_and_persist_vector_db(
+    documents: List[Document],
+    embedding_model: OpenAIEmbeddings,
+    collection_name: str,
+    persist_directory: str
+):
+    """Creates/updates a Chroma vector store and persists it."""
+    # Ensure IDs are strings as required by ChromaDB
+    doc_ids = [str(d.id) for d in documents] if documents else []
+    if not doc_ids:
+        st.warning(f"No documents to add to collection '{collection_name}'.")
+        return
+
+    # Initialize or connect to Chroma vectorstore
+    vector_store = Chroma.from_documents(
+        documents,
+        embedding_model,
+        collection_name=collection_name,
+        persist_directory=persist_directory
+    )
+    st.info(f"Initialized ChromaDB with collection '{collection_name}' at '{persist_directory}'. "
+            f"Documents count: {len(documents)}")
+    return vector_store
+
+def load_vector_db(
+    embedding_model: OpenAIEmbeddings,
+    collection_name: str,
+    persist_directory: str
+) -> Chroma:
+    """Loads an existing Chroma vector store."""
+    try:
+        # Check if the directory exists and contains ChromaDB files
+        if not os.path.exists(persist_directory) or not os.listdir(persist_directory):
+            st.error(f"Vector DB directory '{persist_directory}' is empty or does not exist.") or print(f"Vector DB directory '{persist_directory}' is empty or does not exist.")
+            st.warning("Please ensure the 'nutritional_db' folder is correctly placed/mounted.") or print("Please ensure the 'nutritional_db' folder is correctly placed/mounted.")
+            return None
+
+        vector_store = Chroma(
+            collection_name=collection_name,
+            persist_directory=persist_directory,
+            embedding_function=embedding_model
+        )
+        st.success(f"Successfully loaded ChromaDB collection '{collection_name}' from '{persist_directory}'.") or print(f"Successfully loaded ChromaDB collection '{collection_name}' from '{persist_directory}'.")
+        # You can add a check for the number of documents loaded for verification
+        # Example: print(vector_store._collection.count())
+        return vector_store
+    except Exception as e:
+        st.error(f"Error loading ChromaDB from '{persist_directory}': {e}") or print(f"Error loading ChromaDB from '{persist_directory}': {e}")
+        return None
+
+# --- 3. Agent Workflow Definition (LangGraph) ---
 
 class AgentState(TypedDict):
-    query: str  # The current user query
-    expanded_query: str  # The expanded version of the user query
-    context: List[Dict[str, Any]]  # Retrieved documents (content and metadata)
-    response: str  # The generated response to the user query
-    precision_score: float  # The precision score of the response
-    groundedness_score: float  # The groundedness score of the response
-    groundedness_loop_count: int  # Counter for groundedness refinement loops
-    precision_loop_count: int  # Counter for precision refinement loops
+    """Represents the state of the AI agent at different stages of the workflow."""
+    query: str
+    expanded_query: str
+    context: List[Dict[str, Any]]
+    response: str
+    precision_score: float
+    groundedness_score: float
+    groundedness_loop_count: int
+    precision_loop_count: int
     feedback: str
     query_feedback: str
-    groundedness_check: bool
+    groundedness_check: bool # This field isn't used in should_continue_groundedness, can be removed
     loop_max_iter: int
 
-def expand_query(state):
-    """
-    Expands the user query to improve retrieval of nutrition disorder-related information.
-
-    Args:
-        state (Dict): The current state of the workflow, containing the user query.
 
-    Returns:
-        Dict: The updated state with the expanded query.
-    """
-    print("---------Expanding Query---------")
+# Node functions for LangGraph
+def expand_query(state: AgentState) -> AgentState:
+    st.write("---Expanding Query---")
     system_message = '''
     You are a domain expert assisting in answering questions related to research papers.
     Convert the user query into something that a nutritionist would understand. Use domain related words.
-    Return 3 related search queries based on the user's request seperated by newline.
-    Return only 3 versions of the question as a list.
+    Return three (3) related search queries based on the user's request separated by newline.
+    Return only three (3) versions of the question as a list.
     Perform query expansion on the question received. If there are multiple common ways of phrasing a user question \
     or common synonyms for key words in the question, make sure to return multiple versions \
     of the query with the different phrasings.
-    If the query has multiple parts, split them into separate simpler queries. This is the only case where you can generate more than 3 queries.
+    If the query has multiple parts, split them into separate simpler queries. This is the only case where you can generate more than three (3) queries.
     If there are acronyms or words you are not familiar with, do not try to rephrase them.
     Generate only a list of questions. Do not mention anything before or after the list.
-    Use the query feeback if provided to craft the search queries.
+    Use the query feedback if provided to craft the search queries.
     '''
-
     expand_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
         ("user", "Expand this query: {query} using the feedback: {query_feedback}")
-
     ])
-
-    chain = expand_prompt | llm | StrOutputParser()
+    chain = expand_prompt | st.session_state.llm | StrOutputParser() # Use llm from session state
     expanded_query = chain.invoke({"query": state['query'], "query_feedback":state["query_feedback"]})
-    print("expanded_query", expanded_query)
+    st.write(f"Expanded query:\n{expanded_query}")
     state["expanded_query"] = expanded_query
     return state
 
+def retrieve_context(state: AgentState) -> AgentState:
+    st.write("---Retrieving Context---")
+    query = f"{state['query']}; {state['expanded_query']}"
+    st.write(f"Query used for retrieval:\n{query}")
 
-# Initialize the Chroma vector store for retrieving documents
-vector_store = Chroma(
-    collection_name="nutritional_hypotheticals",
-    persist_directory="./nutritional_db",
-    embedding_function=embedding_model
-)
-
-# Create a retriever from the vector store
-retriever = vector_store.as_retriever(
-    search_type='similarity',
-    search_kwargs={'k': 3}
-)
-
-def retrieve_context(state):
-    """
-    Retrieves context from the vector store using the expanded or original query.
+    # Ensure vector_store is loaded and available in session_state
+    if 'vector_store' not in st.session_state or st.session_state.vector_store is None:
+        st.error("Vector store not initialized. Cannot retrieve context.")
+        state['context'] = []
+        return state
 
-    Args:
-        state (Dict): The current state of the workflow, containing the query and expanded query.
-
-    Returns:
-        Dict: The updated state with the retrieved context.
-    """
-    print("---------retrieve_context---------")
-    # Add original query to the state to improve result.
-    query = f"{state['query']}; {state['expanded_query']}" # Complete the code to define the key for the expanded query
-    #print("Query used for retrieval:", query)  # Debugging: Print the query
-
-    # Retrieve documents from the vector store
+    retriever = st.session_state.vector_store.as_retriever(
+        search_type='similarity',
+        search_kwargs={'k': 3}
+    )
     docs = retriever.invoke(query)
-    print("Retrieved documents:", docs)  # Debugging: Print the raw docs object
+    st.write(f"Retrieved documents (first 100 chars each):\n{[doc.page_content[:100] for doc in docs]}")
 
-    # Extract both page_content and metadata from each document
-    context= [
-        {
-            "content": doc.page_content,  # The actual content of the document
-            "metadata": doc.metadata  # The metadata (e.g., source, page number, etc.)
-        }
+    context = [
+        {"content": doc.page_content, "metadata": doc.metadata}
         for doc in docs
     ]
-    state['context'] = context # Complete the code to define the key for storing the context
-    print("Extracted context with metadata:", context)  # Debugging: Print the extracted context
-    #print(f"Groundedness loop count: {state['groundedness_loop_count']}")
+    state['context'] = context
+    #st.write(f"Extracted context with metadata:\n{context}") # Too verbose for production UI
     return state
 
-
-def craft_response(state: Dict) -> Dict:
-    """
-    Generates a response using the retrieved context, focusing on nutrition disorders.
-
-    Args:
-        state (Dict): The current state of the workflow, containing the query and retrieved context.
-
-    Returns:
-        Dict: The updated state with the generated response.
-    """
-    print("---------craft_response---------")
+def craft_response(state: AgentState) -> AgentState:
+    st.write("---Crafting Response---")
     system_message = '''
     Generates a response to a user query and context provided.
 
@@ -222,41 +332,28 @@ def craft_response(state: Dict) -> Dict:
 
     The answer you provide must come from the user queries with context provided.
     If feedback is provided, use it to craft the response.
-    If information provided is not enough to answer the query respons with 'I don't know the answer. Not in my records.'
+    If information provided is not enough to answer the query respond with 'I don't know the answer. Not in my records.'
     '''
-
     response_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
-        ("user", "Query: {query}\nContext: {context}\n\nfeedback: {feedback}")
+        ("user", "Query:\n{query}\nContext:\n{context}\n\nfeedback:\n{feedback}")
     ])
-
-    chain = response_prompt | llm
+    chain = response_prompt | st.session_state.llm # Use llm from session state
     response = chain.invoke({
         "query": state['query'],
         "context": "\n".join([doc["content"] for doc in state['context']]),
-        "feedback": state["feedback"] # add feedback to the prompt
+        "feedback": state["feedback"]
     })
-    state['response'] = response
-    print("intermediate response: ", response)
-
+    state['response'] = response.content # Access content from AIMessage
+    st.write(f"Intermediate response:\n{state['response']}")
     return state
 
-
-
-def score_groundedness(state: Dict) -> Dict:
-    """
-    Checks whether the response is grounded in the retrieved context.
-
-    Args:
-        state (Dict): The current state of the workflow, containing the response and context.
-
-    Returns:
-        Dict: The updated state with the groundedness score.
-    """
-    print("---------check_groundedness---------")
+def score_groundedness(state: AgentState) -> AgentState:
+    st.write("---Checking Groundedness---")
     system_message = '''
     You are tasked with rating AI generated answers to questions posed by users.
     Please act as an impartial judge and evaluate the quality of the provided answer which attempts to answer the provided question based on a provided context.
+
     In the input, the context is {context}, while the AI generated response is {response}.
 
     Evaluation criteria:
@@ -272,37 +369,24 @@ def score_groundedness(state: Dict) -> Dict:
     Do not show any instructions for deriving your answer.
 
     Output your result as a float number between 0 and 1 using the evaluation criteria.
-    The better the criteria, the cloase it is to 1 and the worse the criteria, the closer it is to 0.
+    The better the criteria, the closer it is to 1 and the worse the criteria, the closer it is to 0.
     '''
-
     groundedness_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
         ("user", "Context: {context}\nResponse: {response}\n\nGroundedness score:")
     ])
-
-    chain = groundedness_prompt | llm | StrOutputParser()
+    chain = groundedness_prompt | st.session_state.llm | StrOutputParser() # Use llm from session state
     groundedness_score = float(chain.invoke({
         "context": "\n".join([doc["content"] for doc in state['context']]),
-        "response": state['response'] # Complete the code to define the response
+        "response": state['response']
     }))
-    print("groundedness_score: ", groundedness_score)
-    state['groundedness_loop_count'] += 1
-    print("#########Groundedness Incremented###########")
     state['groundedness_score'] = groundedness_score
-
+    state['groundedness_loop_count'] += 1
+    st.write(f"Groundedness score: {groundedness_score}")
     return state
 
-def check_precision(state: Dict) -> Dict:
-    """
-    Checks whether the response precisely addresses the user’s query.
-
-    Args:
-        state (Dict): The current state of the workflow, containing the query and response.
-
-    Returns:
-        Dict: The updated state with the precision score.
-    """
-    print("---------check_precision---------")
+def check_precision(state: AgentState) -> AgentState:
+    st.write("---Checking Precision---")
     system_message = '''
     Given question, answer and context verify if the context was useful in arriving at the given answer.
     Give verdict as "1" if useful and "0" if not useful.
@@ -311,277 +395,191 @@ def check_precision(state: Dict) -> Dict:
     0 or near 0 if it is least useful, 0.5 or near 0.5 if retry is warranted, and 1 or close to 1 is most useful.
     Do not show any instructions for deriving your answer.
     '''
-
     precision_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
         ("user", "Query: {query}\nResponse: {response}\n\nPrecision score:")
     ])
-
-    chain = precision_prompt | llm | StrOutputParser() # Complete the code to define the chain of processing
+    chain = precision_prompt | st.session_state.llm | StrOutputParser() # Use llm from session state
     precision_score = float(chain.invoke({
         "query": state['query'],
-        "response": state['response'] # Complete the code to access the response from the state
+        "response": state['response']
     }))
     state['precision_score'] = precision_score
-    print("precision_score:", precision_score)
     state['precision_loop_count'] +=1
-    print("#########Precision Incremented###########")
+    st.write(f"Precision score: {precision_score}")
     return state
 
-def refine_response(state: Dict) -> Dict:
-    """
-    Suggests improvements for the generated response.
-
-    Args:
-        state (Dict): The current state of the workflow, containing the query and response.
-
-    Returns:
-        Dict: The updated state with response refinement suggestions.
-    """
-    print("---------refine_response---------")
-
+def refine_response(state: AgentState) -> AgentState:
+    st.write("---Refining Response---")
     system_message = '''
-    Since the last response failded the groundedness test, and is deemed not satisfactory,
+    Since the last response failed the groundedness test, and is deemed not satisfactory,
     use the feedback in terms of the query, context and the last response
     to identify potential gaps, ambiguities, or missing details, and
     to suggest improvements to enhance accuracy and completeness of the response.
     '''
-
     refine_response_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
         ("user", "Query: {query}\nResponse: {response}\n\n"
                  "What improvements can be made to enhance accuracy and completeness?")
     ])
-
-    chain = refine_response_prompt | llm| StrOutputParser()
-
-    # Store response suggestions in a structured format
+    chain = refine_response_prompt | st.session_state.llm | StrOutputParser() # Use llm from session state
     feedback = f"Previous Response: {state['response']}\nSuggestions: {chain.invoke({'query': state['query'], 'response': state['response']})}"
-    print("feedback: ", feedback)
-    print(f"State: {state}")
     state['feedback'] = feedback
+    st.write(f"Refinement feedback:\n{feedback}")
     return state
 
-
-
-def refine_query(state: Dict) -> Dict:
-    """
-    Suggests improvements for the expanded query.
-
-    Args:
-        state (Dict): The current state of the workflow, containing the query and expanded query.
-
-    Returns:
-        Dict: The updated state with query refinement suggestions.
-    """
-    print("---------refine_query---------")
+def refine_query(state: AgentState) -> AgentState:
+    st.write("---Refining Query---")
     system_message = '''
-    Since the last response failded the precision test, and is deemed not satisfactory,
+    Since the last response failed the precision test, and is deemed not satisfactory,
     use the feedback in terms of the query, context and re-generate extended queries
     to identify specific keywords, scope refinements, or missing details, and
     to provides structured suggestions for improvement to enhance accuracy and completeness of the response.
     '''
-
     refine_query_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
         ("user", "Original Query: {query}\nExpanded Query: {expanded_query}\n\n"
                  "What improvements can be made for a better search?")
     ])
-
-    chain = refine_query_prompt | llm | StrOutputParser()
-
-    # Store refinement suggestions without modifying the original expanded query
+    chain = refine_query_prompt | st.session_state.llm | StrOutputParser() # Use llm from session state
     query_feedback = f"Previous Expanded Query: {state['expanded_query']}\nSuggestions: {chain.invoke({'query': state['query'], 'expanded_query': state['expanded_query']})}"
-    print("query_feedback: ", query_feedback)
-    print(f"Groundedness loop count: {state['groundedness_loop_count']}")
     state['query_feedback'] = query_feedback
+    st.write(f"Query refinement feedback:\n{query_feedback}")
     return state
 
-
-
-def should_continue_groundedness(state):
-  """Decides if groundedness is sufficient or needs improvement."""
-  print("---------should_continue_groundedness---------")
-  print("groundedness loop count: ", state['groundedness_loop_count'])
-  if state['groundedness_score'] >= 0.8: # Complete the code to define the threshold for groundedness
-      print("Moving to precision")
-      return "check_precision"
-  else:
-      if state["groundedness_loop_count"] > state['loop_max_iter']:
-        return "max_iterations_reached"
-      else:
-        print(f"---------Groundedness Score Threshold Not met. Refining Response-----------")
-        return "refine_response"
-
-
-def should_continue_precision(state: Dict) -> str:
-    """Decides if precision is sufficient or needs improvement."""
-    print("---------should_continue_precision---------")
-    print("precision loop count: ", state["precision_loop_count"])
-    if state['precision_score'] > 0.8: # Threshold for precision
-        return "pass"  # Complete the workflow
+def should_continue_groundedness(state: AgentState) -> str:
+    st.write("---Deciding Groundedness Continuation---")
+    st.write(f"Groundedness loop count: {state['groundedness_loop_count']}")
+    if state['groundedness_score'] >= 0.8:
+        st.write("Moving to precision check.")
+        return "check_precision"
     else:
-        if state["precision_loop_count"] >= state['loop_max_iter']: # Maximum allowed loops
+        if state["groundedness_loop_count"] >= state['loop_max_iter']:
+            st.write("Max iterations reached for groundedness.")
             return "max_iterations_reached"
         else:
-            print(f"---------Precision Score Threshold Not met. Refining Query-----------")  # Debugging
-            return "refine_query" # Refine the query
-
-
-
+            st.write("Groundedness score not met. Refining response.")
+            return "refine_response"
+
+def should_continue_precision(state: AgentState) -> str:
+    st.write("---Deciding Precision Continuation---")
+    st.write(f"Precision loop count: {state['precision_loop_count']}")
+    if state['precision_score'] > 0.8:
+        st.write("Precision sufficient. Ending workflow.")
+        return "pass"
+    else:
+        if state["precision_loop_count"] >= state['loop_max_iter']:
+            st.write("Max iterations reached for precision.")
+            return "max_iterations_reached"
+        else:
+            st.write("Precision score not met. Refining query.")
+            return "refine_query"
 
-def max_iterations_reached(state: Dict) -> Dict:
-    """Handles the case when the maximum number of iterations is reached."""
-    print("---------max_iterations_reached---------")
-    """Handles the case when the maximum number of iterations is reached."""
+def max_iterations_reached(state: AgentState) -> AgentState:
+    st.write("---Max Iterations Reached---")
     response = "I'm unable to refine the response further. Please provide more context or clarify your question."
     state['response'] = response
     return state
 
+@tool
+def agentic_rag_tool(query: str) -> Dict[str, Any]:
+    """
+    Runs the RAG-based agent workflow to generate context-aware responses.
+    This function is exposed as a tool for the overall chatbot.
+    """
+    # Initialize state for the LangGraph workflow
+    inputs = {
+        "query": query,
+        "expanded_query": "",
+        "context": [],
+        "response": "",
+        "precision_score": 0.0,
+        "groundedness_score": 0.0,
+        "groundedness_loop_count": 0,
+        "precision_loop_count": 0,
+        "feedback": "",
+        "query_feedback": "",
+        "loop_max_iter": 3
+    }
+    # Compile the workflow once and store it in session state if not already done
+    if 'workflow_app' not in st.session_state:
+        st.session_state.workflow_app = create_rag_workflow().compile()
 
+    # Invoke the compiled LangGraph workflow
+    output = st.session_state.workflow_app.invoke(inputs)
+    return output
 
-from langgraph.graph import END, StateGraph, START
-
-def create_workflow() -> StateGraph:
-    """Creates the updated workflow for the AI nutrition agent."""
-    workflow = StateGraph(AgentState) # Complete the code to define the initial state of the agent
+def create_rag_workflow() -> StateGraph:
+    """Creates the LangGraph workflow for the RAG agent."""
+    workflow = StateGraph(AgentState)
 
-    # Add processing nodes
-    workflow.add_node("expand_query", expand_query) # Step 1: Expand user query. Complete with the function to expand the query
-    workflow.add_node("retrieve_context", retrieve_context) # Step 2: Retrieve relevant documents. Complete with the function to retrieve context
-    workflow.add_node("craft_response", craft_response) # Step 3: Generate a response based on retrieved data. Complete with the function to craft a response
-    workflow.add_node("score_groundedness", score_groundedness) # Step 4: Evaluate response grounding. Complete with the function to score groundedness
-    workflow.add_node("refine_response", refine_response) # Step 5: Improve response if it's weakly grounded. Complete with the function to refine the response
-    workflow.add_node("check_precision", check_precision) # Step 6: Evaluate response precision. Complete with the function to check precision
-    workflow.add_node("refine_query", refine_query) # Step 7: Improve query if response lacks precision. Complete with the function to refine the query
-    workflow.add_node("max_iterations_reached", max_iterations_reached) # Step 8: Handle max iterations. Complete with the function to handle max iterations
+    workflow.add_node("expand_query", expand_query)
+    workflow.add_node("retrieve_context", retrieve_context)
+    workflow.add_node("craft_response", craft_response)
+    workflow.add_node("score_groundedness", score_groundedness)
+    workflow.add_node("refine_response", refine_response)
+    workflow.add_node("check_precision", check_precision)
+    workflow.add_node("refine_query", refine_query)
+    workflow.add_node("max_iterations_reached", max_iterations_reached)
 
-    # Main flow edges
     workflow.add_edge(START, "expand_query")
     workflow.add_edge("expand_query", "retrieve_context")
     workflow.add_edge("retrieve_context", "craft_response")
     workflow.add_edge("craft_response", "score_groundedness")
 
-    # Conditional edges based on groundedness check
     workflow.add_conditional_edges(
         "score_groundedness",
-        should_continue_groundedness, # Use the conditional function
+        should_continue_groundedness,
         {
-            "check_precision": "check_precision", # If well-grounded, proceed to precision check.
-            "refine_response": "refine_response", # If not, refine the response.
-            "max_iterations_reached": "max_iterations_reached" # If max loops reached, exit.
+            "check_precision": "check_precision",
+            "refine_response": "refine_response",
+            "max_iterations_reached": "max_iterations_reached"
         }
     )
+    workflow.add_edge("refine_response", "craft_response")
 
-    workflow.add_edge("refine_response", "craft_response") # Refined responses are reprocessed.
-
-    # Conditional edges based on precision check
     workflow.add_conditional_edges(
         "check_precision",
-        should_continue_precision, # Use the conditional function
+        should_continue_precision,
         {
-            "pass": END, # If precise, complete the workflow.
-            "refine_query": "refine_query", # If imprecise, refine the query.
-            "max_iterations_reached": "max_iterations_reached" # If max loops reached, exit.
+            "pass": END,
+            "refine_query": "refine_query",
+            "max_iterations_reached": "max_iterations_reached"
         }
     )
-
-    workflow.add_edge("refine_query", "expand_query") # Refined queries go through expansion again.
-
+    workflow.add_edge("refine_query", "expand_query")
     workflow.add_edge("max_iterations_reached", END)
 
     return workflow
 
-
-
-
-#=========================== Defining the agentic rag tool ====================#
-
-WORKFLOW_APP = create_workflow().compile()
-
-@tool
-def agentic_rag(query: str):
-    """
-    Runs the RAG-based agent with conversation history for context-aware responses.
-
-    Args:
-        query (str): The current user query.
-
-    Returns:
-        Dict[str, Any]: The updated state with the generated response and conversation history.
-    """
-    # Initialize state with necessary parameters
-    inputs = {
-        "query": query,  # Current user query
-        "expanded_query": "", # Complete the code to define the expanded version of the query
-        "context": [],  # Retrieved documents (initially empty)
-        "response": "", # Complete the code to define the AI-generated response
-        "precision_score": 0.0, # Complete the code to define the precision score of the response
-        "groundedness_score": 0.0, # Complete the code to define the groundedness score of the response
-        "groundedness_loop_count": 0, # Complete the code to define the counter for groundedness loops
-        "precision_loop_count": 0, # Complete the code to define the counter for precision loops
-        "feedback": "", # Complete the code to define the feedback
-        "query_feedback": "", # Complete the code to define the query feedback
-        "loop_max_iter": 3 # Complete the code to define the maximum number of iterations for loops
-    }
-
-    output = WORKFLOW_APP.invoke(inputs)
-
-    return output
-
-#================================ Guardrails ===========================#
-llama_guard_client = Groq(api_key=groq_api_key) # Groq(api_key=llama_api_key)
-# Function to filter user input with Llama Guard
-#def filter_input_with_llama_guard(user_input, model="llama-guard-3-8b"):
-def filter_input_with_llama_guard(user_input, model="meta-llama/llama-guard-4-12b"):
-    """
-    Filters user input using Llama Guard to ensure it is safe.
-
-    Parameters:
-    - user_input: The input provided by the user.
-    - model: The Llama Guard model to be used for filtering (default is "llama-guard-3-8b").
-
-    Returns:
-    - The filtered and safe input.
-    """
-    try:
-        # Create a request to Llama Guard to filter the user input
-        response = llama_guard_client.chat.completions.create(
-            messages=[{"role": "user", "content": user_input}],
-            model=model,
-        )
-        # Return the filtered input
-        return response.choices[0].message.content.strip()
-    except Exception as e:
-        print(f"Error with Llama Guard: {e}")
-        return None
-
-
-#============================= Adding Memory to the agent using mem0 ===============================#
+# --- 4. Main Chatbot Class (Integrating Memory & Agent) ---
 
 class NutritionBot:
-    def __init__(self):
+    def __init__(self, config: Dict):
         """
         Initialize the NutritionBot class, setting up memory, the LLM client, tools, and the agent executor.
         """
+        mem0_api_key = config["MEM0_API_KEY"]
+        openai_api_key = config["AZURE_OPENAI_API_KEY"]
+        openai_api_base = config["AZURE_OPENAI_API_BASE"]
 
         # Initialize a memory client to store and retrieve customer interactions
-        self.memory = MemoryClient(api_key=MEM0_api_key) # userdata.get("mem0"))  # Complete the code to define the memory client API key
+        self.memory = MemoryClient(api_key=mem0_api_key)
 
+        # Initialize the OpenAI client (LangChain ChatOpenAI)
         self.client = ChatOpenAI(
-            #model_name="gpt-4o",  # Specify the model to use (e.g., GPT-4 optimized version)
             model="gpt-4o-mini",
-            #api_key = api_key, # config.get("API_KEY"),  # API key for authentication
-            openai_api_key=api_key,  # Fill in the API key
-            #base_url = endpoint, # config.get("OPENAI_API_BASE"),
-            openai_api_base=endpoint,  # Fill in the endpoint
-            temperature=0  # Controls randomness in responses; 0 ensures deterministic results
+            openai_api_key=openai_api_key,
+            openai_api_base=openai_api_base,
+            temperature=0
         )
+        # Store LLM in session state for use in graph nodes
+        st.session_state.llm = self.client
 
-        # Define tools available to the chatbot, such as web search
-        tools = [agentic_rag]
+        # Define tools available to the chatbot
+        tools = [agentic_rag_tool]
 
-        # Define the system prompt to set the behavior of the chatbot
+        # Define the system prompt for the agent
         system_prompt = """You are a caring and knowledgeable Medical Support Agent, specializing in nutrition disorder-related guidance. Your goal is to provide accurate, empathetic, and tailored nutritional recommendations while ensuring a seamless customer experience.
                           Guidelines for Interaction:
                           Maintain a polite, professional, and reassuring tone.
@@ -590,16 +588,16 @@ class NutritionBot:
                           Engage with the customer by asking about their food preferences, dietary restrictions, and lifestyle before offering recommendations.
                           Ensure consistent and accurate information across conversations.
                           If any detail is unclear or missing, proactively ask for clarification.
-                          Always use the agentic_rag tool to retrieve up-to-date and evidence-based nutrition insights.
+                          Always use the agentic_rag_tool to retrieve up-to-date and evidence-based nutrition insights.
                           Keep track of ongoing issues and follow-ups to ensure continuity in support.
                           Your primary goal is to help customers make informed nutrition decisions that align with their health conditions and personal preferences.
         """
 
         # Build the prompt template for the agent
         prompt = ChatPromptTemplate.from_messages([
-            ("system", system_prompt),  # System instructions
-            ("human", "{input}"),  # Placeholder for human input
-            ("placeholder", "{agent_scratchpad}")  # Placeholder for intermediate reasoning steps
+            ("system", system_prompt),
+            ("human", "{input}"),
+            ("placeholder", "{agent_scratchpad}")
         ])
 
         # Create an agent capable of interacting with tools and executing tasks
@@ -609,80 +607,35 @@ class NutritionBot:
         self.agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
     def store_customer_interaction(self, user_id: str, message: str, response: str, metadata: Dict = None):
-        """
-        Store customer interaction in memory for future reference.
-
-        Args:
-            user_id (str): Unique identifier for the customer.
-            message (str): Customer's query or message.
-            response (str): Chatbot's response.
-            metadata (Dict, optional): Additional metadata for the interaction.
-        """
+        """Store customer interaction in memory for future reference."""
         if metadata is None:
             metadata = {}
-
-        # Add a timestamp to the metadata for tracking purposes
         metadata["timestamp"] = datetime.now().isoformat()
-
-        # Format the conversation for storage
         conversation = [
             {"role": "user", "content": message},
             {"role": "assistant", "content": response}
         ]
-
-        # Store the interaction in the memory client
-        self.memory.add(
-            conversation,
-            user_id=user_id,
-            output_format="v1.1",
-            metadata=metadata
-        )
-
+        self.memory.add(conversation, user_id=user_id, output_format="v1.1", metadata=metadata)
 
     def get_relevant_history(self, user_id: str, query: str) -> List[Dict]:
-        """
-        Retrieve past interactions relevant to the current query.
-
-        Args:
-            user_id (str): Unique identifier for the customer.
-            query (str): The customer's current query.
-
-        Returns:
-            List[Dict]: A list of relevant past interactions.
-        """
-        return self.memory.search(
-            query=query,  # Search for interactions related to the query
-            user_id=user_id,  # Restrict search to the specific user
-            limit=5  # Complete the code to define the limit for retrieved interactions
-        )
-
+        """Retrieve past interactions relevant to the current query."""
+        return self.memory.search(query=query, user_id=user_id, limit=5)
 
     def handle_customer_query(self, user_id: str, query: str) -> str:
-        """
-        Process a customer's query and provide a response, taking into account past interactions.
-
-        Args:
-            user_id (str): Unique identifier for the customer.
-            query (str): Customer's query.
-
-        Returns:
-            str: Chatbot's response.
-        """
-
-        # Retrieve relevant past interactions for context
+        """Process a customer's query and provide a response, taking into account past interactions."""
         relevant_history = self.get_relevant_history(user_id, query)
 
-        # Build a context string from the relevant history
         context = "Previous relevant interactions:\n"
-        for memory in relevant_history:
-            context += f"Customer: {memory['memory']}\n"  # Customer's past messages
-            context += f"Support: {memory['memory']}\n"  # Chatbot's past responses
+        for memory_item in relevant_history:
+            # Mem0 'memory' field is typically a list of dicts or a string.
+            # Assuming 'v1.1' output format from `memory.add` means `memory_item['memory']` is structured.
+            if isinstance(memory_item.get('memory'), list):
+                for part in memory_item['memory']:
+                    context += f"{part['role'].capitalize()}: {part['content']}\n"
+            else: # Fallback for older formats or if it's a simple string
+                 context += f"History: {memory_item.get('memory', 'N/A')}\n"
             context += "---\n"
 
-        # Print context for debugging purposes
-        print("Context: ", context)
-
-        # Prepare a prompt combining past context and the current query
         prompt = f"""
         Context:
         {context}
@@ -691,62 +644,108 @@ class NutritionBot:
 
         Provide a helpful response that takes into account any relevant past interactions.
         """
-        #st.write("Context: ", prompt)
-
-        # Generate a response using the agent
-        response = self.agent_executor.invoke({"input": prompt})
-        st.write("Context: ", response)
-
-        # Store the current interaction for future reference
-        self.store_customer_interaction(
-            user_id=user_id,
-            message=query,
-            response=response["output"],
-            metadata={"type": "support_query"}
-        )
+        st.write("Prompt sent to agent executor:", prompt) # Debugging
 
-        # Return the chatbot's response
-        return response['output']
+        try:
+            response_dict = self.agent_executor.invoke({"input": prompt})
+            response_content = response_dict.get('output', "I'm sorry, I couldn't generate a response.")
+        except Exception as e:
+            st.error(f"Error during agent execution: {e}")
+            response_content = f"I'm sorry, I encountered an internal error: {e}"
 
+        self.store_customer_interaction(user_id=user_id, message=query, response=response_content, metadata={"type": "support_query"})
+        return response_content
 
-#=====================User Interface using streamlit ===========================#
-def nutrition_disorder_streamlit():
-    """
-    A Streamlit-based UI for the Nutrition Disorder Specialist Agent.
-    """
-    st.title("Nutrition Disorder Specialist")
-    st.write("Ask me anything about nutrition disorders, symptoms, causes, treatments, and more.")
-    st.write("Type 'exit' to end the conversation.")
+# --- 5. Streamlit UI ---
+
+def nutrition_disorder_streamlit_app():
+    """Streamlit-based UI for the Nutrition Disorder Specialist Agent."""
+    st.set_page_config(page_title="Nutrition Disorder Specialist", layout="centered")
+
+    st.title("👨‍⚕️ Nutrition Disorder Specialist")
+    st.markdown("Ask me anything about nutrition disorders, symptoms, causes, treatments, and more.")
+    st.markdown("---")
 
-    # Initialize session state for chat history and user_id if they don't exist
+    # Initialize session state variables
     if 'chat_history' not in st.session_state:
         st.session_state.chat_history = []
     if 'user_id' not in st.session_state:
         st.session_state.user_id = None
-
-    # Login form: Only if user is not logged in
+    if 'chatbot' not in st.session_state:
+        st.session_state.chatbot = None
+    if 'config_loaded' not in st.session_state:
+        st.session_state.config_loaded = False
+    if 'vector_store_loaded' not in st.session_state:
+        st.session_state.vector_store_loaded = False
+
+    # --- Configuration Loading and Model Initialization ---
+    if not st.session_state.config_loaded:
+        with st.spinner("Loading configurations and initializing models..."):
+            config = load_config_from_env()
+            if not all(config.values()):
+                st.error("Some environment variables are missing. Please set them up for the app to function.")
+                st.stop() # Stop execution if critical configs are missing
+
+            # Step 1.
+            embedding_model, llm_instance, embedding_function, llama_guard_client_instance = initialize_llms_and_embeddings(config)
+
+            # Step 2. Store initialized objects in session state
+            st.session_state.config = config
+            st.session_state.embedding_model = embedding_model
+            st.session_state.llm = llm_instance
+            st.session_state.embedding_function = embedding_function # Used during vector_store creation/loading
+            st.session_state.llama_guard_client = llama_guard_client_instance
+            st.session_state.config_loaded = True
+            st.rerun() # Rerun to update UI after loading
+
+    # --- Vector Store Loading ---
+    if st.session_state.config_loaded and not st.session_state.vector_store_loaded:
+        with st.spinner("Loading nutrition knowledge base (vector database)..."):
+            # Ensure the nutritional_db directory exists relative to the app.py
+            # In Docker, this means the folder should be copied into /app
+            persist_dir = "./nutritional_db"
+            if not os.path.exists(persist_dir):
+                st.error(f"Required data directory '{persist_dir}' not found. Please ensure it's copied into the Docker image.")
+                st.stop()
+
+            st.session_state.vector_store = load_vector_db(
+                embedding_model=st.session_state.embedding_model,
+                collection_name="nutritional_hypotheticals",
+                persist_directory=persist_dir
+            )
+            if st.session_state.vector_store is None:
+                st.error("Failed to load vector database. Chat functionality will be limited.")
+            st.session_state.vector_store_loaded = True
+            st.rerun() # Rerun to update UI after loading
+
+
+    # --- Login Form ---
     if st.session_state.user_id is None:
         with st.form("login_form", clear_on_submit=True):
-            user_id = st.text_input("Please enter your name to begin:")
+            user_id_input = st.text_input("Please enter your name to begin:", key="user_id_input")
             submit_button = st.form_submit_button("Login")
-            if submit_button and user_id:
-                st.session_state.user_id = user_id
+            if submit_button and user_id_input:
+                st.session_state.user_id = user_id_input.strip()
                 st.session_state.chat_history.append({
                     "role": "assistant",
-                    "content": f"Welcome, {user_id}! How can I help you with nutrition disorders today?"
+                    "content": f"Welcome, {st.session_state.user_id}! How can I help you with nutrition disorders today?"
                 })
-                st.session_state.login_submitted = True  # Set flag to trigger rerun
-        if st.session_state.get("login_submitted", False):
-            st.session_state.pop("login_submitted")
-            st.rerun()
-    else:
+                # Initialize chatbot only after config and vector store are ready
+                if st.session_state.config_loaded and st.session_state.vector_store_loaded:
+                    st.session_state.chatbot = NutritionBot(st.session_state.config)
+                else:
+                    st.warning("Chatbot initialization delayed as configurations or vector store are still loading.")
+                st.rerun()
+
+    # --- Chat Interface ---
+    elif st.session_state.chatbot: # Only show chat if chatbot is initialized
         # Display chat history
         for message in st.session_state.chat_history:
             with st.chat_message(message["role"]):
                 st.write(message["content"])
 
-        # Chat input with custom placeholder text
-        user_query = st.chat_input("Type your question here (or 'exit' to end)...") # Blank #1: Fill in the chat input prompt (e.g., "Type your question here (or 'exit' to end)...")
+        user_query = st.chat_input("Type your question here (e.g., 'What are dietary deficiencies?')")
+
         if user_query:
             if user_query.lower() == "exit":
                 st.session_state.chat_history.append({"role": "user", "content": "exit"})
@@ -756,7 +755,9 @@ def nutrition_disorder_streamlit():
                 st.session_state.chat_history.append({"role": "assistant", "content": goodbye_msg})
                 with st.chat_message("assistant"):
                     st.write(goodbye_msg)
-                st.session_state.user_id = None
+                st.session_state.user_id = None # Log out
+                st.session_state.chatbot = None # Clear chatbot instance
+                st.session_state.chat_history = [] # Clear history on exit
                 st.rerun()
                 return
 
@@ -765,44 +766,35 @@ def nutrition_disorder_streamlit():
                 st.write(user_query)
 
             # Filter input using Llama Guard
-            filtered_result = filter_input_with_llama_guard(user_query) # Blank #2: Fill in with the function name for filtering input (e.g., filter_input_with_llama_guard)
-            #filtered_result = filtered_result.replace("\n", " ")  # Normalize the result
-            #print(f"filtered_result (1): {filtered_result}")
-            st.write(f"filtered_result (1): {filtered_result}")
-            if filtered_result is None:
-                #print("Agent: Sorry, I encountered an error while filtering your input. Please try again.")
-                st.write("Agent: Sorry, I encountered an error while filtering your input. Please try again.")
-                #continue
-            else:
-                filtered_result = filtered_result.replace("\n", " ")  # Normalize the result
-                #print(f"filtered_result (2): {filtered_result}")
-                st.write(f"filtered_result (2): {filtered_result}")
-
-            # Check if input is safe based on allowed statuses
-            if filtered_result in ["safe", "safe S7", "safe S6"]: # Blanks #3, #4, #5: Fill in with allowed safe statuses (e.g., "safe", "unsafe S7", "unsafe S6")
-                st.write("Input is safe.")
-                try:
-
-                    if 'chatbot' not in st.session_state:
-                        st.session_state.chatbot = NutritionBot() # Blank #6: Fill in with the chatbot class initialization (e.g., NutritionBot)
-                    
-                    st.write("chatbot is calling handle_customer_query...")
-                    st.write("user_id: ", st.session_state.user_id)
-                    st.write("user_query: ", user_query)
-                    
-                    response = st.session_state.chatbot.handle_customer_query(st.session_state.user_id, user_query)
-                    st.write("response is returned.")
-                    # Blank #7: Fill in with the method to handle queries (e.g., handle_customer_query)
-
-                    st.write(response)
-                    st.session_state.chat_history.append({"role": "assistant", "content": response})
-                except Exception as e:
-                    st.write(f"Sorry, I encountered an error while processing your query. Please try again. Error: {e}", str(e))
-                    st.session_state.chat_history.append({"role": "assistant", "content": f"Sorry, I encountered an error while processing your query. Please try again. Error: {e}"})
+            with st.spinner("Filtering input for safety..."):
+                filtered_result = filter_input_with_llama_guard(user_query, st.session_state.llama_guard_client)
+                if filtered_result:
+                    filtered_result = filtered_result.replace("\n", " ").strip()
+                st.info(f"Llama Guard says: {filtered_result}") # Show Llama Guard's verdict
+
+            # Process the user query if safe
+            if filtered_result and ("safe" in filtered_result.lower() or "s7" in filtered_result.lower()): # Allow "safe S7" etc.
+                with st.spinner("Thinking..."):
+                    try:
+                        response = st.session_state.chatbot.handle_customer_query(st.session_state.user_id, user_query)
+                        st.session_state.chat_history.append({"role": "assistant", "content": response})
+                        with st.chat_message("assistant"):
+                            st.write(response)
+                    except Exception as e:
+                        error_msg = f"Sorry, I encountered an error while processing your query. Please try again. Error: {e}"
+                        st.error(error_msg)
+                        st.session_state.chat_history.append({"role": "assistant", "content": error_msg})
             else:
-                inappropriate_msg = "I apologize, but I cannot process that input as it may be inappropriate. Please try again."
-                st.write(inappropriate_msg)
+                inappropriate_msg = "I apologize, but I cannot process that input as it may be inappropriate or unsafe. Please try again."
                 st.session_state.chat_history.append({"role": "assistant", "content": inappropriate_msg})
+                with st.chat_message("assistant"):
+                    st.write(inappropriate_msg)
+            st.rerun() # Rerun to update chat history instantly
+
+    elif st.session_state.user_id: # User is logged in but chatbot not ready
+        st.info("Initializing chatbot. Please wait...")
+
 
+# --- Main entry point for Streamlit App ---
 if __name__ == "__main__":
-    nutrition_disorder_streamlit()
+    nutrition_disorder_streamlit_app()