3 files changed, 54 insertions, 60 deletions

diff --git a/gpt4free/theb/__init__.py b/gpt4free/theb/__init__.py
index 726e025e..96053877 100644
--- a/gpt4free/theb/__init__.py
+++ b/gpt4free/theb/__init__.py
@@ -1,8 +1,12 @@
-from re import findall
 from json import loads
 from queue import Queue, Empty
+from re import findall
 from threading import Thread
+from typing import Generator
+
 from curl_cffi import requests
+from fake_useragent import UserAgent
+
 
 class Completion:
     # experimental
@@ -14,29 +18,29 @@ class Completion:
     message_queue = Queue()
     stream_completed = False
 
+    @staticmethod
     def request(prompt: str):
         headers = {
             'authority': 'chatbot.theb.ai',
             'content-type': 'application/json',
             'origin': 'https://chatbot.theb.ai',
-            'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/112.0.0.0 Safari/537.36',
+            'user-agent': UserAgent().random,
         }
 
-        requests.post('https://chatbot.theb.ai/api/chat-process', headers=headers,
-            content_callback = Completion.handle_stream_response,
-            json = {
-                'prompt': prompt,
-                'options': {}
-            }
+        requests.post(
+            'https://chatbot.theb.ai/api/chat-process',
+            headers=headers,
+            content_callback=Completion.handle_stream_response,
+            json={'prompt': prompt, 'options': {}},
         )
 
         Completion.stream_completed = True
 
     @staticmethod
-    def create(prompt: str):
+    def create(prompt: str) -> Generator[str, None, None]:
         Thread(target=Completion.request, args=[prompt]).start()
 
-        while Completion.stream_completed != True or not Completion.message_queue.empty():
+        while not Completion.stream_completed or not Completion.message_queue.empty():
             try:
                 message = Completion.message_queue.get(timeout=0.01)
                 for message in findall(Completion.regex, message):
diff --git a/gui/query_methods.py b/gui/query_methods.py
index adb4ae91..6225453b 100644
--- a/gui/query_methods.py
+++ b/gui/query_methods.py
@@ -7,7 +7,6 @@ from gpt4free import quora, forefront, theb, you
 import random
 
 
-
 def query_forefront(question: str) -> str:
     # create an account
     token = forefront.Account.create(logging=False)
@@ -15,65 +14,59 @@ def query_forefront(question: str) -> str:
     response = ""
     # get a response
     try:
-        for i in forefront.StreamingCompletion.create(token = token, prompt ='hello world', model='gpt-4'):
-            response += i.completion.choices[0].text
-        
-        return response
-    
+        return forefront.Completion.create(token=token, prompt='hello world', model='gpt-4').text
     except Exception as e:
         # Return error message if an exception occurs
-        return f'An error occurred: {e}. Please make sure you are using a valid cloudflare clearance token and user agent.'
+        return (
+            f'An error occurred: {e}. Please make sure you are using a valid cloudflare clearance token and user agent.'
+        )
 
 
 def query_quora(question: str) -> str:
     token = quora.Account.create(logging=False, enable_bot_creation=True)
-    response = quora.Completion.create(
-        model='gpt-4',
-        prompt=question,
-        token=token
-    )
-
-    return response.completion.choices[0].tex
+    return quora.Completion.create(model='gpt-4', prompt=question, token=token).text
 
 
 def query_theb(question: str) -> str:
     # Set cloudflare clearance cookie and get answer from GPT-4 model
     response = ""
     try:
-        result = theb.Completion.create(
-            prompt = question)
-        return result
-    
+        return ''.join(theb.Completion.create(prompt=question))
+
     except Exception as e:
         # Return error message if an exception occurs
-        return f'An error occurred: {e}. Please make sure you are using a valid cloudflare clearance token and user agent.'
+        return (
+            f'An error occurred: {e}. Please make sure you are using a valid cloudflare clearance token and user agent.'
+        )
 
 
 def query_you(question: str) -> str:
     # Set cloudflare clearance cookie and get answer from GPT-4 model
     try:
-        result = you.Completion.create(
-            prompt = question)
+        result = you.Completion.create(prompt=question)
         return result["response"]
-    
+
     except Exception as e:
         # Return error message if an exception occurs
-        return f'An error occurred: {e}. Please make sure you are using a valid cloudflare clearance token and user agent.'
+        return (
+            f'An error occurred: {e}. Please make sure you are using a valid cloudflare clearance token and user agent.'
+        )
+
 
 # Define a dictionary containing all query methods
 avail_query_methods = {
-        "Forefront": query_forefront,
-        "Poe": query_quora,
-        "Theb": query_theb,
-        "You": query_you,
-        # "Writesonic": query_writesonic,
-        # "T3nsor": query_t3nsor,
-        # "Phind": query_phind,
-        # "Ora": query_ora,
-    }
+    "Forefront": query_forefront,
+    "Poe": query_quora,
+    "Theb": query_theb,
+    "You": query_you,
+    # "Writesonic": query_writesonic,
+    # "T3nsor": query_t3nsor,
+    # "Phind": query_phind,
+    # "Ora": query_ora,
+}
 
