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

promptlayer/promptlayer.py

@@ -1,66 +1,64 @@
 import asyncio
-import datetime
+import json
+import logging
 import os
-from copy import deepcopy
-from functools import wraps
-from typing import Any, Dict, List, Literal, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
-from opentelemetry import trace
-from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from opentelemetry.semconv.resource import ResourceAttributes
+import nest_asyncio
 
-from promptlayer.groups import GroupManager
+from promptlayer import exceptions as _exceptions
+from promptlayer.groups import AsyncGroupManager, GroupManager
 from promptlayer.promptlayer_base import PromptLayerBase
-from promptlayer.span_exporter import PromptLayerSpanExporter
-from promptlayer.templates import TemplateManager
-from promptlayer.track import TrackManager
+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 (
-    anthropic_request,
-    anthropic_stream_completion,
-    anthropic_stream_message,
-    openai_request,
-    openai_stream_chat,
-    openai_stream_completion,
-    stream_response,
+    RERAISE_ORIGINAL_EXCEPTION,
+    _get_workflow_workflow_id_or_name,
+    arun_workflow_request,
+    atrack_request,
+    autil_log_request,
     track_request,
+    util_log_request,
 )
 
-MAP_PROVIDER_TO_FUNCTION_NAME = {
-    "openai": {
-        "chat": {
-            "function_name": "openai.chat.completions.create",
-            "stream_function": openai_stream_chat,
-        },
-        "completion": {
-            "function_name": "openai.completions.create",
-            "stream_function": openai_stream_completion,
-        },
-    },
-    "anthropic": {
-        "chat": {
-            "function_name": "anthropic.messages.create",
-            "stream_function": anthropic_stream_message,
-        },
-        "completion": {
-            "function_name": "anthropic.completions.create",
-            "stream_function": anthropic_stream_completion,
-        },
-    },
-}
-
-MAP_PROVIDER_TO_FUNCTION = {
-    "openai": openai_request,
-    "anthropic": anthropic_request,
-}
-
-
-class PromptLayer:
+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")
@@ -71,11 +69,15 @@ class PromptLayer:
                 "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 = self._initialize_tracer(api_key, enable_tracing)
-        self.track = TrackManager(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, self.base_url, self.throw_on_error, enable_tracing
+        )
+        self.track = TrackManager(api_key, self.base_url, self.throw_on_error)
 
     def __getattr__(
         self,
@@ -84,24 +86,20 @@ class PromptLayer:
         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}")
 
@@ -112,105 +110,24 @@ class PromptLayer:
         tags,
         input_variables,
         group_id,
