import asyncio
from pprint import pprint
from typing import List
from autogen_core.models import UserMessage, AssistantMessage
from autogen_ext.models.replay import ReplayChatCompletionClient
from autogen_ext.models.anthropic import AnthropicChatCompletionClient
from autogen_agentchat.agents import AssistantAgent
from autogen_core import CancellationToken 
from autogen_core.model_context import BufferedChatCompletionContext
from autogen_contextplus.conditions import MaxMessageCondition
from autogen_contextplus import ContextPlusChatCompletionContext
from autogen_core.models import LLMMessage

# Buffer summary function to limit messages

def buffered_summary(messages: List[LLMMessage], non_summarized_messages: List[LLMMessage]) -> List[LLMMessage]:
    if len(messages) > 3:
        return messages[-3:]
    return messages

async def main():
    client = ReplayChatCompletionClient(chat_completions=["paris", "seoul", "paris", "seoul"])
    context = ContextPlusChatCompletionContext(modifier_func=buffered_summary, modifier_condition=MaxMessageCondition(max_messages=2))
    agent = AssistantAgent(
        "helper",
        model_client=client,
        system_message="You are a helpful agent",
        model_context=context
    )
    await agent.run(task="What is the capital of France?")
    res = await context.get_messages()
    pprint(res)

if __name__ == "__main__":
    asyncio.run(main())