@embedding_span

Wraps an embedding function with an OpenInference EMBEDDING span. Captures the input text, model name, and resulting embedding vector as span attributes.

---

Decorator / Wrapper

Signature

PythonTypeScript

def embedding_span(
    name: str | None = None,
    model_name: str | None = None,
    query_arg: str = "text",
) -> Callable

interface EmbeddingSpanOptions {
    name?: string;
    modelName?: string;
}

function embeddingSpan<A extends unknown[], R>(
    options: EmbeddingSpanOptions,
    fn: (...args: A) => R | Promise<R>,
): (...args: A) => Promise<R>

Parameters

Parameter	Type	Default	Description
name	str \| None	None	Span name. Defaults to the decorated function's name.
model_name	str \| None	None	Embedding model name to record as a span attribute (e.g., "text-embedding-3-small").
query_arg	str	"text"	Name of the keyword argument containing the text to embed. Falls back to the first positional str argument.

Span Attributes

Attribute	Value
openinference.span.kind	"EMBEDDING"
embedding.model_name	Model name (set when model_name parameter is provided)
input.value	The text being embedded (from query_arg kwarg or first str positional argument)
embedding.embeddings	JSON-encoded embedding vector (set when the function returns list[float])

Embedding vector format

The embedding.embeddings attribute is a JSON string with the format:

[{"embedding.vector": [0.1, 0.2, ...], "embedding.text": "the input text"}]

Example

PythonTypeScript

from coalex.ext.embedding import embedding_span
from openai import OpenAI

client = OpenAI()

@embedding_span(name="embed_query", model_name="text-embedding-3-small")
def embed(text: str) -> list[float]:
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=text,
    )
    return response.data[0].embedding

# Usage
vector = embed(text="What is the refund policy?")
# Span "embed_query" is emitted with:
#   openinference.span.kind = "EMBEDDING"
#   embedding.model_name = "text-embedding-3-small"
#   input.value = "What is the refund policy?"
#   embedding.embeddings = '[{"embedding.vector": [0.1, ...], "embedding.text": "What is the refund policy?"}]'

import { embeddingSpan } from "@coalex-ai/sdk/ext";
import OpenAI from "openai";

const client = new OpenAI();

const embed = embeddingSpan(
    { name: "embed_query", modelName: "text-embedding-3-small" },
    async (text: string): Promise<number[]> => {
        const response = await client.embeddings.create({
            model: "text-embedding-3-small",
            input: text,
        });
        return response.data[0].embedding;
    },
);

// Usage
const vector = await embed("What is the refund policy?");
// Span "embed_query" is emitted with:
//   openinference.span.kind = "EMBEDDING"
//   embedding.model_name = "text-embedding-3-small"
//   input.value = "What is the refund policy?"
//   embedding.embeddings = '[{"embedding.vector": [0.1, ...], "embedding.text": "What is the refund policy?"}]'

Async Support

PythonTypeScript

@embedding_span(name="async_embed", model_name="text-embedding-3-small")
async def embed_async(text: str) -> list[float]:
    response = await async_client.embeddings.create(
        model="text-embedding-3-small",
        input=text,
    )
    return response.data[0].embedding

const embedAsync = embeddingSpan(
    { name: "async_embed", modelName: "text-embedding-3-small" },
    async (text: string): Promise<number[]> => {
        const response = await asyncClient.embeddings.create({
            model: "text-embedding-3-small",
            input: text,
        });
        return response.data[0].embedding;
    },
);

---

EmbeddingSpan Context Manager / Class

For fine-grained control over when attributes are set on the span, use the EmbeddingSpan context manager (Python) or class (TypeScript).

Signature

PythonTypeScript

class EmbeddingSpan:
    def __init__(
        self,
        name: str = "embedding",
        model_name: str | None = None,
    ) -> None: ...

    def set_text(self, text: str) -> None: ...
    def set_vector(self, vector: list[float], text: str = "") -> None: ...

