memtrace-ai 0.2.0__py3-none-any.whl

memtrace/__init__.py

@@ -0,0 +1,31 @@
+from memtrace import prompts
+from memtrace._version import __version__
+from memtrace.adapters.inbound.manual import record_retrieved_chunks, trace_llm_call, trace_step, trace_step_context
+from memtrace.application.context import get_current_run_id, session, user
+from memtrace.dependency_container import flush, init_tracer, shutdown
+from memtrace.langchain import MemTraceCallbackHandler, enable_langchain_instrumentation
+from memtrace.openai_agents import enable_openai_agents_instrumentation
+from memtrace.pydantic_ai import enable_pydantic_ai_instrumentation
+from memtrace.user_feedback import current_trace_id, feedback, retract_feedback
+
+__all__ = [
+    "__version__",
+    "init_tracer",
+    "flush",
+    "shutdown",
+    "session",
+    "user",
+    "trace_step",
+    "trace_step_context",
+    "trace_llm_call",
+    "record_retrieved_chunks",
+    "get_current_run_id",
+    "current_trace_id",
+    "feedback",
+    "retract_feedback",
+    "prompts",
+    "MemTraceCallbackHandler",
+    "enable_langchain_instrumentation",
+    "enable_pydantic_ai_instrumentation",
+    "enable_openai_agents_instrumentation",
+]

memtrace/_version.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

memtrace/adapters/inbound/asgi/__init__.py

@@ -0,0 +1,3 @@
+from memtrace.adapters.inbound.asgi.prompt_override import PromptOverrideMiddleware
+
+__all__ = ["PromptOverrideMiddleware"]

memtrace/adapters/inbound/asgi/prompt_override.py

@@ -0,0 +1,39 @@
+"""ASGI middleware that lets MemTrace's playground test a prompt version in this agent (ADR-071).
+
+MemTrace calls the agent's chat with the token of a short-lived override in the `x-memtrace-prompt-override` header. This
+middleware reads it and, for the duration of that request only, makes `handle.compile()` ask MemTrace for the version the
+token grants. It does nothing unless the agent opted in with `MEMTRACE_ALLOW_PROMPT_OVERRIDE=true`, and a token MemTrace
+does not recognize changes nothing.
+"""
+from typing import Any, Awaitable, Callable, MutableMapping
+
+from memtrace.application.prompt_override import use_override
+from memtrace.config import settings
+from memtrace.domain.prompt import PROMPT_OVERRIDE_HEADER, is_override_token
+
+Scope = MutableMapping[str, Any]
+Receive = Callable[[], Awaitable[MutableMapping[str, Any]]]
+Send = Callable[[MutableMapping[str, Any]], Awaitable[None]]
+
+
+class PromptOverrideMiddleware:
+    """Pure ASGI (FastAPI, Starlette, Quart…): `app.add_middleware(PromptOverrideMiddleware)`."""
+
+    def __init__(self, app: Callable[[Scope, Receive, Send], Awaitable[None]]) -> None:
+        self.app = app
+        self._header = PROMPT_OVERRIDE_HEADER.encode("latin-1")
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        token = self._token(scope) if scope.get("type") in ("http", "websocket") and settings.allow_prompt_override else None
+        if token is None:
+            await self.app(scope, receive, send)
+            return
+        with use_override(token):
+            await self.app(scope, receive, send)
+
+    def _token(self, scope: Scope) -> "str | None":
+        for name, value in scope.get("headers") or ():
+            if name.lower() == self._header:
+                candidate = value.decode("latin-1")
+                return candidate if is_override_token(candidate) else None
+        return None

memtrace/adapters/inbound/auto_instrumentation.py

@@ -0,0 +1,22 @@
+"""Shared plumbing for "auto" instrumentors (LangChain, Pydantic AI…): each just points a
+third-party library at the tracer provider configured by `init_tracer`, no callbacks needed.
+"""
+import logging
+from typing import Any, Optional, Tuple
+
+from memtrace.application.tracing_service import TracingService
+from memtrace.dependency_container import get_service
+
+logger = logging.getLogger("memtrace")
+
+
+def resolve_tracer_provider(feature: str, service: Optional[TracingService] = None) -> Tuple[TracingService, Optional[Any]]:
+    """The `(service, tracer_provider)` to instrument with, or `(service, None)` if tracing is off.
+
+    `feature`: human-readable name for the log line (e.g. "LangChain", "Pydantic AI").
+    """
+    service = service or get_service()
+    provider = service.tracer_provider
+    if provider is None:
+        logger.info("[MemTrace] Tracing is off; %s instrumentation not enabled.", feature)
+    return service, provider

