[fix]

README.md

@@ -33,12 +33,14 @@ pip install free_llms
 ```python
 
 from free_llms.models import GPTChrome
-driver_config = ["--disable-gpu", "--window-size=1920,1080"] # pass in selnium driver config
-with GPTChrome(driver_config,'[email protected]','') as session: # A single session started with ChatGPT. Put in your email and password for ChatGPT account.
-    data = session.send_prompt('Tell me a horror story in 150 words') # First Message
-    data1 = session.send_prompt('Now make it funny') # Second message
+driver_config = [] # pass in selnium driver config except for the following ["--disable-gpu", f"--window-size=1920,1080"]
+with GPTChrome(driver_config=driver_config,
+               email = 'email',
+               password = 'password') as session: # A single session started with ChartGPT
+    data = session.send_prompt("""Write an SQL Query which shows how to get third highest salary
+    """) # First Message
+    data1 = session.send_prompt('Now convert it into python') # Second message
     print(session.messages) # Messages in the current session in pair of <Human,AI>
-
 ```
 
 ## Note:

examples/gpt.py

@@ -1,8 +1,8 @@
 from free_llms.models import GPTChrome
 driver_config = [] # pass in selnium driver config except for the following ["--disable-gpu", f"--window-size=1920,1080"]
 with GPTChrome(driver_config=driver_config,
-               email = '[email protected]',
-               password = '') as session: # A single session started with ChartGPT
+               email = 'email',
+               password = 'password') as session: # A single session started with ChartGPT
     data = session.send_prompt("""Write an SQL Query which shows how to get third highest salary
     """) # First Message
     data1 = session.send_prompt('Now convert it into python') # Second message

src/free_llms/constants.py

@@ -1,2 +1,2 @@
-DEFAULT_WINDOW_SIZE = (1920,1080)
-DRIVERS_DEFAULT_CONFIG = ["--disable-gpu", f"--window-size={DEFAULT_WINDOW_SIZE[0]},{DEFAULT_WINDOW_SIZE[1]}"]
+DEFAULT_WINDOW_SIZE = (1920, 1080)
+DRIVERS_DEFAULT_CONFIG = ["--disable-gpu", f"--window-size={DEFAULT_WINDOW_SIZE[0]},{DEFAULT_WINDOW_SIZE[1]}"]

src/free_llms/models.py

@@ -1,71 +1,77 @@
+import time
 from abc import ABC, abstractmethod
-from typing import Dict, List, Tuple, Optional
+from typing import Any, Dict, List, Tuple
+
 import undetected_chromedriver as uc
-from free_llms.utils import configure_options
-from free_llms.constants import DRIVERS_DEFAULT_CONFIG
+from langchain_core.messages import AIMessage, HumanMessage
+from pydantic import BaseModel, model_validator
 from selenium.common.exceptions import TimeoutException
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.ui import WebDriverWait
-from pydantic import BaseModel, model_validator
-from langchain_core.messages import AIMessage, HumanMessage
-import time
+
+from free_llms.constants import DRIVERS_DEFAULT_CONFIG
+from free_llms.utils import configure_options
+
 
 class LLMChrome(BaseModel, ABC):
     """
     Abstract Class for establishing a single session with an LLM through a Chrome browser.
 
-    This class defines the interface for creating a Chrome-based interaction with a language model for a single session. 
+    This class defines the interface for creating a Chrome-based interaction with a language model for a single session.
 
     Methods:
     login(email: str, password: str, waiting_time: int = 10) -> bool:
         Logs into the language model interface using the provided email and password.
-        
+
     send_prompt(query: str, waiting_time: int = 10) -> str:
         Sends a query prompt to the language model and returns the response as a string.
-        
+
     __enter__():
         Enters the runtime context related to this object (for use in 'with' statements).
-        
+
     __exit__(*args, **kwargs):
         Exits the runtime context related to this object.
     """
+
     email: str
     """LLM Provider Account Email"""
     password: str
     """LLM Proiver Account Password"""
     waiting_time: int = 10
     """How much time do you need to wait for elements to appear. Client cannot change this. This depends on LLM providers"""
-    driver: Optional[uc.Chrome] = None
+    driver: uc.Chrome
     """UnDetected Chrome Driver. This is started automatically and client do not need to provide."""
-    driver_config: List[str] = []
+    driver_config: List[str]
     """Configuration for UnDetected Chrome Driver."""
     messages: List[Tuple[HumanMessage, AIMessage]] = []
     """Messages in the current session"""
