from langchain_core.messages import SystemMessage

from alayaflow.component.chat_model import mk_chat_model_deepseek
from ..state import WorkflowState

chat_model = mk_chat_model_deepseek()

def wrap_deepseek_chat(state: WorkflowState) -> WorkflowState:
    messages = state["messages"].copy()
    updated_state = state.copy()

    retrieved_docs = state.get("retrieved_docs", [])
    if retrieved_docs:
        context_text = "\n\n".join([str(doc) for doc in retrieved_docs])
        context_message = SystemMessage(
            content=f"以下是相关的参考资料，请基于这些资料回答用户的问题：\n\n{context_text}"
        )
        messages.insert(0, context_message)

    response = chat_model.invoke(messages)
    updated_state['chat_response'] = response
    return updated_state
