from typing import Literal
from dataclasses import dataclass
from utils.groq import query_groq

@dataclass
class ACPMessage:
    sender: str
    receiver: str
    performative: Literal["request", "inform"]
    content: str

    def to_dict(self):
        return {
            "sender": self.sender,
            "receiver": self.receiver,
            "performative": self.performative,
            "content": self.content
        }

class BaseAgent:
    def __init__(self, name: str, role: str, model: str = "llama3-8b-8192"):
        self.name = name
        self.role = role
        self.model = model

    def create_message(self, receiver, performative, content):
        return ACPMessage(
            sender=self.name,
            receiver=receiver,
            performative=performative,
            content=content
        )

    def respond(self, content: str) -> str:
        prompt = f"You are {self.name}, a {self.role}.\nUser: {content}\nYou:"
        reply = query_groq(prompt, model=self.model)
        return reply

    def receive_message(self, message: ACPMessage) -> ACPMessage:
        reply = self.respond(message.content)
        return ACPMessage(
            sender=self.name,
            receiver=message.sender,
            performative="inform",
            content=reply
        )