Refactor PromptLayer API integration to include base URL parameter

Author: abubakarsohail

.devcontainer/devcontainer.json

@@ -5,7 +5,9 @@
   // Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile
   "image": "mcr.microsoft.com/devcontainers/python:1-3.9-bullseye",
   "features": {
-    "ghcr.io/devcontainers-extra/features/poetry:2": {}
+    "ghcr.io/devcontainers-extra/features/poetry:2": {
+      "version": "2.2.1"
+    }
   },
   "customizations": {
     "vscode": {
@@ -24,19 +26,14 @@
       }
     }
   },
-
   // Features to add to the dev container. More info: https://containers.dev/features.
   // "features": {},
-
   // Use 'forwardPorts' to make a list of ports inside the container available locally.
   // "forwardPorts": [],
-
   // Use 'postCreateCommand' to run commands after the container is created.
   "postCreateCommand": "poetry install"
-
   // Configure tool-specific properties.
   // "customizations": {},
-
   // Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
   // "remoteUser": "root",
 }

.github/workflows/integration-tests.yml

@@ -1,29 +1,31 @@
 name: Integration Tests
 
 on:
-  push:
-    branches: [$default-branch]
-
-  workflow_dispatch:
+  pull_request:
+    branches: ["master"]
 
 env:
-  PROMPTLAYER_API_KEY: ${{ secrets.PROMPTLAYER_API_KEY }}
-  OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+  POETRY_VIRTUALENVS_CREATE: "false"
+  PYTHON_VERSION: "3.9"
+  POETRY_VERSION: "2.2.1"
+  PIPX_DEFAULT_PYTHON: "python3.9"
 
 jobs:
   integration-tests:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v5
       - name: Set up Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v6
         with:
-          python-version: "3.x"
+          python-version: ${{ env.PYTHON_VERSION }}
+
+      - name: Install Poetry
+        run: pipx install poetry==${{ env.POETRY_VERSION }}
+
       - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install behave openai
-          if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-      - name: Run Integration Tests
-        run: behave
+        run: poetry install
+
+      - name: Run integration tests
+        run: poetry run pytest

poetry.lock

@@ -1,4 +1,4 @@
 from .promptlayer import AsyncPromptLayer, PromptLayer
 
-__version__ = "1.0.71"
+__version__ = "1.0.72"
 __all__ = ["PromptLayer", "AsyncPromptLayer", "__version__"]

promptlayer/__init__.py

@@ -2,19 +2,21 @@
 
 
 class GroupManager:
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str, base_url: str):
         self.api_key = api_key
+        self.base_url = base_url
 
     def create(self):
-        return create(self.api_key)
+        return create(self.api_key, self.base_url)
 
 
 class AsyncGroupManager:
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str, base_url: str):
         self.api_key = api_key
+        self.base_url = base_url
 
-    async def create(self) -> str:
-        return await acreate(self.api_key)
+    async def create(self):
+        return await acreate(self.api_key, self.base_url)
 
 
 __all__ = ["GroupManager", "AsyncGroupManager"]

promptlayer/groups/__init__.py

@@ -1,11 +1,9 @@
 from promptlayer.utils import apromptlayer_create_group, promptlayer_create_group
 
 
-def create(api_key: str = None):
-    """Create a new group."""
-    return promptlayer_create_group(api_key)
+def create(api_key: str, base_url: str):
+    return promptlayer_create_group(api_key, base_url)
 
 
-async def acreate(api_key: str = None) -> str:
-    """Asynchronously create a new group."""
-    return await apromptlayer_create_group(api_key)
+async def acreate(api_key: str, base_url: str):
+    return await apromptlayer_create_group(api_key, base_url)

promptlayer/groups/groups.py

@@ -26,6 +26,10 @@
 logger = logging.getLogger(__name__)
 
 
+def get_base_url(base_url: Union[str, None]):
+    return base_url or os.environ.get("PROMPTLAYER_BASE_URL", "https://api.promptlayer.com")
+
+
 def is_workflow_results_dict(obj: Any) -> bool:
     if not isinstance(obj, dict):
         return False
@@ -49,9 +53,7 @@ def is_workflow_results_dict(obj: Any) -> bool:
 
 class PromptLayer(PromptLayerMixin):
     def __init__(
-        self,
-        api_key: str = None,
-        enable_tracing: bool = False,
+        self, api_key: Union[str, None] = None, enable_tracing: bool = False, base_url: Union[str, None] = None
     ):
         if api_key is None:
             api_key = os.environ.get("PROMPTLAYER_API_KEY")
@@ -62,11 +64,12 @@ def __init__(
                 "Please set the PROMPTLAYER_API_KEY environment variable or pass the api_key parameter."
             )
 
+        self.base_url = get_base_url(base_url)
         self.api_key = api_key
