Skip to content

coalex_context()

Context manager that creates a coalex.invocation parent span. All spans created inside this block are tagged with the agent metadata, making them queryable by agent_id, request_id, and version in the Coalex dashboard.

Signature

@contextlib.contextmanager
def coalex_context(
    *,
    agent_id: str,
    request_id: str | None = None,
    version: str | None = None,
) -> Generator[None, None, None]

interface CoalexContextOptions {
    agentId: string;
    requestId?: string;
    version?: string;
}

async function coalexContext<T>(
    options: CoalexContextOptions,
    fn: () => Promise<T> | T,
): Promise<T>

Callback pattern

The TypeScript SDK uses a callback pattern instead of Python's with statement. Pass your logic as an async function to coalexContext(). A synchronous variant coalexContextSync() is also available.

Parameters

Parameter Type Default Description
agent_id str required Identifier for the AI agent. Used to group all traces from the same agent in the dashboard.
request_id str \| None None Unique request or invocation ID. Use this to correlate a single user request across multiple spans.
version str \| None None Agent version string (e.g., "1.2.0", "canary"). Useful for A/B testing and rollback tracking.

All parameters are keyword-only (enforced by *).

Returns

A context manager that yields None. Use it with with statements.

Span Attributes

The coalex.invocation parent span carries these attributes:

Attribute Set When Value
coalex.agent_id Always Value of agent_id parameter
coalex.request_id request_id is not None Value of request_id parameter
coalex.agent_version version is not None Value of version parameter

Attribute propagation

The CoalexAttributePropagator (configured by register()) automatically copies coalex.* attributes from the parent span to all child spans. This means every LLM call, retrieval, and tool invocation inside the context block inherits the agent metadata without manual tagging.

How It Works

  1. Obtains a tracer named "coalex" from the global TracerProvider.
  2. Starts a new span named "coalex.invocation" as the current span.
  3. Sets the coalex.* attributes on that span.
  4. All spans created inside the with block become children of the coalex.invocation span.
  5. When the block exits, the span ends automatically.
graph TD
    A["coalex.invocation<br/>(coalex.agent_id=support-bot)"] --> B["openai.chat.completions<br/>(auto-instrumented)"]
    A --> C["knowledge_base<br/>(@retrieval_span)"]
    C --> D["embed<br/>(@embedding_span)"]

Examples

Basic usage

import coalex
from openai import OpenAI

coalex.register(endpoint="https://your-org.coalex.ai", api_key="your-key")
coalex.auto_instrument()

client = OpenAI()

with coalex.coalex_context(agent_id="support-bot", request_id="req-001"):
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "What is the refund policy?"}],
    )
    print(response.choices[0].message.content)

import { register, coalexContext, autoInstrument } from "@coalex-ai/sdk";
import OpenAI from "openai";

register({ endpoint: "https://your-org.coalex.ai", apiKey: "your-key" });
autoInstrument();

const client = new OpenAI();

await coalexContext({ agentId: "support-bot", requestId: "req-001" }, async () => {
    const response = await client.chat.completions.create({
        model: "gpt-4o",
        messages: [{ role: "user", content: "What is the refund policy?" }],
    });
    console.log(response.choices[0].message.content);
});

With version tracking

with coalex.coalex_context(
    agent_id="claims-processor",
    request_id="claim-2024-0042",
    version="2.1.0",
):
    # All spans inside here carry:
    # - coalex.agent_id = "claims-processor"
    # - coalex.request_id = "claim-2024-0042"
    # - coalex.agent_version = "2.1.0"
    result = process_claim(claim_data)

await coalexContext({
    agentId: "claims-processor",
    requestId: "claim-2024-0042",
    version: "2.1.0",
}, async () => {
    // All spans inside here carry:
    // - coalex.agent_id = "claims-processor"
    // - coalex.request_id = "claim-2024-0042"
    // - coalex.agent_version = "2.1.0"
    const result = await processClaim(claimData);
});

Nested contexts

You can nest coalex_context blocks. The inner context creates a new parent span as a child of the outer one:

with coalex.coalex_context(agent_id="orchestrator", request_id="req-100"):
    # Outer span: coalex.invocation (agent_id=orchestrator)

    with coalex.coalex_context(agent_id="sub-agent-a", request_id="req-100-a"):
        # Inner span: coalex.invocation (agent_id=sub-agent-a)
        # This is a child of the outer coalex.invocation span
        run_sub_agent_a()

    with coalex.coalex_context(agent_id="sub-agent-b", request_id="req-100-b"):
        run_sub_agent_b()

await coalexContext({ agentId: "orchestrator", requestId: "req-100" }, async () => {
    // Outer span: coalex.invocation (agentId=orchestrator)

    await coalexContext({ agentId: "sub-agent-a", requestId: "req-100-a" }, async () => {
        // Inner span: coalex.invocation (agentId=sub-agent-a)
        // This is a child of the outer coalex.invocation span
        await runSubAgentA();
    });

    await coalexContext({ agentId: "sub-agent-b", requestId: "req-100-b" }, async () => {
        await runSubAgentB();
    });
});

Generating request IDs

Use uuid4 to generate unique request IDs:

import uuid

with coalex.coalex_context(
    agent_id="my-agent",
    request_id=str(uuid.uuid4()),
):
    ...

Use crypto.randomUUID() to generate unique request IDs:

import { randomUUID } from "node:crypto";

await coalexContext({
    agentId: "my-agent",
    requestId: randomUUID(),
}, async () => {
    // ...
});

