matrix_chatgpt_bot/pandora.py
2023-05-30 09:59:25 +08:00

107 lines
3.3 KiB
Python

# https://github.com/pengzhile/pandora/blob/master/doc/HTTP-API.md
import uuid
import aiohttp
import asyncio
class Pandora:
def __init__(self, api_endpoint: str, clientSession: aiohttp.ClientSession) -> None:
self.api_endpoint = api_endpoint.rstrip("/")
self.session = clientSession
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.session.close()
async def gen_title(self, data: dict, conversation_id: str) -> None:
"""
data = {
"model": "",
"message_id": "",
}
:param data: dict
:param conversation_id: str
:return: None
"""
api_endpoint = (
self.api_endpoint + f"/api/conversation/gen_title/{conversation_id}"
)
async with self.session.post(api_endpoint, json=data) as resp:
return await resp.json()
async def talk(self, data: dict) -> None:
api_endpoint = self.api_endpoint + "/api/conversation/talk"
"""
data = {
"prompt": "",
"model": "",
"parent_message_id": "",
"conversation_id": "", # ignore at the first time
"stream": True,
}
:param data: dict
:return: None
"""
data["message_id"] = str(uuid.uuid4())
async with self.session.post(api_endpoint, json=data) as resp:
return await resp.json()
async def goon(self, data: dict) -> None:
"""
data = {
"model": "",
"parent_message_id": "",
"conversation_id": "",
"stream": True,
}
"""
api_endpoint = self.api_endpoint + "/api/conversation/goon"
async with self.session.post(api_endpoint, json=data) as resp:
return await resp.json()
async def test():
model = "text-davinci-002-render-sha-mobile"
api_endpoint = "http://127.0.0.1:8008"
async with aiohttp.ClientSession() as session:
client = Pandora(api_endpoint, session)
conversation_id = None
parent_message_id = str(uuid.uuid4())
first_time = True
async with client:
while True:
prompt = input("BobMaster: ")
if conversation_id:
data = {
"prompt": prompt,
"model": model,
"parent_message_id": parent_message_id,
"conversation_id": conversation_id,
"stream": False,
}
else:
data = {
"prompt": prompt,
"model": model,
"parent_message_id": parent_message_id,
"stream": False,
}
response = await client.talk(data)
conversation_id = response["conversation_id"]
parent_message_id = response["message"]["id"]
content = response["message"]["content"]["parts"][0]
print("ChatGPT: " + content + "\n")
if first_time:
first_time = False
data = {
"model": model,
"message_id": parent_message_id,
}
response = await client.gen_title(data, conversation_id)
if __name__ == "__main__":
asyncio.run(test())