from environs import Env
from pathlib import Path
from langchain.chains import LLMChain
from langchain.chains.base import Chain
from example import EmbeddingExampleSelector
from langchain.base_language import BaseLanguageModel
from langchain.output_parsers.boolean import BooleanOutputParser
from langchain.prompts import FewShotPromptTemplate, PromptTemplate


def load_toxicity_check_chain(env: Env, llm: BaseLanguageModel) -> Chain:
    e = EmbeddingExampleSelector(
        env.str('EMBEDDING_SERVICE_ENDPOINT'),
        'chatbot_toxicity',
        input_fn=lambda x: {'query': x['instruction']},
        output_fn=lambda x: {'instruction': x['instruction'], 'toxicity': x['toxicity']},
    )

    prompt = FewShotPromptTemplate(
        example_selector=e,
        example_prompt=PromptTemplate.from_template(Path("prompts/toxicity_example.tpl").read_text()),
        prefix=Path("prompts/toxicity_prefix.tpl").read_text(),
        suffix=Path("prompts/toxicity_suffix.tpl").read_text(), 
        input_variables=["instruction"],
        output_parser=BooleanOutputParser(),
    )
    return LLMChain(llm=llm, prompt=prompt)