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