Upload 5 files

agent/composer_agent.py

@@ -0,0 +1,47 @@
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def compose_response(response: str) -> str:
+    """
+    Process and enhance the final response
+    """
+    try:
+        # Remove any system artifacts or unwanted patterns
+        print("*********** in composer agent *************")
+        print("response input received : ", response)
+        response = remove_system_artifacts(response)
+        
+        # Apply standard formatting
+        response = format_response(response)
+        
+        return response
+        
+    except Exception as e:
+        logger.error(f"Error in composition: {str(e)}")
+        return response  # Fallback to original
+
+def remove_system_artifacts(text: str) -> str:
+    """Remove any system artifacts or unwanted patterns"""
+    artifacts = ["Assistant:", "AI:", "Human:", "User:"]
+    cleaned = text
+    for artifact in artifacts:
+        cleaned = cleaned.replace(artifact, "")
+    return cleaned.strip()
+
+def format_response(text: str) -> str:
+    """Apply standard formatting"""
+    # Add proper spacing
+    formatted = text.replace("\n\n\n", "\n\n")
+    
+    # Ensure proper capitalization
+    formatted = ". ".join(s.strip().capitalize() for s in formatted.split(". "))
+    
+    # Ensure proper ending punctuation
+    if formatted and not formatted[-1] in ['.', '!', '?']:
+        formatted += '.'
+        
+    return formatted
+    

agent/generic_agent.py

@@ -0,0 +1,111 @@
+from langchain_openai import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate
+from langchain.memory import ConversationBufferMemory
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables
+llm = None
+memory = None
+prompt = None
+
+system_prompt = """
+Role
+You are a knowledgeable and compassionate customer support chatbot specializing in various
+products available in Amazon product catalogue. Your goal is to provide accurate, detailed 
+and empathetic information in response to the customer queries on various issues, challenges
+faced by customer strictly related to the products available in Amazon catalogue. 
+Your tone is warm, professional, and supportive, ensuring customers feel informed and reassured 
+during every interaction. 
+
+Instructions
+Shipment Tracking: When a customer asks about their shipment, request the tracking number and 
+tell them you will call back in 1 hour and provide the status on customer's callback number.
+Issue Resolution: For issues such as delays, incorrect addresses, or lost shipments, respond with
+empathy. Explain next steps clearly, including any proactive measures taken to resolve or escalate
+the issue.
+Proactive Alerts: Offer customers the option to receive notifications about key updates, such as 
+when shipments reach major checkpints or encounter delays.
+FAQ Handling: Address frequently asked questions about handling products, special packaging 
+requirements, and preferred delivery times with clarity and simplicity.
+Tone and Language: Maintain a professional and caring tone, particularly when discussing delays or
+challenges. Show understanding and reassurance.
+
+Constraints
+Privacy: Never disclose personal information beyond what has been verified and confirmed by the 
+customer. Always ask for consent before discussing details about shipments.
+Conciseness: Ensure responses are clear and detailed, avoiding jargon unless necessary for conext.
+Empathy in Communication: When addressing delays or challenges, prioritize empathy and acknowledge
+the customer's concern. Provide next steps and resasssurance.
+Accuracy: Ensure all information shared with customer are accurate and up-to-date. If the query is
+outside Amazon's products and services, clearly say I do not know.
+Jargon-Free Language: Use simple language to explain logistics terms or processes to customers, 
+particularly when dealing with customer on sensitive matter.
+
+Examples
+
+Greetings
+
+User: "Hi, I am John."
+AI: "Hi John. How can I assist you today?
+
+Issue Resolution for Delayed product Shipment
+
+User: "I am worried about the  delayed Amazon shipment."
+AI: "I undersatnd your concern, and I'm here to help. Let me check the
+status of your shipment. If needed, we'll coordinate with the carrier to ensure
+your product's safety and provide you with updates along the way."
+
+Proactive Update Offer
+
+User: "Can I get updates on my product shipment's address."
+AI: "Absolutely! I can send you notification whenever your product's shipment
+reaches a checkpoint or if there are any major updates. Would you like to set that
+up ?"
+
+Out of conext question 
+
+User: "What is the capital city of Nigeria ?"
+AI: "Sorry, I do not know. I know only about Amazon products. In case you haave any furter 
+qiestions on the products and services of Amazon, I can help you."
+
+Closure 
+
+User: "No Thank you."
+AI: "Thank you for contacting Amazon. Have a nice day!"
+"""
+
+
+def initialize_generic_agent(llm_instance, memory_instance):
+    global llm, memory, prompt
+    llm = llm_instance
+    memory = memory_instance
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", system_prompt),
+        ("human", "{input}")
+    ])
+    logger.info("generic agent initialized successfully")
+
+def process(query):
+    chain = prompt | llm
+    response = chain.invoke({"input": query})
+
+    # Update memory if available
+    if memory:
+        memory.save_context({"input": query}, {"output": response.content})
+    return response.content
+
+def clear_context():
+    """Clear the conversation memory"""
+    try:
+        if memory:
+            memory.clear()
+            logger.info("Conversation context cleared successfully")
+        else:
+            logger.warning("No memory instance available to clear")
+    except Exception as e:
+        logger.error(f"Error clearing context: {str(e)}")
+        raise

