... )
        >>> chat = ChatOpenAI(
        ...     temperature=0,
        ...     headers={"user_id": "123"},
        ...     callbacks=[context_callback],
        ...     openai_api_key="API_KEY_HERE",
        ... )
        >>> messages = [
        ...     SystemMessage(content="You translate English to French."),
        ...     HumanMessage(content="I love programming with LangChain."),
        ... ]
        >>> chat(messages)
    Chain Example:
        >>> from langchain.chains import LLMChain
        >>> from langchain_community.chat_models import ChatOpenAI
        >>> from langchain_community.callbacks import ContextCallbackHandler
        >>> context_callback = ContextCallbackHandler(
        ...     token="<CONTEXT_TOKEN_HERE>",
        ... )
        >>> human_message_prompt = HumanMessagePromptTemplate(
        ...     prompt=PromptTemplate(
        ...         template="What is a good name for a company that makes {product}?",
        ...         input_variables=["product"],
        ...    ),
        ... )
        >>> chat_prompt_template = ChatPromptTemplate.from_messages(
        ...   [human_message_prompt]
        ... )
        >>> callback = ContextCallbackHandler(token)
        >>> # Note: the same callback object must be shared between the
        ...   LLM and the chain.
        >>> chat = ChatOpenAI(temperature=0.9, callbacks=[callback])
        >>> chain = LLMChain(
        ...   llm=chat,
        ...   prompt=chat_prompt_template,
        ...   callbacks=[callback]
        ... )
        >>> chain.run("colorful socks")
    """
[docs]    def __init__(self, token: str = "", verbose: bool = False, **kwargs: Any) -> None:
        (
            self.context,


(
            self.context,
            self.credential,
            self.conversation_model,
            self.message_model,
            self.message_role_model,
            self.rating_model,
        ) = import_context()
        token = token or os.environ.get("CONTEXT_TOKEN") or ""
        self.client = self.context.ContextAPI(credential=self.credential(token))
        self.chain_run_id = None
        self.llm_model = None
        self.messages: List[Any] = []
        self.metadata: Dict[str, str] = {}
[docs]    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> Any:
        """Run when the chat model is started."""
        llm_model = kwargs.get("invocation_params", {}).get("model", None)
        if llm_model is not None:
            self.metadata["model"] = llm_model
        if len(messages) == 0:
            return
        for message in messages[0]:
            role = self.message_role_model.SYSTEM
            if message.type == "human":
                role = self.message_role_model.USER
            elif message.type == "system":
                role = self.message_role_model.SYSTEM
            elif message.type == "ai":
                role = self.message_role_model.ASSISTANT
            self.messages.append(
                self.message_model(
                    message=message.content,
                    role=role,
                )
            )
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Run when LLM ends."""


"""Run when LLM ends."""
        if len(response.generations) == 0 or len(response.generations[0]) == 0:
            return
        if not self.chain_run_id:
            generation = response.generations[0][0]
            self.messages.append(
                self.message_model(
                    message=generation.text,
                    role=self.message_role_model.ASSISTANT,
                )
            )
            self._log_conversation()
[docs]    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Run when chain starts."""
        self.chain_run_id = kwargs.get("run_id", None)
[docs]    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """Run when chain ends."""
        self.messages.append(
            self.message_model(
                message=outputs["text"],
                role=self.message_role_model.ASSISTANT,
            )
        )
        self._log_conversation()
        self.chain_run_id = None
    def _log_conversation(self) -> None:
        """Log the conversation to the context API."""
        if len(self.messages) == 0:
            return
        self.client.log.conversation_upsert(
            body={
                "conversation": self.conversation_model(
                    messages=self.messages,
                    metadata=self.metadata,
                )
            }
        )
        self.messages = []
        self.metadata = {}


Source code for langchain_community.callbacks.mlflow_callback
import os
import random
import string
import tempfile
import traceback
from copy import deepcopy
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_core.utils import get_from_dict_or_env
from langchain_community.callbacks.utils import (
    BaseMetadataCallbackHandler,
    flatten_dict,
    hash_string,
    import_pandas,
    import_spacy,
    import_textstat,
)
[docs]def import_mlflow() -> Any:
    """Import the mlflow python package and raise an error if it is not installed."""
    try:
        import mlflow
    except ImportError:
        raise ImportError(
            "To use the mlflow callback manager you need to have the `mlflow` python "
            "package installed. Please install it with `pip install mlflow>=2.3.0`"
        )
    return mlflow
[docs]def analyze_text(
    text: str,
    nlp: Any = None,
) -> dict:
    """Analyze text using textstat and spacy.
    Parameters:
        text (str): The text to analyze.
        nlp (spacy.lang): The spacy language model to use for visualization.
    Returns:
        (dict): A dictionary containing the complexity metrics and visualization
            files serialized to  HTML string.
    """
    resp: Dict[str, Any] = {}
    textstat = import_textstat()
    spacy = import_spacy()
    text_complexity_metrics = {


spacy = import_spacy()
    text_complexity_metrics = {
        "flesch_reading_ease": textstat.flesch_reading_ease(text),
        "flesch_kincaid_grade": textstat.flesch_kincaid_grade(text),
        "smog_index": textstat.smog_index(text),
        "coleman_liau_index": textstat.coleman_liau_index(text),
        "automated_readability_index": textstat.automated_readability_index(text),
        "dale_chall_readability_score": textstat.dale_chall_readability_score(text),
        "difficult_words": textstat.difficult_words(text),
        "linsear_write_formula": textstat.linsear_write_formula(text),
        "gunning_fog": textstat.gunning_fog(text),
        # "text_standard": textstat.text_standard(text),
        "fernandez_huerta": textstat.fernandez_huerta(text),
        "szigriszt_pazos": textstat.szigriszt_pazos(text),
        "gutierrez_polini": textstat.gutierrez_polini(text),
        "crawford": textstat.crawford(text),
        "gulpease_index": textstat.gulpease_index(text),
        "osman": textstat.osman(text),
    }
    resp.update({"text_complexity_metrics": text_complexity_metrics})
    resp.update(text_complexity_metrics)
    if nlp is not None:
        doc = nlp(text)
        dep_out = spacy.displacy.render(  # type: ignore
            doc, style="dep", jupyter=False, page=True
        )
        ent_out = spacy.displacy.render(  # type: ignore


)
        ent_out = spacy.displacy.render(  # type: ignore
            doc, style="ent", jupyter=False, page=True
        )
        text_visualizations = {
            "dependency_tree": dep_out,
            "entities": ent_out,
        }
        resp.update(text_visualizations)
    return resp
[docs]def construct_html_from_prompt_and_generation(prompt: str, generation: str) -> Any:
    """Construct an html element from a prompt and a generation.
    Parameters:
        prompt (str): The prompt.
        generation (str): The generation.
    Returns:
        (str): The html string."""
    formatted_prompt = prompt.replace("\n", "
")
    formatted_generation = generation.replace("\n", "
")
    return f"""

{formatted_prompt}:

        {formatted_generation}

    """
[docs]class MlflowLogger:
    """Callback Handler that logs metrics and artifacts to mlflow server.
    Parameters:
        name (str): Name of the run.
        experiment (str): Name of the experiment.
        tags (dict): Tags to be attached for the run.
        tracking_uri (str): MLflow tracking server uri.
    This handler implements the helper functions to initialize,
    log metrics and artifacts to the mlflow server.
    """
[docs]    def __init__(self, **kwargs: Any):
        self.mlflow = import_mlflow()
        if "DATABRICKS_RUNTIME_VERSION" in os.environ:
            self.mlflow.set_tracking_uri("databricks")


self.mlflow.set_tracking_uri("databricks")
            self.mlf_expid = self.mlflow.tracking.fluent._get_experiment_id()
            self.mlf_exp = self.mlflow.get_experiment(self.mlf_expid)
        else:
            tracking_uri = get_from_dict_or_env(
                kwargs, "tracking_uri", "MLFLOW_TRACKING_URI", ""
            )
            self.mlflow.set_tracking_uri(tracking_uri)
            # User can set other env variables described here
            # > https://www.mlflow.org/docs/latest/tracking.html#logging-to-a-tracking-server
            experiment_name = get_from_dict_or_env(
                kwargs, "experiment_name", "MLFLOW_EXPERIMENT_NAME"
            )
            self.mlf_exp = self.mlflow.get_experiment_by_name(experiment_name)
            if self.mlf_exp is not None:
                self.mlf_expid = self.mlf_exp.experiment_id
            else:
                self.mlf_expid = self.mlflow.create_experiment(experiment_name)
        self.start_run(kwargs["run_name"], kwargs["run_tags"])
[docs]    def start_run(self, name: str, tags: Dict[str, str]) -> None:
        """To start a new run, auto generates the random suffix for name"""
        if name.endswith("-%"):
            rname = "".join(random.choices(string.ascii_uppercase + string.digits, k=7))
            name = name.replace("%", rname)
        self.run = self.mlflow.MlflowClient().create_run(
            self.mlf_expid, run_name=name, tags=tags
        )
[docs]    def finish_run(self) -> None:
        """To finish the run."""
        with self.mlflow.start_run(


"""To finish the run."""
        with self.mlflow.start_run(
            run_id=self.run.info.run_id, experiment_id=self.mlf_expid
        ):
            self.mlflow.end_run()
[docs]    def metric(self, key: str, value: float) -> None:
        """To log metric to mlflow server."""
        with self.mlflow.start_run(
            run_id=self.run.info.run_id, experiment_id=self.mlf_expid
        ):
            self.mlflow.log_metric(key, value)
[docs]    def metrics(
        self, data: Union[Dict[str, float], Dict[str, int]], step: Optional[int] = 0
    ) -> None:
        """To log all metrics in the input dict."""
        with self.mlflow.start_run(
            run_id=self.run.info.run_id, experiment_id=self.mlf_expid
        ):
            self.mlflow.log_metrics(data)
[docs]    def jsonf(self, data: Dict[str, Any], filename: str) -> None:
        """To log the input data as json file artifact."""
        with self.mlflow.start_run(
            run_id=self.run.info.run_id, experiment_id=self.mlf_expid
        ):
            self.mlflow.log_dict(data, f"{filename}.json")
[docs]    def table(self, name: str, dataframe) -> None:  # type: ignore
        """To log the input pandas dataframe as a html table"""
        self.html(dataframe.to_html(), f"table_{name}")
[docs]    def html(self, html: str, filename: str) -> None:
        """To log the input html string as html file artifact."""
        with self.mlflow.start_run(
            run_id=self.run.info.run_id, experiment_id=self.mlf_expid


run_id=self.run.info.run_id, experiment_id=self.mlf_expid
        ):
            self.mlflow.log_text(html, f"{filename}.html")
[docs]    def text(self, text: str, filename: str) -> None:
        """To log the input text as text file artifact."""
        with self.mlflow.start_run(
            run_id=self.run.info.run_id, experiment_id=self.mlf_expid
        ):
            self.mlflow.log_text(text, f"{filename}.txt")
[docs]    def artifact(self, path: str) -> None:
        """To upload the file from given path as artifact."""
        with self.mlflow.start_run(
            run_id=self.run.info.run_id, experiment_id=self.mlf_expid
        ):
            self.mlflow.log_artifact(path)
[docs]    def langchain_artifact(self, chain: Any) -> None:
        with self.mlflow.start_run(
            run_id=self.run.info.run_id, experiment_id=self.mlf_expid
        ):
            self.mlflow.langchain.log_model(chain, "langchain-model")
[docs]class MlflowCallbackHandler(BaseMetadataCallbackHandler, BaseCallbackHandler):
    """Callback Handler that logs metrics and artifacts to mlflow server.
    Parameters:
        name (str): Name of the run.
        experiment (str): Name of the experiment.
        tags (dict): Tags to be attached for the run.
        tracking_uri (str): MLflow tracking server uri.
    This handler will utilize the associated callback method called and formats
    the input of each callback function with metadata regarding the state of LLM run,
    and adds the response to the list of records for both the {method}_records and
    action. It then logs the response to mlflow server.
    """


action. It then logs the response to mlflow server.
    """
