from marqo import Client
import pandas as pd
import numpy as np

from langchain_openai import OpenAI
from langchain.docstore.document import Document
from langchain.chains import LLMChain

from dotenv import load_dotenv

from utilities import (
    load_data,
    extract_text_from_highlights,
    qna_prompt,
    predict_ce,
    get_sorted_inds
)

load_dotenv()

if __name__ == "__main__":

    #############################################################
    #       STEP 0: Install Marqo
    #############################################################

    # run the following docker commands from the terminal to start marqo
    # docker rm -f marqo
    # docker pull marqoai/marqo:2.0.0
    # docker run --name marqo -it -p 8882:8882 --add-host host.docker.internal:host-gateway marqoai/marqo:2.0.0


    #############################################################
    #       STEP 1: Setup Marqo
    #############################################################


    mq = Client()
    index_name = "iron-docs"

    # (optinally) delete if it already exists
    try:
        mq.index(index_name).delete()
    except:
        pass

    # we can set some specific settings for the index. if they are not provided, sensible defaults are used
    index_settings = {
        "model": "flax-sentence-embeddings/all_datasets_v4_MiniLM-L6",
        "normalizeEmbeddings": True,
        "textPreprocessing": {
            "splitLength": 3,
            "splitOverlap": 1,
            "splitMethod": "sentence"
        },
    }

    # create the index with custom settings
    mq.create_index(index_name, settings_dict=index_settings)

    #############################################################
    #       STEP 2: Load the data
    #############################################################

    df = load_data()

    # turn the data into a dict for indexing
    documents = df.to_dict(orient='records')

    #############################################################
    #       STEP 3: Index the data
    #############################################################

    # index the documents
    indexing = mq.index(index_name).add_documents(documents, tensor_fields=["cleaned_text"], client_batch_size=64)

    #############################################################
    #       STEP 4: Search the data
    #############################################################

    # try a generic search
    q = "what is the rated voltage"

    results = mq.index(index_name).search(q)
    print(results['hits'][0])

    #############################################################
    #       STEP 5: Make it chatty
    #############################################################

    highlights, texts = extract_text_from_highlights(results, token_limit=150)
    docs = [Document(page_content=f"Source [{ind}]:" + t) for ind, t in enumerate(texts)]
    llm = OpenAI(temperature=0.9)
    chain_qa = LLMChain(llm=llm, prompt=qna_prompt())
    llm_results = chain_qa.invoke({"summaries": docs, "question": results['query']}, return_only_outputs=True)
    print(llm_results['text'])

    #############################################################
    #       STEP 6: Score the references
    #############################################################

    score_threshold = 0.20
    top_k = 3
    scores = predict_ce(llm_results['text'], texts)
    inds = get_sorted_inds(scores)
    scores = scores.cpu().numpy()
    scores = [np.round(s[0], 2) for s in scores]
    references = [(str(np.round(scores[i], 2)), texts[i]) for i in inds[:top_k] if scores[i] > score_threshold]
    df_ref = pd.DataFrame(references, columns=['score', 'sources'])
    print(df_ref)

langchain_openai.OpenAI langchain.docstore.document.Document
[((349, 362), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (360, 362), False, 'from dotenv import load_dotenv\n'), ((984, 992), 'marqo.Client', 'Client', ([], {}), '()\n', (990, 992), False, 'from marqo import Client\n'), ((1812, 1823), 'utilities.load_data', 'load_data', ([], {}), '()\n', (1821, 1823), False, 'from utilities import load_data, extract_text_from_highlights, qna_prompt, predict_ce, get_sorted_inds\n'), ((2727, 2781), 'utilities.extract_text_from_highlights', 'extract_text_from_highlights', (['results'], {'token_limit': '(150)'}), '(results, token_limit=150)\n', (2755, 2781), False, 'from utilities import load_data, extract_text_from_highlights, qna_prompt, predict_ce, get_sorted_inds\n'), ((2882, 2905), 'langchain_openai.OpenAI', 'OpenAI', ([], {'temperature': '(0.9)'}), '(temperature=0.9)\n', (2888, 2905), False, 'from langchain_openai import OpenAI\n'), ((3331, 3369), 'utilities.predict_ce', 'predict_ce', (["llm_results['text']", 'texts'], {}), "(llm_results['text'], texts)\n", (3341, 3369), False, 'from utilities import load_data, extract_text_from_highlights, qna_prompt, predict_ce, get_sorted_inds\n'), ((3381, 3404), 'utilities.get_sorted_inds', 'get_sorted_inds', (['scores'], {}), '(scores)\n', (3396, 3404), False, 'from utilities import load_data, extract_text_from_highlights, qna_prompt, predict_ce, get_sorted_inds\n'), ((3613, 3667), 'pandas.DataFrame', 'pd.DataFrame', (['references'], {'columns': "['score', 'sources']"}), "(references, columns=['score', 'sources'])\n", (3625, 3667), True, 'import pandas as pd\n'), ((2794, 2839), 'langchain.docstore.document.Document', 'Document', ([], {'page_content': "(f'Source [{ind}]:' + t)"}), "(page_content=f'Source [{ind}]:' + t)\n", (2802, 2839), False, 'from langchain.docstore.document import Document\n'), ((3453, 3470), 'numpy.round', 'np.round', (['s[0]', '(2)'], {}), '(s[0], 2)\n', (3461, 3470), True, 'import numpy as np\n'), ((2946, 2958), 'utilities.qna_prompt', 'qna_prompt', ([], {}), '()\n', (2956, 2958), False, 'from utilities import load_data, extract_text_from_highlights, qna_prompt, predict_ce, get_sorted_inds\n'), ((3511, 3533), 'numpy.round', 'np.round', (['scores[i]', '(2)'], {}), '(scores[i], 2)\n', (3519, 3533), True, 'import numpy as np\n')]