main.py

import warnings
import uuid
from agent import create_my_agent, ContextSchema

# Suppress Pydantic serialization warnings
warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")


def parse_response(response):
    """Parse agent response and extract clean message with tools used"""
    result = {"message": "", "tools_used": []}

    if isinstance(response, dict):
        messages = response.get("messages", [])

        # Extract tools used from messages
        for msg in messages:
            if hasattr(msg, "type") and msg.type == "tool":
                tool_name = getattr(msg, "name", "unknown")
                if tool_name not in result["tools_used"]:
                    result["tools_used"].append(tool_name)
            elif isinstance(msg, dict) and msg.get("type") == "tool":
                tool_name = msg.get("name", "unknown")
                if tool_name not in result["tools_used"]:
                    result["tools_used"].append(tool_name)

        # Get final message
        if messages:
            last_message = messages[-1]
            if hasattr(last_message, "content"):
                result["message"] = last_message.content
            elif isinstance(last_message, dict):
                result["message"] = last_message.get("content", str(last_message))
        else:
            result["message"] = str(response)
    else:
        result["message"] = str(response)

    return result


def main():
    print("🤖 LangChain Deep Agent with Memory started!\n")

    agent = create_my_agent()

    # Create a thread ID for this conversation session
    thread_id = str(uuid.uuid4())
    print(f"🧵 Thread ID: {thread_id}\n")

    # Test questions with context and memory
    test_cases = [
        {
            "question": "Hi, my name is Ali",
            "context": ContextSchema(),
        },
        {
            "question": "Get supplier information",
            "context": ContextSchema(supplierId="SUP-001"),
        },
        {
            "question": "What is my name?",  # Test memory recall
            "context": ContextSchema(),
        },
    ]

    for test in test_cases:
        print(f"❓ Question: {test['question']}")
        print(f"📋 Context: {test['context']}")
        try:
            # Pass thread_id in config to enable memory persistence
            response = agent.invoke(
                {"messages": [{"role": "user", "content": test["question"]}]},
                config={"configurable": {"thread_id": thread_id}},
                context=test["context"],
            )
            parsed = parse_response(response)

            if parsed["tools_used"]:
                print(f"🔧 Tools used: {', '.join(parsed['tools_used'])}")

            print(f"✅ Response: {parsed['message']}\n")
        except Exception as e:
            print(f"❌ Error: {e}\n")


if __name__ == "__main__":
    main()