Source code for langchain.callbacks.promptlayer_callback
"""Callback handler for promptlayer."""
from __future__ import annotations
import datetime
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
from uuid import UUID
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import (
    ChatGeneration,
    LLMResult,
)
from langchain.schema.messages import (
    AIMessage,
    BaseMessage,
    ChatMessage,
    HumanMessage,
    SystemMessage,
)
if TYPE_CHECKING:
    import promptlayer
def _lazy_import_promptlayer() -> promptlayer:
    """Lazy import promptlayer to avoid circular imports."""
    try:
        import promptlayer
    except ImportError:
        raise ImportError(
            "The PromptLayerCallbackHandler requires the promptlayer package. "
            " Please install it with `pip install promptlayer`."
        )
    return promptlayer
[docs]class PromptLayerCallbackHandler(BaseCallbackHandler):
    """Callback handler for promptlayer."""
[docs]    def __init__(
        self,
        pl_id_callback: Optional[Callable[..., Any]] = None,
        pl_tags: Optional[List[str]] = None,
    ) -> None:
        """Initialize the PromptLayerCallbackHandler."""
        _lazy_import_promptlayer()
        self.pl_id_callback = pl_id_callback
        self.pl_tags = pl_tags or []
        self.runs: Dict[UUID, Dict[str, Any]] = {}
[docs]    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,