-        pl_run_span_id: str | None = None,
+        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(
-                request_params, tags, input_variables, group_id, pl_run_span_id, **body
+                self.api_key,
+                request_params,
+                tags,
+                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
 
-    @staticmethod
-    def _initialize_tracer(api_key: str = None, enable_tracing: bool = False):
-        if enable_tracing:
-            resource = Resource(
-                attributes={ResourceAttributes.SERVICE_NAME: "prompt-layer-library"}
-            )
-            tracer_provider = TracerProvider(resource=resource)
-            promptlayer_exporter = PromptLayerSpanExporter(api_key=api_key)
-            span_processor = BatchSpanProcessor(promptlayer_exporter)
-            tracer_provider.add_span_processor(span_processor)
-            trace.set_tracer_provider(tracer_provider)
-            return trace.get_tracer(__name__)
-        else:
-            return None
-
-    @staticmethod
-    def _prepare_get_prompt_template_params(
-        *, prompt_version, prompt_release_label, input_variables, metadata
-    ):
-        params = {}
-
-        if prompt_version:
-            params["version"] = prompt_version
-        if prompt_release_label:
-            params["label"] = prompt_release_label
-        if input_variables:
-            params["input_variables"] = input_variables
-        if metadata:
-            params["metadata_filters"] = metadata
-
-        return params
-
-    @staticmethod
-    def _prepare_llm_request_params(
-        *, prompt_blueprint, prompt_template, prompt_blueprint_model, stream
-    ):
-        provider = prompt_blueprint_model["provider"]
-        kwargs = deepcopy(prompt_blueprint["llm_kwargs"])
-        config = MAP_PROVIDER_TO_FUNCTION_NAME[provider][prompt_template["type"]]
-
-        if provider_base_url := prompt_blueprint.get("provider_base_url"):
-            kwargs["base_url"] = provider_base_url["url"]
-
-        kwargs["stream"] = stream
-        if stream and provider == "openai":
-            kwargs["stream_options"] = {"include_usage": True}
-
-        return {
-            "provider": provider,
-            "function_name": config["function_name"],
-            "stream_function": config["stream_function"],
-            "request_function": MAP_PROVIDER_TO_FUNCTION[provider],
-            "kwargs": kwargs,
-            "prompt_blueprint": prompt_blueprint,
-        }
-
-    def _prepare_track_request_kwargs(
-        self,
-        request_params,
-        tags,
-        input_variables,
-        group_id,
-        pl_run_span_id: str | None = None,
-        metadata: Dict[str, str] | None = None,
-        **body,
-    ):
-        return {
-            "function_name": request_params["function_name"],
-            "provider_type": request_params["provider"],
-            "args": [],
-            "kwargs": request_params["kwargs"],
-            "tags": tags,
-            "request_start_time": datetime.datetime.now(
-                datetime.timezone.utc
-            ).timestamp(),
-            "request_end_time": datetime.datetime.now(
-                datetime.timezone.utc
-            ).timestamp(),
-            "api_key": self.api_key,
-            "metadata": metadata,
-            "prompt_id": request_params["prompt_blueprint"]["id"],
-            "prompt_version": request_params["prompt_blueprint"]["version"],
-            "prompt_input_variables": input_variables,
-            "group_id": group_id,
-            "return_prompt_blueprint": True,
-            "span_id": pl_run_span_id,
-            **body,
-        }
-
     def _run_internal(
         self,
         *,
@@ -218,54 +135,88 @@ class PromptLayer:
         prompt_version: Union[int, None] = None,
         prompt_release_label: Union[str, None] = None,
         input_variables: Union[Dict[str, Any], None] = None,
+        model_parameter_overrides: Union[Dict[str, Any], None] = None,
         tags: Union[List[str], None] = None,
         metadata: Union[Dict[str, str], None] = None,
         group_id: Union[int, None] = None,
         stream: bool = False,
-        pl_run_span_id: str | None = None,
+        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,
             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 {
@@ -280,11 +231,12 @@ class PromptLayer:
         tags,
         input_variables,
         group_id,
-        pl_run_span_id: str | None = None,
-        metadata: Dict[str, str] | None = None,
+        pl_run_span_id: Union[str, None] = None,
+        metadata: Union[Dict[str, str], None] = None,
         **body,
     ):
         track_request_kwargs = self._prepare_track_request_kwargs(
+            self.api_key,
             request_params,
             tags,
             input_variables,
@@ -293,53 +245,268 @@ class PromptLayer:
             metadata=metadata,
             **body,
         )
-        return track_request(**track_request_kwargs)
+        return track_request(self.base_url, self.throw_on_error, **track_request_kwargs)
 
-    @staticmethod
-    def _validate_and_extract_model_from_prompt_blueprint(
-        *, prompt_blueprint, prompt_name
-    ):
-        if not prompt_blueprint["llm_kwargs"]:
-            raise ValueError(
-                f"Prompt '{prompt_name}' does not have any LLM kwargs associated with it."
+    def run(
+        self,
+        prompt_name: str,
+        prompt_version: Union[int, None] = None,
+        prompt_release_label: Union[str, None] = None,
+        input_variables: Union[Dict[str, Any], None] = None,
+        model_parameter_overrides: Union[Dict[str, Any], None] = None,
+        tags: Union[List[str], None] = None,
+        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 or {},
+            "model_parameter_overrides": model_parameter_overrides,
+            "tags": tags,
+            "metadata": metadata,
+            "group_id": group_id,
+            "stream": stream,
+            "provider": provider,
+            "model": model,
+        }
+
+        if self.tracer:
+            with self.tracer.start_as_current_span("PromptLayer Run") as span:
+                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)
+                span.set_attribute("function_output", str(result))
+                return result
+        else:
+            return self._run_internal(**_run_internal_kwargs)
+
+    def run_workflow(
+        self,
+        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,
+        return_all_outputs: Optional[bool] = False,
+        # `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:
+                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()
+
+            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,
+                )
             )
 
-        prompt_blueprint_metadata = prompt_blueprint.get("metadata")
+            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
+                    )
+
+                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,
+        *,
+        provider: str,
+        model: str,
+        input: PromptTemplate,
+        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] = {},
+        prompt_name: Union[str, None] = None,
+        prompt_version_number: Union[int, None] = None,
+        prompt_input_variables: Dict[str, Any] = {},
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        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,
+            output=output,
+            request_start_time=request_start_time,
+            request_end_time=request_end_time,
+            parameters=parameters,
+            tags=tags,
+            metadata=metadata,
+            prompt_name=prompt_name,
+            prompt_version_number=prompt_version_number,
+            prompt_input_variables=prompt_input_variables,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            price=price,
+            function_name=function_name,
+            score=score,
+            prompt_id=prompt_id,
+            score_name=score_name,
+        )
+
+
+class AsyncPromptLayer(PromptLayerMixin):
+    def __init__(
+        self,
+        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")
 
-        if not prompt_blueprint_metadata:
+        if api_key is None:
             raise ValueError(
-                f"Prompt '{prompt_name}' does not have any metadata associated with it."
+                "PromptLayer API key not provided. "
+                "Please set the PROMPTLAYER_API_KEY environment variable or pass the api_key parameter."
             )
 
-        prompt_blueprint_model = prompt_blueprint_metadata.get("model")
+        self.base_url = get_base_url(base_url)
+        self.api_key = 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, self.base_url, self.throw_on_error, enable_tracing
+        )
+        self.track = AsyncTrackManager(api_key, self.base_url, self.throw_on_error)
 
