Upload app.py with huggingface_hub

app.py

@@ -1,611 +1,785 @@
-import streamlit as st
-import json
-import os
-import uuid
 
-import pandas as pd
-from datetime import datetime
-import sqlite3
-import weave
+# 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
+from langchain_core.prompts import ChatPromptTemplate
 
-from langchain.memory import ConversationSummaryBufferMemory
-from openai import OpenAI
-from langchain_openai import ChatOpenAI
-#from langchain.embeddings import OpenAIEmbeddings
-from langchain_community.embeddings import OpenAIEmbeddings
-from langchain_openai import OpenAIEmbeddings
-from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
-from langchain_community.vectorstores import Chroma
-from langchain_core.caches import BaseCache
+# LangChain OpenAI imports
+from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI  # OpenAI embeddings and models
+from langchain.embeddings.openai import OpenAIEmbeddings  # OpenAI embeddings for text vectors
 
+# LlamaParse & LlamaIndex imports
+from llama_parse import LlamaParse  # Document parsing library
+from llama_index.core import Settings, SimpleDirectoryReader  # Core functionalities of the LlamaIndex
 
+# LangGraph import
+from langgraph.graph import StateGraph, END, START  # State graph for managing states in LangChain
 
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_community.utilities.sql_database import SQLDatabase
-from langchain_community.agent_toolkits import create_sql_agent
-from langchain.agents import create_tool_calling_agent, AgentExecutor
-from langchain_core.tools import tool
+# Pydantic import
+from pydantic import BaseModel  # Pydantic for data validation
 
+# Typing imports
+from typing import Dict, List, Tuple, Any, TypedDict  # Python typing for function annotations
 
-from huggingface_hub import CommitScheduler
-from pathlib import Path
+# 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
 
-#BaseCache.register_cache_type("memory", lambda: None)
-#ChatOpenAI.model_rebuild()
 #====================================SETUP=====================================#
 # Fetch secrets from Hugging Face Spaces
