from abc import ABC, abstractmethod
from openai import OpenAI
from typing import Generator

class Ai:

    class AiWrapper(ABC):
        
        @abstractmethod
        def call(self, input: str, **kwargs) -> Generator[str, None, None]:
            """
            Any implmentation of this method should take an input string and return a string,
            optionally taking keyword arguments that can be used to pass additional information.
            """
            pass


    class OpenAIWrapper(AiWrapper):
        # overriding abstract method
        def __init__(self, api_key: str) -> None:
            self.api_key = api_key
            self.client = OpenAI()
            self.client.api_key = api_key
            self.model = "gpt-3.5-turbo"

        def call(self, input:str) -> Generator[str, None, None]: # type: ignore
            client = OpenAI()
            client.api_key = self.api_key
            stream = client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": input}],
                stream=True,
            )
            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    yield chunk.choices[0].delta.content


