src/test_adaptive_chatbot.py

import os
import asyncio
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage
from src.agents.adaptive_chatbot.flow import adaptive_chatbot_agent
from src.utils.logger import logger

# Load environment variables
load_dotenv()

async def test_adaptive_chatbot_flow():
    """
    Test the adaptive chatbot flow with different scenarios
    """
    print("Testing adaptive chatbot flow...")
    
    # Print the flow diagram
    print("\n=== Flow Diagram ===")
    print(adaptive_chatbot_agent.get_graph().draw_mermaid())
    
    # Create test cases
    test_cases = [
        {
            "name": "New user query with no previous info",
            "state": {
                "user_message": "Tôi muốn tìm hiểu về machine learning",
                "session_id": "",
                "messages_history": [],
                "current_system_prompt": "",
                "user_profile": {}
            }
        },
        {
            "name": "Technical query from user",
            "state": {
                "user_message": "Có thể giải thích về cách hoạt động của transformer model trong NLP không?",
                "session_id": "test-session-1",
                "messages_history": [
                    {"content": "Xin chào", "type": "human"},
                    {"content": "Xin chào! Tôi có thể giúp gì cho bạn?", "type": "ai"}
                ],
                "current_system_prompt": "Bạn là một trợ lý AI hữu ích, trả lời câu hỏi người dùng một cách chính xác và đầy đủ.",
                "user_profile": {
                    "technical_level": "intermediate",
                    "preferred_style": "professional",
                    "interests": ["AI", "machine learning"]
                }
            }
        },
        {
            "name": "Simple, non-technical query",
            "state": {
                "user_message": "Thời tiết thế nào ở Hà Nội?",
                "session_id": "test-session-2",
                "messages_history": [],
                "current_system_prompt": "Bạn là một trợ lý AI thân thiện và hữu ích.",
                "user_profile": {}
            }
        }
    ]
    
    # Run test cases
    for i, test_case in enumerate(test_cases):
        print(f"\n\n=== Test Case {i+1}: {test_case['name']} ===")
        print(f"Input state: {test_case['state']}")
        
        try:
            # Add messages tuple if not present
            if "messages" not in test_case["state"]:
                test_case["state"]["messages"] = [("human", test_case["state"]["user_message"])]
                
            # Execute the flow with the test state
            result = await adaptive_chatbot_agent.ainvoke(test_case['state'])
            
            # Display the result
            print("\nResult:")
            for key, value in result.items():
                if key == "messages_history":
                    print(f"  {key}: [... {len(value)} messages ...]")
                elif key == "analysis_result":
                    print(f"  {key}: {value}")
                elif key == "user_profile":
                    print(f"  {key}: {value}")
                elif key == "bot_message" and value:
                    print(f"  {key}: {value[:100]}... (truncated)")
                else:
                    print(f"  {key}: {value}")
                    
            print("\nFlow execution successful!")
        except Exception as e:
            print(f"Error executing flow: {e}")
    
    print("\n=== All tests completed ===")

async def test_specific_node():
    """
    Test a specific node in the flow
    """
    from src.agents.adaptive_chatbot.func import analyze_user_request, generate_probing_questions, update_system_prompt
    
    # Test the analyze_user_request node
    print("\n=== Testing analyze_user_request node ===")
    state = {
        "user_message": "Tôi muốn hiểu sâu hơn về machine learning và các thuật toán",
        "session_id": "test-session",
        "messages_history": [],
        "current_system_prompt": "Bạn là một trợ lý AI hữu ích.",
        "user_profile": {},
        "messages": [("human", "Tôi muốn hiểu sâu hơn về machine learning và các thuật toán")]
    }
    
    try:
        analysis_result = await analyze_user_request(state)
        print("analyze_user_request node result:")
        for key, value in analysis_result.items():
            if key == "analysis_result":
                print(f"  {key}: {value}")
            elif key == "prompt_needs_update":
                print(f"  {key}: {value}")
            elif key == "probing_questions_needed":
                print(f"  {key}: {value}")
    except Exception as e:
        print(f"Error in analyze_user_request node: {e}")
    
    # Test the generate_probing_questions node
    print("\n=== Testing generate_probing_questions node ===")
    state = {
        "user_message": "Tôi muốn học machine learning",
        "session_id": "test-session",
        "messages_history": [],
        "current_system_prompt": "Bạn là một trợ lý AI hữu ích.",
        "user_profile": {},
        "analysis_result": {
            "intent": "learn_about_topic",
            "keywords": ["machine learning"],
            "prompt_needs_update": True,
            "probing_questions_needed": True
        },
        "messages": [("human", "Tôi muốn học machine learning")]
    }
    
    try:
        probing_result = await generate_probing_questions(state)
        print("generate_probing_questions node result:")
        if "probing_questions" in probing_result:
            print(f"  probing_questions: {probing_result['probing_questions']}")
    except Exception as e:
        print(f"Error in generate_probing_questions node: {e}")
        
    # Test the update_system_prompt node
    print("\n=== Testing update_system_prompt node ===")
    state = {
        "user_message": "Tôi muốn học machine learning",
        "session_id": "test-session",
        "messages_history": [],
        "current_system_prompt": "Bạn là một trợ lý AI hữu ích.",
        "user_profile": {
            "technical_level": "beginner",
            "preferred_style": "friendly",
            "interests": ["AI", "machine learning"]
        },
        "analysis_result": {
            "intent": "learn_about_topic",
            "keywords": ["machine learning"],
            "prompt_needs_update": True,
            "probing_questions_needed": False
        },
        "messages": [("human", "Tôi muốn học machine learning")]
    }
    
    try:
        prompt_result = await update_system_prompt(state)
        print("update_system_prompt node result:")
        if "updated_system_prompt" in prompt_result:
            print(f"  updated_system_prompt: {prompt_result['updated_system_prompt']}")
    except Exception as e:
        print(f"Error in update_system_prompt node: {e}")

if __name__ == "__main__":
    # You can choose which test to run:
    # 1. Test the entire flow
    # 2. Test specific nodes
    
    # Run the async tests
    print("Choose a test to run:")
    print("1. Test the entire flow")
    print("2. Test specific nodes")
    
    choice = input("Enter your choice (1 or 2): ")
    
    if choice == "1":
        asyncio.run(test_adaptive_chatbot_flow())
    elif choice == "2":
        asyncio.run(test_specific_node())
    else:
        print("Invalid choice. Please run again with a valid option.")