[docs]    def __init__(
        self,
        name: Optional[str] = "langchainrun-%",
        experiment: Optional[str] = "langchain",
        tags: Optional[Dict] = None,
        tracking_uri: Optional[str] = None,
    ) -> None:
        """Initialize callback handler."""
        import_pandas()
        import_textstat()
        import_mlflow()
        spacy = import_spacy()
        super().__init__()
        self.name = name
        self.experiment = experiment
        self.tags = tags or {}
        self.tracking_uri = tracking_uri
        self.temp_dir = tempfile.TemporaryDirectory()
        self.mlflg = MlflowLogger(
            tracking_uri=self.tracking_uri,
            experiment_name=self.experiment,
            run_name=self.name,
            run_tags=self.tags,
        )
        self.action_records: list = []
        self.nlp = spacy.load("en_core_web_sm")
        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,
        }
        self.records: Dict[str, Any] = {
            "on_llm_start_records": [],
            "on_llm_token_records": [],


"on_llm_start_records": [],
            "on_llm_token_records": [],
            "on_llm_end_records": [],
            "on_chain_start_records": [],
            "on_chain_end_records": [],
            "on_tool_start_records": [],
            "on_tool_end_records": [],
            "on_text_records": [],
            "on_agent_finish_records": [],
            "on_agent_action_records": [],
            "action_records": [],
        }
    def _reset(self) -> None:
        for k, v in self.metrics.items():
            self.metrics[k] = 0
        for k, v in self.records.items():
            self.records[k] = []
[docs]    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Run when LLM starts."""
        self.metrics["step"] += 1
        self.metrics["llm_starts"] += 1
        self.metrics["starts"] += 1
        llm_starts = self.metrics["llm_starts"]
        resp: Dict[str, Any] = {}
        resp.update({"action": "on_llm_start"})
        resp.update(flatten_dict(serialized))
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        for idx, prompt in enumerate(prompts):
            prompt_resp = deepcopy(resp)
            prompt_resp["prompt"] = prompt
            self.records["on_llm_start_records"].append(prompt_resp)
            self.records["action_records"].append(prompt_resp)
            self.mlflg.jsonf(prompt_resp, f"llm_start_{llm_starts}_prompt_{idx}")


[docs]    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Run when LLM generates a new token."""
        self.metrics["step"] += 1
        self.metrics["llm_streams"] += 1
        llm_streams = self.metrics["llm_streams"]
        resp: Dict[str, Any] = {}
        resp.update({"action": "on_llm_new_token", "token": token})
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_llm_token_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"llm_new_tokens_{llm_streams}")
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Run when LLM ends running."""
        self.metrics["step"] += 1
        self.metrics["llm_ends"] += 1
        self.metrics["ends"] += 1
        llm_ends = self.metrics["llm_ends"]
        resp: Dict[str, Any] = {}
        resp.update({"action": "on_llm_end"})
        resp.update(flatten_dict(response.llm_output or {}))
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        for generations in response.generations:
            for idx, generation in enumerate(generations):
                generation_resp = deepcopy(resp)
                generation_resp.update(flatten_dict(generation.dict()))
                generation_resp.update(
                    analyze_text(
                        generation.text,
                        nlp=self.nlp,
                    )
                )
                complexity_metrics: Dict[str, float] = generation_resp.pop(


)
                complexity_metrics: Dict[str, float] = generation_resp.pop(
                    "text_complexity_metrics"
                )  # type: ignore  # noqa: E501
                self.mlflg.metrics(
                    complexity_metrics,
                    step=self.metrics["step"],
                )
                self.records["on_llm_end_records"].append(generation_resp)
                self.records["action_records"].append(generation_resp)
                self.mlflg.jsonf(resp, f"llm_end_{llm_ends}_generation_{idx}")
                dependency_tree = generation_resp["dependency_tree"]
                entities = generation_resp["entities"]
                self.mlflg.html(dependency_tree, "dep-" + hash_string(generation.text))
                self.mlflg.html(entities, "ent-" + hash_string(generation.text))
[docs]    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when LLM errors."""
        self.metrics["step"] += 1
        self.metrics["errors"] += 1
[docs]    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Run when chain starts running."""
        self.metrics["step"] += 1
        self.metrics["chain_starts"] += 1
        self.metrics["starts"] += 1
        chain_starts = self.metrics["chain_starts"]
        resp: Dict[str, Any] = {}
        resp.update({"action": "on_chain_start"})
        resp.update(flatten_dict(serialized))
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])


self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        chain_input = ",".join([f"{k}={v}" for k, v in inputs.items()])
        input_resp = deepcopy(resp)
        input_resp["inputs"] = chain_input
        self.records["on_chain_start_records"].append(input_resp)
        self.records["action_records"].append(input_resp)
        self.mlflg.jsonf(input_resp, f"chain_start_{chain_starts}")
[docs]    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """Run when chain ends running."""
        self.metrics["step"] += 1
        self.metrics["chain_ends"] += 1
        self.metrics["ends"] += 1
        chain_ends = self.metrics["chain_ends"]
        resp: Dict[str, Any] = {}
        chain_output = ",".join([f"{k}={v}" for k, v in outputs.items()])
        resp.update({"action": "on_chain_end", "outputs": chain_output})
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_chain_end_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"chain_end_{chain_ends}")
[docs]    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when chain errors."""
        self.metrics["step"] += 1
        self.metrics["errors"] += 1
