From cd6ccdd03b0a29de4a86b68a711be3930919080f Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <hlohaus@users.noreply.github.com>
Date: Sat, 2 Dec 2023 00:34:05 +0100
Subject: Add proxy support / fix streaming

---
 g4f/Provider/Pi.py | 28 +++++++++++++---------------
 1 file changed, 13 insertions(+), 15 deletions(-)

(limited to 'g4f/Provider')

diff --git a/g4f/Provider/Pi.py b/g4f/Provider/Pi.py
index 9ecebafb..8acf3ec6 100644
--- a/g4f/Provider/Pi.py
+++ b/g4f/Provider/Pi.py
@@ -23,32 +23,30 @@ class Pi(BaseProvider):
         **kwargs
     ) -> CreateResult:
         if not scraper:
-            scraper = cls.get_scraper()
+            scraper = cls.get_scraper(proxy)
         if not conversation:
             conversation = cls.start_conversation(scraper)
         answer = cls.ask(scraper, messages, conversation)
-
-        last_answer = 0
         for line in answer:
             if "text" in line:
-                yield line["text"][last_answer:]
-                last_answer = len(line["text"])
-        
-    def get_scraper():
-        scraper = create_scraper(
+                yield line["text"]
+                        
+    def get_scraper(proxy: str):
+        return create_scraper(
             browser={
                 'browser': 'chrome',
                 'platform': 'windows',
                 'desktop': True
             },
-            sess=session()
+            headers={
+                'Accept': '*/*',
+                'Accept-Encoding': 'deflate,gzip,br',
+            },
+            proxies={
+                "https": proxy
+            }
         )
-        scraper.headers = {
-            'Accept': '*/*',
-            'Accept-Encoding': 'deflate,gzip,br',
-        }
-        return scraper
-
+        
     def start_conversation(scraper: CloudScraper):
         response = scraper.post('https://pi.ai/api/chat/start', data="{}", headers={
             'accept': 'application/json',
-- 
cgit v1.2.3