[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,6 +243,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
try:
user_id = _get_user_id(metadata) user_id = _get_user_id(metadata)
user_props = _get_user_props(metadata) user_props = _get_user_props(metadata)
@ -260,6 +262,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
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,6 +276,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
try:
user_id = _get_user_id(metadata) user_id = _get_user_id(metadata)
user_props = _get_user_props(metadata) user_props = _get_user_props(metadata)
@ -290,6 +295,10 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
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,6 +308,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
try:
token_usage = (response.llm_output or {}).get("token_usage", {}) token_usage = (response.llm_output or {}).get("token_usage", {})
parsed_output = [ parsed_output = [
@ -332,6 +342,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
}, },
} }
self.__send_event(event) self.__send_event(event)
except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_llm_end: {e}")
def on_tool_start( def on_tool_start(
self, self,
@ -344,6 +356,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
try:
user_id = _get_user_id(metadata) user_id = _get_user_id(metadata)
user_props = _get_user_props(metadata) user_props = _get_user_props(metadata)
@ -362,6 +375,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
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,6 +387,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
tags: Union[List[str], None] = None, tags: Union[List[str], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
try:
event = { event = {
"event": "end", "event": "end",
"type": "tool", "type": "tool",
@ -380,6 +396,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
"output": output, "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,6 +410,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
try:
name = serialized.get("id", [None, None, None, None])[3] name = serialized.get("id", [None, None, None, None])[3]
type = "chain" type = "chain"
metadata = metadata or {} metadata = metadata or {}
@ -427,6 +446,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
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_chain_start: {e}")
def on_chain_end( def on_chain_end(
self, self,
@ -436,6 +457,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
try:
event = { event = {
"event": "end", "event": "end",
"type": "chain", "type": "chain",
@ -443,6 +465,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
"output": _parse_output(outputs), "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,6 +476,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
try:
event = { event = {
"event": "start", "event": "start",
"type": "tool", "type": "tool",
@ -461,6 +486,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
"input": _parse_input(action.tool_input), "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,6 +497,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
try:
event = { event = {
"event": "end", "event": "end",
"type": "agent", "type": "agent",
@ -478,6 +506,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
"output": _parse_output(finish.return_values), "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,6 +517,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
try:
event = { event = {
"event": "error", "event": "error",
"type": "chain", "type": "chain",
@ -495,6 +526,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
"error": {"message": str(error), "stack": traceback.format_exc()}, "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,6 +537,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
try:
event = { event = {
"event": "error", "event": "error",
"type": "tool", "type": "tool",
@ -512,6 +546,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
"error": {"message": str(error), "stack": traceback.format_exc()}, "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,6 +557,7 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
try:
event = { event = {
"event": "error", "event": "error",
"type": "llm", "type": "llm",
@ -529,6 +566,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
"error": {"message": str(error), "stack": traceback.format_exc()}, "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"]