[LLMonitorCallbackHandler] Refactor + add llmonitor-py dependency (#11948)

We now require uses to have the pip package `llmonitor` installed. It
allows us to have cleaner code and avoid duplicates between our library
and our code in Langchain.
This commit is contained in:
Hugues Chocart 2023-10-19 08:54:10 +02:00 committed by GitHub
parent 77fc2f7644
commit 008c7df80d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,12 +1,14 @@
import importlib.metadata
import logging import logging
import os import os
import traceback import traceback
import warnings
from contextvars import ContextVar from contextvars import ContextVar
from datetime import datetime
from typing import Any, Dict, List, Literal, Union from typing import Any, Dict, List, Literal, Union
from uuid import UUID from uuid import UUID
import requests import requests
from packaging.version import parse
from langchain.callbacks.base import BaseCallbackHandler from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema.agent import AgentAction, AgentFinish from langchain.schema.agent import AgentAction, AgentFinish
@ -142,7 +144,7 @@ def _get_user_props(metadata: Any) -> Any:
return user_props_ctx.get() return user_props_ctx.get()
metadata = metadata or {} metadata = metadata or {}
return metadata.get("user_props") return metadata.get("user_props", None)
def _parse_lc_message(message: BaseMessage) -> Dict[str, Any]: def _parse_lc_message(message: BaseMessage) -> Dict[str, Any]:
@ -191,6 +193,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
__api_url: str __api_url: str
__app_id: str __app_id: str
__verbose: bool __verbose: bool
__llmonitor_version: str
__has_valid_config: bool
def __init__( def __init__(
self, self,
@ -200,37 +204,58 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
) -> None: ) -> None:
super().__init__() super().__init__()
self.__api_url = api_url or os.getenv("LLMONITOR_API_URL") or DEFAULT_API_URL self.__has_valid_config = True
try:
import llmonitor
self.__llmonitor_version = importlib.metadata.version("llmonitor")
self.__track_event = llmonitor.track_event
except ImportError:
warnings.warn(
"""[LLMonitor] To use the LLMonitor callback handler you need to
have the `llmonitor` Python package installed. Please install it
with `pip install llmonitor`"""
)
self.__has_valid_config = False
if parse(self.__llmonitor_version) < parse("0.0.20"):
warnings.warn(
f"""[LLMonitor] The installed `llmonitor` version is
{self.__llmonitor_version} but `LLMonitorCallbackHandler` requires
at least version 0.0.20 upgrade `llmonitor` with `pip install
--upgrade llmonitor`"""
)
self.__has_valid_config = False
self.__has_valid_config = True
self.__api_url = api_url or os.getenv("LLMONITOR_API_URL") or DEFAULT_API_URL
self.__verbose = verbose or bool(os.getenv("LLMONITOR_VERBOSE")) self.__verbose = verbose or bool(os.getenv("LLMONITOR_VERBOSE"))
_app_id = app_id or os.getenv("LLMONITOR_APP_ID") _app_id = app_id or os.getenv("LLMONITOR_APP_ID")
if _app_id is None: if _app_id is None:
raise ValueError( warnings.warn(
"""app_id must be provided either as an argument or as """[LLMonitor] app_id must be provided either as an argument or as
an environment variable""" an environment variable"""
) )
self.__app_id = _app_id self.__has_valid_config = False
else:
self.__app_id = _app_id
if self.__has_valid_config is False:
return None
try: try:
res = requests.get(f"{self.__api_url}/api/app/{self.__app_id}") res = requests.get(f"{self.__api_url}/api/app/{self.__app_id}")
if not res.ok: if not res.ok:
raise ConnectionError() raise ConnectionError()
except Exception as e: except Exception:
raise ConnectionError( warnings.warn(
f"Could not connect to the LLMonitor API at {self.__api_url}" f"""[LLMonitor] Could not connect to the LLMonitor API at
) from e {self.__api_url}"""
)
def __send_event(self, event: Dict[str, Any]) -> None:
headers = {"Content-Type": "application/json"}
event = {**event, "app": self.__app_id, "timestamp": str(datetime.utcnow())}
if self.__verbose:
print("llmonitor_callback", event)
data = {"events": event}
requests.post(headers=headers, url=f"{self.__api_url}/api/report", json=data)
def on_llm_start( def on_llm_start(
self, self,
@ -243,27 +268,28 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
if self.__has_valid_config is False:
return
try: 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)
name = kwargs.get("invocation_params", {}).get("model_name")
input = _parse_input(prompts)
event = { self.__track_event(
"event": "start", "llm",
"type": "llm", "start",
"userId": user_id, user_id=user_id,
"runId": str(run_id), run_id=str(run_id),
"parentRunId": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"input": _parse_input(prompts), name=name,
"name": kwargs.get("invocation_params", {}).get("model_name"), input=input,
"tags": tags, tags=tags,
"metadata": metadata, metadata=metadata,
} user_props=user_props,
if user_props: )
event["userProps"] = user_props
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_llm_start: {e}") warnings.warn(f"[LLMonitor] An error occurred in on_llm_start: {e}")
def on_chat_model_start( def on_chat_model_start(
self, self,
@ -276,28 +302,29 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if self.__has_valid_config is False:
return
try: 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)
name = kwargs.get("invocation_params", {}).get("model_name")
input = _parse_lc_messages(messages[0])
event = { self.__track_event(
"event": "start", "llm",
"type": "llm", "start",
"userId": user_id, user_id=user_id,
"runId": str(run_id), run_id=str(run_id),
"parentRunId": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"input": _parse_lc_messages(messages[0]), name=name,
"name": kwargs.get("invocation_params", {}).get("model_name"), input=input,
"tags": tags, tags=tags,
"metadata": metadata, metadata=metadata,
} user_props=user_props,
if user_props: )
event["userProps"] = user_props
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning( logging.warning(
f"[LLMonitor] An error occurred in on_chat_model_start: " f"{e}" f"[LLMonitor] An error occurred in on_chat_model_start: {e}"
) )
def on_llm_end( def on_llm_end(
@ -308,9 +335,11 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
if self.__has_valid_config is False:
return
try: try:
token_usage = (response.llm_output or {}).get("token_usage", {}) token_usage = (response.llm_output or {}).get("token_usage", {})
parsed_output = [ parsed_output = [
{ {
"text": generation.text, "text": generation.text,
@ -330,20 +359,19 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
for generation in response.generations[0] for generation in response.generations[0]
] ]
event = { self.__track_event(
"event": "end", "llm",
"type": "llm", "end",
"runId": str(run_id), run_id=str(run_id),
"parent_run_id": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"output": parsed_output, output=parsed_output,
"tokensUsage": { token_usage={
"prompt": token_usage.get("prompt_tokens"), "prompt": token_usage.get("prompt_tokens"),
"completion": token_usage.get("completion_tokens"), "completion": token_usage.get("completion_tokens"),
}, },
} )
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_llm_end: {e}") warnings.warn(f"[LLMonitor] An error occurred in on_llm_end: {e}")
def on_tool_start( def on_tool_start(
self, self,
@ -356,27 +384,27 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
if self.__has_valid_config is False:
return
try: 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)
name = serialized.get("name")
event = { self.__track_event(
"event": "start", "tool",
"type": "tool", "start",
"userId": user_id, user_id=user_id,
"runId": str(run_id), run_id=str(run_id),
"parentRunId": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"name": serialized.get("name"), name=name,
"input": input_str, input=input_str,
"tags": tags, tags=tags,
"metadata": metadata, metadata=metadata,
} user_props=user_props,
if user_props: )
event["userProps"] = user_props
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_tool_start: {e}") warnings.warn(f"[LLMonitor] An error occurred in on_tool_start: {e}")
def on_tool_end( def on_tool_end(
self, self,
@ -387,17 +415,18 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
tags: Union[List[str], None] = None, tags: Union[List[str], None] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
if self.__has_valid_config is False:
return
try: try:
event = { self.__track_event(
"event": "end", "tool",
"type": "tool", "end",
"runId": str(run_id), run_id=str(run_id),
"parent_run_id": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"output": output, output=output,
} )
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_tool_end: {e}") warnings.warn(f"[LLMonitor] An error occurred in on_tool_end: {e}")
def on_chain_start( def on_chain_start(
self, self,
@ -410,6 +439,8 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
metadata: Union[Dict[str, Any], None] = None, metadata: Union[Dict[str, Any], None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if self.__has_valid_config is False:
return
try: try:
name = serialized.get("id", [None, None, None, None])[3] name = serialized.get("id", [None, None, None, None])[3]
type = "chain" type = "chain"
@ -419,35 +450,32 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
if agentName is None: if agentName is None:
agentName = metadata.get("agentName") agentName = metadata.get("agentName")
if name == "AgentExecutor" or name == "PlanAndExecute":
type = "agent"
if agentName is not None: if agentName is not None:
type = "agent" type = "agent"
name = agentName name = agentName
if name == "AgentExecutor" or name == "PlanAndExecute":
type = "agent"
if parent_run_id is not None: if parent_run_id is not None:
type = "chain" type = "chain"
user_id = _get_user_id(metadata) user_id = _get_user_id(metadata)
user_props = _get_user_props(metadata) user_props = _get_user_props(metadata)
input = _parse_input(inputs)
event = { self.__track_event(
"event": "start", type,
"type": type, "start",
"userId": user_id, user_id=user_id,
"runId": str(run_id), run_id=str(run_id),
"parentRunId": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"input": _parse_input(inputs), name=name,
"tags": tags, input=input,
"metadata": metadata, tags=tags,
"name": name, metadata=metadata,
} user_props=user_props,
if user_props: )
event["userProps"] = user_props
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_chain_start: {e}") warnings.warn(f"[LLMonitor] An error occurred in on_chain_start: {e}")
def on_chain_end( def on_chain_end(
self, self,
@ -457,14 +485,18 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if self.__has_valid_config is False:
return
try: try:
event = { output = _parse_output(outputs)
"event": "end",
"type": "chain", self.__track_event(
"runId": str(run_id), "chain",
"output": _parse_output(outputs), "end",
} run_id=str(run_id),
self.__send_event(event) parent_run_id=str(parent_run_id) if parent_run_id else None,
output=output,
)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_chain_end: {e}") logging.warning(f"[LLMonitor] An error occurred in on_chain_end: {e}")
@ -476,16 +508,20 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if self.__has_valid_config is False:
return
try: try:
event = { name = action.tool
"event": "start", input = _parse_input(action.tool_input)
"type": "tool",
"runId": str(run_id), self.__track_event(
"parentRunId": str(parent_run_id) if parent_run_id else None, "tool",
"name": action.tool, "start",
"input": _parse_input(action.tool_input), run_id=str(run_id),
} parent_run_id=str(parent_run_id) if parent_run_id else None,
self.__send_event(event) name=name,
input=input,
)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_agent_action: {e}") logging.warning(f"[LLMonitor] An error occurred in on_agent_action: {e}")
@ -497,15 +533,18 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if self.__has_valid_config is False:
return
try: try:
event = { output = _parse_output(finish.return_values)
"event": "end",
"type": "agent", self.__track_event(
"runId": str(run_id), "agent",
"parentRunId": str(parent_run_id) if parent_run_id else None, "end",
"output": _parse_output(finish.return_values), run_id=str(run_id),
} parent_run_id=str(parent_run_id) if parent_run_id else None,
self.__send_event(event) output=output,
)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_agent_finish: {e}") logging.warning(f"[LLMonitor] An error occurred in on_agent_finish: {e}")
@ -517,15 +556,16 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if self.__has_valid_config is False:
return
try: try:
event = { self.__track_event(
"event": "error", "chain",
"type": "chain", "error",
"runId": str(run_id), run_id=str(run_id),
"parent_run_id": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"error": {"message": str(error), "stack": traceback.format_exc()}, error={"message": str(error), "stack": traceback.format_exc()},
} )
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_chain_error: {e}") logging.warning(f"[LLMonitor] An error occurred in on_chain_error: {e}")
@ -537,15 +577,16 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if self.__has_valid_config is False:
return
try: try:
event = { self.__track_event(
"event": "error", "tool",
"type": "tool", "error",
"runId": str(run_id), run_id=str(run_id),
"parent_run_id": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"error": {"message": str(error), "stack": traceback.format_exc()}, error={"message": str(error), "stack": traceback.format_exc()},
} )
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_tool_error: {e}") logging.warning(f"[LLMonitor] An error occurred in on_tool_error: {e}")
@ -557,15 +598,16 @@ class LLMonitorCallbackHandler(BaseCallbackHandler):
parent_run_id: Union[UUID, None] = None, parent_run_id: Union[UUID, None] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
if self.__has_valid_config is False:
return
try: try:
event = { self.__track_event(
"event": "error", "llm",
"type": "llm", "error",
"runId": str(run_id), run_id=str(run_id),
"parent_run_id": str(parent_run_id) if parent_run_id else None, parent_run_id=str(parent_run_id) if parent_run_id else None,
"error": {"message": str(error), "stack": traceback.format_exc()}, error={"message": str(error), "stack": traceback.format_exc()},
} )
self.__send_event(event)
except Exception as e: except Exception as e:
logging.warning(f"[LLMonitor] An error occurred in on_llm_error: {e}") logging.warning(f"[LLMonitor] An error occurred in on_llm_error: {e}")