48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
from zhipuai import ZhipuAI
|
|
from ahserver.globalEnv import ServerEnv
|
|
from appPublic.worker import awaitify
|
|
from key import apikey
|
|
from aiohttp import web
|
|
|
|
def get_user_apikey(appname, user):
|
|
return apikey
|
|
|
|
class ZhiPuAi:
|
|
def __init__(self, user):
|
|
apikey = get_user_apikey('zhipuai', user)
|
|
if apikey is None:
|
|
raise Exception(f'{user} not registered zhipuAI appkey')
|
|
self.client = ZhipuAI(api_key=apikey)
|
|
|
|
def _generate(self, p, model, h, stream):
|
|
h.append({
|
|
'role':'user',
|
|
'content':p
|
|
})
|
|
resp = self.client.chat.completions.create(model=model, messages=h, stream=stream)
|
|
return resp
|
|
|
|
def generate(self, p, model='glm-4', history=[], stream=True, chunk_func=None):
|
|
reco = awaitify(self._generate)
|
|
resp = await reco(p, model, history, stream)
|
|
if stream:
|
|
for chunk in resp:
|
|
await chunk_func(chunk.choices[0].delta.content)
|
|
else:
|
|
return resp.choices[0].message.content
|
|
|
|
def zhipuai_chat():
|
|
user = get_user()
|
|
resp = web.StreamResponse()
|
|
resp.prepare(request)
|
|
p = params_kw.get('prompt')
|
|
model = params_kw.get('model', 'glm-4')
|
|
h = params_kw.get('history', [])
|
|
stream = params_kw.get('stream', True)
|
|
ai = ZhiPuAi(user)
|
|
content = await ai.generate(p, model, h, stream, resp.write)
|
|
if stream:
|
|
return resp
|
|
return content
|
|
|