# main.py

#####################################################################
# Amazon Bedrock - boto3
#####################################################################

import boto3

# Setup bedrock
bedrock_runtime = boto3.client(
    service_name="bedrock-runtime",
    region_name="us-east-1",
)

#####################################################################
# LLM - Amazon Bedrock LLM using LangChain
#####################################################################

from llama_index.llms import LangChainLLM
from langchain.llms import Bedrock

model_id = "anthropic.claude-v2"
model_kwargs =  { 
    "max_tokens_to_sample": 4096,
    "temperature": 0.7,
    "top_k": 250,
    "top_p": 1,
    "stop_sequences": ["\n\nHuman:"],
}

llm = Bedrock(
    client=bedrock_runtime,
    model_id=model_id,
    model_kwargs=model_kwargs
)

#####################################################################
# Embedding Model - Amazon Titan Embeddings Model using LangChain
#####################################################################

# from llama_index import LangchainEmbedding -> from llama_index.embeddings import LangchainEmbedding
# Source code - https://github.com/run-llama/llama_index/blob/main/llama_index/embeddings/__init__.py
from llama_index.embeddings import LangchainEmbedding
from langchain.embeddings import BedrockEmbeddings

# create embeddings
bedrock_embedding = BedrockEmbeddings(
    client=bedrock_runtime,
    model_id="amazon.titan-embed-text-v1",
)

# load in Bedrock embedding model from langchain
embed_model = LangchainEmbedding(bedrock_embedding)

#####################################################################
# Service Context
#####################################################################

from llama_index import ServiceContext, set_global_service_context

service_context = ServiceContext.from_defaults(
  llm=llm,
  embed_model=embed_model,
  system_prompt="You are an AI assistant answering questions."
)

set_global_service_context(service_context)

#####################################################################
# Streamlit
#####################################################################

import streamlit as st
from llama_index import SimpleDirectoryReader
from llama_index import VectorStoreIndex

st.set_page_config(
  page_title="Qlik Product Documentation 📗 Vector Embedding Index Q&A over you data 😃 ",
  page_icon="📗",
  layout="centered",
  initial_sidebar_state="auto",
  menu_items=None)

st.title("Qlik Product Documentation 📗 Vector Index Q&A over your data 😃")

@st.cache_resource(show_spinner=False)
def load_data():
  """
    Loads and indexes the data using the VectorStoreIndex.
    
    Returns:
    - VectorStoreIndex: Indexed representation of your data.
  """
  with st.spinner(
    text="Loading and indexing your data. This may take a while..."):
    reader=SimpleDirectoryReader(input_dir="./data", recursive=True)
    docs=reader.load_data()

    index=VectorStoreIndex.from_documents(docs)
    return index

# Create Index
index=load_data()

# Create Query Engine
query_engine=index.as_query_engine(similarity_top_k=3)

# Take input from the user
user_input=st.text_input("Enter Your Query", "")

# Display the input
if st.button("Submit"):
  st.write(f"Your Query: {user_input}")

  with st.spinner("Thinking..."):
    # Query the index
    result=query_engine.query(f"\n\nHuman:{user_input}\n\nAssistant:")

    # Display the results
    st.write(f"Answer: {str(result)}")