-
-model_name = "gpt-4o"
-
-# Extract the OpenAI key and endpoint from the configuration
-
-api_key = os.environ["API_KEY"]
-endpoint = os.environ["OPENAI_API_BASE"]
-
-
-# Define the location of the SQLite database
-db_loc = 'ecomm.db'
-
-# Create a SQLDatabase instance from the SQLite database URI
-db = SQLDatabase.from_uri(f"sqlite:///{db_loc}")
-
-# Retrieve the schema information of the database tables
-database_schema = db.get_table_info()
-
-
-# Let's initiate w&b weave with a project name - this will automatically save all the llm calls made using openai or gemini
-# Make sure to save your w&b api key in secrets as WANDB_API_KEY
-weave.init('ecomm_support')
-# <--------------------------------------------------------- Uncomment to log to WANDB
-
-#=================================Setup Logging=====================================#
-
-log_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
-log_folder = log_file.parent
-
-log_scheduler = CommitScheduler(
-    repo_id="chatbot-logs", #Dataset name where we want to save the logs.
-    repo_type="dataset",
-    folder_path=log_folder,
-    path_in_repo="data",
-    every=5 # Saves data every x minute
+api_key = config.get("API_KEY")
+endpoint = config.get("OPENAI_API_BASE")
+groq_api_key = config.get('LLAMA_API_KEY') # llama_api_key = os.environ['GROQ_API_KEY']
+MEM0_api_key = config.get('mem0') # MEM0_api_key = os.environ['mem0']
+
+# 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.
 
-history_file = Path("history/")/f"data_{uuid.uuid4()}.json"
-history_folder = history_file.parent
-
-history_scheduler = CommitScheduler(
-    repo_id="chatbot-history", #Dataset name where we want to save the logs.
-    repo_type="dataset",
-    folder_path=history_folder,
-    path_in_repo="data",
-    every=5 # Saves data every x minute
-)
-
-#=================================SQL_AGENT=====================================#
-
-# Define the system message for the agent, including instructions and available tables
-system_message = f"""You are a SQLite expert agent designed to interact with a SQLite database.
-Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
-Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 100 results using the LIMIT clause as per SQLite. You can order the results to return the most informative data in the database..
-You can order the results by a relevant column to return the most interesting examples in the database.
-You must query only the columns that are needed to answer the question. Wrap each column name in double quotes (") to denote them as delimited identifiers.
-You have access to tools for interacting with the database.
-Only use the given tools. Only use the information returned by the tools to construct your final answer.
-You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.
-DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.
-You are not allowed to make dummy data.
-If the question does not seem related to the database, just return "I don't know" as the answer.
-Before you execute the query, tell us why you are executing it and what you expect to find briefly.
-Only use the following tables:
-{database_schema}
-"""
-
-# Create a full prompt template for the agent using the system message and placeholders
-full_prompt = ChatPromptTemplate.from_messages(
-    [
-        ("system", system_message),
-        ("human", '{input}'),
-        MessagesPlaceholder("agent_scratchpad")
-    ]
+# Initialize the OpenAI Embeddings
+embedding_model = OpenAIEmbeddings(
+    openai_api_base=endpoint,
+    openai_api_key=api_key,
+    model='text-embedding-ada-002'
 )
 
-# Initialize the ChatOpenAI model with the extracted configuration
+# Initialize the Chat OpenAI model
 llm = ChatOpenAI(
     openai_api_base=endpoint,
     openai_api_key=api_key,
-    model="gpt-4o",
-    streaming=False  # Explicitly disabling streaming
-)
-
-# Create the SQL agent using the ChatOpenAI model, database, and prompt template
-sqlite_agent = create_sql_agent(
-    llm=llm,
-    db=db,
-    prompt=full_prompt,
-    agent_type="openai-tools",
-    agent_executor_kwargs={'handle_parsing_errors': True},
-    max_iterations=5,
-    verbose=True
+    model="gpt-4o", # used gpt4o instead of gpt-4o-mini to get improved results
+    streaming=False
 )
-#### Let's convert the sql agent into a tool that our fin agent can use.
-
-@tool
-def sql_tool(user_input):
+# 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======================#
+
+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
+    feedback: str
+    query_feedback: str
+    groundedness_check: bool
+    loop_max_iter: int
+
+def expand_query(state):
     """
-    Gathers information regarding purchases, transactions, returns, refunds, etc.
-    Executes a SQL query using the sqlite_agent and returns the result.
+    Expands the user query to improve retrieval of nutrition disorder-related information.
+
     Args:
-        user_input (str): a natural language query string explaining what information is required while also providing the necessary details to get the information.
+        state (Dict): The current state of the workflow, containing the user query.
+
     Returns:
-        str: The result of the SQL query execution. If an error occurs, the exception is returned as a string.
+        Dict: The updated state with the expanded query.
     """
-    try:
-        # Invoke the sqlite_agent with the user input (SQL query)
-        response = sqlite_agent.invoke(user_input)
+    print("---------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.
+    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 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.
+    '''
+
+    expand_prompt = ChatPromptTemplate.from_messages([
+        ("system", system_message),
+        ("user", "Expand this query: {query} using the feedback: {query_feedback}")
 
-        # Extract the output from the response
-        prediction = response['output']
+    ])
 
-    except Exception as e:
-        # If an exception occurs, capture the exception message
-        prediction = e
-
-    # Return the result or the exception message
-    return prediction
-
-#=================================== RAG TOOL======================================#
-qna_system_message = """
-You are an assistant to a support agent. Your task is to provide relevant information about the Python package Streamlit.
-User input will include the necessary context for you to answer their questions. This context will begin with the token: ###Context.
-The context contains references to specific portions of documents relevant to the user's query, along with source links.
-The source for a context will begin with the token ###Source
-When crafting your response:
-1. Select only context relevant to answer the question.
-2. User questions will begin with the token: ###Question.
-3. If the context provided doesn't answer the question respond with - "I do not have sufficient information to answer that"
-4. If user asks for product - list all the products that are relevant to his query. If you don't have that product try to cross sell with one of the products we have that is related to what they are interested in.
-You should get information about similar products in the context.
-Please adhere to the following guidelines:
-- Your response should only be about the question asked and nothing else.
-- Answer only using the context provided.
-- Do not mention anything about the context in your final answer.
-- If the answer is not found in the context, it is very very important for you to respond with "I don't know."
-- Always quote the source when you use the context. Cite the relevant source at the end of your response under the section - Source:
-- Do not make up sources. Use the links provided in the sources section of the context and nothing else. You are prohibited from providing other links/sources.
-Here is an example of how to structure your response:
-Answer:
-[Answer]
-Source:
-[Source]
-"""
-
-qna_user_message_template = """
-###Context
-Here are some documents and their source that may be relevant to the question mentioned below.
-{context}
-###Question
-{question}
-"""
-# Load the persisted DB
-persisted_vectordb_location = 'policy_docs/policy_docs'
-#Create a Colelction Name
-collection_name = 'policy_docs'
-
-embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
-# Load the persisted DB
+    chain = expand_prompt | llm | StrOutputParser()
+    expanded_query = chain.invoke({"query": state['query'], "query_feedback":state["query_feedback"]})
+    print("expanded_query", expanded_query)
+    state["expanded_query"] = expanded_query
+    return state
+
+
+# Initialize the Chroma vector store for retrieving documents
 vector_store = Chroma(
-    collection_name=collection_name,
-    persist_directory=persisted_vectordb_location,
+    collection_name="nutritional_hypotheticals",
+    persist_directory="./nutritional_db2",
     embedding_function=embedding_model
-
 )
 
+# Create a retriever from the vector store
 retriever = vector_store.as_retriever(
     search_type='similarity',
-    search_kwargs={'k': 5}
+    search_kwargs={'k': 3}
 )
 
+def retrieve_context(state):
+    """
+    Retrieves context from the vector store using the expanded or original query.
 
-client = OpenAI(
-    api_key=api_key,
-    base_url=endpoint
-)
+    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---------")
+    query = state['query'] # state['_____']  # 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
+    docs = retriever.invoke(query)
+    print("Retrieved documents:", docs)  # Debugging: Print the raw docs object
+
+    # 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.)
+        }
+        for doc in docs
+    ]
+    state['context'] = context # state['_____'] = 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']}")
+    return state
 
-@tool
-def rag(user_input: str) -> str:
 
+def craft_response(state: Dict) -> Dict:
     """
-      Answers questions regarding products, and policies using product descriptions, product policies, and general policies of business using RAG.
-      Args:
-          user_input (str): The input question or query from the user.
-      Returns:
-          response (str): Return the generated response or an error message if an exception occurs.
+    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---------")
+    system_message = # ________________________
+    '''
+    Generates a response to a user query and context provided.
 
-    relevant_document_chunks = retriever.invoke(user_input)
-    context_list = [d.page_content + "\n ###Source: " + d.metadata['source'] + "\n\n " for d in relevant_document_chunks]
+    Parameters:
+    query (str): The user's query and expanded queries based on user's query.
+    context (str): The documents retrieved relevant to the queries.
 
-    context_for_query = ". ".join(context_list)
+    Returns:
+    response (str): The response generated by the model.
 
-    prompt = [
-    {'role':'system', 'content': qna_system_message},
-    {'role': 'user', 'content': qna_user_message_template.format(
-        context=context_for_query,
-        question=user_input
-        )
-    }
-    ]
-    try:
-        response = client.chat.completions.create(
-        model="gpt-4o",
-        messages=prompt
-        )
+    The function performs the following steps:
+    1. Constructs a prompt containing system and user prompts.
+    2. Sends the prompt containing user queries with context provided to the GPT model to generate a response.
+    3. Displays the response.
 
-        prediction = response.choices[0].message.content
-    except Exception as e:
-        prediction = f'Sorry, I encountered the following error: \n {e}'
+    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.'
+    '''
 
+    response_prompt = ChatPromptTemplate.from_messages([
+        ("system", system_message),
+        ("user", "Query: {query}\nContext: {context}\n\nfeedback: {feedback}")
+    ])
 
-    return prediction
+    chain = response_prompt | llm
+    response = chain.invoke({
+        "query": state['query'],
+        "context": "\n".join([doc["content"] for doc in state['context']]),
+        "feedback": state["feedback"] # ________________ # add feedback to the prompt
+    })
+    state['response'] = response
+    print("intermediate response: ", response)
 
+    return state
 
-#=================================== Other TOOLS======================================#
 
-# Function to log actions
-def log_history(email: str,chat_history: list) -> None:
-    # Save the log to the file
-    with history_scheduler.lock:
-        # Open the log file in append mode
-        with history_file.open("a") as f:
-            f.write(json.dumps({
-                "email": email,
-                "chat_history": chat_history,
-                "timestamp": str(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
-            }))
 
-    #st.write("chat_recorded")
+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---------")
+    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:
+    The task is to judge the extent to which the metric is followed by the answer.
+    1 - The metric is not followed at all
+    2 - The metric is followed only to a limited extent
+    3 - The metric is followed to a good extent
+    4 - The metric is followed mostly
+    5 - The metric is followed completely
+
+    The answer should be derived only from the information presented in the context
+
+    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.
+    '''
+
+    groundedness_prompt = ChatPromptTemplate.from_messages([
+        ("system", system_message),
+        ("user", "Context: {context}\nResponse: {response}\n\nGroundedness score:")
+    ])
 
-def log_action(customer_id: str,task: str, details: str) -> None:
-    # Save the log to the file
-    with log_scheduler.lock:
-        # Open the log file in append mode
-        with log_file.open("a") as f:
-            f.write(json.dumps({
-                "customer_id": customer_id,
-                "task": task,
-                "details": details
-            }))
+    chain = groundedness_prompt | llm | StrOutputParser()
+    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
+    }))
+    print("groundedness_score: ", groundedness_score)
+    state['groundedness_loop_count'] += 1
+    print("#########Groundedness Incremented###########")
+    state['groundedness_score'] = groundedness_score
 
