patch: Minor fixes

Simatwa · May 2, 2024 · e918aa3 · e918aa3
1 parent a30ae59
commit e918aa3
Show file tree

Hide file tree

Showing 14 changed files with 59 additions and 45 deletions.
diff --git a/src/pytgpt/async_providers.py b/src/pytgpt/async_providers.py
@@ -11,23 +11,23 @@
 
 mapper: dict[str, object] = {
     "phind": AsyncPHIND,
-    "yepchat": AsyncYEPCHAT,
     "opengpt": AsyncOPENGPT,
-    "openai": AsyncOPENAI,
+    "koboldai": AsyncKOBOLDAI,
+    "blackboxai": AsyncBLACKBOXAI,
+    "gpt4free": AsyncGPT4FREE,
     "llama2": AsyncLLAMA2,
+    "yepchat": AsyncYEPCHAT,
     "leo": AsyncLEO,
-    "koboldai": AsyncKOBOLDAI,
     "groq": AsyncGROQ,
-    "gpt4free": AsyncGPT4FREE,
-    "blackboxai": AsyncBLACKBOXAI,
+    "openai": AsyncOPENAI,
 }
 
 tgpt_mapper: dict[str, object] = {
     "phind": AsyncPHIND,
-    "yepchat": AsyncYEPCHAT,
     "opengpt": AsyncOPENGPT,
-    "llama2": AsyncLLAMA2,
     "koboldai": AsyncKOBOLDAI,
     # "gpt4free": AsyncGPT4FREE,
     "blackboxai": AsyncBLACKBOXAI,
+    "llama2": AsyncLLAMA2,
+    "yepchat": AsyncYEPCHAT,
 }
diff --git a/src/pytgpt/auto/main.py b/src/pytgpt/auto/main.py
@@ -321,7 +321,7 @@ async def ask(
         }
 
         # tgpt-based providers
-        for provider_name, provider_obj in self.async_provider_map.items():
+        for provider_name, provider_obj in async_provider_map.items():
             if provider_name in self.exclude:
                 continue
             try:

diff --git a/src/pytgpt/blackboxai/main.py b/src/pytgpt/blackboxai/main.py
@@ -1,6 +1,7 @@
 import json
 import httpx
 import requests
+import pytgpt.exceptions as exceptions
 from pytgpt.utils import Optimizers
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
@@ -140,8 +141,8 @@ def for_stream():
                 or not response.headers.get("Content-Type")
                 == "text/plain; charset=utf-8"
             ):
-                raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                raise exceptions.FailedToGenerateResponseError(
+                    f"Failed to generate response - ({response.status_code}, {response.reason})"
                 )
             streaming_text = ""
             for value in response.iter_lines(
@@ -343,8 +344,8 @@ async def for_stream():
                     or not response.headers.get("Content-Type")
                     == "text/plain; charset=utf-8"
                 ):
-                    raise Exception(
-                        f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                    raise exceptions.FailedToGenerateResponseError(
+                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                     )
                 streaming_text = ""
                 async for value in response.aiter_lines():

diff --git a/src/pytgpt/exceptions.py b/src/pytgpt/exceptions.py
@@ -0,0 +1,4 @@
+class FailedToGenerateResponseError(Exception):
+    """Provider failed to fetch response"""
+
+    pass
diff --git a/src/pytgpt/gpt4free/main.py b/src/pytgpt/gpt4free/main.py
@@ -4,6 +4,7 @@
 from pytgpt.base import Provider, AsyncProvider
 from pytgpt import available_providers
 import g4f
+import pytgpt.exceptions as exceptions
 from typing import AsyncGenerator
 
 g4f.debug.version_check = False

diff --git a/src/pytgpt/groq/main.py b/src/pytgpt/groq/main.py
@@ -1,6 +1,7 @@
 import requests
 import httpx
 import json
+import pytgpt.exceptions as exceptions
 from pytgpt.utils import Optimizers
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
@@ -159,8 +160,8 @@ def for_stream():
                 self.chat_endpoint, json=payload, stream=True, timeout=self.timeout
             )
             if not response.ok:
-                raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                raise exceptions.FailedToGenerateResponseError(
+                    f"Failed to generate response - ({response.status_code}, {response.reason})"
                 )
 
             message_load = ""
@@ -190,8 +191,8 @@ def for_non_stream():
                 self.chat_endpoint, json=payload, stream=False, timeout=self.timeout
             )
             if not response.ok:
-                raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                raise exceptions.FailedToGenerateResponseError(
+                    f"Failed to generate response - ({response.status_code}, {response.reason})"
                 )
             resp = response.json()
             self.last_response.update(resp)
@@ -401,7 +402,7 @@ async def for_stream():
             ) as response:
                 if not response.is_success:
                     raise Exception(
-                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
+                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                     )
 
                 message_load = ""
@@ -431,7 +432,7 @@ async def for_non_stream():
             )
             if not response.is_success:
                 raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
+                    f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                 )
             resp = response.json()
             self.last_response.update(resp)

diff --git a/src/pytgpt/koboldai/main.py b/src/pytgpt/koboldai/main.py
@@ -1,6 +1,7 @@
 import requests
 import json
 import httpx
+import pytgpt.exceptions as exceptions
 from pytgpt.utils import Optimizers
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
@@ -122,8 +123,8 @@ def for_stream():
                 self.chat_endpoint, json=payload, stream=True, timeout=self.timeout
             )
             if not response.ok:
-                raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                raise exceptions.FailedToGenerateResponseError(
+                    f"Failed to generate response - ({response.status_code}, {response.reason})"
                 )
 
             message_load = ""
@@ -309,8 +310,8 @@ async def for_stream():
                 "POST", self.chat_endpoint, json=payload, timeout=self.timeout
             ) as response:
                 if not response.is_success:
-                    raise Exception(
-                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
+                    raise exceptions.FailedToGenerateResponseError(
+                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                     )
 
                 message_load = ""

diff --git a/src/pytgpt/leo/main.py b/src/pytgpt/leo/main.py
@@ -1,6 +1,7 @@
 import requests
 import httpx
 import json
+import pytgpt.exceptions as exceptions
 from pytgpt.utils import Optimizers
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
@@ -157,8 +158,8 @@ def for_stream():
                 or not response.headers.get("Content-Type")
                 == "text/event-stream; charset=utf-8"
             ):
-                raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                raise exceptions.FailedToGenerateResponseError(
+                    f"Failed to generate response - ({response.status_code}, {response.reason})"
                 )
 
             for value in response.iter_lines(
@@ -371,8 +372,8 @@ async def for_stream():
                     or not response.headers.get("Content-Type")
                     == "text/event-stream; charset=utf-8"
                 ):
-                    raise Exception(
-                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
+                    raise exceptions.FailedToGenerateResponseError(
+                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                     )
                 async for value in response.aiter_lines():
                     try:

diff --git a/src/pytgpt/llama2/main.py b/src/pytgpt/llama2/main.py
@@ -1,6 +1,7 @@
 import requests
 import json
 import httpx
+import pytgpt.exceptions as exceptions
 from pytgpt.utils import Optimizers
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
@@ -141,8 +142,8 @@ def for_stream():
                 or not response.headers.get("Content-Type")
                 == "text/plain; charset=utf-8"
             ):
-                raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                raise exceptions.FailedToGenerateResponseError(
+                    f"Failed to generate response - ({response.status_code}, {response.reason})"
                 )
 
             message_load: str = ""
@@ -348,8 +349,8 @@ async def for_stream():
                     or not response.headers.get("Content-Type")
                     == "text/plain; charset=utf-8"
                 ):
-                    raise Exception(
-                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
+                    raise exceptions.FailedToGenerateResponseError(
+                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                     )
                 message_load: str = ""
                 async for value in response.aiter_lines():

diff --git a/src/pytgpt/openai/main.py b/src/pytgpt/openai/main.py
@@ -1,6 +1,7 @@
 import requests
 import json
 import httpx
+import pytgpt.exceptions as exceptions
 from pytgpt.utils import Optimizers
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
@@ -136,7 +137,7 @@ def ask(
                     conversation_prompt if conversationally else prompt
                 )
             else:
-                raise Exception(
+                raise exceptions.FailedToGenerateResponseError(
                     f"Optimizer is not one of {self.__available_optimizers}"
                 )
         session.headers.update(self.headers)
@@ -155,7 +156,7 @@ def for_stream():
                 self.chat_endpoint, json=payload, stream=True, timeout=self.timeout
             )
             if not response.ok:
-                raise Exception(
+                raise exceptions.FailedToGenerateResponseError(
                     f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
                 )
 
@@ -189,7 +190,7 @@ def for_non_stream():
                 not response.ok
                 or not response.headers.get("Content-Type", "") == "application/json"
             ):
-                raise Exception(
+                raise exceptions.FailedToGenerateResponseError(
                     f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
                 )
             resp = response.json()
@@ -398,7 +399,7 @@ async def for_stream():
             ) as response:
                 if not response.is_success:
                     raise Exception(
-                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
+                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                     )
 
                 message_load = ""
@@ -432,7 +433,7 @@ async def for_non_stream():
                 or not response.headers.get("Content-Type", "") == "application/json"
             ):
                 raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