memtrace/adapters/inbound/langchain/__init__.py

@@ -0,0 +1,5 @@
+from memtrace.adapters.inbound.langchain.auto import enable_langchain_instrumentation
+from memtrace.adapters.inbound.langchain.callback import MemTraceCallbackHandler
+from memtrace.adapters.inbound.langchain.prompt_middleware import prompt_middleware
+
+__all__ = ["MemTraceCallbackHandler", "enable_langchain_instrumentation", "prompt_middleware"]

memtrace/adapters/inbound/langchain/auto.py

@@ -0,0 +1,43 @@
+"""Automatic LangChain / LangGraph instrumentation through OpenTelemetry's instrumentor."""
+import logging
+import os
+from typing import Optional
+
+from memtrace.adapters.inbound.auto_instrumentation import resolve_tracer_provider
+from memtrace.application.tracing_service import TracingService
+
+logger = logging.getLogger("memtrace")
+
+_INSTALL_HINT = "pip install 'memtrace-ai[langchain]'"
+_CONTENT_ENV = "TRACELOOP_TRACE_CONTENT"  # honored by opentelemetry-instrumentation-langchain
+
+
+def enable_langchain_instrumentation(service: Optional[TracingService] = None) -> None:
+    """Enables automatic LangChain / LangGraph instrumentation (no callbacks needed).
+
+    Requires: pip install 'memtrace-ai[langchain]'. Raises ImportError if it is missing.
+
+    Spans go to the tracer configured by `init_tracer`, and prompts/completions are only
+    recorded when MEMTRACE_CAPTURE_CONTENT=true. Safe to call more than once. It is an alternative to
+    MemTraceCallbackHandler; use one of them, not both, or every run is traced twice.
+    """
+    try:
+        from opentelemetry.instrumentation.langchain import LangchainInstrumentor
+    except ImportError as exc:
+        raise ImportError(f"LangChain instrumentation requires: {_INSTALL_HINT}") from exc
+
+    # provider is process-wide, survives shutdown()/init_tracer()
+    service, provider = resolve_tracer_provider("LangChain", service)
+    if provider is None:
+        return
+
+    # Respect the privacy default of ADR-004 unless the user set the instrumentor's own switch
+    os.environ.setdefault(_CONTENT_ENV, "true" if service.captures_content else "false")
+
+    instrumentor = LangchainInstrumentor()
+    if instrumentor.is_instrumented_by_opentelemetry:
+        # Already bound to MemTrace's switchable provider, which follows init_tracer()/shutdown()
+        logger.debug("[MemTrace] LangChain auto-instrumentation was already enabled")
+        return
+    instrumentor.instrument(tracer_provider=provider)
+    logger.info("[MemTrace] LangChain auto-instrumentation enabled")

