import streamlit as st
import requests
import json
from langchain.vectorstores import Vectara
from sentence_transformers import CrossEncoder

# Input your API keys
vectara_customer_id = "1668875253"
vectara_corpus_id = 2
vectara_api_key = 'zqt_Y3kD9bueJq3QO5t_FISVQLmgTWMDhzgMgK9Isw'

# Initialize Vectara
vectara_instance = Vectara(
    vectara_customer_id='1668875253',
    vectara_corpus_id=2,
    vectara_api_key='zqt_Y3kD9bueJq3QO5t_FISVQLmgTWMDhzgMgK9Isw',
)

# Model initialization
model = CrossEncoder('vectara/hallucination_evaluation_model')

# Streamlit app
st.title('Fact-Checked Finance RAG-Based App Using Vectara HHEM')

# Input message from the user
message = st.text_input('Enter your message')

# Button to trigger the processing
if st.button('Process'):
    # Processing logic
    corpus_key = [
        {
            "customerId": vectara_customer_id,
            "corpusId": vectara_corpus_id,
            "lexicalInterpolationConfig": {"lambda": 0.025},
        }
    ]
    data = {
        "query": [
            {
                "query": message,
                "start": 0,
                "numResults": 10,
                "contextConfig": {
                    "sentencesBefore": 2,
                    "sentencesAfter": 2,
                },
                "corpusKey": corpus_key,
                "summary": [
                    {
                        "responseLang": "eng",
                        "maxSummarizedResults": 5,
                    }
                ]
            }
        ]
    }

    headers = {
        "x-api-key": vectara_api_key,
        "customer-id": vectara_customer_id,
        "Content-Type": "application/json",
    }
    response = requests.post(
        headers=headers,
        url="https://api.vectara.io/v1/query",
        data=json.dumps(data),
    )
    if response.status_code != 200:
        st.error("Query failed")
    else:
        result = response.json()
        responses = result["responseSet"][0]["response"]
        summary = result["responseSet"][0]["summary"][0]["text"]

        res = [[r['text'], r['score']] for r in responses]
        texts = [r[0] for r in res[:5]]
        scores = [model.predict([text, summary]) for text in texts]

        # text_elements = []
        # docs = vectara_instance.similarity_search(message)
        # for source_idx, source_doc in enumerate(docs[:5]):
        #     text_elements.append(source_doc.page_content)

        ans = f"{summary}\n HHEM Scores: {scores}"
        st.text(ans)
        # st.text("Sources:")
        # for text in text_elements:
        #     st.text(text)