Tests

makefile

@@ -6,3 +6,6 @@ lint:
 format:
 	poetry run ruff format $(PYTHON_FILES)
 	poetry run ruff --select I --fix $(PYTHON_FILES)
+
+test:
+	poetry run pytest

pyproject.toml

@@ -7,7 +7,7 @@ license = "MIT"
 readme = "README.md"
 
 [tool.poetry.dependencies]
-python = ">=3.8.1,<4.0"
+python = ">=3.9,<4.0"
 langchain = "^0.2.1"
 beautifulsoup4 = "^4.12.3"
 selenium = "^4.21.0"
@@ -20,6 +20,7 @@ langchain-community = "^0.2.3"
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.4.7"
 mypy = "^1.10.0"
+pytest = "^8.2.2"
 
 [tool.ruff]
 line-length = 150
@@ -36,6 +37,8 @@ select = [
 disallow_untyped_defs = "True"
 disable_error_code = ["import-untyped"]
 
+[tool.pytest.ini_options]
+pythonpath = ["./src"]
 
 [build-system]
 requires = ["poetry-core"]

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]}"]
+DRIVERS_DEFAULT_CONFIG = ["--disable-gpu", f"--window-size={DEFAULT_WINDOW_SIZE[0]},{DEFAULT_WINDOW_SIZE[1]}", "--no-sandbox"]

src/free_llms/models.py

@@ -25,6 +25,8 @@ class LLMChrome(BaseModel, ABC):
 
     This class defines the interface for creating a Chrome-based interaction with a language model for a single session.
 
+    You can also explicitly set up your chrome browser version by the env variable `CHROME_VERSION`.
+
     Methods:
     login(email: str, password: str, waiting_time: int = 10) -> bool:
         Logs into the language model interface using the provided email and password.
@@ -348,6 +350,9 @@ def login(self, retries_attempt: int) -> bool:
                     EC.element_to_be_clickable((By.XPATH, self._elements_identifier["Login_Button"]))
                 )
                 login_button.click()
+                WebDriverWait(self.driver, self.waiting_time).until(
+                    EC.presence_of_element_located((By.XPATH, self._elements_identifier["Prompt_Text_Area"]))
+                )
                 self.run_manager.on_text(text=f"Login succeed on attempt no. {i+1}", verbose=self.verbose)
                 return True
             except TimeoutException:

tests/unit_tests/test_models.py

@@ -0,0 +1,72 @@
+import pytest
+
+from free_llms.models import PreplexityChrome,GPTChrome,MistralChrome,AIMessage
+
+
+
+def test_gpt_chrome():
+    with pytest.raises(ValueError, match="Cannot Login given the credentials"):
+        with GPTChrome(driver_config=[],
+                email = 'wrong_email',
+                password = 'wrong_password',
+                retries_attempt=1
+                ) as session:
+            pass
+    assert GPTChrome(
+            driver_config=[],
+            email = 'wrong_email',
+            password = 'wrong_password')._elements_identifier == {
+            "Login": '//*[@id="__next"]/div[1]/div[2]/div[1]/div/div/button[1]',  # noqa: E501
+            "Email": "username",
+            "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
+        }
+    assert GPTChrome(driver_config=[], email = 'wrong_email', password = 'wrong_password')._model_url == 'https://chatgpt.com/auth/login?sso='
+    chrome_instance = GPTChrome(driver_config=[],email = 'wrong_email',password = 'wrong_password',retries_attempt=1)
+    chrome_instance.driver.get('https://chatgpt.com/')
+    ans = chrome_instance.send_prompt('How are you doing?')
+    assert isinstance(ans,AIMessage)
+
+
+def test_preplexity_chrome():
+    assert PreplexityChrome(
+            driver_config=[],
+            email = 'email',
+            password = 'password')._elements_identifier == {
+            "Prompt_Text_Area": "/html/body/div/main/div/div/div/div/div/div/div[1]/div[2]/div/div/span/div/div/textarea",
+            "Prompt_Text_Area_Submit": "#__next > main > div > div > div.grow.lg\:pr-sm.lg\:pb-sm.lg\:pt-sm > div > div > div > div.relative.flex.h-full.flex-col > div.mt-lg.w-full.grow.items-center.md\:mt-0.md\:flex.border-borderMain\/50.ring-borderMain\/50.divide-borderMain\/50.dark\:divide-borderMainDark\/50.dark\:ring-borderMainDark\/50.dark\:border-borderMainDark\/50.bg-transparent > div > div > span > div > div > div.bg-background.dark\:bg-offsetDark.flex.items-center.space-x-2.justify-self-end.rounded-full.col-start-3.row-start-2.-mr-2 > button",  # noqa: E501
+            "Prompt_Text_Area_Output": "/html/body/div/main/div/div/div/div/div/div[2]/div[1]/div/div/div[1]/div/div/div[3]/div/div[1]/div[2]/div/div[2]",  # noqa: E501
+            "Prompt_Text_Area_Output_Related": "/html/body/div/main/div/div/div/div/div/div[2]/div[1]/div/div/div[1]/div/div/div[3]/div/div[1]/div[3]/div/div",  # noqa: E501
+            "App_Download_Button": "/html/body/div[1]/main/div[3]/div/div/div/div[2]/div[1]/div/div/button",
+        }
+    assert PreplexityChrome(driver_config=[], email = 'email', password = 'password')._model_url == 'https://www.perplexity.ai/'
+    with PreplexityChrome(driver_config=[],
+               email = '', 
+               password = '',
+               ) as session:
+        ans = session.send_prompt("How are you doing?")
+        assert isinstance(ans,AIMessage)
+
+
+def test_mistral_chrome():
+    with pytest.raises(ValueError, match="Cannot Login given the credentials"):
+        with MistralChrome(driver_config=[],
+                email = 'wrong_email',
+                password = 'wrong_password',
+                retries_attempt=1
+                ) as session: # A single session started with ChartGPT
+            pass
+    assert MistralChrome(
+            driver_config=[],
+            email = 'wrong_email',
+            password = 'wrong_password')._elements_identifier == {
+            "Email": ":Rclkn:",
+            "Password": ":Rklkn:",
+            "Login_Button": "/html/body/main/div/div[1]/div/div/div[2]/div/form[2]/div[3]/div[2]/div/button",
+            "Prompt_Text_Area": "/html/body/div[1]/div[2]/div[2]/div/div[2]/div/div[1]/div/textarea",
+            "Prompt_Text_Area_Submit": "/html/body/div[1]/div[2]/div[2]/div/div[2]/div/div[1]/div/button",
+            "Prompt_Text_Area_Output": "/html/body/div[1]/div[2]/div[2]/div/div[1]/div[1]/div[{current}]/div[2]/div[1]",
+        }
+    assert MistralChrome(driver_config=[], email = 'wrong_email', password = 'wrong_password')._model_url == 'https://chat.mistral.ai/chat'