From b12cb49c5337b65d52b2012e18feb5441b1f8754 Mon Sep 17 00:00:00 2001 From: Arslan Saleem Date: Mon, 9 Dec 2024 22:12:36 +0100 Subject: [PATCH] feat(BambooLLM): adding implementation for bamboollm updated version (#1462) * feat(BambooLLM): adding implementation for bamboollm updated version * remove exta print statement * fix: exception message * Update pandasai/__init__.py Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> --------- Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> --- pandasai/__init__.py | 6 +++--- pandasai/dataframe/base.py | 4 ++-- pandasai/helpers/request.py | 8 ++++---- pandasai/llm/bamboo_llm.py | 7 ++++--- 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/pandasai/__init__.py b/pandasai/__init__.py index b639eb1fd..8ccd2fe21 100644 --- a/pandasai/__init__.py +++ b/pandasai/__init__.py @@ -83,11 +83,11 @@ def load(dataset_path: str, virtualized=False) -> DataFrame: global _dataset_loader dataset_full_path = os.path.join(find_project_root(), "datasets", dataset_path) if not os.path.exists(dataset_full_path): - api_key = os.environ.get("PANDAAI_API_KEY", None) - api_url = os.environ.get("PANDAAI_API_URL", None) + api_key = os.environ.get("PANDASAI_API_KEY", None) + api_url = os.environ.get("PANDASAI_API_URL", None) if not api_url or not api_key: raise PandasAIApiKeyError( - "Set PANDAAI_API_URL and PANDAAI_API_KEY in environment to pull dataset from the remote server" + "Please set the PANDASAI_API_URL and PANDASAI_API_KEY environment variables to pull the dataset from the remote server." ) request_session = get_pandaai_session() diff --git a/pandasai/dataframe/base.py b/pandasai/dataframe/base.py index 12ce616e3..1fa940104 100644 --- a/pandasai/dataframe/base.py +++ b/pandasai/dataframe/base.py @@ -257,11 +257,11 @@ def push(self): ) def pull(self): - api_key = os.environ.get("PANDAAI_API_KEY", None) + api_key = os.environ.get("PANDASAI_API_KEY", None) if not api_key: raise PandasAIApiKeyError( - "Set PANDAAI_API_URL and PANDAAI_API_KEY in environment to pull dataset to the remote server" + "Set PANDASAI_API_URL and PANDASAI_API_KEY in environment to pull dataset to the remote server" ) request_session = get_pandaai_session() diff --git a/pandasai/helpers/request.py b/pandasai/helpers/request.py index fcb4bc0de..730a723c2 100644 --- a/pandasai/helpers/request.py +++ b/pandasai/helpers/request.py @@ -64,7 +64,7 @@ def make_request( url = urljoin(self._endpoint_url, self._version_path + path) if headers is None: headers = { - "Authorization": f"Bearer {self._api_key}", + "x-authorization": f"Bearer {self._api_key}", "Content-Type": "application/json", # or any other headers you need } @@ -99,11 +99,11 @@ def make_request( def get_pandaai_session(): - api_url = os.environ.get("PANDAAI_API_URL", None) - api_key = os.environ.get("PANDAAI_API_KEY", None) + api_url = os.environ.get("PANDASAI_API_URL", None) + api_key = os.environ.get("PANDASAI_API_KEY", None) if not api_url or not api_key: raise PandasAIApiKeyError( - "Set PANDAAI_API_URL and PANDAAI_API_KEY in environment to push/pull dataset to the remote server" + "Set PANDASAI_API_URL and PANDASAI_API_KEY in environment to push/pull dataset to the remote server" ) return Session(endpoint_url=api_url, api_key=api_key) diff --git a/pandasai/llm/bamboo_llm.py b/pandasai/llm/bamboo_llm.py index c7cddc6df..db2420791 100644 --- a/pandasai/llm/bamboo_llm.py +++ b/pandasai/llm/bamboo_llm.py @@ -18,9 +18,10 @@ def __init__( self._session = Session(endpoint_url=endpoint_url, api_key=api_key) def call(self, instruction: BasePrompt, _context=None) -> str: - data = instruction.to_json() - response = self._session.post("/llm/chat", json=data) - return response["data"] + response = self._session.post( + "/query", json={"prompt": instruction.to_string()} + ) + return response["answer"] @property def type(self) -> str: