from langchain.chains import LLMChain
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate)
from chains.azure_openai import CustomAzureOpenAI
from prompts.related_question import system_template, human_template
from config import OPENAI_API_TYPE, OPENAI_API_VERSION, OPENAI_API_KEY, OPENAI_API_BASE, DEPLOYMENT_ID, API_KEY
from langchain.chat_models import ChatOpenAI

class RelatedQuestion(LLMChain):
    prompt = ChatPromptTemplate.from_messages(
        [SystemMessagePromptTemplate.from_template(
            system_template),
         HumanMessagePromptTemplate.from_template(human_template)
         ])
    llm = CustomAzureOpenAI(deployment_name=DEPLOYMENT_ID,
                    openai_api_type=OPENAI_API_TYPE,
                    openai_api_base=OPENAI_API_BASE,
                    openai_api_version=OPENAI_API_VERSION,
                    openai_api_key=OPENAI_API_KEY,
                    temperature=0.0)

if __name__ == "__main__":
    rel = RelatedQuestion()
    inputs = "Hello bot"
    outputs = "Hello! How can I assist you today?"
    import re
    pattern = "\d. {*.?}"
    
    res = rel.predict(inputs=inputs, outputs=outputs)
    print(res)
    out = list(map(lambda x: x.split('. ')[-1], res.split('\n')))
    results = [[a] for a in out]
    print(results)