From 8baef948a415ad6163ba385d443560a253e619f4 Mon Sep 17 00:00:00 2001
From: Shawn Qiu <siyangqiu@gmail.com>
Date: Fri, 17 May 2024 19:02:10 -0700
Subject: [PATCH] fix some typing

---
 agentops/client.py | 38 ++++++++++++++++++++++----------------
 agentops/event.py  |  2 +-
 2 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/agentops/client.py b/agentops/client.py
index a6ffabc4..2f2a3646 100644
--- a/agentops/client.py
+++ b/agentops/client.py
@@ -5,7 +5,16 @@
         Client: Provides methods to interact with the AgentOps service.
 """
 import os
-import uuid
+import inspect
+import atexit
+import signal
+import sys
+import threading
+import traceback
+import logging
+from decimal import Decimal
+from uuid import UUID, uuid4
+from typing import Optional, List, Union
 
 from .event import ActionEvent, ErrorEvent, Event
 from .enums import EndState
@@ -13,19 +22,7 @@
 from .session import Session
 from .worker import Worker
 from .host_env import get_host_env
-from uuid import uuid4
-from typing import Optional, List, Union
-import traceback
-import logging
 from .log_config import logger
-from decimal import Decimal
-import inspect
-import atexit
-import signal
-import sys
-import threading
-
-
 from .meta_client import MetaClient
 from .config import Configuration, ConfigurationError
 from .llm_tracker import LlmTracker
@@ -92,6 +89,11 @@ def __init__(self,
                                         endpoint=endpoint,
                                         max_wait_time=max_wait_time,
                                         max_queue_size=max_queue_size)
+
+            if inherited_session_id is not None:
+                # Check if inherited_session_id is valid
+                UUID(inherited_session_id)
+
         except ConfigurationError:
             return
 
@@ -289,8 +291,12 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi
         if not config and not self.config:
             return logger.warning("Cannot start session - missing configuration")
 
-        self._session = Session(inherited_session_id or uuid4(),
-                                tags or self._tags_for_future_session, host_env=get_host_env(self._env_data_opt_out))
+        session_id = UUID(
+            inherited_session_id) if inherited_session_id is not None else uuid4()
+
+        self._session = Session(session_id=session_id,
+                                tags=tags or self._tags_for_future_session,
+                                host_env=get_host_env(self._env_data_opt_out))
         self._worker = Worker(config or self.config)
         start_session_result = self._worker.start_session(self._session)
         if not start_session_result:
@@ -338,7 +344,7 @@ def end_session(self,
 
     def create_agent(self, name: str, agent_id: Optional[str] = None) -> str:
         if agent_id is None:
-            agent_id = str(uuid.uuid4())
+            agent_id = str(uuid4())
         if self._worker:
             self._worker.create_agent(name=name, agent_id=agent_id)
             return agent_id
diff --git a/agentops/event.py b/agentops/event.py
index 4f507692..203f5766 100644
--- a/agentops/event.py
+++ b/agentops/event.py
@@ -39,7 +39,7 @@ class Event:
 
     event_type: EventType
     params: Optional[dict] = None
-    returns: Optional[str] = None
+    returns: Optional[str | List[str]] = None
     init_timestamp: str = field(default_factory=get_ISO_time)
     end_timestamp: Optional[str] = None
     agent_id: Optional[UUID] = field(