%pip install llama-index > /dev/null
%pip install llama-index-postprocessor-mixedbreadai-rerank > /dev/null
%pip install llama-index-llms-openai > /dev/null

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.response.pprint_utils import pprint_response

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

import os
from llama_index.embeddings.mixedbreadai import MixedbreadAIEmbedding

# You can visit https://www.mixedbread.ai/api-reference#quick-start-guide
# to get an api key
mixedbread_api_key = os.environ.get("MXBAI_API_KEY", "your-api-key")
model_name = "mixedbread-ai/mxbai-embed-large-v1"

mixbreadai_embeddings = MixedbreadAIEmbedding(
    api_key=mixedbread_api_key, model_name=model_name
)

# load documents
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

# build index
index = VectorStoreIndex.from_documents(
    documents=documents, embed_model=mixbreadai_embeddings
)

from llama_index.postprocessor.mixedbreadai_rerank import MixedbreadAIRerank

mixedbreadai_rerank = MixedbreadAIRerank(
    api_key=mixedbread_api_key,
    top_n=2,
    model="mixedbread-ai/mxbai-rerank-large-v1",
)

query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[mixedbreadai_rerank],
)
response = query_engine.query(
    "What did Sam Altman do in this essay?",
)

pprint_response(response, show_source=True)

query_engine = index.as_query_engine(
    similarity_top_k=2,
)
response = query_engine.query(
    "What did Sam Altman do in this essay?",
)

pprint_response(response, show_source=True)