1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
|
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"
models = [default_model]
@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])
|
