import flask
from flask import Flask, request, jsonify
from langchain import PromptTemplate, LLMChain
#from langchain.memory import StreamlitChatMessageHistory
#import numpy as np
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
#from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
from langchain import HuggingFaceHub
import os
from dotenv import load_dotenv
load_dotenv()
#from pathlib import Path
from huggingface_hub import InferenceClient
from langchain import HuggingFaceHub
import requests
import uuid
import sys

# 初始化Chatbot
HUGGINGFACEHUB_API_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')
repo_id = os.environ.get('repo_id')
#port = os.getenv('port')

llm = HuggingFaceHub(repo_id=repo_id,
                     huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
                     model_kwargs={"min_length":1024,
                                   "max_new_tokens":5632, "do_sample":True,
                                   "temperature":0.1,
                                   "top_k":50,
                                   "top_p":0.95, "eos_token_id":49155}) 

prompt_template = """
<<SYS>>You are a helpful, respectful and honest assistant. If you don't know the answer to a question, please don't share false information.In each conversation, question is placed after [INST] while your answer should be placed after [/INST].<</SYS>>
[INST] {user_question} [/INST]
assistant:
"""

llm_chain = LLMChain(llm=llm, prompt=PromptTemplate.from_template(prompt_template))

# 定义API端点
app = Flask(__name__)
@app.route('/api/chat', methods=['POST'])
def chat():
    data = request.get_json()
    #user_query = data['query']
#此处的['query']中的query可以自定义名称，例如修改为user_question，那么调用API的代码中，需要相应的使用data = {'user_question': user_query}，user_question需一致
    user_query = data['user_question'] 
    print("API Input:"+user_query)
    #temp_user_query=user_query
    # 调用Chatbot
    initial_response = llm_chain.run(user_query)
    print("API Call Output:"+initial_response)
    #st.write("AI Response")
    #st.write(initial_response)
    return jsonify({'response': initial_response})

#if __name__ == "__main__":
#    uvicorn.run(app, host='0.0.0.0', port=5000)