## Setup
# Import the necessary Libraries

import os
import uuid
import json
import tiktoken
from datasets import load_dataset

import gradio as gr
import pandas as pd

from openai import OpenAI

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_community.vectorstores import Chroma

from huggingface_hub import CommitScheduler
from pathlib import Path

# Create Client
client = OpenAI(
    base_url="https://api.endpoints.anyscale.com/v1",
    api_key=os.environ['ANYSCALE_API_KEY']
)

# Define the embedding model and the vectorstore
embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')

collection_name = '10-K_reports'


# Load the persisted vectorDB
vectorstore_persisted = Chroma(
    collection_name=collection_name,
    persist_directory='./10k_reports_db',
    embedding_function=embedding_model
)

# Prepare the logging functionality

log_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
log_folder = log_file.parent

scheduler = CommitScheduler(
    repo_id="reports_qna-rag-logs",
    repo_type="dataset",
    folder_path=log_folder,
    path_in_repo="data",
    every=2
)

# Define the Q&A system message
qna_system_message = """
You are a financial analyst in a financial technology firm. Your task is to provide actionable insights from queries on annual reports.

User input will have the context necessary for you to answer their questions.
This context will begin with the token: ###Context.
The context contains references to specific portions of a document relevant to the user query.

When crafting your response:
1. Select only context relevant to answer the question.
2. User questions will begin with the token: ###Question.
3. If the question is irrelevant to your role, respond with "I am not sure how to respond to that"

Adhere to the following guidelines
- Please answer only using the context provided in the input. Do not mention anything about the context in your final answer.
- Your response should only be about the question asked and nothing else.
- Answer only using the context provided.
- Do not mention anything about the context in your final answer.
- If the answer is not found in the context, it is very very important for you to respond with "I don't know."
- Do not make up sources.
"""




# Define the user message template
qna_user_message_template = """
###Context
Here are some documents that are relevant to the question mentioned below.
{context}

###Question
{question}
"""



# Define the predict function that runs when 'Submit' is clicked or when a API request is made
def predict(user_input,company):

    match company:
        case "IBM":
            comp = "IBM"
        case "Meta":
            comp = "Meta"
        case "Amazon Workspace":
            comp = "aws"
        case "Google":
            comp = "google"
        case "Microsoft":
            comp = "msft"
    
    filter = "dataset/"+comp+"-10-k-2023.pdf"
    retriever = vectorstore_persisted.as_retriever(
            search_type='similarity',
            search_kwargs={'k': 5, "filter": {"source":filter}})

    # Create context_for_query
    relevant_document_chunks = retriever.get_relevant_documents(user_input)
    context_list = [d.page_content + str(d.metadata['page']) for d in relevant_document_chunks]
    context_for_query = ". ".join(context_list)

    # Create messages
    prompt = [
      {'role':'system', 'content': qna_system_message},
      {'role': 'user', 'content': qna_user_message_template.format(
           context=context_for_query,
           question=user_input
          )
      }
    ]

    # Get response from the LLM
    try:
        response = client.chat.completions.create(
            model='mistralai/Mistral-7B-Instruct-v0.1',
            messages=prompt,
            temperature=0
        )

        prediction = response.choices[0].message.content.strip()
    except Exception as e:
        prediction = f'Error:\n {e}'


    # While the prediction is made, log both the inputs and outputs to a local log file
    # While writing to the log file, ensure that the commit scheduler is locked to avoid parallel
    # access

    with scheduler.lock:
        with log_file.open("a") as f:
            f.write(json.dumps(
                {
                    'user_input': user_input,
                    'retrieved_context': context_for_query,
                    'model_response': prediction
                }
            ))
            f.write("\n")

    return prediction

# Set-up the Gradio UI
# Add text box and radio button to the interface
# The radio button is used to select the company 10k report in which the context needs to be retrieved.

textbox = gr.Textbox(placeholder="Enter your query here", label="Input query")
company = gr.Radio(["IBM", "Meta", "Amazon Workspace", "Google", "Microsoft"], label="Which company?")

model_output = gr.Label(label="Response")

# Create the interface
# For the inputs parameter of Interface provide [textbox,company]
demo = gr.Interface(
    fn=predict,
    inputs=[textbox,company],
    outputs=model_output,
    title="Company Financial Analyzer",
    description="This API allows you to analyze key information from 10-K reports of companies",
    allow_flagging="auto",
    concurrency_limit=8
)

demo.queue()
demo.launch()