from pymongo import MongoClient
import os
import time
import gradio as gr
import requests
import traceback
import google.generativeai as genai




genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
try:
    # Initialize MongoDB python client
    MONGODB_URI = os.getenv("MONGODB_ATLAS_URI")
    client = MongoClient(MONGODB_URI, appname="devrel.content.python")

    DB_NAME = "google-ai"
    COLLECTION_NAME = "embedded_docs"
    ATLAS_VECTOR_SEARCH_INDEX_NAME = "vector_index"
    collection = client[DB_NAME][COLLECTION_NAME]

    ### Insert data about 5 individual employees
    collection.delete_many({})
    collection.insert_many([
        {
            '_id' : '54633',
            'content' : 'Employee number 54633, name  John Doe, department  Sales, location  New York, salary  100000'
        },
        {
            '_id' : '54634',
            'content' : 'Employee number 54634, name  Jane Doe, department  Marketing, location  Los Angeles, salary  120000',

        },
        {
            '_id' : '54635',
            'content' : 'Employee number 54635, name  John Smith, department  Engineering, location  San Francisco, salary  150000'
        },
        {
            '_id' : '54636',
            'content' : 'Employee number 54636, name  Jane Smith, department  Finance, location  Chicago, salary  130000'
        },
        {
            '_id' : '54637',
            'content' : 'Employee number 54637, name  John Johnson, department  HR, location  Miami, salary  110000'
        },
        {
            '_id' : '54638',
            'content' : 'Employee number 54638, name  Jane Johnson, department  Operations, location  Seattle, salary  140000'
        }
    ])


    # Exception handling to catch and display errors during the pipeline execution.
except Exception as erorr_message:
    print("An error occurred: \n" + erorr_message)

gemini_pro = genai.GenerativeModel('gemini-pro')

def embed_text(text):
    result = genai.embed_content(
    model="models/embedding-001",
    content=text,
    task_type="retrieval_document",
    title="Embedding of single string")

    return result['embedding']

def get_rag_output(context, question):
   
    template = f""" You are an hr assistant, answer in detail. Answer the question based only on the following context:
        ```
        {context}
        ```
        Question: {question}
        """
    response = gemini_pro.generate_content([template], stream=False)
    return response.text

def mongodb_vector_query(message):
    docs = collection.aggregate([
        {
            '$vectorSearch' : {
                'index' : 'vector_index',
                'queryVector' : embed_text(message),
                'path' : 'embedding',
                'numCandidates' : 10,
                'limit' : 5
            }
        },
        {
            '$project': {
                'embedding': 0
            }
        }
    ])

    return list(docs)




def get_rag(message, history):

    try:
        context = mongodb_vector_query(message)
        result = get_rag_output(context, message)
        
        # print(result)
        print_llm_text = result
        for i in range(len(print_llm_text)):
            time.sleep(0.03)
            yield print_llm_text[: i+1]
        
    
    except Exception as e:
        error_message = traceback.format_exc()
        print("An error occurred: \n" + error_message)
        yield error_message
        



def fetch_url_data(url):
    try:
        response = requests.get(url)
        response.raise_for_status()  # Raises an HTTPError if the HTTP request returned an unsuccessful status code
        return response.text
    except requests.RequestException as e:
        return f"Error: {e}"
        
# Setup Gradio interface
with gr.Blocks() as demo:
    with gr.Tab("Demo"):
       
        gr.ChatInterface(get_rag,examples=["List all employees", "Where does jane work?", "Who has the highest salary? List it"], title="Atlas Vector Search Chat",description="This small chat uses a similarity search to find relevant plots as listed above, it uses MongoDB Atlas and Google Gemini.",submit_btn="Search").queue()

    with gr.Tab("Code"):
        gr.Code(label="Code", language="python", value=fetch_url_data('https://huggingface.co/spaces/MongoDB/mongodb-gemini-rag/raw/main/app.py'))

        

    

if __name__ == "__main__":
    demo.launch()