from langchain_community.callbacks import get_openai_callback
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

from prompts.accountability_agent import template


class AccountabilityAgent:
    def __init__(self, model: str = "gpt-4-1106-preview", temperature: float = 0.3) -> None:
        self._prompt = PromptTemplate(input_variables=["query", "synthesis"], template=template)

        self._llm = ChatOpenAI(model=model, temperature=temperature)

        self._chain = self._prompt | self._llm

    def run(self, query: str, synthesis: str) -> tuple[str, dict]:
        with get_openai_callback() as cb:
            accountability = self._chain.invoke({"query": query, "synthesis": synthesis}).content.strip()

            tokens = cb.total_tokens
            cost = cb.total_cost

        return accountability, {"tokens": tokens, "cost": cost}