app.py

######################## WRITE YOUR CODE HERE  #########################

# Import necessary libraries
import os  # Operating system library for file handling
from llama_index.core import Settings
from typing import Dict, List, Tuple, Any  # Python typing for function annotations
from langchain_core.runnables import RunnablePassthrough  # LangChain core library for running pipelines
from langchain.text_splitter import RecursiveCharacterTextSplitter  # For splitting text documents
from langchain_community.document_loaders import PyPDFLoader  # PDF document loader
from langchain_community.vectorstores import FAISS  # FAISS vector store
from langchain.embeddings.openai import OpenAIEmbeddings  # OpenAI embeddings for text vectors
from langchain.prompts import ChatPromptTemplate  # Template for chat prompts
from langchain_core.output_parsers import StrOutputParser  # String output parser
from langgraph.graph import StateGraph, END  # State graph for managing states in LangChain
from pydantic import BaseModel  # Pydantic for data validation
import numpy as np  # Numpy for numerical operations
from typing import TypedDict  # Typing for structured data

from datetime import datetime
from langchain_core.tools import tool
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from groq import Groq

from chromadb.utils.embedding_functions import HuggingFaceEmbeddingFunction
from langchain_groq import ChatGroq
from langchain_community.vectorstores import Chroma 

os.environ["MEM0_HOME"] = "./.mem0"
from mem0 import MemoryClient

groq_api_key = os.environ.get('GROQ_API_KEY')
hf_api_key = os.environ.get('HF_API_KEY')
llama_api_key = os.environ.get('LLAMA_API_KEY')
mem0 = os.environ.get('MEM0_API_KEY')

# Initialize the Huggingface embedding embedding function for Chroma
embedding_function = HuggingFaceEmbeddingFunction(
    api_key=hf_api_key,
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)

#Initialize the Huggingface Sentence transformers Embeddings
from langchain.embeddings import HuggingFaceEmbeddings
embedding_model = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={'device': 'cuda'},
    encode_kwargs={'normalize_embeddings': True}
)

#initialize the Groq llama 3.1 model
llm = ChatGroq(
    api_key=groq_api_key,
    model_name="llama-3.1-8b-instant",  # Groq supports various models like llama3-70b-8192, mixtral-8x7b, etc.
    temperature=0.0  # Set to 0 for deterministic responses
)

# set the LLM and embedding model in the LlamaIndex settings.
Settings.llm = llm
Settings.embedding = embedding_model

#Defines the AgentState class using TypedDict. It represents the state of the AI agent at different stages of the workflow.
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

#Define all te agents states that is to be used.

def expand_query(state: AgentState) -> AgentState:
    """
    Expands the user query to improve retrieval of nutrition disorder-related information using few-shot prompting.

    Args:
        state (Dict): The current state of the workflow, containing the user query.

    Returns:
        Dict: The updated state with the expanded query.
    """
    system_message = '''You are a helpful research assistant that is well versed with crafting searching queries \
            for nutrition disorders that includes Related symptoms and manifestations, Common diagnostic criteria, Related health conditions \
            and many more to get accurate and relevant results.
            Return 3 related search queries based on the user's request seperated by newline.
            Use the feeback if provided to craft the search query'''

    expand_prompt = ChatPromptTemplate.from_messages([
        ("system", system_message),
        ("user", "{query}")
    ])

    chain = expand_prompt | llm | StrOutputParser()
    state['expanded_query'] = chain.invoke({"query": state['query']})
    return state

#Retrieve contex from vector store
# Initialize the Chroma vector store for retrieving documents
vector_store = Chroma(
    collection_name="semantic_chunks",
    persist_directory="./research_db",
    embedding_function=embedding_model
)

# Create a retriever from the vector store
retriever = vector_store.as_retriever(
    search_type='similarity',
    search_kwargs={'k': 3}
)

