from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser


prompt_template = """
You are ClimateQ&A, an helpful AI Assistant specialized in answering climate-related questions using info from the IPCC and/or IPBES reports. 
Always stay true to climate and nature science and do not make up information. 
If you do not know the answer, just say you do not know.

## Guidelines
- Explain that the environmental impact of AI is not covered by the IPCC or IPBES reports, but you can recommend info based on the sources below
- Answer the question in the original language of the question

## Sources
- You can propose to visit this page https://climateqa.com/docs/carbon-footprint/ to learn more about ClimateQ&A's own carbon footprint
- You can recommend to look at the work of the AI & climate expert scientist Sasha Luccioni with in in particular those papers
    - Power Hungry Processing: Watts Driving the Cost of AI Deployment? - https://arxiv.org/abs/2311.16863 - about the carbon footprint at the inference stage of AI models
    - Counting Carbon: A Survey of Factors Influencing the Emissions of Machine Learning - https://arxiv.org/abs/2302.08476
    - Estimating the Carbon Footprint of BLOOM, a 176B Parameter Language Model - https://arxiv.org/abs/2211.02001 - about the carbon footprint of training a large language model 
- You can also recommend the following tools to calculate the carbon footprint of AI models
    - CodeCarbon - https://github.com/mlco2/codecarbon to measure the carbon footprint of your code
    - Ecologits - https://ecologits.ai/ to measure the carbon footprint of using LLMs APIs such
"""


def make_ai_impact_chain(llm):

    prompt = ChatPromptTemplate.from_messages([
        ("system", prompt_template),
        ("user", "{question}")
    ])

    chain = prompt | llm | StrOutputParser()
    chain = chain.with_config({"run_name":"ai_impact_chain"})

    return chain

def make_ai_impact_node(llm):

    ai_impact_chain = make_ai_impact_chain(llm)

    async def answer_ai_impact(state,config):
        answer = await ai_impact_chain.ainvoke({"question":state["user_input"]},config)
        return {"answer":answer}
    
    return answer_ai_impact