import sys
import os
from http import HTTPStatus
import dashscope
from dashscope.api_entities.dashscope_response import Role
from appPublic.worker import awaitify

from key import apikey

def get_user_apikey(apiname, user):
	return apikey

class QianWen:
	def __init__(self, user):
		self.key = get_user_apikey('qianwen', user)
	
	def _generate(self, p, model, h, stream):
		h.append({
			'role':'user',
			'content':p
		})
		dashscope.api_key = self.key
		response = dashscope.Generation.call(
			model,
			messages=h,
			stream=stream,
			incremental_output=True,
			result_format='message')
		if stream:
			return response

		if response.status_code == HTTPStatus.OK:
			return response
		raise Exception(f'Qianwen Error:{response.status_code}:{resposne.message}')
	async def generate(self, p, model='qwen-turbo', h=[], chunk_func=None):
		reco = awaitify(self._generate)
		stream = True
		if chunk_func is None:
			stream = False
		resp = await reco(p, model, h, stream)
		if stream:
			for r in resp:
				chunk_func(r.output.choices[0].message.content)
		else:
			return resp.choices[0].message.content
	
a_content = ''
async def main():
	global a_content
	def f(t):
		global a_content
		a_content += t
		sys.stdout.write(t)
		sys.stdout.flush()

	h = []
	ai = QianWen('aaa')
	while True:
		print('prompt:')
		p = input()
		if p == '':
			continue
		a_content = ''
		print('answer:')
		await ai.generate(p, h=h, chunk_func=f)
		h.append({'role': 'assistant', 'content':a_content})
		print('\n')
		
if __name__ == '__main__':
	import asyncio
	asyncio.get_event_loop().run_until_complete(main())