class EmbeddingSpan {
    constructor(name?: string, modelName?: string);
    start(): void;
    setText(text: string): void;
    setVector(vector: number[], text?: string): void;
    end(): void;
}

Constructor Parameters

Parameter	Type	Default	Description
name	str / string	"embedding"	Span name.
model_name / modelName	str \| None / string?	None / undefined	Embedding model name to record.

Methods

Python	TypeScript	Description
set_text(text)	setText(text)	Set the input.value attribute on the span.
set_vector(vector, text="")	setVector(vector, text?)	Set the embedding.embeddings attribute with the vector and associated text.
--	start()	Start the span (TypeScript only; Python uses with statement).
--	end()	End the span (TypeScript only; Python uses with statement).

Example

PythonTypeScript

from coalex.ext.embedding import EmbeddingSpan

with EmbeddingSpan("embed_document", model_name="text-embedding-3-small") as e:
    text = "Metformin is used for type 2 diabetes management."
    e.set_text(text)

    # Call your embedding API
    vector = embedding_api.embed(text)
    e.set_vector(vector, text)

import { EmbeddingSpan } from "@coalex-ai/sdk/ext";

const e = new EmbeddingSpan("embed_document", "text-embedding-3-small");
e.start();
try {
    const text = "Metformin is used for type 2 diabetes management.";
    e.setText(text);

    // Call your embedding API
    const vector = await embeddingApi.embed(text);
    e.setVector(vector, text);
} finally {
    e.end();
}

---

Batch Embedding Example

from coalex.ext.embedding import embedding_span, EmbeddingSpan

@embedding_span(name="embed_single", model_name="text-embedding-3-small")
def embed_text(text: str) -> list[float]:
    response = client.embeddings.create(model="text-embedding-3-small", input=text)
    return response.data[0].embedding

# For batch embedding, use the context manager for more control
def embed_batch(texts: list[str]) -> list[list[float]]:
    vectors = []
    for text in texts:
        with EmbeddingSpan("embed_batch_item", model_name="text-embedding-3-small") as e:
            e.set_text(text)
            response = client.embeddings.create(model="text-embedding-3-small", input=text)
            vector = response.data[0].embedding
            e.set_vector(vector, text)
            vectors.append(vector)
    return vectors

---

With Retrieval Pipeline

A common pattern is to embed queries before searching a vector store:

import coalex
from coalex.ext.embedding import embedding_span
from coalex.ext.retrieval import retrieval_span, Document

@embedding_span(name="embed_query", model_name="text-embedding-3-small")
def embed_query(text: str) -> list[float]:
    response = client.embeddings.create(model="text-embedding-3-small", input=text)
    return response.data[0].embedding

@retrieval_span(name="vector_search", query_arg="query")
def search(query: str) -> list[Document]:
    vector = embed_query(text=query)
    results = vector_store.search(vector, top_k=10)
    return [Document(content=r.text, id=r.id, score=r.score) for r in results]

with coalex.coalex_context(agent_id="search-agent"):
    docs = search(query="diabetes treatment guidelines")
    # Trace contains:
    #   coalex.invocation
    #     └── vector_search (RETRIEVER)
    #           └── embed_query (EMBEDDING)

---

API Reference

coalex.ext.embedding

Embedding span decorator and context manager.

Classes

EmbeddingSpan

Context manager for instrumenting an embedding call as an EMBEDDING span.

Example::

with EmbeddingSpan("embed", model_name="text-embedding-3-small") as e:
    e.set_text(text)
    vector = embed(text)
    e.set_vector(vector, text)

Source code in coalex/ext/embedding.py

