from langchain_ollama.chat_models import ChatOllama
from langchain_core.messages import SystemMessage
from langchain_core.prompts import HumanMessagePromptTemplate, ChatPromptTemplate
from langchain.tools import Tool

def query_fits_to_answer(query: str, answer: str) -> bool:

    def rate(rating: bool) -> None:
        """Rate answer as correct (True) or as incorrect (False)."""

    prompt = ChatPromptTemplate.from_messages([
        SystemMessage(content="""You are a rating machine. You rate answers as correct if they are
        1. factually correct (every statement made)
        2. fitting response to the query answering all questions prompted
        
        if the answer does not mach these criteria, rate the answer as incorrect. **Only use the rate tool. Do not answer conversationally**.
        Do not answer with <I'm sorry but I do not have the capability to perform this task for you, I am happy to help you with any other queries you may have.> or anything like it. Just use the `rate` tool."""),
        HumanMessagePromptTemplate.from_template(template="""Query:
        {query}

        Answer:
        {answer}
        """)
    ]).invoke({"query": query, "answer": answer})

    llm = ChatOllama(model="llama3-groq-tool-use:70b").bind_tools([rate])

    ai_msg = llm.invoke(prompt)

    try:
        return ai_msg.tool_calls[0]['args']['rating']
    except IndexError as e:
        print(f"\rValidation Error of <{ai_msg.content}> Retrying...")
        return query_fits_to_answer(query=query, answer=answer)

if __name__ == "__main__":
    # print(query_fits_to_answer(
    #     query="Who is Obama?",
    #     answer="Barack Obama was the 44th President of the United States, serving two terms from January 2009 to January 2017. He was a significant figure in American politics and made history as the first African American to hold the office."
    # ))
    # print(query_fits_to_answer(
    #     query="Who is Obama?",
    #     answer="Quantum computing is a revolutionary technology that uses the principles of quantum mechanics to perform calculations and operations on data. It's a fundamentally different approach from classical computing, which is based on bits (0s and 1s) and transistors."
    # ))
    # print(query_fits_to_answer(
    #     query="Who is Obama?",
    #     answer="Barack Obama was the 72th President of the United States, serving two terms from January 2005 to January 2013. He was a significant figure in American politics and made history as the first Chinese American to hold the office."
    # ))
    print(query_fits_to_answer(
        query="Who is Obama?",
        answer="Barack Obama was the 45th President of the United States, serving two terms from January 2009 to January 2017. He was a significant figure in American politics and made history as the first Chinese American to hold the office."
    ))