import streamlit as st
from dotenv import load_dotenv

load_dotenv()

import os
import tempfile
from llama_index import SimpleDirectoryReader, StorageContext, LLMPredictor
from llama_index import VectorStoreIndex
from llama_index import ServiceContext
from llama_index.embeddings.langchain import LangchainEmbedding
from langchain.chat_models import ChatOpenAI
import tiktoken
from langchain.embeddings import CohereEmbeddings
import openai

os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
openai.api_key = st.secrets["OPENAI_API_KEY"]
os.environ["COHERE_API_KEY"] = st.secrets["COHERE_API_KEY"]

llm_predictor = LLMPredictor(llm = ChatOpenAI(temperature = 0, model_name = 'gpt-3.5-turbo', max_tokens = -1, openai_api_key = openai.api_key))

embed_model = LangchainEmbedding(CohereEmbeddings(model = "embed-english-light-v2.0"))

storage_context = StorageContext.from_defaults()
service_context = ServiceContext.from_defaults(llm_predictor = llm_predictor, embed_model = embed_model)

def num_tokens_from_string(string: str, encoding_name: str) -> int:
    encoding = tiktoken.encoding_for_model(encoding_name)
    num_tokens = len(encoding.encode(string))
    return num_tokens

@st.cache_resource
def preprocessing(uploaded_file):
    if uploaded_file:
        temp_dir = tempfile.TemporaryDirectory()
        file_path = os.path.join(temp_dir.name, uploaded_file.name)
        with open(file_path, "wb") as f:
            f.write(uploaded_file.read())
        document = SimpleDirectoryReader(input_files = [file_path]).load_data()
        tokens = num_tokens_from_string(document[0].text, 'gpt-3.5-turbo')
        global context
        context = document[0].text
        if tokens <= 4000:
            print('Case - A')
            return context
        else:
            print('Case - B')
            index = VectorStoreIndex.from_documents(document, service_context = service_context, storage_context = storage_context)
            global engine
            engine = index.as_query_engine(similarity_top_k = 3)
            return engine

@st.cache_resource
def run(_query_engine, query):
    if type(_query_engine) == str:
        print('Executing Case - A')
        response = openai.ChatCompletion.create(
                model = "gpt-3.5-turbo",
                messages=[
                        {"role": "system", "content": "You are a helpful assistant who answers questions given context."},
                        {"role": "user", "content": f"The question is - {query}\nThe provided context is - {_query_engine}\nAnswer the question to the best of your abilities."},
                    ]
                )
        st.write(response['choices'][0]['message']['content'])
    else:
        print('Executing Case - B')
        st.write(query_engine.query(query).response)
        return True

st.set_page_config(layout = "wide")

st.title("Document Querying")

uploaded_file = st.file_uploader('Upload your file')

query_engine = preprocessing(uploaded_file)

if query_engine:
    query = st.text_input('Enter your Query.', key = 'query_input')
    if query:
        run(query_engine, query)
langchain.embeddings.CohereEmbeddings langchain.chat_models.ChatOpenAI
[((55, 68), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (66, 68), False, 'from dotenv import load_dotenv\n'), ((860, 890), 'llama_index.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {}), '()\n', (888, 890), False, 'from llama_index import SimpleDirectoryReader, StorageContext, LLMPredictor\n'), ((909, 996), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm_predictor': 'llm_predictor', 'embed_model': 'embed_model'}), '(llm_predictor=llm_predictor, embed_model=\n    embed_model)\n', (937, 996), False, 'from llama_index import ServiceContext\n'), ((2820, 2853), 'streamlit.set_page_config', 'st.set_page_config', ([], {'layout': '"""wide"""'}), "(layout='wide')\n", (2838, 2853), True, 'import streamlit as st\n'), ((2857, 2886), 'streamlit.title', 'st.title', (['"""Document Querying"""'], {}), "('Document Querying')\n", (2865, 2886), True, 'import streamlit as st\n'), ((2904, 2940), 'streamlit.file_uploader', 'st.file_uploader', (['"""Upload your file"""'], {}), "('Upload your file')\n", (2920, 2940), True, 'import streamlit as st\n'), ((787, 837), 'langchain.embeddings.CohereEmbeddings', 'CohereEmbeddings', ([], {'model': '"""embed-english-light-v2.0"""'}), "(model='embed-english-light-v2.0')\n", (803, 837), False, 'from langchain.embeddings import CohereEmbeddings\n'), ((1080, 1122), 'tiktoken.encoding_for_model', 'tiktoken.encoding_for_model', (['encoding_name'], {}), '(encoding_name)\n', (1107, 1122), False, 'import tiktoken\n'), ((3016, 3069), 'streamlit.text_input', 'st.text_input', (['"""Enter your Query."""'], {'key': '"""query_input"""'}), "('Enter your Query.', key='query_input')\n", (3029, 3069), True, 'import streamlit as st\n'), ((644, 747), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0)', 'model_name': '"""gpt-3.5-turbo"""', 'max_tokens': '(-1)', 'openai_api_key': 'openai.api_key'}), "(temperature=0, model_name='gpt-3.5-turbo', max_tokens=-1,\n    openai_api_key=openai.api_key)\n", (654, 747), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1286, 1315), 'tempfile.TemporaryDirectory', 'tempfile.TemporaryDirectory', ([], {}), '()\n', (1313, 1315), False, 'import tempfile\n'), ((1336, 1383), 'os.path.join', 'os.path.join', (['temp_dir.name', 'uploaded_file.name'], {}), '(temp_dir.name, uploaded_file.name)\n', (1348, 1383), False, 'import os\n'), ((2198, 2537), 'openai.ChatCompletion.create', 'openai.ChatCompletion.create', ([], {'model': '"""gpt-3.5-turbo"""', 'messages': '[{\'role\': \'system\', \'content\':\n    \'You are a helpful assistant who answers questions given context.\'}, {\n    \'role\': \'user\', \'content\':\n    f"""The question is - {query}\nThe provided context is - {_query_engine}\nAnswer the question to the best of your abilities."""\n    }]'}), '(model=\'gpt-3.5-turbo\', messages=[{\'role\':\n    \'system\', \'content\':\n    \'You are a helpful assistant who answers questions given context.\'}, {\n    \'role\': \'user\', \'content\':\n    f"""The question is - {query}\nThe provided context is - {_query_engine}\nAnswer the question to the best of your abilities."""\n    }])\n', (2226, 2537), False, 'import openai\n'), ((2645, 2699), 'streamlit.write', 'st.write', (["response['choices'][0]['message']['content']"], {}), "(response['choices'][0]['message']['content'])\n", (2653, 2699), True, 'import streamlit as st\n'), ((1828, 1939), 'llama_index.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['document'], {'service_context': 'service_context', 'storage_context': 'storage_context'}), '(document, service_context=service_context,\n    storage_context=storage_context)\n', (1859, 1939), False, 'from llama_index import VectorStoreIndex\n'), ((1486, 1532), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', ([], {'input_files': '[file_path]'}), '(input_files=[file_path])\n', (1507, 1532), False, 'from llama_index import SimpleDirectoryReader, StorageContext, LLMPredictor\n')]