promptlayer 1.0.35__py3-none-any.whl → 1.0.78__py3-none-any.whl

promptlayer/promptlayer.py

@@ -1,31 +1,64 @@
 import asyncio
+import json
+import logging
 import os
 from typing import Any, Dict, List, Literal, Optional, Union
 
 import nest_asyncio
 
+from promptlayer import exceptions as _exceptions
 from promptlayer.groups import AsyncGroupManager, GroupManager
 from promptlayer.promptlayer_base import PromptLayerBase
 from promptlayer.promptlayer_mixins import PromptLayerMixin
+from promptlayer.streaming import astream_response, stream_response
 from promptlayer.templates import AsyncTemplateManager, TemplateManager
 from promptlayer.track import AsyncTrackManager, TrackManager
 from promptlayer.types.prompt_template import PromptTemplate
 from promptlayer.utils import (
+    RERAISE_ORIGINAL_EXCEPTION,
+    _get_workflow_workflow_id_or_name,
     arun_workflow_request,
-    astream_response,
     atrack_request,
     autil_log_request,
-    stream_response,
     track_request,
     util_log_request,
 )
 
+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
+
+    required_keys = {
+        "status",
+        "value",
+        "error_message",
+        "raw_error_message",
+        "is_output_node",
+    }
+
+    for val in obj.values():
+        if not isinstance(val, dict):
+            return False
+        if not required_keys.issubset(val.keys()):
+            return False
+
+    return True
+
 
 class PromptLayer(PromptLayerMixin):
     def __init__(
         self,
-        api_key: str = None,
+        api_key: Union[str, None] = None,
         enable_tracing: bool = False,
+        base_url: Union[str, None] = None,
+        throw_on_error: bool = True,
     ):
         if api_key is None:
             api_key = os.environ.get("PROMPTLAYER_API_KEY")
@@ -36,13 +69,15 @@ class PromptLayer(PromptLayerMixin):
                 "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.throw_on_error = throw_on_error
+        self.templates = TemplateManager(api_key, self.base_url, self.throw_on_error)
+        self.group = GroupManager(api_key, self.base_url, self.throw_on_error)
         self.tracer_provider, self.tracer = self._initialize_tracer(
-            api_key, enable_tracing
+            api_key, self.base_url, self.throw_on_error, enable_tracing
         )
-        self.track = TrackManager(api_key)
+        self.track = TrackManager(api_key, self.base_url, self.throw_on_error)
 
     def __getattr__(
         self,
@@ -51,24 +86,20 @@ class PromptLayer(PromptLayerMixin):
         if name == "openai":
             import openai as openai_module
 
-            openai = 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
             )
-            return openai
         elif name == "anthropic":
             import anthropic as anthropic_module
 
-            anthropic = PromptLayerBase(
+            return 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
         else:
             raise AttributeError(f"module {__name__} has no attribute {name}")
 
@@ -80,6 +111,7 @@ class PromptLayer(PromptLayerMixin):
         input_variables,
         group_id,
         pl_run_span_id: Union[str, None] = None,
+        request_start_time: Union[float, None] = None,
     ):
         def _track_request(**body):
             track_request_kwargs = self._prepare_track_request_kwargs(
@@ -89,9 +121,10 @@ class PromptLayer(PromptLayerMixin):
                 input_variables,
                 group_id,
                 pl_run_span_id,
+                request_start_time=request_start_time,
                 **body,
             )
-            return track_request(**track_request_kwargs)
+            return track_request(self.base_url, self.throw_on_error, **track_request_kwargs)
 
         return _track_request
 
@@ -108,50 +141,82 @@ class PromptLayer(PromptLayerMixin):
         group_id: Union[int, None] = None,
         stream: bool = False,
         pl_run_span_id: Union[str, None] = None,
+        provider: Union[str, None] = None,
+        model: Union[str, None] = None,
     ) -> Dict[str, Any]:
+        import datetime
+
         get_prompt_template_params = self._prepare_get_prompt_template_params(
             prompt_version=prompt_version,
             prompt_release_label=prompt_release_label,
             input_variables=input_variables,
             metadata=metadata,
+            provider=provider,
+            model=model,
+            model_parameter_overrides=model_parameter_overrides,
         )
         prompt_blueprint = self.templates.get(prompt_name, get_prompt_template_params)
