import asyncio

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.core.llms import ChatMessage, MessageRole, MockFunctionCallingLLM
from llama_index.core.memory import Memory


def memory_reply(messages, **kwargs):
    memory_text = "\n".join(str(message.content or "") for message in messages)
    answer = "Ticket 7421 belongs to Maya." if "Maya" in memory_text else "No ticket owner found."
    return ChatMessage(role=MessageRole.ASSISTANT, content=answer)


async def main():
    memory = Memory.from_defaults(session_id="support-ticket-7421", token_limit=1000)
    await memory.aput_messages(
        [
            ChatMessage(role=MessageRole.USER, content="Remember that ticket 7421 belongs to Maya."),
            ChatMessage(role=MessageRole.ASSISTANT, content="Ticket 7421 belongs to Maya."),
        ]
    )

    llm = MockFunctionCallingLLM(response_generator=memory_reply, is_chat_model=True)
    agent = FunctionAgent(
        tools=[],
        llm=llm,
        system_prompt="Answer from chat history when possible.",
    )

    response = await agent.run("Who owns ticket 7421?", memory=memory)
    print(f"agent response: {response}")

    print("stored memory:")
    for message in await memory.aget_all():
        print(f"{message.role.value} | {message.content}")


asyncio.run(main())