Skip to content

Commit 159df68

Browse files
cannonLCKcannonLCK
authored
Initial upload
1 parent 9978af5 commit 159df68

File tree

3 files changed

+365
-0
lines changed

3 files changed

+365
-0
lines changed

‎main.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
'''
2+
main.py
3+
Reverse proxy for wrtn
4+
support system messages.
5+
'''
6+
from os import path
7+
import time
8+
import json
9+
import wrtn
10+
import random
11+
from aiohttp import web
12+
13+
Wrtn = wrtn.WrtnAPI(debug=False)
14+
15+
from os import path
16+
json_path = path.join(path.dirname(__file__), "./wrtn.json")
17+
18+
async def handle_request(request):
19+
'''login and retrieve refresh_token, and update json'''
20+
with open(json_path, "r", encoding='utf-8') as file:
21+
json_data = json.load(file)
22+
random_account = random.choice(json_data)
23+
random_account['key'] = await Wrtn.login(random_account)
24+
# Write the updated JSON data back to the file
25+
26+
with open(json_path, "w", encoding='utf-8') as file:
27+
json.dump(json_data, file, indent=4)
28+
29+
model = lambda id : {"id": id, "object": "model", "created": int(time.time() * 1000),
30+
"owned_by": "", "permission": [], "root": "", "parent": None}
31+
if request.method != 'POST':
32+
data = {
33+
"object": "list",
34+
"data": [
35+
model("gpt-3.5-turbo"),
36+
model("gpt-4"),
37+
]
38+
}
39+
# print(data)
40+
return web.json_response(data)
41+
else:
42+
# Handle POST requests here
43+
data = await request.json()
44+
preprompt = [
45+
"[Use proper language regarding on a context.]"
46+
]
47+
system = {
48+
"role":"system",
49+
"content": ""
50+
}
51+
system['content'] += ' '.join(preprompt)
52+
oldmsg = [system] + data['messages'][:-1]
53+
msg = data['messages'][-1]['content']
54+
print(oldmsg)
55+
response = web.StreamResponse()
56+
response.headers['Content-Type'] = 'application/json'
57+
await response.prepare(request)
58+
temp_res_text = ""
59+
try:
60+
chunk = {
61+
"id": f"chatcmpl-{(str(random.random())[2:])}",
62+
"created": int(time.time() * 1000),
63+
"object": "chat.completion.chunk",
64+
"model": data["model"],
65+
"choices": [{ "delta": {"role": "assistant"}, "finish_reason": None, "index": 0 }]
66+
}
67+
await response.write(f"data: {json.dumps(chunk)}\n\n".encode())
68+
print(chunk)
69+
async for res_text in Wrtn.chat_by_json(await Wrtn.make_chatbot(),
70+
msg=msg,
71+
oldmsg=oldmsg,
72+
model=data['model']
73+
):
74+
temp_res_text = res_text
75+
chunk = {
76+
"id": f"chatcmpl-{(str(random.random())[2:])}",
77+
"created": int(time.time() * 1000),
78+
"object": "chat.completion.chunk",
79+
"model": data["model"],
80+
"choices": [{ "delta": {"content": res_text}, "finish_reason": None, "index": 0 }]
81+
}
82+
await response.write(f"data: {json.dumps(chunk)}\n\n".encode())
83+
# print(chunk)
84+
except Exception as ex:
85+
print(f"Error: {str(ex)}\n Chunk: {temp_res_text}")
86+
finally:
87+
chunk = {
88+
"id": f"chatcmpl-{(str(random.random())[2:])}",
89+
"created": int(time.time() * 1000),
90+
"object": "chat.completion.chunk",
91+
"model": data["model"],
92+
"choices": [{ "delta": {}, "finish_reason": "stop", "index": 0 }]
93+
}
94+
await response.write(f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode())
95+
await response.write_eof()
96+
return response
97+
98+
app = web.Application()
99+
app.router.add_route('*', '/{path:.*}', handle_request)
100+
web.run_app(app, host='127.0.0.1', port = 41323)

