Skip to main content
`Ctrl`+`K`

#

# Source code for langchain.agents.agent_toolkits.powerbi.toolkit

    """Toolkit for interacting with a Power BI dataset."""
    from typing import List, Optional

    from pydantic import Field

    from langchain.agents.agent_toolkits.base import BaseToolkit
    from langchain.callbacks.base import BaseCallbackManager
    from langchain.chains.llm import LLMChain
    from langchain.prompts import PromptTemplate
    from langchain.schema import BaseLanguageModel
    from langchain.tools import BaseTool
    from langchain.tools.powerbi.prompt import QUESTION_TO_QUERY
    from langchain.tools.powerbi.tool import (
        InfoPowerBITool,
        InputToQueryTool,
        ListPowerBITool,
        QueryPowerBITool,
    )
    from langchain.utilities.powerbi import PowerBIDataset

    class PowerBIToolkit(BaseToolkit):
        """Toolkit for interacting with PowerBI dataset."""

        powerbi: PowerBIDataset = Field(exclude=True)
        llm: BaseLanguageModel = Field(exclude=True)
        examples: Optional[str] = None
        callback_manager: Optional[BaseCallbackManager] = None

        class Config:
            """Configuration for this pydantic object."""

            arbitrary_types_allowed = True

        def get_tools(self) -> List[BaseTool]:
            """Get the tools in the toolkit."""
            if self.callback_manager:
                chain = (
                    LLMChain(
                        llm=self.llm,
                        callback_manager=self.callback_manager,
                        prompt=PromptTemplate(
                            template=QUESTION_TO_QUERY,
                            input_variables=["tool_input", "tables", "schemas", "examples"],
                        ),
                    ),
                )
            else:
                chain = (
                    LLMChain(
                        llm=self.llm,
                        prompt=PromptTemplate(
                            template=QUESTION_TO_QUERY,
                            input_variables=["tool_input", "tables", "schemas", "examples"],
                        ),
                    ),
                )
            return [
                QueryPowerBITool(powerbi=self.powerbi),
                InfoPowerBITool(powerbi=self.powerbi),
                ListPowerBITool(powerbi=self.powerbi),
                InputToQueryTool(
                    powerbi=self.powerbi,
                    llm_chain=chain,
                    examples=self.examples,
                ),
            ]