Skip to content

Commit eb072db

Browse files
committed
refactor: reorganize session management into dedicated components
Split session logic into dedicated components for better separation of concerns and maintainability: - Move session code to dedicated session/ module - Split Session class into: - Session: Data container with minimal public API - SessionManager: Handles lifecycle and state management - SessionApi: Handles API communication - SessionTelemetry: Manages event recording and OTEL integration Key fixes: - Proper UUID and timestamp serialization in events - Consistent API key header handling - Correct token cost formatting in analytics - Proper session ID inheritance - Tags conversion and validation - Event counts type handling This refactor improves code organization while maintaining backward compatibility through the session/__init__.py module. Signed-off-by: Teo <[email protected]>
1 parent 8f92278 commit eb072db

File tree

10 files changed

+664
-697
lines changed

10 files changed

+664
-697
lines changed

agentops/client.py

Lines changed: 16 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -198,47 +198,31 @@ def start_session(
198198
self,
199199
tags: Optional[List[str]] = None,
200200
inherited_session_id: Optional[str] = None,
201-
) -> Union[Session, None]:
202-
"""
203-
Start a new session for recording events.
204-
205-
Args:
206-
tags (List[str], optional): Tags that can be used for grouping or sorting later.
207-
e.g. ["test_run"].
208-
config: (Configuration, optional): Client configuration object
209-
inherited_session_id (optional, str): assign session id to match existing Session
210-
"""
201+
) -> Optional[Session]:
202+
"""Start a new session"""
211203
if not self.is_initialized:
212-
return
213-
214-
if inherited_session_id is not None:
215-
try:
216-
session_id = UUID(inherited_session_id)
217-
except ValueError:
218-
return logger.warning(f"Invalid session id: {inherited_session_id}")
219-
else:
220-
session_id = uuid4()
204+
return None
221205

222-
session_tags = self._config.default_tags.copy()
223-
if tags is not None:
224-
session_tags.update(tags)
206+
try:
207+
session_id = UUID(inherited_session_id) if inherited_session_id else uuid4()
208+
except ValueError:
209+
return logger.warning(f"Invalid session id: {inherited_session_id}")
225210

211+
default_tags = list(self._config.default_tags) if self._config.default_tags else []
226212
session = Session(
227213
session_id=session_id,
228-
tags=list(session_tags),
229-
host_env=self.host_env,
230214
config=self._config,
215+
tags=tags or default_tags,
216+
host_env=self.host_env,
231217
)
232218

233-
if not session.is_running:
234-
return logger.error("Failed to start session")
235-
236-
if self._pre_init_queue["agents"] and len(self._pre_init_queue["agents"]) > 0:
237-
for agent_args in self._pre_init_queue["agents"]:
238-
session.create_agent(name=agent_args["name"], agent_id=agent_args["agent_id"])
239-
self._pre_init_queue["agents"] = []
219+
if session.is_running:
220+
# Process any queued agents
221+
if self._pre_init_queue["agents"]:
222+
for agent_args in self._pre_init_queue["agents"]:
223+
session.create_agent(name=agent_args["name"], agent_id=agent_args["agent_id"])
224+
self._pre_init_queue["agents"].clear()
240225

241-
self._sessions.append(session)
242226
return session
243227

244228
def end_session(

agentops/event.py

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ class Event:
2525
end_timestamp(str): A timestamp indicating when the event ended. Defaults to the time when this Event was instantiated.
2626
agent_id(UUID, optional): The unique identifier of the agent that triggered the event.
2727
id(UUID): A unique identifier for the event. Defaults to a new UUID.
28+
session_id(UUID, optional): The unique identifier of the session that the event belongs to.
2829
2930
foo(x=1) {
3031
...
@@ -43,6 +44,7 @@ class Event:
4344
end_timestamp: Optional[str] = None
4445
agent_id: Optional[UUID] = field(default_factory=check_call_stack_for_agent_id)
4546
id: UUID = field(default_factory=uuid4)
47+
session_id: Optional[UUID] = None
4648

4749

4850
@dataclass
@@ -105,7 +107,7 @@ class ToolEvent(Event):
105107

106108

107109
@dataclass
108-
class ErrorEvent:
110+
class ErrorEvent(Event):
109111
"""
110112
For recording any errors e.g. ones related to agent execution
111113
@@ -115,21 +117,31 @@ class ErrorEvent:
115117
code(str, optional): A code that can be used to identify the error e.g. 501.
116118
details(str, optional): Detailed information about the error.
117119
logs(str, optional): For detailed information/logging related to the error.
118-
timestamp(str): A timestamp indicating when the error occurred. Defaults to the time when this ErrorEvent was instantiated.
119-
120120
"""
121121

122+
# Inherit common Event fields
123+
event_type: str = field(default=EventType.ERROR.value)
124+
125+
# Error-specific fields
122126
trigger_event: Optional[Event] = None
123127
exception: Optional[BaseException] = None
124128
error_type: Optional[str] = None
125129
code: Optional[str] = None
126130
details: Optional[Union[str, Dict[str, str]]] = None
127131
logs: Optional[str] = field(default_factory=traceback.format_exc)
128-
timestamp: str = field(default_factory=get_ISO_time)
129132

130133
def __post_init__(self):
131-
self.event_type = EventType.ERROR.value
134+
"""Process exception if provided"""
132135
if self.exception:
133136
self.error_type = self.error_type or type(self.exception).__name__
134137
self.details = self.details or str(self.exception)
135138
self.exception = None # removes exception from serialization
139+
140+
# Ensure end timestamp is set
141+
if not self.end_timestamp:
142+
self.end_timestamp = get_ISO_time()
143+
144+
@property
145+
def timestamp(self) -> str:
146+
"""Maintain backward compatibility with old code expecting timestamp"""
147+
return self.init_timestamp

0 commit comments

Comments
 (0)