"""Activity and session operations for AWM workflows."""
import json
from collections.abc import Generator
from datetime import datetime
from typing import Any
import chariot_api._openapi.awm.models as api_models
from chariot import _apis
from chariot.awm import _utils, models
__all__ = [
"send_message",
"send_clear_context",
"send_tool_call_approvals",
"get_activities",
"get_sessions",
]
def _require_data(response: Any, operation: str) -> Any:
if response.data is None:
raise RuntimeError(f"Received malformed response (missing `data`) from {operation}")
return response.data
def _require_raw_data(response: Any, operation: str) -> Any:
payload = _require_raw_payload(response, operation)
return payload.get("data")
def _require_raw_payload(response: Any, operation: str) -> dict[str, Any]:
payload_bytes = getattr(response, "data", None)
if payload_bytes is None and hasattr(response, "read"):
payload_bytes = response.read()
if payload_bytes is None:
raise RuntimeError(f"Received malformed response (missing response body) from {operation}")
payload_text = (
payload_bytes.decode("utf-8") if isinstance(payload_bytes, bytes) else str(payload_bytes)
)
payload = json.loads(payload_text)
if not isinstance(payload, dict):
raise RuntimeError(f"Received malformed response (invalid JSON object) from {operation}")
return payload
def _send_activity_raw(workflow_id: str, payload: dict[str, Any]) -> Any:
# `send_activity` drops the raw activity payload in this generated client. We serialize and
# execute the request directly so we can parse and return the created activity reliably.
request = _apis.awm.activity_api._send_activity_serialize(
workflow_id=workflow_id,
body=payload,
_request_auth=None,
_content_type=None,
_headers=None,
_host_index=0,
)
response_data = _apis.awm.activity_api.api_client.call_api(*request)
return response_data.response
def _to_activity(data: dict[str, Any]) -> models.Activity:
return models.Activity(
id=data["id"],
workflow_id=data["workflowID"],
type=data["type"],
body=data.get("body"),
session_id=data.get("sessionID", ""),
final=bool(data.get("final", False)),
created_at=data["createdAt"],
from_agent=data.get("fromAgent", ""),
to_agent=data.get("toAgent", ""),
user_id=data.get("userID", ""),
processed_at=data.get("processedAt"),
)
def _to_session(data: api_models.ModelSessionDetails) -> models.SessionDetails:
raw = data.model_dump()
return models.SessionDetails(
session_id=raw["session_id"],
status=raw["status"],
agents=raw.get("agents") or [],
user_ids=raw.get("user_ids") or [],
total_activity_count=int(raw.get("total_activity_count", 0)),
user_activity_count=int(raw.get("user_activity_count", 0)),
tool_activity_count=int(raw.get("tool_activity_count", 0)),
created_at=raw["created_at"],
updated_at=raw["updated_at"],
schedule_id=raw.get("schedule_id"),
total_request_duration_micros=raw.get("total_request_duration_micros"),
)
[docs]
@_apis.login_required
def send_message(
workflow_id: str,
agent_name: str,
body: str,
session_id: str | None = None,
) -> models.Activity:
"""Send a user-to-agent message activity to a workflow.
:param workflow_id: Workflow id
:param agent_name: Target agent name
:param body: Message text
:param session_id: Existing session ID; a new KSUID is generated server-side if omitted
:rtype: models.Activity
"""
payload: dict[str, Any] = {
"type": "user_to_agent_message",
"toAgent": agent_name,
"body": body,
}
if session_id:
payload["sessionID"] = session_id
response = _send_activity_raw(workflow_id=workflow_id, payload=payload)
raw_item = _require_raw_data(response, "send_message")
if not isinstance(raw_item, dict):
raise RuntimeError(
"Received malformed response (`data` should be an object) from send_message"
)
return _to_activity(raw_item)
[docs]
@_apis.login_required
def send_clear_context(workflow_id: str, agent_name: str, session_id: str) -> models.Activity:
"""Send a clear_context activity to reset an agent's context within a session."""
response = _send_activity_raw(
workflow_id=workflow_id,
payload={
"type": "clear_context",
"toAgent": agent_name,
"body": None,
"sessionID": session_id,
},
)
raw_item = _require_raw_data(response, "send_clear_context")
if not isinstance(raw_item, dict):
raise RuntimeError(
"Received malformed response (`data` should be an object) from send_clear_context"
)
return _to_activity(raw_item)
[docs]
@_apis.login_required
def get_activities(
workflow_id: str,
session_id: str | None = None,
types: list[str] | None = None,
from_time: datetime | None = None,
sort_order: str = "asc",
limit: int | None = None,
after_key: str | None = None,
) -> Generator[models.Activity, None, None]:
"""List activities for a workflow, optionally filtered by session and type.
:param workflow_id: Workflow id
:param session_id: Filter to a specific session
:param types: Activity types to include (e.g. ["agent_to_user_message"])
:param from_time: Only include activities created at or after this time
:param sort_order: "asc" or "desc"
:param limit: Maximum number of results to return
:param after_key: The pagination key to fetch more items
:rtype: Generator[models.Activity, None, None]
"""
next_after_key = after_key
yielded = 0
while True:
remaining = limit - yielded if limit is not None else None
if remaining is not None and remaining <= 0:
return
request_limit = (
min(remaining, _utils.PAGINATION_PAGE_SIZE_CAP)
if remaining is not None
else _utils.PAGINATION_PAGE_SIZE_CAP
)
response = _apis.awm.activity_api.list_activity_without_preload_content(
workflow_id=workflow_id,
session_id=[session_id] if session_id else None,
type=types,
var_from=from_time.isoformat() if from_time else None,
sort_order=sort_order,
limit=request_limit,
after_key=next_after_key,
)
payload = _require_raw_payload(response, "list_activities")
raw_items = payload.get("data")
if raw_items is None:
return
if not isinstance(raw_items, list):
raise RuntimeError(
"Received malformed response (`data` should be a list or null) from list_activities"
)
for item in raw_items:
yield _to_activity(item)
yielded += 1
if limit is not None and yielded >= limit:
return
next_after_key = payload.get("after_key") or payload.get("afterKey")
if len(raw_items) == 0 or not next_after_key:
return
[docs]
@_apis.login_required
def get_sessions(
workflow_id: str,
session_ids: list[str] | None = None,
user_ids: list[str] | None = None,
agents: list[str] | None = None,
from_time: datetime | None = None,
to_time: datetime | None = None,
limit: int | None = None,
sort_order: str = "desc",
after_key: str | None = None,
) -> Generator[models.SessionDetails, None, None]:
"""List sessions for a workflow.
:param workflow_id: Workflow id
:param session_ids: Filter to specific session IDs
:param user_ids: Filter to sessions for one or more users
:param agents: Filter to sessions containing one or more agents
:param from_time: Only include sessions created at or after this time
:param to_time: Only include sessions created before this time
:param limit: Maximum number of results to return
:param sort_order: "asc" or "desc"
:param after_key: The pagination key to fetch more items
:rtype: Generator[models.SessionDetails, None, None]
"""
params = {
k: v
for k, v in {
"session_id": session_ids,
"user_id": user_ids,
"agent": agents,
"var_from": from_time.isoformat() if from_time else None,
"to": to_time.isoformat() if to_time else None,
"limit": limit,
"sort_order": sort_order,
"after_key": after_key,
}.items()
if v is not None
}
def _paginate_fn(**params: Any) -> Any:
return _apis.awm.activity_api.list_sessions(
workflow_id=workflow_id,
**params,
)
return _utils.paginate(
func=_paginate_fn,
params=params,
transform_item=_to_session,
)