class EmbeddingSpan:
    """Context manager for instrumenting an embedding call as an EMBEDDING span.

    Example::

        with EmbeddingSpan("embed", model_name="text-embedding-3-small") as e:
            e.set_text(text)
            vector = embed(text)
            e.set_vector(vector, text)
    """

    def __init__(self, name: str = "embedding", model_name: str | None = None) -> None:
        self._name = name
        self._model_name = model_name
        self._span: Any = None
        self._ctx: Any = None

    def __enter__(self) -> EmbeddingSpan:
        tracer = _get_tracer()
        parent_ctx = _get_parent_context()
        self._ctx = tracer.start_as_current_span(self._name, context=parent_ctx)
        self._span = self._ctx.__enter__()
        self._span.set_attribute("openinference.span.kind", "EMBEDDING")
        if self._model_name:
            self._span.set_attribute("embedding.model_name", self._model_name)
        return self

    def set_text(self, text: str) -> None:
        self._span.set_attribute("input.value", text)

    def set_vector(self, vector: list[float], text: str = "") -> None:
        encoded = json.dumps([{"embedding.vector": vector, "embedding.text": text}])
        self._span.set_attribute("embedding.embeddings", encoded)

    def __exit__(self, *exc_info: Any) -> bool | None:
        return self._ctx.__exit__(*exc_info)

Functions

embedding_span

embedding_span(
    name: str | None = None,
    model_name: str | None = None,
    query_arg: str = "text",
) -> Callable[[F], F]

Decorator that wraps an embedding function with an OpenInference EMBEDDING span.

Supports both sync and async functions. The span captures: - openinference.span.kind = EMBEDDING - embedding.model_name — if provided - input.value — text being embedded - embedding.embeddings — JSON-encoded vector if function returns list[float]

Parameters:

Name	Type	Description	Default
name	str | None	Span name. Defaults to the function name.	None
model_name	str | None	Embedding model name to record.	None
query_arg	str	Name of the kwarg containing the text to embed.	'text'

Source code in coalex/ext/embedding.py

def embedding_span(
    name: str | None = None,
    model_name: str | None = None,
    query_arg: str = "text",
) -> Callable[[F], F]:
    """Decorator that wraps an embedding function with an OpenInference EMBEDDING span.

    Supports both sync and async functions. The span captures:
    - ``openinference.span.kind`` = ``EMBEDDING``
    - ``embedding.model_name`` — if provided
    - ``input.value`` — text being embedded
    - ``embedding.embeddings`` — JSON-encoded vector if function returns ``list[float]``

    Args:
        name: Span name. Defaults to the function name.
        model_name: Embedding model name to record.
        query_arg: Name of the kwarg containing the text to embed.
    """

    def decorator(fn: F) -> F:
        span_name = name or fn.__name__

        if asyncio.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _get_tracer()
                parent_ctx = _get_parent_context()
                with tracer.start_as_current_span(span_name, context=parent_ctx) as span:
                    span.set_attribute("openinference.span.kind", "EMBEDDING")
                    if model_name:
                        span.set_attribute("embedding.model_name", model_name)
                    text = kwargs.get(query_arg) or next((a for a in args if isinstance(a, str)), None)
                    if text:
                        span.set_attribute("input.value", str(text))
                    result = await fn(*args, **kwargs)
                    if isinstance(result, list) and result and isinstance(result[0], float):
                        encoded = json.dumps([{"embedding.vector": result, "embedding.text": str(text or "")}])
                        span.set_attribute("embedding.embeddings", encoded)
                    return result

            return async_wrapper  # type: ignore[return-value]

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            tracer = _get_tracer()
            parent_ctx = _get_parent_context()
            with tracer.start_as_current_span(span_name, context=parent_ctx) as span:
                span.set_attribute("openinference.span.kind", "EMBEDDING")
                if model_name:
                    span.set_attribute("embedding.model_name", model_name)
                text = kwargs.get(query_arg) or next((a for a in args if isinstance(a, str)), None)
                if text:
                    span.set_attribute("input.value", str(text))
                result = fn(*args, **kwargs)
                if isinstance(result, list) and result and isinstance(result[0], float):
                    encoded = json.dumps([{"embedding.vector": result, "embedding.text": str(text or "")}])
                    span.set_attribute("embedding.embeddings", encoded)
                return result

        return wrapper  # type: ignore[return-value]

    return decorator