from langchain.retrievers import AmazonKendraRetriever
from langchain.chains import ConversationalRetrievalChain
from langchain import SagemakerEndpoint
from langchain.llms.sagemaker_endpoint import LLMContentHandler
from langchain.prompts import PromptTemplate
import sys
import json
import os


class bcolors:
    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKCYAN = "\033[96m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


MAX_HISTORY_LENGTH = 5


def build_chain():
    region = os.environ["AWS_REGION"]
    kendra_index_id = os.environ["KENDRA_INDEX_ID"]
    endpoint_name = os.environ["FALCON_40B_ENDPOINT"]
    language_code = os.environ["LANGUAGE_CODE"]

    class ContentHandler(LLMContentHandler):
        content_type = "application/json"
        accepts = "application/json"

        def transform_input(self, prompt: str, model_kwargs: dict) -> bytes:
            input_str = json.dumps({"inputs": prompt, "parameters": model_kwargs})
            return input_str.encode("utf-8")

        def transform_output(self, output: bytes) -> str:
            response_json = json.loads(output.read().decode("utf-8"))
            return response_json[0]["generated_text"]

    content_handler = ContentHandler()

    llm = SagemakerEndpoint(
        endpoint_name=endpoint_name,
        region_name=region,
        model_kwargs={
            "temperature": 0.8,
            "max_new_tokens": 512,
            "do_sample": True,
            "top_p": 0.9,
            "repetition_penalty": 1.03,
            "stop": ["\nUser:", "<|endoftext|>", "</s>"],
        },
        content_handler=content_handler,
    )

    retriever = AmazonKendraRetriever(
        index_id=kendra_index_id,
        region_name=region,
        top_k=1,
        attribute_filter={
            "EqualsTo": {
                "Key": "_language_code",
                "Value": {"StringValue": language_code},
            }
        },
    )

    prompt_template = """
  システム: システムは資料から抜粋して質問に答えます。資料にない内容には答えず、正直に「わかりません」と答えます。

  {context}

  上記の資料に基づいて以下の質問について資料から抜粋して回答を生成します。資料にない内容には答えず「わかりません」と答えます。
  ユーザー: {question}
  システム:
  """
    PROMPT = PromptTemplate(
        template=prompt_template, input_variables=["context", "question"]
    )

    condense_qa_template = """
  次のような会話とフォローアップの質問に基づいて、フォローアップの質問を独立した質問に言い換えてください。

  フォローアップの質問: {question}
  独立した質問:"""
    standalone_question_prompt = PromptTemplate.from_template(condense_qa_template)

    qa = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        condense_question_prompt=standalone_question_prompt,
        return_source_documents=True,
        verbose=True,
        combine_docs_chain_kwargs={"prompt": PROMPT},
    )
    return qa


def run_chain(chain, prompt: str, history=[]):
    return chain({"question": prompt, "chat_history": history})


if __name__ == "__main__":
    chat_history = []
    qa = build_chain()
    print(bcolors.OKBLUE + "Hello! How can I help you?" + bcolors.ENDC)
    print(
        bcolors.OKCYAN
        + "Ask a question, start a New search: or CTRL-D to exit."
        + bcolors.ENDC
    )
    print(">", end=" ", flush=True)
    for query in sys.stdin:
        if query.strip().lower().startswith("new search:"):
            query = query.strip().lower().replace("new search:", "")
            chat_history = []
        elif len(chat_history) == MAX_HISTORY_LENGTH:
            chat_history.pop(0)
        result = run_chain(qa, query, chat_history)
        chat_history.append((query, result["answer"]))
        print(bcolors.OKGREEN + result["answer"] + bcolors.ENDC)
        if "source_documents" in result:
            print(bcolors.OKGREEN + "Sources:")
            for d in result["source_documents"]:
                print(d.metadata["source"])
        print(bcolors.ENDC)
        print(
            bcolors.OKCYAN
            + "Ask a question, start a New search: or CTRL-D to exit."
            + bcolors.ENDC
        )
        print(">", end=" ", flush=True)
    print(bcolors.OKBLUE + "Bye" + bcolors.ENDC)

langchain.SagemakerEndpoint langchain.chains.ConversationalRetrievalChain.from_llm langchain.retrievers.AmazonKendraRetriever langchain.prompts.PromptTemplate.from_template langchain.prompts.PromptTemplate
[((1327, 1604), 'langchain.SagemakerEndpoint', 'SagemakerEndpoint', ([], {'endpoint_name': 'endpoint_name', 'region_name': 'region', 'model_kwargs': "{'temperature': 0.8, 'max_new_tokens': 512, 'do_sample': True, 'top_p': 0.9,\n    'repetition_penalty': 1.03, 'stop': ['\\nUser:', '<|endoftext|>', '</s>']}", 'content_handler': 'content_handler'}), "(endpoint_name=endpoint_name, region_name=region,\n    model_kwargs={'temperature': 0.8, 'max_new_tokens': 512, 'do_sample': \n    True, 'top_p': 0.9, 'repetition_penalty': 1.03, 'stop': ['\\nUser:',\n    '<|endoftext|>', '</s>']}, content_handler=content_handler)\n", (1344, 1604), False, 'from langchain import SagemakerEndpoint\n'), ((1731, 1915), 'langchain.retrievers.AmazonKendraRetriever', 'AmazonKendraRetriever', ([], {'index_id': 'kendra_index_id', 'region_name': 'region', 'top_k': '(1)', 'attribute_filter': "{'EqualsTo': {'Key': '_language_code', 'Value': {'StringValue': language_code}}\n    }"}), "(index_id=kendra_index_id, region_name=region, top_k=1,\n    attribute_filter={'EqualsTo': {'Key': '_language_code', 'Value': {\n    'StringValue': language_code}}})\n", (1752, 1915), False, 'from langchain.retrievers import AmazonKendraRetriever\n'), ((2227, 2312), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'template': 'prompt_template', 'input_variables': "['context', 'question']"}), "(template=prompt_template, input_variables=['context',\n    'question'])\n", (2241, 2312), False, 'from langchain.prompts import PromptTemplate\n'), ((2482, 2532), 'langchain.prompts.PromptTemplate.from_template', 'PromptTemplate.from_template', (['condense_qa_template'], {}), '(condense_qa_template)\n', (2510, 2532), False, 'from langchain.prompts import PromptTemplate\n'), ((2543, 2766), 'langchain.chains.ConversationalRetrievalChain.from_llm', 'ConversationalRetrievalChain.from_llm', ([], {'llm': 'llm', 'retriever': 'retriever', 'condense_question_prompt': 'standalone_question_prompt', 'return_source_documents': '(True)', 'verbose': '(True)', 'combine_docs_chain_kwargs': "{'prompt': PROMPT}"}), "(llm=llm, retriever=retriever,\n    condense_question_prompt=standalone_question_prompt,\n    return_source_documents=True, verbose=True, combine_docs_chain_kwargs={\n    'prompt': PROMPT})\n", (2580, 2766), False, 'from langchain.chains import ConversationalRetrievalChain\n'), ((989, 1047), 'json.dumps', 'json.dumps', (["{'inputs': prompt, 'parameters': model_kwargs}"], {}), "({'inputs': prompt, 'parameters': model_kwargs})\n", (999, 1047), False, 'import json\n')]