fix some typing

agentops/client.py

@@ -5,27 +5,24 @@
         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
 from .helpers import get_ISO_time, singleton, check_call_stack_for_agent_id, get_partner_frameworks
 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

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(