[docs]    def on_tool_start(
        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
    ) -> None:
        """Run when tool starts running."""


) -> None:
        """Run when tool starts running."""
        self.metrics["step"] += 1
        self.metrics["tool_starts"] += 1
        self.metrics["starts"] += 1
        tool_starts = self.metrics["tool_starts"]
        resp: Dict[str, Any] = {}
        resp.update({"action": "on_tool_start", "input_str": input_str})
        resp.update(flatten_dict(serialized))
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_tool_start_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"tool_start_{tool_starts}")
[docs]    def on_tool_end(self, output: str, **kwargs: Any) -> None:
        """Run when tool ends running."""
        self.metrics["step"] += 1
        self.metrics["tool_ends"] += 1
        self.metrics["ends"] += 1
        tool_ends = self.metrics["tool_ends"]
        resp: Dict[str, Any] = {}
        resp.update({"action": "on_tool_end", "output": output})
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_tool_end_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"tool_end_{tool_ends}")
[docs]    def on_tool_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when tool errors."""
        self.metrics["step"] += 1
        self.metrics["errors"] += 1


self.metrics["step"] += 1
        self.metrics["errors"] += 1
[docs]    def on_text(self, text: str, **kwargs: Any) -> None:
        """
        Run when agent is ending.
        """
        self.metrics["step"] += 1
        self.metrics["text_ctr"] += 1
        text_ctr = self.metrics["text_ctr"]
        resp: Dict[str, Any] = {}
        resp.update({"action": "on_text", "text": text})
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_text_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"on_text_{text_ctr}")
[docs]    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
        """Run when agent ends running."""
        self.metrics["step"] += 1
        self.metrics["agent_ends"] += 1
        self.metrics["ends"] += 1
        agent_ends = self.metrics["agent_ends"]
        resp: Dict[str, Any] = {}
        resp.update(
            {
                "action": "on_agent_finish",
                "output": finish.return_values["output"],
                "log": finish.log,
            }
        )
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_agent_finish_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"agent_finish_{agent_ends}")
[docs]    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:


"""Run on agent action."""
        self.metrics["step"] += 1
        self.metrics["tool_starts"] += 1
        self.metrics["starts"] += 1
        tool_starts = self.metrics["tool_starts"]
        resp: Dict[str, Any] = {}
        resp.update(
            {
                "action": "on_agent_action",
                "tool": action.tool,
                "tool_input": action.tool_input,
                "log": action.log,
            }
        )
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_agent_action_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"agent_action_{tool_starts}")
    def _create_session_analysis_df(self) -> Any:
        """Create a dataframe with all the information from the session."""
        pd = import_pandas()
        on_llm_start_records_df = pd.DataFrame(self.records["on_llm_start_records"])
        on_llm_end_records_df = pd.DataFrame(self.records["on_llm_end_records"])
        llm_input_columns = ["step", "prompt"]
        if "name" in on_llm_start_records_df.columns:
            llm_input_columns.append("name")
        elif "id" in on_llm_start_records_df.columns:
            # id is llm class's full import path. For example:
            # ["langchain", "llms", "openai", "AzureOpenAI"]
            on_llm_start_records_df["name"] = on_llm_start_records_df["id"].apply(
                lambda id_: id_[-1]
            )
            llm_input_columns.append("name")
        llm_input_prompts_df = (


llm_input_columns.append("name")
        llm_input_prompts_df = (
            on_llm_start_records_df[llm_input_columns]
            .dropna(axis=1)
            .rename({"step": "prompt_step"}, axis=1)
        )
        complexity_metrics_columns = []
        visualizations_columns = []
        complexity_metrics_columns = [
            "flesch_reading_ease",
            "flesch_kincaid_grade",
            "smog_index",
            "coleman_liau_index",
            "automated_readability_index",
            "dale_chall_readability_score",
            "difficult_words",
            "linsear_write_formula",
            "gunning_fog",
            # "text_standard",
            "fernandez_huerta",
            "szigriszt_pazos",
            "gutierrez_polini",
            "crawford",
            "gulpease_index",
            "osman",
        ]
        visualizations_columns = ["dependency_tree", "entities"]
        llm_outputs_df = (
            on_llm_end_records_df[
                [
                    "step",
                    "text",
                    "token_usage_total_tokens",
                    "token_usage_prompt_tokens",
                    "token_usage_completion_tokens",
                ]
                + complexity_metrics_columns
                + visualizations_columns
            ]
            .dropna(axis=1)
            .rename({"step": "output_step", "text": "output"}, axis=1)
        )
        session_analysis_df = pd.concat([llm_input_prompts_df, llm_outputs_df], axis=1)
        session_analysis_df["chat_html"] = session_analysis_df[
            ["prompt", "output"]
        ].apply(


["prompt", "output"]
        ].apply(
            lambda row: construct_html_from_prompt_and_generation(
                row["prompt"], row["output"]
            ),
            axis=1,
        )
        return session_analysis_df
[docs]    def flush_tracker(self, langchain_asset: Any = None, finish: bool = False) -> None:
        pd = import_pandas()
        self.mlflg.table("action_records", pd.DataFrame(self.records["action_records"]))
        session_analysis_df = self._create_session_analysis_df()
        chat_html = session_analysis_df.pop("chat_html")
        chat_html = chat_html.replace("\n", "", regex=True)
        self.mlflg.table("session_analysis", pd.DataFrame(session_analysis_df))
        self.mlflg.html("".join(chat_html.tolist()), "chat_html")
        if langchain_asset:
            # To avoid circular import error
            # mlflow only supports LLMChain asset
            if "langchain.chains.llm.LLMChain" in str(type(langchain_asset)):
                self.mlflg.langchain_artifact(langchain_asset)
            else:
                langchain_asset_path = str(Path(self.temp_dir.name, "model.json"))
                try:
                    langchain_asset.save(langchain_asset_path)
                    self.mlflg.artifact(langchain_asset_path)
                except ValueError:
                    try:
                        langchain_asset.save_agent(langchain_asset_path)
                        self.mlflg.artifact(langchain_asset_path)
                    except AttributeError:
                        print("Could not save model.")
                        traceback.print_exc()
                        pass
                    except NotImplementedError:
                        print("Could not save model.")
                        traceback.print_exc()
                        pass
                except NotImplementedError:
                    print("Could not save model.")
                    traceback.print_exc()
                    pass


print("Could not save model.")
                    traceback.print_exc()
                    pass
        if finish:
            self.mlflg.finish_run()
            self._reset()


Source code for langchain_community.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_core.callbacks import BaseCallbackHandler
from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    ChatMessage,
    HumanMessage,
    SystemMessage,
)
from langchain_core.outputs import (
    ChatGeneration,
    LLMResult,
)
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,


tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Any:
        self.runs[run_id] = {
            "messages": [self._create_message_dicts(m)[0] for m in messages],
            "invocation_params": kwargs.get("invocation_params", {}),
            "name": ".".join(serialized["id"]),
            "request_start_time": datetime.datetime.now().timestamp(),
            "tags": tags,
        }
[docs]    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Any:
        self.runs[run_id] = {
            "prompts": prompts,
            "invocation_params": kwargs.get("invocation_params", {}),
            "name": ".".join(serialized["id"]),
            "request_start_time": datetime.datetime.now().timestamp(),
            "tags": tags,
        }
[docs]    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        from promptlayer.utils import get_api_key, promptlayer_api_request
        run_info = self.runs.get(run_id, {})
        if not run_info:
            return
        run_info["request_end_time"] = datetime.datetime.now().timestamp()
        for i in range(len(response.generations)):
            generation = response.generations[i][0]
            resp = {


generation = response.generations[i][0]
            resp = {
                "text": generation.text,
                "llm_output": response.llm_output,
            }
            model_params = run_info.get("invocation_params", {})
            is_chat_model = run_info.get("messages", None) is not None
            model_input = (
                run_info.get("messages", [])[i]
                if is_chat_model
                else [run_info.get("prompts", [])[i]]
            )
            model_response = (
                [self._convert_message_to_dict(generation.message)]
                if is_chat_model and isinstance(generation, ChatGeneration)
                else resp
            )
            pl_request_id = promptlayer_api_request(
                run_info.get("name"),
                "langchain",
                model_input,
                model_params,
                self.pl_tags,
                model_response,
                run_info.get("request_start_time"),
                run_info.get("request_end_time"),
                get_api_key(),
                return_pl_id=bool(self.pl_id_callback is not None),
                metadata={
                    "_langchain_run_id": str(run_id),
                    "_langchain_parent_run_id": str(parent_run_id),
                    "_langchain_tags": str(run_info.get("tags", [])),
                },
            )
            if self.pl_id_callback:
                self.pl_id_callback(pl_request_id)
    def _convert_message_to_dict(self, message: BaseMessage) -> Dict[str, Any]:
        if isinstance(message, HumanMessage):
            message_dict = {"role": "user", "content": message.content}
        elif isinstance(message, AIMessage):
            message_dict = {"role": "assistant", "content": message.content}
        elif isinstance(message, SystemMessage):


elif isinstance(message, SystemMessage):
            message_dict = {"role": "system", "content": message.content}
        elif isinstance(message, ChatMessage):
            message_dict = {"role": message.role, "content": message.content}
        else:
            raise ValueError(f"Got unknown type {message}")
        if "name" in message.additional_kwargs:
            message_dict["name"] = message.additional_kwargs["name"]
        return message_dict
    def _create_message_dicts(
        self, messages: List[BaseMessage]
    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
        params: Dict[str, Any] = {}
        message_dicts = [self._convert_message_to_dict(m) for m in messages]
        return message_dicts, params


Source code for langchain_community.callbacks.argilla_callback
import os
import warnings
from typing import Any, Dict, List, Optional
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from packaging.version import parse
[docs]class ArgillaCallbackHandler(BaseCallbackHandler):
    """Callback Handler that logs into Argilla.
    Args:
        dataset_name: name of the `FeedbackDataset` in Argilla. Note that it must
            exist in advance. If you need help on how to create a `FeedbackDataset` in
            Argilla, please visit
            https://docs.argilla.io/en/latest/guides/llms/practical_guides/use_argilla_callback_in_langchain.html.
        workspace_name: name of the workspace in Argilla where the specified
            `FeedbackDataset` lives in. Defaults to `None`, which means that the
            default workspace will be used.
        api_url: URL of the Argilla Server that we want to use, and where the
            `FeedbackDataset` lives in. Defaults to `None`, which means that either
            `ARGILLA_API_URL` environment variable or the default will be used.
        api_key: API Key to connect to the Argilla Server. Defaults to `None`, which
            means that either `ARGILLA_API_KEY` environment variable or the default
            will be used.
    Raises:
        ImportError: if the `argilla` package is not installed.
        ConnectionError: if the connection to Argilla fails.
        FileNotFoundError: if the `FeedbackDataset` retrieval from Argilla fails.
    Examples:
        >>> from langchain_community.llms import OpenAI
        >>> from langchain_community.callbacks import ArgillaCallbackHandler
        >>> argilla_callback = ArgillaCallbackHandler(


>>> argilla_callback = ArgillaCallbackHandler(
        ...     dataset_name="my-dataset",
        ...     workspace_name="my-workspace",
        ...     api_url="http://localhost:6900",
        ...     api_key="argilla.apikey",
        ... )
        >>> llm = OpenAI(
        ...     temperature=0,
        ...     callbacks=[argilla_callback],
        ...     verbose=True,
        ...     openai_api_key="API_KEY_HERE",
        ... )
        >>> llm.generate([
        ...     "What is the best NLP-annotation tool out there? (no bias at all)",
        ... ])
        "Argilla, no doubt about it."
    """
    REPO_URL: str = "https://github.com/argilla-io/argilla"
    ISSUES_URL: str = f"{REPO_URL}/issues"
    BLOG_URL: str = "https://docs.argilla.io/en/latest/guides/llms/practical_guides/use_argilla_callback_in_langchain.html"  # noqa: E501
    DEFAULT_API_URL: str = "http://localhost:6900"
[docs]    def __init__(
        self,
        dataset_name: str,
        workspace_name: Optional[str] = None,
        api_url: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> None:
        """Initializes the `ArgillaCallbackHandler`.
        Args:
            dataset_name: name of the `FeedbackDataset` in Argilla. Note that it must
                exist in advance. If you need help on how to create a `FeedbackDataset`
                in Argilla, please visit


in Argilla, please visit
                https://docs.argilla.io/en/latest/guides/llms/practical_guides/use_argilla_callback_in_langchain.html.
            workspace_name: name of the workspace in Argilla where the specified
                `FeedbackDataset` lives in. Defaults to `None`, which means that the
                default workspace will be used.
            api_url: URL of the Argilla Server that we want to use, and where the
                `FeedbackDataset` lives in. Defaults to `None`, which means that either
                `ARGILLA_API_URL` environment variable or the default will be used.
            api_key: API Key to connect to the Argilla Server. Defaults to `None`, which
                means that either `ARGILLA_API_KEY` environment variable or the default
                will be used.
        Raises:
            ImportError: if the `argilla` package is not installed.
            ConnectionError: if the connection to Argilla fails.
            FileNotFoundError: if the `FeedbackDataset` retrieval from Argilla fails.
        """
        super().__init__()
        # Import Argilla (not via `import_argilla` to keep hints in IDEs)
        try:
            import argilla as rg  # noqa: F401
            self.ARGILLA_VERSION = rg.__version__
        except ImportError:
            raise ImportError(
                "To use the Argilla callback manager you need to have the `argilla` "
                "Python package installed. Please install it with `pip install argilla`"
            )
        # Check whether the Argilla version is compatible
        if parse(self.ARGILLA_VERSION) < parse("1.8.0"):
            raise ImportError(
                f"The installed `argilla` version is {self.ARGILLA_VERSION} but "


f"The installed `argilla` version is {self.ARGILLA_VERSION} but "
                "`ArgillaCallbackHandler` requires at least version 1.8.0. Please "
                "upgrade `argilla` with `pip install --upgrade argilla`."
            )
        # Show a warning message if Argilla will assume the default values will be used
        if api_url is None and os.getenv("ARGILLA_API_URL") is None:
            warnings.warn(
                (
                    "Since `api_url` is None, and the env var `ARGILLA_API_URL` is not"
                    f" set, it will default to `{self.DEFAULT_API_URL}`, which is the"
                    " default API URL in Argilla Quickstart."
                ),
            )
            api_url = self.DEFAULT_API_URL
        if api_key is None and os.getenv("ARGILLA_API_KEY") is None:
            self.DEFAULT_API_KEY = (
                "admin.apikey"
                if parse(self.ARGILLA_VERSION) < parse("1.11.0")
                else "owner.apikey"
            )
            warnings.warn(
                (
                    "Since `api_key` is None, and the env var `ARGILLA_API_KEY` is not"
                    f" set, it will default to `{self.DEFAULT_API_KEY}`, which is the"
                    " default API key in Argilla Quickstart."
                ),
            )
            api_url = self.DEFAULT_API_URL
        # Connect to Argilla with the provided credentials, if applicable
        try:
            rg.init(api_key=api_key, api_url=api_url)
        except Exception as e:
            raise ConnectionError(
                f"Could not connect to Argilla with exception: '{e}'.\n"


f"Could not connect to Argilla with exception: '{e}'.\n"
                "Please check your `api_key` and `api_url`, and make sure that "
                "the Argilla server is up and running. If the problem persists "
                f"please report it to {self.ISSUES_URL} as an `integration` issue."
            ) from e
        # Set the Argilla variables
        self.dataset_name = dataset_name
        self.workspace_name = workspace_name or rg.get_workspace()
        # Retrieve the `FeedbackDataset` from Argilla (without existing records)
        try:
            extra_args = {}
            if parse(self.ARGILLA_VERSION) < parse("1.14.0"):
                warnings.warn(
                    f"You have Argilla {self.ARGILLA_VERSION}, but Argilla 1.14.0 or"
                    " higher is recommended.",
                    UserWarning,
                )
                extra_args = {"with_records": False}
            self.dataset = rg.FeedbackDataset.from_argilla(
                name=self.dataset_name,
                workspace=self.workspace_name,
                **extra_args,
            )
        except Exception as e:
            raise FileNotFoundError(
                f"`FeedbackDataset` retrieval from Argilla failed with exception `{e}`."
                f"\nPlease check that the dataset with name={self.dataset_name} in the"
                f" workspace={self.workspace_name} exists in advance. If you need help"
                " on how to create a `langchain`-compatible `FeedbackDataset` in"
                f" Argilla, please visit {self.BLOG_URL}. If the problem persists"
                f" please report it to {self.ISSUES_URL} as an `integration` issue."
            ) from e
        supported_fields = ["prompt", "response"]


) from e
        supported_fields = ["prompt", "response"]
        if supported_fields != [field.name for field in self.dataset.fields]:
            raise ValueError(
                f"`FeedbackDataset` with name={self.dataset_name} in the workspace="
                f"{self.workspace_name} had fields that are not supported yet for the"
                f"`langchain` integration. Supported fields are: {supported_fields},"
                f" and the current `FeedbackDataset` fields are {[field.name for field in self.dataset.fields]}."  # noqa: E501
                " For more information on how to create a `langchain`-compatible"
                f" `FeedbackDataset` in Argilla, please visit {self.BLOG_URL}."
            )
        self.prompts: Dict[str, List[str]] = {}
        warnings.warn(
            (
                "The `ArgillaCallbackHandler` is currently in beta and is subject to"
                " change based on updates to `langchain`. Please report any issues to"
                f" {self.ISSUES_URL} as an `integration` issue."
            ),
        )
[docs]    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Save the prompts in memory when an LLM starts."""
        self.prompts.update({str(kwargs["parent_run_id"] or kwargs["run_id"]): prompts})
[docs]    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Do nothing when a new token is generated."""
        pass
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Log records to Argilla when an LLM ends."""


"""Log records to Argilla when an LLM ends."""
        # Do nothing if there's a parent_run_id, since we will log the records when
        # the chain ends
        if kwargs["parent_run_id"]:
            return
        # Creates the records and adds them to the `FeedbackDataset`
        prompts = self.prompts[str(kwargs["run_id"])]
        for prompt, generations in zip(prompts, response.generations):
            self.dataset.add_records(
                records=[
                    {
                        "fields": {
                            "prompt": prompt,
                            "response": generation.text.strip(),
                        },
                    }
                    for generation in generations
                ]
            )
        # Pop current run from `self.runs`
        self.prompts.pop(str(kwargs["run_id"]))
        if parse(self.ARGILLA_VERSION) < parse("1.14.0"):
            # Push the records to Argilla
            self.dataset.push_to_argilla()
[docs]    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when LLM outputs an error."""
        pass
[docs]    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """If the key `input` is in `inputs`, then save it in `self.prompts` using
        either the `parent_run_id` or the `run_id` as the key. This is done so that
        we don't log the same input prompt twice, once when the LLM starts and once
        when the chain starts.
        """
        if "input" in inputs:
            self.prompts.update(
                {


if "input" in inputs:
            self.prompts.update(
                {
                    str(kwargs["parent_run_id"] or kwargs["run_id"]): (
                        inputs["input"]
                        if isinstance(inputs["input"], list)
                        else [inputs["input"]]
                    )
                }
            )
[docs]    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """If either the `parent_run_id` or the `run_id` is in `self.prompts`, then
        log the outputs to Argilla, and pop the run from `self.prompts`. The behavior
        differs if the output is a list or not.
        """
        if not any(
            key in self.prompts
            for key in [str(kwargs["parent_run_id"]), str(kwargs["run_id"])]
        ):
            return
        prompts = self.prompts.get(str(kwargs["parent_run_id"])) or self.prompts.get(
            str(kwargs["run_id"])
        )
        for chain_output_key, chain_output_val in outputs.items():
            if isinstance(chain_output_val, list):
                # Creates the records and adds them to the `FeedbackDataset`
                self.dataset.add_records(
                    records=[
                        {
                            "fields": {
                                "prompt": prompt,
                                "response": output["text"].strip(),
                            },
                        }
                        for prompt, output in zip(
                            prompts,  # type: ignore
                            chain_output_val,
                        )
                    ]
                )
            else:
                # Creates the records and adds them to the `FeedbackDataset`
                self.dataset.add_records(
                    records=[
                        {
                            "fields": {
                                "prompt": " ".join(prompts),  # type: ignore


"prompt": " ".join(prompts),  # type: ignore
                                "response": chain_output_val.strip(),
                            },
                        }
                    ]
                )
        # Pop current run from `self.runs`
        if str(kwargs["parent_run_id"]) in self.prompts:
            self.prompts.pop(str(kwargs["parent_run_id"]))
        if str(kwargs["run_id"]) in self.prompts:
            self.prompts.pop(str(kwargs["run_id"]))
        if parse(self.ARGILLA_VERSION) < parse("1.14.0"):
            # Push the records to Argilla
            self.dataset.push_to_argilla()
[docs]    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when LLM chain outputs an error."""
        pass
[docs]    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        """Do nothing when tool starts."""
        pass
[docs]    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
        """Do nothing when agent takes a specific action."""
        pass
[docs]    def on_tool_end(
        self,
        output: str,
        observation_prefix: Optional[str] = None,
        llm_prefix: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """Do nothing when tool ends."""
        pass
[docs]    def on_tool_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when tool outputs an error."""
        pass


"""Do nothing when tool outputs an error."""
        pass
[docs]    def on_text(self, text: str, **kwargs: Any) -> None:
        """Do nothing"""
        pass
[docs]    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
        """Do nothing"""
        pass


Source code for langchain_community.callbacks.wandb_callback
import json
import tempfile
from copy import deepcopy
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Union
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_community.callbacks.utils import (
    BaseMetadataCallbackHandler,
    flatten_dict,
    hash_string,
    import_pandas,
    import_spacy,
    import_textstat,
)
[docs]def import_wandb() -> Any:
    """Import the wandb python package and raise an error if it is not installed."""
    try:
        import wandb  # noqa: F401
    except ImportError:
        raise ImportError(
            "To use the wandb callback manager you need to have the `wandb` python "
            "package installed. Please install it with `pip install wandb`"
        )
    return wandb
[docs]def load_json_to_dict(json_path: Union[str, Path]) -> dict:
    """Load json file to a dictionary.
    Parameters:
        json_path (str): The path to the json file.
    Returns:
        (dict): The dictionary representation of the json file.
    """
    with open(json_path, "r") as f:
        data = json.load(f)
    return data
[docs]def analyze_text(
    text: str,
    complexity_metrics: bool = True,
    visualize: bool = True,
    nlp: Any = None,
    output_dir: Optional[Union[str, Path]] = None,
) -> dict:
    """Analyze text using textstat and spacy.
    Parameters:


"""Analyze text using textstat and spacy.
    Parameters:
        text (str): The text to analyze.
        complexity_metrics (bool): Whether to compute complexity metrics.
        visualize (bool): Whether to visualize the text.
        nlp (spacy.lang): The spacy language model to use for visualization.
        output_dir (str): The directory to save the visualization files to.
    Returns:
        (dict): A dictionary containing the complexity metrics and visualization
            files serialized in a wandb.Html element.
    """
    resp = {}
    textstat = import_textstat()
    wandb = import_wandb()
    spacy = import_spacy()
    if complexity_metrics:
        text_complexity_metrics = {
            "flesch_reading_ease": textstat.flesch_reading_ease(text),
            "flesch_kincaid_grade": textstat.flesch_kincaid_grade(text),
            "smog_index": textstat.smog_index(text),
            "coleman_liau_index": textstat.coleman_liau_index(text),
            "automated_readability_index": textstat.automated_readability_index(text),
            "dale_chall_readability_score": textstat.dale_chall_readability_score(text),
            "difficult_words": textstat.difficult_words(text),
            "linsear_write_formula": textstat.linsear_write_formula(text),
            "gunning_fog": textstat.gunning_fog(text),
            "text_standard": textstat.text_standard(text),
            "fernandez_huerta": textstat.fernandez_huerta(text),
            "szigriszt_pazos": textstat.szigriszt_pazos(text),


"gutierrez_polini": textstat.gutierrez_polini(text),
            "crawford": textstat.crawford(text),
            "gulpease_index": textstat.gulpease_index(text),
            "osman": textstat.osman(text),
        }
        resp.update(text_complexity_metrics)
    if visualize and nlp and output_dir is not None:
        doc = nlp(text)
        dep_out = spacy.displacy.render(  # type: ignore
            doc, style="dep", jupyter=False, page=True
        )
        dep_output_path = Path(output_dir, hash_string(f"dep-{text}") + ".html")
        dep_output_path.open("w", encoding="utf-8").write(dep_out)
        ent_out = spacy.displacy.render(  # type: ignore
            doc, style="ent", jupyter=False, page=True
        )
        ent_output_path = Path(output_dir, hash_string(f"ent-{text}") + ".html")
        ent_output_path.open("w", encoding="utf-8").write(ent_out)
        text_visualizations = {
            "dependency_tree": wandb.Html(str(dep_output_path)),
            "entities": wandb.Html(str(ent_output_path)),
        }
        resp.update(text_visualizations)
    return resp
[docs]def construct_html_from_prompt_and_generation(prompt: str, generation: str) -> Any:
    """Construct an html element from a prompt and a generation.
    Parameters:
        prompt (str): The prompt.
        generation (str): The generation.
    Returns:
        (wandb.Html): The html element."""
    wandb = import_wandb()
    formatted_prompt = prompt.replace("\n", "<br>")


formatted_prompt = prompt.replace("\n", "
")
    formatted_generation = generation.replace("\n", "
")
    return wandb.Html(
        f"""

{formatted_prompt}:

        {formatted_generation}

    """,
        inject=False,
    )
[docs]class WandbCallbackHandler(BaseMetadataCallbackHandler, BaseCallbackHandler):
    """Callback Handler that logs to Weights and Biases.
    Parameters:
        job_type (str): The type of job.
        project (str): The project to log to.
        entity (str): The entity to log to.
        tags (list): The tags to log.
        group (str): The group to log to.
        name (str): The name of the run.
        notes (str): The notes to log.
        visualize (bool): Whether to visualize the run.
        complexity_metrics (bool): Whether to log complexity metrics.
        stream_logs (bool): Whether to stream callback actions to W&B
    This handler will utilize the associated callback method called and formats
    the input of each callback function with metadata regarding the state of LLM run,
    and adds the response to the list of records for both the {method}_records and
    action. It then logs the response using the run.log() method to Weights and Biases.
    """
[docs]    def __init__(
        self,
        job_type: Optional[str] = None,
        project: Optional[str] = "langchain_callback_demo",
        entity: Optional[str] = None,
        tags: Optional[Sequence] = None,
        group: Optional[str] = None,


group: Optional[str] = None,
        name: Optional[str] = None,
        notes: Optional[str] = None,
        visualize: bool = False,
        complexity_metrics: bool = False,
        stream_logs: bool = False,
    ) -> None:
        """Initialize callback handler."""
        wandb = import_wandb()
        import_pandas()
        import_textstat()
        spacy = import_spacy()
        super().__init__()
        self.job_type = job_type
        self.project = project
        self.entity = entity
        self.tags = tags
        self.group = group
        self.name = name
        self.notes = notes
        self.visualize = visualize
        self.complexity_metrics = complexity_metrics
        self.stream_logs = stream_logs
        self.temp_dir = tempfile.TemporaryDirectory()
        self.run: wandb.sdk.wandb_run.Run = wandb.init(  # type: ignore
            job_type=self.job_type,
            project=self.project,
            entity=self.entity,
            tags=self.tags,
            group=self.group,
            name=self.name,
            notes=self.notes,
        )
        warning = (
            "DEPRECATION: The `WandbCallbackHandler` will soon be deprecated in favor "
            "of the `WandbTracer`. Please update your code to use the `WandbTracer` "
            "instead."
        )
        wandb.termwarn(
            warning,
            repeat=False,
        )
        self.callback_columns: list = []
        self.action_records: list = []
        self.complexity_metrics = complexity_metrics
        self.visualize = visualize
        self.nlp = spacy.load("en_core_web_sm")
    def _init_resp(self) -> Dict:


def _init_resp(self) -> Dict:
        return {k: None for k in self.callback_columns}
[docs]    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Run when LLM starts."""
        self.step += 1
        self.llm_starts += 1
        self.starts += 1
        resp = self._init_resp()
        resp.update({"action": "on_llm_start"})
        resp.update(flatten_dict(serialized))
        resp.update(self.get_custom_callback_meta())
        for prompt in prompts:
            prompt_resp = deepcopy(resp)
            prompt_resp["prompts"] = prompt
            self.on_llm_start_records.append(prompt_resp)
            self.action_records.append(prompt_resp)
            if self.stream_logs:
                self.run.log(prompt_resp)
[docs]    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Run when LLM generates a new token."""
        self.step += 1
        self.llm_streams += 1
        resp = self._init_resp()
        resp.update({"action": "on_llm_new_token", "token": token})
        resp.update(self.get_custom_callback_meta())
        self.on_llm_token_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.run.log(resp)
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Run when LLM ends running."""
        self.step += 1
        self.llm_ends += 1
        self.ends += 1
        resp = self._init_resp()


self.ends += 1
        resp = self._init_resp()
        resp.update({"action": "on_llm_end"})
        resp.update(flatten_dict(response.llm_output or {}))
        resp.update(self.get_custom_callback_meta())
        for generations in response.generations:
            for generation in generations:
                generation_resp = deepcopy(resp)
                generation_resp.update(flatten_dict(generation.dict()))
                generation_resp.update(
                    analyze_text(
                        generation.text,
                        complexity_metrics=self.complexity_metrics,
                        visualize=self.visualize,
                        nlp=self.nlp,
                        output_dir=self.temp_dir.name,
                    )
                )
                self.on_llm_end_records.append(generation_resp)
                self.action_records.append(generation_resp)
                if self.stream_logs:
                    self.run.log(generation_resp)
[docs]    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when LLM errors."""
        self.step += 1
        self.errors += 1
[docs]    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Run when chain starts running."""
        self.step += 1
        self.chain_starts += 1
        self.starts += 1
        resp = self._init_resp()
        resp.update({"action": "on_chain_start"})
        resp.update(flatten_dict(serialized))
        resp.update(self.get_custom_callback_meta())
        chain_input = inputs["input"]
        if isinstance(chain_input, str):
            input_resp = deepcopy(resp)
            input_resp["input"] = chain_input
            self.on_chain_start_records.append(input_resp)
            self.action_records.append(input_resp)


self.on_chain_start_records.append(input_resp)
            self.action_records.append(input_resp)
            if self.stream_logs:
                self.run.log(input_resp)
        elif isinstance(chain_input, list):
            for inp in chain_input:
                input_resp = deepcopy(resp)
                input_resp.update(inp)
                self.on_chain_start_records.append(input_resp)
                self.action_records.append(input_resp)
                if self.stream_logs:
                    self.run.log(input_resp)
        else:
            raise ValueError("Unexpected data format provided!")
[docs]    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """Run when chain ends running."""
        self.step += 1
        self.chain_ends += 1
        self.ends += 1
        resp = self._init_resp()
        resp.update({"action": "on_chain_end", "outputs": outputs["output"]})
        resp.update(self.get_custom_callback_meta())
        self.on_chain_end_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.run.log(resp)
[docs]    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when chain errors."""
        self.step += 1
        self.errors += 1
[docs]    def on_tool_start(
        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
    ) -> None:
        """Run when tool starts running."""
        self.step += 1
        self.tool_starts += 1
        self.starts += 1
        resp = self._init_resp()
        resp.update({"action": "on_tool_start", "input_str": input_str})
        resp.update(flatten_dict(serialized))


resp.update(flatten_dict(serialized))
        resp.update(self.get_custom_callback_meta())
        self.on_tool_start_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.run.log(resp)
[docs]    def on_tool_end(self, output: str, **kwargs: Any) -> None:
        """Run when tool ends running."""
        self.step += 1
        self.tool_ends += 1
        self.ends += 1
        resp = self._init_resp()
        resp.update({"action": "on_tool_end", "output": output})
        resp.update(self.get_custom_callback_meta())
        self.on_tool_end_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.run.log(resp)
[docs]    def on_tool_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when tool errors."""
        self.step += 1
        self.errors += 1
[docs]    def on_text(self, text: str, **kwargs: Any) -> None:
        """
        Run when agent is ending.
        """
        self.step += 1
        self.text_ctr += 1
        resp = self._init_resp()
        resp.update({"action": "on_text", "text": text})
        resp.update(self.get_custom_callback_meta())
        self.on_text_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.run.log(resp)
[docs]    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
        """Run when agent ends running."""
        self.step += 1
        self.agent_ends += 1
        self.ends += 1


self.agent_ends += 1
        self.ends += 1
        resp = self._init_resp()
        resp.update(
            {
                "action": "on_agent_finish",
                "output": finish.return_values["output"],
                "log": finish.log,
            }
        )
        resp.update(self.get_custom_callback_meta())
        self.on_agent_finish_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.run.log(resp)
[docs]    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
        """Run on agent action."""
        self.step += 1
        self.tool_starts += 1
        self.starts += 1
        resp = self._init_resp()
        resp.update(
            {
                "action": "on_agent_action",
                "tool": action.tool,
                "tool_input": action.tool_input,
                "log": action.log,
            }
        )
        resp.update(self.get_custom_callback_meta())
        self.on_agent_action_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.run.log(resp)
    def _create_session_analysis_df(self) -> Any:
        """Create a dataframe with all the information from the session."""
        pd = import_pandas()
        on_llm_start_records_df = pd.DataFrame(self.on_llm_start_records)
        on_llm_end_records_df = pd.DataFrame(self.on_llm_end_records)
        llm_input_prompts_df = (
            on_llm_start_records_df[["step", "prompts", "name"]]
            .dropna(axis=1)
            .rename({"step": "prompt_step"}, axis=1)
        )
        complexity_metrics_columns = []


)
        complexity_metrics_columns = []
        visualizations_columns = []
        if self.complexity_metrics:
            complexity_metrics_columns = [
                "flesch_reading_ease",
                "flesch_kincaid_grade",
                "smog_index",
                "coleman_liau_index",
                "automated_readability_index",
                "dale_chall_readability_score",
                "difficult_words",
                "linsear_write_formula",
                "gunning_fog",
                "text_standard",
                "fernandez_huerta",
                "szigriszt_pazos",
                "gutierrez_polini",
                "crawford",
                "gulpease_index",
                "osman",
            ]
        if self.visualize:
            visualizations_columns = ["dependency_tree", "entities"]
        llm_outputs_df = (
            on_llm_end_records_df[
                [
                    "step",
                    "text",
                    "token_usage_total_tokens",
                    "token_usage_prompt_tokens",
                    "token_usage_completion_tokens",
                ]
                + complexity_metrics_columns
                + visualizations_columns
            ]
            .dropna(axis=1)
            .rename({"step": "output_step", "text": "output"}, axis=1)
        )
        session_analysis_df = pd.concat([llm_input_prompts_df, llm_outputs_df], axis=1)
        session_analysis_df["chat_html"] = session_analysis_df[
            ["prompts", "output"]
        ].apply(
            lambda row: construct_html_from_prompt_and_generation(
                row["prompts"], row["output"]
            ),
            axis=1,
        )
        return session_analysis_df


),
            axis=1,
        )
        return session_analysis_df
[docs]    def flush_tracker(
        self,
        langchain_asset: Any = None,
        reset: bool = True,
        finish: bool = False,
        job_type: Optional[str] = None,
        project: Optional[str] = None,
        entity: Optional[str] = None,
        tags: Optional[Sequence] = None,
        group: Optional[str] = None,
        name: Optional[str] = None,
        notes: Optional[str] = None,
        visualize: Optional[bool] = None,
        complexity_metrics: Optional[bool] = None,
    ) -> None:
        """Flush the tracker and reset the session.
        Args:
            langchain_asset: The langchain asset to save.
            reset: Whether to reset the session.
            finish: Whether to finish the run.
            job_type: The job type.
            project: The project.
            entity: The entity.
            tags: The tags.
            group: The group.
            name: The name.
            notes: The notes.
            visualize: Whether to visualize.
            complexity_metrics: Whether to compute complexity metrics.
            Returns:
                None
        """
        pd = import_pandas()
        wandb = import_wandb()
        action_records_table = wandb.Table(dataframe=pd.DataFrame(self.action_records))
        session_analysis_table = wandb.Table(
            dataframe=self._create_session_analysis_df()
        )
        self.run.log(
            {
                "action_records": action_records_table,
                "session_analysis": session_analysis_table,
            }
        )
        if langchain_asset:


}
        )
        if langchain_asset:
            langchain_asset_path = Path(self.temp_dir.name, "model.json")
            model_artifact = wandb.Artifact(name="model", type="model")
            model_artifact.add(action_records_table, name="action_records")
            model_artifact.add(session_analysis_table, name="session_analysis")
            try:
                langchain_asset.save(langchain_asset_path)
                model_artifact.add_file(str(langchain_asset_path))
                model_artifact.metadata = load_json_to_dict(langchain_asset_path)
            except ValueError:
                langchain_asset.save_agent(langchain_asset_path)
                model_artifact.add_file(str(langchain_asset_path))
                model_artifact.metadata = load_json_to_dict(langchain_asset_path)
            except NotImplementedError as e:
                print("Could not save model.")
                print(repr(e))
                pass
            self.run.log_artifact(model_artifact)
        if finish or reset:
            self.run.finish()
            self.temp_dir.cleanup()
            self.reset_callback_meta()
        if reset:
            self.__init__(  # type: ignore
                job_type=job_type if job_type else self.job_type,
                project=project if project else self.project,
                entity=entity if entity else self.entity,
                tags=tags if tags else self.tags,
                group=group if group else self.group,
                name=name if name else self.name,
                notes=notes if notes else self.notes,
                visualize=visualize if visualize else self.visualize,
                complexity_metrics=complexity_metrics
                if complexity_metrics
                else self.complexity_metrics,
            )


Source code for langchain_community.callbacks.confident_callback
# flake8: noqa
import os
import warnings
from typing import Any, Dict, List, Optional, Union
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.outputs import LLMResult
[docs]class DeepEvalCallbackHandler(BaseCallbackHandler):
    """Callback Handler that logs into deepeval.
    Args:
        implementation_name: name of the `implementation` in deepeval
        metrics: A list of metrics
    Raises:
        ImportError: if the `deepeval` package is not installed.
    Examples:
        >>> from langchain_community.llms import OpenAI
        >>> from langchain_community.callbacks import DeepEvalCallbackHandler
        >>> from deepeval.metrics import AnswerRelevancy
        >>> metric = AnswerRelevancy(minimum_score=0.3)
        >>> deepeval_callback = DeepEvalCallbackHandler(
        ...     implementation_name="exampleImplementation",
        ...     metrics=[metric],
        ... )
        >>> llm = OpenAI(
        ...     temperature=0,
        ...     callbacks=[deepeval_callback],
        ...     verbose=True,
        ...     openai_api_key="API_KEY_HERE",
        ... )
        >>> llm.generate([
        ...     "What is the best evaluation tool out there? (no bias at all)",
        ... ])
        "Deepeval, no doubt about it."
    """
    REPO_URL: str = "https://github.com/confident-ai/deepeval"
    ISSUES_URL: str = f"{REPO_URL}/issues"
    BLOG_URL: str = "https://docs.confident-ai.com"  # noqa: E501


[docs]    def __init__(
        self,
        metrics: List[Any],
        implementation_name: Optional[str] = None,
    ) -> None:
        """Initializes the `deepevalCallbackHandler`.
        Args:
            implementation_name: Name of the implementation you want.
            metrics: What metrics do you want to track?
        Raises:
            ImportError: if the `deepeval` package is not installed.
            ConnectionError: if the connection to deepeval fails.
        """
        super().__init__()
        # Import deepeval (not via `import_deepeval` to keep hints in IDEs)
        try:
            import deepeval  # ignore: F401,I001
        except ImportError:
            raise ImportError(
                """To use the deepeval callback manager you need to have the
                `deepeval` Python package installed. Please install it with
                `pip install deepeval`"""
            )
        if os.path.exists(".deepeval"):
            warnings.warn(
                """You are currently not logging anything to the dashboard, we
                recommend using `deepeval login`."""
            )
        # Set the deepeval variables
        self.implementation_name = implementation_name
        self.metrics = metrics
        warnings.warn(
            (
                "The `DeepEvalCallbackHandler` is currently in beta and is subject to"
                " change based on updates to `langchain`. Please report any issues to"
                f" {self.ISSUES_URL} as an `integration` issue."
            ),
        )
[docs]    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:


) -> None:
        """Store the prompts"""
        self.prompts = prompts
[docs]    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Do nothing when a new token is generated."""
        pass
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Log records to deepeval when an LLM ends."""
        from deepeval.metrics.answer_relevancy import AnswerRelevancy
        from deepeval.metrics.bias_classifier import UnBiasedMetric
        from deepeval.metrics.metric import Metric
        from deepeval.metrics.toxic_classifier import NonToxicMetric
        for metric in self.metrics:
            for i, generation in enumerate(response.generations):
                # Here, we only measure the first generation's output
                output = generation[0].text
                query = self.prompts[i]
                if isinstance(metric, AnswerRelevancy):
                    result = metric.measure(
                        output=output,
                        query=query,
                    )
                    print(f"Answer Relevancy: {result}")
                elif isinstance(metric, UnBiasedMetric):
                    score = metric.measure(output)
                    print(f"Bias Score: {score}")
                elif isinstance(metric, NonToxicMetric):
                    score = metric.measure(output)
                    print(f"Toxic Score: {score}")
                else:
                    raise ValueError(
                        f"""Metric {metric.__name__} is not supported by deepeval
                        callbacks."""
                    )
[docs]    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when LLM outputs an error."""
        pass
[docs]    def on_chain_start(


pass
[docs]    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Do nothing when chain starts"""
        pass
[docs]    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """Do nothing when chain ends."""
        pass
[docs]    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when LLM chain outputs an error."""
        pass
[docs]    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        """Do nothing when tool starts."""
        pass
[docs]    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
        """Do nothing when agent takes a specific action."""
        pass
[docs]    def on_tool_end(
        self,
        output: str,
        observation_prefix: Optional[str] = None,
        llm_prefix: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """Do nothing when tool ends."""
        pass
[docs]    def on_tool_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when tool outputs an error."""
        pass
[docs]    def on_text(self, text: str, **kwargs: Any) -> None:
        """Do nothing"""
        pass
[docs]    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
        """Do nothing"""
        pass


Source code for langchain_community.callbacks.clearml_callback
from __future__ import annotations
import tempfile
from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Sequence
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_community.callbacks.utils import (
    BaseMetadataCallbackHandler,
    flatten_dict,
    hash_string,
    import_pandas,
    import_spacy,
    import_textstat,
    load_json,
)
if TYPE_CHECKING:
    import pandas as pd
[docs]def import_clearml() -> Any:
    """Import the clearml python package and raise an error if it is not installed."""
    try:
        import clearml  # noqa: F401
    except ImportError:
        raise ImportError(
            "To use the clearml callback manager you need to have the `clearml` python "
            "package installed. Please install it with `pip install clearml`"
        )
    return clearml
[docs]class ClearMLCallbackHandler(BaseMetadataCallbackHandler, BaseCallbackHandler):
    """Callback Handler that logs to ClearML.
    Parameters:
        job_type (str): The type of clearml task such as "inference", "testing" or "qc"
        project_name (str): The clearml project name
        tags (list): Tags to add to the task
        task_name (str): Name of the clearml task
        visualize (bool): Whether to visualize the run.
        complexity_metrics (bool): Whether to log complexity metrics
        stream_logs (bool): Whether to stream callback actions to ClearML
    This handler will utilize the associated callback method and formats


This handler will utilize the associated callback method and formats
    the input of each callback function with metadata regarding the state of LLM run,
    and adds the response to the list of records for both the {method}_records and
    action. It then logs the response to the ClearML console.
    """
[docs]    def __init__(
        self,
        task_type: Optional[str] = "inference",
        project_name: Optional[str] = "langchain_callback_demo",
        tags: Optional[Sequence] = None,
        task_name: Optional[str] = None,
        visualize: bool = False,
        complexity_metrics: bool = False,
        stream_logs: bool = False,
    ) -> None:
        """Initialize callback handler."""
        clearml = import_clearml()
        spacy = import_spacy()
        super().__init__()
        self.task_type = task_type
        self.project_name = project_name
        self.tags = tags
        self.task_name = task_name
        self.visualize = visualize
        self.complexity_metrics = complexity_metrics
        self.stream_logs = stream_logs
        self.temp_dir = tempfile.TemporaryDirectory()
        # Check if ClearML task already exists (e.g. in pipeline)
        if clearml.Task.current_task():
            self.task = clearml.Task.current_task()
        else:
            self.task = clearml.Task.init(  # type: ignore
                task_type=self.task_type,
                project_name=self.project_name,
                tags=self.tags,
                task_name=self.task_name,
                output_uri=True,
            )
        self.logger = self.task.get_logger()
        warning = (
            "The clearml callback is currently in beta and is subject to change "


"The clearml callback is currently in beta and is subject to change "
            "based on updates to `langchain`. Please report any issues to "
            "https://github.com/allegroai/clearml/issues with the tag `langchain`."
        )
        self.logger.report_text(warning, level=30, print_console=True)
        self.callback_columns: list = []
        self.action_records: list = []
        self.complexity_metrics = complexity_metrics
        self.visualize = visualize
        self.nlp = spacy.load("en_core_web_sm")
    def _init_resp(self) -> Dict:
        return {k: None for k in self.callback_columns}
[docs]    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Run when LLM starts."""
        self.step += 1
        self.llm_starts += 1
        self.starts += 1
        resp = self._init_resp()
        resp.update({"action": "on_llm_start"})
        resp.update(flatten_dict(serialized))
        resp.update(self.get_custom_callback_meta())
        for prompt in prompts:
            prompt_resp = deepcopy(resp)
            prompt_resp["prompts"] = prompt
            self.on_llm_start_records.append(prompt_resp)
            self.action_records.append(prompt_resp)
            if self.stream_logs:
                self.logger.report_text(prompt_resp)
[docs]    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Run when LLM generates a new token."""
        self.step += 1
        self.llm_streams += 1
        resp = self._init_resp()


self.llm_streams += 1
        resp = self._init_resp()
        resp.update({"action": "on_llm_new_token", "token": token})
        resp.update(self.get_custom_callback_meta())
        self.on_llm_token_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.logger.report_text(resp)
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Run when LLM ends running."""
        self.step += 1
        self.llm_ends += 1
        self.ends += 1
        resp = self._init_resp()
        resp.update({"action": "on_llm_end"})
        resp.update(flatten_dict(response.llm_output or {}))
        resp.update(self.get_custom_callback_meta())
        for generations in response.generations:
            for generation in generations:
                generation_resp = deepcopy(resp)
                generation_resp.update(flatten_dict(generation.dict()))
                generation_resp.update(self.analyze_text(generation.text))
                self.on_llm_end_records.append(generation_resp)
                self.action_records.append(generation_resp)
                if self.stream_logs:
                    self.logger.report_text(generation_resp)
[docs]    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when LLM errors."""
        self.step += 1
        self.errors += 1
[docs]    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Run when chain starts running."""
        self.step += 1
        self.chain_starts += 1
        self.starts += 1


self.chain_starts += 1
        self.starts += 1
        resp = self._init_resp()
        resp.update({"action": "on_chain_start"})
        resp.update(flatten_dict(serialized))
        resp.update(self.get_custom_callback_meta())
        chain_input = inputs.get("input", inputs.get("human_input"))
        if isinstance(chain_input, str):
            input_resp = deepcopy(resp)
            input_resp["input"] = chain_input
            self.on_chain_start_records.append(input_resp)
            self.action_records.append(input_resp)
            if self.stream_logs:
                self.logger.report_text(input_resp)
        elif isinstance(chain_input, list):
            for inp in chain_input:
                input_resp = deepcopy(resp)
                input_resp.update(inp)
                self.on_chain_start_records.append(input_resp)
                self.action_records.append(input_resp)
                if self.stream_logs:
                    self.logger.report_text(input_resp)
        else:
            raise ValueError("Unexpected data format provided!")
[docs]    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """Run when chain ends running."""
        self.step += 1
        self.chain_ends += 1
        self.ends += 1
        resp = self._init_resp()
        resp.update(
            {
                "action": "on_chain_end",
                "outputs": outputs.get("output", outputs.get("text")),
            }
        )
        resp.update(self.get_custom_callback_meta())
        self.on_chain_end_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.logger.report_text(resp)
[docs]    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when chain errors."""


"""Run when chain errors."""
        self.step += 1
        self.errors += 1
[docs]    def on_tool_start(
        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
    ) -> None:
        """Run when tool starts running."""
        self.step += 1
        self.tool_starts += 1
        self.starts += 1
        resp = self._init_resp()
        resp.update({"action": "on_tool_start", "input_str": input_str})
        resp.update(flatten_dict(serialized))
        resp.update(self.get_custom_callback_meta())
        self.on_tool_start_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.logger.report_text(resp)
[docs]    def on_tool_end(self, output: str, **kwargs: Any) -> None:
        """Run when tool ends running."""
        self.step += 1
        self.tool_ends += 1
        self.ends += 1
        resp = self._init_resp()
        resp.update({"action": "on_tool_end", "output": output})
        resp.update(self.get_custom_callback_meta())
        self.on_tool_end_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.logger.report_text(resp)
[docs]    def on_tool_error(self, error: BaseException, **kwargs: Any) -> None:
        """Run when tool errors."""
        self.step += 1
        self.errors += 1
[docs]    def on_text(self, text: str, **kwargs: Any) -> None:
        """
        Run when agent is ending.
        """
        self.step += 1
        self.text_ctr += 1
        resp = self._init_resp()


self.text_ctr += 1
        resp = self._init_resp()
        resp.update({"action": "on_text", "text": text})
        resp.update(self.get_custom_callback_meta())
        self.on_text_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.logger.report_text(resp)
[docs]    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
        """Run when agent ends running."""
        self.step += 1
        self.agent_ends += 1
        self.ends += 1
        resp = self._init_resp()
        resp.update(
            {
                "action": "on_agent_finish",
                "output": finish.return_values["output"],
                "log": finish.log,
            }
        )
        resp.update(self.get_custom_callback_meta())
        self.on_agent_finish_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.logger.report_text(resp)
[docs]    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
        """Run on agent action."""
        self.step += 1
        self.tool_starts += 1
        self.starts += 1
        resp = self._init_resp()
        resp.update(
            {
                "action": "on_agent_action",
                "tool": action.tool,
                "tool_input": action.tool_input,
                "log": action.log,
            }
        )
        resp.update(self.get_custom_callback_meta())
        self.on_agent_action_records.append(resp)
        self.action_records.append(resp)
        if self.stream_logs:
            self.logger.report_text(resp)
[docs]    def analyze_text(self, text: str) -> dict:


[docs]    def analyze_text(self, text: str) -> dict:
        """Analyze text using textstat and spacy.
        Parameters:
            text (str): The text to analyze.
        Returns:
            (dict): A dictionary containing the complexity metrics.
        """
        resp = {}
        textstat = import_textstat()
        spacy = import_spacy()
        if self.complexity_metrics:
            text_complexity_metrics = {
                "flesch_reading_ease": textstat.flesch_reading_ease(text),
                "flesch_kincaid_grade": textstat.flesch_kincaid_grade(text),
                "smog_index": textstat.smog_index(text),
                "coleman_liau_index": textstat.coleman_liau_index(text),
                "automated_readability_index": textstat.automated_readability_index(
                    text
                ),
                "dale_chall_readability_score": textstat.dale_chall_readability_score(
                    text
                ),
                "difficult_words": textstat.difficult_words(text),
                "linsear_write_formula": textstat.linsear_write_formula(text),
                "gunning_fog": textstat.gunning_fog(text),
                "text_standard": textstat.text_standard(text),
                "fernandez_huerta": textstat.fernandez_huerta(text),
                "szigriszt_pazos": textstat.szigriszt_pazos(text),
                "gutierrez_polini": textstat.gutierrez_polini(text),
                "crawford": textstat.crawford(text),
                "gulpease_index": textstat.gulpease_index(text),
                "osman": textstat.osman(text),
            }


"osman": textstat.osman(text),
            }
            resp.update(text_complexity_metrics)
        if self.visualize and self.nlp and self.temp_dir.name is not None:
            doc = self.nlp(text)
            dep_out = spacy.displacy.render(  # type: ignore
                doc, style="dep", jupyter=False, page=True
            )
            dep_output_path = Path(
                self.temp_dir.name, hash_string(f"dep-{text}") + ".html"
            )
            dep_output_path.open("w", encoding="utf-8").write(dep_out)
            ent_out = spacy.displacy.render(  # type: ignore
                doc, style="ent", jupyter=False, page=True
            )
            ent_output_path = Path(
                self.temp_dir.name, hash_string(f"ent-{text}") + ".html"
            )
            ent_output_path.open("w", encoding="utf-8").write(ent_out)
            self.logger.report_media(
                "Dependencies Plot", text, local_path=dep_output_path
            )
            self.logger.report_media("Entities Plot", text, local_path=ent_output_path)
        return resp
    @staticmethod
    def _build_llm_df(
        base_df: pd.DataFrame, base_df_fields: Sequence, rename_map: Mapping
    ) -> pd.DataFrame:
        base_df_fields = [field for field in base_df_fields if field in base_df]
        rename_map = {
            map_entry_k: map_entry_v
            for map_entry_k, map_entry_v in rename_map.items()
            if map_entry_k in base_df_fields
        }
        llm_df = base_df[base_df_fields].dropna(axis=1)
        if rename_map:


if rename_map:
            llm_df = llm_df.rename(rename_map, axis=1)
        return llm_df
    def _create_session_analysis_df(self) -> Any:
        """Create a dataframe with all the information from the session."""
        pd = import_pandas()
        on_llm_end_records_df = pd.DataFrame(self.on_llm_end_records)
        llm_input_prompts_df = ClearMLCallbackHandler._build_llm_df(
            base_df=on_llm_end_records_df,
            base_df_fields=["step", "prompts"]
            + (["name"] if "name" in on_llm_end_records_df else ["id"]),
            rename_map={"step": "prompt_step"},
        )
        complexity_metrics_columns = []
        visualizations_columns: List = []
        if self.complexity_metrics:
            complexity_metrics_columns = [
                "flesch_reading_ease",
                "flesch_kincaid_grade",
                "smog_index",
                "coleman_liau_index",
                "automated_readability_index",
                "dale_chall_readability_score",
                "difficult_words",
                "linsear_write_formula",
                "gunning_fog",
                "text_standard",
                "fernandez_huerta",
                "szigriszt_pazos",
                "gutierrez_polini",
                "crawford",
                "gulpease_index",
                "osman",
            ]
        llm_outputs_df = ClearMLCallbackHandler._build_llm_df(
            on_llm_end_records_df,
            [
                "step",
                "text",
                "token_usage_total_tokens",
                "token_usage_prompt_tokens",
                "token_usage_completion_tokens",


"token_usage_prompt_tokens",
                "token_usage_completion_tokens",
            ]
            + complexity_metrics_columns
            + visualizations_columns,
            {"step": "output_step", "text": "output"},
        )
        session_analysis_df = pd.concat([llm_input_prompts_df, llm_outputs_df], axis=1)
        return session_analysis_df
[docs]    def flush_tracker(
        self,
        name: Optional[str] = None,
        langchain_asset: Any = None,
        finish: bool = False,
    ) -> None:
        """Flush the tracker and setup the session.
        Everything after this will be a new table.
        Args:
            name: Name of the performed session so far so it is identifiable
            langchain_asset: The langchain asset to save.
            finish: Whether to finish the run.
            Returns:
                None
        """
        pd = import_pandas()
        clearml = import_clearml()
        # Log the action records
        self.logger.report_table(
            "Action Records", name, table_plot=pd.DataFrame(self.action_records)
        )
        # Session analysis
        session_analysis_df = self._create_session_analysis_df()
        self.logger.report_table(
            "Session Analysis", name, table_plot=session_analysis_df
        )
        if self.stream_logs:
            self.logger.report_text(
                {
                    "action_records": pd.DataFrame(self.action_records),
                    "session_analysis": session_analysis_df,
                }
            )
        if langchain_asset:
            langchain_asset_path = Path(self.temp_dir.name, "model.json")
            try:
                langchain_asset.save(langchain_asset_path)
                # Create output model and connect it to the task
                output_model = clearml.OutputModel(


output_model = clearml.OutputModel(
                    task=self.task, config_text=load_json(langchain_asset_path)
                )
                output_model.update_weights(
                    weights_filename=str(langchain_asset_path),
                    auto_delete_file=False,
                    target_filename=name,
                )
            except ValueError:
                langchain_asset.save_agent(langchain_asset_path)
                output_model = clearml.OutputModel(
                    task=self.task, config_text=load_json(langchain_asset_path)
                )
                output_model.update_weights(
                    weights_filename=str(langchain_asset_path),
                    auto_delete_file=False,
                    target_filename=name,
                )
            except NotImplementedError as e:
                print("Could not save model.")
                print(repr(e))
                pass
        # Cleanup after adding everything to ClearML
        self.task.flush(wait_for_uploads=True)
        self.temp_dir.cleanup()
        self.temp_dir = tempfile.TemporaryDirectory()
        self.reset_callback_meta()
        if finish:
            self.task.close()


Source code for langchain_community.callbacks.sagemaker_callback
import json
import os
import shutil
import tempfile
from copy import deepcopy
from typing import Any, Dict, List, Optional
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_community.callbacks.utils import (
    flatten_dict,
)
[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