+
     class Config:
         """Configuration for this pydantic object."""
+
         arbitrary_types_allowed = True
 
-    @model_validator(mode='before')
-    def check_start_driver(cls, data: Dict):
+    @model_validator(mode="before")
+    def check_start_driver(cls, data: Dict) -> Dict:
         """
         Validates and starts the Chrome driver with appropriate configurations.
         Ensures certain configurations are not modified or reused.
         """
-        if 'driver' in data:
-            raise ValueError('You cannot pass in an already started driver')
-        if 'message_jump' in data:
-            raise ValueError('You cannot set message jump')
+        if "driver" in data:
+            raise ValueError("You cannot pass in an already started driver")
+        if "message_jump" in data:
+            raise ValueError("You cannot set message jump")
         for default_config in DRIVERS_DEFAULT_CONFIG:
-            if default_config in data['driver_config']:
-                raise ValueError(f'You cannot put in {default_config} in your provided driver config')
-        for started_config in data['driver_config']:
-            if '--window-size' in started_config:
-                raise ValueError('You cannot change the window size in your provided driver config')
-        options = configure_options(data['driver_config'] + DRIVERS_DEFAULT_CONFIG)
-        data['driver'] = uc.Chrome(options=options, headless=True)
+            if default_config in data["driver_config"]:
+                raise ValueError(f"You cannot put in {default_config} in your provided driver config")
+        for started_config in data["driver_config"]:
+            if "--window-size" in started_config:
+                raise ValueError("You cannot change the window size in your provided driver config")
+        options = configure_options(data["driver_config"] + DRIVERS_DEFAULT_CONFIG)
+        data["driver"] = uc.Chrome(options=options, headless=True)
         return data
-        
+
     @property
     @abstractmethod
     def _model_url(self) -> str:
@@ -77,26 +83,28 @@ def _model_url(self) -> str:
     def _elements_identifier(self) -> Dict[str, str]:
         """A dictionary containing identifiers for various elements on the page"""
         pass
-    
+
     @property
     def session_history(self) -> List[Tuple[HumanMessage, AIMessage]]:
         """All of the messages in the current session in the form of Human and AI pairs."""
         return self.messages
 
-
     @abstractmethod
-    def login(self, email: str, password: str, waiting_time: int = 10) -> bool:
+    def login(self, retries_attempt: int = 3) -> bool:
         """
-        Logs into LLM Provider Browser using the provided email and password. 
+        Logs into LLM Provider Browser using the provided email and password.
         No SSO (Single Sign On)
 
+        Args:
+        retries_attempt (str): The number of attempts to do a login request
+
         Returns:
         bool: True if login is successful, False otherwise.
         """
         pass
 
     @abstractmethod
-    def send_prompt(self, query: str, waiting_time: int = 10) -> str:
+    def send_prompt(self, query: str) -> str:
         """
         Sends a query prompt to LLM Provider and returns the response as a string.
 
@@ -107,18 +115,16 @@ def send_prompt(self, query: str, waiting_time: int = 10) -> str:
         str: The response from LLM Provider.
         """
         pass
-
-
-    def __enter__(self):
+
+    def __enter__(self) -> "LLMChrome":
         """
         Enters the runtime context related to this object (for use in 'with' statements).
         Automatically logs in upon entering the context.
         """
         self.login()
         return self
-
-
-    def __exit__(self,*args,**kwargs):
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
         """
         Exits the runtime context and closes the browser session.
         """
