from alayaflow.component.memory import init_memory, query_message, add_message
from ..state import WorkflowState


def wrap_init_memory(state: WorkflowState) -> WorkflowState:
    user_id = state["user_id"]
    session_id = state["session_id"]
    
    original_result = init_memory(user_id, session_id)
    
    updated_state = state.copy()
    # updated_state["memory_initialized"] = original_result["memory_initialized"]
    
    return updated_state


def wrap_query_message(state: WorkflowState) -> WorkflowState:
    user_id = state["user_id"]
    session_id = state["session_id"]
    messages = state.get("messages", [])
    query_message(user_id, session_id, messages)
    return state.copy()


def wrap_add_message(state: WorkflowState) -> WorkflowState:
    user_id = state["user_id"]
    session_id = state["session_id"]
    messages = state.get("messages", [])
    add_message(user_id, session_id, messages)
    return state.copy()
