import streamlit as st
import os
from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient
from PyPDF2 import PdfReader
# Import
#import textwrap
import openai
from langchain.llms import AzureOpenAI, OpenAI
from langchain.embeddings import OpenAIEmbeddings
from llama_index.vector_stores import RedisVectorStore
from llama_index import LangchainEmbedding
from llama_index import (
    GPTVectorStoreIndex,
    SimpleDirectoryReader,
    LLMPredictor,
    PromptHelper,
    ServiceContext,
    StorageContext
)

import sys

import logging
logging.basicConfig(stream=sys.stdout, level=logging.INFO) # logging.DEBUG for more verbose output
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
REDIS_PORT = os.getenv("REDIS_PORT", "6379")
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "")

OPENAI_API_TYPE = os.getenv("OPENAI_API_TYPE", "")
OPENAI_COMPLETIONS_ENGINE = os.getenv("OPENAI_COMPLETIONS_ENGINE", "text-davinci-003")
OPENAI_EMBEDDINGS_ENGINE = os.getenv("OPENAI_EMBEDDINGS_ENGINE", "text-embedding-ada-002")


STORAGE_CONNECTION_STRING=os.getenv("STORAGE_CONNECTION_STRING", "")
CONTAINER_NAME=os.getenv("CONTAINER_NAME", "data")

def get_embeddings():
    if OPENAI_API_TYPE=="azure":
        #currently Azure OpenAI embeddings require request for service limit increase to be useful
        #using build-in HuggingFace instead
        #from langchain.embeddings import HuggingFaceEmbeddings
        #embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

        from langchain.embeddings import OpenAIEmbeddings
        embeddings = OpenAIEmbeddings(deployment=OPENAI_EMBEDDINGS_ENGINE, chunk_size=1 )
    else:
        from langchain.embeddings import OpenAIEmbeddings
        # Init OpenAI Embeddings
        embeddings = OpenAIEmbeddings()
    return embeddings 

def get_llm():
    if OPENAI_API_TYPE=="azure":
        openai.api_type = "azure"
        openai.api_base = os.getenv("OPENAI_API_BASE")
        openai.api_version = os.getenv("OPENAI_API_VERSION")
        openai.api_key = os.getenv("OPENAI_API_KEY")
        text_model_deployment = OPENAI_COMPLETIONS_ENGINE
        from langchain.llms import AzureOpenAI
        llm = AzureOpenAI(deployment_name=text_model_deployment, model_kwargs={
            "api_key": openai.api_key,
            "api_base": openai.api_base,
            "api_type": openai.api_type,
            "api_version": openai.api_version,
        })
        #llm_predictor = LLMPredictor(llm=llm)
    else:
        from langchain.llms import OpenAI
        llm=OpenAI()
    return llm


@st.cache_resource
def get_query_engine():

    blob_service_client = BlobServiceClient.from_connection_string(STORAGE_CONNECTION_STRING)
    container_client = blob_service_client.get_container_client(container=CONTAINER_NAME)
    download_file_path = "/tmp/docs"
    isExist = os.path.exists(download_file_path)
    if not isExist:
        os.makedirs(download_file_path)

    # List the blobs in the container
    blob_list = container_client.list_blobs()
    for blob in blob_list:
        print("\t" + blob.name)
        if not os.path.exists( download_file_path+ "/" + blob.name): 
            print("\nDownloading blob to \n\t" + download_file_path+ "/" + blob.name)
            with open(file=download_file_path + "/" + blob.name, mode="wb") as download_file:
                download_file.write(container_client.download_blob(blob.name).readall())
        else:
            print("\nSkipping \n\t" + download_file_path+ "/" + blob.name)

    # load documents
    documents = SimpleDirectoryReader(download_file_path).load_data()
    print('Document ID:', documents[0].doc_id)


    from llama_index.storage.storage_context import StorageContext

    vector_store = RedisVectorStore(
        index_name="chevy_docs",
        index_prefix="llama",
        redis_url="rediss://default:{}@{}:{}".format(REDIS_PASSWORD,REDIS_HOST,REDIS_PORT),
        overwrite=True
    )

    llm_predictor = LLMPredictor(llm=get_llm())
    llm_embedding = LangchainEmbedding(get_embeddings())
    service_context = ServiceContext.from_defaults(
        llm_predictor=llm_predictor,
        embed_model=llm_embedding,
    )
    storage_context = StorageContext.from_defaults(
        vector_store=vector_store
    )
    index = GPTVectorStoreIndex.from_documents(
        documents, 
        storage_context=storage_context, 
        service_context=service_context
    )


    return index.as_query_engine()

file = open("assets/app-info.md", "r")
st.markdown(file.read())
query_engine = get_query_engine()
user_query = st.text_input("Query:", 'What types of variants are available for the Chevrolet Colorado?')
try:
    response = query_engine.query(user_query)
except Exception as e:
   response = "Error: %s" % str(e) 