@@ -134,11 +140,11 @@ class GPTChrome(LLMChrome):
         response = browser.send_prompt("What is the capital of France?")
         print(response)
     """
-    
+
     @property
     def _model_url(self) -> str:
         return "https://chatgpt.com/auth/login?sso="
-    
+
     @property
     def _elements_identifier(self) -> Dict[str, str]:
         return {
@@ -147,35 +153,37 @@ def _elements_identifier(self) -> Dict[str, str]:
             "Email_Continue": "action",
             "Password": '//*[@id="password"]',
             "Prompt_Text_Area": "prompt-textarea",
-            "Prompt_Text_Output": '//*[@id="__next"]/div[1]/div[2]/main/div[2]/div[1]/div/div/div/div/div[{current}]/div/div/div[2]/div[2]/div[1]/div/div', # noqa: E501
+            "Prompt_Text_Output": '//*[@id="__next"]/div[1]/div[2]/main/div[2]/div[1]/div/div/div/div/div[{current}]/div/div/div[2]/div[2]/div[1]/div/div',  # noqa: E501
         }
-    
-    def login(self) -> bool:
+
+    def login(self, retries_attempt: int = 3) -> bool:
         self.driver.get(self._model_url)
-        try:
-            login_button = WebDriverWait(self.driver, self.waiting_time).until(
-                EC.element_to_be_clickable((By.XPATH, self._elements_identifier["Login"]))
-            )
-            login_button.click()
-            email_input = WebDriverWait(self.driver, self.waiting_time).until(
-                EC.presence_of_element_located((By.ID, self._elements_identifier["Email"]))
-            )
-            email_input.click()
-            email_input.send_keys(self.email)
-            continue_button = WebDriverWait(self.driver, self.waiting_time).until(
-                EC.presence_of_element_located((By.NAME, self._elements_identifier["Email_Continue"]))
-            )
-            continue_button.click()
-            password_button = WebDriverWait(self.driver, self.waiting_time).until(
-                EC.element_to_be_clickable((By.XPATH, self._elements_identifier["Password"]))
-            )
-            password_button.clear()
-            password_button.click()
-            password_button.send_keys(self.password)
-            password_button.submit()
-        except TimeoutException:
-            self.login()
-        return True
+        for _ in range(retries_attempt):
+            try:
+                login_button = WebDriverWait(self.driver, self.waiting_time).until(
+                    EC.element_to_be_clickable((By.XPATH, self._elements_identifier["Login"]))
+                )
+                login_button.click()
+                email_input = WebDriverWait(self.driver, self.waiting_time).until(
+                    EC.presence_of_element_located((By.ID, self._elements_identifier["Email"]))
+                )
+                email_input.click()
+                email_input.send_keys(self.email)
+                continue_button = WebDriverWait(self.driver, self.waiting_time).until(
+                    EC.presence_of_element_located((By.NAME, self._elements_identifier["Email_Continue"]))
+                )
+                continue_button.click()
+                password_button = WebDriverWait(self.driver, self.waiting_time).until(
+                    EC.element_to_be_clickable((By.XPATH, self._elements_identifier["Password"]))
+                )
+                password_button.clear()
+                password_button.click()
+                password_button.send_keys(self.password)
+                password_button.submit()
+                return True
+            except TimeoutException:
+                continue
+        return False
 
     def send_prompt(self, query: str) -> str:
         while True:
@@ -189,20 +197,19 @@ def send_prompt(self, query: str) -> str:
                 self.driver.quit()
                 self.driver = uc.Chrome(options=configure_options(self.driver_config + DRIVERS_DEFAULT_CONFIG), headless=True)
                 self.driver.get(current_url)
-            
+
         text_area.click()
         text_area.send_keys(query)
         text_area.submit()
-        raw_message = ''
-        streaming = ''
+        raw_message = ""
         time.sleep(self.waiting_time)  # Wait for the query to be processed
         current_n, prev_n = 0, -1
         message_jump = 3
         while current_n != prev_n:
             prev_n = current_n
-            streaming = self.driver.find_element(By.XPATH, self._elements_identifier['Prompt_Text_Output'].format(current=message_jump))
-            raw_message = streaming.get_attribute('innerHTML')
-            current_n = len(raw_message)
+            streaming = self.driver.find_element(By.XPATH, self._elements_identifier["Prompt_Text_Output"].format(current=message_jump))
+            raw_message = streaming.get_attribute("innerHTML")
+            current_n = len(raw_message) if raw_message is not None else 0
         message_jump += 2
-        self.messages.append([HumanMessage(content=query), AIMessage(content=raw_message)])
-        return raw_message
+        self.messages.append((HumanMessage(content=query), AIMessage(content=raw_message)))
+        return raw_message

src/free_llms/utils.py

@@ -1,7 +1,11 @@
 """Helper Functions for all of our models"""
+
 from typing import List
+
 import undetected_chromedriver as uc
 from fake_useragent import UserAgent
+
+
 def configure_options(driver_config: List[str]) -> uc.ChromeOptions:
     """
     Configures the Chrome options.
@@ -13,8 +17,8 @@ def configure_options(driver_config: List[str]) -> uc.ChromeOptions:
     uc.ChromeOptions: Configured Chrome options.
     """
     chrome_options = uc.ChromeOptions()
-    userAgent = UserAgent(browsers='chrome').random
+    userAgent = UserAgent(browsers="chrome").random
     for arg in driver_config:
         chrome_options.add_argument(arg)
     chrome_options.add_argument(f"--user-agent={userAgent}")
-    return chrome_options
+    return chrome_options