Skip to content

Commit

Permalink
Instrument litellm (#163)
Browse files Browse the repository at this point in the history 
* Setting session to none if server does not return 200 for /sessions

* WIP. Changing override logic

* Working

* Tidying

* WIP

* Finished litellm

* Added back override with deprecation message. Docstring fixes

* Indent

* Updated minimum supported litellm version

* syntax
  • Loading branch information
@HowieG
HowieG authored Apr 26, 2024
1 parent 500b024 commit bdb7a41
Show file tree
Hide file tree
Showing 3 changed files with 183 additions and 81 deletions.
51 changes: 47 additions & 4 deletions agentops/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,8 @@ def init(api_key: Optional[str] = None,
max_wait_time: Optional[int] = None,
max_queue_size: Optional[int] = None,
tags: Optional[List[str]] = None,
override=True,
override: Optional[bool] = None, # Deprecated
instrument_llm_calls=True,
auto_start_session=True,
inherited_session_id: Optional[str] = None
):
Expand All @@ -37,11 +38,12 @@ def init(api_key: Optional[str] = None,
max_queue_size (int, optional): The maximum size of the event queue. Defaults to 100.
tags (List[str], optional): Tags for the sessions that can be used for grouping or
sorting later (e.g. ["GPT-4"]).
override (bool): Whether to override and LLM calls to emit as events.
override (bool, optional): [Deprecated] Use `instrument_llm_calls` instead. Whether to instrument LLM calls and emit LLMEvents..
instrument_llm_calls (bool): Whether to instrument LLM calls and emit LLMEvents..
auto_start_session (bool): Whether to start a session automatically when the client is created.
inherited_session_id (optional, str): Init Agentops with an existing Session
Attributes:
"""
"""
set_logging_level_info()
c = Client(api_key=api_key,
parent_key=parent_key,
Expand All @@ -50,32 +52,67 @@ def init(api_key: Optional[str] = None,
max_queue_size=max_queue_size,
tags=tags,
override=override,
instrument_llm_calls=instrument_llm_calls,
auto_start_session=auto_start_session,
inherited_session_id=inherited_session_id
)

return inherited_session_id or c.current_session_id


def end_session(end_state: str,
end_state_reason: Optional[str] = None,
video: Optional[str] = None):
"""
End the current session with the AgentOps service.
Args:
end_state (str): The final state of the session. Options: Success, Fail, or Indeterminate.
end_state_reason (str, optional): The reason for ending the session.
video (str, optional): URL to a video recording of the session
"""
Client().end_session(end_state, end_state_reason, video)


def start_session(tags: Optional[List[str]] = None, config: Optional[Configuration] = None, inherited_session_id: Optional[str] = None):
"""
Start a new session for recording events.
Args:
tags (List[str], optional): Tags that can be used for grouping or sorting later.
e.g. ["test_run"].
config: (Configuration, optional): Client configuration object
"""
return Client().start_session(tags, config, inherited_session_id)


def record(event: Event | ErrorEvent):
"""
Record an event with the AgentOps service.
Args:
event (Event): The event to record.
"""
Client().record(event)


def add_tags(tags: List[str]):
"""
Append to session tags at runtime.
Args:
tags (List[str]): The list of tags to append.
"""
Client().add_tags(tags)


def set_tags(tags: List[str]):
"""
Replace session tags at runtime.
Args:
tags (List[str]): The list of tags to set.
"""
Client().set_tags(tags)


Expand All @@ -84,4 +121,10 @@ def get_api_key() -> str:


def set_parent_key(parent_key):
"""
Set the parent API key which has visibility to projects it is parent to.
Args:
parent_key (str): The API key of the parent organization to set.
"""
Client().set_parent_key(parent_key)
136 changes: 80 additions & 56 deletions agentops/client.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
"""
AgentOps client module that provides a client class with public interfaces and configuration.
AgentOps client module that provides a client class with public interfaces and configuration.
Classes:
Client: Provides methods to interact with the AgentOps service.
Classes:
Client: Provides methods to interact with the AgentOps service.
"""

from .event import ActionEvent, ErrorEvent, Event
Expand All @@ -29,27 +29,28 @@
@singleton
class Client(metaclass=MetaClient):
"""
Client for AgentOps service.
Args:
api_key (str, optional): API Key for AgentOps services. If none is provided, key will
be read from the AGENTOPS_API_KEY environment variable.
parent_key (str, optional): Organization key to give visibility of all user sessions the user's organization. If none is provided, key will
be read from the AGENTOPS_PARENT_KEY environment variable.
endpoint (str, optional): The endpoint for the AgentOps service. If none is provided, key will
be read from the AGENTOPS_API_ENDPOINT environment variable. Defaults to 'https://api.agentops.ai'.
max_wait_time (int, optional): The maximum time to wait in milliseconds before flushing the queue.
Defaults to 30,000 (30 seconds)
max_queue_size (int, optional): The maximum size of the event queue. Defaults to 100.
tags (List[str], optional): Tags for the sessions that can be used for grouping or
sorting later (e.g. ["GPT-4"]).
override (bool): Whether to override and LLM calls to emit as events.
auto_start_session (bool): Whether to start a session automatically when the client is created.
inherited_session_id (optional, str): Init Agentops with an existing Session
Attributes:
_session (Session, optional): A Session is a grouping of events (e.g. a run of your agent).
_worker (Worker, optional): A Worker manages the event queue and sends session updates to the AgentOps api server
Client for AgentOps service.
Args:
api_key (str, optional): API Key for AgentOps services. If none is provided, key will
be read from the AGENTOPS_API_KEY environment variable.
parent_key (str, optional): Organization key to give visibility of all user sessions the user's organization. If none is provided, key will
be read from the AGENTOPS_PARENT_KEY environment variable.
endpoint (str, optional): The endpoint for the AgentOps service. If none is provided, key will
be read from the AGENTOPS_API_ENDPOINT environment variable. Defaults to 'https://api.agentops.ai'.
max_wait_time (int, optional): The maximum time to wait in milliseconds before flushing the queue.
Defaults to 30,000 (30 seconds)
max_queue_size (int, optional): The maximum size of the event queue. Defaults to 100.
tags (List[str], optional): Tags for the sessions that can be used for grouping or
sorting later (e.g. ["GPT-4"]).
override (bool, optional): [Deprecated] Use `instrument_llm_calls` instead. Whether to instrument LLM calls and emit LLMEvents..
instrument_llm_calls (bool): Whether to instrument LLM calls and emit LLMEvents..
auto_start_session (bool): Whether to start a session automatically when the client is created.
inherited_session_id (optional, str): Init Agentops with an existing Session
Attributes:
_session (Session, optional): A Session is a grouping of events (e.g. a run of your agent).
_worker (Worker, optional): A Worker manages the event queue and sends session updates to the AgentOps api server
"""

def __init__(self,
Expand All @@ -59,11 +60,17 @@ def __init__(self,
max_wait_time: Optional[int] = None,
max_queue_size: Optional[int] = None,
tags: Optional[List[str]] = None,
override=True,
override: Optional[bool] = None, # Deprecated
instrument_llm_calls=True,
auto_start_session=True,
inherited_session_id: Optional[str] = None
):

if override is not None:
logger.warning("🖇 AgentOps: The 'override' parameter is deprecated. Use 'instrument_llm_calls' instead.",
DeprecationWarning, stacklevel=2)
instrument_llm_calls = instrument_llm_calls or override

self._session = None
self._worker = None
self._tags = tags
Expand All @@ -82,12 +89,17 @@ def __init__(self,
if auto_start_session:
self.start_session(tags, self.config, inherited_session_id)

if override:
if 'openai' in sys.modules:
self.llm_tracker = LlmTracker(self)
self.llm_tracker.override_api('openai')
if instrument_llm_calls:
self.llm_tracker = LlmTracker(self)
self.llm_tracker.override_api()

def add_tags(self, tags: List[str]):
"""
Append to session tags at runtime.
Args:
tags (List[str]): The list of tags to append.
"""
if self._tags is not None:
self._tags.extend(tags)
else:
Expand All @@ -98,6 +110,12 @@ def add_tags(self, tags: List[str]):
self._worker.update_session(self._session)

def set_tags(self, tags: List[str]):
"""
Replace session tags at runtime.
Args:
tags (List[str]): The list of tags to set.
"""
self._tags = tags

if self._session is not None:
Expand All @@ -106,10 +124,10 @@ def set_tags(self, tags: List[str]):

def record(self, event: Event | ErrorEvent):
"""
Record an event with the AgentOps service.
Record an event with the AgentOps service.
Args:
event (Event): The event to record.
Args:
event (Event): The event to record.
"""

if self._session is not None and not self._session.has_ended:
Expand Down Expand Up @@ -198,16 +216,16 @@ async def _record_event_async(self, func, event_name, *args, **kwargs):

def start_session(self, tags: Optional[List[str]] = None, config: Optional[Configuration] = None, inherited_session_id: Optional[str] = None):
"""
Start a new session for recording events.
Start a new session for recording events.
Args:
tags (List[str], optional): Tags that can be used for grouping or sorting later.
e.g. ["test_run"].
config: (Configuration, optional): Client configuration object
inherited_session_id (optional, str): assign session id to match existing Session
Args:
tags (List[str], optional): Tags that can be used for grouping or sorting later.
e.g. ["test_run"].
config: (Configuration, optional): Client configuration object
inherited_session_id (optional, str): assign session id to match existing Session
"""
set_logging_level_info()

if self._session is not None:
return logger.warning("🖇 AgentOps: Cannot start session - session already started")

Expand All @@ -222,7 +240,7 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi
return logger.warning("🖇 AgentOps: Cannot start session")

logger.info('View info on this session at https://app.agentops.ai/drilldown?session_id={}'
.format(self._session.session_id))
.format(self._session.session_id))

return self._session.session_id

Expand All @@ -231,12 +249,12 @@ def end_session(self,
end_state_reason: Optional[str] = None,
video: Optional[str] = None):
"""
End the current session with the AgentOps service.
End the current session with the AgentOps service.
Args:
end_state (str): The final state of the session. Options: Success, Fail, or Indeterminate.
end_state_reason (str, optional): The reason for ending the session.
video (str, optional): The video screen recording of the session
Args:
end_state (str): The final state of the session. Options: Success, Fail, or Indeterminate.
end_state_reason (str, optional): The reason for ending the session.
video (str, optional): The video screen recording of the session
"""
if self._session is None or self._session.has_ended:
return logger.warning("🖇 AgentOps: Cannot end session - no current session")
Expand Down Expand Up @@ -267,11 +285,11 @@ def cleanup(end_state: Optional[str] = 'Fail', end_state_reason: Optional[str] =

def signal_handler(signum, frame):
"""
Signal handler for SIGINT (Ctrl+C) and SIGTERM. Ends the session and exits the program.
Signal handler for SIGINT (Ctrl+C) and SIGTERM. Ends the session and exits the program.
Args:
signum (int): The signal number.
frame: The current stack frame.
Args:
signum (int): The signal number.
frame: The current stack frame.
"""
signal_name = 'SIGINT' if signum == signal.SIGINT else 'SIGTERM'
logger.info(
Expand All @@ -282,13 +300,13 @@ def signal_handler(signum, frame):

def handle_exception(exc_type, exc_value, exc_traceback):
"""
Handle uncaught exceptions before they result in program termination.
Handle uncaught exceptions before they result in program termination.
Args:
exc_type (Type[BaseException]): The type of the exception.
exc_value (BaseException): The exception instance.
exc_traceback (TracebackType): A traceback object encapsulating the call stack at the
point where the exception originally occurred.
Args:
exc_type (Type[BaseException]): The type of the exception.
exc_value (BaseException): The exception instance.
exc_traceback (TracebackType): A traceback object encapsulating the call stack at the
point where the exception originally occurred.
"""
formatted_traceback = ''.join(traceback.format_exception(exc_type, exc_value,
exc_traceback))
Expand All @@ -315,7 +333,13 @@ def current_session_id(self):
def api_key(self):
return self.config.api_key

def set_parent_key(self, parent_key):
def set_parent_key(self, parent_key: str):
"""
Set the parent API key which has visibility to projects it is parent to.
Args:
parent_key (str): The API key of the parent organization to set.
"""
if self._worker:
self._worker.config.parent_key = parent_key

Expand Down
Loading

0 comments on commit bdb7a41

Please sign in to comment.