mirror of
https://github.com/hwchase17/langchain
synced 2024-11-08 07:10:35 +00:00
Fix for KeyError in MlflowCallbackHandler (#7051)
- Description: `MlflowCallbackHandler` fails with `KeyError: "['name']
not in index"`. See https://github.com/hwchase17/langchain/issues/5770
for more details. Root cause is that LangChain does not pass "name" as a
part of `serialized` argument to `on_llm_start()` callback method. The
commit where this change was made is probably this:
18af149e91
.
My bug fix derives "name" from "id" field.
- Issue: https://github.com/hwchase17/langchain/issues/5770
---------
Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
parent
df746ad821
commit
2179d4eef8
@ -551,8 +551,18 @@ class MlflowCallbackHandler(BaseMetadataCallbackHandler, BaseCallbackHandler):
|
|||||||
on_llm_start_records_df = pd.DataFrame(self.records["on_llm_start_records"])
|
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"])
|
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_prompts_df = (
|
||||||
on_llm_start_records_df[["step", "prompt", "name"]]
|
on_llm_start_records_df[llm_input_columns]
|
||||||
.dropna(axis=1)
|
.dropna(axis=1)
|
||||||
.rename({"step": "prompt_step"}, axis=1)
|
.rename({"step": "prompt_step"}, axis=1)
|
||||||
)
|
)
|
||||||
|
Loading…
Reference in New Issue
Block a user