from smolagents import Tool
from langchain_community.llms import HuggingFaceHub


class MindMapTool(Tool):
    name = "mind_map"
    description = "Generate a mind map text structure based on a PDF summary."
    inputs = {
        "pdf_summary": {"type": "string", "description": "A concise summary of the paper or key concepts."}
    }
    output_type = "string"

    def forward(self, pdf_summary: str) -> str:
        llm = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-alpha", task="text-generation")
        prompt = (
            "Create a detailed mind map in text format from the following summary. "
            "Use bullet points and indentation to represent main topics and subtopics:\n\n"
            f"{pdf_summary}\n\nMind map:"
        )
        return llm(prompt)