from langchain_core.messages import AIMessageChunk, HumanMessage

from agent.study import study_agent
from models.agent import QueryRequest


def query_agent(query: QueryRequest):
    try:
        user_msg = f"学习内容：\n{query.message}"

        # 流式调用Agent
        for chunk, metadata in study_agent.stream(
                {"messages": [HumanMessage(content=user_msg)]},
                stream_mode="messages"
        ):
            if isinstance(chunk, AIMessageChunk):
                if isinstance(chunk, AIMessageChunk) and chunk.content:
                    yield chunk.content
    except Exception as e:
        print(f"\n[错误]: {str(e)}")
        yield "信息检索失败，请重新输入问题提问"