from llama_index.core import (
    VectorStoreIndex, 
    SimpleDirectoryReader, 
    StorageContext, 
    ServiceContext, 
    load_index_from_storage
)
import streamlit as st
from llama_index.core.node_parser import SemanticSplitterNodeParser
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.core import Settings
import pandas as pd
from llama_index.core import Document
from llama_index.core.schema import TextNode
from llama_index.vector_stores.qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from llama_index.core import VectorStoreIndex
from llama_index.core.postprocessor import (
    NERPIINodePostprocessor,
    SentenceEmbeddingOptimizer,
)
from llama_index.core import QueryBundle
from llama_index.core.schema import NodeWithScore, TextNode
from llama_index.core import PromptTemplate

embed_model = HuggingFaceEmbedding(
    model_name="nomic-ai/nomic-embed-text-v1.5",trust_remote_code=True
)
Settings.embed_model = embed_model

llm = HuggingFaceLLM(
    tokenizer_name="AgentPublic/guillaumetell-7b",
    model_name="AgentPublic/guillaumetell-7b",
)
Settings.llm = llm 

client = QdrantClient(url='https://8dd8ef0e-5b57-4276-ad0c-cfd557b7327b.us-east4-0.gcp.cloud.qdrant.io:6333',api_key='qoLzdopcAx7R1245fQcs1xQk3kWoekZnP08RzED5qRkuLC2a5xF6YA')
vector_store = QdrantVectorStore(client=client, collection_name='chat_violence_vf')
sentence_index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
query_engine = sentence_index.as_query_engine(similarity_top_k=5, llm=llm)

# use this for testing
vector_retriever = sentence_index.as_retriever(similarity_top_k=5)
pii_processor = NERPIINodePostprocessor(llm=llm)
def filter_pii_fn(**kwargs):
    # run optimizer
    query_bundle = QueryBundle(query_str=kwargs["query_str"], llm=llm)

    new_nodes = pii_processor.postprocess_nodes(
        [NodeWithScore(node=TextNode(text=kwargs["context_str"]))],
        query_bundle=query_bundle,
    )
    new_node = new_nodes[0]
    return new_node.get_content()

qa_prompt_tmpl_str = (
    '''Les informations contextuelles sont ci-dessous.\n
    ---------------------\n
    {context_str}\n
    ---------------------\n
    Vous êtes un chatbot assistant en santé mentale nommé "EmpowerHer". Votre expertise consiste exclusivement à fournir des informations et des conseils sur tout ce qui concerne la violence contre les femmes. Cela inclut leur donner un soutien psychologique et les aider à contacter des personnes avec lesquelles elles peuvent se relier. Vous ne fournissez pas d'informations en dehors de ce domaine. Si une question ne concerne pas la violence contre les femmes, répondez avec : "Je me spécialise uniquement dans l'aide aux femmes pour qu'elles se sentent bien dans leur peau. Étant donné les informations contextuelles et non les connaissances préalables,"
    répondez à la question.\n
    Question : {query_str}\n
    Réponse : '''
)

qa_prompt_tmpl = PromptTemplate(
    qa_prompt_tmpl_str, function_mappings={"context_str": filter_pii_fn}
)

query_engine.update_prompts(
    {"response_synthesizer:text_qa_template": qa_prompt_tmpl}
)
# Define a simple Streamlit app
st.title('EmpowerHer at Your service my Lady')
query = st.text_input("Enter your query")
# If the 'Submit' button is clicked
if st.button("Submit"):
    if not query.strip():
        st.error(f"Please provide the search query.")
    else:
        try:
            response = query_engine.query(query)
            st.success(response)
        except Exception as e:
            st.error(f"An error occurred: {e}")