summaryrefslogblamecommitdiffstats
path: root/g4f/Provider/Pi.py
blob: 266647ba4ce80bd0e05ab3f0e248898f2a78c82c (plain) (tree)
1
2
3
4
5
6
7
8
9

                                  
           


                                                          
                                                                          
 
                           
                                          

                          
                   
                        







                           

                                    
                                    

                      

                                                                                              
                               
                                                                  


                                            
                                                               

                              
                                  



                                                                                    


                                         
                                  
                                                         
        
                                                                 
                  
                                            
         
                                                                               
                                  

                              
                                                                 
                     
                           
                                            

                           
                                                                                      
                                  
                                          



                                                         
        
from __future__ import annotations

import json

from ..typing import CreateResult, Messages
from .base_provider import AbstractProvider, format_prompt
from ..requests import Session, get_session_from_browser, raise_for_status

class Pi(AbstractProvider):
    url             = "https://pi.ai/talk"
    working         = True
    supports_stream = True
    _session = None
    default_model = "pi"

    @classmethod
    def create_completion(
        cls,
        model: str,
        messages: Messages,
        stream: bool,
        proxy: str = None,
        timeout: int = 180,
        conversation_id: str = None,
        webdriver: WebDriver = None,
        **kwargs
    ) -> CreateResult:
        if cls._session is None:
            cls._session = get_session_from_browser(url=cls.url, proxy=proxy, timeout=timeout)
        if not conversation_id:
            conversation_id = cls.start_conversation(cls._session)
            prompt = format_prompt(messages)
        else:
            prompt = messages[-1]["content"]
        answer = cls.ask(cls._session, prompt, conversation_id)
        for line in answer:
            if "text" in line:
                yield line["text"]
    
    @classmethod
    def start_conversation(cls, session: Session) -> str:
        response = session.post('https://pi.ai/api/chat/start', data="{}", headers={
            'accept': 'application/json',
            'x-api-version': '3'
        })
        raise_for_status(response)
        return response.json()['conversations'][0]['sid']
        
    def get_chat_history(session: Session, conversation_id: str):
        params = {
            'conversation': conversation_id,
        }
        response = session.get('https://pi.ai/api/chat/history', params=params)
        raise_for_status(response)
        return response.json()

    def ask(session: Session, prompt: str, conversation_id: str):
        json_data = {
            'text': prompt,
            'conversation': conversation_id,
            'mode': 'BASE',
        }
        response = session.post('https://pi.ai/api/chat', json=json_data, stream=True)
        raise_for_status(response)
        for line in response.iter_lines():
            if line.startswith(b'data: {"text":'):
               yield json.loads(line.split(b'data: ')[1])
            elif line.startswith(b'data: {"title":'):
               yield json.loads(line.split(b'data: ')[1])