[LLMonitor Callback Handler]: Add error handling (#11563)

Wraps every callback handler method in error handlers to avoid breaking
users' programs when an error occurs inside the handler.

Thanks @valdo99 for the suggestion 🙂
This commit is contained in:
Hugues Chocart 2023-10-09 22:26:35 +02:00 committed by GitHub
parent 2aabfafe1e
commit 258ae1ba5f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,3 +1,4 @@
import logging
import os import os
import traceback import traceback
from contextvars import ContextVar from contextvars import ContextVar
@ -242,24 +243,27 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
user_id = _get_user_id(metadata) try:
user_props = _get_user_props(metadata) user_id = _get_user_id(metadata)
user_props = _get_user_props(metadata)
event = { event = {
"event": "start", "event": "start",
"type": "llm", "type": "llm",
"userId": user_id, "userId": user_id,
"runId": str(run_id), "runId": str(run_id),
"parentRunId": str(parent_run_id) if parent_run_id else None, "parentRunId": str(parent_run_id) if parent_run_id else None,
"input": _parse_input(prompts), "input": _parse_input(prompts),
"name": kwargs.get("invocation_params", {}).get("model_name"), "name": kwargs.get("invocation_params", {}).get("model_name"),
"tags": tags, "tags": tags,
"metadata": metadata, "metadata": metadata,
} }
if user_props: if user_props:
event["userProps"] = user_props event["userProps"] = user_props
self.__send_event(event) self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_llm_start: {e}")
def on_chat_model_start( def on_chat_model_start(
self, self,
@ -272,24 +276,29 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
user_id = _get_user_id(metadata) try:
user_props = _get_user_props(metadata) user_id = _get_user_id(metadata)
user_props = _get_user_props(metadata)
event = { event = {
"event": "start", "event": "start",
"type": "llm", "type": "llm",
"userId": user_id, "userId": user_id,
"runId": str(run_id), "runId": str(run_id),
"parentRunId": str(parent_run_id) if parent_run_id else None, "parentRunId": str(parent_run_id) if parent_run_id else None,
"input": _parse_lc_messages(messages[0]), "input": _parse_lc_messages(messages[0]),
"name": kwargs.get("invocation_params", {}).get("model_name"), "name": kwargs.get("invocation_params", {}).get("model_name"),
"tags": tags, "tags": tags,
"metadata": metadata, "metadata": metadata,
} }
if user_props: if user_props:
event["userProps"] = user_props event["userProps"] = user_props
self.__send_event(event) self.__send_event(event)
except Exception as e:
logging.warning(
f"[LLMonitor] An error occurred in on_chat_model_start: " f"{e}"
)
def on_llm_end( def on_llm_end(
self, self,
@ -299,39 +308,42 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
token_usage = (response.llm_output or {}).get("token_usage", {}) try:
token_usage = (response.llm_output or {}).get("token_usage", {})
parsed_output = [ parsed_output = [
{ {
"text": generation.text, "text": generation.text,
"role": "ai", "role": "ai",
**( **(
{ {
"functionCall": generation.message.additional_kwargs[ "functionCall": generation.message.additional_kwargs[
"function_call" "function_call"
] ]
} }
if hasattr(generation, "message") if hasattr(generation, "message")
and hasattr(generation.message, "additional_kwargs") and hasattr(generation.message, "additional_kwargs")
and "function_call" in generation.message.additional_kwargs and "function_call" in generation.message.additional_kwargs
else {} else {}
), ),
}
for generation in response.generations[0]
]
event = {
"event": "end",
"type": "llm",
"runId": str(run_id),
"parent_run_id": str(parent_run_id) if parent_run_id else None,
"output": parsed_output,
"tokensUsage": {
"prompt": token_usage.get("prompt_tokens"),
"completion": token_usage.get("completion_tokens"),
},
} }
for generation in response.generations[0] self.__send_event(event)
] except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_llm_end: {e}")
event = {
"event": "end",
"type": "llm",
"runId": str(run_id),
"parent_run_id": str(parent_run_id) if parent_run_id else None,
"output": parsed_output,
"tokensUsage": {
"prompt": token_usage.get("prompt_tokens"),
"completion": token_usage.get("completion_tokens"),
},
}
self.__send_event(event)
def on_tool_start( def on_tool_start(
self, self,
@ -344,24 +356,27 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
user_id = _get_user_id(metadata) try:
user_props = _get_user_props(metadata) user_id = _get_user_id(metadata)
user_props = _get_user_props(metadata)
event = { event = {
"event": "start", "event": "start",
"type": "tool", "type": "tool",
"userId": user_id, "userId": user_id,
"runId": str(run_id), "runId": str(run_id),
"parentRunId": str(parent_run_id) if parent_run_id else None, "parentRunId": str(parent_run_id) if parent_run_id else None,
"name": serialized.get("name"), "name": serialized.get("name"),
"input": input_str, "input": input_str,
"tags": tags, "tags": tags,
"metadata": metadata, "metadata": metadata,
} }
if user_props: if user_props:
event["userProps"] = user_props event["userProps"] = user_props
self.__send_event(event) self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_tool_start: {e}")
def on_tool_end( def on_tool_end(
self, self,
@ -372,14 +387,17 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
tags: Union[List[str], None] = None, tags: Union[List[str], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
event = { try:
"event": "end", event = {
"type": "tool", "event": "end",
"runId": str(run_id), "type": "tool",
"parent_run_id": str(parent_run_id) if parent_run_id else None, "runId": str(run_id),
"output": output, "parent_run_id": str(parent_run_id) if parent_run_id else None,
} "output": output,
self.__send_event(event) }
self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_tool_end: {e}")
def on_chain_start( def on_chain_start(
self, self,
@ -392,41 +410,44 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
name = serialized.get("id", [None, None, None, None])[3] try:
type = "chain" name = serialized.get("id", [None, None, None, None])[3]
metadata = metadata or {}
agentName = metadata.get("agent_name")
if agentName is None:
agentName = metadata.get("agentName")
if agentName is not None:
type = "agent"
name = agentName
if name == "AgentExecutor" or name == "PlanAndExecute":
type = "agent"
if parent_run_id is not None:
type = "chain" type = "chain"
metadata = metadata or {}
user_id = _get_user_id(metadata) agentName = metadata.get("agent_name")
user_props = _get_user_props(metadata) if agentName is None:
agentName = metadata.get("agentName")
event = { if agentName is not None:
"event": "start", type = "agent"
"type": type, name = agentName
"userId": user_id, if name == "AgentExecutor" or name == "PlanAndExecute":
"runId": str(run_id), type = "agent"
"parentRunId": str(parent_run_id) if parent_run_id else None,
"input": _parse_input(inputs),
"tags": tags,
"metadata": metadata,
"name": name,
}
if user_props:
event["userProps"] = user_props
self.__send_event(event) if parent_run_id is not None:
type = "chain"
user_id = _get_user_id(metadata)
user_props = _get_user_props(metadata)
event = {
"event": "start",
"type": type,
"userId": user_id,
"runId": str(run_id),
"parentRunId": str(parent_run_id) if parent_run_id else None,
"input": _parse_input(inputs),
"tags": tags,
"metadata": metadata,
"name": name,
}
if user_props:
event["userProps"] = user_props
self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_chain_start: {e}")
def on_chain_end( def on_chain_end(
self, self,
@ -436,13 +457,16 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
event = { try:
"event": "end", event = {
"type": "chain", "event": "end",
"runId": str(run_id), "type": "chain",
"output": _parse_output(outputs), "runId": str(run_id),
} "output": _parse_output(outputs),
self.__send_event(event) }
self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_chain_end: {e}")
def on_agent_action( def on_agent_action(
self, self,
@ -452,15 +476,18 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
event = { try:
"event": "start", event = {
"type": "tool", "event": "start",
"runId": str(run_id), "type": "tool",
"parentRunId": str(parent_run_id) if parent_run_id else None, "runId": str(run_id),
"name": action.tool, "parentRunId": str(parent_run_id) if parent_run_id else None,
"input": _parse_input(action.tool_input), "name": action.tool,
} "input": _parse_input(action.tool_input),
self.__send_event(event) }
self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_agent_action: {e}")
def on_agent_finish( def on_agent_finish(
self, self,
@ -470,14 +497,17 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
event = { try:
"event": "end", event = {
"type": "agent", "event": "end",
"runId": str(run_id), "type": "agent",
"parentRunId": str(parent_run_id) if parent_run_id else None, "runId": str(run_id),
"output": _parse_output(finish.return_values), "parentRunId": str(parent_run_id) if parent_run_id else None,
} "output": _parse_output(finish.return_values),
self.__send_event(event) }
self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_agent_finish: {e}")
def on_chain_error( def on_chain_error(
self, self,
@ -487,14 +517,17 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
event = { try:
"event": "error", event = {
"type": "chain", "event": "error",
"runId": str(run_id), "type": "chain",
"parent_run_id": str(parent_run_id) if parent_run_id else None, "runId": str(run_id),
"error": {"message": str(error), "stack": traceback.format_exc()}, "parent_run_id": str(parent_run_id) if parent_run_id else None,
} "error": {"message": str(error), "stack": traceback.format_exc()},
self.__send_event(event) }
self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_chain_error: {e}")
def on_tool_error( def on_tool_error(
self, self,
@ -504,14 +537,17 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
event = { try:
"event": "error", event = {
"type": "tool", "event": "error",
"runId": str(run_id), "type": "tool",
"parent_run_id": str(parent_run_id) if parent_run_id else None, "runId": str(run_id),
"error": {"message": str(error), "stack": traceback.format_exc()}, "parent_run_id": str(parent_run_id) if parent_run_id else None,
} "error": {"message": str(error), "stack": traceback.format_exc()},
self.__send_event(event) }
self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_tool_error: {e}")
def on_llm_error( def on_llm_error(
self, self,
@ -521,14 +557,17 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
event = { try:
"event": "error", event = {
"type": "llm", "event": "error",
"runId": str(run_id), "type": "llm",
"parent_run_id": str(parent_run_id) if parent_run_id else None, "runId": str(run_id),
"error": {"message": str(error), "stack": traceback.format_exc()}, "parent_run_id": str(parent_run_id) if parent_run_id else None,
} "error": {"message": str(error), "stack": traceback.format_exc()},
self.__send_event(event) }
self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_llm_error: {e}")
__all__ = ["LLMonitorCallbackHandler", "identify"] __all__ = ["LLMonitorCallbackHandler", "identify"]