-def query(user_input: str, selected_method: str = "Random") -> str:
 
+def query(user_input: str, selected_method: str = "Random") -> str:
     # If a specific query method is selected (not "Random") and the method is in the dictionary, try to call it
     if selected_method != "Random" and selected_method in avail_query_methods:
         try:
@@ -104,4 +97,3 @@ def query(user_input: str, selected_method: str = "Random") -> str:
             query_methods_list.remove(chosen_query)
 
     return result
-
diff --git a/gui/streamlit_chat_app.py b/gui/streamlit_chat_app.py
index 1c0dd92a..68011229 100644
--- a/gui/streamlit_chat_app.py
+++ b/gui/streamlit_chat_app.py
@@ -1,6 +1,6 @@
+import atexit
 import os
 import sys
-import atexit
 
 sys.path.append(os.path.join(os.path.dirname(__file__), os.path.pardir))
 
@@ -9,9 +9,9 @@ from streamlit_chat import message
 from query_methods import query, avail_query_methods
 import pickle
 
-
 conversations_file = "conversations.pkl"
 
+
 def load_conversations():
     try:
         with open(conversations_file, "rb") as f:
@@ -31,11 +31,11 @@ def save_conversations(conversations, current_conversation):
             break
     if not updated:
         conversations.append(current_conversation)
-    
+
     temp_conversations_file = "temp_" + conversations_file
     with open(temp_conversations_file, "wb") as f:
         pickle.dump(conversations, f)
-    
+
     os.replace(temp_conversations_file, conversations_file)
 
 
@@ -44,10 +44,10 @@ def exit_handler():
     # Perform cleanup operations here, like saving data or closing open files.
     save_conversations(st.session_state.conversations, st.session_state.current_conversation)
 
+
 # Register the exit_handler function to be called when the program is closing.
 atexit.register(exit_handler)
 
-
 st.header("Chat Placeholder")
 
 if 'conversations' not in st.session_state:
@@ -61,7 +61,7 @@ if 'selected_conversation' not in st.session_state:
 
 if 'input_field_key' not in st.session_state:
     st.session_state['input_field_key'] = 0
-    
+
 if 'query_method' not in st.session_state:
     st.session_state['query_method'] = query
 
@@ -69,20 +69,22 @@ if 'query_method' not in st.session_state:
 if 'current_conversation' not in st.session_state or st.session_state['current_conversation'] is None:
     st.session_state['current_conversation'] = {'user_inputs': [], 'generated_responses': []}
 
-
 input_placeholder = st.empty()
-user_input = input_placeholder.text_input('You:', key=f'input_text_{len(st.session_state["current_conversation"]["user_inputs"])}')
+user_input = input_placeholder.text_input(
+    'You:', key=f'input_text_{len(st.session_state["current_conversation"]["user_inputs"])}'
+)
 submit_button = st.button("Submit")
 
 if user_input or submit_button:
     output = query(user_input, st.session_state['query_method'])
     escaped_output = output.encode('utf-8').decode('unicode-escape')
-    
+
     st.session_state.current_conversation['user_inputs'].append(user_input)
     st.session_state.current_conversation['generated_responses'].append(escaped_output)
     save_conversations(st.session_state.conversations, st.session_state.current_conversation)
-    user_input = input_placeholder.text_input('You:', value='', key=f'input_text_{len(st.session_state["current_conversation"]["user_inputs"])}')  # Clear the input field
-
+    user_input = input_placeholder.text_input(
+        'You:', value='', key=f'input_text_{len(st.session_state["current_conversation"]["user_inputs"])}'
+    )  # Clear the input field
 
 # Add a button to create a new conversation
 if st.sidebar.button("New Conversation"):
@@ -90,11 +92,7 @@ if st.sidebar.button("New Conversation"):
     st.session_state['current_conversation'] = {'user_inputs': [], 'generated_responses': []}
     st.session_state['input_field_key'] += 1
 
-st.session_state['query_method'] = st.sidebar.selectbox(
-    "Select API:",
-    options=avail_query_methods,
-    index=0
-)
+st.session_state['query_method'] = st.sidebar.selectbox("Select API:", options=avail_query_methods, index=0)
 
 # Sidebar
 st.sidebar.header("Conversation History")