from langchain_openai import AzureChatOpenAI

from src.prompts import refiner_prompt
from src.state import AgentState


def refiner_node(state: AgentState):
    """Refines the final answer to a query by extracting the essential
    information from the verbose answer.

    Args:
        state (AgentState): The current state of the agent, containing
                            the original query and results.

    Returns:
        str: A refined answer extracted from the verbose answer.
    """
    # Create the evaluation chain
    chain = (
        refiner_prompt |
        AzureChatOpenAI(deployment_name="gpt-4o"))

    # Invoke the evaluation chain
    response = chain.invoke({
        "query": state["original_query"],
        "answer": state["results"][-1]
    })

    return {
        **state,
        "final_answer": response.content
    }