+        if not prompt_blueprint:
+            raise _exceptions.PromptLayerNotFoundError(
+                f"Prompt template '{prompt_name}' not found.",
+                response=None,
+                body=None,
+            )
         prompt_blueprint_model = self._validate_and_extract_model_from_prompt_blueprint(
             prompt_blueprint=prompt_blueprint, prompt_name=prompt_name
         )
-        llm_request_params = self._prepare_llm_request_params(
+        llm_data = self._prepare_llm_data(
             prompt_blueprint=prompt_blueprint,
             prompt_template=prompt_blueprint["prompt_template"],
             prompt_blueprint_model=prompt_blueprint_model,
-            model_parameter_overrides=model_parameter_overrides,
             stream=stream,
         )
 
-        response = llm_request_params["request_function"](
-            llm_request_params["prompt_blueprint"], **llm_request_params["kwargs"]
+        # Capture start time before making the LLM request
+        request_start_time = datetime.datetime.now(datetime.timezone.utc).timestamp()
+
+        # response is just whatever the LLM call returns
+        # streaming=False > Pydantic model instance
+        # streaming=True > generator that yields ChatCompletionChunk pieces as they arrive
+        response = llm_data["request_function"](
+            prompt_blueprint=llm_data["prompt_blueprint"],
+            client_kwargs=llm_data["client_kwargs"],
+            function_kwargs=llm_data["function_kwargs"],
         )
 
+        # Capture end time after the LLM request completes
+        request_end_time = datetime.datetime.now(datetime.timezone.utc).timestamp()
+
         if stream:
             return stream_response(
-                response,
-                self._create_track_request_callable(
-                    request_params=llm_request_params,
+                generator=response,
+                after_stream=self._create_track_request_callable(
+                    request_params=llm_data,
                     tags=tags,
                     input_variables=input_variables,
                     group_id=group_id,
                     pl_run_span_id=pl_run_span_id,
+                    request_start_time=request_start_time,
                 ),
-                llm_request_params["stream_function"],
+                map_results=llm_data["stream_function"],
+                metadata=llm_data["prompt_blueprint"]["metadata"],
             )
 
+        if isinstance(response, dict):
+            request_response = response
+        else:
+            request_response = response.model_dump(mode="json")
+
         request_log = self._track_request_log(
-            llm_request_params,
+            llm_data,
             tags,
             input_variables,
             group_id,
             pl_run_span_id,
             metadata=metadata,
-            request_response=response.model_dump(),
+            request_response=request_response,
+            request_start_time=request_start_time,
+            request_end_time=request_end_time,
         )
 
         return {
@@ -180,7 +245,7 @@ class PromptLayer(PromptLayerMixin):
             metadata=metadata,
             **body,
         )
-        return track_request(**track_request_kwargs)
+        return track_request(self.base_url, self.throw_on_error, **track_request_kwargs)
 
     def run(
         self,
@@ -193,17 +258,21 @@ class PromptLayer(PromptLayerMixin):
         metadata: Union[Dict[str, str], None] = None,
         group_id: Union[int, None] = None,
         stream: bool = False,
+        provider: Union[str, None] = None,
+        model: Union[str, None] = None,
     ) -> Dict[str, Any]:
         _run_internal_kwargs = {
             "prompt_name": prompt_name,
             "prompt_version": prompt_version,
             "prompt_release_label": prompt_release_label,
-            "input_variables": input_variables,
+            "input_variables": input_variables or {},
             "model_parameter_overrides": model_parameter_overrides,
             "tags": tags,
             "metadata": metadata,
             "group_id": group_id,
             "stream": stream,
+            "provider": provider,
+            "model": model,
         }
 
         if self.tracer:
@@ -211,9 +280,7 @@ class PromptLayer(PromptLayerMixin):
                 span.set_attribute("prompt_name", prompt_name)
                 span.set_attribute("function_input", str(_run_internal_kwargs))
                 pl_run_span_id = hex(span.context.span_id)[2:].zfill(16)
-                result = self._run_internal(
-                    **_run_internal_kwargs, pl_run_span_id=pl_run_span_id
-                )
+                result = self._run_internal(**_run_internal_kwargs, pl_run_span_id=pl_run_span_id)
                 span.set_attribute("function_output", str(result))
                 return result
         else:
@@ -221,51 +288,63 @@ class PromptLayer(PromptLayerMixin):
 
     def run_workflow(
         self,
-        workflow_name: str,
+        workflow_id_or_name: Optional[Union[int, str]] = None,
         input_variables: Optional[Dict[str, Any]] = None,
         metadata: Optional[Dict[str, str]] = None,
         workflow_label_name: Optional[str] = None,
-        workflow_version: Optional[
-            int
-        ] = None,  # This is the version number, not the version ID
+        workflow_version: Optional[int] = None,
         return_all_outputs: Optional[bool] = False,
-    ) -> Dict[str, Any]:
+        # `workflow_name` deprecated, kept for backward compatibility only.
+        # Allows `workflow_name` to be passed both as keyword and positional argument
+        # (virtually identical to `workflow_id_or_name`)
+        workflow_name: Optional[str] = None,
+    ) -> Union[Dict[str, Any], Any]:
         try:
             try:
-                # Check if we're inside a running event loop
-                loop = asyncio.get_running_loop()
+                loop = asyncio.get_running_loop()  # Check if we're inside a running event loop
             except RuntimeError:
                 loop = None
 
             if loop and loop.is_running():
                 nest_asyncio.apply()
-                # If there's an active event loop, use `await` directly
-                return asyncio.run(
-                    arun_workflow_request(
-                        workflow_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,
-                    )
+
+            results = asyncio.run(
+                arun_workflow_request(
+                    api_key=self.api_key,
+                    base_url=self.base_url,
+                    throw_on_error=self.throw_on_error,
+                    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,
+                    return_all_outputs=return_all_outputs,
                 )
-            else:
-                # If there's no active event loop, use `asyncio.run()`
-                return asyncio.run(
-                    arun_workflow_request(
-                        workflow_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,
+            )
+
+            if not return_all_outputs and is_workflow_results_dict(results):
+                output_nodes = [node_data for node_data in results.values() if node_data.get("is_output_node")]
+                if not output_nodes:
+                    raise _exceptions.PromptLayerNotFoundError(
+                        f"Output nodes not found: {json.dumps(results, indent=4)}", response=None, body=results
                     )
-                )
-        except Exception as e:
-            raise Exception(f"Error running workflow: {str(e)}")
+
+                if not any(node.get("status") == "SUCCESS" for node in output_nodes):
+                    raise _exceptions.PromptLayerAPIError(
+                        f"None of the output nodes have succeeded: {json.dumps(results, indent=4)}",
+                        response=None,
+                        body=results,
+                    )
+
+            return results
+        except Exception as ex:
+            logger.exception("Error running workflow")
+            if RERAISE_ORIGINAL_EXCEPTION:
+                raise
+            else:
+                raise _exceptions.PromptLayerAPIError(
+                    f"Error running workflow: {str(ex)}", response=None, body=None
+                ) from ex
 
     def log_request(
         self,
@@ -276,6 +355,8 @@ class PromptLayer(PromptLayerMixin):
         output: PromptTemplate,
         request_start_time: float,
         request_end_time: float,
+        # TODO(dmu) MEDIUM: Avoid using mutable defaults
+        # TODO(dmu) MEDIUM: Deprecate and remove this wrapper function?
         parameters: Dict[str, Any] = {},
         tags: List[str] = [],
         metadata: Dict[str, str] = {},
@@ -287,9 +368,13 @@ class PromptLayer(PromptLayerMixin):
         price: float = 0.0,
         function_name: str = "",
         score: int = 0,
+        prompt_id: Union[int, None] = None,
+        score_name: Union[str, None] = None,
     ):
         return util_log_request(
             self.api_key,
+            self.base_url,
+            throw_on_error=self.throw_on_error,
             provider=provider,
             model=model,
             input=input,
@@ -307,14 +392,18 @@ class PromptLayer(PromptLayerMixin):
             price=price,
             function_name=function_name,
             score=score,
+            prompt_id=prompt_id,
+            score_name=score_name,
         )
 
 
 class AsyncPromptLayer(PromptLayerMixin):
     def __init__(
         self,
-        api_key: str = None,
+        api_key: Union[str, None] = None,
         enable_tracing: bool = False,
+        base_url: Union[str, None] = None,
+        throw_on_error: bool = True,
     ):
         if api_key is None:
             api_key = os.environ.get("PROMPTLAYER_API_KEY")
@@ -325,34 +414,34 @@ class AsyncPromptLayer(PromptLayerMixin):
                 "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.throw_on_error = throw_on_error
+        self.templates = AsyncTemplateManager(api_key, self.base_url, self.throw_on_error)
+        self.group = AsyncGroupManager(api_key, self.base_url, self.throw_on_error)
         self.tracer_provider, self.tracer = self._initialize_tracer(
-            api_key, enable_tracing
+            api_key, self.base_url, self.throw_on_error, enable_tracing
         )
-        self.track = AsyncTrackManager(api_key)
+        self.track = AsyncTrackManager(api_key, self.base_url, self.throw_on_error)
 
-    def __getattr__(
-        self, name: Union[Literal["openai"], Literal["anthropic"], Literal["prompts"]]
-    ):
+    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,
+                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
         else:
@@ -360,28 +449,37 @@ class AsyncPromptLayer(PromptLayerMixin):
 
     async def run_workflow(
         self,
-        workflow_name: str,
+        workflow_id_or_name: Optional[Union[int, str]] = None,
         input_variables: Optional[Dict[str, Any]] = None,
         metadata: Optional[Dict[str, str]] = None,
         workflow_label_name: Optional[str] = None,
-        workflow_version: Optional[
-            int
-        ] = None,  # This is the version number, not the version ID
+        workflow_version: Optional[int] = None,  # This is the version number, not the version ID
         return_all_outputs: Optional[bool] = False,
-    ) -> Dict[str, Any]:
+        # `workflow_name` deprecated, kept for backward compatibility only.
+        # Allows `workflow_name` to be passed both as keyword and positional argument
+        # (virtually identical to `workflow_id_or_name`)
+        workflow_name: Optional[str] = None,
+    ) -> Union[Dict[str, Any], Any]:
         try:
-            result = await arun_workflow_request(
-                workflow_name=workflow_name,
+            return await arun_workflow_request(
+                api_key=self.api_key,
+                base_url=self.base_url,
+                throw_on_error=self.throw_on_error,
+                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,
             )
-            return result
-        except Exception as e:
-            raise Exception(f"Error running workflow: {str(e)}")
+        except Exception as ex:
+            logger.exception("Error running workflow")
+            if RERAISE_ORIGINAL_EXCEPTION:
+                raise
+            else:
+                raise _exceptions.PromptLayerAPIError(
+                    f"Error running workflow: {str(ex)}", response=None, body=None
+                ) from ex
 
     async def run(
         self,
@@ -394,6 +492,8 @@ class AsyncPromptLayer(PromptLayerMixin):
         metadata: Union[Dict[str, str], None] = None,
         group_id: Union[int, None] = None,
         stream: bool = False,
+        provider: Union[str, None] = None,
+        model: Union[str, None] = None,
     ) -> Dict[str, Any]:
         _run_internal_kwargs = {
             "prompt_name": prompt_name,
@@ -405,6 +505,8 @@ class AsyncPromptLayer(PromptLayerMixin):
             "metadata": metadata,
             "group_id": group_id,
             "stream": stream,
+            "provider": provider,
+            "model": model,
         }
 
         if self.tracer:
@@ -412,9 +514,7 @@ class AsyncPromptLayer(PromptLayerMixin):
                 span.set_attribute("prompt_name", prompt_name)
                 span.set_attribute("function_input", str(_run_internal_kwargs))
                 pl_run_span_id = hex(span.context.span_id)[2:].zfill(16)
-                result = await self._run_internal(
-                    **_run_internal_kwargs, pl_run_span_id=pl_run_span_id
-                )
+                result = await self._run_internal(**_run_internal_kwargs, pl_run_span_id=pl_run_span_id)
                 span.set_attribute("function_output", str(result))
                 return result
         else:
@@ -440,9 +540,12 @@ class AsyncPromptLayer(PromptLayerMixin):
         price: float = 0.0,
         function_name: str = "",
         score: int = 0,
+        prompt_id: Union[int, None] = None,
     ):
         return await autil_log_request(
             self.api_key,
+            self.base_url,
+            throw_on_error=self.throw_on_error,
             provider=provider,
             model=model,
             input=input,
@@ -460,6 +563,7 @@ class AsyncPromptLayer(PromptLayerMixin):
             price=price,
             function_name=function_name,
             score=score,
+            prompt_id=prompt_id,
         )
 
     async def _create_track_request_callable(
@@ -470,6 +574,7 @@ class AsyncPromptLayer(PromptLayerMixin):
         input_variables,
         group_id,
         pl_run_span_id: Union[str, None] = None,
+        request_start_time: Union[float, None] = None,
     ):
         async def _track_request(**body):
             track_request_kwargs = self._prepare_track_request_kwargs(
@@ -479,9 +584,10 @@ class AsyncPromptLayer(PromptLayerMixin):
                 input_variables,
                 group_id,
                 pl_run_span_id,
+                request_start_time=request_start_time,
                 **body,
             )
-            return await atrack_request(**track_request_kwargs)
+            return await atrack_request(self.base_url, self.throw_on_error, **track_request_kwargs)
 
         return _track_request
 
@@ -505,7 +611,7 @@ class AsyncPromptLayer(PromptLayerMixin):
             metadata=metadata,
             **body,
         )
-        return await atrack_request(**track_request_kwargs)
+        return await atrack_request(self.base_url, self.throw_on_error, **track_request_kwargs)
 
     async def _run_internal(
         self,
@@ -520,54 +626,81 @@ class AsyncPromptLayer(PromptLayerMixin):
         group_id: Union[int, None] = None,
         stream: bool = False,
         pl_run_span_id: Union[str, None] = None,
+        provider: Union[str, None] = None,
+        model: Union[str, None] = None,
     ) -> Dict[str, Any]:
+        import datetime
+
         get_prompt_template_params = self._prepare_get_prompt_template_params(
             prompt_version=prompt_version,
             prompt_release_label=prompt_release_label,
             input_variables=input_variables,
             metadata=metadata,
+            provider=provider,
+            model=model,
+            model_parameter_overrides=model_parameter_overrides,
         )
-        prompt_blueprint = await self.templates.get(
-            prompt_name, get_prompt_template_params
-        )
+        prompt_blueprint = await self.templates.get(prompt_name, get_prompt_template_params)
+        if not prompt_blueprint:
+            raise _exceptions.PromptLayerNotFoundError(
+                f"Prompt template '{prompt_name}' not found.",
+                response=None,
+                body=None,
+            )
         prompt_blueprint_model = self._validate_and_extract_model_from_prompt_blueprint(
             prompt_blueprint=prompt_blueprint, prompt_name=prompt_name
         )
-        llm_request_params = self._prepare_llm_request_params(
+        llm_data = self._prepare_llm_data(
             prompt_blueprint=prompt_blueprint,
             prompt_template=prompt_blueprint["prompt_template"],
             prompt_blueprint_model=prompt_blueprint_model,
-            model_parameter_overrides=model_parameter_overrides,
             stream=stream,
             is_async=True,
         )
 
-        response = await llm_request_params["request_function"](
-            llm_request_params["prompt_blueprint"], **llm_request_params["kwargs"]
+        # Capture start time before making the LLM request
+        request_start_time = datetime.datetime.now(datetime.timezone.utc).timestamp()
+
+        response = await llm_data["request_function"](
+            prompt_blueprint=llm_data["prompt_blueprint"],
+            client_kwargs=llm_data["client_kwargs"],
+            function_kwargs=llm_data["function_kwargs"],
         )
 
+        # Capture end time after the LLM request completes
+        request_end_time = datetime.datetime.now(datetime.timezone.utc).timestamp()
+
+        if hasattr(response, "model_dump"):
+            request_response = response.model_dump(mode="json")
+        else:
+            request_response = response
+
         if stream:
             track_request_callable = await self._create_track_request_callable(
-                request_params=llm_request_params,
+                request_params=llm_data,
                 tags=tags,
                 input_variables=input_variables,
                 group_id=group_id,
                 pl_run_span_id=pl_run_span_id,
+                request_start_time=request_start_time,
             )
             return astream_response(
-                response,
+                request_response,
                 track_request_callable,
-                llm_request_params["stream_function"],
+                llm_data["stream_function"],
+                llm_data["prompt_blueprint"]["metadata"],
             )
 
         request_log = await self._track_request_log(
-            llm_request_params,
+            llm_data,
             tags,
             input_variables,
             group_id,
             pl_run_span_id,
             metadata=metadata,
-            request_response=response.model_dump(),
+            request_response=request_response,
+            request_start_time=request_start_time,
+            request_end_time=request_end_time,
         )
 
         return {