from langchain.chains.combine_documents.refine import RefineDocumentsChain
from langchain.chains.qa_with_sources import load_qa_with_sources_chain

from template import template, template_many, template_many_out_put
from langchain import LLMChain
from langchain.llms import OpenAI, AzureOpenAI
from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from langchain.prompts import PromptTemplate

def LLMChain_test_outputs(context, result_count):

    llm = OpenAI(model_name="text-davinci-003", engine="text-davinci-003",temperature=0)
    prompt = PromptTemplate(template=template, input_variables=['context'])
    chain = LLMChain(llm=llm, prompt=prompt)
    return chain.run(
        {
            'context': context,
        }
    )

def LLMChain_test(context):
    llm = OpenAI(model_name="text-davinci-003", engine="text-davinci-003",temperature=0)
    prompt = PromptTemplate(template=template_many, input_variables=['context','positive', 'negative','indifference'])
    chain = LLMChain(llm=llm, prompt=prompt)
    return chain.run(
        {
            'context': context,
            'positive': '正面',
            'negative': '负面',
            'indifference': '无关'
        }
    )

def LLMChain_test_many(context,count):
    llm = OpenAI(model_name="text-davinci-003", engine="text-davinci-003", temperature=0,max_tokens=3000,verbose=True)
    prompt = PromptTemplate(template=template_many,
                            input_variables=['count','context', 'positive', 'negative', 'indifference'],
                            )
    chain = LLMChain(llm=llm, prompt=prompt)
    return chain.run(
        {
            'count': count,
            'context': context,
            'positive': '正面',
            'negative': '负面',
            'indifference': '无关'
        })

def LLMChain_test_many_output(context,count):
    llm = OpenAI(model_name="text-davinci-003", engine="text-davinci-003", temperature=0,max_tokens=3000,verbose=True)

    response_schemas = [
        ResponseSchema(name="index", description="句子的序号作为数组的下标"),
        ResponseSchema(name="emotion", description="每个句子的情绪，枚举值：正面、负面、无关"),
        ResponseSchema(name="description", description="对每个句子的总结"),
        ResponseSchema(name="count", description="单独key整体的正面、负面、无关的条数，并且只对整体"),
        ResponseSchema(name="summary", description="给出所有句子整体的中文简体总结")
    ]

    # 初始化解析器
    output_parser = StructuredOutputParser.from_response_schemas(response_schemas)

    format_instructions = output_parser.get_format_instructions()

    prompt = PromptTemplate(template=template_many_out_put,
                            input_variables=['count','context', 'positive', 'negative', 'indifference'],
                            partial_variables={"format_instructions": format_instructions},)

    promptValue = prompt.format(**{
            'count': count,
            'context': context,
            'positive': '正面',
            'negative': '负面',
            'indifference': '无关'
        })
    llm_output = llm(promptValue)
    # return output_parser.parse(llm_output)
    return llm_output