-        if not prompt_blueprint_model:
-            raise ValueError(
-                f"Prompt '{prompt_name}' does not have a model parameters associated with it."
+    def __getattr__(self, name: Union[Literal["openai"], Literal["anthropic"], Literal["prompts"]]):
+        if name == "openai":
+            import openai as openai_module
+
+            openai = 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
 
-        return prompt_blueprint_model
+            anthropic = PromptLayerBase(
+                self.api_key,
+                self.base_url,
+                anthropic_module,
+                function_name="anthropic",
+                provider_type="anthropic",
+                tracer=self.tracer,
+            )
+            return anthropic
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
 
-    def run(
+    async def run_workflow(
+        self,
+        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
+        return_all_outputs: Optional[bool] = False,
+        # `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:
+            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,
+                return_all_outputs=return_all_outputs,
+            )
+        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,
         prompt_name: str,
         prompt_version: Union[int, None] = None,
         prompt_release_label: Union[str, None] = None,
         input_variables: Union[Dict[str, Any], None] = None,
+        model_parameter_overrides: Union[Dict[str, Any], None] = None,
         tags: Union[List[str], None] = None,
         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,
+            "model_parameter_overrides": model_parameter_overrides,
             "tags": tags,
             "metadata": metadata,
             "group_id": group_id,
             "stream": stream,
+            "provider": provider,
+            "model": model,
         }
 
         if self.tracer:
@@ -347,52 +514,197 @@ class PromptLayer:
                 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 = 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:
-            return self._run_internal(**_run_internal_kwargs)
+            return await self._run_internal(**_run_internal_kwargs)
+
+    async def log_request(
+        self,
+        *,
+        provider: str,
+        model: str,
+        input: PromptTemplate,
+        output: PromptTemplate,
+        request_start_time: float,
+        request_end_time: float,
+        parameters: Dict[str, Any] = {},
+        tags: List[str] = [],
+        metadata: Dict[str, str] = {},
+        prompt_name: Union[str, None] = None,
+        prompt_version_number: Union[int, None] = None,
+        prompt_input_variables: Dict[str, Any] = {},
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        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,
+            output=output,
+            request_start_time=request_start_time,
+            request_end_time=request_end_time,
+            parameters=parameters,
+            tags=tags,
+            metadata=metadata,
+            prompt_name=prompt_name,
+            prompt_version_number=prompt_version_number,
+            prompt_input_variables=prompt_input_variables,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            price=price,
+            function_name=function_name,
+            score=score,
+            prompt_id=prompt_id,
+        )
 
-    def traceable(self, attributes=None):
-        def decorator(func):
-            @wraps(func)
-            def sync_wrapper(*args, **kwargs):
-                if self.tracer:
-                    with self.tracer.start_as_current_span(func.__name__) as span:
-                        if attributes:
-                            for key, value in attributes.items():
-                                span.set_attribute(key, value)
-
-                        span.set_attribute(
-                            "function_input", str({"args": args, "kwargs": kwargs})
-                        )
-                        result = func(*args, **kwargs)
-                        span.set_attribute("function_output", str(result))
-
-                        return result
-                else:
-                    return func(*args, **kwargs)
-
-            @wraps(func)
-            async def async_wrapper(*args, **kwargs):
-                if self.tracer:
-                    with self.tracer.start_as_current_span(func.__name__) as span:
-                        if attributes:
-                            for key, value in attributes.items():
-                                span.set_attribute(key, value)
-
-                        span.set_attribute(
-                            "function_input", str({"args": args, "kwargs": kwargs})
-                        )
-                        result = await func(*args, **kwargs)
-                        span.set_attribute("function_output", str(result))
-
-                        return result
-                else:
-                    return await func(*args, **kwargs)
-
-            return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
-
-        return decorator
+    async def _create_track_request_callable(
+        self,
+        *,
+        request_params,
+        tags,
+        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(
+                self.api_key,
+                request_params,
+                tags,
+                input_variables,
+                group_id,
+                pl_run_span_id,
+                request_start_time=request_start_time,
+                **body,
+            )
+            return await atrack_request(self.base_url, self.throw_on_error, **track_request_kwargs)
+
+        return _track_request
+
+    async def _track_request_log(
+        self,
+        request_params,
+        tags,
+        input_variables,
+        group_id,
+        pl_run_span_id: Union[str, None] = None,
+        metadata: Union[Dict[str, str], None] = None,
+        **body,
+    ):
+        track_request_kwargs = self._prepare_track_request_kwargs(
+            self.api_key,
+            request_params,
+            tags,
+            input_variables,
+            group_id,
+            pl_run_span_id,
+            metadata=metadata,
+            **body,
+        )
+        return await atrack_request(self.base_url, self.throw_on_error, **track_request_kwargs)
+
+    async def _run_internal(
+        self,
+        *,
+        prompt_name: str,
+        prompt_version: Union[int, None] = None,
+        prompt_release_label: Union[str, None] = None,
+        input_variables: Union[Dict[str, Any], None] = None,
+        model_parameter_overrides: Union[Dict[str, Any], None] = None,
+        tags: Union[List[str], None] = None,
+        metadata: Union[Dict[str, str], None] = None,
+        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)
+        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_data = self._prepare_llm_data(
+            prompt_blueprint=prompt_blueprint,
+            prompt_template=prompt_blueprint["prompt_template"],
+            prompt_blueprint_model=prompt_blueprint_model,
+            stream=stream,
+            is_async=True,
+        )
+
+        # 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_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(
+                request_response,
+                track_request_callable,
+                llm_data["stream_function"],
+                llm_data["prompt_blueprint"]["metadata"],
+            )
+
+        request_log = await self._track_request_log(
+            llm_data,
+            tags,
+            input_variables,
+            group_id,
+            pl_run_span_id,
+            metadata=metadata,
+            request_response=request_response,
+            request_start_time=request_start_time,
+            request_end_time=request_end_time,
+        )
+
+        return {
+            "request_id": request_log.get("request_id", None),
+            "raw_response": response,
+            "prompt_blueprint": request_log.get("prompt_blueprint", None),
+        }