+    return state
 
-@tool
-def register_feedback(intent, customer_id, feedback, rating):
+def check_precision(state: Dict) -> Dict:
     """
-    Logs customer feedback into the feedback log.
+    Checks whether the response precisely addresses the user’s query.
+
     Args:
-        intent (str): The category of the support query (e.g., "cancel_order", "get_refund").
-        customer_id (int): The unique ID of the customer.
-        feedback (str): The feedback provided by the customer.
-        rating(int): The rating provided by the customer out of 5
+        state (Dict): The current state of the workflow, containing the query and response.
+
     Returns:
-        str: Success message.
+        Dict: The updated state with the precision score.
     """
-    details = {
-        "intent": intent,
-        "customer_id": customer_id,
-        "feedback": feedback,
-        "rating": rating
-    }
-    log_action(customer_id,"register_feedback", details)
-    #print("register_feedback success")
-    #return "Feedback registered successfully!"
+    print("---------check_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.
+    Output your result as a float number between 0 and 1 
+    Give verdict as a scaled numeric value of type float between 0 and 1, such that
+    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:")
+    ])
 
-@tool
-def defer_to_human(customer_id, query, intent, reason):
+    chain = precision_prompt | llm | StrOutputParser() # _____________ | llm | StrOutputParser() # Complete the code to define the chain of processing
+    precision_score = float(chain.invoke({
+        "query": state['query'],
+        "response": state['response'] # ______________ # Complete the code to access the response from the state
+    }))
+    state['precision_score'] = precision_score
+    print("precision_score:", precision_score)
+    state['precision_loop_count'] +=1
+    print("#########Precision Incremented###########")
+    return state
+
+def refine_response(state: Dict) -> Dict:
     """
-    Logs customer details and the reason for deferring to a human agent.
+    Suggests improvements for the generated response.
+
     Args:
-        customer_id (int): The unique ID of the customer whose query is being deferred.
-        query (str): The customer's query or issue that needs human intervention.
-        reason (str): The reason why the query cannot be resolved by the chatbot.
+        state (Dict): The current state of the workflow, containing the query and response.
+
     Returns:
-        str: Success message indicating the deferral was logged.
+        Dict: The updated state with response refinement suggestions.
     """
