import numpy as np
from langchain.prompts import PromptTemplate
from langchain.schema import StrOutputParser, BaseRetriever
from langchain.schema.runnable import RunnablePassthrough

from langchain_google_genai import ChatGoogleGenerativeAI
from trulens_eval.feedback.provider.langchain import Langchain
from trulens_eval import Tru, Feedback, TruChain
from trulens_eval.feedback import Groundedness
from trulens_eval.schema import Select
from trulens_eval.utils.serial import all_queries
from trulens_eval.utils.json import jsonify

from src.embeddings import build_base_embeddings
from src.vectordb import load_chroma
from src.reranker import build_reranker
from src.retrieval_qa import build_rerank_retriever
from src.llms import build_llm


# Setup RAG
embedding_function = build_base_embeddings()
vectordb = load_chroma(embedding_function)
reranker = build_reranker()
retriever = build_rerank_retriever(vectordb, reranker)
llm = build_llm()

QA_TEMPLATE = """You are an assistant for question-answering tasks. \
Use the following pieces of retrieved context to answer the question. \
If you don't know the answer, just say you don't know.
Question: {question}
Context: {context}
Answer:"""

prompt = PromptTemplate.from_template(QA_TEMPLATE)


def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough}
    | prompt
    | llm
    | StrOutputParser
)


# Evaluate with trulens-eval

# Define provider and database
_llm = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0)
provider = Langchain(chain=_llm)

database_url = "sqlite:///data/trulens.db"
tru = Tru(database_url=database_url, database_redact_keys=True)
# tru.reset_database()

# Using TruChain
app_json = jsonify(rag_chain)
retrievers = []
for lens in all_queries(app_json):
    try:
        comp = lens.get_sole_item(rag_chain)
        if isinstance(comp, BaseRetriever):
            retrievers.append((lens, comp))
    except Exception:
        pass

context = (
    (Select.RecordCalls + retrievers[0][0]).get_relevant_documents.rets[:].page_content
)

f_qa_relevance = Feedback(
    provider.relevance_with_cot_reasonse, name="Answer Relevance"
).on_input_output()

f_context_relevance = (
    Feedback(provider.qs_relevance_with_cot_reasons, name="Context Relevance")
    .on_input()
    .on(context)
    .aggregate(np.mean)
)

grounded = Groundedness(groundedness_provider=provider)
f_groundedness = (
    Feedback(grounded.groundedness_measure_with_cot_reasons, name="Groundedness")
    .on(context.collect())
    .on_output()
    .aggregate(grounded.grounded_statements_aggregator)
)

app_id = "Chain1"

tru_recorder = TruChain(
    rag_chain,
    app_id=app_id,
    feedbacks=[
        f_qa_relevance,
        f_context_relevance,
        f_groundedness,
    ],
)

qns = ...
for qn in qns:
    with tru_recorder as recording:
        res = rag_chain.invoke(qn)

# Results
# dashboard
tru.run_dashboard(port=8601)

# # dataframe
# records_df, feednack = tru.get_records_and_feednack(app_ids=[app_id])
# records_df.head()

langchain.prompts.PromptTemplate.from_template langchain_google_genai.ChatGoogleGenerativeAI
[((778, 801), 'src.embeddings.build_base_embeddings', 'build_base_embeddings', ([], {}), '()\n', (799, 801), False, 'from src.embeddings import build_base_embeddings\n'), ((813, 844), 'src.vectordb.load_chroma', 'load_chroma', (['embedding_function'], {}), '(embedding_function)\n', (824, 844), False, 'from src.vectordb import load_chroma\n'), ((856, 872), 'src.reranker.build_reranker', 'build_reranker', ([], {}), '()\n', (870, 872), False, 'from src.reranker import build_reranker\n'), ((885, 927), 'src.retrieval_qa.build_rerank_retriever', 'build_rerank_retriever', (['vectordb', 'reranker'], {}), '(vectordb, reranker)\n', (907, 927), False, 'from src.retrieval_qa import build_rerank_retriever\n'), ((934, 945), 'src.llms.build_llm', 'build_llm', ([], {}), '()\n', (943, 945), False, 'from src.llms import build_llm\n'), ((1205, 1246), 'langchain.prompts.PromptTemplate.from_template', 'PromptTemplate.from_template', (['QA_TEMPLATE'], {}), '(QA_TEMPLATE)\n', (1233, 1246), False, 'from langchain.prompts import PromptTemplate\n'), ((1536, 1593), 'langchain_google_genai.ChatGoogleGenerativeAI', 'ChatGoogleGenerativeAI', ([], {'model': '"""gemini-pro"""', 'temperature': '(0)'}), "(model='gemini-pro', temperature=0)\n", (1558, 1593), False, 'from langchain_google_genai import ChatGoogleGenerativeAI\n'), ((1605, 1626), 'trulens_eval.feedback.provider.langchain.Langchain', 'Langchain', ([], {'chain': '_llm'}), '(chain=_llm)\n', (1614, 1626), False, 'from trulens_eval.feedback.provider.langchain import Langchain\n'), ((1677, 1734), 'trulens_eval.Tru', 'Tru', ([], {'database_url': 'database_url', 'database_redact_keys': '(True)'}), '(database_url=database_url, database_redact_keys=True)\n', (1680, 1734), False, 'from trulens_eval import Tru, Feedback, TruChain\n'), ((1787, 1805), 'trulens_eval.utils.json.jsonify', 'jsonify', (['rag_chain'], {}), '(rag_chain)\n', (1794, 1805), False, 'from trulens_eval.utils.json import jsonify\n'), ((1834, 1855), 'trulens_eval.utils.serial.all_queries', 'all_queries', (['app_json'], {}), '(app_json)\n', (1845, 1855), False, 'from trulens_eval.utils.serial import all_queries\n'), ((2426, 2470), 'trulens_eval.feedback.Groundedness', 'Groundedness', ([], {'groundedness_provider': 'provider'}), '(groundedness_provider=provider)\n', (2438, 2470), False, 'from trulens_eval.feedback import Groundedness\n'), ((2709, 2812), 'trulens_eval.TruChain', 'TruChain', (['rag_chain'], {'app_id': 'app_id', 'feedbacks': '[f_qa_relevance, f_context_relevance, f_groundedness]'}), '(rag_chain, app_id=app_id, feedbacks=[f_qa_relevance,\n    f_context_relevance, f_groundedness])\n', (2717, 2812), False, 'from trulens_eval import Tru, Feedback, TruChain\n'), ((2155, 2226), 'trulens_eval.Feedback', 'Feedback', (['provider.relevance_with_cot_reasonse'], {'name': '"""Answer Relevance"""'}), "(provider.relevance_with_cot_reasonse, name='Answer Relevance')\n", (2163, 2226), False, 'from trulens_eval import Tru, Feedback, TruChain\n'), ((2280, 2354), 'trulens_eval.Feedback', 'Feedback', (['provider.qs_relevance_with_cot_reasons'], {'name': '"""Context Relevance"""'}), "(provider.qs_relevance_with_cot_reasons, name='Context Relevance')\n", (2288, 2354), False, 'from trulens_eval import Tru, Feedback, TruChain\n'), ((2494, 2571), 'trulens_eval.Feedback', 'Feedback', (['grounded.groundedness_measure_with_cot_reasons'], {'name': '"""Groundedness"""'}), "(grounded.groundedness_measure_with_cot_reasons, name='Groundedness')\n", (2502, 2571), False, 'from trulens_eval import Tru, Feedback, TruChain\n')]