KaniumCogs/reginaldCog/llm_clients/llm_client.py

from abc import ABC, abstractmethod
from enum import Enum
from openai import OpenAI
from discord import Message
from reginaldCog.messenger_clients.messenger_client import ClientMessage, DiscordMessageAdapter


class ILLMContent(ABC):
    pass


class OpenAIContent(ILLMContent):
    def __init__(self):
        self._content_items: list[dict[str, str]] = []

    @property
    def content_items(self) -> list[dict]:
        return self._content_items

    @content_items.setter
    def content_items(self, value: list[dict[str, str]]):
        self._content_items = value


class ILLMContentBuilder(ABC):
    pass


class OpenAIContentBuilder(ILLMContentBuilder):
    def __init__(self, content: OpenAIContent):
        self.content = content

    def add_output_text(self, text: str):
        item = {"type": "output_text", "text": text}
        self.content.content_items.append(item)
        return self

    def add_input_text(self, text: str):
        item = {"type": "input_text", "text": text}
        self.content.content_items.append(item)
        return self

    def add_input_image(self, image_url: str):
        item = {"type": "input_image", "image_url": image_url}
        self.content.content_items.append(item)
        return self

    def add_from_dict(self, item: dict):
        self.content.content_items.append(item)
        return self


class ILLMMessage(ABC):
    pass


class OpenAIMessage(ILLMMessage):
    def __init__(self):
        self.content = OpenAIContent()
        self.role = ""

    @property
    def to_dict(self):
        return {"role": self.role, "content": self.content.content_items}


class ILLMMessageBuilder(ABC):
    pass


class OpenAIMessageBuilder(ILLMMessageBuilder):
    def __init__(self, message: OpenAIMessage):
        self.message = message

    def set_role(self, role: str):
        self.message.role = role
        return self

    def set_content(self, content: OpenAIContent):
        self.message.content = content
        return self


class ILLMPrompt(ABC):
    pass


class OpenAIPrompt(ILLMPrompt):
    def __init__(self):
        self.messages = []

    @property
    def to_list(self):
        return [i_message.to_dict for i_message in self.messages]


class ILLMPromptBuilder(ABC):
    @abstractmethod
    def __init__(self, prompt: ILLMPrompt):
        pass

    @abstractmethod
    def add_message(self, message: ILLMMessage):
        pass


class OpenAIPromptBuilder(ILLMPromptBuilder):
    def __init__(self, prompt: OpenAIPrompt):
        self.prompt = prompt

    def add_message(self, message: OpenAIMessage):
        self.prompt.messages.append(message)


class ILLMClient(ABC):
    @abstractmethod
    def get_response(self, prompt: ILLMPrompt):
        pass


class OpenAIClient(ILLMClient):
    content_class = OpenAIContent
    content_builder_class = OpenAIContentBuilder
    message_class = OpenAIMessage
    message_builder_class = OpenAIMessageBuilder
    prompt_class = OpenAIPrompt
    prompt_builder_class = OpenAIPromptBuilder

    def __init__(self):
        self.model = 'gpt-4.1-mini'
        self.client = OpenAI()

    def get_response(self, prompt: OpenAIPrompt):
        response_input = {"model": self.model, "input": prompt.to_list}
        return self.client.responses.create(**response_input)


class IMessageAdapter(ABC):
    @abstractmethod
    def to_message(self) -> ILLMMessage:
        pass


class OpenAIResponseAdapter(IMessageAdapter):
    def __init__(self, response):
        self.response = response
        self.response_output = response.output[0]

    def to_message(self) -> OpenAIMessage:
        content = OpenAIContent()
        content_builder = OpenAIContentBuilder(content)
        message = OpenAIMessage()
        message_builder = OpenAIMessageBuilder(message)

        message_builder.set_role(self.response_output.role)\
                       .set_content(content)

        for i_content_item in self.response_output.content:
            item = i_content_item.to_dict()
            content_builder.add_from_dict(item)

        return message


class LMMClientType(Enum):
    OPENAI = OpenAIClient


class MessengerClientMessageAdapter(IMessageAdapter):
    def __init__(self, message: ClientMessage, llm_client: LMMClientType):
        self.message = message
        self.llm_client = llm_client

    def to_message(self) -> ILLMMessage:
        content = self.llm_client.value.content_class()
        content_builder = self.llm_client.value.content_builder_class(content)
        message = self.llm_client.value.message_class()
        message_builder = self.llm_client.value.message_builder_class(message)

        message_builder.set_role("user")\
                       .set_content(content)

        if self.message:
            content_builder.add_input_text(self.message.content)
        for i_image_url in self.message.image_urls:
            content_builder.add_input_image(i_image_url)

        return message


if __name__ == "__main__":
    pass