‎util.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
'''Simple wrapper util class for API requesting'''
2+
import aiohttp
3+
from dataclasses import dataclass, field
4+
from typing import Callable, Optional
5+
6+
@dataclass
7+
class Http:
8+
url: str = field(default_factory=str)
9+
header: dict = field(default_factory=dict)
10+
query: dict = field(default_factory=dict)
11+
payload: dict = field(default_factory=dict)
12+
13+
class APIRequester:
14+
'''Simple class for get, post, delete'''
15+
# session = requests.Session()
16+
# async_session = aiohttp.ClientSession()
17+
def __init__(self, base_url:str, debug=False):
18+
self.debug = debug
19+
self.base_url = base_url
20+
21+
async def _handle_response(self, response, callback:Callable):
22+
if self.debug:
23+
print(response.status)
24+
if 200 <= response.status <= 299:
25+
if "application/json" in response.headers["Content-Type"]:
26+
data = await response.json()
27+
return callback(data) if callback is not None else data
28+
else:
29+
raise Exception(response.status)
30+
31+
async def _handle_event_stream(self, response, callback:Callable):
32+
if self.debug:
33+
print(response.status)
34+
if 200 <= response.status <= 299:
35+
async for data in response.content.iter_chunks():
36+
yield callback(data)
37+
# yield callback(data) if callback is not None else data
38+
else:
39+
raise Exception(response.status)
40+
41+
async def _make_request(self, method:str, info:Http, callback = None, timeout:int = 5, **kwargs):
42+
if self.debug:
43+
print(f"{method.upper()} request to URL: {self.base_url + info.url}\nParams: {info}\nAdditional: {kwargs}")
44+
45+
async with aiohttp.ClientSession() as session:
46+
async with session.request(
47+
method=method,
48+
url=self.base_url + info.url,
49+
headers=info.header,
50+
params=info.query,
51+
json=info.payload,
52+
timeout=timeout, **kwargs
53+
) as response:
54+
return await self._handle_response(response, callback)
55+
56+
async def get(self, info:Http, callback:Optional[Callable] = None):
57+
return await self._make_request("get", info=info, callback=callback)
58+
59+
async def post(self, info:Http, callback:Optional[Callable] = None):
60+
return await self._make_request("post", info=info, callback=callback, timeout=6000000)
61+
62+
async def delete(self, info:Http, callback:Optional[Callable] = None):
63+
return await self._make_request("delete", info=info, callback=callback)
64+
65+
async def stream(self, info:Http, callback:Callable):
66+
'''stream: params = query, json = payload'''
67+
async with aiohttp.ClientSession() as session:
68+
async with session.post(self.base_url + info.url, json=info.payload, params=info.query, headers=info.header, timeout=6000000) as response:
69+
# return await self._handle_response_stream(response)
70+
async for res_text in self._handle_event_stream(response, callback=callback):
71+
yield res_text

