from agent.callbacks import LoggingCallback
import logging

agent = ComputerAgent(
    model="anthropic/claude-3-5-sonnet-20241022",
    tools=[computer],
    callbacks=[
        LoggingCallback(
            logger=logging.getLogger("cua"), 
            level=logging.INFO
        )
    ]
)

agent = ComputerAgent(
    model="anthropic/claude-3-5-sonnet-20241022",
    tools=[computer],
    verbosity=logging.INFO  # Auto-adds LoggingCallback
)

from agent.callbacks.base import AsyncCallbackHandler
import logging

class CustomLogger(AsyncCallbackHandler):
    def __init__(self, logger_name="agent"):
        self.logger = logging.getLogger(logger_name)
        self.logger.setLevel(logging.INFO)
        
        # Add console handler
        handler = logging.StreamHandler()
        formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)
    
    async def on_run_start(self, kwargs, old_items):
        self.logger.info(f"Agent run started with model: {kwargs.get('model')}")
    
    async def on_computer_call_start(self, item):
        action = item.get('action', {})
        self.logger.info(f"Computer action: {action.get('type')}")
    
    async def on_usage(self, usage):
        cost = usage.get('response_cost', 0)
        self.logger.info(f"API call cost: ${cost:.4f}")
    
    async def on_run_end(self, kwargs, old_items, new_items):
        self.logger.info("Agent run completed")

# Use custom logger
agent = ComputerAgent(
    model="anthropic/claude-3-5-sonnet-20241022",
    tools=[computer],
    callbacks=[CustomLogger("my_agent")]
)