diff --git a/agentx/tracing/tracer.py b/agentx/tracing/tracer.py index 20e38ca..b4e389d 100644 --- a/agentx/tracing/tracer.py +++ b/agentx/tracing/tracer.py @@ -581,6 +581,14 @@ def __init__(self) -> None: self.output: Any = None +class _MemoryOpRecorder: + """Handle yielded by ``Tracer.trace_memory()`` - set ``output`` (what was recalled or + stored) inside the block.""" + + def __init__(self) -> None: + self.output: Any = None + + class _ToolCallRecorder: """Handle yielded by ``Tracer.trace_tool_call()`` - set ``output`` inside the block. ``success``/``error`` may be set manually; an exception escaping the block sets them @@ -832,6 +840,66 @@ def record_retrieval( span_kind="retrieval", ) + def record_memory( + self, + name: str = "Memory", + *, + operation: Optional[str] = None, + query: Optional[str] = None, + output: Any = None, + duration_ms: Optional[float] = None, + start_time: Optional[float] = None, + end_time: Optional[float] = None, + ) -> None: + """ + Manually record a long-term-memory operation (a Mem0/Zep/Letta-style recall or store) + as a ``span_kind="memory"`` child span of the active span. ``operation`` is free text - + conventionally ``"read"`` or ``"write"`` - carried in the span's metadata; the kind + itself stays one value so dashboards and scorers can select all memory activity at once. + Deliberately NOT a retrieval: retrieval spans feed the RAG judges' ``{context}`` + (knowledge grounding), while memory is recalled state - see the engine's spanKind.ts. + """ + active_span = self.current_span + if active_span is None: + return + active_span.child_span( + name, + start_time=start_time, + end_time=end_time, + duration_ms=duration_ms, + input=query, + output=output, + metadata={"kind": "memory", **({"operation": operation} if operation else {})}, + span_kind="memory", + ) + + @contextmanager + def trace_memory( + self, name: str = "Memory", *, operation: Optional[str] = None, query: Optional[str] = None + ) -> Iterator["_MemoryOpRecorder"]: + """ + Context manager that times a memory operation and records it via + :meth:`record_memory` on exit:: + + with tracer.trace_memory("user prefs", operation="read", query=user_id) as m: + m.output = memory.search(user_id, question) + """ + start_t = time.time() + recorder = _MemoryOpRecorder() + try: + yield recorder + finally: + end_t = time.time() + self.record_memory( + name, + operation=operation, + query=query, + output=recorder.output, + duration_ms=(end_t - start_t) * 1000, + start_time=start_t, + end_time=end_t, + ) + @contextmanager def trace_retrieval(self, name: str = "Retrieval", *, query: Optional[str] = None) -> Iterator["_RetrievalRecorder"]: """ diff --git a/tests/test_span_tree.py b/tests/test_span_tree.py index 2dc3218..fc8b424 100644 --- a/tests/test_span_tree.py +++ b/tests/test_span_tree.py @@ -594,6 +594,22 @@ def test_trace_tool_call_emits_real_child_span(): assert root["tool_calls"][0]["output"] == "digital purchases are final" +def test_trace_memory_emits_a_memory_kind_child_span(): + tracer = make_tracer() + with tracer.trace("agent") as span: + with tracer.trace_memory("user prefs", operation="read", query="u-42") as m: + m.output = ["prefers window seats"] + + wires = enqueued_wires(tracer) + child = wires[0] + assert child["name"] == "user prefs" + assert child["span_kind"] == "memory" + assert child["metadata"]["kind"] == "memory" + assert child["metadata"]["operation"] == "read" + assert child["input"] == "u-42" + assert child["output"] == ["prefers window seats"] + + def test_trace_retrieval_emits_real_child_span(): tracer = make_tracer() with tracer.trace("agent") as span: