webscout 4.7__py3-none-any.whl → 4.9__py3-none-any.whl

webscout/Provider/OLLAMA.py

@@ -13,10 +13,10 @@ import io
 import re
 import json
 import yaml
-from ..AIutel import Optimizers
-from ..AIutel import Conversation
-from ..AIutel import AwesomePrompts, sanitize_stream
-from ..AIbase import Provider, AsyncProvider
+from webscout.AIutel import Optimizers
+from webscout.AIutel import Conversation
+from webscout.AIutel import AwesomePrompts, sanitize_stream
+from webscout.AIbase import Provider, AsyncProvider
 from webscout import exceptions
 from typing import Any, AsyncGenerator, Dict
 import logging
@@ -182,6 +182,7 @@ class OLLAMA(Provider):
         assert isinstance(response, dict), "Response should be of dict data-type only"
         return response["text"]
 if __name__ == "__main__":
-    ollama_provider = OLLAMA(model="qwen2:0.5b")
-    response = ollama_provider.chat("What is the meaning of life?")
-    print(response)
+    ollama_provider = OLLAMA(model="qwen:0.5b")
+    response = ollama_provider.chat("hi")
+    for r in response:
+        print(r, end="", flush=True)

webscout/Provider/{Geminiflash.py → PI.py}

@@ -1,13 +1,16 @@
+import cloudscraper
+import json
+
 import requests
-from ..AIbase import Provider
-from ..AIutel import Conversation
-from ..AIutel import Optimizers
-from ..AIutel import AwesomePrompts
-from webscout import exceptions
+from webscout.AIutel import Optimizers
+from webscout.AIutel import Conversation
+from webscout.AIutel import AwesomePrompts
+from webscout.AIbase import Provider
 
-class GEMINIFLASH(Provider):
+class PiAI(Provider):
     def __init__(
         self,
+        conversation_id: str,
         is_conversation: bool = True,
         max_tokens: int = 600,
         timeout: int = 30,
@@ -18,9 +21,10 @@ class GEMINIFLASH(Provider):
         history_offset: int = 10250,
         act: str = None,
     ):
-        """Initializes GEMINI
+        """Instantiates PiAI
 
         Args:
+            conversation_id (str): The conversation ID for the Pi.ai chat.
             is_conversation (bool, optional): Flag for chatting conversationally. Defaults to True.
             max_tokens (int, optional): Maximum number of tokens to be generated upon completion. Defaults to 600.
             timeout (int, optional): Http request timeout. Defaults to 30.
@@ -31,17 +35,44 @@ class GEMINIFLASH(Provider):
             history_offset (int, optional): Limit conversation history to this number of last texts. Defaults to 10250.
             act (str|int, optional): Awesome prompt key or index. (Used as intro). Defaults to None.
         """
+        self.conversation_id = conversation_id
+        self.scraper = cloudscraper.create_scraper()
+        self.url = 'https://pi.ai/api/chat'
+        self.headers = {
+            'Accept': 'text/event-stream',
+            'Accept-Encoding': 'gzip, deflate, br, zstd',
+            'Accept-Language': 'en-US,en;q=0.9,en-IN;q=0.8',
+            'Content-Type': 'application/json',
+            'DNT': '1',
+            'Origin': 'https://pi.ai',
+            'Referer': 'https://pi.ai/talk',
+            'Sec-CH-UA': '"Not)A;Brand";v="99", "Microsoft Edge";v="127", "Chromium";v="127"',
+            'Sec-CH-UA-Mobile': '?0',
+            'Sec-CH-UA-Platform': '"Windows"',
+            'Sec-Fetch-Dest': 'empty',
+            'Sec-Fetch-Mode': 'cors',
+            'Sec-Fetch-Site': 'same-origin',
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36 Edg/127.0.0.0',
+            'X-Api-Version': '3'
+        }
+        self.cookies = {
+                '__Host-session': 'Ca5SoyAMJEaaB79jj1T69',
+                '__cf_bm': 'g07oaL0jcstNfKDyZv7_YFjN0jnuBZjbMiXOWhy7V7A-1723536536-1.0.1.1-xwukd03L7oIAUqPG.OHbFNatDdHGZ28mRGsbsqfjBlpuy.b8w6UZIk8F3knMhhtNzwo4JQhBVdtYOlG0MvAw8A'
+            }
+
         self.session = requests.Session()
         self.is_conversation = is_conversation
         self.max_tokens_to_sample = max_tokens
-        self.chat_endpoint = "https://gemini-flash.developer-house.workers.dev"
+        self.stream_chunk_size = 64
         self.timeout = timeout
         self.last_response = {}
+
         self.__available_optimizers = (
             method
             for method in dir(Optimizers)
             if callable(getattr(Optimizers, method)) and not method.startswith("__")
         )
+        self.session.headers.update(self.headers)
         Conversation.intro = (
             AwesomePrompts().get_act(
                 act, raise_not_found=True, default=None, case_insensitive=True
@@ -75,7 +106,7 @@ class GEMINIFLASH(Provider):
            dict : {}
         ```json
         {
-           "text" : "How may I assist you today?"
+            "text" : "How may I assist you today?"
         }
         ```
         """
