from langchain_core.messages import AIMessageChunk, HumanMessage

from agent.study import question_agent, knowledge_agent
from schemas.agent import QueryQuestionRequest, QueryKnowledgeRequest


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

        # 流式调用Agent
        for chunk, metadata in question_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 "信息检索失败，请重新输入问题提问"


def query_knowledge_agent(query: QueryKnowledgeRequest):
    try:
        user_msg = f"学科：{query.subject}，模块：{query.module}， 名称：{query.name}"

        # 流式调用Agent
        for chunk, metadata in knowledge_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 "信息检索失败，请重新输入问题提问"