from datetime import datetime

from langchain_core.messages import HumanMessage, AIMessageChunk

from agent.bill import bill_agent
from modes.agent import QueryRequest


async def query_bill_agent(query: QueryRequest):
    try:
        print(query)
        today = datetime.today().strftime("%Y-%m-%d")

        message_content = (
            f"[当前系统日期: {today}]\n\n"
            f"{query.message}"
        )

        message = HumanMessage(content=message_content)

        tool_called = False
        content_generate = False

        # 流式调用Agent
        for chunk, metadata in bill_agent.stream(
                {"messages": [message]},
                stream_mode="messages"
        ):
            # print(chunk)
            if isinstance(chunk, AIMessageChunk):
                if chunk.content:
                    if not content_generate:
                        content_generate = True
                        yield "\n 📋 **分析结果**：\n"
                    yield chunk.content

                for tc in chunk.tool_call_chunks:
                    if tc.get("args"):
                        if not tool_called:
                            tool_called = True
                            yield "🛠️ **开始查询数据**：查询日期:"
                        yield tc["args"]

    except Exception as e:
        print(f"\n[错误]: {str(e)}")
        yield "信息检索失败，请重新输入问题提问"