Notes

  • coalex_context() requires register() to have been called first (so a global TracerProvider exists). If register() has not been called, spans are created on a no-op tracer and silently discarded.
  • The context manager is not reentrant -- do not reuse a single coalex_context instance. Create a new one for each invocation.
  • The request_id does not need to be globally unique, but it should be unique within a reasonable time window (e.g., per user session or per API request).

API Reference

coalex.context.coalex_context 

coalex_context(
    *,
    agent_id: str,
    request_id: str | None = None,
    version: str | None = None,
    prompt_version: int | None = None,
    prompt_type: str | None = None,
) -> Generator[None, None, None]

Create a parent span with Coalex attributes for all child spans.

All spans created inside this context are descendants of a coalex.invocation span that carries the agent metadata. This makes every child span queryable by agent_id, request_id, and version in the Bronze layer.

Context vars are set before the span is created so that CoalexAttributePropagator can propagate attributes even across trace boundaries (e.g. when LangChain auto-instrumentors create new root spans with separate trace IDs).

Do not wrap a generator that is consumed after the enclosing function returns. A streaming HTTP response is the usual case: the body generator runs once the handler has already returned, so this block exits in a different Context from the one it entered, and the variables cannot be restored (COA-1980). Enter the context inside the generator instead, so its lifetime matches the work it describes::

async def body():
    with coalex_context(agent_id=..., request_id=...):
        async for token in agent.stream(...):
            yield token

Getting this wrong no longer raises — it leaks a context var and logs at debug — but the spans will be attributed to whatever context happens to be current, which is rarely what you meant.

prompt_version / prompt_type let an agent that fetches its prompt once at startup (outside any request trace) attribute each request to a prompt version. get_prompt() only stamps the trace it is called in, so the cache-at-startup pattern needs the version passed here per request (COA-1339). When omitted, any value already set by get_prompt() is left untouched.

Parameters:

Name Type Description Default
agent_id str

Identifier for the AI agent.

 required
request_id str | None

Unique request/invocation ID.

 None
version str | None

Agent version string.

 None
prompt_version int | None

Prompt Vault version driving this invocation.

 None
prompt_type str | None

Prompt type (e.g. "system").

 None
Source code in coalex/context.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
@contextlib.contextmanager
def coalex_context(
    *,
    agent_id: str,
    request_id: str | None = None,
    version: str | None = None,
    prompt_version: int | None = None,
    prompt_type: str | None = None,
) -> Generator[None, None, None]:
    """Create a parent span with Coalex attributes for all child spans.

    All spans created inside this context are descendants of a
    ``coalex.invocation`` span that carries the agent metadata.
    This makes every child span queryable by agent_id, request_id,
    and version in the Bronze layer.

    Context vars are set before the span is created so that
    ``CoalexAttributePropagator`` can propagate attributes even
    across trace boundaries (e.g. when LangChain auto-instrumentors
    create new root spans with separate trace IDs).

    **Do not wrap a generator that is consumed after the enclosing function returns.** A streaming
    HTTP response is the usual case: the body generator runs once the handler has already returned,
    so this block exits in a different Context from the one it entered, and the variables cannot be
    restored (COA-1980). Enter the context *inside* the generator instead, so its lifetime matches
    the work it describes::

        async def body():
            with coalex_context(agent_id=..., request_id=...):
                async for token in agent.stream(...):
                    yield token

    Getting this wrong no longer raises — it leaks a context var and logs at debug — but the spans
    will be attributed to whatever context happens to be current, which is rarely what you meant.

    ``prompt_version`` / ``prompt_type`` let an agent that fetches its prompt
    once at startup (outside any request trace) attribute each request to a
    prompt version. ``get_prompt()`` only stamps the trace it is called in, so
    the cache-at-startup pattern needs the version passed here per request
    (COA-1339). When omitted, any value already set by ``get_prompt()`` is left
    untouched.

    Args:
        agent_id: Identifier for the AI agent.
        request_id: Unique request/invocation ID.
        version: Agent version string.
        prompt_version: Prompt Vault version driving this invocation.
        prompt_type: Prompt type (e.g. ``"system"``).
    """
    token_a = _coalex_agent_id.set(agent_id)
    token_r = _coalex_request_id.set(request_id)
    token_v = _coalex_agent_version.set(version)
    # Only touch the prompt context vars when explicitly given, so a value set by
    # get_prompt() isn't clobbered when the caller doesn't pass one.
    token_pv = _coalex_prompt_version.set(prompt_version) if prompt_version is not None else None
    token_pt = _coalex_prompt_type.set(prompt_type) if prompt_type is not None else None

    try:
        tracer = trace.get_tracer("coalex")

        attrs: dict[str, str | int] = {"coalex.agent_id": agent_id}
        if request_id is not None:
            attrs["coalex.request_id"] = request_id
        if version is not None:
            attrs["coalex.agent_version"] = version
        if prompt_version is not None:
            attrs["coalex.prompt_version"] = prompt_version
        if prompt_type is not None:
            attrs["coalex.prompt_type"] = prompt_type

        with tracer.start_as_current_span("coalex.invocation", attributes=attrs) as span:
            yield
            span.set_status(Status(StatusCode.OK))
    finally:
        _restore(_coalex_agent_id, token_a)
        _restore(_coalex_request_id, token_r)
        _restore(_coalex_agent_version, token_v)
        _restore(_coalex_prompt_version, token_pv)
        _restore(_coalex_prompt_type, token_pt)