webscout 8.2.4__py3-none-any.whl → 8.2.6__py3-none-any.whl

webscout/AIauto.py

@@ -1,3 +1,9 @@
+"""
+This module provides the AUTO provider, which automatically selects and uses
+an available LLM provider from the webscout library that doesn't require
+API keys or cookies.
+"""
+
 from webscout.AIbase import Provider
 from webscout.exceptions import AllProvidersFailure
 from typing import Union, Any, Dict, Generator
@@ -7,8 +13,20 @@ import random
 import inspect
 
 def load_providers():
+    """
+    Dynamically loads all Provider classes from the webscout.Provider package.
+
+    Identifies providers that require special authentication parameters like
+    'api_key', 'cookie_file', or 'cookie_path'.
+
+    Returns:
+        tuple: A tuple containing:
+            - provider_map (dict): A dictionary mapping uppercase provider names to their classes.
+            - api_key_providers (set): A set of uppercase provider names requiring special authentication.
+    """
     provider_map = {}
     api_key_providers = set()
+    cookie_providers = set()
     provider_package = importlib.import_module("webscout.Provider")
     
     for _, module_name, _ in pkgutil.iter_modules(provider_package.__path__):
@@ -18,16 +36,24 @@ def load_providers():
                 attr = getattr(module, attr_name)
                 if isinstance(attr, type) and issubclass(attr, Provider) and attr != Provider:
                     provider_map[attr_name.upper()] = attr
-                    # Check if the provider needs an API key
-                    if 'api_key' in inspect.signature(attr.__init__).parameters:
+                    # Check if the provider needs special parameters
+                    sig = inspect.signature(attr.__init__).parameters
+                    if 'api_key' in sig:
                         api_key_providers.add(attr_name.upper())
+                    if 'cookie_file' in sig or 'cookie_path' in sig:
+                        cookie_providers.add(attr_name.upper())
         except Exception:
             pass
-    return provider_map, api_key_providers
+    return provider_map, api_key_providers.union(cookie_providers)
 
 provider_map, api_key_providers = load_providers()
 
 class AUTO(Provider):
+    """
+    An automatic provider that cycles through available free providers
+    until one successfully processes the request. Excludes providers
+    requiring API keys or cookies by default.
+    """
     def __init__(
         self,
         is_conversation: bool = True,
@@ -40,26 +66,51 @@ class AUTO(Provider):
         history_offset: int = 10250,
         act: str = None,
         exclude: list[str] = [],
+        print_provider_info: bool = False,
     ):
-        self.provider = None
-        self.provider_name = None
-        self.is_conversation = is_conversation
-        self.max_tokens = max_tokens
-        self.timeout = timeout
-        self.intro = intro
-        self.filepath = filepath
-        self.update_file = update_file
-        self.proxies = proxies
-        self.history_offset = history_offset
-        self.act = act
-        self.exclude = [e.upper() for e in exclude]
+        """
+        Initializes the AUTO provider.
+
+        Args:
+            is_conversation (bool): Flag for conversational mode. Defaults to True.
+            max_tokens (int): Maximum tokens for the response. Defaults to 600.
+            timeout (int): Request timeout in seconds. Defaults to 30.
+            intro (str, optional): Introductory prompt. Defaults to None.
+            filepath (str, optional): Path for conversation history. Defaults to None.
+            update_file (bool): Whether to update the history file. Defaults to True.
+            proxies (dict): Proxies for requests. Defaults to {}.
+            history_offset (int): History character offset limit. Defaults to 10250.
+            act (str, optional): Awesome prompt key. Defaults to None.
+            exclude (list[str]): List of provider names (uppercase) to exclude. Defaults to [].
+            print_provider_info (bool): Whether to print the name of the successful provider. Defaults to False.
+        """
+        self.provider = None  # type: Provider
+        self.provider_name = None  # type: str
+        self.is_conversation: bool = is_conversation
+        self.max_tokens: int = max_tokens
+        self.timeout: int = timeout
+        self.intro: str = intro
+        self.filepath: str = filepath
+        self.update_file: bool = update_file
+        self.proxies: dict = proxies
+        self.history_offset: int = history_offset
+        self.act: str = act
+        self.exclude: list[str] = [e.upper() for e in exclude]
+        self.print_provider_info: bool = print_provider_info
+
 
     @property
     def last_response(self) -> dict[str, Any]:
+        """
+        Returns the last response dictionary from the successful provider.
+        """
         return self.provider.last_response if self.provider else {}
 
     @property
     def conversation(self) -> object:
+        """
+        Returns the conversation object from the successful provider.
+        """
         return self.provider.conversation if self.provider else None
 
     def ask(
@@ -69,8 +120,21 @@ class AUTO(Provider):
         raw: bool = False,
         optimizer: str = None,
         conversationally: bool = False,
-        run_new_test: bool = False,
     ) -> Union[Dict, Generator]:
+        """
+        Sends the prompt to providers, trying each available free provider
+        in a random order until one succeeds.
+
+        Args:
+            prompt (str): The user's prompt.
+            stream (bool): Whether to stream the response. Defaults to False.
+            raw (bool): Whether to return the raw response format. Defaults to False.
+            optimizer (str, optional): Name of the optimizer to use. Defaults to None.
+            conversationally (bool): Whether to apply optimizer conversationally. Defaults to False.
+
+        Returns:
+            Union[Dict, Generator]: The response dictionary or generator from the successful provider.
+        """
         ask_kwargs = {
             "prompt": prompt,
             "stream": stream,
@@ -105,6 +169,9 @@ class AUTO(Provider):
                     act=self.act,
                 )
                 response = self.provider.ask(**ask_kwargs)
+                # Print provider info if enabled
+                if self.print_provider_info:
+                    print(f"\033[1;34m{self.provider_name}\033[0m\n")
                 return response
             except Exception:
                 continue
@@ -118,26 +185,49 @@ class AUTO(Provider):
         stream: bool = False,
         optimizer: str = None,
         conversationally: bool = False,
-        run_new_test: bool = False,
     ) -> Union[str, Generator[str, None, None]]: 
+        """
+        Provides a simplified chat interface, returning the message string(s).
+
+        Args:
+            prompt (str): The user's prompt.
+            stream (bool): Whether to stream the response. Defaults to False.
+            optimizer (str, optional): Name of the optimizer to use. Defaults to None.
+            conversationally (bool): Whether to apply optimizer conversationally. Defaults to False.
+
+        Returns:
+            Union[str, Generator[str, None, None]]: The response string or a generator yielding
+                                                     response chunks.
+        """
+        # run_new_test is not used in the current implementation, but we keep it for API compatibility
         response = self.ask(
             prompt,
             stream,
             optimizer=optimizer,
             conversationally=conversationally,
-            run_new_test=run_new_test,
         )
         
         if stream:
-            return (self.get_message(chunk) for chunk in response)
+            for chunk in response:
+                yield self.get_message(chunk)
         else:
             return self.get_message(response)
 
     def get_message(self, response: dict) -> str:
+        """
+        Extracts the message text from the provider's response dictionary.
+
+        Args:
+            response (dict): The response dictionary from the ask method.
+
+        Returns:
+            str: The extracted message string.
+        """
         assert self.provider is not None, "Chat with AI first"
         return self.provider.get_message(response)
     
 if __name__ == "__main__":
-    auto = AUTO()
-    response = auto.chat("Hello, how are you?")
-    print(response)
+    auto = AUTO(print_provider_info=True)
+    response = auto.chat("Hello, how are you?", stream=True)
+    for chunk in response:
+        print(chunk, end="", flush=True)