webscout 7.0__py3-none-any.whl → 7.2__py3-none-any.whl

webscout/Bard.py

@@ -1,51 +1,145 @@
-import argparse
+#########################################
+# Code Generated by o3-mini-high
+#########################################
 import asyncio
 import json
 import os
 import random
 import re
 import string
-import sys
-from typing import Dict, List, Tuple
+from enum import Enum
+from pathlib import Path
+from datetime import datetime
+from typing import Dict, List, Tuple, Union, Optional
 
 import httpx
-from prompt_toolkit import prompt
-from prompt_toolkit import PromptSession
-from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
-from prompt_toolkit.completion import WordCompleter
-from prompt_toolkit.history import InMemoryHistory
-from prompt_toolkit.key_binding import KeyBindings
-from rich.console import Console
-from rich.markdown import Markdown
+from httpx import AsyncClient, HTTPStatusError
 
+# For image models using validation. Adjust based on organization internal pydantic.
+from pydantic import BaseModel, field_validator
 
-def __create_session() -> PromptSession:
-    return PromptSession(history=InMemoryHistory())
+# Rich is retained for logging within image methods.
+from rich.console import Console
+from rich.markdown import Markdown
 
+console = Console()
+
+#########################################
+# New Enums and functions for endpoints,
+# headers, models, file upload and images.
+#########################################
+
+class Endpoint(Enum):
+    INIT = "https://gemini.google.com/app"
+    GENERATE = "https://gemini.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate"
+    ROTATE_COOKIES = "https://accounts.google.com/RotateCookies"
+    UPLOAD = "https://content-push.googleapis.com/upload"
+
+class Headers(Enum):
+    GEMINI = {
+        "Content-Type": "application/x-www-form-urlencoded;charset=utf-8",
+        "Host": "gemini.google.com",
+        "Origin": "https://gemini.google.com",
+        "Referer": "https://gemini.google.com/",
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
+        "X-Same-Domain": "1",
+    }
+    ROTATE_COOKIES = {
+        "Content-Type": "application/json",
+    }
+    UPLOAD = {
+        "Push-ID": "feeds/mcudyrk2a4khkz",
+    }
+
+class Model(Enum):
+    UNSPECIFIED = ("unspecified", {}, False)
+    G_2_0_FLASH = (
+        "gemini-2.0-flash",
+        {"x-goog-ext-525001261-jspb": '[null,null,null,null,"f299729663a2343f"]'},
+        False,
+    )
+    G_2_0_FLASH_EXP = (
+        "gemini-2.0-flash-exp",
+        {"x-goog-ext-525001261-jspb": '[null,null,null,null,"f299729663a2343f"]'},
+        False,
+    )  # Deprecated, should be removed in the future
+    G_2_0_FLASH_THINKING = (
+        "gemini-2.0-flash-thinking",
+        {"x-goog-ext-525001261-jspb": '[null,null,null,null,"9c17b1863f581b8a"]'},
+        False,
+    )
+    G_2_0_FLASH_THINKING_WITH_APPS = (
+        "gemini-2.0-flash-thinking-with-apps",
+        {"x-goog-ext-525001261-jspb": '[null,null,null,null,"f8f8f5ea629f5d37"]'},
+        False,
+    )
+    G_2_0_EXP_ADVANCED = (
+        "gemini-2.0-exp-advanced",
+        {"x-goog-ext-525001261-jspb": '[null,null,null,null,"b1e46a6037e6aa9f"]'},
+        True,
+    )
+    G_1_5_FLASH = (
+        "gemini-1.5-flash",
+        {"x-goog-ext-525001261-jspb": '[null,null,null,null,"418ab5ea040b5c43"]'},
+        False,
+    )
+    G_1_5_PRO = (
+        "gemini-1.5-pro",
+        {"x-goog-ext-525001261-jspb": '[null,null,null,null,"9d60dfae93c9ff1f"]'},
+        True,
+    )
+    G_1_5_PRO_RESEARCH = (
+        "gemini-1.5-pro-research",
+        {"x-goog-ext-525001261-jspb": '[null,null,null,null,"e5a44cb1dae2b489"]'},
+        True,
+    )
 
