from steamship.agents.llms.openai import OpenAI
from steamship.agents.react import ReACTAgent
from steamship.agents.schema.message_selectors import MessageWindowMessageSelector
from steamship.agents.service.agent_service import AgentService
from steamship.agents.tools.image_generation import DalleTool
from steamship.agents.tools.search import SearchTool
from steamship.utils.repl import AgentREPL


class MyAssistant(AgentService):
    """MyAssistant is an example AgentService that exposes a single test endpoint
    for trying out Agent-based invocations. It is configured with two simple Tools
    to provide an overview of the types of tasks it can accomplish (here, search
    and image generation)."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_default_agent(
            ReACTAgent(
                tools=[
                    SearchTool(),
                    DalleTool(),
                ],
                llm=OpenAI(
                    self.client, temperature=0, model_name="gpt-4-0314"
                ),  # MUST PIN to older model
                message_selector=MessageWindowMessageSelector(k=2),
            )
        )


if __name__ == "__main__":
    # AgentREPL provides a mechanism for local execution of an AgentService method.
    # This is used for simplified debugging as agents and tools are developed and
    # added.
    AgentREPL(MyAssistant, agent_package_config={}).run()
