Upload planning_agent.py

agent/planning_agent.py

@@ -0,0 +1,223 @@
+from langchain_openai import ChatOpenAI
+from langchain.agents import initialize_agent, AgentType
+from langchain.tools import Tool
+from langchain.memory import ConversationBufferMemory, SimpleMemory
+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 agent.order_tracking_agent as order_tracking_agent #new 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
+order_tracking_agent_instance = None  # Global variable
+
+
+def initialize_planning_agent(llm_instance, chat_memory_instance, query_memory_instance):
+    global llm, chat_memory, query_memory, agent, order_tracking_agent_instance
+    
+    llm = llm_instance
+    chat_memory = chat_memory_instance
+    query_memory = query_memory_instance
+    
+    # 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)
+    order_tracking_agent_instance = order_tracking_agent.OrderTrackingAgent() # Initialize Order Tracking Agent
+
+    # composer_agent.initialize_composer_agent(llm, memory)
+    
+    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"
+        ),
+
+        Tool(
+            name="handle_order_tracking_query",
+            func=handle_order_tracking_query,
+            description="Use this to handle order tracking, refunds, or replacement queries based on the order number or tracking number."
+        )
+
+    ]
+    
+    
+    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,
+        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 handle_order_tracking_query(query):
+    # Summarize chat history before handling order tracking queries
+    summarize_chat_history()
+    # Get the original query or context from memory
+    original_query = query_memory.memories.get('original_query', query)
+    response = order_tracking_agent_instance.process_query(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
+        return agent.run(
+            f"Process this user query: {query}"
+        )
+    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()
+