webscout 6.2b0__py3-none-any.whl → 6.4__py3-none-any.whl

webscout/AIbase.py

@@ -1,240 +1,113 @@
-from abc import ABC
-from abc import abstractmethod
-from typing import AsyncGenerator, List, Union, Generator
-
-class Provider(ABC):
-    """Base class for text-based AI providers"""
-
-    @abstractmethod
-    def ask(
-        self,
-        prompt: str,
-        stream: bool = False,
-        raw: bool = False,
-        optimizer: str = None,
-        conversationally: bool = False,
-    ) -> dict:
-        """Chat with AI
-
-        Args:
-            prompt (str): Prompt to be sent
-            stream (bool, optional): Flag for streaming response. Defaults to False.
-            raw (bool, optional): Stream back raw response as received
-            optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
-            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
-        Returns:
-           dict : {}
-        ```json
-        {
-            "completion": "\nNext: domestic cat breeds with short hair >>",
-            "stop_reason": null,
-            "truncated": false,
-            "stop": null,
-            "model": "llama-2-13b-chat",
-            "log_id": "cmpl-3kYiYxSNDvgMShSzFooz6t",
-            "exception": null
-        }
-        ```
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-    @abstractmethod
-    def chat(
-        self,
-        prompt: str,
-        stream: bool = False,
-        optimizer: str = None,
-        conversationally: bool = False,
-    ) -> str:
-        """Generate response `str`
-        Args:
-            prompt (str): Prompt to be sent
-            stream (bool, optional): Flag for streaming response. Defaults to False.
-            optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
-            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
-        Returns:
-            str: Response generated
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-    @abstractmethod
-    def get_message(self, response: dict) -> str:
-        """Retrieves message only from response
-
-        Args:
-            response (dict): Response generated by `self.ask`
-
-        Returns:
-            str: Message extracted
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-class AsyncProvider(ABC):
-    """Asynchronous base class for text-based AI providers"""
-
-    @abstractmethod
-    async def ask(
-        self,
-        prompt: str,
-        stream: bool = False,
-        raw: bool = False,
-        optimizer: str = None,
-        conversationally: bool = False,
-    ) -> dict:
-        """Asynchronously chat with AI
-
-        Args:
-            prompt (str): Prompt to be sent
-            stream (bool, optional): Flag for streaming response. Defaults to False.
-            raw (bool, optional): Stream back raw response as received
-            optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
-            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
-        Returns:
-           dict : {}
-        ```json
-        {
-            "completion": "\nNext: domestic cat breeds with short hair >>",
-            "stop_reason": null,
-            "truncated": false,
-            "stop": null,
-            "model": "llama-2-13b-chat",
-            "log_id": "cmpl-3kYiYxSNDvgMShSzFooz6t",
-            "exception": null
-        }
-        ```
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-    @abstractmethod
-    async def chat(
-        self,
-        prompt: str,
-        stream: bool = False,
-        optimizer: str = None,
-        conversationally: bool = False,
-    ) -> str:
-        """Asynchronously generate response `str`
-        Args:
-            prompt (str): Prompt to be sent
-            stream (bool, optional): Flag for streaming response. Defaults to False.
-            optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`
-            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
-        Returns:
-            str: Response generated
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-    @abstractmethod
-    async def get_message(self, response: dict) -> str:
-        """Asynchronously retrieves message only from response
-
-        Args:
-            response (dict): Response generated by `self.ask`
-
-        Returns:
-            str: Message extracted
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-class TTSProvider(ABC):
-    """Base class for text-to-speech providers"""
-
-    @abstractmethod
-    def tts(self, text: str) -> Union[bytes, str, Generator[bytes, None, None]]:
-        """
-        Converts text to speech.
-
-        Args:
-            text (str): The text to be converted to speech.
-
-        Returns:
-            Union[bytes, str, Generator[bytes, None, None]]: 
-                - The raw audio bytes if the provider supports returning audio directly.
-                - The filename of the saved audio file if the provider saves the audio to a file.
-                - A generator yielding chunks of audio bytes if the provider supports streaming audio.
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-class AsyncTTSProvider(ABC):
-    """Asynchronous base class for text-to-speech providers"""
-
-    @abstractmethod
-    async def tts(self, text: str) -> Union[bytes, str, AsyncGenerator[bytes, None]]:
-        """
-        Asynchronously converts text to speech.
-
-        Args:
-            text (str): The text to be converted to speech.
-
-        Returns:
-            Union[bytes, str, AsyncGenerator[bytes, None]]: 
-                - The raw audio bytes if the provider supports returning audio directly.
-                - The filename of the saved audio file if the provider saves the audio to a file.
-                - An asynchronous generator yielding chunks of audio bytes if the provider supports streaming audio.
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-class ImageProvider(ABC):
-    """Base class for text-to-image providers"""
-
-    @abstractmethod
-    def generate(self, prompt: str, amount: int = 1) -> List[bytes]:
-        """Generates images from a text prompt.
-
-        Args:
-            prompt (str): The text prompt describing the desired image.
-            amount (int, optional): The number of images to generate. Defaults to 1.
-
-        Returns:
-            List[bytes]: A list of generated images in bytes format.
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-    @abstractmethod
-    def save(self, response: List[bytes], name: str = None, dir: str = None) -> List[str]:
-        """Saves the generated images to files.
-
-        Args:
-            response (List[bytes]): The list of generated images in bytes format.
-            name (str, optional): The base filename for the images. Defaults to None.
-            dir (str, optional): The directory to save the images. Defaults to None.
-
-        Returns:
-            List[str]: A list of the saved filenames.
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-class AsyncImageProvider(ABC):
-    """Base class for asynchronous text-to-image providers"""
-
-    @abstractmethod
-    async def generate(self, prompt: str, amount: int = 1) -> Union[AsyncGenerator[bytes, None], List[bytes]]:
-        """Asynchronously generates images from a text prompt.
-
-        Args:
-            prompt (str): The text prompt describing the desired image.
-            amount (int, optional): The number of images to generate. Defaults to 1.
-
-        Returns:
-            Union[AsyncGenerator[bytes, None], List[bytes]]: 
-                - An asynchronous generator yielding image bytes.
-                - A list of image bytes if not streaming.
-        """
-        raise NotImplementedError("Method needs to be implemented in subclass")
-
-    @abstractmethod
-    async def save(self, response: Union[AsyncGenerator[bytes, None], List[bytes]], 
-                   name: str = None, dir: str = None) -> List[str]:
-        """Asynchronously saves the generated images to files.
-
-        Args:
-            response (Union[AsyncGenerator[bytes, None], List[bytes]]): 
-                - The asynchronous generator yielding images in bytes format (if streaming).
-                - The list of generated images in bytes format (if not streaming).
-            name (str, optional): The base filename for the images. Defaults to None.
-            dir (str, optional): The directory to save the images. Defaults to None.
-
-        Returns:
-            List[str]: A list of the saved filenames.
-        """
+from abc import ABC
+from abc import abstractmethod
+from pathlib import Path
+from typing import AsyncGenerator, List, Union, Generator, Optional
+from typing_extensions import TypeAlias
+
+# Type aliases for better readability
+Response: TypeAlias = dict[str, Union[str, bool, None]]
+ImageData: TypeAlias = Union[bytes, str, Generator[bytes, None, None]]
+AsyncImageData: TypeAlias = Union[bytes, str, AsyncGenerator[bytes, None]]
+
+class AIProviderError(Exception):
+    pass
+
+class Provider(ABC):
+
+    @abstractmethod
+    def ask(
+        self,
+        prompt: str,
+        stream: bool = False,
+        raw: bool = False,
+        optimizer: Optional[str] = None,
+        conversationally: bool = False,
+    ) -> Response:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+    @abstractmethod
+    def chat(
+        self,
+        prompt: str,
+        stream: bool = False,
+        optimizer: Optional[str] = None,
+        conversationally: bool = False,
+    ) -> str:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+    @abstractmethod
+    def get_message(self, response: Response) -> str:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+class AsyncProvider(ABC):
+
+    @abstractmethod
+    async def ask(
+        self,
+        prompt: str,
+        stream: bool = False,
+        raw: bool = False,
+        optimizer: Optional[str] = None,
+        conversationally: bool = False,
+    ) -> Response:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+    @abstractmethod
+    async def chat(
+        self,
+        prompt: str,
+        stream: bool = False,
+        optimizer: Optional[str] = None,
+        conversationally: bool = False,
+    ) -> str:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+    @abstractmethod
+    async def get_message(self, response: Response) -> str:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+class TTSProvider(ABC):
+
+    @abstractmethod
+    def tts(self, text: str) -> ImageData:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+class AsyncTTSProvider(ABC):
+
+    @abstractmethod
+    async def tts(self, text: str) -> AsyncImageData:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+class ImageProvider(ABC):
+
+    @abstractmethod
+    def generate(self, prompt: str, amount: int = 1) -> List[bytes]:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+    @abstractmethod
+    def save(
+        self, 
+        response: List[bytes], 
+        name: Optional[str] = None, 
+        dir: Optional[Union[str, Path]] = None
+    ) -> List[str]:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+class AsyncImageProvider(ABC):
+
+    @abstractmethod
+    async def generate(
+        self, 
+        prompt: str, 
+        amount: int = 1
+    ) -> Union[AsyncGenerator[bytes, None], List[bytes]]:
+        raise NotImplementedError("Method needs to be implemented in subclass")
+
+    @abstractmethod
+    async def save(
+        self,
+        response: Union[AsyncGenerator[bytes, None], List[bytes]],
+        name: Optional[str] = None,
+        dir: Optional[Union[str, Path]] = None
+    ) -> List[str]:
         raise NotImplementedError("Method needs to be implemented in subclass")