Skip to content

Commit a2dcb70

Browse files
91pavan91pavan
authored
feat: emit real-time lifecycle events for LLM calls (#126)
* feat: emit real-time lifecycle events for LLM calls Signed-off-by: Pavan Sudheendra <pavan0591@gmail.com> * test: add missing tests Signed-off-by: Pavan Sudheendra <pavan0591@gmail.com> * refactor: ruff format Signed-off-by: Pavan Sudheendra <pavan0591@gmail.com> * feat: updates Signed-off-by: Pavan Sudheendra <pavan0591@gmail.com> * feat: bump version to 1.0.45 Signed-off-by: Pavan Sudheendra <pavan0591@gmail.com> * feat: pin ruff to a version Signed-off-by: Pavan Sudheendra <pavan0591@gmail.com> --------- Signed-off-by: Pavan Sudheendra <pavan0591@gmail.com>
1 parent ac5af6b commit a2dcb70

9 files changed

Lines changed: 398 additions & 6 deletions

File tree

‎.github/workflows/linter_and_formatter.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,11 @@ jobs:
1616

1717
- name: Check code has no lint errors
1818
uses: astral-sh/ruff-action@4919ec5cf1f49eff0871dbcea0da843445b837e6
19+
with:
20+
version: "0.15.14"
1921

2022
- name: Check code is correctly formatted
2123
uses: astral-sh/ruff-action@4919ec5cf1f49eff0871dbcea0da843445b837e6
2224
with:
2325
args: "format --check"
26+
version: "0.15.14"

‎ioa_observe/materializer/session_state.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,21 @@ class SessionToolState:
108108
last_started_at: datetime | None = None
109109
last_completed_at: datetime | None = None
110110
last_input: str | None = None
111+
last_output: str | None = None
112+
113+
114+
@dataclass
115+
class SessionLLMState:
116+
name: str
117+
active_count: int = 0
118+
started_count: int = 0
119+
completed_count: int = 0
120+
status: str = "idle"
121+
version: int = 0
122+
last_started_at: datetime | None = None
123+
last_completed_at: datetime | None = None
124+
last_input: str | None = None
125+
last_output: str | None = None
111126

112127

113128
@dataclass
@@ -121,6 +136,7 @@ class SessionState:
121136
nodes: dict[str, SessionNodeState] = field(default_factory=dict)
122137
edges: dict[str, SessionEdgeState] = field(default_factory=dict)
123138
tools: dict[str, SessionToolState] = field(default_factory=dict)
139+
llms: dict[str, SessionLLMState] = field(default_factory=dict)
124140

125141
def snapshot(self) -> dict[str, Any]:
126142
return {
@@ -142,6 +158,10 @@ def snapshot(self) -> dict[str, Any]:
142158
_serialize_state(tool)
143159
for tool in sorted(self.tools.values(), key=lambda item: item.name)
144160
],
161+
"llms": [
162+
_serialize_state(llm)
163+
for llm in sorted(self.llms.values(), key=lambda item: item.name)
164+
],
145165
}
146166

147167

@@ -205,6 +225,11 @@ def apply_event(
205225
RuntimeEventName.TOOL_COMPLETED,
206226
}:
207227
self._apply_tool_event(session, record, event_name)
228+
elif event_name in {
229+
RuntimeEventName.LLM_STARTED,
230+
RuntimeEventName.LLM_COMPLETED,
231+
}:
232+
self._apply_llm_event(session, record, event_name)
208233

209234
return session
210235

