from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

tldr_prompt_template = """
Create a mind map of the given research paper along the given lines:
1. Background: A brief overview of what's being studied in what context
2. Justification: Why the researchers conducted the study
3. Method: How the researchers arrived at the result
4. Major Findings: The main findings of the study
5. Key Results: More details about the results of the study
6. Conclusion: Significance of the findings and what they mean for future research

The above sections may differ from paper to paper, hence you may need to adjust the structure accordingly by dropping / merging one or more sections.

Here is the research paper abstract: ####{paper}####

Ensure that the outline is structured in Markdown format for clarity, facilitating its integration into documents or presentations.
"""


tldr_output_parser = StrOutputParser()
tldr_prompt = PromptTemplate(
    template=tldr_prompt_template,
    input_variables=["paper"],
)
tldr_chain = lambda model: tldr_prompt | model | tldr_output_parser