+    print("---------refine_response---------")
 
-    details = {
-        "customer_id": customer_id,
-        "query": query,
-        "reason": reason,
-        "intent": intent
-    }
+    system_message = # ________________________
+    '''
+    Since the last response failded 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.
+    '''
 
-    log_action(customer_id,"defer_to_human", details)
-    #return "Case deferred to human agent and logged successfully!"
+    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()
 
-@tool
-def days_since(delivered_date: str) ->str:
+    # Store response suggestions in a structured format
+    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
+    return state
+
+
+
+def refine_query(state: Dict) -> Dict:
     """
-    Calculates the number of days since the product was delivered. This helps in determining whether the product is within return period or not.
+    Suggests improvements for the expanded query.
+
     Args:
-        delivered_date (str): The date when the product was delivered in the format 'YYYY-MM-DD'.
-    """
-    try:
-        # Convert the delivered_date string to a datetime object
-        delivered_date = datetime.strptime(delivered_date, '%Y-%m-%d')
-        today = datetime.today()
-
-        # Calculate the difference in days
-        days_difference = (today - delivered_date).days
-
-        return str(days_difference)
-    except ValueError as e:
-        return f"Error: {e}"
-
-def build_prompt(df):
-
-    system_message = f"""
-
-        You are an intelligent e-commerce chatbot designed to assist users with pre-order and post-order queries. Your job is to
-
-        Gather necessary information from the user to help them with their query.
-        If at any point you cannot determine the next steps - defer to human. you do not have clearance to go beyond the scope the following flow.
-        Do not provide sql inputs to the sql tool - you only need to ask in natural language what information you need.
-        You are only allowed to provide information relevant to the particular customer and the customer information is provided below. you can provide information of this customer only. Following is the information about the customer from the last 2 weeks:
-
-        {df}
-
-        If this information is not enough to answer question, identify the customer from data above and fetch necessary information usign the sql_tool or rag tool - do not fetch information of other customers.
-        use the details provided in the above file to fetch information from sql tool - like customer id, email and phone. Refrain from asking customers details unless necessary.
-        If customer asks about a product, you should act as a sales representative and help them understand the product as much as possible and provide all the necessary information for them. You should also provide them the link to the product which you can get from the source of the information.
-        If a customer asks a query about a policy, be grounded to the context provided to you. if at any point you don't the right thing to say, politely tell the customer that you are not the right person to answer this and defer it to a human.
-        Any time you defer it to a human, you should tell the customer why you did it in a polite manner.
-        MANDATORY STEP:
-        After helping the customer with their concern,
-        - Ask if the customer needs help with anything else. If they ask for anything from the above list help them and along with that,
-        1. Ask for their feedback and rating out of 5.
-        2. then, Use the `register_feedback` tool to log it.  - you MUST ask customer feedback along with asking customer what else they need help with.
-        3. After receving customer feedback exit the chat by responding with 'Bye'.
-
-        ---
-        ### **Handling Out-of-Scope Queries:**
-        If the user's query, at any point is not covered by the workflows above:
-        - Respond:
-          > "This is beyond my skill. Let me connect you to a customer service agent" and get necessary details from the customer and use the defer_to_human tool.
-        - Get customer feedback and rating out of 5.
-        - After getting feedback, end the conversation by saying 'Bye'.
-        ---
-        ### **IMPORTANT Notes for the Model:**
-        - Always fetch additional required details from the database and do not blindly believe details provided by the customer like customer id, email and phone number. You should get the customer id from the system prompt. Cross check with the database and stay loyal to the database.
-        - Be empathetic to the customer but loyal to the instructions provided to you. Try to deescalate a situation before deferring it to human and defer to human only once.
-        - Always aim to minimize the number of questions asked by retrieving as much information as possible from `sql_tool` and `rag` tool.
-        - Follow the exact workflows for each query category.
-        - You will always confirm the order id even if the customer has only one order before you fetch any details.
-        """
+        state (Dict): The current state of the workflow, containing the query and expanded query.
 
-    #st.write(system_message)
-    prompt = ChatPromptTemplate.from_messages([
+    Returns:
+        Dict: The updated state with query refinement suggestions.
+    """
+    print("---------refine_query---------")
+    system_message = # ________________________
+    '''
+    Since the last response failded 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),
-        ("human", "{input}"),
-        ("placeholder", "{agent_scratchpad}"),
+        ("user", "Original Query: {query}\nExpanded Query: {expanded_query}\n\n"
+                 "What improvements can be made for a better search?")
     ])
 
-    return prompt
+    chain = refine_query_prompt | llm | StrOutputParser()
+
+    # Store refinement suggestions without modifying the original expanded query
+    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
+    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
+    else:
+        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_query" # ____________  # Refine the query
 
 
-#===============================================Streamlit=========================================#
 
 
-def login_page():
-    st.title("Login Page")
+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."""
+    response = "I'm unable to refine the response further. Please provide more context or clarify your question."
+    state['response'] = response
+    return state
 
-    email = st.text_input("Email")
-    password = st.text_input("Password", type="password")
 
-    login_button = st.button("Login")
 
-    if login_button:
-        if authenticate_user(email, password):
-            st.session_state.logged_in = True
-            st.session_state.email = email
-            st.success("Login successful! Redirecting to Chatbot...")
-            st.rerun()
-        else:
-            st.error("Invalid email or password.")
+from langgraph.graph import END, StateGraph, START
 
-def authenticate_user(email, phone):
-    connection = sqlite3.connect("ecomm.db")  # Replace with your .db file path
-    cursor = connection.cursor()
+def create_workflow() -> StateGraph:
+    """Creates the updated workflow for the AI nutrition agent."""
+    workflow = StateGraph(State) # StateGraph(_____ )  # Complete the code to define the initial state of the agent
 
-    query = "SELECT first_name FROM customers WHERE email = ? AND phone = ?"
-    cursor.execute(query, (email, phone))
-    user = cursor.fetchone()
+    # 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
 
-    if user:
-        return True  # Login successful
-    return False  # Login failed
+    # 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")
 
-### Prefetch details
+    # Conditional edges based on groundedness check
+    workflow.add_conditional_edges(
+        "score_groundedness",
+        should_continue_groundedness, # ___________,  # Use the conditional function
+        {
+            "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.
+        }
+    )
 
-def fetch_details(email):
-    try:
+    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
+        {
+            "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.
+        }
+    )
 
-        # Connect to the SQLite database
-        connection = sqlite3.connect("ecomm.db")  # Replace with your .db file path
-        cursor = connection.cursor()
-
-        query = f"""
-        SELECT
-            c.customer_id,
-            c.first_name || ' ' || c.last_name AS customer_name,
-            c.email,
-            c.phone,
-            c.address AS customer_address,
-            o.order_id,
-            o.order_date,
-            o.status AS order_status,
-            o.price AS order_price,
-            p.name AS product_name,
-            p.price AS product_price,
-            i.invoice_date,
-            i.amount AS invoice_amount,
-            i.invoice_url,
-            s.delivery_date,
-            s.shipping_status,
-            s.shipping_address,
-            r.refund_amount,
-            r.refund_status
-        FROM Customers c
-        LEFT JOIN Orders o ON c.customer_id = o.customer_id
-        LEFT JOIN Products p ON o.product_id = p.product_id
-        LEFT JOIN Invoices i ON o.order_id = i.order_id
-        LEFT JOIN Shipping s ON o.order_id = s.order_id
-        LEFT JOIN Refund r ON o.order_id = r.order_id
-        WHERE o.order_date >= datetime('now', '-30 days')
-          AND c.email = ?
-        ORDER BY o.order_date DESC;
-        """
+    workflow.add_edge("refine_query", "expand_query") # __________, ___________)  # Refined queries go through expansion again.
 
-        cursor.execute(query, (email,))
-        columns = [description[0] for description in cursor.description]  # Extract column names
-        results = cursor.fetchall()  # Fetch all rows
-        #st.write(results)
-        # Convert results into a list of dictionaries
-        details = [dict(zip(columns, row)) for row in results]
-        #st.write(details)
-        return str(details).replace("{","/").replace("}","/")
+    workflow.add_edge("max_iterations_reached", END)
 
-    except Exception as e:
-        st.write(f"Error: {e}")
-    finally:
-        # Close the connection
-        if connection:
-            cursor.close()
-            connection.close()
+    return workflow
 
-# Function to process user input and generate a chatbot response
 
-def chatbot_interface():
-    st.title("E-Commerce Chatbot")
 
-    if 'conversation_history' not in st.session_state:
-        st.session_state.conversation_history = [{"role": "assistant", "content": "welcome! I am Raha, how can I help you on this beautiful day?"}]
 
+#=========================== 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.
 
-    details = fetch_details(st.session_state.email)
-    # st.write(details)
-    prompt = build_prompt(details)
-    tools = []
-     #[sql_tool,defer_to_human, rag, register_feedback, days_since]
+    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)
 
