Some more PowerBI pydantic and import fixes (#3461)

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 = (
+            chain = LLMChain(
-                LLMChain(
+                llm=self.llm,
-                    llm=self.llm,
+                callback_manager=self.callback_manager,
-                    callback_manager=self.callback_manager,
+                prompt=PromptTemplate(
-                    prompt=PromptTemplate(
+                    template=QUESTION_TO_QUERY,
-                        template=QUESTION_TO_QUERY,
+                    input_variables=["tool_input", "tables", "schemas", "examples"],
-                        input_variables=["tool_input", "tables", "schemas", "examples"],
-                    ),
                 ),
             )
         else:
-            chain = (
+            chain = LLMChain(
-                LLMChain(
+                llm=self.llm,
-                    llm=self.llm,
+                prompt=PromptTemplate(
-                    prompt=PromptTemplate(
+                    template=QUESTION_TO_QUERY,
-                        template=QUESTION_TO_QUERY,
+                    input_variables=["tool_input", "tables", "schemas", "examples"],
-                        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,
             ),
         ]

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."

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
+    template: Optional[str] = QUESTION_TO_QUERY
-    examples: str = DEFAULT_FEWSHOT_EXAMPLES
+    examples: Optional[str] = DEFAULT_FEWSHOT_EXAMPLES
 
     class Config:
         """Configuration for this pydantic object."""

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)

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"

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}