memtrace/adapters/inbound/langchain/callback.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from memtrace.adapters.inbound.langchain.mapping import (
+    generation_dicts,
+    message_dicts,
+    request_to_call,
+    response_to_call,
+    step_name,
+)
+from memtrace.application.retrieval_capture import note_retrieved
+from memtrace.application.tracing_service import TracingService
+from memtrace.dependency_container import get_service
+from memtrace.domain import semconv as sc
+from memtrace.domain.attributes import llm_attributes, retrieved_chunks
+from memtrace.domain.model import StepType
+from memtrace.domain.serialization import to_json
+
+logger = logging.getLogger("memtrace")
+
+_INSTALL_HINT = "pip install 'memtrace-ai[langchain]'"
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+except ImportError:
+    class BaseCallbackHandler:  # type: ignore[no-redef]
+        """Stand-in so importing this module never requires LangChain; instantiating does."""
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            raise ImportError(f"MemTraceCallbackHandler requires LangChain: {_INSTALL_HINT}")
+
+
+class MemTraceCallbackHandler(BaseCallbackHandler):
+    """LangChain / LangGraph callback handler.
+
+    Usage: `chain.invoke(x, config={"callbacks": [MemTraceCallbackHandler()]})`.
+    Inside a `@trace_step`, its spans hang from it.
+    """
+
+    def __init__(self, service: Optional[TracingService] = None) -> None:
+        super().__init__()
+        self._service = service or get_service()
+
+    # ----- helpers -----
+
+    def _context_attrs(self, tags: Optional[List[str]], metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+        md = metadata or {}
+        return {
+            sc.MEMTRACE_TAGS: [str(t) for t in tags] if tags else None,
+            sc.MEMTRACE_METADATA: to_json(md, self._service.max_content_length) if md else None,
+            # `thread_id` is the LangGraph conversation identifier
+            sc.GEN_AI_CONVERSATION_ID: md.get("session_id") or md.get("conversation_id") or md.get("thread_id"),
+            sc.ENDUSER_ID: md.get("user_id"),
+            # LangGraph injects `langgraph_node`/`langgraph_step` into each run's metadata (ADR-011)
+            # the graph's own root run has no `langgraph_*` keys, only `ls_integration` (ADR-011)
+            sc.MEMTRACE_FRAMEWORK: "langgraph" if md.get("ls_integration") == "langgraph" or any(k.startswith("langgraph_") for k in md) else "langchain",
+        }
+
+    def _start(self, name, step_type, run_id, parent_run_id, tags, metadata, extra=None) -> None:
+        attrs = {**self._context_attrs(tags, metadata), **(extra or {})}
+        self._service.start_run(name, step_type, run_id=run_id, parent_run_id=parent_run_id, attributes=attrs)
+
+    def _end(self, run_id, attrs=None, error=None) -> None:
+        self._service.end_run(run_id, error=error, attributes=attrs)
+
+    # ----- chains / agents -----
+
+    def on_chain_start(self, serialized, inputs, *, run_id, parent_run_id=None, tags=None, metadata=None, **kwargs) -> Any:
+        self._start(step_name(serialized, kwargs, "Chain"), StepType.CHAIN, run_id, parent_run_id, tags, metadata,
+                    {sc.MEMTRACE_INPUT: self._service.capture(inputs)})
+
+    def on_chain_end(self, outputs, *, run_id, parent_run_id=None, **kwargs) -> Any:
+        self._end(run_id, {sc.MEMTRACE_OUTPUT: self._service.capture(outputs)})
+
+    def on_chain_error(self, error, *, run_id, parent_run_id=None, **kwargs) -> Any:
+        self._end(run_id, error=error)
+
+    # ----- tools -----
+
+    def on_tool_start(self, serialized, input_str, *, run_id, parent_run_id=None, tags=None, metadata=None, inputs=None, **kwargs) -> Any:
+        name = step_name(serialized, kwargs, "Tool")
+        self._start(name, StepType.TOOL, run_id, parent_run_id, tags, metadata, {
+            sc.GEN_AI_TOOL_NAME: name,
+            sc.GEN_AI_TOOL_CALL_ID: kwargs.get("tool_call_id"),
+            sc.GEN_AI_TOOL_CALL_ARGUMENTS: self._service.capture(inputs if inputs is not None else input_str),
+        })
+
+    def on_tool_end(self, output, *, run_id, parent_run_id=None, **kwargs) -> Any:
+        result = getattr(output, "content", output)  # ToolMessage or plain value
+        self._end(run_id, {sc.GEN_AI_TOOL_CALL_RESULT: self._service.capture(result)})
+
+    def on_tool_error(self, error, *, run_id, parent_run_id=None, **kwargs) -> Any:
+        self._end(run_id, error=error)
+
+    # ----- retrievers -----
+
+    def on_retriever_start(self, serialized, query, *, run_id, parent_run_id=None, tags=None, metadata=None, **kwargs) -> Any:
+        self._start(step_name(serialized, kwargs, "Retriever"), StepType.RETRIEVER, run_id, parent_run_id, tags, metadata,
+                    {sc.MEMTRACE_INPUT: self._service.capture(query)})
+
+    def on_retriever_end(self, documents, *, run_id, parent_run_id=None, **kwargs) -> Any:
+        chunks = retrieved_chunks(documents)
+        note_retrieved(chunks)
+        self._end(run_id, {
+            sc.MEMTRACE_RETRIEVER_DOCUMENTS: len(documents or []),
+            sc.MEMTRACE_RETRIEVER_CHUNKS: self._service.capture(chunks),
+        })
+
+    def on_retriever_error(self, error, *, run_id, parent_run_id=None, **kwargs) -> Any:
+        self._end(run_id, error=error)
+
+    # ----- LLMs -----
+
+    def _llm_start(self, serialized, run_id, parent_run_id, tags, metadata, invocation_params, kwargs, input_messages) -> None:
+        call = request_to_call(metadata, invocation_params or kwargs.get("invocation_params") or {})
+        base = step_name(serialized, {}, "llm")
+        self._start(
+            f"{base}: {call.model}" if call.model else base, StepType.LLM, run_id, parent_run_id, tags, metadata,
+            llm_attributes(call, self._service.capture(input_messages)),
+        )
+
+    def on_chat_model_start(self, serialized, messages, *, run_id, parent_run_id=None, tags=None, metadata=None, invocation_params=None, **kwargs) -> Any:
+        self._llm_start(serialized, run_id, parent_run_id, tags, metadata, invocation_params, kwargs,
+                        message_dicts(messages) if self._service.captures_content else None)
+
+    def on_llm_start(self, serialized, prompts, *, run_id, parent_run_id=None, tags=None, metadata=None, invocation_params=None, **kwargs) -> Any:
+        self._llm_start(serialized, run_id, parent_run_id, tags, metadata, invocation_params, kwargs,
+                        [{"role": "user", "content": p} for p in prompts])
+
+    def on_llm_end(self, response, *, run_id, parent_run_id=None, **kwargs) -> Any:
+        try:
+            output = generation_dicts(response) if self._service.captures_content else None
+            attrs = llm_attributes(response_to_call(response), output_messages=self._service.capture(output))
+        except Exception as exc:  # an unexpectedly shaped response must not lose the span
+            logger.warning("[MemTrace] Error extracting data from the LLM response: %s", exc)
+            attrs = None
+        self._end(run_id, attrs)
+
+    def on_llm_error(self, error, *, run_id, parent_run_id=None, **kwargs) -> Any:
+        self._end(run_id, error=error)

memtrace/adapters/inbound/langchain/mapping.py

@@ -0,0 +1,110 @@
+"""Pure translation of LangChain structures into domain concepts. No `TracingService`, no OTel."""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+
+from memtrace.domain.model import LlmCall
+
+
+def _get(obj: Any, key: str) -> Any:
+    return obj.get(key) if isinstance(obj, dict) else getattr(obj, key, None)
+
+
+def _as_int(val: Any) -> Optional[int]:
+    return val if isinstance(val, int) and not isinstance(val, bool) else None
+
+
+def _tokens(usage: Any) -> Tuple[Optional[int], Optional[int], Optional[int]]:
+    """(input, output, total) from `usage_metadata` (LangChain), `token_usage` (OpenAI) or `usage` (Anthropic)."""
+    if not usage:
+        return None, None, None
+    inp = _as_int(_get(usage, "input_tokens"))
+    inp = inp if inp is not None else _as_int(_get(usage, "prompt_tokens"))
+    out = _as_int(_get(usage, "output_tokens"))
+    out = out if out is not None else _as_int(_get(usage, "completion_tokens"))
+    return inp, out, _as_int(_get(usage, "total_tokens"))
+
+
+def response_to_call(response: Any) -> LlmCall:
+    """Response data of an `LLMResult` (duck-typed): messages first, `llm_output` as a fallback."""
+    inp = out = total = None
+    model = resp_id = None
+    finish: List[str] = []
+
+    for gens in getattr(response, "generations", None) or []:
+        for gen in gens or []:
+            msg = _get(gen, "message")
+            meta = _get(msg, "response_metadata") or {}
+            info = _get(gen, "generation_info") or {}
+            if inp is None and out is None:
+                inp, out, total = _tokens(_get(msg, "usage_metadata"))
+            model = model or _get(meta, "model_name") or _get(meta, "model")
+            resp_id = resp_id or _get(msg, "id") or _get(meta, "id")
+            reason = _get(info, "finish_reason") or _get(meta, "finish_reason") or _get(meta, "stop_reason")
+            if reason:
+                finish.append(str(reason))
+
+    llm_output = getattr(response, "llm_output", None) or {}
+    if inp is None and out is None:
+        inp, out, total = _tokens(_get(llm_output, "token_usage") or _get(llm_output, "usage"))
+    model = model or _get(llm_output, "model_name") or _get(llm_output, "model")
+    return LlmCall(
+        response_model=model, response_id=resp_id, finish_reasons=finish,
+        input_tokens=inp, output_tokens=out, total_tokens=total,
+    )
+
+
+def request_to_call(metadata: Optional[Dict[str, Any]], params: Dict[str, Any]) -> LlmCall:
+    md = metadata or {}
+    return LlmCall(
+        provider=md.get("ls_provider") or params.get("_type"),
+        model=md.get("ls_model_name") or params.get("model_name") or params.get("model") or params.get("model_id"),
+        temperature=md.get("ls_temperature", params.get("temperature")),
+        max_tokens=md.get("ls_max_tokens", params.get("max_tokens") or params.get("max_tokens_to_sample")),
+        top_p=params.get("top_p"),
+        frequency_penalty=params.get("frequency_penalty"),
+        presence_penalty=params.get("presence_penalty"),
+    )
+
+
+def step_name(serialized: Optional[Dict[str, Any]], kwargs: Dict[str, Any], default: str) -> str:
+    """`serialized` may be None (LangGraph, runnables); LangChain passes `name` as a kwarg."""
+    ser = serialized or {}
+    ident = ser.get("id")
+    return kwargs.get("name") or ser.get("name") or (ident[-1] if ident else None) or default
+
+
+def _tool_calls_of(msg: Any) -> Any:
+    """`tool_calls` of the model message (or of `additional_kwargs` for some providers)."""
+    calls = getattr(msg, "tool_calls", None)
+    if not calls:
+        calls = (getattr(msg, "additional_kwargs", None) or {}).get("tool_calls")
+    return calls or None
+
+
+def message_dicts(messages: List[List[Any]]) -> List[Dict[str, Any]]:
+    out = []
+    for batch in messages:
+        for m in batch:
+            d: Dict[str, Any] = {
+                "role": getattr(m, "type", None) or getattr(m, "role", None) or "user",
+                "content": getattr(m, "content", str(m)),
+            }
+            calls = _tool_calls_of(m)
+            if calls:
+                d["tool_calls"] = calls
+            out.append(d)
+    return out
+
+
+def generation_dicts(response: Any) -> List[Dict[str, Any]]:
+    out = []
+    for gens in getattr(response, "generations", None) or []:
+        for g in gens:
+            msg = getattr(g, "message", None)
+            d: Dict[str, Any] = {"role": getattr(msg, "type", "assistant"), "content": getattr(g, "text", "")}
+            calls = _tool_calls_of(msg)
+            if calls:
+                d["tool_calls"] = calls
+            out.append(d)
+    return out

memtrace/adapters/inbound/langchain/prompt_middleware.py

@@ -0,0 +1,45 @@
+"""Gives a registry prompt to a LangChain agent without freezing its version (ADR-068).
+
+`create_agent(system_prompt="...")` takes a string, so a prompt compiled when the agent is built would stay on the
+version it had at that moment. LangChain's middleware runs on every model call, which is where the version has to be
+decided: `prompt_middleware(handle)` replaces the system message with `handle.compile(...)` each time.
+"""
+from typing import Any, Callable, Mapping
+
+from memtrace.application.prompt_registry import PromptHandle
+
+
+def prompt_middleware(handle: PromptHandle, **variables: Any) -> Any:
+    """Middleware for `create_agent(..., middleware=[...])` that uses `handle` as the agent's system prompt.
+
+    Every model call asks the handle for the version in use at that moment (in memory, no network), so moving the tag
+    in MemTrace changes the agent without rebuilding it. The trace is linked to the version used (the prompt name and
+    version are written on the current span).
+
+    `variables` fill the `{{placeholders}}` of the prompt. A value that is a function is called on every model call with
+    LangChain's `ModelRequest` (`request.state`, `request.runtime.context`…) and its result is used::
+
+        agent = create_agent(
+            model,
+            tools,
+            middleware=[prompt_middleware(weather, city=lambda request: request.runtime.context["city"], tone="friendly")],
+        )
+
+    It *replaces* the system message: do not also pass `system_prompt=` to `create_agent`.
+    Needs `langchain>=1.0` (`pip install langchain`).
+    """
+    try:
+        from langchain.agents.middleware import dynamic_prompt
+    except ImportError as exc:
+        raise ImportError("prompt_middleware needs langchain>=1.0: pip install 'langchain>=1.0'") from exc
+
+    def resolve(request: Any) -> str:
+        return handle.compile(**_values(variables, request))
+
+    # LangChain names a middleware after its function and refuses two with the same name in one agent
+    resolve.__name__ = "memtrace_prompt_" + "".join(c if c.isalnum() else "_" for c in handle.name)
+    return dynamic_prompt(resolve)
+
+
+def _values(variables: Mapping[str, Any], request: Any) -> Mapping[str, Any]:
+    return {name: value(request) if isinstance(value, Callable) else value for name, value in variables.items()}  # type: ignore[arg-type]

memtrace/adapters/inbound/manual/__init__.py

@@ -0,0 +1,5 @@
+from memtrace.adapters.inbound.manual.llm import trace_llm_call
+from memtrace.adapters.inbound.manual.retrieval import record_retrieved_chunks
+from memtrace.adapters.inbound.manual.step import trace_step, trace_step_context
+
+__all__ = ["trace_step", "trace_step_context", "trace_llm_call", "record_retrieved_chunks"]

memtrace/adapters/inbound/manual/llm.py

@@ -0,0 +1,47 @@
+from typing import Any, Dict, List, Optional, Sequence
+
+from memtrace.application.tracing_service import TracingService
+from memtrace.dependency_container import get_service
+from memtrace.domain.model import LlmCall
+
+
+def trace_llm_call(
+    provider: str,
+    model: str,
+    *,
+    operation: str = "chat",
+    input_messages: Optional[List[Dict[str, Any]]] = None,
+    output_messages: Optional[List[Dict[str, Any]]] = None,
+    input_tokens: Optional[int] = None,
+    output_tokens: Optional[int] = None,
+    response_model: Optional[str] = None,
+    response_id: Optional[str] = None,
+    finish_reasons: Sequence[str] = (),
+    temperature: Optional[float] = None,
+    max_tokens: Optional[int] = None,
+    top_p: Optional[float] = None,
+    frequency_penalty: Optional[float] = None,
+    presence_penalty: Optional[float] = None,
+    attributes: Optional[Dict[str, Any]] = None,
+    service: Optional[TracingService] = None,
+) -> None:
+    """Records the standardized GenAI attributes on the current span. Never raises.
+
+    Everything after `provider` and `model` is keyword-only.
+    """
+    call = LlmCall(
+        provider=provider,
+        model=model,
+        operation=operation,
+        response_model=response_model,
+        response_id=response_id,
+        finish_reasons=finish_reasons,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        top_p=top_p,
+        frequency_penalty=frequency_penalty,
+        presence_penalty=presence_penalty,
+    )
+    (service or get_service()).record_llm_call(call, input_messages, output_messages, attributes)

memtrace/adapters/inbound/manual/retrieval.py

@@ -0,0 +1,14 @@
+from typing import Any, Optional, Sequence
+
+from memtrace.application.tracing_service import TracingService
+from memtrace.dependency_container import get_service
+
+
+def record_retrieved_chunks(documents: Sequence[Any], *, service: Optional[TracingService] = None) -> None:
+    """Records the chunks a retrieval step returned on the current span (use inside a `retriever` step). Never raises.
+
+    `documents`: strings, dicts (`text`, `id`?, `source`?, `score`?) or LangChain-like documents
+    (`page_content` + `metadata`). The list order is the rank. The chunk text is exported only when
+    content capture is enabled; the count is always recorded.
+    """
+    (service or get_service()).record_retrieved_chunks(documents)

memtrace/adapters/inbound/manual/step.py

@@ -0,0 +1,81 @@
+import functools
+import inspect
+from contextlib import AbstractContextManager
+from typing import Any, Dict, Optional, Union
+
+from memtrace.application.tracing_service import TracingService
+from memtrace.dependency_container import get_service
+from memtrace.domain.attributes import step_input_attributes, step_output_attributes
+from memtrace.domain.model import StepType
+
+
+def trace_step_context(
+    name: str,
+    step_type: Union[StepType, str] = StepType.CHAIN,
+    attributes: Optional[Dict[str, Any]] = None,
+    *,
+    service: Optional[TracingService] = None,
+) -> AbstractContextManager:
+    """Instruments a block; spans created inside it (including those of LangChain or of
+    auto-instrumented libraries) hang from it. Yields the `run_id`.
+
+    `service` selects an explicit tracer (from `init_tracer`); by default the global one is used.
+    """
+    return (service or get_service()).step(name, step_type, attributes)
+
+
+def _describe_call(func: Any, args: tuple, kwargs: dict) -> Any:
+    try:
+        bound = inspect.signature(func).bind_partial(*args, **kwargs)
+        return {k: v for k, v in bound.arguments.items() if k not in ("self", "cls")}
+    except Exception:
+        return {"args": list(args), "kwargs": kwargs}
+
+
+def trace_step(
+    name: Optional[str] = None,
+    step_type: Union[StepType, str] = StepType.CHAIN,
+    *,
+    service: Optional[TracingService] = None,
+):
+    """Decorator for sync or async functions and methods of agents and tools.
+
+    With MEMTRACE_CAPTURE_CONTENT=true it also records the arguments and the return value.
+    `service` selects an explicit tracer; by default the global one is used.
+    """
+
+    def decorator(func):
+        span_name = name or func.__name__
+
+        def record_input(svc, run_id, args, kwargs) -> None:
+            payload = _describe_call(func, args, kwargs) if svc.captures_content else None
+            svc.annotate_run(run_id, step_input_attributes(step_type, span_name, svc.capture(payload)))
+
+        def record_output(svc, run_id, result) -> None:
+            svc.annotate_run(run_id, step_output_attributes(step_type, svc.capture(result)))
+
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                svc = service or get_service()
+                with svc.step(span_name, step_type) as run_id:
+                    record_input(svc, run_id, args, kwargs)
+                    result = await func(*args, **kwargs)
+                    record_output(svc, run_id, result)
+                    return result
+
+            return async_wrapper
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            svc = service or get_service()
+            with svc.step(span_name, step_type) as run_id:
+                record_input(svc, run_id, args, kwargs)
+                result = func(*args, **kwargs)
+                record_output(svc, run_id, result)
+                return result
+
+        return wrapper
+
+    return decorator

memtrace/adapters/inbound/openai_agents/__init__.py

@@ -0,0 +1,4 @@
+from memtrace.adapters.inbound.openai_agents.auto import enable_openai_agents_instrumentation
+from memtrace.adapters.inbound.openai_agents.callback import MemTraceTracingProcessor
+
+__all__ = ["MemTraceTracingProcessor", "enable_openai_agents_instrumentation"]

memtrace/adapters/inbound/openai_agents/auto.py

@@ -0,0 +1,35 @@
+"""One call to send every OpenAI Agents SDK run to MemTrace."""
+import logging
+from typing import Optional
+
+from memtrace.application.tracing_service import TracingService
+
+logger = logging.getLogger("memtrace")
+
+_INSTALL_HINT = "pip install 'memtrace-ai[openai-agents]'"
+_installed: Optional[object] = None
+
+
+def enable_openai_agents_instrumentation(service: Optional[TracingService] = None, *, keep_openai_export: bool = True) -> None:
+    """Registers a MemTrace tracing processor with the OpenAI Agents SDK, for all agents.
+
+    Requires: pip install 'memtrace-ai[openai-agents]'. Raises ImportError if it is missing. Safe to call more than once.
+
+    `keep_openai_export`: the Agents SDK also uploads its traces to the OpenAI dashboard by default. `True` (default) keeps
+    that; `False` replaces it, so traces only go to MemTrace. Prompts and completions are recorded only with
+    MEMTRACE_CAPTURE_CONTENT=true.
+    """
+    global _installed
+    try:
+        from agents import add_trace_processor, set_trace_processors
+    except ImportError as exc:
+        raise ImportError(f"OpenAI Agents instrumentation requires: {_INSTALL_HINT}") from exc
+    from memtrace.adapters.inbound.openai_agents.callback import MemTraceTracingProcessor
+
+    if _installed is not None:
+        logger.debug("[MemTrace] OpenAI Agents instrumentation was already enabled")
+        return
+    processor = MemTraceTracingProcessor(service)
+    (add_trace_processor if keep_openai_export else lambda p: set_trace_processors([p]))(processor)
+    _installed = processor
+    logger.info("[MemTrace] OpenAI Agents instrumentation enabled")