Tony20100/langchain_full huggingface dataset:
"on_chain_end_records": [],
            "on_tool_start_records": [],
            "on_tool_end_records": [],
            "on_text_records": [],
            "on_agent_finish_records": [],
            "on_agent_action_records": [],
            "on_retriever_start_records": [],
            "on_retriever_end_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."""
"""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,
                        textstat=self.textstat,
                    )
                )
                if "text_complexity_metrics" in generation_resp:
                    complexity_metrics: Dict[str, float] = generation_resp.pop(
complexity_metrics: Dict[str, float] = generation_resp.pop(
                        "text_complexity_metrics"
                    )
                    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}")
                if "dependency_tree" in generation_resp:
                    dependency_tree = generation_resp["dependency_tree"]
                    self.mlflg.html(
                        dependency_tree, "dep-" + hash_string(generation.text)
                    )
                if "entities" in generation_resp:
                    entities = generation_resp["entities"]
                    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"])
        if isinstance(inputs, dict):
if isinstance(inputs, dict):
            chain_input = ",".join([f"{k}={v}" for k, v in inputs.items()])
        elif isinstance(inputs, list):
            chain_input = ",".join([str(input) for input in inputs])
        else:
            chain_input = str(inputs)
        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: Union[Dict[str, Any], str, List[str]], **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] = {}
        if isinstance(outputs, dict):
            chain_output = ",".join([f"{k}={v}" for k, v in outputs.items()])
        elif isinstance(outputs, list):
            chain_output = ",".join(map(str, outputs))
        else:
            chain_output = str(outputs)
        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}")
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."""
        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: Any, **kwargs: Any) -> None:
        """Run when tool ends running."""
        output = str(output)
        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.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
[docs]    def on_text(self, text: str, **kwargs: Any) -> None:
        """
        Run when text is received.
        """
        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,
            }
        )
"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}")
[docs]    def on_retriever_start(
        self,
        serialized: Dict[str, Any],
        query: str,
        **kwargs: Any,
    ) -> Any:
        """Run when Retriever starts running."""
        self.metrics["step"] += 1
        self.metrics["retriever_starts"] += 1
        self.metrics["starts"] += 1
        retriever_starts = self.metrics["retriever_starts"]
        resp: Dict[str, Any] = {}
resp: Dict[str, Any] = {}
        resp.update({"action": "on_retriever_start", "query": query})
        resp.update(flatten_dict(serialized))
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_retriever_start_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"retriever_start_{retriever_starts}")
[docs]    def on_retriever_end(
        self,
        documents: Sequence[Document],
        **kwargs: Any,
    ) -> Any:
        """Run when Retriever ends running."""
        self.metrics["step"] += 1
        self.metrics["retriever_ends"] += 1
        self.metrics["ends"] += 1
        retriever_ends = self.metrics["retriever_ends"]
        resp: Dict[str, Any] = {}
        retriever_documents = [
            {
                "page_content": doc.page_content,
                "metadata": {
                    k: (
                        str(v)
                        if not isinstance(v, list)
                        else ",".join(str(x) for x in v)
                    )
                    for k, v in doc.metadata.items()
                },
            }
            for doc in documents
        ]
        resp.update({"action": "on_retriever_end", "documents": retriever_documents})
        resp.update(self.metrics)
        self.mlflg.metrics(self.metrics, step=self.metrics["step"])
        self.records["on_retriever_end_records"].append(resp)
        self.records["action_records"].append(resp)
        self.mlflg.jsonf(resp, f"retriever_end_{retriever_ends}")
[docs]    def on_retriever_error(self, error: BaseException, **kwargs: Any) -> Any:
        """Run when Retriever errors."""
        self.metrics["step"] += 1
        self.metrics["errors"] += 1
    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 = (
            on_llm_start_records_df[llm_input_columns]
            .dropna(axis=1)
            .rename({"step": "prompt_step"}, axis=1)
        )
        complexity_metrics_columns = (
            get_text_complexity_metrics() if self.textstat is not None else []
        )
        visualizations_columns = (
            ["dependency_tree", "entities"] if self.nlp is not None else []
        )
        token_usage_columns = [
            "token_usage_total_tokens",
            "token_usage_prompt_tokens",
"token_usage_total_tokens",
            "token_usage_prompt_tokens",
            "token_usage_completion_tokens",
        ]
        token_usage_columns = [
            x for x in token_usage_columns if x in on_llm_end_records_df.columns
        ]
        llm_outputs_df = (
            on_llm_end_records_df[
                [
                    "step",
                    "text",
                ]
                + token_usage_columns
                + 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(
            lambda row: construct_html_from_prompt_and_generation(
                row["prompt"], row["output"]
            ),
            axis=1,
        )
        return session_analysis_df
    def _contain_llm_records(self) -> bool:
        return bool(self.records["on_llm_start_records"])
[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"]))
        if self._contain_llm_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.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.")  # noqa: T201
                        traceback.print_exc()
                        pass
                    except NotImplementedError:
                        print("Could not save model.")  # noqa: T201
                        traceback.print_exc()
                        pass
                except NotImplementedError:
                    print("Could not save model.")  # noqa: T201
                    traceback.print_exc()
                    pass
        if finish:
            self.mlflg.finish_run()
            self._reset()
else:
                arthur = ArthurAI(
                    url=arthur_url, login=arthur_login, password=arthur_password
                )
        # get model from Arthur by the provided model ID
        try:
            arthur_model = arthur.get_model(model_id)
        except ResponseClientError:
            raise ValueError(
                f"Was unable to retrieve model with id {model_id} from Arthur."
                " Make sure the ID corresponds to a model that is currently"
                " registered with your Arthur account."
            )
        return cls(arthur_model)
[docs]    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """On LLM start, save the input prompts"""
        run_id = kwargs["run_id"]
        self.run_map[run_id]["input_texts"] = prompts
        self.run_map[run_id]["start_time"] = time()
[docs]    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """On LLM end, send data to Arthur."""
        try:
            import pytz
        except ImportError as e:
            raise ImportError(
                "Could not import pytz. Please install it with 'pip install pytz'."
            ) from e
        run_id = kwargs["run_id"]
        # get the run params from this run ID,
        # or raise an error if this run ID has no corresponding metadata in self.run_map
        try:
            run_map_data = self.run_map[run_id]
        except KeyError as e:
            raise KeyError(
                "This function has been called with a run_id"
                " that was never registered in on_llm_start()."
" that was never registered in on_llm_start()."
                " Restart and try running the LLM again"
            ) from e
        # mark the duration time between on_llm_start() and on_llm_end()
        time_from_start_to_end = time() - run_map_data["start_time"]
        # create inferences to log to Arthur
        inferences = []
        for i, generations in enumerate(response.generations):
            for generation in generations:
                inference = {
                    "partner_inference_id": str(uuid.uuid4()),
                    "inference_timestamp": datetime.now(tz=pytz.UTC),
                    self.input_attr: run_map_data["input_texts"][i],
                    self.output_attr: generation.text,
                }
                if generation.generation_info is not None:
                    # add finish reason to the inference
                    # if generation info contains a finish reason and
                    # if the ArthurModel was registered to monitor finish_reason
                    if (
                        FINISH_REASON in generation.generation_info
                        and FINISH_REASON in self.attr_names
                    ):
                        inference[FINISH_REASON] = generation.generation_info[
                            FINISH_REASON
                        ]
                    # add token likelihoods data to the inference if the ArthurModel
                    # was registered to monitor token likelihoods
                    logprobs_data = generation.generation_info["logprobs"]
                    if (
                        logprobs_data is not None
                        and self.token_likelihood_attr is not None
                    ):
                        logprobs = logprobs_data["top_logprobs"]
                        likelihoods = [
                            {k: np.exp(v) for k, v in logprobs[i].items()}
                            for i in range(len(logprobs))
                        ]
                        inference[self.token_likelihood_attr] = likelihoods
]
                        inference[self.token_likelihood_attr] = likelihoods
                # add token usage counts to the inference if the
                # ArthurModel was registered to monitor token usage
                if (
                    isinstance(response.llm_output, dict)
                    and TOKEN_USAGE in response.llm_output
                ):
                    token_usage = response.llm_output[TOKEN_USAGE]
                    if (
                        PROMPT_TOKENS in token_usage
                        and PROMPT_TOKENS in self.attr_names
                    ):
                        inference[PROMPT_TOKENS] = token_usage[PROMPT_TOKENS]
                    if (
                        COMPLETION_TOKENS in token_usage
                        and COMPLETION_TOKENS in self.attr_names
                    ):
                        inference[COMPLETION_TOKENS] = token_usage[COMPLETION_TOKENS]
                # add inference duration to the inference if the ArthurModel
                # was registered to monitor inference duration
                if DURATION in self.attr_names:
                    inference[DURATION] = time_from_start_to_end
                inferences.append(inference)
        # send inferences to arthur
        self.arthur_model.send_inferences(inferences)
[docs]    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """On chain start, do nothing."""
[docs]    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """On chain end, do nothing."""
[docs]    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when LLM outputs an error."""
[docs]    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """On new token, pass."""
"""On new token, pass."""
[docs]    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when LLM chain outputs an error."""
[docs]    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        """Do nothing when tool starts."""
[docs]    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
        """Do nothing when agent takes a specific action."""
[docs]    def on_tool_end(
        self,
        output: Any,
        observation_prefix: Optional[str] = None,
        llm_prefix: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """Do nothing when tool ends."""
[docs]    def on_tool_error(self, error: BaseException, **kwargs: Any) -> None:
        """Do nothing when tool outputs an error."""
[docs]    def on_text(self, text: str, **kwargs: Any) -> None:
        """Do nothing"""
[docs]    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
        """Do nothing"""