promptlayer 1.0.71__py3-none-any.whl → 1.0.73__py3-none-any.whl

promptlayer/utils.py

@@ -5,19 +5,35 @@ import functools
 import json
 import logging
 import os
-import sys
 import types
+from contextlib import asynccontextmanager
 from copy import deepcopy
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Coroutine, Dict, List, Optional, Union
 from uuid import uuid4
 
 import httpx
 import requests
+import urllib3
+import urllib3.util
 from ably import AblyRealtime
 from ably.types.message import Message
+from centrifuge import (
+    Client,
+    PublicationContext,
+    SubscriptionEventHandler,
+    SubscriptionState,
+)
 from opentelemetry import context, trace
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+)
 
+from promptlayer import exceptions as _exceptions
 from promptlayer.types import RequestLog
 from promptlayer.types.prompt_template import (
     GetPromptTemplate,
@@ -28,8 +44,6 @@ from promptlayer.types.prompt_template import (
 )
 
 # Configuration
-# TODO(dmu) MEDIUM: Use `PROMPTLAYER_` prefix instead of `_PROMPTLAYER` suffix
-URL_API_PROMPTLAYER = os.environ.setdefault("URL_API_PROMPTLAYER", "https://api.promptlayer.com")
 RERAISE_ORIGINAL_EXCEPTION = os.getenv("PROMPTLAYER_RE_RAISE_ORIGINAL_EXCEPTION", "False").lower() == "true"
 RAISE_FOR_STATUS = os.getenv("PROMPTLAYER_RAISE_FOR_STATUS", "False").lower() == "true"
 DEFAULT_HTTP_TIMEOUT = 5
@@ -37,7 +51,9 @@ DEFAULT_HTTP_TIMEOUT = 5
 WORKFLOW_RUN_URL_TEMPLATE = "{base_url}/workflows/{workflow_id}/run"
 WORKFLOW_RUN_CHANNEL_NAME_TEMPLATE = "workflows:{workflow_id}:run:{channel_name_suffix}"
 SET_WORKFLOW_COMPLETE_MESSAGE = "SET_WORKFLOW_COMPLETE"
-WS_TOKEN_REQUEST_LIBRARY_URL = URL_API_PROMPTLAYER + "/ws-token-request-library"
+WS_TOKEN_REQUEST_LIBRARY_URL = (
+    f"{os.getenv('PROMPTLAYER_BASE_URL', 'https://api.promptlayer.com')}/ws-token-request-library"
+)
 
 
 logger = logging.getLogger(__name__)
@@ -48,6 +64,34 @@ class FinalOutputCode(Enum):
     EXCEEDS_SIZE_LIMIT = "EXCEEDS_SIZE_LIMIT"
 
 
+def should_retry_error(exception):
+    """Check if an exception should trigger a retry.
+
+    Only retries on server errors (5xx) and rate limits (429).
+    """
+    if hasattr(exception, "response"):
+        response = exception.response
+        if hasattr(response, "status_code"):
+            status_code = response.status_code
+            if status_code >= 500 or status_code == 429:
+                return True
+
+    if isinstance(exception, (_exceptions.PromptLayerInternalServerError, _exceptions.PromptLayerRateLimitError)):
+        return True
+
+    return False
+
+
+def retry_on_api_error(func):
+    return retry(
+        retry=retry_if_exception(should_retry_error),
+        stop=stop_after_attempt(4),  # 4 total attempts (1 initial + 3 retries)
+        wait=wait_exponential(multiplier=2, max=15),  # 2s, 4s, 8s
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        reraise=True,
+    )(func)
+
+
 def _get_http_timeout():
     try:
         return float(os.getenv("PROMPTLAYER_HTTP_TIMEOUT", DEFAULT_HTTP_TIMEOUT))
@@ -71,62 +115,60 @@ def _get_workflow_workflow_id_or_name(workflow_id_or_name, workflow_name):
     return workflow_id_or_name
 
 
-async def _get_final_output(execution_id: int, return_all_outputs: bool, *, headers: Dict[str, str]) -> Dict[str, Any]:
+async def _get_final_output(
+    base_url: str, execution_id: int, return_all_outputs: bool, *, headers: Dict[str, str]
+) -> Dict[str, Any]:
     async with httpx.AsyncClient() as client:
         response = await client.get(
-            f"{URL_API_PROMPTLAYER}/workflow-version-execution-results",
+            f"{base_url}/workflow-version-execution-results",
             headers=headers,
             params={"workflow_version_execution_id": execution_id, "return_all_outputs": return_all_outputs},
         )
-        response.raise_for_status()
+        if response.status_code != 200:
+            raise_on_bad_response(response, "PromptLayer had the following error while getting workflow results")
         return response.json()
 
 
 # TODO(dmu) MEDIUM: Consider putting all these functions into a class, so we do not have to pass
 #                   `authorization_headers` into each function
-async def _resolve_workflow_id(workflow_id_or_name: Union[int, str], headers):
+async def _resolve_workflow_id(base_url: str, workflow_id_or_name: Union[int, str], headers):
     if isinstance(workflow_id_or_name, int):
         return workflow_id_or_name
 
     # TODO(dmu) LOW: Should we warn user here to avoid using workflow names in favor of workflow id?
     async with _make_httpx_client() as client:
         # TODO(dmu) MEDIUM: Generalize the way we make async calls to PromptLayer API and reuse it everywhere
-        response = await client.get(f"{URL_API_PROMPTLAYER}/workflows/{workflow_id_or_name}", headers=headers)
-        if RAISE_FOR_STATUS:
-            response.raise_for_status()
-        elif response.status_code != 200:
-            raise_on_bad_response(response, "PromptLayer had the following error while running your workflow")
+        response = await client.get(f"{base_url}/workflows/{workflow_id_or_name}", headers=headers)
+        if response.status_code != 200:
+            raise_on_bad_response(response, "PromptLayer had the following error while resolving workflow")
 
         return response.json()["workflow"]["id"]
 
 
-async def _get_ably_token(channel_name, authentication_headers):
+async def _get_ably_token(base_url: str, channel_name, authentication_headers):
     try:
         async with _make_httpx_client() as client:
             response = await client.post(
-                f"{URL_API_PROMPTLAYER}/ws-token-request-library",
+                f"{base_url}/ws-token-request-library",
                 headers=authentication_headers,
                 params={"capability": channel_name},
             )
-            if RAISE_FOR_STATUS:
-                response.raise_for_status()
-            elif response.status_code != 201:
+            if response.status_code != 201:
                 raise_on_bad_response(
                     response,
                     "PromptLayer had the following error while getting WebSocket token",
                 )
-            return response.json()["token_details"]["token"]
+            return response.json()
     except Exception as ex:
         error_message = f"Failed to get WebSocket token: {ex}"
-        print(error_message)  # TODO(dmu) MEDIUM: Remove prints in favor of logging
         logger.exception(error_message)
         if RERAISE_ORIGINAL_EXCEPTION:
             raise
         else:
-            raise Exception(error_message)
+            raise _exceptions.PromptLayerAPIError(error_message, response=None, body=None) from ex
 
 
-def _make_message_listener(results_future, execution_id_future, return_all_outputs, headers):
+def _make_message_listener(base_url: str, results_future, execution_id_future, return_all_outputs, headers):
     # We need this function to be mocked by unittests
     async def message_listener(message: Message):
         if results_future.cancelled() or message.name != SET_WORKFLOW_COMPLETE_MESSAGE:
@@ -140,7 +182,7 @@ def _make_message_listener(results_future, execution_id_future, return_all_outpu
         if (result_code := message_data.get("result_code")) in (FinalOutputCode.OK.value, None):
             results = message_data["final_output"]
         elif result_code == FinalOutputCode.EXCEEDS_SIZE_LIMIT.value:
-            results = await _get_final_output(execution_id, return_all_outputs, headers=headers)
+            results = await _get_final_output(base_url, execution_id, return_all_outputs, headers=headers)
         else:
             raise NotImplementedError(f"Unsupported final output code: {result_code}")
 
@@ -149,15 +191,20 @@ def _make_message_listener(results_future, execution_id_future, return_all_outpu
     return message_listener
 
 
-async def _subscribe_to_workflow_completion_channel(channel, execution_id_future, return_all_outputs, headers):
+async def _subscribe_to_workflow_completion_channel(
+    base_url: str, channel, execution_id_future, return_all_outputs, headers
+):
     results_future = asyncio.Future()
-    message_listener = _make_message_listener(results_future, execution_id_future, return_all_outputs, headers)
+    message_listener = _make_message_listener(
+        base_url, results_future, execution_id_future, return_all_outputs, headers
+    )
     await channel.subscribe(SET_WORKFLOW_COMPLETE_MESSAGE, message_listener)
     return results_future, message_listener
 
 
 async def _post_workflow_id_run(
     *,
+    base_url: str,
     authentication_headers,
     workflow_id,
     input_variables: Dict[str, Any],
@@ -168,7 +215,7 @@ async def _post_workflow_id_run(
     channel_name_suffix: str,
     _url_template: str = WORKFLOW_RUN_URL_TEMPLATE,
 ):
-    url = _url_template.format(base_url=URL_API_PROMPTLAYER, workflow_id=workflow_id)
+    url = _url_template.format(base_url=base_url, workflow_id=workflow_id)
     payload = {
         "input_variables": input_variables,
         "metadata": metadata,
@@ -180,22 +227,19 @@ async def _post_workflow_id_run(
     try:
         async with _make_httpx_client() as client:
             response = await client.post(url, json=payload, headers=authentication_headers)
-            if RAISE_FOR_STATUS:
-                response.raise_for_status()
-            elif response.status_code != 201:
+            if response.status_code != 201:
                 raise_on_bad_response(response, "PromptLayer had the following error while running your workflow")
 
             result = response.json()
             if warning := result.get("warning"):
-                print(f"WARNING: {warning}")
+                logger.warning(f"{warning}")
     except Exception as ex:
         error_message = f"Failed to run workflow: {str(ex)}"
-        print(error_message)  # TODO(dmu) MEDIUM: Remove prints in favor of logging
         logger.exception(error_message)
         if RERAISE_ORIGINAL_EXCEPTION:
             raise
         else:
-            raise Exception(error_message)
+            raise _exceptions.PromptLayerAPIError(error_message, response=None, body=None) from ex
 
     return result.get("workflow_version_execution_id")
 
@@ -205,7 +249,9 @@ async def _wait_for_workflow_completion(channel, results_future, message_listene
     try:
         return await asyncio.wait_for(results_future, timeout)
     except asyncio.TimeoutError:
-        raise Exception("Workflow execution did not complete properly")
+        raise _exceptions.PromptLayerAPITimeoutError(
+            "Workflow execution did not complete properly", response=None, body=None
+        )
     finally:
         channel.unsubscribe(SET_WORKFLOW_COMPLETE_MESSAGE, message_listener)
 
@@ -215,14 +261,55 @@ def _make_channel_name_suffix():
     return uuid4().hex
 
 
+MessageCallback = Callable[[Message], Coroutine[None, None, None]]
+
+
+class SubscriptionEventLoggerHandler(SubscriptionEventHandler):
+    def __init__(self, callback: MessageCallback):
+        self.callback = callback
+
+    async def on_publication(self, ctx: PublicationContext):
+        message_name = ctx.pub.data.get("message_name", "unknown")
+        data = ctx.pub.data.get("data", "")
+        message = Message(name=message_name, data=data)
+        await self.callback(message)
+
+
+@asynccontextmanager
+async def centrifugo_client(address: str, token: str):
+    client = Client(address, token=token)
+    try:
+        await client.connect()
+        yield client
+    finally:
+        await client.disconnect()
+
+
+@asynccontextmanager
+async def centrifugo_subscription(client: Client, topic: str, message_listener: MessageCallback):
+    subscription = client.new_subscription(
+        topic,
+        events=SubscriptionEventLoggerHandler(message_listener),
+    )
+    try:
+        await subscription.subscribe()
+        yield
+    finally:
+        if subscription.state == SubscriptionState.SUBSCRIBED:
+            await subscription.unsubscribe()
+
+
+@retry_on_api_error
 async def arun_workflow_request(
     *,
+    api_key: str,
+    base_url: str,
+    throw_on_error: bool,
     workflow_id_or_name: Optional[Union[int, str]] = None,
     input_variables: Dict[str, Any],
     metadata: Optional[Dict[str, Any]] = None,
     workflow_label_name: Optional[str] = None,
     workflow_version_number: Optional[int] = None,
-    api_key: str,
     return_all_outputs: Optional[bool] = False,
     timeout: Optional[int] = 3600,
     # `workflow_name` deprecated, kept for backward compatibility only.
@@ -230,22 +317,50 @@ async def arun_workflow_request(
 ):
     headers = {"X-API-KEY": api_key}
     workflow_id = await _resolve_workflow_id(
-        _get_workflow_workflow_id_or_name(workflow_id_or_name, workflow_name), headers
+        base_url, _get_workflow_workflow_id_or_name(workflow_id_or_name, workflow_name), headers
     )
     channel_name_suffix = _make_channel_name_suffix()
     channel_name = WORKFLOW_RUN_CHANNEL_NAME_TEMPLATE.format(
         workflow_id=workflow_id, channel_name_suffix=channel_name_suffix
     )
-    ably_token = await _get_ably_token(channel_name, headers)
-    async with AblyRealtime(token=ably_token) as ably_client:
+    ably_token = await _get_ably_token(base_url, channel_name, headers)
+    token = ably_token["token_details"]["token"]
+
+    execution_id_future = asyncio.Future[int]()
+
+    if ably_token.get("messaging_backend") == "centrifugo":
+        address = urllib3.util.parse_url(base_url)._replace(scheme="wss", path="/connection/websocket").url
+        async with centrifugo_client(address, token) as client:
+            results_future = asyncio.Future[dict[str, Any]]()
+            async with centrifugo_subscription(
+                client,
+                channel_name,
+                _make_message_listener(base_url, results_future, execution_id_future, return_all_outputs, headers),
+            ):
+                execution_id = await _post_workflow_id_run(
+                    base_url=base_url,
+                    authentication_headers=headers,
+                    workflow_id=workflow_id,
+                    input_variables=input_variables,
+                    metadata=metadata,
+                    workflow_label_name=workflow_label_name,
+                    workflow_version_number=workflow_version_number,
+                    return_all_outputs=return_all_outputs,
+                    channel_name_suffix=channel_name_suffix,
+                )
+                execution_id_future.set_result(execution_id)
+                await asyncio.wait_for(results_future, timeout)
+                return results_future.result()
+
+    async with AblyRealtime(token=token) as ably_client:
         # It is crucial to subscribe before running a workflow, otherwise we may miss a completion message
         channel = ably_client.channels.get(channel_name)
-        execution_id_future = asyncio.Future()
         results_future, message_listener = await _subscribe_to_workflow_completion_channel(
-            channel, execution_id_future, return_all_outputs, headers
+            base_url, channel, execution_id_future, return_all_outputs, headers
         )
 
         execution_id = await _post_workflow_id_run(
+            base_url=base_url,
             authentication_headers=headers,
             workflow_id=workflow_id,
             input_variables=input_variables,
@@ -261,6 +376,8 @@ async def arun_workflow_request(
 
 
 def promptlayer_api_handler(
+    api_key: str,
+    base_url: str,
     function_name,
     provider_type,
     args,
@@ -269,7 +386,6 @@ def promptlayer_api_handler(
     response,
     request_start_time,
     request_end_time,
-    api_key,
     return_pl_id=False,
     llm_request_span_id=None,
 ):
@@ -292,9 +408,11 @@ def promptlayer_api_handler(
                 "llm_request_span_id": llm_request_span_id,
             },
             api_key=api_key,
+            base_url=base_url,
         )
     else:
         request_id = promptlayer_api_request(
+            base_url=base_url,
             function_name=function_name,
             provider_type=provider_type,
             args=args,
@@ -313,6 +431,8 @@ def promptlayer_api_handler(
 
 
 async def promptlayer_api_handler_async(
+    api_key: str,
+    base_url: str,
     function_name,
     provider_type,
     args,
@@ -321,13 +441,14 @@ async def promptlayer_api_handler_async(
     response,
     request_start_time,
     request_end_time,
-    api_key,
     return_pl_id=False,
     llm_request_span_id=None,
 ):
     return await run_in_thread_async(
         None,
         promptlayer_api_handler,
+        api_key,
+        base_url,
         function_name,
         provider_type,
         args,
@@ -336,7 +457,6 @@ async def promptlayer_api_handler_async(
         response,
         request_start_time,
         request_end_time,
-        api_key,
         return_pl_id=return_pl_id,
         llm_request_span_id=llm_request_span_id,
     )
@@ -356,6 +476,7 @@ def convert_native_object_to_dict(native_object):
 
 def promptlayer_api_request(
     *,
+    base_url: str,
     function_name,
     provider_type,
     args,
@@ -376,7 +497,7 @@ def promptlayer_api_request(
         response = response.dict()
     try:
         request_response = requests.post(
-            f"{URL_API_PROMPTLAYER}/track-request",
+            f"{base_url}/track-request",
             json={
                 "function_name": function_name,
                 "provider_type": provider_type,
@@ -400,43 +521,57 @@ def promptlayer_api_request(
                 request_response, "WARNING: While logging your request PromptLayer had the following error"
             )
     except Exception as e:
-        print(f"WARNING: While logging your request PromptLayer had the following error: {e}", file=sys.stderr)
+        logger.warning(f"While logging your request PromptLayer had the following error: {e}")
     if request_response is not None and return_pl_id:
         return request_response.json().get("request_id")
 
 
-def track_request(**body):
+@retry_on_api_error
+def track_request(base_url: str, throw_on_error: bool, **body):
     try:
         response = requests.post(
-            f"{URL_API_PROMPTLAYER}/track-request",
+            f"{base_url}/track-request",
             json=body,
         )
         if response.status_code != 200:
-            warn_on_bad_response(
-                response, f"PromptLayer had the following error while tracking your request: {response.text}"
-            )
+            if throw_on_error:
+                raise_on_bad_response(response, "PromptLayer had the following error while tracking your request")
+            else:
+                warn_on_bad_response(
+                    response, f"PromptLayer had the following error while tracking your request: {response.text}"
+                )
         return response.json()
     except requests.exceptions.RequestException as e:
-        print(f"WARNING: While logging your request PromptLayer had the following error: {e}", file=sys.stderr)
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while tracking your request: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While logging your request PromptLayer had the following error: {e}")
         return {}
 
 
-async def atrack_request(**body: Any) -> Dict[str, Any]:
+@retry_on_api_error
+async def atrack_request(base_url: str, throw_on_error: bool, **body: Any) -> Dict[str, Any]:
     try:
         async with _make_httpx_client() as client:
             response = await client.post(
-                f"{URL_API_PROMPTLAYER}/track-request",
+                f"{base_url}/track-request",
                 json=body,
             )
-            if RAISE_FOR_STATUS:
-                response.raise_for_status()
-            elif response.status_code != 200:
-                warn_on_bad_response(
-                    response, f"PromptLayer had the following error while tracking your request: {response.text}"
-                )
+            if response.status_code != 200:
+                if throw_on_error:
+                    raise_on_bad_response(response, "PromptLayer had the following error while tracking your request")
+                else:
+                    warn_on_bad_response(
+                        response, f"PromptLayer had the following error while tracking your request: {response.text}"
+                    )
         return response.json()
     except httpx.RequestError as e:
-        print(f"WARNING: While logging your request PromptLayer had the following error: {e}", file=sys.stderr)
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while tracking your request: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While logging your request PromptLayer had the following error: {e}")
         return {}
 
 
@@ -468,7 +603,10 @@ def promptlayer_api_request_async(
     )
 
 
-def promptlayer_get_prompt(prompt_name, api_key, version: int = None, label: str = None):
+@retry_on_api_error
+def promptlayer_get_prompt(
+    api_key: str, base_url: str, throw_on_error: bool, prompt_name, version: int = None, label: str = None
+):
     """
     Get a prompt from the PromptLayer library
     version: version of the prompt to get, None for latest
@@ -476,25 +614,40 @@ def promptlayer_get_prompt(prompt_name, api_key, version: int = None, label: str
     """
     try:
         request_response = requests.get(
-            f"{URL_API_PROMPTLAYER}/library-get-prompt-template",
+            f"{base_url}/library-get-prompt-template",
             headers={"X-API-KEY": api_key},
             params={"prompt_name": prompt_name, "version": version, "label": label},
         )
     except Exception as e:
-        raise Exception(f"PromptLayer had the following error while getting your prompt: {e}")
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIError(
+                f"PromptLayer had the following error while getting your prompt: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"PromptLayer had the following error while getting your prompt: {e}")
+        return None
     if request_response.status_code != 200:
-        raise_on_bad_response(
-            request_response,
-            "PromptLayer had the following error while getting your prompt",
-        )
+        if throw_on_error:
+            raise_on_bad_response(
+                request_response,
+                "PromptLayer had the following error while getting your prompt",
+            )
+        else:
+            warn_on_bad_response(
+                request_response,
+                "WARNING: PromptLayer had the following error while getting your prompt",
+            )
+            return None
 
     return request_response.json()
 
 
-def promptlayer_publish_prompt(prompt_name, prompt_template, commit_message, tags, api_key, metadata=None):
+@retry_on_api_error
+def promptlayer_publish_prompt(
+    api_key: str, base_url: str, throw_on_error: bool, prompt_name, prompt_template, commit_message, tags, metadata=None
+):
     try:
         request_response = requests.post(
-            f"{URL_API_PROMPTLAYER}/library-publish-prompt-template",
+            f"{base_url}/library-publish-prompt-template",
             json={
                 "prompt_name": prompt_name,
                 "prompt_template": prompt_template,
@@ -505,19 +658,34 @@ def promptlayer_publish_prompt(prompt_name, prompt_template, commit_message, tag
             },
         )
     except Exception as e:
-        raise Exception(f"PromptLayer had the following error while publishing your prompt: {e}")
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIError(
+                f"PromptLayer had the following error while publishing your prompt: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"PromptLayer had the following error while publishing your prompt: {e}")
+        return False
     if request_response.status_code != 200:
-        raise_on_bad_response(
-            request_response,
-            "PromptLayer had the following error while publishing your prompt",
-        )
+        if throw_on_error:
+            raise_on_bad_response(
+                request_response,
+                "PromptLayer had the following error while publishing your prompt",
+            )
+        else:
+            warn_on_bad_response(
+                request_response,
+                "WARNING: PromptLayer had the following error while publishing your prompt",
+            )
+            return False
     return True
 
 
-def promptlayer_track_prompt(request_id, prompt_name, input_variables, api_key, version, label):
+@retry_on_api_error
+def promptlayer_track_prompt(
+    api_key: str, base_url: str, throw_on_error: bool, request_id, prompt_name, input_variables, version, label
+):
     try:
         request_response = requests.post(
-            f"{URL_API_PROMPTLAYER}/library-track-prompt",
+            f"{base_url}/library-track-prompt",
             json={
                 "request_id": request_id,
                 "prompt_name": prompt_name,
@@ -528,29 +696,39 @@ def promptlayer_track_prompt(request_id, prompt_name, input_variables, api_key,
             },
         )
         if request_response.status_code != 200:
-            warn_on_bad_response(
-                request_response,
-                "WARNING: While tracking your prompt PromptLayer had the following error",
-            )
-            return False
+            if throw_on_error:
+                raise_on_bad_response(
+                    request_response,
+                    "While tracking your prompt PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    request_response,
+                    "WARNING: While tracking your prompt PromptLayer had the following error",
+                )
+                return False
     except Exception as e:
-        print(
-            f"WARNING: While tracking your prompt PromptLayer had the following error: {e}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIError(
+                f"While tracking your prompt PromptLayer had the following error: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your prompt PromptLayer had the following error: {e}")
         return False
     return True
 
 
+@retry_on_api_error
 async def apromptlayer_track_prompt(
+    api_key: str,
+    base_url: str,
     request_id: str,
     prompt_name: str,
     input_variables: Dict[str, Any],
-    api_key: Optional[str] = None,
     version: Optional[int] = None,
     label: Optional[str] = None,
+    throw_on_error: bool = True,
 ) -> bool:
-    url = f"{URL_API_PROMPTLAYER}/library-track-prompt"
+    url = f"{base_url}/library-track-prompt"
     payload = {
         "request_id": request_id,
         "prompt_name": prompt_name,
@@ -563,28 +741,31 @@ async def apromptlayer_track_prompt(
         async with _make_httpx_client() as client:
             response = await client.post(url, json=payload)
 
-        if RAISE_FOR_STATUS:
-            response.raise_for_status()
-        elif response.status_code != 200:
-            warn_on_bad_response(
-                response,
-                "WARNING: While tracking your prompt, PromptLayer had the following error",
-            )
-            return False
+        if response.status_code != 200:
+            if throw_on_error:
+                raise_on_bad_response(response, "While tracking your prompt, PromptLayer had the following error")
+            else:
+                warn_on_bad_response(
+                    response,
+                    "WARNING: While tracking your prompt, PromptLayer had the following error",
+                )
+                return False
     except httpx.RequestError as e:
-        print(
-            f"WARNING: While tracking your prompt PromptLayer had the following error: {e}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"While tracking your prompt PromptLayer had the following error: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your prompt PromptLayer had the following error: {e}")
         return False
 
     return True
 
 
-def promptlayer_track_metadata(request_id, metadata, api_key):
+@retry_on_api_error
+def promptlayer_track_metadata(api_key: str, base_url: str, throw_on_error: bool, request_id, metadata):
     try:
         request_response = requests.post(
-            f"{URL_API_PROMPTLAYER}/library-track-metadata",
+            f"{base_url}/library-track-metadata",
             json={
                 "request_id": request_id,
                 "metadata": metadata,
@@ -592,22 +773,32 @@ def promptlayer_track_metadata(request_id, metadata, api_key):
             },
         )
         if request_response.status_code != 200:
-            warn_on_bad_response(
-                request_response,
-                "WARNING: While tracking your metadata PromptLayer had the following error",
-            )
-            return False
+            if throw_on_error:
+                raise_on_bad_response(
+                    request_response,
+                    "While tracking your metadata PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    request_response,
+                    "WARNING: While tracking your metadata PromptLayer had the following error",
+                )
+                return False
     except Exception as e:
-        print(
-            f"WARNING: While tracking your metadata PromptLayer had the following error: {e}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIError(
+                f"While tracking your metadata PromptLayer had the following error: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your metadata PromptLayer had the following error: {e}")
         return False
     return True
 
 
-async def apromptlayer_track_metadata(request_id: str, metadata: Dict[str, Any], api_key: Optional[str] = None) -> bool:
-    url = f"{URL_API_PROMPTLAYER}/library-track-metadata"
+@retry_on_api_error
+async def apromptlayer_track_metadata(
+    api_key: str, base_url: str, throw_on_error: bool, request_id: str, metadata: Dict[str, Any]
+) -> bool:
+    url = f"{base_url}/library-track-metadata"
     payload = {
         "request_id": request_id,
         "metadata": metadata,
@@ -617,55 +808,71 @@ async def apromptlayer_track_metadata(request_id: str, metadata: Dict[str, Any],
         async with _make_httpx_client() as client:
             response = await client.post(url, json=payload)
 
-        if RAISE_FOR_STATUS:
-            response.raise_for_status()
-        elif response.status_code != 200:
-            warn_on_bad_response(
-                response,
-                "WARNING: While tracking your metadata, PromptLayer had the following error",
-            )
-            return False
+        if response.status_code != 200:
+            if throw_on_error:
+                raise_on_bad_response(
+                    response,
+                    "While tracking your metadata, PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    response,
+                    "WARNING: While tracking your metadata, PromptLayer had the following error",
+                )
+                return False
     except httpx.RequestError as e:
-        print(
-            f"WARNING: While tracking your metadata PromptLayer had the following error: {e}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"While tracking your metadata PromptLayer had the following error: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your metadata PromptLayer had the following error: {e}")
         return False
 
     return True
 
 
-def promptlayer_track_score(request_id, score, score_name, api_key):
+@retry_on_api_error
+def promptlayer_track_score(api_key: str, base_url: str, throw_on_error: bool, request_id, score, score_name):
     try:
         data = {"request_id": request_id, "score": score, "api_key": api_key}
         if score_name is not None:
             data["name"] = score_name
         request_response = requests.post(
-            f"{URL_API_PROMPTLAYER}/library-track-score",
+            f"{base_url}/library-track-score",
             json=data,
         )
         if request_response.status_code != 200:
-            warn_on_bad_response(
-                request_response,
-                "WARNING: While tracking your score PromptLayer had the following error",
-            )
-            return False
+            if throw_on_error:
+                raise_on_bad_response(
+                    request_response,
+                    "While tracking your score PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    request_response,
+                    "WARNING: While tracking your score PromptLayer had the following error",
+                )
+                return False
     except Exception as e:
-        print(
-            f"WARNING: While tracking your score PromptLayer had the following error: {e}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIError(
+                f"While tracking your score PromptLayer had the following error: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your score PromptLayer had the following error: {e}")
         return False
     return True
 
 
+@retry_on_api_error
 async def apromptlayer_track_score(
+    api_key: str,
+    base_url: str,
+    throw_on_error: bool,
     request_id: str,
     score: float,
     score_name: Optional[str],
-    api_key: Optional[str] = None,
 ) -> bool:
-    url = f"{URL_API_PROMPTLAYER}/library-track-score"
+    url = f"{base_url}/library-track-score"
     data = {
         "request_id": request_id,
         "score": score,
@@ -677,19 +884,24 @@ async def apromptlayer_track_score(
         async with _make_httpx_client() as client:
             response = await client.post(url, json=data)
 
-        if RAISE_FOR_STATUS:
-            response.raise_for_status()
-        elif response.status_code != 200:
-            warn_on_bad_response(
-                response,
-                "WARNING: While tracking your score, PromptLayer had the following error",
-            )
-            return False
+        if response.status_code != 200:
+            if throw_on_error:
+                raise_on_bad_response(
+                    response,
+                    "While tracking your score, PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    response,
+                    "WARNING: While tracking your score, PromptLayer had the following error",
+                )
+                return False
     except httpx.RequestError as e:
-        print(
-            f"WARNING: While tracking your score PromptLayer had the following error: {str(e)}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while tracking your score: {str(e)}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your score PromptLayer had the following error: {str(e)}")
         return False
 
     return True
@@ -753,11 +965,12 @@ def build_anthropic_content_blocks(events):
 
 
 class GeneratorProxy:
-    def __init__(self, generator, api_request_arguments, api_key):
+    def __init__(self, generator, api_request_arguments, api_key, base_url):
         self.generator = generator
         self.results = []
         self.api_request_arugments = api_request_arguments
         self.api_key = api_key
+        self.base_url = base_url
 
     def __iter__(self):
         return self
@@ -772,6 +985,7 @@ class GeneratorProxy:
                 await self.generator._AsyncMessageStreamManager__api_request,
                 api_request_arguments,
                 self.api_key,
+                self.base_url,
             )
 
     def __enter__(self):
@@ -782,6 +996,7 @@ class GeneratorProxy:
                 stream,
                 api_request_arguments,
                 self.api_key,
+                self.base_url,
             )
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -800,7 +1015,7 @@ class GeneratorProxy:
 
     def __getattr__(self, name):
         if name == "text_stream":  # anthropic async stream
-            return GeneratorProxy(self.generator.text_stream, self.api_request_arugments, self.api_key)
+            return GeneratorProxy(self.generator.text_stream, self.api_request_arugments, self.api_key, self.base_url)
         return getattr(self.generator, name)
 
     def _abstracted_next(self, result):
@@ -822,6 +1037,7 @@ class GeneratorProxy:
 
         if end_anthropic or end_openai:
             request_id = promptlayer_api_request(
+                base_url=self.base_url,
                 function_name=self.api_request_arugments["function_name"],
                 provider_type=self.api_request_arugments["provider_type"],
                 args=self.api_request_arugments["args"],
@@ -912,39 +1128,71 @@ async def run_in_thread_async(executor, func, *args, **kwargs):
 def warn_on_bad_response(request_response, main_message):
     if hasattr(request_response, "json"):
         try:
-            print(
-                f"{main_message}: {request_response.json().get('message')}",
-                file=sys.stderr,
-            )
+            logger.warning(f"{main_message}: {request_response.json().get('message')}")
         except json.JSONDecodeError:
-            print(
-                f"{main_message}: {request_response}",
-                file=sys.stderr,
-            )
+            logger.warning(f"{main_message}: {request_response}")
     else:
-        print(f"{main_message}: {request_response}", file=sys.stderr)
+        logger.warning(f"{main_message}: {request_response}")
 
 
 def raise_on_bad_response(request_response, main_message):
+    """Raise an appropriate exception based on the HTTP status code."""
+    status_code = getattr(request_response, "status_code", None)
+
+    body = None
+    error_detail = None
     if hasattr(request_response, "json"):
         try:
-            raise Exception(
-                f"{main_message}: {request_response.json().get('message') or request_response.json().get('error')}"
-            )
-        except json.JSONDecodeError:
-            raise Exception(f"{main_message}: {request_response}")
+            body = request_response.json()
+            error_detail = body.get("message") or body.get("error") or body.get("detail")
+        except (json.JSONDecodeError, AttributeError):
+            body = getattr(request_response, "text", str(request_response))
+            error_detail = body
+    else:
+        body = str(request_response)
+        error_detail = body
+
+    if error_detail:
+        err_msg = f"{main_message}: {error_detail}"
     else:
-        raise Exception(f"{main_message}: {request_response}")
+        err_msg = main_message
+
+    if status_code == 400:
+        raise _exceptions.PromptLayerBadRequestError(err_msg, response=request_response, body=body)
+
+    if status_code == 401:
+        raise _exceptions.PromptLayerAuthenticationError(err_msg, response=request_response, body=body)
+
+    if status_code == 403:
+        raise _exceptions.PromptLayerPermissionDeniedError(err_msg, response=request_response, body=body)
+
+    if status_code == 404:
+        raise _exceptions.PromptLayerNotFoundError(err_msg, response=request_response, body=body)
+
+    if status_code == 409:
+        raise _exceptions.PromptLayerConflictError(err_msg, response=request_response, body=body)
+
+    if status_code == 422:
+        raise _exceptions.PromptLayerUnprocessableEntityError(err_msg, response=request_response, body=body)
+
+    if status_code == 429:
+        raise _exceptions.PromptLayerRateLimitError(err_msg, response=request_response, body=body)
+
+    if status_code and status_code >= 500:
+        raise _exceptions.PromptLayerInternalServerError(err_msg, response=request_response, body=body)
+
+    raise _exceptions.PromptLayerAPIStatusError(err_msg, response=request_response, body=body)
 
 
 async def async_wrapper(
+    api_key: str,
+    base_url: str,
     coroutine_obj,
     return_pl_id,
     request_start_time,
     function_name,
     provider_type,
     tags,
-    api_key: str = None,
     llm_request_span_id: str = None,
     tracer=None,
     *args,
@@ -957,6 +1205,8 @@ async def async_wrapper(
         response = await coroutine_obj
         request_end_time = datetime.datetime.now().timestamp()
         result = await promptlayer_api_handler_async(
+            api_key,
+            base_url,
             function_name,
             provider_type,
             args,
@@ -965,7 +1215,6 @@ async def async_wrapper(
             response,
             request_start_time,
             request_end_time,
-            api_key,
             return_pl_id=return_pl_id,
             llm_request_span_id=llm_request_span_id,
         )
@@ -980,53 +1229,75 @@ async def async_wrapper(
         context.detach(token)
 
 
-def promptlayer_create_group(api_key: str = None):
+@retry_on_api_error
+def promptlayer_create_group(api_key: str, base_url: str, throw_on_error: bool):
     try:
         request_response = requests.post(
-            f"{URL_API_PROMPTLAYER}/create-group",
+            f"{base_url}/create-group",
             json={
                 "api_key": api_key,
             },
         )
         if request_response.status_code != 200:
-            warn_on_bad_response(
-                request_response,
-                "WARNING: While creating your group PromptLayer had the following error",
-            )
-            return False
+            if throw_on_error:
+                raise_on_bad_response(
+                    request_response,
+                    "While creating your group PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    request_response,
+                    "WARNING: While creating your group PromptLayer had the following error",
+                )
+                return False
     except requests.exceptions.RequestException as e:
-        # I'm aiming for a more specific exception catch here
-        raise Exception(f"PromptLayer had the following error while creating your group: {e}")
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while creating your group: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While creating your group PromptLayer had the following error: {e}")
+        return False
     return request_response.json()["id"]
 
 
-async def apromptlayer_create_group(api_key: Optional[str] = None) -> str:
+@retry_on_api_error
+async def apromptlayer_create_group(api_key: str, base_url: str, throw_on_error: bool):
     try:
         async with _make_httpx_client() as client:
             response = await client.post(
-                f"{URL_API_PROMPTLAYER}/create-group",
+                f"{base_url}/create-group",
                 json={
                     "api_key": api_key,
                 },
             )
 
-        if RAISE_FOR_STATUS:
-            response.raise_for_status()
-        elif response.status_code != 200:
-            warn_on_bad_response(
-                response,
-                "WARNING: While creating your group, PromptLayer had the following error",
-            )
-            return False
+        if response.status_code != 200:
+            if throw_on_error:
+                raise_on_bad_response(
+                    response,
+                    "While creating your group, PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    response,
+                    "WARNING: While creating your group, PromptLayer had the following error",
+                )
+                return False
         return response.json()["id"]
     except httpx.RequestError as e:
-        raise Exception(f"PromptLayer had the following error while creating your group: {str(e)}") from e
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while creating your group: {str(e)}", response=None, body=None
+            ) from e
+        logger.warning(f"While creating your group PromptLayer had the following error: {e}")
+        return False
 
 
-def promptlayer_track_group(request_id, group_id, api_key: str = None):
+@retry_on_api_error
+def promptlayer_track_group(api_key: str, base_url: str, throw_on_error: bool, request_id, group_id):
     try:
         request_response = requests.post(
-            f"{URL_API_PROMPTLAYER}/track-group",
+            f"{base_url}/track-group",
             json={
                 "api_key": api_key,
                 "request_id": request_id,
@@ -1034,18 +1305,29 @@ def promptlayer_track_group(request_id, group_id, api_key: str = None):
             },
         )
         if request_response.status_code != 200:
-            warn_on_bad_response(
-                request_response,
-                "WARNING: While tracking your group PromptLayer had the following error",
-            )
-            return False
+            if throw_on_error:
+                raise_on_bad_response(
+                    request_response,
+                    "While tracking your group PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    request_response,
+                    "WARNING: While tracking your group PromptLayer had the following error",
+                )
+                return False
     except requests.exceptions.RequestException as e:
-        # I'm aiming for a more specific exception catch here
-        raise Exception(f"PromptLayer had the following error while tracking your group: {e}")
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while tracking your group: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your group PromptLayer had the following error: {e}")
+        return False
     return True
 
 
-async def apromptlayer_track_group(request_id, group_id, api_key: str = None):
+@retry_on_api_error
+async def apromptlayer_track_group(api_key: str, base_url: str, throw_on_error: bool, request_id, group_id):
     try:
         payload = {
             "api_key": api_key,
@@ -1054,59 +1336,86 @@ async def apromptlayer_track_group(request_id, group_id, api_key: str = None):
         }
         async with _make_httpx_client() as client:
             response = await client.post(
-                f"{URL_API_PROMPTLAYER}/track-group",
+                f"{base_url}/track-group",
                 headers={"X-API-KEY": api_key},
                 json=payload,
             )
 
-        if RAISE_FOR_STATUS:
-            response.raise_for_status()
-        elif response.status_code != 200:
-            warn_on_bad_response(
-                response,
-                "WARNING: While tracking your group, PromptLayer had the following error",
-            )
-            return False
+        if response.status_code != 200:
+            if throw_on_error:
+                raise_on_bad_response(
+                    response,
+                    "While tracking your group, PromptLayer had the following error",
+                )
+            else:
+                warn_on_bad_response(
+                    response,
+                    "WARNING: While tracking your group, PromptLayer had the following error",
+                )
+                return False
     except httpx.RequestError as e:
-        print(
-            f"WARNING: While tracking your group PromptLayer had the following error: {e}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while tracking your group: {str(e)}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your group PromptLayer had the following error: {e}")
         return False
 
     return True
 
 
+@retry_on_api_error
 def get_prompt_template(
-    prompt_name: str, params: Union[GetPromptTemplate, None] = None, api_key: str = None
+    api_key: str, base_url: str, throw_on_error: bool, prompt_name: str, params: Union[GetPromptTemplate, None] = None
 ) -> GetPromptTemplateResponse:
     try:
         json_body = {"api_key": api_key}
         if params:
             json_body = {**json_body, **params}
         response = requests.post(
-            f"{URL_API_PROMPTLAYER}/prompt-templates/{prompt_name}",
+            f"{base_url}/prompt-templates/{prompt_name}",
             headers={"X-API-KEY": api_key},
             json=json_body,
         )
         if response.status_code != 200:
-            raise Exception(f"PromptLayer had the following error while getting your prompt template: {response.text}")
+            if throw_on_error:
+                raise_on_bad_response(
+                    response, "PromptLayer had the following error while getting your prompt template"
+                )
+            else:
+                warn_on_bad_response(
+                    response, "WARNING: PromptLayer had the following error while getting your prompt template"
+                )
+                return None
 
-        warning = response.json().get("warning", None)
-        if warning is not None:
-            warn_on_bad_response(
-                warning,
-                "WARNING: While getting your prompt template",
-            )
         return response.json()
+    except requests.exceptions.ConnectionError as e:
+        err_msg = f"PromptLayer had the following error while getting your prompt template: {e}"
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(err_msg, response=None, body=None) from e
+        logger.warning(err_msg)
+        return None
+    except requests.exceptions.Timeout as e:
+        err_msg = f"PromptLayer had the following error while getting your prompt template: {e}"
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPITimeoutError(err_msg, response=None, body=None) from e
+        logger.warning(err_msg)
+        return None
     except requests.exceptions.RequestException as e:
-        raise Exception(f"PromptLayer had the following error while getting your prompt template: {e}")
+        err_msg = f"PromptLayer had the following error while getting your prompt template: {e}"
+        if throw_on_error:
+            raise _exceptions.PromptLayerError(err_msg, response=None, body=None) from e
+        logger.warning(err_msg)
+        return None
 
 
+@retry_on_api_error
 async def aget_prompt_template(
+    api_key: str,
+    base_url: str,
+    throw_on_error: bool,
     prompt_name: str,
     params: Union[GetPromptTemplate, None] = None,
-    api_key: str = None,
 ) -> GetPromptTemplateResponse:
     try:
         json_body = {"api_key": api_key}
@@ -1114,36 +1423,53 @@ async def aget_prompt_template(
             json_body.update(params)
         async with _make_httpx_client() as client:
             response = await client.post(
-                f"{URL_API_PROMPTLAYER}/prompt-templates/{prompt_name}",
+                f"{base_url}/prompt-templates/{prompt_name}",
                 headers={"X-API-KEY": api_key},
                 json=json_body,
             )
 
-            if RAISE_FOR_STATUS:
-                response.raise_for_status()
-            elif response.status_code != 200:
-                raise_on_bad_response(
-                    response,
-                    "PromptLayer had the following error while getting your prompt template",
-                )
-        warning = response.json().get("warning", None)
-        if warning:
-            warn_on_bad_response(
-                warning,
-                "WARNING: While getting your prompt template",
-            )
+            if response.status_code != 200:
+                if throw_on_error:
+                    raise_on_bad_response(
+                        response,
+                        "PromptLayer had the following error while getting your prompt template",
+                    )
+                else:
+                    warn_on_bad_response(
+                        response, "WARNING: While getting your prompt template PromptLayer had the following error"
+                    )
+                    return None
         return response.json()
+    except (httpx.ConnectError, httpx.NetworkError) as e:
+        err_msg = f"PromptLayer had the following error while getting your prompt template: {str(e)}"
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(err_msg, response=None, body=None) from e
+        logger.warning(err_msg)
+        return None
+    except httpx.TimeoutException as e:
+        err_msg = f"PromptLayer had the following error while getting your prompt template: {str(e)}"
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPITimeoutError(err_msg, response=None, body=None) from e
+        logger.warning(err_msg)
+        return None
     except httpx.RequestError as e:
-        raise Exception(f"PromptLayer had the following error while getting your prompt template: {str(e)}") from e
+        err_msg = f"PromptLayer had the following error while getting your prompt template: {str(e)}"
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(err_msg, response=None, body=None) from e
+        logger.warning(err_msg)
+        return None
 
 
+@retry_on_api_error
 def publish_prompt_template(
+    api_key: str,
+    base_url: str,
+    throw_on_error: bool,
     body: PublishPromptTemplate,
-    api_key: str = None,
 ) -> PublishPromptTemplateResponse:
     try:
         response = requests.post(
-            f"{URL_API_PROMPTLAYER}/rest/prompt-templates",
+            f"{base_url}/rest/prompt-templates",
             headers={"X-API-KEY": api_key},
             json={
                 "prompt_template": {**body},
@@ -1152,22 +1478,38 @@ def publish_prompt_template(
             },
         )
         if response.status_code == 400:
-            raise Exception(
-                f"PromptLayer had the following error while publishing your prompt template: {response.text}"
-            )
+            if throw_on_error:
+                raise_on_bad_response(
+                    response, "PromptLayer had the following error while publishing your prompt template"
+                )
+            else:
+                warn_on_bad_response(
+                    response, "WARNING: PromptLayer had the following error while publishing your prompt template"
+                )
+                return None
         return response.json()
     except requests.exceptions.RequestException as e:
-        raise Exception(f"PromptLayer had the following error while publishing your prompt template: {e}")
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while publishing your prompt template: {e}",
+                response=None,
+                body=None,
+            ) from e
+        logger.warning(f"PromptLayer had the following error while publishing your prompt template: {e}")
+        return None
 
 
+@retry_on_api_error
 async def apublish_prompt_template(
+    api_key: str,
+    base_url: str,
+    throw_on_error: bool,
     body: PublishPromptTemplate,
-    api_key: str = None,
 ) -> PublishPromptTemplateResponse:
     try:
         async with _make_httpx_client() as client:
             response = await client.post(
-                f"{URL_API_PROMPTLAYER}/rest/prompt-templates",
+                f"{base_url}/rest/prompt-templates",
                 headers={"X-API-KEY": api_key},
                 json={
                     "prompt_template": {**body},
@@ -1176,46 +1518,68 @@ async def apublish_prompt_template(
                 },
             )
 
-        if RAISE_FOR_STATUS:
-            response.raise_for_status()
-        elif response.status_code == 400:
-            raise Exception(
-                f"PromptLayer had the following error while publishing your prompt template: {response.text}"
-            )
-        if response.status_code != 201:
-            raise_on_bad_response(
-                response,
-                "PromptLayer had the following error while publishing your prompt template",
-            )
+        if response.status_code == 400 or response.status_code != 201:
+            if throw_on_error:
+                raise_on_bad_response(
+                    response,
+                    "PromptLayer had the following error while publishing your prompt template",
+                )
+            else:
+                warn_on_bad_response(
+                    response, "WARNING: PromptLayer had the following error while publishing your prompt template"
+                )
+                return None
         return response.json()
     except httpx.RequestError as e:
-        raise Exception(f"PromptLayer had the following error while publishing your prompt template: {str(e)}") from e
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while publishing your prompt template: {str(e)}",
+                response=None,
+                body=None,
+            ) from e
+        logger.warning(f"PromptLayer had the following error while publishing your prompt template: {e}")
+        return None
 
 
+@retry_on_api_error
 def get_all_prompt_templates(
-    page: int = 1, per_page: int = 30, api_key: str = None, label: str = None
+    api_key: str, base_url: str, throw_on_error: bool, page: int = 1, per_page: int = 30, label: str = None
 ) -> List[ListPromptTemplateResponse]:
     try:
         params = {"page": page, "per_page": per_page}
         if label:
             params["label"] = label
         response = requests.get(
-            f"{URL_API_PROMPTLAYER}/prompt-templates",
+            f"{base_url}/prompt-templates",
             headers={"X-API-KEY": api_key},
             params=params,
         )
         if response.status_code != 200:
-            raise Exception(
-                f"PromptLayer had the following error while getting all your prompt templates: {response.text}"
-            )
+            if throw_on_error:
+                raise_on_bad_response(
+                    response, "PromptLayer had the following error while getting all your prompt templates"
+                )
+            else:
+                warn_on_bad_response(
+                    response, "WARNING: PromptLayer had the following error while getting all your prompt templates"
+                )
+                return []
         items = response.json().get("items", [])
         return items
     except requests.exceptions.RequestException as e:
-        raise Exception(f"PromptLayer had the following error while getting all your prompt templates: {e}")
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while getting all your prompt templates: {e}",
+                response=None,
+                body=None,
+            ) from e
+        logger.warning(f"PromptLayer had the following error while getting all your prompt templates: {e}")
+        return []
 
 
+@retry_on_api_error
 async def aget_all_prompt_templates(
-    page: int = 1, per_page: int = 30, api_key: str = None, label: str = None
+    api_key: str, base_url: str, throw_on_error: bool, page: int = 1, per_page: int = 30, label: str = None
 ) -> List[ListPromptTemplateResponse]:
     try:
         params = {"page": page, "per_page": per_page}
@@ -1223,22 +1587,33 @@ async def aget_all_prompt_templates(
             params["label"] = label
         async with _make_httpx_client() as client:
             response = await client.get(
-                f"{URL_API_PROMPTLAYER}/prompt-templates",
+                f"{base_url}/prompt-templates",
                 headers={"X-API-KEY": api_key},
                 params=params,
             )
 
-        if RAISE_FOR_STATUS:
-            response.raise_for_status()
-        elif response.status_code != 200:
-            raise_on_bad_response(
-                response,
-                "PromptLayer had the following error while getting all your prompt templates",
-            )
+        if response.status_code != 200:
+            if throw_on_error:
+                raise_on_bad_response(
+                    response,
+                    "PromptLayer had the following error while getting all your prompt templates",
+                )
+            else:
+                warn_on_bad_response(
+                    response, "WARNING: PromptLayer had the following error while getting all your prompt templates"
+                )
+                return []
         items = response.json().get("items", [])
         return items
     except httpx.RequestError as e:
-        raise Exception(f"PromptLayer had the following error while getting all your prompt templates: {str(e)}") from e
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIConnectionError(
+                f"PromptLayer had the following error while getting all your prompt templates: {str(e)}",
+                response=None,
+                body=None,
+            ) from e
+        logger.warning(f"PromptLayer had the following error while getting all your prompt templates: {e}")
+        return []
 
 
 def openai_chat_request(client, **kwargs):
@@ -1259,7 +1634,7 @@ def openai_request(prompt_blueprint: GetPromptTemplateResponse, client_kwargs: d
     from openai import OpenAI
 
     client = OpenAI(**client_kwargs)
-    api_type = prompt_blueprint["metadata"]["model"]["api_type"]
+    api_type = prompt_blueprint["metadata"]["model"].get("api_type", "chat-completions")
 
     if api_type == "chat-completions":
         request_to_make = MAP_TYPE_TO_OPENAI_FUNCTION[prompt_blueprint["prompt_template"]["type"]]
@@ -1286,7 +1661,7 @@ async def aopenai_request(prompt_blueprint: GetPromptTemplateResponse, client_kw
     from openai import AsyncOpenAI
 
     client = AsyncOpenAI(**client_kwargs)
-    api_type = prompt_blueprint["metadata"]["model"]["api_type"]
+    api_type = prompt_blueprint["metadata"]["model"].get("api_type", "chat-completions")
 
     if api_type == "chat-completions":
         request_to_make = AMAP_TYPE_TO_OPENAI_FUNCTION[prompt_blueprint["prompt_template"]["type"]]
@@ -1299,7 +1674,7 @@ def azure_openai_request(prompt_blueprint: GetPromptTemplateResponse, client_kwa
     from openai import AzureOpenAI
 
     client = AzureOpenAI(azure_endpoint=client_kwargs.pop("base_url", None))
-    api_type = prompt_blueprint["metadata"]["model"]["api_type"]
+    api_type = prompt_blueprint["metadata"]["model"].get("api_type", "chat-completions")
 
     if api_type == "chat-completions":
         request_to_make = MAP_TYPE_TO_OPENAI_FUNCTION[prompt_blueprint["prompt_template"]["type"]]
@@ -1314,7 +1689,7 @@ async def aazure_openai_request(
     from openai import AsyncAzureOpenAI
 
     client = AsyncAzureOpenAI(azure_endpoint=client_kwargs.pop("base_url", None))
-    api_type = prompt_blueprint["metadata"]["model"]["api_type"]
+    api_type = prompt_blueprint["metadata"]["model"].get("api_type", "chat-completions")
 
     if api_type == "chat-completions":
         request_to_make = AMAP_TYPE_TO_OPENAI_FUNCTION[prompt_blueprint["prompt_template"]["type"]]
@@ -1372,54 +1747,66 @@ def get_api_key():
     # raise an error if the api key is not set
     api_key = os.environ.get("PROMPTLAYER_API_KEY")
     if not api_key:
-        raise Exception(
-            "Please set your PROMPTLAYER_API_KEY environment variable or set API KEY in code using 'promptlayer.api_key = <your_api_key>' "
+        raise _exceptions.PromptLayerAuthenticationError(
+            "Please set your PROMPTLAYER_API_KEY environment variable or set API KEY in code using 'promptlayer.api_key = <your_api_key>'",
+            response=None,
+            body=None,
         )
     return api_key
 
 
-def util_log_request(api_key: str, **kwargs) -> Union[RequestLog, None]:
+@retry_on_api_error
+def util_log_request(api_key: str, base_url: str, throw_on_error: bool, **kwargs) -> Union[RequestLog, None]:
     try:
         response = requests.post(
-            f"{URL_API_PROMPTLAYER}/log-request",
+            f"{base_url}/log-request",
             headers={"X-API-KEY": api_key},
             json=kwargs,
         )
         if response.status_code != 201:
-            warn_on_bad_response(
-                response,
-                "WARNING: While logging your request PromptLayer had the following error",
-            )
-            return None
+            if throw_on_error:
+                raise_on_bad_response(response, "PromptLayer had the following error while logging your request")
+            else:
+                warn_on_bad_response(
+                    response,
+                    "WARNING: While logging your request PromptLayer had the following error",
+                )
+                return None
         return response.json()
     except Exception as e:
-        print(
-            f"WARNING: While tracking your prompt PromptLayer had the following error: {e}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIError(
+                f"While logging your request PromptLayer had the following error: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your prompt PromptLayer had the following error: {e}")
         return None
 
 
-async def autil_log_request(api_key: str, **kwargs) -> Union[RequestLog, None]:
+@retry_on_api_error
+async def autil_log_request(api_key: str, base_url: str, throw_on_error: bool, **kwargs) -> Union[RequestLog, None]:
     try:
         async with _make_httpx_client() as client:
             response = await client.post(
-                f"{URL_API_PROMPTLAYER}/log-request",
+                f"{base_url}/log-request",
                 headers={"X-API-KEY": api_key},
                 json=kwargs,
             )
         if response.status_code != 201:
-            warn_on_bad_response(
-                response,
-                "WARNING: While logging your request PromptLayer had the following error",
-            )
-            return None
+            if throw_on_error:
+                raise_on_bad_response(response, "PromptLayer had the following error while logging your request")
+            else:
+                warn_on_bad_response(
+                    response,
+                    "WARNING: While logging your request PromptLayer had the following error",
+                )
+                return None
         return response.json()
     except Exception as e:
-        print(
-            f"WARNING: While tracking your prompt PromptLayer had the following error: {e}",
-            file=sys.stderr,
-        )
+        if throw_on_error:
+            raise _exceptions.PromptLayerAPIError(
+                f"While logging your request PromptLayer had the following error: {e}", response=None, body=None
+            ) from e
+        logger.warning(f"While tracking your prompt PromptLayer had the following error: {e}")
         return None
 
 
@@ -1448,6 +1835,26 @@ async def amistral_request(
     return await client.chat.complete_async(**function_kwargs)
 
 
+class _GoogleStreamWrapper:
+    """Wrapper to keep Google client alive during streaming."""
+
+    def __init__(self, stream_generator, client):
+        self._stream = stream_generator
+        self._client = client  # Keep client alive
+
+    def __iter__(self):
+        return self._stream.__iter__()
+
+    def __next__(self):
+        return next(self._stream)
+
+    def __aiter__(self):
+        return self._stream.__aiter__()
+
+    async def __anext__(self):
+        return await self._stream.__anext__()
+
+
 def google_chat_request(client, **kwargs):
     from google.genai.chats import Content
 
@@ -1458,7 +1865,8 @@ def google_chat_request(client, **kwargs):
     chat = client.chats.create(model=model, history=history, config=generation_config)
     last_message = history[-1].parts if history else ""
     if stream:
-        return chat.send_message_stream(message=last_message)
+        stream_gen = chat.send_message_stream(message=last_message)
+        return _GoogleStreamWrapper(stream_gen, client)
     return chat.send_message(message=last_message)
 
 
@@ -1468,7 +1876,8 @@ def google_completions_request(client, **kwargs):
     contents = kwargs.get("contents", [])
     stream = kwargs.pop("stream", False)
     if stream:
-        return client.models.generate_content_stream(model=model, contents=contents, config=config)
+        stream_gen = client.models.generate_content_stream(model=model, contents=contents, config=config)
+        return _GoogleStreamWrapper(stream_gen, client)
     return client.models.generate_content(model=model, contents=contents, config=config)
 
 
@@ -1503,7 +1912,8 @@ async def agoogle_chat_request(client, **kwargs):
     chat = client.aio.chats.create(model=model, history=history, config=generation_config)
     last_message = history[-1].parts[0] if history else ""
     if stream:
-        return await chat.send_message_stream(message=last_message)
+        stream_gen = await chat.send_message_stream(message=last_message)
+        return _GoogleStreamWrapper(stream_gen, client)
     return await chat.send_message(message=last_message)
 
 
@@ -1513,8 +1923,9 @@ async def agoogle_completions_request(client, **kwargs):
     contents = kwargs.get("contents", [])
     stream = kwargs.pop("stream", False)
     if stream:
-        return await client.aio.models.generate_content_stream(model=model, contents=contents, config=config)
-        return await client.aio.models.generate_content(model=model, contents=contents, config=config)
+        stream_gen = await client.aio.models.generate_content_stream(model=model, contents=contents, config=config)
+        return _GoogleStreamWrapper(stream_gen, client)
+    return await client.aio.models.generate_content(model=model, contents=contents, config=config)
 
 
 AMAP_TYPE_TO_GOOGLE_FUNCTION = {