-    chatbot = ChatOpenAI(
-    openai_api_base=endpoint,
-    openai_api_key=api_key,
-    model="gpt-4o",
-    streaming=False, # Explicitly disabling streaming
-    temperature=0
-    )
+    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"):
+    """
+    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:
-      st.write("Attempting direct LLM test...")
-      test_response = chatbot.invoke("Hello, can you hear me?")
-      st.success(f"Direct LLM Test OK: Received response.")
-      # Optionally display part of the response if needed for confirmation
-      # st.write(test_response.content[:100])
+        # 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:
-      st.error(f"Direct LLM Test FAILED: {e}")
-      st.error("The basic connection to the LLM endpoint might be failing. Check API Key, Endpoint URL, and Network.")
-      # You might want to stop execution here if the basic test fails
-      st.stop()
+        print(f"Error with Llama Guard: {e}")
+        return None
+
+
+#============================= Adding Memory to the agent using mem0 ===============================#
 
-    agent = create_tool_calling_agent(chatbot, tools, prompt)
-    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+class NutritionBot:
+    def __init__(self):
+        """
+        Initialize the NutritionBot class, setting up memory, the LLM client, tools, and the agent executor.
+        """
 
+        # 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
 
+        # Initialize the OpenAI client using the provided credentials
+        self.client = ChatOpenAI(
+            model_name="gpt-4o",  # Used gpt-4o to get improved results; Specify the model to use (e.g., GPT-4 optimized version)
+            api_key=config.get("API_KEY"),  # API key for authentication
+            endpoint = config.get("OPENAI_API_BASE"),
+            temperature=0  # Controls randomness in responses; 0 ensures deterministic results
+        )
 
-    # Display chat messages from history on app rerun
-    for message in st.session_state.conversation_history:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
+        # Define tools available to the chatbot, such as web search
+        tools = [agentic_rag]
+
+        # Define the system prompt to set the behavior of the chatbot
+        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.
+                          Show genuine empathy for customer concerns and health challenges.
+                          Reference past interactions to provide personalized and consistent advice.
+                          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.
+                          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.
 
-    # React to user input
-    if user_input := st.chat_input("You: ", key="chat_input"):
-        # Display user message in chat message container
-        st.chat_message("user").markdown(user_input)
-        with st.spinner("Processing..."):
+        """
 
