Source code for langchain.callbacks.sagemaker_callback
import json
import os
import shutil
import tempfile
from copy import deepcopy
from typing import Any, Dict, List, Optional
from langchain.callbacks.base import BaseCallbackHandler
from langchain.callbacks.utils import (
    flatten_dict,
)
from langchain.schema import AgentAction, AgentFinish, LLMResult
[docs]def save_json(data: dict, file_path: str) -> None:
    """Save dict to local file path.
    Parameters:
        data (dict): The dictionary to be saved.
        file_path (str): Local file path.
    """
    with open(file_path, "w") as outfile:
        json.dump(data, outfile)
[docs]class SageMakerCallbackHandler(BaseCallbackHandler):
    """Callback Handler that logs prompt artifacts and metrics to SageMaker Experiments.
    Parameters:
        run (sagemaker.experiments.run.Run): Run object where the experiment is logged.
    """
[docs]    def __init__(self, run: Any) -> None:
        """Initialize callback handler."""
        super().__init__()
        self.run = run
        self.metrics = {
            "step": 0,
            "starts": 0,
            "ends": 0,
            "errors": 0,
            "text_ctr": 0,
            "chain_starts": 0,
            "chain_ends": 0,
            "llm_starts": 0,
            "llm_ends": 0,
            "llm_streams": 0,
            "tool_starts": 0,
            "tool_ends": 0,
            "agent_ends": 0,
        }
        # Create a temporary directory
        self.temp_dir = tempfile.mkdtemp()