from dataclasses import dataclass, field
from agentS.agentscollaboration import AgentsManager
from agentS.consts import MODEL_NAME
from browsergym.experiments import AbstractAgentArgs

from demo_agent.agents.legacy import Flags


@dataclass
class AgentSArgs(AbstractAgentArgs):
    model_name: str = MODEL_NAME
    llm_type: str = ""
    sync: bool = True
    architecture: str = 'general'
    flags: Flags = field(default_factory=lambda: Flags())
    max_retry: int = 4
    tools = None
    env_policies: list = field(default_factory=list)

    def make_agent(self):
        return AgentsManager(
            llm_name=self.model_name,
            llm_type=self.llm_type,
            sync=self.sync,
            architecture=self.architecture,
            flags=self.flags,
            max_retry=self.max_retry,
            env_policies=self.env_policies,
        )
