from langchain.llms import LlamaCpp
from langchain.chat_models import ChatOpenAI
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.cache import SQLiteCache
import langchain

import itertools

from utils import setup_logger

from dotenv import load_dotenv

import os

# Load the .env file
load_dotenv()

OPEN_AI_KEY = os.getenv("OPEN_AI_KEY")

os.environ["OPENAI_API_KEY"] = OPEN_AI_KEY

logger = setup_logger('contr_detector_logger', 'app.log')
langchain.llm_cache = SQLiteCache(database_path=".langchain.db")


callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
llm_llama = LlamaCpp(
        # model_path="llama-2-7b.Q4_K_M.gguf",
        model_path="models/OpenOrca-Platypus2-13B-Q4_K_M.gguf",
        temperature=0,
        max_tokens=1000,
        top_p=3,
        callback_manager=callback_manager,
        verbose=True,  # Verbose is required to pass to the callback manager
    )

# TODO: move the prompt to a file to be configured
prompt_template = """
            Statement 1: {doc1}
            Statement 2: {doc2}

            Question: Are these two statements contradictory? Answer "yes" or "no".
        """

prompt = PromptTemplate.from_template(prompt_template)

llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_key=OPEN_AI_KEY)
llm_chain = LLMChain(llm=llm, prompt=prompt)

def detect_contradictions(documents, metadatas, model_type: str):
    contrs = []
    for doc1, doc2 in itertools.combinations(zip(documents, metadatas), 2):
        # print(doc1)

        doc1, meta1 = doc1
        doc2, meta2 = doc2

        if model_type == "openAI":
            llm = llm_chain
            result = llm_chain({"doc1": doc1, "doc2": doc2}, return_only_outputs=True)
            print(result)
            if "yes" in result['text'].lower():
                logger.info(f"Contradiction: {doc1} {doc2}")
                print(f"Contradiction: {doc1} {doc2}")
                contrs.append(((doc1, meta1), (doc2, meta2)))
                # break # TODO: remove
            else:
                logger.info(f"No contradiction: {doc1} {doc2}")
                print(f"No contradiction: {doc1} {doc2}")
        else: 
            llm = llm_llama

            prompt = f"""
                Statement 1: {doc1}
                Statement 2: {doc2}

                Question: Are these two statements contradictory? Answer "yes" or "no".
            """

            if "yes" in llm(prompt).lower():
                logger.info(f"Contradiction: {doc1} {doc2}")
                print(f"Contradiction: {doc1} {doc2}")
                contrs.append(((doc1, meta1), (doc2, meta2)))
            else:
                logger.info(f"No contradiction: {doc1} {doc2}")
                print(f"No contradiction: {doc1} {doc2}")
    
    print("Done with checking for contradictions")
    print(contrs)
    return contrs
langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler langchain.chains.llm.LLMChain langchain.chat_models.ChatOpenAI langchain.llms.LlamaCpp langchain.cache.SQLiteCache langchain.prompts.PromptTemplate.from_template
[((476, 489), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (487, 489), False, 'from dotenv import load_dotenv\n'), ((505, 529), 'os.getenv', 'os.getenv', (['"""OPEN_AI_KEY"""'], {}), "('OPEN_AI_KEY')\n", (514, 529), False, 'import os\n'), ((584, 632), 'utils.setup_logger', 'setup_logger', (['"""contr_detector_logger"""', '"""app.log"""'], {}), "('contr_detector_logger', 'app.log')\n", (596, 632), False, 'from utils import setup_logger\n'), ((655, 697), 'langchain.cache.SQLiteCache', 'SQLiteCache', ([], {'database_path': '""".langchain.db"""'}), "(database_path='.langchain.db')\n", (666, 697), False, 'from langchain.cache import SQLiteCache\n'), ((783, 946), 'langchain.llms.LlamaCpp', 'LlamaCpp', ([], {'model_path': '"""models/OpenOrca-Platypus2-13B-Q4_K_M.gguf"""', 'temperature': '(0)', 'max_tokens': '(1000)', 'top_p': '(3)', 'callback_manager': 'callback_manager', 'verbose': '(True)'}), "(model_path='models/OpenOrca-Platypus2-13B-Q4_K_M.gguf',\n    temperature=0, max_tokens=1000, top_p=3, callback_manager=\n    callback_manager, verbose=True)\n", (791, 946), False, 'from langchain.llms import LlamaCpp\n'), ((1340, 1385), 'langchain.prompts.PromptTemplate.from_template', 'PromptTemplate.from_template', (['prompt_template'], {}), '(prompt_template)\n', (1368, 1385), False, 'from langchain.prompts import PromptTemplate\n'), ((1393, 1479), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0)', 'model_name': '"""gpt-3.5-turbo"""', 'openai_api_key': 'OPEN_AI_KEY'}), "(temperature=0, model_name='gpt-3.5-turbo', openai_api_key=\n    OPEN_AI_KEY)\n", (1403, 1479), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1487, 1519), 'langchain.chains.llm.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt'}), '(llm=llm, prompt=prompt)\n', (1495, 1519), False, 'from langchain.chains.llm import LLMChain\n'), ((736, 768), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (766, 768), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n')]