from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import CharacterTextSplitter
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.memory import ConversationBufferMemor
from langchain.schema import AIMessage, HumanMessage
from langchain import PromptTemplate, LLMChain
from langchain.llms import TextGen
from langchain.cache import InMemoryCache


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)


chain = load_qa_chain(OpenAI(temperature=0,model_name="gpt-3.5-turbo", verbose=True), chain_type="stuff",verbose=True)

prompt = "您是回答所有ANSYS软件使用查询的得力助手,如果所问的内容不在范围内，请回答“您提的问题不在本知识库内，请重新提问”，所有问题必需用中文回答"

def predict(message, history):
    history_openai_format = []
    for human, assistant in history:
        history_openai_format.append({"role": "system", "content": prompt })
        history_openai_format.append({"role": "user", "content": human })
        history_openai_format.append({"role": "assistant", "content":assistant})
    history_openai_format.append({"role": "user", "content": message})
    docs = docsearch.similarity_search(message)
    response = chain.run(input_documents=docs, question=message)
   
    
    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="请输入您的问题", container=False, scale=7),
    title="欢迎使用ANSYS软件AI机器人",
    examples=["你是谁?", "请介绍一下Fluent 软件的用户界面说明", "请用关于春天写一首100字的诗","数学题：小红有3元钱，小红买了2斤香蕉，香蕉的价格是每斤1元。问小红一共花了多少钱?","请用表格做一份学生课程表"],
                description="🦊请避免输入有违公序良俗的问题，模型可能无法回答不合适的问题🐇",).queue().launch()