agent/planning_agent.py

@@ -0,0 +1,210 @@
+from langchain_openai import ChatOpenAI
+from langchain.agents import initialize_agent, AgentType
+from clearml import Logger, Task
+from langchain.tools import Tool
+from langchain.memory import ConversationBufferMemory, SimpleMemory
+from langchain_community.callbacks import ClearMLCallbackHandler
+from langchain_core.callbacks import StdOutCallbackHandler
+from langchain_community.utilities import SerpAPIWrapper
+import agent.router_agent as router_agent
+import agent.product_review_agent as product_review_agent
+import agent.generic_agent as generic_agent
+import agent.composer_agent as composer_agent
+import logging
+
+# Set httpx (HTTP request) logging to WARNING or ERROR level
+# This will hide the HTTP request logs while keeping agent thoughts visible
+logging.getLogger("httpx").setLevel(logging.WARNING)   # added on 23-Nob
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables
+llm = None
+chat_memory = None
+query_memory = None
+agent = None
+clearml_callback = None
+
+def initialize_planning_agent(llm_instance, chat_memory_instance, query_memory_instance, clearml_instance):
+    global llm, chat_memory, query_memory, agent, clearml_callback
+    
+    llm = llm_instance
+    chat_memory = chat_memory_instance
+    query_memory = query_memory_instance
+    clearml_callback = clearml_instance
+    SERPAPI_API_KEY = "619f2302253fbe56448bcf82565caf2a3263d845944682533f10b09a0d1650e6"
+    
+    # Initialize agents
+    router_agent.initialize_router_agent(llm, chat_memory)
+    product_review_agent.initialize_product_review_agent(llm, chat_memory)
+    generic_agent.initialize_generic_agent(llm, chat_memory)
+    # composer_agent.initialize_composer_agent(llm, memory)
+
+    #ltool = Tool("serpapi", llm=llm, callbacks=clearml_callback)
+    #ltool2 = load_tools("llm-math", llm=llm, callbacks=clearml_callback)
+    
+    tools = [
+        Tool(
+            name="route_query",
+            func=route_query,
+            description="Determine query type. Returns either 'product_review' or 'generic'"
+        ),
+        Tool(
+            name="get_product_info",
+            func=get_product_info,
+            description="Use this to get product-related data such as features, prices, availability, or reviews"
+        ),
+        Tool(
+            name="handle_generic_query",
+            func=handle_generic_query,
+            description="Use this to get response to user queries which are generic and where the retrieval of product details are not required"
+        ),
+        Tool(
+            name="compose_response",
+            func=compose_response,
+            description="Use this to only format the response. After this step, return the formatted response to main.py"
+        )
+    ]
+    
+    system_prompt = """You are an efficient AI planning agent. Follow these rules strictly:
+
+    CRITICAL INSTRUCTION:
+    For simple queries listed below, skip the route_query and directly go to handle_generic_query.
+
+    SIMPLE QUERIES (NEVER use tools):
+    1. Greetings: "hi", "hello", "hey", "good morning", "good evening", "good afternoon"
+    2. Farewells: "bye", "goodbye", "see you", "take care"
+    3. Thank you messages: "thanks", "thank you", "thanks a lot", "appreciate it"
+    4. Simple confirmations: "okay", "yes", "no", "sure", "alright"
+    5. Basic courtesy: "how are you?", "how are you doing?", "what's up?", "what are you doing?"
+    6. Simple acknowledgments: "got it", "understood", "I see"
+
+    FOR ALL OTHER QUERIES:
+    1. Use route_query to determine if query is product_review or generic
+    2. If route_query returns 'generic', use handle_generic_query and STOP
+    3. If route_query returns 'product_review', use get_product_info and STOP
+
+    EXAMPLES:
+    
+    User: "Hi"
+    Thought: This is a Simple greeting, I will use handle_generic_query to get appropriate response
+    Action: handle_generic_query
+    Observation: "Hi! How can I help you today?"
+    Thought: I have got the final answer. I will use compose_responses to format the response.
+    Action: compose_responses
+    Final Answer: "Hi! How can I help you today?"
+
+    User: "I got my package delivered yesterday. It was delivered very late. I want to file a complaint."
+    Thought: This is a generic query that does not require product details. I will use handle_generic_query to get appropriate response.
+    Action: handle_generic_query
+    Action Input: User query: I got my package delivered yesterday. It was delivered very late. I want to file a complaint.
+    Observation: {'intermediate_steps': [], 'output': "I'm sorry to hear about the delay in your package delivery. I understand your frustration and I'm here to assist you with filing a complaint. To better assist you, could you please provide me with the tracking number of your package? Once I have that information, I will look into the matter and ensure that your feedback is addressed appropriately.", 'action': 'Final Answer', 'action_input': "I'm sorry to hear about the delay in your package delivery. I understand your frustration and I'm here to assist you with filing a complaint. To better assist you, could you please provide me with the tracking number of your package? Once I have that information, I will look into the matter and ensure that your feedback is addressed appropriately."}
+    Thought:I have got the final answer. I will use compose_responses to format the response.
+    Action: compose_responses
+    Final Answer: I'm sorry to hear about the delay in your package delivery. I understand your frustration and I'm here to assist you with filing a complaint. To better assist you, could you please provide me with the tracking number of your package? Once I have that information, I will look into the matter and ensure that your feedback is addressed appropriately.
+
+    Remember: For simple queries listed above, respond immediately with Final Answer WITHOUT using tools.
+    """
+
+    agent = initialize_agent(
+        tools,
+        llm,
+        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+        callbacks=clearml_callback,
+        verbose=True,
+        memory=chat_memory,
+        system_message=system_prompt,
+        early_stopping_method="generate",
+        max_iterations=2
+    )
+    logger.info("Planning agent initialized successfully")
+
+def summarize_chat_history():
+    """Summarize the chat history to retain context without overwhelming memory."""
+    try:
+        if chat_memory:
+            # Retrieve the chat history as a list of messages
+            chat_history = chat_memory.buffer  # Assuming this is a list of messages
+            
+            logger.info(f"Chat history type: {type(chat_history)}")
+            
+            if chat_history:
+                # Extract text from AIMessage objects
+                text_messages = [msg.content for msg in chat_history if hasattr(msg, 'content')]
+                
+                logger.info(f"Extracted messages: {text_messages}")
+                
+                # Get the last 5 messages for the summary
+                summary = "\n".join(text_messages[-5:])  # Adjust as needed
+                logger.info(f"Generated summary: {summary}")
+                
+                # Clear the old history
+                chat_memory.clear()  # Clear the old history
+                
+                # If the memory allows appending, do that
+                chat_memory.buffer.append(summary)  # Append the summary
+                # Or if there's a method to set the buffer, do that:
+                # chat_memory.set_memory([summary])  # If such a method exists
+    except Exception as e:
+        logger.error(f"Error summarizing chat history: {str(e)}")
+
+def route_query(query):
+    # Summarize chat history before routing
+    summarize_chat_history()
+    # Get original query from memory if needed
+    original_query = query_memory.memories.get('original_query', query)
+    return router_agent.classify_query(original_query)
+
+def get_product_info(query):
+    # Summarize chat history before retrieving product info
+    summarize_chat_history()
+    # Get original query from memory if needed
+    original_query = query_memory.memories.get('original_query', query)
+    response = product_review_agent.process(original_query)
+
+    return {
+        "intermediate_steps": [],
+        "output": response,
+        "action": "Final Answer",
+        "action_input": response
+    }
+
+def handle_generic_query(query):
+    # Summarize chat history before handling generic queries
+    summarize_chat_history()
+    # Get original query from memory if needed
+    original_query = query_memory.memories.get('original_query', query)
+    response = generic_agent.process(original_query)
+    return {
+        "intermediate_steps": [],
+        "output": response,
+        "action": "Final Answer",
+        "action_input": response
+    }
+
+def compose_response(response):
+    return composer_agent.compose_response(response)
+
+def execute(query):
+    try:
+        # Store original query
+        query_memory.memories['original_query'] = query 
+        response = agent.run(
+            f"Process this user query: {query}")
+        # clearml_callback.flush_tracker(langchain_asset=agent,name="Planning agent", finish=True)
+        #clearml_callback.flush_tracker(langchain_asset=llm,name="Planning agent")
+        return response
+    except Exception as e:
+        logger.error(f"Error in planning agent: {str(e)}")
+        return f"Error in planning agent: {str(e)}"
+
+def clear_context():
+    if chat_memory:
+        chat_memory.clear()
+    if query_memory:
+        query_memory.memories.clear()
+    product_review_agent.clear_context()
+    generic_agent.clear_context()
+    

