import utils
import os

import openai
from llama_index import SimpleDirectoryReader
from llama_index import Document
from llama_index import VectorStoreIndex
from llama_index import ServiceContext
from llama_index.llms import OpenAI

from llama_index.embeddings import HuggingFaceEmbedding


openai.api_key = utils.get_openai_api_key()

if __name__ == "__main__":

    documents = SimpleDirectoryReader(
        input_files=["./raw_documents/HI_knowledge_base.pdf"]
    ).load_data()

    document = Document(text="\n\n".join([doc.text for doc in documents]))

    ### gpt-4-1106-preview
    ### gpt-3.5-turbo-1106 / gpt-3.5-turbo
    llm = OpenAI(model="gpt-3.5-turbo-1106", temperature=0.1)
    embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

    service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)
    index = VectorStoreIndex.from_documents([document], service_context=service_context)

    query_engine = index.as_query_engine()

    response = query_engine.query(
        ("Intermediate and Long Term Care (ILTC) services are for those who need further care and"
         "treatment after discharge from the hospital, who may need assistance with their activities of"
         "daily living. This can be through"
        )
    )
    print(str(response))