from openai import OpenAI

from trim.utils import TrimOpenAIResponses, TrimOpenAICompletions


class TrimClient(OpenAI):
    def __init__(self, api_key: str = None, **kwargs):
        super().__init__(api_key=api_key, **kwargs)
        # Wrap the original responses
        self.responses.create = TrimOpenAIResponses(self.responses.create)
        self.chat.completions.create = TrimOpenAICompletions(self.chat.completions.create)
