#!/usr/bin/python
# -*- coding: UTF-8 -*-
import os

import jieba as jb
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import DirectoryLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import TokenTextSplitter
from langchain.vectorstores import Chroma


def init():
    files = ['2022.txt']
    for file in files:
        with open(f"./data/{file}", 'r', encoding='utf-8') as f:
            data = f.read()

        cut_data = " ".join([w for w in list(jb.cut(data))])
        cut_file = f"./data/cut/cut_{file}"
        with open(cut_file, 'w') as f:
            f.write(cut_data)


def load():
    loader = DirectoryLoader('./data/cut', glob='**/*.txt')
    docs = loader.load()
    text_splitter = TokenTextSplitter(chunk_size=1000, chunk_overlap=0)
    docs_texts = text_splitter.split_documents(docs)
    api_key = os.environ.get('OPENAI_API_KEY')
    embeddings = OpenAIEmbeddings(openai_api_key=api_key)
    vectordb = Chroma.from_documents(docs_texts, embeddings, persist_directory='./data/cut/')
    vectordb.persist()
    openai_ojb = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
    chain = ConversationalRetrievalChain.from_llm(openai_ojb, vectordb.as_retriever())
    return chain

chain = load()

def get_ans(question):
    chat_history = []
    result = chain({
        'chat_history': chat_history,
        'question': question,
    })
    return result['answer']

if __name__ == '__main__':
    s = input('please input:')
    while s != 'exit':
        ans = get_ans(s)
        print(ans)
        s = input('please input:')