Source code for langgraph_agent_toolkit.core.observability.base

import asyncio
import os
from abc import ABC, abstractmethod
from contextlib import contextmanager
from functools import wraps
from typing import Any, Callable, Dict, List, Literal, Optional, TypeVar, Union, cast

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.prompts.chat import (
    AIMessagePromptTemplate,
    BaseMessage,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
)

from langgraph_agent_toolkit.core.observability.types import ChatMessageDict, MessageRole


T = TypeVar("T")

# Prompt templates accepted by `push_prompt`.
PromptTemplateType = Union[str, List[ChatMessageDict]]

# Return type for `pull_prompt`.
PromptReturnType = Union[ChatPromptTemplate, str, dict, None]




class BaseObservabilityPlatform(ABC):
    """Base class for observability platforms.

    Provides shared utilities without disk I/O.
    This prevents blocking in asynchronous contexts.
    """



    def __init__(self, remote_first: bool = False):
        """Initialize the observability platform.

        Args:
            remote_first: Prioritize remote prompts over the local cache.

        """
        self._required_vars: List[str] = []
        self._remote_first = remote_first


    @property
    def remote_first(self) -> bool:
        return self._remote_first

    @property
    def required_vars(self) -> List[str]:
        return self._required_vars

    @required_vars.setter
    def required_vars(self, value: List[str]) -> None:
        self._required_vars = value



    def validate_environment(self) -> bool:
        """Validate that required environment variables are set."""
        missing_vars = [var for var in self._required_vars if not os.environ.get(var)]
        if missing_vars:
            raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}")
        return True




    @staticmethod
    def requires_env_vars(func: Callable[..., T]) -> Callable[..., T]:
        """Validate environment variables before calling a method."""

        @wraps(func)
        def wrapper(self, *args, **kwargs):
            self.validate_environment()
            return func(self, *args, **kwargs)

        return wrapper




    @abstractmethod
    def get_callback_handler(self, **kwargs) -> Any:
        """Get the callback handler for the observability platform."""
        pass




    @abstractmethod
    def before_shutdown(self) -> None:
        """Perform any necessary cleanup before shutdown."""
        pass




    @abstractmethod
    def record_feedback(self, run_id: str, key: str, score: float, **kwargs) -> None:
        """Record feedback for a run."""
        pass




    @abstractmethod
    def push_prompt(
        self,
        name: str,
        prompt_template: PromptTemplateType,
        metadata: Optional[Dict[str, Any]] = None,
        force_create_new_version: bool = True,
    ) -> None:
        """Push a prompt to the observability platform."""
        pass




    @abstractmethod
    def pull_prompt(
        self,
        name: str,
        template_format: Literal["f-string", "mustache", "jinja2"] = "f-string",
        **kwargs,
    ) -> PromptReturnType:
        """Pull a prompt from the observability platform."""
        pass




    async def apull_prompt(
        self,
        name: str,
        template_format: Literal["f-string", "mustache", "jinja2"] = "f-string",
        **kwargs,
    ) -> PromptReturnType:
        """Asynchronously run `pull_prompt` in a thread pool."""
        return await asyncio.to_thread(self.pull_prompt, name, template_format=template_format, **kwargs)




    @abstractmethod
    def delete_prompt(self, name: str) -> None:
        """Delete a prompt from the observability platform."""
        pass




    @contextmanager
    def trace_context(self, run_id: str, **kwargs):
        """Create an execution trace context.

        Subclasses provide platform-specific implementations.

        Args:
            run_id: Run ID to use as trace ID.
            **kwargs: Context parameters, such as `user_id` and `input`.

        Yields:
            `None` or a platform-specific context object.

        """
        yield




    def update_trace(self, trace: Any, **attributes: Any) -> None:
        """Set attributes on a trace from `trace_context`.

        Call this after the run to record output.
        Platforms without a trace object do nothing.
        """
        pass


    def _convert_to_chat_prompt(self, prompt_template: PromptTemplateType) -> ChatPromptTemplate:
        if isinstance(prompt_template, str):
            return ChatPromptTemplate.from_template(prompt_template)
        elif isinstance(prompt_template, list) and all(isinstance(msg, dict) for msg in prompt_template):
            messages = []
            for msg in prompt_template:
                role = msg.get("role", "")
                content = msg.get("content", "")

                match role.lower():
                    case MessageRole.SYSTEM:
                        messages.append(SystemMessage(content=content))
                    case MessageRole.HUMAN | MessageRole.USER:
                        messages.append(HumanMessage(content=content))
                    case MessageRole.AI | MessageRole.ASSISTANT:
                        messages.append(AIMessage(content=content))
                    case MessageRole.PLACEHOLDER | MessageRole.MESSAGES_PLACEHOLDER:
                        messages.append(MessagesPlaceholder(variable_name=content))
                    case _:
                        raise ValueError(f"Unknown message role: {role}")

            return ChatPromptTemplate.from_messages(messages)
        else:
            return cast(ChatPromptTemplate, prompt_template)

    def _process_messages_from_prompt(
        self,
        messages: List[Any],
        template_format: Literal["f-string", "mustache", "jinja2"] = "f-string",
    ) -> List[Any]:
        MESSAGE_TYPE_MAP = {
            MessageRole.SYSTEM: SystemMessagePromptTemplate,
            MessageRole.HUMAN: HumanMessagePromptTemplate,
            MessageRole.USER: HumanMessagePromptTemplate,
            MessageRole.AI: AIMessagePromptTemplate,
            MessageRole.ASSISTANT: AIMessagePromptTemplate,
        }

        processed_messages = []

        for msg in messages:
            if isinstance(msg, MessagesPlaceholder):
                processed_messages.append(msg)
                continue

            if isinstance(msg, BaseMessage):
                msg_type = MessageRole(msg.type)
                if msg_type in MESSAGE_TYPE_MAP:
                    template_class = MESSAGE_TYPE_MAP[msg_type]
                    processed_messages.append(
                        template_class.from_template(msg.content, template_format=template_format)
                    )
                continue

            if isinstance(msg, dict) and "role" in msg and "content" in msg:
                role, content = msg["role"], msg["content"]

                if role in MESSAGE_TYPE_MAP:
                    template_class = MESSAGE_TYPE_MAP[role]
                    processed_messages.append(template_class.from_template(content, template_format=template_format))
                    continue

                if role.lower() in (MessageRole.PLACEHOLDER, MessageRole.MESSAGES_PLACEHOLDER):
                    processed_messages.append(MessagesPlaceholder(variable_name=content))
                    continue

            if isinstance(msg, tuple) and len(msg) == 2:
                role, content = msg
                if role in MESSAGE_TYPE_MAP:
                    template_class = MESSAGE_TYPE_MAP[role]
                    processed_messages.append(template_class.from_template(content, template_format=template_format))
                    continue

                if role.lower() in (MessageRole.PLACEHOLDER, MessageRole.MESSAGES_PLACEHOLDER):
                    processed_messages.append(MessagesPlaceholder(variable_name=content))
                    continue

            processed_messages.append(msg)

        return processed_messages

    def _process_prompt_object(
        self,
        prompt_obj: Any,
        template_format: Literal["f-string", "mustache", "jinja2"] = "f-string",
    ) -> ChatPromptTemplate:
        if isinstance(prompt_obj, ChatPromptTemplate):
            return prompt_obj

        if hasattr(prompt_obj, "messages") and isinstance(prompt_obj.messages, list):
            processed_messages = self._process_messages_from_prompt(
                prompt_obj.messages, template_format=template_format
            )
            if processed_messages:
                return ChatPromptTemplate.from_messages(processed_messages)

        elif isinstance(prompt_obj, list):
            if all(isinstance(item, dict) and "role" in item and "content" in item for item in prompt_obj):
                processed_messages = self._process_messages_from_prompt(prompt_obj, template_format=template_format)
                if processed_messages:
                    return ChatPromptTemplate.from_messages(processed_messages)

        if isinstance(prompt_obj, str):
            return ChatPromptTemplate.from_template(prompt_obj, template_format=template_format)

        raise ValueError(f"Could not process prompt object of type {type(prompt_obj)}")