‎wrtn.py‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
'''
2+
wrtn.py
3+
reverse proxy api for wrtn
4+
'''
5+
import json
6+
from util import Http, APIRequester
7+
from fake_useragent import UserAgent
8+
import inspect
9+
10+
class ChunkDecoder:
11+
def __init__(self):
12+
self.buffer = ""
13+
14+
def decode(self, chunk_tuple):
15+
for chunk in chunk_tuple:
16+
if(isinstance(chunk, bytes)):
17+
return self.decode_chunks(chunk)
18+
19+
def decode_chunks(self, chunk_byte):
20+
chunk_str = self.buffer + chunk_byte.decode(encoding="utf-8", errors="ignore")
21+
chunks = chunk_str.split('\n\n')
22+
last_chunk = chunks[-1]
23+
self.buffer = "" if last_chunk == "" else last_chunk
24+
json_objects = [json.loads((line.split(": ", 1)[1]))
25+
for line in chunks[:-1] if line.strip().startswith('data: {')]
26+
res_text = "".join(obj['chunk'] for obj in json_objects if 'chunk' in obj if obj['chunk'] is not None)
27+
return res_text
28+
29+
#TODO: Divide WrtnAPI to Chat, Studio, Tool/ChatBot
30+
class WrtnAPI:
31+
'''
32+
WrtnAPI: needs json path
33+
'''
34+
def __init__(self, debug = False) -> None:
35+
self.debug = debug
36+
self.header = {
37+
"Init": {
38+
"Refresh": "",
39+
"Content-Type": "application/json"
40+
},
41+
"Auth": {
42+
"User-Agent": UserAgent().safari
43+
},
44+
}
45+
self.wrtn_requester = APIRequester("https://api.wow.wrtn.ai", debug=debug)
46+
self.wrtn_studio_requester = APIRequester("https://studio-api.wow.wrtn.ai", debug=debug)
47+
48+
async def login(self, account):
49+
response = await self.wrtn_requester.post(
50+
info=Http(
51+
url="/auth/local",
52+
header=self.header["Init"],
53+
payload={
54+
"email": account['id'],
55+
"password": account['pw']
56+
},
57+
),
58+
)
59+
self.header["Init"] |= {"Refresh": response['data']['refreshToken']}
60+
return response['data']['refreshToken']
61+
62+
async def access_token(self):
63+
'''refresh access_token with refresh_key'''
64+
response = await self.wrtn_requester.post(
65+
info=Http(
66+
url="/auth/refresh",
67+
header=self.header["Init"],
68+
),
69+
)
70+
return {"Authorization": "Bearer " + response['data']['accessToken']}
71+
72+
async def chat_by_json(self, _id: str, msg: str, oldmsg: list[dict], model:str='gpt-3.5-turbo'):
73+
'''chat method with json. currently: gpt-3.5 and gpt-4, context length is unknown'''
74+
75+
self.header["Auth"] |= await self.access_token()
76+
decoder = ChunkDecoder()
77+
async for response in self.wrtn_studio_requester.stream(
78+
info=Http(
79+
url=f"/store/chat-bot/{_id}/generate",
80+
header=self.header["Auth"],
81+
payload={
82+
"content": msg,
83+
"model": model,
84+
"oldMessages": oldmsg,
85+
},
86+
), callback=decoder.decode
87+
):
88+
yield response
89+
await self.delete_chatbot(_id)
90+
91+
async def get_chatbot(self) -> list:
92+
'''loads all chatbots from this user'''
93+
94+
95+
self.header["Auth"] |= await self.access_token()
96+
response = await self.wrtn_studio_requester.get(
97+
info=Http(
98+
url="/studio/chat-bot",
99+
header=self.header["Auth"]
100+
),
101+
)
102+
return response['data']
103+
104+
async def delete_chatbot(self, _id: str):
105+
'''delete a chatbot with specific id. You must delete your chatbot after use.'''
106+
107+
108+
self.header["Auth"] |= await self.access_token()
109+
response = await self.wrtn_studio_requester.delete(
110+
info=Http(
111+
url=f"/studio/chat-bot/{_id}",
112+
header=self.header["Auth"],
113+
),
114+
)
115+
return response
116+
117+
async def get_tool(self) -> list:
118+
'''loads all tools from this user'''
119+
120+
121+
self.header["Auth"] |= await self.access_token()
122+
response = await self.wrtn_studio_requester.get(
123+
info=Http(
124+
url="/studio/tool",
125+
header=self.header["Auth"],
126+
),
127+
)
128+
return response['data']
129+
130+
async def delete_tool(self, _id: str):
131+
'''delete a tool with specific id. You must delete your chatbot after use.'''
132+
133+
134+
self.header["Auth"] |= await self.access_token()
135+
response = await self.wrtn_studio_requester.delete(
136+
info=Http(
137+
url=f"/studio/tool/{_id}",
138+
header=self.header["Auth"],
139+
),
140+
)
141+
return response
142+
143+
async def make_chatbot(self):
144+
'''
145+
make a chatbot with specific option.
146+
Must register a chatbot with passing an toxicity test, to use a json prompt with system role.
147+
'''
148+
# toolList = await self.get_tool()
149+
# for tool in toolList['toolList']:
150+
# await self.delete_tool(tool['id'])
151+
# chatBotList = await self.get_chatbot()
152+
# for chatBot in chatBotList['chatBotList']:
153+
# await self.delete_chatbot(chatBot['id'])
154+
155+
156+
self.header["Auth"] |= await self.access_token()
157+
payload = {
158+
"difficulty":"hard",
159+
"icon":"faceSmile",
160+
"title":" ",
161+
"description":" ",
162+
"category":[" "],
163+
"firstMessage":" ",
164+
"selectTypeForExampleQuestion":" ",
165+
"exampleQuestion":[""],
166+
"promptForEasy":{
167+
"role":"","personality":"","requirement":""
168+
},
169+
"promptForDifficult":"",
170+
"userName":" ",
171+
"isDeleted":False,
172+
# "additionalInformation": "",
173+
"isTemporarySave":False,
174+
"openType":"비공개",
175+
"priceType":"무료",
176+
}
177+
response = await self.wrtn_studio_requester.post(
178+
info=Http(
179+
url="/studio/chat-bot",
180+
header=self.header["Auth"],
181+
),
182+
)
183+
184+
_id = response['data']['chatBotList'][0]['id']
185+
payload["userId"] = response['data']['chatBotList'][0]['userId']
186+
payload["chatBotId"] = _id
187+
chatbot = await self.wrtn_studio_requester.post(
188+
info=Http(
189+
url=f"/studio/chat-bot/{_id}",
190+
payload=payload,
191+
header=self.header["Auth"],
192+
),
193+
)
194+
return chatbot['data']['chatBotList'][0]['id']

0 commit comments

Comments
 (0)