@@ -422,6 +447,55 @@ def _apply_tool_event(
422447
tool.completed_count += 1
423448
tool.status = "idle" if tool.active_count == 0 else "running"
424449
tool.last_completed_at = record.event_time
450+
tool_output = _optional_attribute(
451+
record, RuntimeEventAttribute.TOOL_OUTPUT.value
452+
)
453+
if tool_output is not None:
454+
tool.last_output = tool_output
455+
456+
session.status = "active"
457+
458+
def _apply_llm_event(
459+
self,
460+
session: SessionState,
461+
record: RuntimeEventRecord,
462+
event_name: RuntimeEventName,
463+
) -> None:
464+
llm_name = _required_attribute(record, RuntimeEventAttribute.LLM_NAME.value)
465+
llm = session.llms.get(llm_name)
466+
if llm is None:
467+
llm = SessionLLMState(name=llm_name)
468+
session.llms[llm_name] = llm
469+
470+
if record.snapshot_version and record.snapshot_version < llm.version:
471+
return
472+
llm.version = max(llm.version, record.snapshot_version)
473+
474+
if event_name == RuntimeEventName.LLM_STARTED:
475+
llm.active_count += 1
476+
llm.started_count += 1
477+
llm.status = "running"
478+
llm.last_started_at = record.event_time
479+
llm_input = _optional_attribute(
480+
record, RuntimeEventAttribute.LLM_INPUT.value
481+
)
482+
if llm_input is not None:
483+
llm.last_input = llm_input
484+
else:
485+
llm.active_count = max(0, llm.active_count - 1)
486+
llm.completed_count += 1
487+
llm.status = "idle" if llm.active_count == 0 else "running"
488+
llm.last_completed_at = record.event_time
489+
llm_input = _optional_attribute(
490+
record, RuntimeEventAttribute.LLM_INPUT.value
491+
)
492+
if llm_input is not None:
493+
llm.last_input = llm_input
494+
llm_output = _optional_attribute(
495+
record, RuntimeEventAttribute.LLM_OUTPUT.value
496+
)
497+
if llm_output is not None:
498+
llm.last_output = llm_output
425499

426500
session.status = "active"
427501

‎ioa_observe/sdk/decorators/base.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -625,12 +625,16 @@ def _cleanup_span(span, ctx_token):
625625

626626
tool_name = getattr(span, "_ioa_tool_name", None)
627627
if session_id and tool_name:
628+
tool_output = span.attributes.get(OBSERVE_ENTITY_OUTPUT)
628629
emit_runtime_event(
629630
build_runtime_event_attributes(
630631
RuntimeEventName.TOOL_COMPLETED,
631632
session_id=session_id,
632633
snapshot_version=next_session_event_version(session_id),
633-
**{RuntimeEventAttribute.TOOL_NAME.value: tool_name},
634+
**{
635+
RuntimeEventAttribute.TOOL_NAME.value: tool_name,
636+
RuntimeEventAttribute.TOOL_OUTPUT.value: tool_output,
637+
},
634638
)
635639
)
636640

‎ioa_observe/sdk/tracing/runtime_events.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ class RuntimeEventName(str, Enum):
1717
TOPOLOGY_EDGE_UPDATED = "topology.edge.updated"
1818
TOOL_STARTED = "tool.started"
1919
TOOL_COMPLETED = "tool.completed"
20+
LLM_STARTED = "llm.started"
21+
LLM_COMPLETED = "llm.completed"
2022
A2A_MESSAGE_SENT = "a2a.message.sent"
2123
A2A_MESSAGE_RECEIVED = "a2a.message.received"
2224
SLIM_MESSAGE_SENT = "slim.message.sent"
@@ -34,6 +36,11 @@ class RuntimeEventAttribute(str, Enum):
3436
AGENT_INPUT = "agent.input"
3537
TOOL_NAME = "tool.name"
3638
TOOL_INPUT = "tool.input"
39+
TOOL_OUTPUT = "tool.output"
40+
LLM_NAME = "llm.name"
41+
LLM_INPUT = "llm.input"
42+
LLM_OUTPUT = "llm.output"
43+
LLM_CALL_ID = "llm.call.id"
3744
SOURCE_AGENT = "source.agent"
3845
TARGET_AGENT = "target.agent"
3946
MESSAGE_ID = "message.id"
@@ -67,6 +74,10 @@ class RuntimeEventAttribute(str, Enum):
6774
| {RuntimeEventAttribute.TOOL_NAME.value},
6875
RuntimeEventName.TOOL_COMPLETED.value: COMMON_REQUIRED_ATTRIBUTES
6976
| {RuntimeEventAttribute.TOOL_NAME.value},
77+
RuntimeEventName.LLM_STARTED.value: COMMON_REQUIRED_ATTRIBUTES
78+
| {RuntimeEventAttribute.LLM_NAME.value},
79+
RuntimeEventName.LLM_COMPLETED.value: COMMON_REQUIRED_ATTRIBUTES
80+
| {RuntimeEventAttribute.LLM_NAME.value},
7081
RuntimeEventName.A2A_MESSAGE_SENT.value: COMMON_REQUIRED_ATTRIBUTES
7182
| {
7283
RuntimeEventAttribute.SOURCE_AGENT.value,

‎ioa_observe/sdk/tracing/tracing.py‎

Lines changed: 166 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,16 @@
5050
validate_transformer_rules,
5151
)
5252
from ioa_observe.sdk.tracing.topology import (
53+
next_session_event_version,
5354
record_session_completed,
5455
record_session_started,
5556
)
57+
from ioa_observe.sdk.tracing.runtime_event_emitter import emit_runtime_event
58+
from ioa_observe.sdk.tracing.runtime_events import (
59+
RuntimeEventAttribute,
60+
RuntimeEventName,
61+
build_runtime_event_attributes,
62+
)
5663
from ioa_observe.sdk.utils import is_notebook
5764
from ioa_observe.sdk.client import kv_store
5865

@@ -160,6 +167,10 @@ def __new__(
160167
# Track spans that have been processed to avoid duplicates
161168
obj._processed_spans = set()
162169
obj._processed_spans_lock = threading.Lock()
170+
obj._active_llm_spans: dict[
171+
int, tuple[str, str, str | None, str | None]
172+
] = {}
173+
obj._active_llm_spans_lock = threading.Lock()
163174
TracerWrapper.app_name = TracerWrapper.resource_attributes.get(
164175
"service.name", "observe"
165176
)
@@ -438,6 +449,19 @@ def _span_processor_on_start(self, span, parent_context):
438449

439450
if is_llm_span(span):
440451
self.llm_call_counter.add(1, attributes=span.attributes)
452+
if session_id is not None:
453+
llm_runtime_context = _llm_runtime_context(
454+
span,
455+
session_id=str(session_id),
456+
agent_name=workflow_name,
457+
)
458+
with self._active_llm_spans_lock:
459+
self._active_llm_spans[span.context.span_id] = llm_runtime_context
460+
_emit_llm_started_runtime_event(
461+
span,
462+
llm_runtime_context,
463+
capture_content=self._should_capture_content(),
464+
)
441465

442466
span.set_attribute("ioa_start_time", time.time()) # Record start time
443467

@@ -517,6 +541,54 @@ def span_processor_on_ending(self, span):
517541
)
518542

519543
determine_reliability_score(span)
544+
with self._active_llm_spans_lock:
545+
llm_runtime_context = self._active_llm_spans.pop(span.context.span_id, None)
546+
if llm_runtime_context is None and is_llm_span(span):
547+
session_id = span.attributes.get("session.id")
548+
if session_id:
549+
self.llm_call_counter.add(1, attributes=span.attributes)
550+
llm_runtime_context = _llm_runtime_context(
551+
span,
552+
session_id=str(session_id),
553+
agent_name=span.attributes.get(OBSERVE_WORKFLOW_NAME),
554+
)
555+
_emit_llm_started_runtime_event(
556+
span,
557+
llm_runtime_context,
558+
capture_content=self._should_capture_content(),
559+
)
560+
if llm_runtime_context:
561+
llm_name, llm_session_id, llm_agent_name, llm_operation = (
562+
llm_runtime_context
563+
)
564+
capture_content = self._should_capture_content()
565+
llm_input = (
566+
_llm_span_content(span, SpanAttributes.LLM_PROMPTS)
567+
if capture_content
568+
else None
569+
)
570+
llm_output = (
571+
_llm_span_content(span, SpanAttributes.LLM_COMPLETIONS)
572+
if capture_content
573+
else None
574+
)
575+
emit_runtime_event(
576+
build_runtime_event_attributes(
577+
RuntimeEventName.LLM_COMPLETED,
578+
session_id=llm_session_id,
579+
snapshot_version=next_session_event_version(llm_session_id),
580+
**{
581+
RuntimeEventAttribute.LLM_NAME.value: llm_name,
582+
RuntimeEventAttribute.LLM_CALL_ID.value: format(
583+
span.context.span_id, "016x"
584+
),
585+
RuntimeEventAttribute.AGENT_NAME.value: llm_agent_name,
586+
RuntimeEventAttribute.LLM_INPUT.value: llm_input,
587+
RuntimeEventAttribute.LLM_OUTPUT.value: llm_output,
588+
"operation.name": llm_operation,
589+
},
590+
)
591+
)
520592
# start_time = span.attributes.get("ioa_start_time")
521593

522594
# Apply transformations if enabled
@@ -649,6 +721,11 @@ def flush(self):
649721
def get_tracer(self):
650722
return self.__tracer_provider.get_tracer(TRACER_NAME)
651723

724+
def _should_capture_content(self) -> bool:
725+
return self.enable_content_tracing or bool(
726+
get_value("override_enable_content_tracing")
727+
)
728+
652729
def record_agent_execution(self, agent_name: str, success: bool):
653730
with self._agent_execution_counts_lock:
654731
counts = self._agent_execution_counts.setdefault(agent_name, [0, 0])
@@ -931,7 +1008,95 @@ def set_external_prompt_tracing_context(
9311008

9321009

9331010
def is_llm_span(span) -> bool:
934-
return span.attributes.get(SpanAttributes.LLM_REQUEST_TYPE) is not None
1011+
attributes = span.attributes
1012+
if attributes.get(SpanAttributes.LLM_REQUEST_TYPE) is not None:
1013+
return True
1014+
if attributes.get("gen_ai.operation.name") is not None:
1015+
return True
1016+
return attributes.get(SpanAttributes.LLM_SYSTEM) is not None and (
1017+
attributes.get(SpanAttributes.LLM_REQUEST_MODEL) is not None
1018+
or attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) is not None
1019+
)
1020+
1021+
1022+
def _llm_span_name(span) -> str:
1023+
return str(
1024+
span.attributes.get(SpanAttributes.LLM_REQUEST_MODEL)
1025+
or span.attributes.get(SpanAttributes.LLM_RESPONSE_MODEL)
1026+
or span.attributes.get(SpanAttributes.LLM_SYSTEM)
1027+
or span.name
1028+
)
1029+
1030+
1031+
def _llm_span_operation(span) -> str | None:
1032+
operation = span.attributes.get(
1033+
SpanAttributes.LLM_REQUEST_TYPE
1034+
) or span.attributes.get("gen_ai.operation.name")
1035+
if operation is not None:
1036+
return str(operation)
1037+
if "." in span.name:
1038+
return span.name.rsplit(".", 1)[-1]
1039+
return None
1040+
1041+
1042+
def _llm_runtime_context(
1043+
span,
1044+
*,
1045+
session_id: str,
1046+
agent_name,
1047+
) -> tuple[str, str, str | None, str | None]:
1048+
return (
1049+
_llm_span_name(span),
1050+
session_id,
1051+
str(agent_name) if agent_name is not None else None,
1052+
_llm_span_operation(span),
1053+
)
1054+
1055+
1056+
def _emit_llm_started_runtime_event(
1057+
span,
1058+
llm_runtime_context: tuple[str, str, str | None, str | None],
1059+
*,
1060+
capture_content: bool,
1061+
) -> None:
1062+
llm_name, session_id, agent_name, operation = llm_runtime_context
1063+
llm_input = (
1064+
_llm_span_content(span, SpanAttributes.LLM_PROMPTS) if capture_content else None
1065+
)
1066+
emit_runtime_event(
1067+
build_runtime_event_attributes(
1068+
RuntimeEventName.LLM_STARTED,
1069+
session_id=session_id,
1070+
snapshot_version=next_session_event_version(session_id),
1071+
**{
1072+
RuntimeEventAttribute.LLM_NAME.value: llm_name,
1073+
RuntimeEventAttribute.LLM_CALL_ID.value: format(
1074+
span.context.span_id, "016x"
1075+
),
1076+
RuntimeEventAttribute.AGENT_NAME.value: agent_name,
1077+
RuntimeEventAttribute.LLM_INPUT.value: llm_input,
1078+
"operation.name": operation,
1079+
},
1080+
)
1081+
)
1082+
1083+
1084+
def _llm_span_content(span, prefix: str) -> str | None:
1085+
pattern = re.compile(rf"^{re.escape(prefix)}\.(\d+)(?:\.(.+))?$")
1086+
messages: dict[int, dict[str, object]] = {}
1087+
for key, value in span.attributes.items():
1088+
match = pattern.match(key)
1089+
if match is None:
1090+
continue
1091+
index = int(match.group(1))
1092+
field = match.group(2) or "content"
1093+
messages.setdefault(index, {})[field] = value
1094+
if not messages:
1095+
return None
1096+
ordered_messages = [
1097+
dict(sorted(messages[index].items())) for index in sorted(messages)
1098+
]
1099+
return json.dumps(ordered_messages, separators=(",", ":"), default=str)
9351100

9361101

9371102
def init_spans_exporter(api_endpoint: str, headers: Dict[str, str]) -> SpanExporter:

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
55

66
[project]
77
name = "ioa-observe-sdk"
8-
version = "1.0.44"
8+
version = "1.0.45"
99
license = "Apache-2.0"
1010
description = "IOA Observability SDK"
1111
readme = "README.md"

0 commit comments

Comments
 (0)