-def __create_completer(commands: list, pattern_str: str = "$") -> WordCompleter:
-    return WordCompleter(words=commands, pattern=re.compile(pattern_str))
+    def __init__(self, name: str, header: dict, advanced_only: bool):
+        self.model_name = name
+        self.model_header = header
+        self.advanced_only = advanced_only
 
+    @classmethod
+    def from_name(cls, name: str) -> "Model":
+        for model in cls:
+            if model.model_name == name:
+                return model
+        raise ValueError(
+            f"Unknown model name: {name}. Available models: {', '.join([m.model_name for m in cls])}"
+        )
 
-def __get_input(
-    prompt_sess: PromptSession = None,
-    completer: WordCompleter = None,
-    key_bindings: KeyBindings = None,
-) -> str:
+async def upload_file(file: Union[bytes, str, Path], proxy: Optional[str] = None) -> str:
     """
-    Multiline input function.
+    Upload a file to Google's server and return its identifier.
+
+    Parameters:
+        file: bytes | str | Path
+            File data in bytes, or path to the file to be uploaded.
+        proxy: str, optional
+            Proxy URL.
+
+    Returns:
+        str: Identifier of the uploaded file.
+    Raises:
+        httpx.HTTPStatusError: If the upload request failed.
     """
-    return (
-        prompt_sess.prompt(
-            completer=completer,
-            multiline=True,
-            auto_suggest=AutoSuggestFromHistory(),
-            key_bindings=key_bindings,
+    if not isinstance(file, bytes):
+        with open(file, "rb") as f:
+            file = f.read()
+
+    async with AsyncClient(http2=True, proxies=proxy) as client:
+        response = await client.post(
+            url=Endpoint.UPLOAD.value,
+            headers=Headers.UPLOAD.value,
+            files={"file": file},
+            follow_redirects=True,
         )
-        if prompt_sess
-        else prompt(multiline=True)
-    )
+        response.raise_for_status()
+        return response.text
 
+#########################################
+# Cookie loading and Chatbot classes
+#########################################
 
 def load_cookies(cookie_path: str) -> Tuple[str, str]:
     """Loads cookies from the provided JSON file."""
@@ -62,56 +156,56 @@ def load_cookies(cookie_path: str) -> Tuple[str, str]:
     except StopIteration:
         raise Exception("Required cookies not found in the cookie file.")
 
-
 class Chatbot:
     """
     Synchronous wrapper for the AsyncChatbot class.
     """
-
     def __init__(
         self,
         cookie_path: str,
         proxy: dict = None,
         timeout: int = 20,
+        model: Model = Model.UNSPECIFIED
     ):
         self.loop = asyncio.get_event_loop()
         self.secure_1psid, self.secure_1psidts = load_cookies(cookie_path)
         self.async_chatbot = self.loop.run_until_complete(
-            AsyncChatbot.create(self.secure_1psid, self.secure_1psidts, proxy, timeout),
+            AsyncChatbot.create(self.secure_1psid, self.secure_1psidts, proxy, timeout, model)
         )
 
     def save_conversation(self, file_path: str, conversation_name: str):
         return self.loop.run_until_complete(
-            self.async_chatbot.save_conversation(file_path, conversation_name),
+            self.async_chatbot.save_conversation(file_path, conversation_name)
         )
 
     def load_conversations(self, file_path: str) -> List[Dict]:
         return self.loop.run_until_complete(
-            self.async_chatbot.load_conversations(file_path),
+            self.async_chatbot.load_conversations(file_path)
         )
 
     def load_conversation(self, file_path: str, conversation_name: str) -> bool:
         return self.loop.run_until_complete(
-            self.async_chatbot.load_conversation(file_path, conversation_name),
+            self.async_chatbot.load_conversation(file_path, conversation_name)
         )
 
     def ask(self, message: str) -> dict:
         return self.loop.run_until_complete(self.async_chatbot.ask(message))
 
-
 class AsyncChatbot:
     """
     A class to interact with Google Gemini.
-    Parameters
-        session: str
-            The __Secure_1PSID cookie.
-        session_ts: str
-            The __secure_1psidts cookie.
-        proxy: str
+    Parameters:
+        secure_1psid: str
+            The __Secure-1PSID cookie.
+        secure_1psidts: str
+            The __Secure-1PSIDTS cookie.
+        proxy: dict
+            Http request proxy.
         timeout: int
             Request timeout in seconds.
+        model: Model
+            Selected model for the session.
     """
-
     __slots__ = [
         "headers",
         "_reqid",
@@ -124,6 +218,7 @@ class AsyncChatbot:
         "secure_1psid",
         "session",
         "timeout",
+        "model",
     ]
 
     def __init__(
@@ -132,23 +227,11 @@ class AsyncChatbot:
         secure_1psidts: str,
         proxy: dict = None,
         timeout: int = 20,
+        model: Model = Model.UNSPECIFIED,
     ):
-        """Constructor
-
-        Args:
-            secure_1psid (str): __Secure-1PSID cookie value
-            secure_1psidts (str): __Secure-1PSIDTS cookie value
-            proxy (dict, optional): Http request proxy. Defaults to None.
-            timeout (int, optional): htpp request timeout. Defaults to 20.
-        """
-        headers = {
-            "Content-Type": "application/x-www-form-urlencoded;charset=utf-8",
-            "Host": "gemini.google.com",
-            "Origin": "https://gemini.google.com",
-            "Referer": "https://gemini.google.com/",
-            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
-            "X-Same-Domain": "1",
-        }
+        headers = Headers.GEMINI.value.copy()
+        if model != Model.UNSPECIFIED:
+            headers.update(model.model_header)
         self._reqid = int("".join(random.choices(string.digits, k=4)))
         self.proxy = proxy
         self.conversation_id = ""
@@ -161,6 +244,7 @@ class AsyncChatbot:
         self.session.cookies.set("__Secure-1PSID", secure_1psid)
         self.session.cookies.set("__Secure-1PSIDTS", secure_1psidts)
         self.timeout = timeout
+        self.model = model
 
     @classmethod
     async def create(
@@ -169,25 +253,14 @@ class AsyncChatbot:
         secure_1psidts: str,
         proxy: dict = None,
         timeout: int = 20,
+        model: Model = Model.UNSPECIFIED,
     ) -> "AsyncChatbot":
-        """
-        Async constructor.
-        """
-        instance = cls(secure_1psid, secure_1psidts, proxy, timeout)
+        instance = cls(secure_1psid, secure_1psidts, proxy, timeout, model)
         instance.SNlM0e = await instance.__get_snlm0e()
         return instance
 
     async def save_conversation(self, file_path: str, conversation_name: str) -> None:
-        """
-        Saves conversation to the file
-        :param file_path: file to save (json)
-        :param conversation_name: any name of current conversation (unique one)
-        :return: None
-        """
-        # Load conversations from file
         conversations = await self.load_conversations(file_path)
-
-        # Update existing one
         conversation_exists = False
         for conversation in conversations:
             if conversation["conversation_name"] == conversation_name:
@@ -198,8 +271,6 @@ class AsyncChatbot:
                 conversation["choice_id"] = self.choice_id
                 conversation["SNlM0e"] = self.SNlM0e
                 conversation_exists = True
-
-        # Create conversation object
         if not conversation_exists:
             conversation = {
                 "conversation_name": conversation_name,
@@ -210,25 +281,16 @@ class AsyncChatbot:
                 "SNlM0e": self.SNlM0e,
             }
             conversations.append(conversation)
-
-        # Save to the file
         with open(file_path, "w", encoding="utf-8") as f:
             json.dump(conversations, f, indent=4)
 
     async def load_conversations(self, file_path: str) -> List[Dict]:
-        # Check if file exists
         if not os.path.isfile(file_path):
             return []
         with open(file_path, encoding="utf-8") as f:
             return json.load(f)
 
     async def load_conversation(self, file_path: str, conversation_name: str) -> bool:
-        """
-        Loads a conversation from history file. Returns whether the conversation was found
-        :param file_path: File with conversations (json)
-        :param conversation_name: unique conversation name
-        :return: True if the conversation was found
-        """
         conversations = await self.load_conversations(file_path)
         for conversation in conversations:
             if conversation["conversation_name"] == conversation_name:
@@ -241,46 +303,24 @@ class AsyncChatbot:
         return False
 
     async def __get_snlm0e(self):
-        # Find "SNlM0e":"<ID>"
-        if (
-            not (self.secure_1psid and self.secure_1psidts)
-            or self.secure_1psid[:2] != "g."
-        ):
-            raise Exception(
-                "Enter correct __Secure_1PSID and __Secure_1PSIDTS value. __Secure_1PSID value must start with a g dot (g.). ",
-            )
-        resp = await self.session.get(
-            "https://gemini.google.com/app",
-            timeout=10,
-            follow_redirects=True,
-        )
+        if not (self.secure_1psid and self.secure_1psidts) or self.secure_1psid[:2] != "g.":
+            raise Exception("Enter correct __Secure_1PSID and __Secure_1PSIDTS value. __Secure_1PSID value must start with a g.")
+        resp = await self.session.get(Endpoint.INIT.value, timeout=10, follow_redirects=True)
         if resp.status_code != 200:
-            raise Exception(
-                f"Response code not 200. Response Status is {resp.status_code}",
-            )
-        SNlM0e = re.search(r'"SNlM0e":"(.*?)"', resp.text)
-        if not SNlM0e:
-            raise Exception(
-                "SNlM0e value not found in response. Check __Secure_1PSID value."
-                "\nNOTE : The cookies expire after a short period; ensure you update them as frequent as possible."
-                f" Failed with status {resp.status_code} - {resp.reason_phrase}",
-            )
-        return SNlM0e.group(1)
+            raise Exception(f"Response code not 200. Response Status is {resp.status_code}")
+        snlm0e_match = re.search(r'"SNlM0e":"(.*?)"', resp.text)
+        if not snlm0e_match:
+            raise Exception("SNlM0e value not found in response. Check __Secure_1PSID value."
+                            "\nNOTE: The cookies expire after a short period; ensure you update them frequently."
+                            f" Failed with status {resp.status_code} - {resp.reason_phrase}")
+        return snlm0e_match.group(1)
 
     async def ask(self, message: str) -> dict:
-        """
-        Send a message to Google Gemini and return the response.
-        :param message: The message to send to Google Gemini.
-        :return: A dict containing the response from Google Gemini.
-        """
-        # url params
         params = {
             "bl": "boq_assistant-bard-web-server_20230713.13_p0",
             "_reqid": str(self._reqid),
             "rt": "c",
         }
-
-        # message arr -> data["f.req"]. Message is double json stringified
         message_struct = [
             [message],
             None,
@@ -291,21 +331,24 @@ class AsyncChatbot:
             "at": self.SNlM0e,
         }
         resp = await self.session.post(
-            "https://gemini.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate",
+            Endpoint.GENERATE.value,
             params=params,
             data=data,
             timeout=self.timeout,
         )
-        chat_data = json.loads(resp.content.splitlines()[3])[0][2]
-        if not chat_data:
+        try:
+            chat_data_line = resp.content.splitlines()[3]
+            chat_data = json.loads(chat_data_line)[0][2]
+        except (IndexError, json.JSONDecodeError):
             return {"content": f"Gemini encountered an error: {resp.content}."}
+        if not chat_data:
+            return {"content": f"Gemini returned empty response: {resp.content}."}
         json_chat_data = json.loads(chat_data)
         images = []
         if len(json_chat_data) >= 3:
-            if len(json_chat_data[4][0]) >= 4:
-                if json_chat_data[4][0][4]:
-                    for img in json_chat_data[4][0][4]:
-                        images.append(img[0][0][0])
+            if len(json_chat_data[4][0]) >= 4 and json_chat_data[4][0][4]:
+                for img in json_chat_data[4][0][4]:
+                    images.append(img[0][0][0])
         results = {
             "content": json_chat_data[4][0][1][0],
             "conversation_id": json_chat_data[1][0],
@@ -321,45 +364,184 @@ class AsyncChatbot:
         self._reqid += 100000
         return results
 
+#########################################
+# New Image classes
+#########################################
 
-if __name__ == "__main__":
-    import sys
-    sys.exit(0)
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--session",
-        help="__Secure-1PSID cookie",
-        type=str,
-        required=True,
-    )
-    parser.add_argument(
-        "--session_ts",
-        help="__secure_1psidts cookie.",
-        type=str,
-        required=True,
-    )
-    args = parser.parse_args()
+class Image(BaseModel):
+    """
+    A single image object returned from Gemini.
+    Parameters:
+        url: str
+            URL of the image.
+        title: str, optional
+            Title of the image (default: "[Image]").
+        alt: str, optional
+            Optional description.
+        proxy: str, optional
+            Proxy used when saving the image.
+    """
+    url: str
+    title: str = "[Image]"
+    alt: str = ""
+    proxy: Optional[str] = None
+
+    def __str__(self):
+        return f"{self.title}({self.url}) - {self.alt}"
+
+    def __repr__(self):
+        short_url = self.url if len(self.url) <= 20 else self.url[:8] + "..." + self.url[-12:]
+        return f"Image(title='{self.title}', url='{short_url}', alt='{self.alt}')"
+
+    async def save(
+        self,
+        path: str = "temp",
+        filename: Optional[str] = None,
+        cookies: Optional[dict] = None,
+        verbose: bool = False,
+        skip_invalid_filename: bool = False,
+    ) -> Optional[str]:
+        """
+        Save the image to disk.
+        Parameters:
+            path: str, optional
+                Directory to save the image (default "./temp").
+            filename: str, optional
+                Filename to use; if not provided, inferred from URL.
+            cookies: dict, optional
+                Cookies used for the image request.
+            verbose: bool, optional
+                If True, outputs status messages (default False).
+            skip_invalid_filename: bool, optional
+                If True, skips saving if the filename is invalid.
+        Returns:
+            Absolute path of the saved image if successful; None if skipped.
+        Raises:
+            httpx.HTTPError if the network request fails.
+        """
+        filename = filename or self.url.split("/")[-1].split("?")[0]
+        try:
+            filename = re.search(r"^(.*\.\w+)", filename).group()
+        except AttributeError:
+            if verbose:
+                console.log(f"Invalid filename: {filename}")
+            if skip_invalid_filename:
+                return None
+        async with AsyncClient(http2=True, follow_redirects=True, cookies=cookies, proxies=self.proxy) as client:
+            response = await client.get(self.url)
+            if response.status_code == 200:
+                content_type = response.headers.get("content-type")
+                if content_type and "image" not in content_type:
+                    console.log(f"Warning: Content type of {filename} is {content_type}, not an image.")
+                dest_path = Path(path)
+                dest_path.mkdir(parents=True, exist_ok=True)
+                dest = dest_path / filename
+                dest.write_bytes(response.content)
+                if verbose:
+                    console.log(f"Image saved as {dest.resolve()}")
+                return str(dest.resolve())
+            else:
+                raise HTTPStatusError(
+                    f"Error downloading image: {response.status_code} {response.reason_phrase}",
+                    request=response.request,
+                    response=response,
+                )
+
+class WebImage(Image):
+    """
+    Image retrieved from web.
+    Returned when asking Gemini to "SEND an image of [something]".
+    """
+    pass
 
-    chatbot = Chatbot(args.session, args.session_ts)
-    prompt_session = __create_session()
-    completions = __create_completer(["!exit", "!reset"])
+class GeneratedImage(Image):
+    """
+    Image generated by ImageFX (Google's AI image generator).
+    Parameters:
+        cookies: dict[str, str]
+            Cookies used from the GeminiClient.
+    """
+    cookies: Dict[str, str]
+
+    @field_validator("cookies")
+    def validate_cookies(cls, v):
+        if not v:
+            raise ValueError("GeneratedImage requires cookies from GeminiClient.")
+        return v
+
+    async def save(self, **kwargs) -> Optional[str]:
+        """
+        Save the generated image to disk.
+        Parameters:
+            filename: str, optional
+                Filename to use; generated images are in .png format.
+            Additional arguments are passed to Image.save.
+        Returns:
+            Absolute path of the saved image if successful.
+        """
+        filename = kwargs.pop("filename", None) or f"{datetime.now().strftime('%Y%m%d%H%M%S')}_{self.url[-10:]}.png"
+        return await super().save(filename=filename, cookies=self.cookies, **kwargs)
 
+#########################################
+# Main usage demonstration
+#########################################
+
+if __name__ == "__main__":
+    """
+    Usage demonstration:
+    - Reads cookies from 'cookies.json'
+    - Initializes the synchronous Chatbot wrapper.
+    - Performs a text query.
+    - Performs an image generation query and downloads the generated image.
+    - Demonstrates saving a conversation.
+    """
+    # Define the path to cookies file
+    cookies_file = r"C:\Users\hp\Desktop\Webscout\cookies.json"
+    
+    # Create Chatbot instance with a chosen model
+    try:
+        bot = Chatbot(cookie_path=cookies_file, model=Model.G_2_0_FLASH_THINKING_WITH_APPS)
+    except Exception as e:
+        console.log(f"[red]Error initializing Chatbot: {e}[/red]")
+        exit(1)
+
+    # Sample text query
+    text_message = "How many r's in word strawberry?"
+    console.log("[green]Sending text query to Gemini...[/green]")
+    try:
+        response_text = bot.ask(text_message)
+        console.log("[blue]Text Response:[/blue]")
+        console.print(Markdown(response_text.get("content", "No content received.")))
+    except Exception as e:
+        console.log(f"[red]Error sending text query: {e}[/red]")
+
+    # Image generation query
+    image_message = "Generate an image of a scenic view."
+    console.log("[green]Requesting image generation from Gemini...[/green]")
+    try:
+        response_image = bot.ask(image_message)
+        # Check if any image URL is returned in the response
+        image_urls = response_image.get("images", [])
+        if not image_urls:
+            console.log("[red]No image URLs returned in response.[/red]")
+        else:
+            image_url = image_urls[0]
+            console.log(f"[blue]Image URL received: {image_url}[/blue]")
+            # Use GeneratedImage class to download the generated image
+            generated_img = GeneratedImage(
+                url=image_url,
+                cookies={"__Secure-1PSID": bot.secure_1psid, "__Secure-1PSIDTS": bot.secure_1psidts}
+            )
+            saved_path = asyncio.run(generated_img.save(path="downloaded_images", verbose=True))
+            console.log(f"[blue]Generated image saved at: {saved_path}[/blue]")
+    except Exception as e:
+        console.log(f"[red]Error processing image generation: {e}[/red]")
+    
+    # Demonstrate saving a conversation
+    conversation_file = "conversations.json"
+    conversation_name = "Sample Conversation"
     try:
-        while True:
-            console.print("You:")
-            user_prompt = __get_input(prompt_sess=prompt_session, completer=completions)
-            console.print()
-            if user_prompt == "!exit":
-                break
-            elif user_prompt == "!reset":
-                chatbot.conversation_id = ""
-                chatbot.response_id = ""
-                chatbot.choice_id = ""
-                continue
-            print("Google Gemini:")
-            response = chatbot.ask(user_prompt)
-            console.print(Markdown(response["content"]))
-            console.print(response["images"] if response.get("images") else "")
-            print()
-    except KeyboardInterrupt:
-        print("Exiting...")
+        bot.save_conversation(conversation_file, conversation_name)
+        console.log(f"[green]Conversation saved to {conversation_file} under the name '{conversation_name}'.[/green]")
+    except Exception as e:
+        console.log(f"[red]Error saving conversation: {e}[/red]")