inherit session ids (#154)

* inherited session ids

* inherited session ids

* return session id in start_session

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):

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,22 +198,23 @@ 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.
 
         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
         """
         if self._session is not None:
             return logging.warning("🖇 AgentOps: Cannot start session - session already started")
 
         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):

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
+        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