Skip to content

Commit 6a8e7af

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 6a8e7af

File tree

10 files changed

+665
-698
lines changed

10 files changed

+665
-698
lines changed

agentops/client.py

Lines changed: 19 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -198,47 +198,34 @@ 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(
224+
name=agent_args["name"],
225+
agent_id=agent_args["agent_id"]
226+
)
227+
self._pre_init_queue["agents"].clear()
240228

241-
self._sessions.append(session)
242229
return session
243230

244231
def end_session(

agentops/event.py

Lines changed: 17 additions & 6 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,30 @@ 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
"""
121-
121+
# Inherit common Event fields
122+
event_type: str = field(default=EventType.ERROR.value)
123+
124+
# Error-specific fields
122125
trigger_event: Optional[Event] = None
123126
exception: Optional[BaseException] = None
124127
error_type: Optional[str] = None
125128
code: Optional[str] = None
126129
details: Optional[Union[str, Dict[str, str]]] = None
127130
logs: Optional[str] = field(default_factory=traceback.format_exc)
128-
timestamp: str = field(default_factory=get_ISO_time)
129131

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

0 commit comments

Comments
 (0)