import inference
from agent import Agent, AgentField

CLASS_DESCRIPTION = """
Expert AI with customizable system prompt.
Will get guidelines, documents, and other advanced features in the future.
""".strip()

class ExpertAgent(Agent):
    identifier = "expert"
    description = CLASS_DESCRIPTION
    schema_fields = [
        AgentField('system_prompt', 'string', default="You are an expert AI assistant."),
        AgentField('model', 'model_id', default=None),
        AgentField('temperature', 'number', default=None),
        AgentField('repetition_penalty', 'number', default=None),
    ]

    def __init__(self, app, data):
        Agent.__init__(self, app, data)

    async def run(self, run):
        opts = {}

        if self.config.get("temperature") is not None:
            opts["temperature"] = self.config["temperature"]
        if self.config.get("repetition_penalty") is not None:
            opts["repetition_penalty"] = self.config["repetition_penalty"]

        # 1. Build message list based on input
        messages = inference.build_messages(
            system_prompt=self.config.get("system_prompt", ""),
            instruction=run.input.get("instruction"),
            input=run.input.get("input"),
            messages=run.input.get("messages"),
        )

        # 2. Call inference, passing the run to feed run.stream and tool calls
        await inference.chat_completion(
            self.app,
            model_id=self.config.get("model"),
            messages=messages,
            tools=run.input.get("tools", []),
            opts=opts,
            agent_run=run,
        )

        run.stream.listen(lambda chunk: None, lambda: run.close())