@@ -89,28 +120,38 @@ class GEMINIFLASH(Provider):
                 raise Exception(
                     f"Optimizer is not one of {self.__available_optimizers}"
                 )
-        self.session.headers.update(
-            {
-                "Content-Type": "application/json",
-                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
-            }
-        )
-        payload = {"question": conversation_prompt}
 
-        response = self.session.get(
-            self.chat_endpoint, params=payload, stream=True, timeout=self.timeout
-        )
-        if not response.ok:
-            raise exceptions.FailedToGenerateResponseError(
-                f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
+        data = {
+            'text': conversation_prompt,
+            'conversation': self.conversation_id
+        }
+
+        def for_stream():
+            response = self.scraper.post(self.url, headers=self.headers, cookies=self.cookies, json=data, stream=True, timeout=self.timeout)
+            
+            streaming_text = ""
+            for line in response.iter_lines(decode_unicode=True):
+                if line.startswith("data: "):
+                    json_data = line[6:]
+                    try:
+                        parsed_data = json.loads(json_data)
+                        if 'text' in parsed_data:
+                            streaming_text += parsed_data['text']
+                            resp = dict(text=streaming_text)
+                            self.last_response.update(resp)
+                            yield parsed_data if raw else resp
+                    except json.JSONDecodeError:
+                        continue
+            self.conversation.update_chat_history(
+                prompt, self.get_message(self.last_response)
             )
-        
-        resp = response.json()
-        self.last_response.update(resp)
-        self.conversation.update_chat_history(
-            prompt, self.get_message(self.last_response)
-        )
-        return resp
+
+        def for_non_stream():
+            for _ in for_stream():
+                pass
+            return self.last_response
+
+        return for_stream() if stream else for_non_stream()
 
     def chat(
         self,
@@ -128,13 +169,24 @@ class GEMINIFLASH(Provider):
         Returns:
             str: Response generated
         """
-        return self.get_message(
-            self.ask(
-                prompt,
-                optimizer=optimizer,
-                conversationally=conversationally,
-            )
-        )
+
+        def for_stream():
+            for response in self.ask(
+                prompt, True, optimizer=optimizer, conversationally=conversationally
+            ):
+                yield self.get_message(response).encode('utf-8').decode('utf-8')
+
+        def for_non_stream():
+            return self.get_message(
+                self.ask(
+                    prompt,
+                    False,
+                    optimizer=optimizer,
+                    conversationally=conversationally,
+                )
+            ).encode('utf-8').decode('utf-8')
+
+        return for_stream() if stream else for_non_stream()
 
     def get_message(self, response: dict) -> str:
         """Retrieves message only from response
@@ -146,7 +198,11 @@ class GEMINIFLASH(Provider):
             str: Message extracted
         """
         assert isinstance(response, dict), "Response should be of dict data-type only"
-        try:
-            return response["content"]
-        except KeyError:
-            return ""
+        return response["text"]
+
+if __name__ == '__main__':
+    from rich import print
+    ai = PiAI(conversation_id="6kti6HPbUKKWUAEpeD7vQ")  
+    response = ai.chat(input(">>> "))
+    for chunk in response:
+        print(chunk, end="", flush=True)