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



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

#分割文档
text_splitter = CharacterTextSplitter(
     separator="\n",
     chunk_size=1000,
     chunk_overlap=200,
     length_function=len
 )

texts = text_splitter.split_text("./output_1.txt")

# 嵌入模型
#embeddings = OpenAIEmbeddings()
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

# 加载数据
docsearch = FAISS.from_texts(texts, embeddings)

model_url = "http://36.103.234.50:5000"

llm = TextGen(model_url=model_url)

chain = load_qa_chain(llm)

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)
    

    partial_message = ""
    for chunk in response:
        if len(chunk) != 0:
            partial_message = partial_message + chunk
            yield partial_message 

    
    return response

gr.ChatInterface(predict).queue().launch()