langchain.llms.Bedrock langchain.embeddings.BedrockEmbeddings
[((225, 294), 'boto3.client', 'boto3.client', ([], {'service_name': '"""bedrock-runtime"""', 'region_name': '"""us-east-1"""'}), "(service_name='bedrock-runtime', region_name='us-east-1')\n", (237, 294), False, 'import boto3\n'), ((760, 837), 'langchain.llms.Bedrock', 'Bedrock', ([], {'client': 'bedrock_runtime', 'model_id': 'model_id', 'model_kwargs': 'model_kwargs'}), '(client=bedrock_runtime, model_id=model_id, model_kwargs=model_kwargs)\n', (767, 837), False, 'from langchain.llms import Bedrock\n'), ((1410, 1495), 'langchain.embeddings.BedrockEmbeddings', 'BedrockEmbeddings', ([], {'client': 'bedrock_runtime', 'model_id': '"""amazon.titan-embed-text-v1"""'}), "(client=bedrock_runtime, model_id='amazon.titan-embed-text-v1'\n    )\n", (1427, 1495), False, 'from langchain.embeddings import BedrockEmbeddings\n'), ((1566, 1603), 'llama_index.embeddings.LangchainEmbedding', 'LangchainEmbedding', (['bedrock_embedding'], {}), '(bedrock_embedding)\n', (1584, 1603), False, 'from llama_index.embeddings import LangchainEmbedding\n'), ((1850, 1978), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm': 'llm', 'embed_model': 'embed_model', 'system_prompt': '"""You are an AI assistant answering questions."""'}), "(llm=llm, embed_model=embed_model,\n    system_prompt='You are an AI assistant answering questions.')\n", (1878, 1978), False, 'from llama_index import ServiceContext, set_global_service_context\n'), ((1984, 2027), 'llama_index.set_global_service_context', 'set_global_service_context', (['service_context'], {}), '(service_context)\n', (2010, 2027), False, 'from llama_index import ServiceContext, set_global_service_context\n'), ((2293, 2492), 'streamlit.set_page_config', 'st.set_page_config', ([], {'page_title': '"""Qlik Product Documentation 📗 Vector Embedding Index Q&A over you data 😃 """', 'page_icon': '"""📗"""', 'layout': '"""centered"""', 'initial_sidebar_state': '"""auto"""', 'menu_items': 'None'}), "(page_title=\n    'Qlik Product Documentation 📗 Vector Embedding Index Q&A over you data 😃 ',\n    page_icon='📗', layout='centered', initial_sidebar_state='auto',\n    menu_items=None)\n", (2311, 2492), True, 'import streamlit as st\n'), ((2492, 2566), 'streamlit.title', 'st.title', (['"""Qlik Product Documentation 📗 Vector Index Q&A over your data 😃"""'], {}), "('Qlik Product Documentation 📗 Vector Index Q&A over your data 😃')\n", (2500, 2566), True, 'import streamlit as st\n'), ((2569, 2606), 'streamlit.cache_resource', 'st.cache_resource', ([], {'show_spinner': '(False)'}), '(show_spinner=False)\n', (2586, 2606), True, 'import streamlit as st\n'), ((3177, 3214), 'streamlit.text_input', 'st.text_input', (['"""Enter Your Query"""', '""""""'], {}), "('Enter Your Query', '')\n", (3190, 3214), True, 'import streamlit as st\n'), ((3239, 3258), 'streamlit.button', 'st.button', (['"""Submit"""'], {}), "('Submit')\n", (3248, 3258), True, 'import streamlit as st\n'), ((3262, 3299), 'streamlit.write', 'st.write', (['f"""Your Query: {user_input}"""'], {}), "(f'Your Query: {user_input}')\n", (3270, 3299), True, 'import streamlit as st\n'), ((2781, 2856), 'streamlit.spinner', 'st.spinner', ([], {'text': '"""Loading and indexing your data. This may take a while..."""'}), "(text='Loading and indexing your data. This may take a while...')\n", (2791, 2856), True, 'import streamlit as st\n'), ((2874, 2931), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', ([], {'input_dir': '"""./data"""', 'recursive': '(True)'}), "(input_dir='./data', recursive=True)\n", (2895, 2931), False, 'from llama_index import SimpleDirectoryReader\n'), ((2971, 3008), 'llama_index.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['docs'], {}), '(docs)\n', (3002, 3008), False, 'from llama_index import VectorStoreIndex\n'), ((3308, 3333), 'streamlit.spinner', 'st.spinner', (['"""Thinking..."""'], {}), "('Thinking...')\n", (3318, 3333), True, 'import streamlit as st\n')]