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 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)

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



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(SystemMessage(content=system))
        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="请输入您的问题", container=False, scale=7),
    title="欢迎使用ANSYS软件AI机器人",
    examples=["你是谁?", "请介绍一下Fluent 软件的用户界面说明", "create-bounding-box","ANSYS Fluent Architecture"],
                description="🦊请避免输入有违公序良俗的问题，模型可能无法回答不合适的问题🐇",).queue().launch()