def retrieve_context(state: AgentState) -> AgentState:
    """
    Retrieves context from the vector store using the expanded or original query.

    Args:
        state (Dict): The current state of the workflow, containing the query and expanded query.

    Returns:
        Dict: The updated state with the retrieved context.
    """
    query = state['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
    state['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
    ]

    print("Extracted context with metadata:", state['context'])  # Debugging: Print the extracted context
    return state

def craft_response(state: Dict) -> Dict:
    """
    Generates a response using the retrieved context, focusing on nutrition disorders.

    Args:
        state (Dict): The current state of the workflow, containing the query and retrieved context.

    Returns:
        Dict: The updated state with the generated response.
    """
    print("---------craft_response---------")
    system_message = '''You are a helpful assistant specialized in identifying and explaining nutrition-related disorders. 
Use the provided context to craft an accurate, informative, and concise response to the user's query. 
Incorporate user feedback if available.'''

    response_prompt = ChatPromptTemplate.from_messages([
        ("system", system_message),
        ("user", "Query: {query}\nContext: {context}\n\nfeedback: {feedback}")
    ])

    chain = response_prompt | llm
    response = chain.invoke({
        "query": state['query'],
        "context": "\n".join([doc["content"] for doc in state['context']]),
        "feedback": state.get('feedback', '') # add feedback to the prompt
    })
    state['response'] = response
    print("intermediate response: ", response)

    return state

def score_groundedness(state: Dict) -> Dict:
    """
    Checks whether the response is grounded in the retrieved context.

    Args:
        state (Dict): The current state of the workflow, containing the response and context.

    Returns:
        Dict: The updated state with the groundedness score.
    """
    print("---------check_groundedness---------")
    system_message = '''You are a helpful evaluator. Given a response and the supporting context, 
assign a groundedness score between 0 and 1 based on how well the response is supported by the context.
- 1.0 means the response is completely supported by the context.
- 0.0 means the response is not grounded at all in the context.
Only return the score as a number.'''

    groundedness_prompt = ChatPromptTemplate.from_messages([
        ("system", system_message),
        ("user", "Context: {context}\nResponse: {response}\n\nGroundedness score:")
    ])

    chain = groundedness_prompt | llm | StrOutputParser()
    groundedness_score = float(chain.invoke({
        "context": "\n".join([doc["content"] for doc in state['context']]),
        "response": state["response"] #
    }))
    print("groundedness_score: ", groundedness_score)
    state['groundedness_loop_count'] += 1
    print("#########Groundedness Incremented###########")
    state['groundedness_score'] = groundedness_score

    return state

def check_precision(state: Dict) -> Dict:
    """
    Checks whether the response precisely addresses the user’s query.

    Args:
        state (Dict): The current state of the workflow, containing the query and response.

    Returns:
        Dict: The updated state with the precision score.
    """
    print("---------check_precision---------")
    system_message = '''You are an expert evaluator. Given a user query and the assistant's response, 
assign a precision score between 0 and 1 based on how directly and accurately the response answers the query.
- A score of 1.0 means the response is completely relevant, focused, and precise.
- A score of 0.0 means the response is unrelated or vague.
Only return the score as a single float.'''

    precision_prompt = ChatPromptTemplate.from_messages([
        ("system", system_message),
        ("user", "Query: {query}\nResponse: {response}\n\nPrecision score:")
    ])

    chain = precision_prompt | llm | StrOutputParser()
    precision_score = float(chain.invoke({
        "query": state['query'],
        "response": state['response']
    }))
    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:
    """
    Suggests improvements for the generated response.

    Args:
        state (Dict): The current state of the workflow, containing the query and response.

    Returns:
        Dict: The updated state with response refinement suggestions.
    """
    print("---------refine_response---------")

    system_message = '''You are an expert assistant helping to refine AI-generated responses. 
Given a user query and a corresponding response, suggest actionable improvements to increase the response's accuracy, completeness, and clarity.
Focus on fixing missing information, correcting inaccuracies, or improving phrasing.
Keep the suggestions concise and practical.'''

    refine_response_prompt = ChatPromptTemplate.from_messages([
        ("system", system_message),
        ("user", "Query: {query}\nResponse: {response}\n\n"
                 "What improvements can be made to enhance accuracy and completeness?")
    ])

    chain = refine_response_prompt | llm | StrOutputParser()

    # Store response suggestions in a structured format
    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:
    """
    Suggests improvements for the expanded query.

    Args:
        state (Dict): The current state of the workflow, containing the query and expanded query.

    Returns:
        Dict: The updated state with query refinement suggestions.
    """
    print("---------refine_query---------")
    system_message = '''You are an expert in query optimization for information retrieval.
Your task is to evaluate the original query and its expanded form, and suggest improvements to enhance clarity, specificity, and search effectiveness.
Focus on refining medical or technical terminology if needed, and ensure the intent is clearly captured in the expanded version.'''

    refine_query_prompt = ChatPromptTemplate.from_messages([
        ("system", system_message),
        ("user", "Original Query: {query}\nExpanded Query: {expanded_query}\n\n"
                 "What improvements can be made for a better search?")
    ])

    chain = refine_query_prompt | llm | StrOutputParser()

    # Store refinement suggestions without modifying the original expanded query
    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.75:  # 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.75:  # 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

def max_iterations_reached(state: AgentState) -> AgentState:
    """Handles the case where max iterations are reached."""
    state['response'] = "We need more context to provide an accurate answer."
    return state

#AI work flow
from langgraph.graph import END, StateGraph, START

def create_workflow() -> StateGraph:
    """Creates the updated workflow for the AI nutrition agent."""
    workflow = StateGraph(AgentState)

    # Add processing nodes
    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.

    # Main flow edges
    workflow.add_edge(START, "expand_query")
    workflow.add_edge("expand_query", "retrieve_context")
    workflow.add_edge("retrieve_context", "craft_response")
    workflow.add_edge("craft_response", "score_groundedness")

    # Conditional edges based on groundedness check
    workflow.add_conditional_edges(
        "score_groundedness",
        should_continue_groundedness,
        {
            "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.
        }
    )
    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,
        {
            "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.
        }
    )

    workflow.add_edge("refine_query", "expand_query")  # Refined queries go through expansion again.
    workflow.add_edge("max_iterations_reached", END)

    return workflow

