diff --git a/langchain/agents/agent_toolkits/powerbi/toolkit.py b/langchain/agents/agent_toolkits/powerbi/toolkit.py index 0f810224..00056563 100644 --- a/langchain/agents/agent_toolkits/powerbi/toolkit.py +++ b/langchain/agents/agent_toolkits/powerbi/toolkit.py @@ -35,24 +35,20 @@ class PowerBIToolkit(BaseToolkit): 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"], - ), + 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"], - ), + chain = LLMChain( + llm=self.llm, + prompt=PromptTemplate( + template=QUESTION_TO_QUERY, + input_variables=["tool_input", "tables", "schemas", "examples"], ), ) return [ @@ -60,8 +56,8 @@ class PowerBIToolkit(BaseToolkit): InfoPowerBITool(powerbi=self.powerbi), ListPowerBITool(powerbi=self.powerbi), InputToQueryTool( - powerbi=self.powerbi, llm_chain=chain, + powerbi=self.powerbi, examples=self.examples, ), ] diff --git a/langchain/tools/powerbi/prompt.py b/langchain/tools/powerbi/prompt.py index 363b937d..c9d8a8dc 100644 --- a/langchain/tools/powerbi/prompt.py +++ b/langchain/tools/powerbi/prompt.py @@ -58,5 +58,5 @@ BAD_REQUEST_RESPONSE = ( "Bad request. Please ask the question_to_query_powerbi tool to provide the query." ) BAD_REQUEST_RESPONSE_ESCALATED = "You already tried this, please try a different query." - +SCHEMA_ERROR_RESPONSE = "Bad request, are you sure the table name is correct?" UNAUTHORIZED_RESPONSE = "Unauthorized. Try changing your authentication, do not retry." diff --git a/langchain/tools/powerbi/tool.py b/langchain/tools/powerbi/tool.py index f095d13b..67efe423 100644 --- a/langchain/tools/powerbi/tool.py +++ b/langchain/tools/powerbi/tool.py @@ -147,8 +147,8 @@ class InputToQueryTool(BaseTool): """ # noqa: E501 llm_chain: LLMChain powerbi: PowerBIDataset = Field(exclude=True) - template: str = QUESTION_TO_QUERY - examples: str = DEFAULT_FEWSHOT_EXAMPLES + template: Optional[str] = QUESTION_TO_QUERY + examples: Optional[str] = DEFAULT_FEWSHOT_EXAMPLES class Config: """Configuration for this pydantic object.""" diff --git a/langchain/utilities/powerbi.py b/langchain/utilities/powerbi.py index 1a60c55a..89247f37 100644 --- a/langchain/utilities/powerbi.py +++ b/langchain/utilities/powerbi.py @@ -4,23 +4,20 @@ from __future__ import annotations import logging import os -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union +from typing import Any, Dict, Iterable, List, Optional, Union import aiohttp import requests from aiohttp import ServerTimeoutError +from azure.core.credentials import TokenCredential +from azure.core.exceptions import ClientAuthenticationError from pydantic import BaseModel, Field, root_validator from requests.exceptions import Timeout -from langchain.tools.powerbi.prompt import BAD_REQUEST_RESPONSE, UNAUTHORIZED_RESPONSE +from langchain.tools.powerbi.prompt import SCHEMA_ERROR_RESPONSE, UNAUTHORIZED_RESPONSE _LOGGER = logging.getLogger(__name__) -if TYPE_CHECKING: - from azure.core.exceptions import ClientAuthenticationError - from azure.identity import ChainedTokenCredential - from azure.identity._internal import InteractiveCredential - BASE_URL = os.getenv("POWERBI_BASE_URL", "https://api.powerbi.com/v1.0/myorg/datasets/") @@ -36,7 +33,7 @@ class PowerBIDataset(BaseModel): dataset_id: str table_names: List[str] group_id: Optional[str] = None - credential: Optional[Union[ChainedTokenCredential, InteractiveCredential]] = None + credential: Optional[TokenCredential] = None token: Optional[str] = None impersonated_user_name: Optional[str] = None sample_rows_in_table_info: int = Field(default=1, gt=0, le=10) @@ -144,7 +141,7 @@ class PowerBIDataset(BaseModel): continue except Exception as exc: # pylint: disable=broad-exception-caught if "bad request" in str(exc).lower(): - return BAD_REQUEST_RESPONSE + return SCHEMA_ERROR_RESPONSE if "unauthorized" in str(exc).lower(): return UNAUTHORIZED_RESPONSE return str(exc) @@ -167,7 +164,7 @@ class PowerBIDataset(BaseModel): continue except Exception as exc: # pylint: disable=broad-exception-caught if "bad request" in str(exc).lower(): - return BAD_REQUEST_RESPONSE + return SCHEMA_ERROR_RESPONSE if "unauthorized" in str(exc).lower(): return UNAUTHORIZED_RESPONSE return str(exc) diff --git a/poetry.lock b/poetry.lock index 95ef1572..0e54edc0 100644 --- a/poetry.lock +++ b/poetry.lock @@ -571,7 +571,7 @@ name = "azure-core" version = "1.26.4" description = "Microsoft Azure Core Library for Python" category = "main" -optional = true +optional = false python-versions = ">=3.7" files = [ {file = "azure-core-1.26.4.zip", hash = "sha256:075fe06b74c3007950dd93d49440c2f3430fd9b4a5a2756ec8c79454afc989c6"}, @@ -9298,4 +9298,4 @@ qdrant = ["qdrant-client"] [metadata] lock-version = "2.0" python-versions = ">=3.8.1,<4.0" -content-hash = "1dd0c2f259c674c2f3f2e2212459bf7a056f72f25e0a271194ad7ac8f70a3ac1" +content-hash = "b629046308d7f32d4f456972ff669a383c6d349fcf1c89e6e167a74b28cbb458" diff --git a/pyproject.toml b/pyproject.toml index 0633b7d7..3d554d26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ SQLAlchemy = ">1.3,<3" requests = "^2" PyYAML = ">=5.4.1" numpy = "^1" +azure-core = {version = "^1.26.4"} tqdm = {version = ">=4.48.0", optional = true} openapi-schema-pydantic = "^1.2" faiss-cpu = {version = "^1", optional = true}