## Setup
# Import the necessary Libraries
import os 
os.system('python -m pip install tiktoken')
os.system('python -m pip install openai')
os.system('python -m pip install langchain')
os.system('python -m pip install langchain_community')
os.system('python -m pip install sentence-transformers')
os.system('python -m pip install chromadb')

import json
import tiktoken

import pandas as pd

from openai import OpenAI

from langchain.text_splitter import RecursiveCharacterTextSplitter

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
import gradio as gr
import uuid

# Create Client
# Initialise the client
client = OpenAI(
    base_url="https://api.endpoints.anyscale.com/v1",
    api_key=os.environ['ANYSCALE_API_KEY']
)
#Provide the model name
model_name = 'mlabonne/NeuralHermes-2.5-Mistral-7B'
# Define the embedding model and the vectorstore
embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')

# Load the persisted vectorDB
persisted_vectordb_location = './report_db'

#Create a Colelction Name
collection_name = 'reports'

# Load the persisted DB
vectorstore_persisted = Chroma(
    collection_name=collection_name,
    persist_directory=persisted_vectordb_location,
    embedding_function=embedding_model
)
retriever = vectorstore_persisted.as_retriever(
    search_type='similarity',
    search_kwargs={'k': 5}
)

# Prepare the logging functionality
log_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
log_folder = log_file.parent

scheduler = CommitScheduler(
    repo_id="Multi-Service-Agent-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 to a financial technology firm who answers client queries on extensive collection of 10-K reports from various industry players, which contain detailed information about financial performance, risk factors, market trends, and strategic initiatives. User input will have the context required by you to answer user questions.
This context will begin with the token: ###Context.
The context contains references to specific portions of a document relevant to the user query.

User questions will begin with the token: ###Question.

Please answer only using the context provided in the input. Do not mention anything about the context in your final answer.

If the answer is not found in the context, respond "I don't know".
"""

# 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):

    filter = "dataset/"+company+"-10-k-2023.pdf"
    relevant_document_chunks = vectorstore_persisted.similarity_search(user_input, k=5, filter={"source":filter})

    # Create context_for_query
    context_list = [document.page_content for document 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=model_name,
      messages=prompt,
      temperature=0
      )
      prediction = response.choices[0].message.content.strip()
    except Exception as e:
      prediction = "I don't know"
      print(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", lines=6)
#company = gr.Radio()
company = gr.Dropdown(
    ['google', 'aws', 'msft', 'IBM','Meta'],
    label='Company'
)

# Create the interface
# For the inputs parameter of Interface provide [textbox,company]
demo = gr.Interface(fn=predict, inputs=[textbox,company], outputs="text",
                 title="10K Reports Q&A System",
                 description="This web API presents an interface to ask questions on 10K Reports of companies",
                 article="Note that questions that are not relevant to 10K Reports or not within the sample documents will be answered with I don't know.",
                 concurrency_limit=16
)

demo.queue()
demo.launch()