st.markdown(str(response))
#print(str(response))
langchain.embeddings.OpenAIEmbeddings langchain.llms.OpenAI langchain.llms.AzureOpenAI
[((558, 616), 'logging.basicConfig', 'logging.basicConfig', ([], {'stream': 'sys.stdout', 'level': 'logging.INFO'}), '(stream=sys.stdout, level=logging.INFO)\n', (577, 616), False, 'import logging\n'), ((744, 780), 'os.getenv', 'os.getenv', (['"""REDIS_HOST"""', '"""localhost"""'], {}), "('REDIS_HOST', 'localhost')\n", (753, 780), False, 'import os\n'), ((794, 825), 'os.getenv', 'os.getenv', (['"""REDIS_PORT"""', '"""6379"""'], {}), "('REDIS_PORT', '6379')\n", (803, 825), False, 'import os\n'), ((843, 874), 'os.getenv', 'os.getenv', (['"""REDIS_PASSWORD"""', '""""""'], {}), "('REDIS_PASSWORD', '')\n", (852, 874), False, 'import os\n'), ((894, 926), 'os.getenv', 'os.getenv', (['"""OPENAI_API_TYPE"""', '""""""'], {}), "('OPENAI_API_TYPE', '')\n", (903, 926), False, 'import os\n'), ((955, 1013), 'os.getenv', 'os.getenv', (['"""OPENAI_COMPLETIONS_ENGINE"""', '"""text-davinci-003"""'], {}), "('OPENAI_COMPLETIONS_ENGINE', 'text-davinci-003')\n", (964, 1013), False, 'import os\n'), ((1041, 1104), 'os.getenv', 'os.getenv', (['"""OPENAI_EMBEDDINGS_ENGINE"""', '"""text-embedding-ada-002"""'], {}), "('OPENAI_EMBEDDINGS_ENGINE', 'text-embedding-ada-002')\n", (1050, 1104), False, 'import os\n'), ((1133, 1175), 'os.getenv', 'os.getenv', (['"""STORAGE_CONNECTION_STRING"""', '""""""'], {}), "('STORAGE_CONNECTION_STRING', '')\n", (1142, 1175), False, 'import os\n'), ((1191, 1226), 'os.getenv', 'os.getenv', (['"""CONTAINER_NAME"""', '"""data"""'], {}), "('CONTAINER_NAME', 'data')\n", (1200, 1226), False, 'import os\n'), ((4640, 4735), 'streamlit.text_input', 'st.text_input', (['"""Query:"""', '"""What types of variants are available for the Chevrolet Colorado?"""'], {}), "('Query:',\n    'What types of variants are available for the Chevrolet Colorado?')\n", (4653, 4735), True, 'import streamlit as st\n'), ((688, 728), 'logging.StreamHandler', 'logging.StreamHandler', ([], {'stream': 'sys.stdout'}), '(stream=sys.stdout)\n', (709, 728), False, 'import logging\n'), ((2701, 2768), 'azure.storage.blob.BlobServiceClient.from_connection_string', 'BlobServiceClient.from_connection_string', (['STORAGE_CONNECTION_STRING'], {}), '(STORAGE_CONNECTION_STRING)\n', (2741, 2768), False, 'from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient\n'), ((2910, 2944), 'os.path.exists', 'os.path.exists', (['download_file_path'], {}), '(download_file_path)\n', (2924, 2944), False, 'import os\n'), ((4135, 4224), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm_predictor': 'llm_predictor', 'embed_model': 'llm_embedding'}), '(llm_predictor=llm_predictor, embed_model=\n    llm_embedding)\n', (4163, 4224), False, 'from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, LLMPredictor, PromptHelper, ServiceContext, StorageContext\n'), ((4265, 4320), 'llama_index.storage.storage_context.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {'vector_store': 'vector_store'}), '(vector_store=vector_store)\n', (4293, 4320), False, 'from llama_index.storage.storage_context import StorageContext\n'), ((4347, 4463), 'llama_index.GPTVectorStoreIndex.from_documents', 'GPTVectorStoreIndex.from_documents', (['documents'], {'storage_context': 'storage_context', 'service_context': 'service_context'}), '(documents, storage_context=\n    storage_context, service_context=service_context)\n', (4381, 4463), False, 'from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, LLMPredictor, PromptHelper, ServiceContext, StorageContext\n'), ((657, 676), 'logging.getLogger', 'logging.getLogger', ([], {}), '()\n', (674, 676), False, 'import logging\n'), ((1645, 1712), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'deployment': 'OPENAI_EMBEDDINGS_ENGINE', 'chunk_size': '(1)'}), '(deployment=OPENAI_EMBEDDINGS_ENGINE, chunk_size=1)\n', (1661, 1712), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((1836, 1854), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (1852, 1854), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((1987, 2015), 'os.getenv', 'os.getenv', (['"""OPENAI_API_BASE"""'], {}), "('OPENAI_API_BASE')\n", (1996, 2015), False, 'import os\n'), ((2045, 2076), 'os.getenv', 'os.getenv', (['"""OPENAI_API_VERSION"""'], {}), "('OPENAI_API_VERSION')\n", (2054, 2076), False, 'import os\n'), ((2102, 2129), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (2111, 2129), False, 'import os\n'), ((2249, 2443), 'langchain.llms.AzureOpenAI', 'AzureOpenAI', ([], {'deployment_name': 'text_model_deployment', 'model_kwargs': "{'api_key': openai.api_key, 'api_base': openai.api_base, 'api_type': openai\n    .api_type, 'api_version': openai.api_version}"}), "(deployment_name=text_model_deployment, model_kwargs={'api_key':\n    openai.api_key, 'api_base': openai.api_base, 'api_type': openai.\n    api_type, 'api_version': openai.api_version})\n", (2260, 2443), False, 'from langchain.llms import AzureOpenAI\n'), ((2605, 2613), 'langchain.llms.OpenAI', 'OpenAI', ([], {}), '()\n', (2611, 2613), False, 'from langchain.llms import OpenAI\n'), ((2973, 3004), 'os.makedirs', 'os.makedirs', (['download_file_path'], {}), '(download_file_path)\n', (2984, 3004), False, 'import os\n'), ((3164, 3216), 'os.path.exists', 'os.path.exists', (["(download_file_path + '/' + blob.name)"], {}), "(download_file_path + '/' + blob.name)\n", (3178, 3216), False, 'import os\n'), ((3615, 3656), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', (['download_file_path'], {}), '(download_file_path)\n', (3636, 3656), False, 'from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, LLMPredictor, PromptHelper, ServiceContext, StorageContext\n')]