Upload app.py with huggingface_hub

app.py

@@ -29,7 +29,7 @@ from langchain.agents import create_tool_calling_agent, AgentExecutor
 from langchain_core.prompts import ChatPromptTemplate
 
 # LangChain OpenAI imports
-from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI, ChatOpenAI # OpenAI embeddings and models
+from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI, ChatOpenAI  # OpenAI embeddings and models
 from langchain.embeddings.openai import OpenAIEmbeddings  # OpenAI embeddings for text vectors
 
 # LlamaParse & LlamaIndex imports
@@ -54,7 +54,7 @@ from datetime import datetime
 
 #====================================SETUP=====================================#
 # Fetch secrets from Hugging Face Spaces
-api_key = os.getenv("API_KEY")
+api_key = os.getenv("API_KEY")  #config.get("API_KEY")
 endpoint = os.getenv("OPENAI_API_BASE")
 llama_api_key = os.environ['GROQ_API_KEY']
 MEM0_api_key = os.environ['mem0']
@@ -66,7 +66,6 @@ embedding_function = chromadb.utils.embedding_functions.OpenAIEmbeddingFunction(
     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(
@@ -83,14 +82,11 @@ llm = ChatOpenAI(
     model="gpt-4o-mini",
     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 = _____ # Complete the code to define the LLM model
-# Settings.embedding = _____ # Complete the code to define the embedding model
-Settings.llm = llm # Complete the code to define the LLM model
-Settings.embedding = embedding_model # Complete the code to define the embedding model
 
+# set the LLM and embedding model in the LlamaIndex settings.
+Settings.llm = llm
+Settings.embedding = embedding_model
 
 #================================Creating Langgraph agent======================#
 
@@ -119,10 +115,29 @@ def expand_query(state):
         Dict: The updated state with the expanded query.
     """
     print("---------Expanding Query---------")
-    # system_message = '''________________________'''
-    system_message = '''You are a nutrition expert and language model specialized in nutritional disorders. Your task is to expand the provided query by incorporating related keywords, synonyms, and additional context that can improve the retrieval of detailed nutrition disorder-related information.'''
+    system_message = '''You are an AI specializing in improving search queries to retrieve the most relevant nutrition disorder-related information.
+Your task is to **refine** and **expand** the given query so that better search results are obtained, while **keeping the original intent** unchanged.
+
+Guidelines:
+- Add **specific details** where needed. Example: If a user asks about "anorexia," specify aspects like symptoms, causes, or treatment options.
+- Include **related terms** to improve retrieval (e.g., “bulimia” → “bulimia nervosa vs binge eating disorder”).
+- If the user provides an unclear query, suggest necessary clarifications.
+- **DO NOT** answer the question. Your job is only to enhance the query.
+
+Examples:
+1. User Query: "Tell me about eating disorders."
+   Expanded Query: "Provide details on eating disorders, including types (e.g., anorexia nervosa, bulimia nervosa), symptoms, causes, and treatment options."
 
+2. User Query: "What is anorexia?"
+   Expanded Query: "Explain anorexia nervosa, including its symptoms, causes, risk factors, and treatment options."
 
+3. User Query: "How to treat bulimia?"
+   Expanded Query: "Describe treatment options for bulimia nervosa, including psychotherapy, medications, and lifestyle changes."
+
+4. User Query: "What are the effects of malnutrition?"
+   Expanded Query: "Explain the effects of malnutrition on physical and mental health, including specific nutrient deficiencies and their consequences."
+
+Now, expand the following query:'''
 
     expand_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
@@ -162,8 +177,7 @@ def retrieve_context(state):
         Dict: The updated state with the retrieved context.
     """
     print("---------retrieve_context---------")
-    # query = state['_____']  # Complete the code to define the key for the expanded query
-    query = state['expanded_query']  # Complete the code to define the key for the expanded query
+    query = state['expanded_query']
     #print("Query used for retrieval:", query)  # Debugging: Print the query
 
     # Retrieve documents from the vector store
@@ -178,9 +192,7 @@ def retrieve_context(state):
         }
         for doc in docs
     ]
-    # state['_____'] = context  # Complete the code to define the key for storing the context
-    state['context'] = context  # Complete the code to define the key for storing the context
-
+    state['context'] = context
     print("Extracted context with metadata:", context)  # Debugging: Print the extracted context
     #print(f"Groundedness loop count: {state['groundedness_loop_count']}")
     return state
@@ -197,27 +209,38 @@ def craft_response(state: Dict) -> Dict:
     Returns:
         Dict: The updated state with the generated response.
     """
-    print("---------craft_response---------")
-    # system_message = '''________________________'''
-    system_message = '''You are a nutrition expert and your responses should be clear, concise, and evidence-based. Use the provided context to accurately address the user's query regarding nutritional disorders.'''
-
+    system_message = '''You are a professional AI nutrition disorder specialist generating responses based on retrieved documents.
+Your task is to use the given **context** to generate a highly accurate, informative, and user-friendly response.
+
+Guidelines:
+- **Be direct and concise** while ensuring completeness.
+- **DO NOT include information that is not present in the context.**
+- If multiple sources exist, synthesize them into a coherent response.
+- If the context does not fully answer the query, state what additional information is needed.
+- Use bullet points when explaining complex concepts.
+
+Example:
+User Query: "What are the symptoms of anorexia nervosa?"
+Context:
+1. Anorexia nervosa is characterized by extreme weight loss and fear of gaining weight.
+2. Common symptoms include restricted eating, distorted body image, and excessive exercise.
+Response:
+"Anorexia nervosa is an eating disorder characterized by extreme weight loss and an intense fear of gaining weight. Common symptoms include:
+- Restricted eating
+- Distorted body image
+- Excessive exercise
+If you or someone you know is experiencing these symptoms, it is important to seek professional help."'''
 
     response_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
-        ("user", "Query: {query}\nContext: {context}\n\nfeedback: {feedback}")
+        ("user", "Query: {query}\nContext: {context}\n\nResponse:")
     ])
 
-    chain = response_prompt | llm
-    response = chain.invoke({
+    chain = response_prompt | llm | StrOutputParser()
+    state['response'] = chain.invoke({
         "query": state['query'],
-        "context": "\n".join([doc["content"] for doc in state['context']]),
-        # "feedback": ________________ # add feedback to the prompt
-        "feedback": state.get("query_feedback", "No additional feedback provided")  # add feedback to the prompt
-
+        "context": "\n".join([doc["content"] for doc in state['context']])  # Extract content from each document
     })
-    state['response'] = response
-    print("intermediate response: ", response)
-
     return state
 
 
@@ -233,9 +256,37 @@ def score_groundedness(state: Dict) -> Dict:
         Dict: The updated state with the groundedness score.
     """
     print("---------check_groundedness---------")
-    # system_message = '''________________________'''
-    system_message = '''You are an evaluator for response groundedness. Given the context and the response related to nutritional disorders, provide a numerical score between 0 and 1 where 0 means the response is not grounded at all, and 1 means it is completely grounded in the context.'''
+    system_message = '''You are an AI tasked with evaluating whether a response is grounded in the provided context and includes proper citations.
+
+Guidelines:
+1. **Groundedness Check**:
+   - Verify that the response accurately reflects the information in the context.
+   - Flag any unsupported claims or deviations from the context.
+
+2. **Citation Check**:
+   - Ensure that the response includes citations to the source material (e.g., "According to [Source], ...").
+   - If citations are missing, suggest adding them.
 
+3. **Scoring**:
+   - Assign a groundedness score between 0 and 1, where 1 means fully grounded and properly cited.
+
+Examples:
+1. Response: "Anorexia nervosa is caused by genetic factors (Source 1)."
+   Context: "Anorexia nervosa is influenced by genetic, environmental, and psychological factors (Source 1)."
+   Evaluation: "The response is grounded and properly cited. Groundedness score: 1.0."
+
+2. Response: "Bulimia nervosa can be cured with diet alone."
+   Context: "Treatment for bulimia nervosa involves psychotherapy and medications (Source 2)."
+   Evaluation: "The response is ungrounded and lacks citations. Groundedness score: 0.2."
+
+3. Response: "Anorexia nervosa has a high mortality rate."
+   Context: "Anorexia nervosa has one of the highest mortality rates among psychiatric disorders (Source 3)."
+   Evaluation: "The response is grounded but lacks a citation. Groundedness score: 0.7. ."
+
+   ****Return only a float score (e.g., 0.9). Do not provide explanations.****
+
+Now, evaluate the following response:
+'''
 
     groundedness_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
@@ -245,14 +296,12 @@ def score_groundedness(state: Dict) -> Dict:
     chain = groundedness_prompt | llm | StrOutputParser()
     groundedness_score = float(chain.invoke({
         "context": "\n".join([doc["content"] for doc in state['context']]),
-        # "response": __________ # Complete the code to define the response
-        "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_score: ",groundedness_score)
+    state['groundedness_loop_count'] +=1
     print("#########Groundedness Incremented###########")
     state['groundedness_score'] = groundedness_score
-
     return state
 
 
@@ -268,23 +317,49 @@ def check_precision(state: Dict) -> Dict:
         Dict: The updated state with the precision score.
     """
     print("---------check_precision---------")
-    # system_message = '''________________________'''
-    system_message = '''You are an evaluator for response precision. Given the query and the response, provide a numerical score between 0 and 1 where 0 indicates that the response does not address the query at all, and 1 indicates that the response precisely addresses the query.'''
-
-
+    system_message = '''You are an AI evaluator assessing the **precision** of the response.
+Your task is to **score** how well the response addresses the user’s original nutrition disorder-related query.
+
+Scoring Criteria:
+- 1.0 → The response is fully precise, directly answering the question.
+- 0.7 → The response is mostly correct but contains some generalization.
+- 0.5 → The response is somewhat relevant but lacks key details.
+- 0.3 → The response is vague or only partially correct.
+- 0.0 → The response is incorrect or misleading.
+
+Examples:
+1. Query: "What are the symptoms of anorexia nervosa?"
+   Response: "The symptoms of anorexia nervosa include extreme weight loss, fear of gaining weight, and a distorted body image."
+   Precision Score: 1.0
+
+2. Query: "How is bulimia nervosa treated?"
+   Response: "Bulimia nervosa is treated with therapy and medications."
+   Precision Score: 0.7
+
+3. Query: "What causes binge eating disorder?"
+   Response: "Binge eating disorder is caused by a combination of genetic, psychological, and environmental factors."
+   Precision Score: 0.5
+
+4. Query: "What are the effects of malnutrition?"
+   Response: "Malnutrition can lead to health problems."
+   Precision Score: 0.3
+
+5. Query: "What is the mortality rate of anorexia nervosa?"
+   Response: "Anorexia nervosa is a type of eating disorder."
+   Precision Score: 0.0
+
+   *****Return only a float score (e.g., 0.9). Do not provide explanations.*****
+Now, evaluate the following query and response:
+'''
     precision_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
         ("user", "Query: {query}\nResponse: {response}\n\nPrecision score:")
     ])
 
-    # chain = _____________ | llm | StrOutputParser() # Complete the code to define the chain of processing
-    chain = precision_prompt | llm | StrOutputParser() # Complete the code to define the chain of processing
-
+    chain = precision_prompt | llm | StrOutputParser()
     precision_score = float(chain.invoke({
         "query": state['query'],
-        # "response":______________ # Complete the code to access the response from the state
-        "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)
@@ -306,9 +381,29 @@ def refine_response(state: Dict) -> Dict:
     """
     print("---------refine_response---------")
 
-    # system_message = '''________________________'''
-    system_message = '''You are an expert editor in nutritional science communications. Your role is to review the response given to a nutritional query and provide clear suggestions to improve its accuracy, clarity, and completeness. Focus on making sure that the response fully addresses the query and is supported by evidence-based nutritional guidelines.'''
+    system_message = '''You are an AI response refinement assistant. Your task is to suggest **improvements** for the given response.
+
+### Guidelines:
+- Identify **gaps in the explanation** (missing key details).
+- Highlight **unclear or vague parts** that need elaboration.
+- Suggest **additional details** that should be included for better accuracy.
+- Ensure the refined response is **precise** and **grounded** in the retrieved context.
+
+### Examples:
+1. Query: "What are the symptoms of anorexia nervosa?"
+   Response: "The symptoms include weight loss and fear of gaining weight."
+   Suggestions: "The response is missing key details about behavioral and emotional symptoms. Add details like 'distorted body image' and 'restrictive eating patterns.'"
+
+2. Query: "How is bulimia nervosa treated?"
+   Response: "Bulimia nervosa is treated with therapy."
+   Suggestions: "The response is too vague. Specify the types of therapy (e.g., cognitive-behavioral therapy) and mention other treatments like nutritional counseling and medications."
+
+3. Query: "What causes binge eating disorder?"
+   Response: "Binge eating disorder is caused by psychological factors."
+   Suggestions: "The response is incomplete. Add details about genetic and environmental factors, and explain how they contribute to the disorder."
 
+Now, suggest improvements for the following response:
+'''
 
     refine_response_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
@@ -338,9 +433,28 @@ def refine_query(state: Dict) -> Dict:
         Dict: The updated state with query refinement suggestions.
     """
     print("---------refine_query---------")
-    # system_message = '''________________________'''
-    system_message = '''You are a search query refinement expert. Given the original and expanded queries related to nutritional disorders, provide suggestions to refine the query further for improved search results.'''
+    system_message = '''You are an AI query refinement assistant. Your task is to suggest **improvements** for the expanded query.
 
+### Guidelines:
+- Add **specific keywords** to improve document retrieval.
+- Identify **missing details** that should be included.
+- Suggest **ways to narrow the scope** for better precision.
+
+### Examples:
+1. Original Query: "Tell me about eating disorders."
+   Expanded Query: "Provide details on eating disorders, including types, symptoms, causes, and treatment options."
+   Suggestions: "Add specific types of eating disorders like 'anorexia nervosa' and 'bulimia nervosa' to improve retrieval."
+
+2. Original Query: "What is anorexia?"
+   Expanded Query: "Explain anorexia nervosa, including its symptoms and causes."
+   Suggestions: "Include details about treatment options and risk factors to make the query more comprehensive."
+
+3. Original Query: "How to treat bulimia?"
+   Expanded Query: "Describe treatment options for bulimia nervosa."
+   Suggestions: "Specify types of treatments like 'cognitive-behavioral therapy' and 'medications' for better precision."
+
+Now, suggest improvements for the following expanded query:
+'''
 
     refine_query_prompt = ChatPromptTemplate.from_messages([
         ("system", system_message),
@@ -363,8 +477,7 @@ 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'] >= _____:  # Complete the code to define the threshold for groundedness
-  if state['groundedness_score'] >= 0.7:  # Complete the code to define the threshold for groundedness
+  if state['groundedness_score'] >= 0.4:  # Threshold for groundedness
       print("Moving to precision")
       return "check_precision"
   else:
@@ -378,24 +491,19 @@ def should_continue_groundedness(state):
 def should_continue_precision(state: Dict) -> str:
     """Decides if precision is sufficient or needs improvement."""
     print("---------should_continue_precision---------")
-    # print("precision loop count: ", ___________)
-    # if ___________:  # Threshold for precision
-
-    print("precision loop count: ", state['precision_loop_count'])
-    if state['precision_score'] >= 0.8:  # Threshold for precision
+    print("precision loop count: ",state['precision_loop_count'])
+    if state['precision_score'] >= 0.7:  # Threshold for precision
         return "pass"  # Complete the workflow
     else:
-        # if ___________:  # Maximum allowed loops
-        if state["precision_loop_count"] > state['loop_max_iter']:  # Maximum allowed loops
+        if state['precision_loop_count'] > state['loop_max_iter']:  # Maximum allowed loops
             return "max_iterations_reached"
         else:
             print(f"---------Precision Score Threshold Not met. Refining Query-----------")  # Debugging
-            # return ____________  # Refine the query
+# Exit the loop
             return "refine_query"  # Refine the query
 
 
 
-
 def max_iterations_reached(state: Dict) -> Dict:
     """Handles the case when the maximum number of iterations is reached."""
     print("---------max_iterations_reached---------")
@@ -406,29 +514,27 @@ def max_iterations_reached(state: Dict) -> Dict:
 
 
 
-from langgraph.graph import END, StateGraph, START
-
 def create_workflow() -> StateGraph:
     """Creates the updated workflow for the AI nutrition agent."""
-    # workflow = StateGraph(_____ )  # Complete the code to define the initial state of the agent
-    workflow = StateGraph(dict)  # Complete the code to define the initial state of the 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)         # Step 1: Expand user query.
+    workflow.add_node("retrieve_context", retrieve_context) # Step 2: Retrieve relevant documents.
+    workflow.add_node("craft_response", craft_response)     # Step 3: Generate a response based on retrieved data.
+    workflow.add_node("score_groundedness", score_groundedness)  # Step 4: Evaluate response grounding.
+    workflow.add_node("refine_response", refine_response)   # Step 5: Improve response if it's weakly grounded.
+    workflow.add_node("check_precision", check_precision)   # Step 6: Evaluate response precision.
+    workflow.add_node("refine_query", refine_query)         # Step 7: Improve query if response lacks precision.
+    workflow.add_node("max_iterations_reached", max_iterations_reached)  # Step 8: Handle max iterations.
+    # workflow.add_node("groundedness_decider",groundedness_decider)
     # 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")
+    # workflow.add_edge("score_groundedness","groundedness_decider")
+
 
     # Conditional edges based on groundedness check
     workflow.add_conditional_edges(
@@ -440,7 +546,6 @@ def create_workflow() -> StateGraph:
             "max_iterations_reached": "max_iterations_reached"  # If max loops reached, exit.
         }
     )
-
     workflow.add_edge("refine_response", "craft_response")  # Refined responses are reprocessed.
 
     # Conditional edges based on precision check
@@ -450,17 +555,19 @@ def create_workflow() -> StateGraph:
         {
             "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.
+            "max_iterations_reached": "max_iterations_reached"  # If max loops reached, exit.
         }
     )
-
     workflow.add_edge("refine_query", "expand_query")  # Refined queries go through expansion again.
 
     workflow.add_edge("max_iterations_reached", END)
+    # Set entry point
+    # workflow.set_entry_point("expand_query")
 
     return workflow
 
 
+
 #=========================== Defining the agentic rag tool ====================#
 WORKFLOW_APP = create_workflow().compile()
 @tool
@@ -477,16 +584,17 @@ def agentic_rag(query: str):
     # Initialize state with necessary parameters
     inputs = {
         "query": query,  # Current user query
-        "expanded_query": "",  # Complete the code to define the expanded version of the query
+        "expanded_query": "",  # 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": 5  # Complete the code to define the maximum number of iterations for loops
+        "response": "",  # AI-generated response
+        "precision_score": 0.0,  # Precision score of the response
+        "groundedness_score": 0.0,  # Groundedness score of the response
+        "groundedness_loop_count": 0,  # Counter for groundedness loops
+        "precision_loop_count": 0,  # Counter for precision loops
+        "feedback": "",
+        "query_feedback":"",
+        "loop_max_iter":2
+
     }
 
     output = WORKFLOW_APP.invoke(inputs)
@@ -530,9 +638,8 @@ class NutritionBot:
         """
 
         # Initialize a memory client to store and retrieve customer interactions
-        self.memory = MemoryClient(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 using the provided credentials
         self.client = ChatOpenAI(
             model_name="gpt-4o-mini",  # Specify the model to use (e.g., GPT-4 optimized version)
             api_key=config.get("API_KEY"),  # API key for authentication
@@ -540,6 +647,7 @@ class NutritionBot:
             temperature=0  # Controls randomness in responses; 0 ensures deterministic results
         )
 
+
         # Define tools available to the chatbot, such as web search
         tools = [agentic_rag]
 
@@ -571,7 +679,6 @@ class NutritionBot:
         # Wrap the agent in an executor to manage tool interactions and execution flow
         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.
@@ -602,7 +709,6 @@ class NutritionBot:
             metadata=metadata
         )
 
-
     def get_relevant_history(self, user_id: str, query: str) -> List[Dict]:
         """
         Retrieve past interactions relevant to the current query.
@@ -617,12 +723,9 @@ class NutritionBot:
         return self.memory.search(
             query=query,  # Search for interactions related to the query
             user_id=user_id,  # Restrict search to the specific user
-            # limit=_____  # Complete the code to define the limit for retrieved interactions
-            limit=5  # Complete the code to define the limit for retrieved interactions
-
+            limit=5  # Retrieve up to 5 relevant interactions
         )
 
-
     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.
@@ -700,6 +803,8 @@ def nutrition_disorder_streamlit():
                     "content": f"Welcome, {user_id}! How can I help you with nutrition disorders today?"
                 })
                 st.session_state.login_submitted = True  # Set flag to trigger rerun
+
+        # Trigger rerun outside the form if login was successful
         if st.session_state.get("login_submitted", False):
             st.session_state.pop("login_submitted")
             st.rerun()
@@ -709,11 +814,11 @@ def nutrition_disorder_streamlit():
             with st.chat_message(message["role"]):
                 st.write(message["content"])
 
-        # Chat input with custom placeholder text
-        # user_query = st.chat_input(__________)  # 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 (or 'exit' to end)...")  # Blank #1
+        # Chat input
+        user_query = st.chat_input("Type your question here (or 'exit' to end)...")
 
         if user_query:
+            # Check if user wants to exit
             if user_query.lower() == "exit":
                 st.session_state.chat_history.append({"role": "user", "content": "exit"})
                 with st.chat_message("user"):
@@ -726,41 +831,41 @@ def nutrition_disorder_streamlit():
                 st.rerun()
                 return
 
+            # Add user message to chat history
             st.session_state.chat_history.append({"role": "user", "content": user_query})
             with st.chat_message("user"):
                 st.write(user_query)
 
-
-
-            # Filter input using Llama Guard
-            # filtered_result = __________(user_query)  # Blank #2: Fill in with the function name for filtering input (e.g., filter_input_with_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 = check_input_safety(user_query, llama_guard_client)
-            filtered_result = filtered_result.replace("\n", " ")  # Normalize the result
-
-            # Check if input is safe based on allowed statuses
-            # if filtered_result in [__________, __________, __________]:  # Blanks #3, #4, #5: Fill in with allowed safe statuses (e.g., "safe", "unsafe S7", "unsafe S6")
-            if filtered_result in ["SAFE", "S6", "S7"]:  # Blanks #3, #4, #5: Fill in with allowed safe statuses (e.g., "safe", "unsafe S7", "unsafe S6")
-
-                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)
-
-                    # response = st.session_state.chatbot.__________(st.session_state.user_id, user_query)
-                    response = st.session_state.chatbot.handle_customer_query(st.session_state.user_id, user_query)
-
-                    # 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:
-                    error_msg = f"Sorry, I encountered an error while processing your query. Please try again. Error: {str(e)}"
-                    st.write(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)
-                st.session_state.chat_history.append({"role": "assistant", "content": inappropriate_msg})
+            # Filter input
+            filtered_result = filter_input_with_llama_guard(user_query)
+            # print(filtered_result)
+            st.write(filtered_result)
+
+            # Process through the agent
+            with st.chat_message("assistant"):
+                # if filtered_result in ["safe", "unsafe S7", "unsafe S6"]:
+                if filtered_result in ["SAFE", "S6", "S7"]:
+                    try:
+                        # Initialize chatbot if not already done
+                        if 'chatbot' not in st.session_state:
+                            st.session_state.chatbot = NutritionBot()
+
+                        # Get response from the chatbot
+                        response = st.session_state.chatbot.handle_customer_query(
+                            st.session_state.user_id,
+                            user_query
+                        )
+
+                        st.write(response)
+                        st.session_state.chat_history.append({"role": "assistant", "content": response})
+                    except Exception as e:
+                        error_msg = f"Sorry, I encountered an error while processing your query. Please try again. Error: {str(e)}"
+                        st.write(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)
+                    st.session_state.chat_history.append({"role": "assistant", "content": inappropriate_msg})
 
 if __name__ == "__main__":
     nutrition_disorder_streamlit()