From e74af803e7add7790b366545f1fe6c80934de798 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <hlohaus@users.noreply.github.com>
Date: Fri, 1 Dec 2023 23:56:12 +0100
Subject: Update pi provider

---
 g4f/Provider/Pi.py | 93 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 93 insertions(+)
 create mode 100644 g4f/Provider/Pi.py

(limited to 'g4f/Provider/Pi.py')

diff --git a/g4f/Provider/Pi.py b/g4f/Provider/Pi.py
new file mode 100644
index 00000000..9ecebafb
--- /dev/null
+++ b/g4f/Provider/Pi.py
@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from ..typing import CreateResult, Messages
+from .base_provider import BaseProvider, format_prompt
+
+import json
+from cloudscraper import CloudScraper, session, create_scraper
+
+class Pi(BaseProvider):
+    url             = "https://chat-gpt.com"
+    working         = True
+    supports_stream = True
+
+    @classmethod
+    def create_completion(
+        cls,
+        model: str,
+        messages: Messages,
+        stream: bool,
+        proxy: str = None,
+        scraper: CloudScraper = None,
+        conversation: dict = None,
+        **kwargs
+    ) -> CreateResult:
+        if not scraper:
+            scraper = cls.get_scraper()
+        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(
+            browser={
+                'browser': 'chrome',
+                'platform': 'windows',
+                'desktop': True
+            },
+            sess=session()
+        )
+        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',
+            'x-api-version': '3'
+        })
+        if 'Just a moment' in response.text:
+            raise RuntimeError('Error: Cloudflare detected')
+        return Conversation(
+            response.json()['conversations'][0]['sid'],
+            response.cookies
+        )
+        
+    def get_chat_history(scraper: CloudScraper, conversation: Conversation):
+        params = {
+            'conversation': conversation.sid,
+        }
+        response = scraper.get('https://pi.ai/api/chat/history', params=params, cookies=conversation.cookies)
+        if 'Just a moment' in response.text:
+            raise RuntimeError('Error: Cloudflare detected')
+        return response.json()
+
+    def ask(scraper: CloudScraper, messages: Messages, conversation: Conversation):
+        json_data = {
+            'text': format_prompt(messages),
+            'conversation': conversation.sid,
+            'mode': 'BASE',
+        }
+        response = scraper.post('https://pi.ai/api/chat', json=json_data, cookies=conversation.cookies, stream=True)
+        
+        for line in response.iter_lines(chunk_size=1024, decode_unicode=True):
+            if 'Just a moment' in line:
+                raise RuntimeError('Error: Cloudflare detected')
+            if line.startswith('data: {"text":'):
+               yield json.loads(line.split('data: ')[1])
+            if line.startswith('data: {"title":'):
+               yield json.loads(line.split('data: ')[1])
+                
+class Conversation():
+    def __init__(self, sid: str, cookies):
+        self.sid = sid
+        self.cookies = cookies
+        
\ No newline at end of file
-- 
cgit v1.2.3