from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains.question_answering import load_qa_chain
from langchain.llms import OpenAI
from gradio import gradio as gr
from langchain.chat_models import ChatOpenAI

from langchain import PromptTemplate, LLMChain
from langchain.llms import TextGen
from langchain.cache import InMemoryCache

from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.schema import (
    AIMessage,
    HumanMessage,
    SystemMessage
)

import time
import langchain
import os
OPENAI_API_KEY=os.getenv('OPENAI_API_KEY')

# 嵌入模型
#embeddings = OpenAIEmbeddings()
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-large-en")

# 加载数据
#docsearch = FAISS.from_texts(texts, embeddings)
docsearch = FAISS.load_local("./faiss_index", embeddings)
template="您是回答ANSYS软件使用查询的得力助手,所有回复必需用中文"
chain = load_qa_chain(OpenAI(temperature=0,model_name="gpt-3.5-turbo"), chain_type="stuff",verbose=True)

def predict(message, history):
    history_langchain_format = []
    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))
    history_langchain_format.append(HumanMessage(content=message))
    docs = docsearch.similarity_search(message)
    response = chain.run(input_documents=docs, question=message + template)
       
    partial_message = ""
    for chunk in response:
        if len(chunk[0]) != 0:
            time.sleep(0.1)
            partial_message = partial_message + chunk[0]
            yield partial_message    

langchain.llm_cache = InMemoryCache() 
   
gr.ChatInterface(predict,
    textbox=gr.Textbox(placeholder="请提问关于ANSYS软件的问题", container=False, scale=7),
    title="欢迎使用ANSYS软件AI机器人",
    examples=["你是谁?", "请介绍一下Fluent 软件", "create-bounding-box","ANSYS Fluent Architecture","ANSYS Fluent 的软件架构是怎么样的"],
                description='本AI助手为并行公司实验性产品，回答的内容由大模型推理，如回复的内容跟实际情况有偏差请理解').queue().launch(debug=True,auth=('paratera', 'paratera@2023'))