from alayaflow.component.memory import query_vdb_message
from ..state import WorkflowState


def wrap_query_vdb_message(state: WorkflowState) -> WorkflowState:
    messages = state.get("messages", [])
    limit = state.get("limit", 5)
    result = query_vdb_message(messages, limit)
    updated_state = state.copy()
    updated_state["retrieved_docs"] = result.get("vdb_results", [])
    return updated_state