-            # Add user message to conversation history
-            st.session_state.conversation_history.append({"role": "user", "content": user_input})
+        # 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
+        ])
 
-            conversation_input = "\n".join(
-                [f"{turn['role'].capitalize()}: {turn['content']}" for turn in st.session_state.conversation_history]
-            )
+        # Create an agent capable of interacting with tools and executing tasks
+        agent = create_tool_calling_agent(self.client, tools, prompt)
 
-            try:
-                # Pass the history to the agent
-                response = agent_executor.invoke({"input": conversation_input})
+        # Wrap the agent in an executor to manage tool interactions and execution flow
+        self.agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
-                # Add the chatbot's response to the history
-                chatbot_response = response['output']
-                st.session_state.conversation_history.append({"role": "assistant", "content": chatbot_response})
-                # Check if the assistant's response contains "exit"
-                if "bye" in chatbot_response.lower():
-                    log_history(st.session_state.email,st.session_state.conversation_history)
 
-                # Display the chatbot's response
-                with st.chat_message("assistant"):
-                    st.markdown(chatbot_response)
-
-            except ValueError as ve:
-                if "No generation chunks were returned" in str(ve):
-                    st.error(f"Agent Error: Failed to get a response from the LLM during agent execution (ValueError: {ve}).")
-                    st.error("This often indicates an issue with the custom API endpoint (compatibility, content filtering, or error) or network problems.")
-                    st.error("Check the Hugging Face Space logs for more details, especially network errors or messages related to content policy.")
-                    # Log the error for debugging
-                    print(f"Caught ValueError: {ve}") # This will print to HF Space logs
-                    # Optionally add the error to chat history for visibility
-                    st.session_state.conversation_history.append({"role": "assistant", "content": f"Sorry, I encountered an internal error (ValueError: No generation chunks). Please check logs or try again later."})
-                    with st.chat_message("assistant"):
-                         st.markdown("Sorry, I encountered an internal error (ValueError: No generation chunks). Please check logs or try again later.")
-
-                else:
-                    # Handle other ValueErrors if necessary
-                    st.error(f"Agent Error: An unexpected ValueError occurred: {ve}")
-                    print(f"Caught other ValueError: {ve}")
-                    st.session_state.conversation_history.append({"role": "assistant", "content": f"Sorry, I encountered an unexpected internal error (ValueError)."})
-                    with st.chat_message("assistant"):
-                         st.markdown("Sorry, I encountered an unexpected internal error (ValueError).")
-
-            except Exception as e:
-                st.error(f"Agent Error: An unexpected error occurred: {e}")
-                st.error("Check Hugging Face Space logs for the full traceback.")
-                # Log the full error and traceback for debugging
-                import traceback
-                print(f"Caught Exception: {e}")
-                print(traceback.format_exc()) # Print full traceback to logs
-                # Check if it looks like a content policy error based on the message
-                if "policy" in str(e).lower() or "safety" in str(e).lower() or "blocked" in str(e).lower():
-                     st.warning("The error message suggests the request might have been blocked by a content policy.")
-                     st.session_state.conversation_history.append({"role": "assistant", "content": f"Sorry, my response might have been blocked by a content policy. Error: {e}"})
-                     with st.chat_message("assistant"):
-                         st.markdown(f"Sorry, my response might have been blocked by a content policy. Error: {e}")
-                else:
-                    st.session_state.conversation_history.append({"role": "assistant", "content": f"Sorry, I encountered an unexpected error: {e}"})
-                    with st.chat_message("assistant"):
-                        st.markdown(f"Sorry, I encountered an unexpected error: {e}")
-
-def main():
-    # Check if the user is logged in
-    if "logged_in" in st.session_state and st.session_state["logged_in"]:
-        # Show chatbot page if logged in
-        chatbot_interface()
+    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.
+        """
+        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
+        )
+
+
+    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=_____  # Complete the code to define the limit for retrieved 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.
+
+        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
+        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
+            context += "---\n"
+
+        # Print context for debugging purposes
+        print("Context: ", context)
+
+        # Prepare a prompt combining past context and the current query
+        prompt = f"""
+        Context:
+        {context}
+
+        Current customer query: {query}
+
+        Provide a helpful response that takes into account any relevant past interactions.
+        """
+
+        # Generate a response using the agent
+        response = self.agent_executor.invoke({"input": prompt})
+
+        # Store the current interaction for future reference
+        self.store_customer_interaction(
+            user_id=user_id,
+            message=query,
+            response=response["output"],
+            metadata={"type": "support_query"}
+        )
+
+        # Return the chatbot's response
+        return response['output']
+
+
+#=====================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.")
+
+    # Initialize session state for chat history and user_id if they don't exist
+    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 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:")
+            submit_button = st.form_submit_button("Login")
+            if submit_button and user_id:
+                st.session_state.user_id = user_id
+                st.session_state.chat_history.append({
+                    "role": "assistant",
+                    "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
+        if st.session_state.get("login_submitted", False):
+            st.session_state.pop("login_submitted")
+            st.rerun()
     else:
-        # Show login page if not logged in
-        login_page()
+        # 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)...")
+        if user_query:
+            if user_query.lower() == "exit":
+                st.session_state.chat_history.append({"role": "user", "content": "exit"})
+                with st.chat_message("user"):
+                    st.write("exit")
+                goodbye_msg = "Goodbye! Feel free to return if you have more questions about nutrition disorders."
+                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.rerun()
+                return
+
+            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 = filter_input_with_llama_guard(user_query) # __________(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
+
+            # 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")
+                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.handle_customer_query(st.session_state.user_id, user_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})
 
 if __name__ == "__main__":
-    main()
+    nutrition_disorder_streamlit()