From 5a39aaadd5c07a0b741f40e7f0443f303234a30e Mon Sep 17 00:00:00 2001 From: Braelyn Boynton Date: Thu, 18 Apr 2024 02:10:39 -0400 Subject: [PATCH] inherit session ids (#154) * inherited session ids * inherited session ids * return session id in start_session --- agentops/__init__.py | 29 ++++++++++++++++++----------- agentops/client.py | 15 ++++++++++----- tests/test_session.py | 27 +++++++++++++++++++++++++-- 3 files changed, 53 insertions(+), 18 deletions(-) diff --git a/agentops/__init__.py b/agentops/__init__.py index ae4c1e58..232b3e90 100755 --- a/agentops/__init__.py +++ b/agentops/__init__.py @@ -17,7 +17,9 @@ def init(api_key: Optional[str] = None, max_queue_size: Optional[int] = None, tags: Optional[List[str]] = None, override=True, - auto_start_session=True): + auto_start_session=True, + inherited_session_id: Optional[str] = None + ): """ Initializes the AgentOps singleton pattern. @@ -36,17 +38,22 @@ def init(api_key: Optional[str] = None, 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: """ - Client(api_key=api_key, - parent_key=parent_key, - endpoint=endpoint, - max_wait_time=max_wait_time, - max_queue_size=max_queue_size, - tags=tags, - override=override, - auto_start_session=auto_start_session) + c = Client(api_key=api_key, + parent_key=parent_key, + endpoint=endpoint, + max_wait_time=max_wait_time, + max_queue_size=max_queue_size, + tags=tags, + override=override, + 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, @@ -55,8 +62,8 @@ def end_session(end_state: str, Client().end_session(end_state, end_state_reason, video) -def start_session(tags: Optional[List[str]] = None, config: Optional[Configuration] = None): - Client().start_session(tags, config) +def start_session(tags: Optional[List[str]] = None, config: Optional[Configuration] = None, inherited_session_id: Optional[str] = None): + return Client().start_session(tags, config, inherited_session_id) def record(event: Event | ErrorEvent): diff --git a/agentops/client.py b/agentops/client.py index 69998765..a4eaedea 100644 --- a/agentops/client.py +++ b/agentops/client.py @@ -46,6 +46,7 @@ class Client(metaclass=MetaClient): 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 @@ -59,7 +60,8 @@ def __init__(self, max_queue_size: Optional[int] = None, tags: Optional[List[str]] = None, override=True, - auto_start_session=True + auto_start_session=True, + inherited_session_id: Optional[str] = None ): self._session = None @@ -78,7 +80,7 @@ def __init__(self, self._handle_unclean_exits() if auto_start_session: - self.start_session(tags, self.config) + self.start_session(tags, self.config, inherited_session_id) if override: if 'openai' in sys.modules: @@ -196,7 +198,7 @@ async def _record_event_async(self, func, event_name, *args, **kwargs): return returns - def start_session(self, tags: Optional[List[str]] = None, config: Optional[Configuration] = None): + 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. @@ -204,6 +206,7 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi 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 """ if self._session is not None: return logging.warning("🖇 AgentOps: Cannot start session - session already started") @@ -211,7 +214,7 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi if not config and not self.config: return logging.warning("🖇 AgentOps: Cannot start session - missing configuration") - self._session = Session(uuid4(), tags or self._tags, host_env=get_host_env()) + self._session = Session(inherited_session_id or uuid4(), tags or self._tags, host_env=get_host_env()) self._worker = Worker(config or self.config) start_session_result = self._worker.start_session(self._session) if not start_session_result: @@ -221,6 +224,8 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi logging.info('View info on this session at https://app.agentops.ai/drilldown?session_id={}' .format(self._session.session_id)) + return self._session.session_id + def end_session(self, end_state: str, end_state_reason: Optional[str] = None, @@ -304,7 +309,7 @@ def handle_exception(exc_type, exc_value, exc_traceback): @property def current_session_id(self): - return self._session.session_id + return self._session.session_id if self._session else None @property def api_key(self): diff --git a/tests/test_session.py b/tests/test_session.py index 339835de..c553e7bd 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -69,9 +69,32 @@ def test_tags(self, mock_req): agentops.end_session(end_state) time.sleep(0.15) - # Assert 3 requets, 1 for session init, 1 for event, 1 for end session + # Assert 3 requests, 1 for session init, 1 for event, 1 for end session assert len(mock_req.request_history) == 3 assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key request_json = mock_req.last_request.json() assert request_json['session']['end_state'] == end_state - assert request_json['session']['tags'] == tags \ No newline at end of file + assert request_json['session']['tags'] == tags + + def test_inherit_session_id(self, mock_req): + # Arrange + inherited_id = '4f72e834-ff26-4802-ba2d-62e7613446f1' + agentops.start_session(tags=['test'], config=self.config, inherited_session_id=inherited_id) + + # Act + agentops.record(ActionEvent(self.event_type)) + agentops.record(ActionEvent(self.event_type)) + time.sleep(0.15) + + # event session_id correct + request_json = mock_req.last_request.json() + assert request_json['session_id'] == inherited_id + + # Act + end_state = 'Success' + agentops.end_session(end_state) + time.sleep(0.15) + + # Assert session ended with correct id + request_json = mock_req.last_request.json() + assert request_json['session']['session_id'] == inherited_id