agent/product_review_agent.py

@@ -0,0 +1,217 @@
+# ***********************************************************************************************
+# Instruction for using the program
+# ***********************************************************************************************
+# Please make sure the embeddings.npy file is available in data folder
+# Please make sure the documents.pkl file is available in data folder
+# Please set the path appropriately inside the program. You will find the below two statements 
+# where you need to mention the correct path name.
+# embedding_path = '/workspaces/IISC_cap_langchain/data/embeddings.npy'
+# documents_path = '/workspaces/IISC_cap_langchain/documents.pkl'
+# ***********************************************************************************************
+
+import openai
+import numpy as np
+import pandas as pd
+from openai import OpenAI
+from langchain_community.chat_models import ChatOpenAI
+from langchain_community.document_loaders import CSVLoader
+from langchain_community.embeddings import OpenAIEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain.prompts import ChatPromptTemplate
+from langchain.schema.output_parser import StrOutputParser
+from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
+from clearml import StorageManager, Dataset
+
+import faiss
+import warnings
+import os
+
+warnings.filterwarnings("ignore")
+import pickle
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables
+llm = None
+chat_memory = None
+# vectorstore = None
+
+def initialize_product_review_agent(llm_instance, memory_instance):
+    """Initialize the product review agent with LLM and memory instances"""
+    global llm, chat_memory
+
+    llm = llm_instance
+    chat_memory = memory_instance
+ 
+
+
+def process(query):
+
+    # Initialize the OpenAIEmbeddings
+    embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
+
+    System_Prompt = """
+    Role and Capabilities:
+You are an AI customer service specialist for Amazon, focusing on the various products available in Amazon. Your primary functions are:
+1. Providing accurate product information including cost, availability, features, top review
+2. Handling delivery-related queries
+3. Addressing product availability
+4. Offering technical support for electronics
+
+Core Instructions:
+1. Product Information:
+   - Provide detailed specifications and features
+   - Compare similar products when relevant
+   - Only discuss products found in the provided context
+   - Highlight key benefits and limitations
+   - Top review from customer
+
+2. Price & Availability:
+   - Quote exact prices from the provided context
+   - Explain any pricing variations or discounts
+   - Provide clear stock availability information
+   - Mention delivery timeframes when available
+
+3. Query Handling:
+   - Address the main query first, then provide additional relevant information
+   - For multi-part questions, structure answers in bullet points
+   - If information is missing from context, explicitly state this
+   - Suggest alternatives when a product is unavailable
+
+Communication Guidelines:
+1. Response Structure:
+   - Start with a direct answer to the query
+   - Provide supporting details and context
+   - End with a clear next step or call to action
+   - Include standard closing: "Thank you for choosing Amazon. Is there anything else I can help you with?"
+
+2. Tone and Style:
+   - Professional yet friendly
+   - Clear and jargon-free language
+   - Empathetic and patient
+   - Concise but comprehensive
+
+Limitations and Restrictions:
+1. Only provide information present in the given context
+2. Clearly state when information is not available
+3. Never share personal or sensitive information
+4. Don't make promises about delivery times unless explicitly stated in context
+
+Error Handling:
+1. Missing Information: "I apologize, but I don't have that [specific information] in my current context. Would you like me to provide related details about [alternative topic]?"
+2. Out of Scope: "While I can't assist with [topic], I'd be happy to help you with electronics or home care products."
+3. Technical Issues: "I apologize for any inconvenience. Could you please rephrase your question or provide more details?"
+
+Response Format:
+1. For product queries:
+   - Product name and model
+   - Price and availability
+   - Key features
+   - Top review
+   - Comparison among similar products (example : cell phone with cell phone)
+   - Recommendations if relevant
+
+2. For service queries:
+   - Current status
+   - Next steps
+   - Timeline (if available)
+   - Contact options
+
+Remember: Always verify information against the provided context before responding. Don't make assumptions or provide speculative information.
+
+    """
+
+    # Get existing chat history from memory
+    chat_history = ""
+    if chat_memory:
+        messages = chat_memory.chat_memory.messages
+        if messages:
+            chat_history = "\nPrevious conversation:\n"
+            for i in range(0, len(messages), 2):
+                if i + 1 < len(messages):
+                    chat_history += f"Human: {messages[i].content}\n"
+                    chat_history += f"Assistant: {messages[i+1].content}\n"
+
+    # Check if embeddings already exist
+    embedding_path = './data/embeddings.npy'
+    documents_path = './documents.pkl'
+
+    # Modify the get_embedding function to use LangChain's OpenAIEmbeddings
+    def get_embedding(text, engine="text-embedding-ada-002"):
+        return embeddings.embed_query(text)
+
+    try:
+        if not os.path.exists(embedding_path):
+            raise FileNotFoundError(f"Embedding file not found at: {embedding_path}")
+        if not os.path.exists(documents_path):
+            raise FileNotFoundError(f"Documents file not found at: {documents_path}")
+    except FileNotFoundError as e:
+        logger.error(str(e))
+        raise
+
+    if os.path.exists(embedding_path) and os.path.exists(documents_path):
+        # Load existing embeddings and documents
+        embeddings_list = np.load(embedding_path)
+        with open(documents_path, 'rb') as f:
+            documents = pickle.load(f)
+
+    dataset = Dataset.create(dataset_name="embeddingsdataset", dataset_project="dataset_embeddings")
+    dataset.add_files(path=documents_path)
+    dataset.upload()
+    dataset.finalize()
+
+    # Create FAISS index with faster search
+    embeddings_np = np.array(embeddings_list).astype('float32')
+    index=faiss.IndexFlatL2(len(embeddings_list[0]))
+    index.add(embeddings_np)
+
+    query_embedding = get_embedding(query)
+    query_embedding_np = np.array([query_embedding]).astype('float32')
+
+    _, indices = index.search(query_embedding_np, 2)
+    retrieved_docs = [documents[i] for i in indices[0]]
+    context = ' '.join(retrieved_docs)
+    print("context retrieved :", context)
+    print('*' * 100)
+
+    # Include chat history in the prompt for context
+    structured_prompt = f"""
+    Context:
+    {context}
+
+    {chat_history}
+    
+    Current Query:
+    {query}
+    """
+
+    print("structured prompt created :", structured_prompt)
+    print('*' * 100)
+    # Create messages for the chat model
+    messages = [
+        {"role": "system", "content": System_Prompt},
+        {"role": "user", "content": structured_prompt}
+    ]
+
+    # For chat completion, you can use LangChain's ChatOpenAI
+    chat_model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.5)
+    response = chat_model.invoke(messages).content
+
+    # Update memory
+    if chat_memory:
+        chat_memory.chat_memory.add_user_message(query)
+        chat_memory.chat_memory.add_ai_message(response)
+
+    logger.info(f"Successfully processed query: {query}")
+    print("response returned by product_review_agent", response)
+    return response
+
+
+def clear_context():
+    """Clear the conversation memory"""
+    if chat_memory:
+        chat_memory.clear()
+        logger.info("Conversation context cleared")

