import sys

import chromadb
import pandas
import sqlite3
from langchain.embeddings import OpenAIEmbeddings
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.document_loaders import TextLoader
from langchain.document_loaders import WikipediaLoader
from langchain.retrievers.multi_query import MultiQueryRetriever
import langchain
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.schema import AIMessage, HumanMessage, SystemMessage
from langchain.cache import InMemoryCache
from langchain import PromptTemplate
import os
import openai
from langchain.prompts import (
    ChatPromptTemplate,
    PromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.schema import AIMessage, HumanMessage, SystemMessage

os.environ["OPENAI_API_KEY"] = "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"
openai.api_key = "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"
api_key = "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"
llm = OpenAI()
# chat = ChatOpenAI(openai_api_key=api_key, temperature=0)
embedding_function = OpenAIEmbeddings()


def us_constitution_helper(question):
    loader = TextLoader("some_data/US_Constitution.txt")
    documents = loader.load()

    text_splitter = CharacterTextSplitter.from_tiktoken_encoder(chunk_size=500)
    docs = text_splitter.split_documents(documents)

    _embedding_function = OpenAIEmbeddings()
    db = Chroma.from_documents(
        docs, _embedding_function, persist_directory="./US_Constitution"
    )
    db.persist()

    chat = ChatOpenAI(openai_api_key=api_key, temperature=0)
    compressor = LLMChainExtractor.from_llm(chat)

    compression_retriever = ContextualCompressionRetriever(
        base_compressor=compressor, base_retriever=db.as_retriever()
    )

    compressed_docs = compression_retriever.get_relevant_documents(question)

    return compressed_docs[0].page_content


print(us_constitution_helper("What is the 13th Amendment?"))

langchain.document_loaders.TextLoader langchain.text_splitter.CharacterTextSplitter.from_tiktoken_encoder langchain.llms.OpenAI langchain.chat_models.ChatOpenAI langchain.vectorstores.Chroma.from_documents langchain.embeddings.OpenAIEmbeddings langchain.retrievers.document_compressors.LLMChainExtractor.from_llm
[((1247, 1255), 'langchain.llms.OpenAI', 'OpenAI', ([], {}), '()\n', (1253, 1255), False, 'from langchain.llms import OpenAI\n'), ((1336, 1354), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (1352, 1354), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((1408, 1451), 'langchain.document_loaders.TextLoader', 'TextLoader', (['"""some_data/US_Constitution.txt"""'], {}), "('some_data/US_Constitution.txt')\n", (1418, 1451), False, 'from langchain.document_loaders import TextLoader\n'), ((1503, 1562), 'langchain.text_splitter.CharacterTextSplitter.from_tiktoken_encoder', 'CharacterTextSplitter.from_tiktoken_encoder', ([], {'chunk_size': '(500)'}), '(chunk_size=500)\n', (1546, 1562), False, 'from langchain.text_splitter import CharacterTextSplitter\n'), ((1642, 1660), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (1658, 1660), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((1670, 1762), 'langchain.vectorstores.Chroma.from_documents', 'Chroma.from_documents', (['docs', '_embedding_function'], {'persist_directory': '"""./US_Constitution"""'}), "(docs, _embedding_function, persist_directory=\n    './US_Constitution')\n", (1691, 1762), False, 'from langchain.vectorstores import Chroma\n'), ((1801, 1850), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'openai_api_key': 'api_key', 'temperature': '(0)'}), '(openai_api_key=api_key, temperature=0)\n', (1811, 1850), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1868, 1900), 'langchain.retrievers.document_compressors.LLMChainExtractor.from_llm', 'LLMChainExtractor.from_llm', (['chat'], {}), '(chat)\n', (1894, 1900), False, 'from langchain.retrievers.document_compressors import LLMChainExtractor\n')]