import abc
import dspy
from abc import abstractmethod


class LLMAgent(abc.ABC):
    def __init__(self, name: str, llm: str, temperature: float, logger):
        self.name = name
        if "lambda" in llm:
            self.lm =dspy.LM("openai/" + llm.split("/")[1], temperature=temperature, api_base="https://api.lambdalabs.com/v1", api_key="ADD_KEY_HERE")
        else:
            self.lm = dspy.LM(llm, temperature=temperature)
        self.temperature = temperature
        self.logger = logger
        dspy.configure(lm=self.lm)

    @abstractmethod
    def __call__(self, context: dict):
        pass

    @abstractmethod
    def restore(self):
        pass

class MemoryAgent(abc.ABC):
    def __init__(self, name: str, llm: str, temperature: float, emb: str, logger):
        self.name = name
        self.lm = dspy.LM(llm)
        self.temperature
        self.emb = dspy.Embedder(emb)
        self.logger = logger
        dspy.configure(self.lm)

    @abstractmethod
    def __call__(self):
        pass

    @abstractmethod
    def restore(self):
        pass