+                    f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                 )
             resp = response.json()
             self.last_response.update(resp)

diff --git a/src/pytgpt/opengpt/main.py b/src/pytgpt/opengpt/main.py
@@ -6,6 +6,7 @@
 from pytgpt.utils import Optimizers
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
+import pytgpt.exceptions as exceptions
 from pytgpt.base import Provider, AsyncProvider
 from typing import AsyncGenerator
 
@@ -133,7 +134,7 @@ def ask(
                     conversation_prompt if conversationally else prompt
                 )
             else:
-                raise Exception(
+                raise exceptions.FailedToGenerateResponseError(
                     f"Optimizer is not one of {self.__available_optimizers}"
                 )
 
@@ -165,8 +166,8 @@ def for_stream():
                 or not response.headers.get("Content-Type")
                 == "text/event-stream; charset=utf-8"
             ):
-                raise Exception(
-                    f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+                raise exceptions.FailedToGenerateResponseError(
+                    f"Failed to generate response - ({response.status_code}, {response.reason})"
                 )
 
             for value in response.iter_lines(
@@ -395,7 +396,7 @@ async def for_stream():
                     or not response.headers.get("Content-Type")
                     == "text/event-stream; charset=utf-8"
                 ):
-                    raise Exception(
+                    raise exceptions.FailedToGenerateResponseError(
                         f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
                     )
 

diff --git a/src/pytgpt/phind/main.py b/src/pytgpt/phind/main.py
@@ -7,6 +7,7 @@
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
 from pytgpt.base import Provider, AsyncProvider
+import pytgpt.exceptions as exceptions
 from typing import AsyncGenerator
 
 session = requests.Session()
@@ -147,7 +148,7 @@ def for_stream():
                 or not response.headers.get("Content-Type")
                 == "text/event-stream; charset=utf-8"
             ):
-                raise Exception(
+                raise exceptions.FailedToGenerateResponseError(
                     f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
                 )
             streaming_text = ""
@@ -391,8 +392,8 @@ async def for_stream():
                     or not response.headers.get("Content-Type")
                     == "text/event-stream; charset=utf-8"
                 ):
-                    raise Exception(
-                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
+                    raise exceptions.FailedToGenerateResponseError(
+                        f"Failed to generate response - ({response.status_code}, {response.reason_phrase})"
                     )
                 streaming_text = ""
                 async for value in response.aiter_lines():

diff --git a/src/pytgpt/yepchat/main.py b/src/pytgpt/yepchat/main.py
@@ -5,6 +5,7 @@
 from pytgpt.utils import Conversation
 from pytgpt.utils import AwesomePrompts
 from pytgpt.utils import sanitize_stream
+import pytgpt.exceptions as exceptions
 from pytgpt.base import Provider, AsyncProvider
 
 session = requests.Session()
@@ -150,7 +151,7 @@ def for_stream():
                 self.chat_endpoint, json=payload, stream=True, timeout=self.timeout
             )
             if not response.ok:
-                raise Exception(
+                raise exceptions.FailedToGenerateResponseError(
                     f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
                 )
 
@@ -375,7 +376,7 @@ async def for_stream():
                 "POST", self.chat_endpoint, json=payload, timeout=self.timeout
             ) as response:
                 if not response.is_success:
-                    raise Exception(
+                    raise exceptions.FailedToGenerateResponseError(
                         f"Failed to generate response - ({response.status_code}, {response.reason_phrase}) - {response.text}"
                     )
 

diff --git a/tests/test_auto.py → tests/test_auto_tgpt.py b/tests/test_auto.py → tests/test_auto_tgpt.py
@@ -12,7 +12,7 @@ def setUp(self):
 
 class TestAsyncAuto(base.AsyncProviderBase):
 
-    async def setUp(self):
+    def setUp(self):
         self.bot = AsyncAUTO()
         self.prompt = base.prompt