agent/router_agent.py

@@ -0,0 +1,151 @@
+from langchain_openai import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate
+from langchain.memory import ConversationBufferMemory, SimpleMemory
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables
+llm = None
+chat_memory = None
+query_memory = None
+prompt = None
+
+def initialize_router_agent(llm_instance, chat_memory_instance):
+    global llm, chat_memory, prompt
+    llm = llm_instance
+    chat_memory = chat_memory_instance
+
+    system_prompt = """You are an intelligent query classification system for an e-commerce platform.
+    Your role is to accurately categorize incoming customer queries into one of two categories:
+
+    1. product_review: 
+       - Queries about product features, specifications, or capabilities
+       - Questions about product prices and availability
+       - Requests for product reviews or comparisons
+       - Questions about product warranties or guarantees
+       - Inquiries about product shipping or delivery
+       - Questions about product compatibility or dimensions
+       - Requests for recommendations between products
+
+    2. generic:
+       - General customer service inquiries
+       - Account-related questions
+       - Technical support issues not related to specific products
+       - Website navigation help
+       - Payment or billing queries
+       - Return policy questions
+       - Company information requests
+       - Non-product related shipping questions
+       - Any other queries not directly related to specific products
+
+    INSTRUCTIONS:
+    - Analyze the input query carefully
+    - Respond ONLY with either "product_review" or "generic"
+    - Do not include any other text in your response
+    - If unsure, classify as "generic"
+
+    EXAMPLES:
+
+    User: "What are the features of the Samsung Galaxy S21?"
+    Assistant: product_review
+
+    User: "How much does the iPhone 13 Pro Max cost?"
+    Assistant: product_review
+
+    User: "Can you compare the Dell XPS 15 with the MacBook Pro?"
+    Assistant: product_review
+
+    User: "Is the Sony WH-1000XM4 headphone available in black?"
+    Assistant: product_review
+
+    User: "What's the battery life of the iPad Pro?"
+    Assistant: product_review
+
+    User: "I need help resetting my password"
+    Assistant: generic
+
+    User: "Where can I view my order history?"
+    Assistant: generic
+
+    User: "How do I update my shipping address?"
+    Assistant: generic
+
+    User: "What are your return policies?"
+    Assistant: generic
+
+    User: "I haven't received my refund yet"
+    Assistant: generic
+
+    User: "Do you ship internationally?"
+    Assistant: generic
+
+    User: "Can you recommend a good gaming laptop under $1000?"
+    Assistant: product_review
+
+    User: "What's the warranty period for electronics?"
+    Assistant: generic
+
+    User: "Is the Instant Pot dishwasher safe?"
+    Assistant: product_review
+
+    User: "How do I track my order?"
+    Assistant: generic
+    """
+
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", system_prompt),
+        ("human", "{input}")
+    ])
+    logger.info("Router agent initialized successfully")
+
+
+def classify_query(query):
+    try:
+        # Create chain with memory
+        chain = prompt | llm
+
+        # Add query to chat history before classification
+        if chat_memory and hasattr(chat_memory, 'chat_memory'):
+            chat_memory.chat_memory.add_user_message(query)
+
+        # Classify the query
+        response = chain.invoke({"input": query})
+        category = response.content.strip().lower()
+        
+        # Validate category
+        if category not in ["product_review", "generic"]:
+            category = "generic"  # Default fallback
+
+
+        # Add classification result to chat history
+        if chat_memory and hasattr(chat_memory, 'chat_memory'):
+            chat_memory.chat_memory.add_ai_message(f"Query classified as: {category}")
+        
+        logger.info(f"Query: {query}")
+        logger.info(f"Classification: {category}")
+        print("**** in router agent****")
+        print("query :", query)
+        print("category :", category)
+
+        return category
+
+    except Exception as e:
+        print(f"Error in routing: {str(e)}")
+        return "generic"  # Default fallback on error
+
+
+def get_classification_history():
+    """Retrieve classification history from memory"""
+    if chat_memory and hasattr(chat_memory, 'chat_memory'):
+        return chat_memory.chat_memory.messages
+    return []
+
+
+def clear_context():
+    """Clear all memory contexts"""
+    if chat_memory:
+        chat_memory.clear()
+    logger.info("Router agent context cleared")