import logging

logging.basicConfig(level=logging.CRITICAL)

import os
from pathlib import Path

import openai
from dotenv import load_dotenv
from langchain.chat_models import ChatOpenAI
from llama_index import (
    GPTVectorStoreIndex,
    LLMPredictor,
    ServiceContext,
    StorageContext,
    download_loader,
    load_index_from_storage,
)
from utils import CACHE, FILES, models, cls, handle_save, handle_exit, initialize, select_file

load_dotenv()
openai.api_key = os.environ["OPENAI_API_KEY"]
history = []

llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.618, model_name=models["gpt-3"], max_tokens=256))
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, chunk_size_limit=1024)


def make_index(file):
    cls()
    print("👀 Loading...")

    PDFReader = download_loader("PDFReader")
    loader = PDFReader()
    documents = loader.load_data(file=Path(FILES) / file)

    if os.path.exists(Path(CACHE) / file):
        print("📚 Index found in cache")
        return
    else:
        print("📚 Index not found in cache, creating it...")
        index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)
        index.storage_context.persist(persist_dir=Path(CACHE) / file)


def chat(file_name, index):
    while True:
        prompt = input("\n😎 Prompt: ")
        if prompt == "exit":
            handle_exit()
        elif prompt == "save":
            handle_save(str(file_name), history)

        query_engine = index.as_query_engine(response_mode="compact")

        response = query_engine.query(prompt)
        print("\n👻 Response: " + str(response))
        history.append({"user": prompt, "response": str(response)})


def ask(file_name):
    try:
        print("👀 Loading...")
        storage_context = StorageContext.from_defaults(persist_dir=Path(CACHE) / file_name)
        index = load_index_from_storage(storage_context, service_context=service_context)
        cls()
        print("✅ Ready! Let's start the conversation")
        print("ℹ️ Press Ctrl+C to exit")
        chat(file_name, index)
    except KeyboardInterrupt:
        handle_exit()


if __name__ == "__main__":
    initialize()
    file = select_file()
    if file:
        file_name = Path(file).name
        make_index(file_name)
        ask(file_name)
    else:
        print("No files found")
        handle_exit()

langchain.chat_models.ChatOpenAI
[((16, 59), 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.CRITICAL'}), '(level=logging.CRITICAL)\n', (35, 59), False, 'import logging\n'), ((444, 457), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (455, 457), False, 'from dotenv import load_dotenv\n'), ((644, 729), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm_predictor': 'llm_predictor', 'chunk_size_limit': '(1024)'}), '(llm_predictor=llm_predictor, chunk_size_limit=1024\n    )\n', (672, 729), False, 'from llama_index import GPTVectorStoreIndex, LLMPredictor, ServiceContext, StorageContext, download_loader, load_index_from_storage\n'), ((753, 758), 'utils.cls', 'cls', ([], {}), '()\n', (756, 758), False, 'from utils import CACHE, FILES, models, cls, handle_save, handle_exit, initialize, select_file\n'), ((802, 830), 'llama_index.download_loader', 'download_loader', (['"""PDFReader"""'], {}), "('PDFReader')\n", (817, 830), False, 'from llama_index import GPTVectorStoreIndex, LLMPredictor, ServiceContext, StorageContext, download_loader, load_index_from_storage\n'), ((2171, 2183), 'utils.initialize', 'initialize', ([], {}), '()\n', (2181, 2183), False, 'from utils import CACHE, FILES, models, cls, handle_save, handle_exit, initialize, select_file\n'), ((2195, 2208), 'utils.select_file', 'select_file', ([], {}), '()\n', (2206, 2208), False, 'from utils import CACHE, FILES, models, cls, handle_save, handle_exit, initialize, select_file\n'), ((551, 624), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0.618)', 'model_name': "models['gpt-3']", 'max_tokens': '(256)'}), "(temperature=0.618, model_name=models['gpt-3'], max_tokens=256)\n", (561, 624), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1099, 1177), 'llama_index.GPTVectorStoreIndex.from_documents', 'GPTVectorStoreIndex.from_documents', (['documents'], {'service_context': 'service_context'}), '(documents, service_context=service_context)\n', (1133, 1177), False, 'from llama_index import GPTVectorStoreIndex, LLMPredictor, ServiceContext, StorageContext, download_loader, load_index_from_storage\n'), ((1871, 1944), 'llama_index.load_index_from_storage', 'load_index_from_storage', (['storage_context'], {'service_context': 'service_context'}), '(storage_context, service_context=service_context)\n', (1894, 1944), False, 'from llama_index import GPTVectorStoreIndex, LLMPredictor, ServiceContext, StorageContext, download_loader, load_index_from_storage\n'), ((1953, 1958), 'utils.cls', 'cls', ([], {}), '()\n', (1956, 1958), False, 'from utils import CACHE, FILES, models, cls, handle_save, handle_exit, initialize, select_file\n'), ((2361, 2374), 'utils.handle_exit', 'handle_exit', ([], {}), '()\n', (2372, 2374), False, 'from utils import CACHE, FILES, models, cls, handle_save, handle_exit, initialize, select_file\n'), ((937, 948), 'pathlib.Path', 'Path', (['CACHE'], {}), '(CACHE)\n', (941, 948), False, 'from pathlib import Path\n'), ((1374, 1387), 'utils.handle_exit', 'handle_exit', ([], {}), '()\n', (1385, 1387), False, 'from utils import CACHE, FILES, models, cls, handle_save, handle_exit, initialize, select_file\n'), ((2124, 2137), 'utils.handle_exit', 'handle_exit', ([], {}), '()\n', (2135, 2137), False, 'from utils import CACHE, FILES, models, cls, handle_save, handle_exit, initialize, select_file\n'), ((2242, 2252), 'pathlib.Path', 'Path', (['file'], {}), '(file)\n', (2246, 2252), False, 'from pathlib import Path\n'), ((894, 905), 'pathlib.Path', 'Path', (['FILES'], {}), '(FILES)\n', (898, 905), False, 'from pathlib import Path\n'), ((1228, 1239), 'pathlib.Path', 'Path', (['CACHE'], {}), '(CACHE)\n', (1232, 1239), False, 'from pathlib import Path\n'), ((1830, 1841), 'pathlib.Path', 'Path', (['CACHE'], {}), '(CACHE)\n', (1834, 1841), False, 'from pathlib import Path\n')]