WORKFLOW_APP = create_workflow().compile()

@tool
def agentic_rag(query: str):
    """
    Runs the RAG-based agent with conversation history for context-aware responses.

    Args:
        query (str): The current user query.

    Returns:
        Dict[str, Any]: The updated state with the generated response and conversation history.
    """
    # Initialize state with necessary parameters
    inputs = {
        "query": query,
        "expanded_query": "",
        "context": [],
        "response": "",
        "precision_score": 0.0,
        "groundedness_score": 0.0,
        "groundedness_loop_count": 0,
        "precision_loop_count": 0,
        "feedback": "",
        "query_feedback": "",
        "loop_max_iter": 3
    }

    output = WORKFLOW_APP.invoke(inputs)

    return output

# llama_api_key = os.environ.get('LLAMA_API_KEY')
llama_guard_client = 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:
        # Create a request to Llama Guard to filter the user input
        response = llama_guard_client.chat.completions.create(
            messages=[{"role": "user", "content": user_input}],
            model=model,
        )
        # Return the filtered input
        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"Error with Llama Guard: {e}")
        return None
    
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=userdata.get("mem0"))  # Complete the code to define the memory client API key
        self.memory = MemoryClient(api_key=mem0)
        # Initialize the Azure OpenAI client using the provided credentials
        # self.client = AzureChatOpenAI(
        #     model_name="_____",  # Specify the model to use (e.g., GPT-4 optimized version)
        #     api_key=config['_____'],  # API key for authentication
        #     azure_endpoint=config['_____'],  # Endpoint URL for Azure OpenAI
        #     api_version=config['_____'],  # API version being used
        #     temperature=_____  # Controls randomness in responses; 0 ensures deterministic results
        # )
        self.client = ChatGroq(
              api_key=llama_api_key,
              model_name="llama-3.1-8b-instant",  # Groq supports various models like llama3-70b-8192, mixtral-8x7b, etc.
              temperature=0.0
        )

        # 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.

        """

        # 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
        ])

        # Create an agent capable of interacting with tools and executing tasks
        agent = create_tool_calling_agent(self.client, tools, prompt)

        # 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.

        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=10  # 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']
    
def nutrition_disorder_agent():
    """
    A conversational agent that answers nutrition disorder-related questions
    using a RAG-based workflow with safety filtering and user session handling.
    """

    print("Welcome to the Nutrition Disorder Specialist Agent!")
    print("You can ask me anything about nutrition disorders, such as symptoms, causes, treatments, and more.")
    print("Type 'exit' to end the conversation.\n")

    chatbot = NutritionBot()  # Initialize chatbot instance
    print("Login by providing customer name")  # This provides a way to initiate a chat as different users.
    user_id = input()  # Get user ID for tracking conversation sessions

    while True:
        # Get user input
        print("How can I help you?")
        user_query = input("You: ")

        # Define the logic for exitting the loop' [if user types in exit]
        if user_query.lower() == "exit":
            print("Agent: Goodbye! Feel free to return if you have more questions.")
            break

        # Filter input through Llama Guard - returns "SAFE" or "UNSAFE"
        filtered_result = filter_input_with_llama_guard(user_query)  # Call function to filter input
        filtered_result = filtered_result.replace("\n", " ")  # Normalize the result
        print(filtered_result)

        if filtered_result.upper() in ["SAFE","S6", "S7"]:  # You need to by pass some cases like "S6" and "S7" so that it can work effectively.
            # Process the user query using the RAG workflow
            try:
                response = chatbot.handle_customer_query(user_id, user_query)  # Call chatbot handler function
                print(f"Agent: {response}\n")

            except Exception as e:
                print("Agent: Sorry, I encountered an error while processing your query. Please try again.")
                print(f"Error: {e}\n")
        else:
            print("Agent: I apologize, but I cannot process that input as it may be inappropriate. Please try again.")


if __name__ == "__main__":
    # Check if we're running in HF Spaces (which sets PORT environment variable)
    if os.environ.get('PORT'):
        import gradio as gr
        
        def process_query(user_name, query):
            try:
                # Initialize chatbot
                chatbot = NutritionBot()
                
                # Filter input
                filtered_result = filter_input_with_llama_guard(query)
                filtered_result = filtered_result.replace("\n", " ")
                
                if filtered_result.upper() in ["SAFE", "S6", "S7"]:
                    # Get response
                    response = chatbot.handle_customer_query(user_name, query)
                    return response
                else:
                    return "I apologize, but I cannot process that input as it may be inappropriate. Please try again."
            except Exception as e:
                return f"Error processing your request: {str(e)}"
        
        # Create Gradio interface
        with gr.Blocks(title="Nutrition Disorder Assistant") as demo:
            gr.Markdown("# Nutrition Disorder Assistant")
            gr.Markdown("Ask questions about nutrition disorders, symptoms, treatments, and more.")
            
            with gr.Row():
                user_name = gr.Textbox(label="Your Name (for session tracking)")
            
            with gr.Row():
                query_input = gr.Textbox(label="Your Question", placeholder="Ask about nutrition disorders...")
                submit_btn = gr.Button("Submit")
            
            with gr.Row():
                output = gr.Textbox(label="Response")
            
            submit_btn.click(fn=process_query, inputs=[user_name, query_input], outputs=output)
        
        # Launch the app
        demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get('PORT', 7860)))
    else:
        # Run the CLI version when executed directly
        nutrition_disorder_agent()