-        self.templates = TemplateManager(api_key)
-        self.group = GroupManager(api_key)
-        self.tracer_provider, self.tracer = self._initialize_tracer(api_key, enable_tracing)
-        self.track = TrackManager(api_key)
+        self.templates = TemplateManager(api_key, self.base_url)
+        self.group = GroupManager(api_key, self.base_url)
+        self.tracer_provider, self.tracer = self._initialize_tracer(api_key, self.base_url, enable_tracing)
+        self.track = TrackManager(api_key, self.base_url)
 
     def __getattr__(
         self,
@@ -75,15 +78,18 @@ def __getattr__(
         if name == "openai":
             import openai as openai_module
 
-            return PromptLayerBase(openai_module, function_name="openai", api_key=self.api_key, tracer=self.tracer)
+            return PromptLayerBase(
+                self.api_key, self.base_url, openai_module, function_name="openai", tracer=self.tracer
+            )
         elif name == "anthropic":
             import anthropic as anthropic_module
 
             return PromptLayerBase(
+                self.api_key,
+                self.base_url,
                 anthropic_module,
                 function_name="anthropic",
                 provider_type="anthropic",
-                api_key=self.api_key,
                 tracer=self.tracer,
             )
         else:
@@ -212,7 +218,7 @@ def _track_request_log(
             metadata=metadata,
             **body,
         )
-        return track_request(**track_request_kwargs)
+        return track_request(self.base_url, **track_request_kwargs)
 
     def run(
         self,
@@ -277,12 +283,13 @@ def run_workflow(
 
             results = asyncio.run(
                 arun_workflow_request(
+                    api_key=self.api_key,
+                    base_url=self.base_url,
                     workflow_id_or_name=_get_workflow_workflow_id_or_name(workflow_id_or_name, workflow_name),
                     input_variables=input_variables or {},
                     metadata=metadata,
                     workflow_label_name=workflow_label_name,
                     workflow_version_number=workflow_version,
-                    api_key=self.api_key,
                     return_all_outputs=return_all_outputs,
                 )
             )
@@ -330,6 +337,7 @@ def log_request(
     ):
         return util_log_request(
             self.api_key,
+            self.base_url,
             provider=provider,
             model=model,
             input=input,
@@ -354,9 +362,7 @@ def log_request(
 
 class AsyncPromptLayer(PromptLayerMixin):
     def __init__(
-        self,
-        api_key: str = None,
-        enable_tracing: bool = False,
+        self, api_key: Union[str, None] = None, enable_tracing: bool = False, base_url: Union[str, None] = None
     ):
         if api_key is None:
             api_key = os.environ.get("PROMPTLAYER_API_KEY")
@@ -367,31 +373,30 @@ def __init__(
                 "Please set the PROMPTLAYER_API_KEY environment variable or pass the api_key parameter."
             )
 
+        self.base_url = get_base_url(base_url)
         self.api_key = api_key
-        self.templates = AsyncTemplateManager(api_key)
-        self.group = AsyncGroupManager(api_key)
-        self.tracer_provider, self.tracer = self._initialize_tracer(api_key, enable_tracing)
-        self.track = AsyncTrackManager(api_key)
+        self.templates = AsyncTemplateManager(api_key, self.base_url)
+        self.group = AsyncGroupManager(api_key, self.base_url)
+        self.tracer_provider, self.tracer = self._initialize_tracer(api_key, self.base_url, enable_tracing)
+        self.track = AsyncTrackManager(api_key, self.base_url)
 
     def __getattr__(self, name: Union[Literal["openai"], Literal["anthropic"], Literal["prompts"]]):
         if name == "openai":
             import openai as openai_module
 
             openai = PromptLayerBase(
-                openai_module,
-                function_name="openai",
-                api_key=self.api_key,
-                tracer=self.tracer,
+                self.api_key, self.base_url, openai_module, function_name="openai", tracer=self.tracer
             )
             return openai
         elif name == "anthropic":
             import anthropic as anthropic_module
 
             anthropic = PromptLayerBase(
+                self.api_key,
+                self.base_url,
                 anthropic_module,
                 function_name="anthropic",
                 provider_type="anthropic",
-                api_key=self.api_key,
                 tracer=self.tracer,
             )
             return anthropic
@@ -413,12 +418,13 @@ async def run_workflow(
     ) -> Union[Dict[str, Any], Any]:
         try:
             return await arun_workflow_request(
+                api_key=self.api_key,
+                base_url=self.base_url,
                 workflow_id_or_name=_get_workflow_workflow_id_or_name(workflow_id_or_name, workflow_name),
                 input_variables=input_variables or {},
                 metadata=metadata,
                 workflow_label_name=workflow_label_name,
                 workflow_version_number=workflow_version,
-                api_key=self.api_key,
                 return_all_outputs=return_all_outputs,
             )
         except Exception as ex:
@@ -491,6 +497,7 @@ async def log_request(
     ):
         return await autil_log_request(
             self.api_key,
+            self.base_url,
             provider=provider,
             model=model,
             input=input,
@@ -530,7 +537,7 @@ async def _track_request(**body):
                 pl_run_span_id,
                 **body,
             )
-            return await atrack_request(**track_request_kwargs)
+            return await atrack_request(self.base_url, **track_request_kwargs)
 
         return _track_request
 
@@ -554,7 +561,7 @@ async def _track_request_log(
             metadata=metadata,
             **body,
         )
-        return await atrack_request(**track_request_kwargs)
+        return await atrack_request(self.base_url, **track_request_kwargs)
 
     async def _run_internal(
         self,
@@ -631,6 +638,6 @@ async def _run_internal(
 
         return {
             "request_id": request_log.get("request_id", None),
-            "raw_response": request_response,
+            "raw_response": response,
             "prompt_blueprint": request_log.get("prompt_blueprint", None),
         }