Skip to content

openarmature.retrieval

The retrieval-provider capability.

The embedding + rerank provider protocols, their response types, and the bundled reference providers: an OpenAI-compatible embedding provider, the TEI embedding + rerank providers, and the Jina and Cohere embedding + rerank providers. Embedding and rerank are sibling surfaces on the same capability.

EmbeddingProvider

Bases: Protocol

The shape of any retrieval-provider embedding implementation.

Implementations are bound to a single embedding model identifier; switching models means constructing a new provider, not passing a different argument per call.

ready async 

ready() -> None

Verify the bound embedding model is reachable and serving.

embed async 

embed(
    input: Sequence[str],
    *,
    config: EmbeddingRuntimeConfig | None = None
) -> EmbeddingResponse

Embed input into one vector per string, in input order.

Parameters:

Name Type Description Default
input Sequence[str]

The strings to embed. Always a list, even for a single-string caller (wrap as a one-element list). Not mutated by the implementation.

 required
config EmbeddingRuntimeConfig | None

Optional per-call request parameters.

 None

Returns an :class:EmbeddingResponse whose vectors[i] is the embedding of input[i]; order is preserved, never permuted.

RerankProvider

Bases: Protocol

The shape of any retrieval-provider rerank implementation.

Implementations are bound to a single rerank model identifier; switching models means constructing a new provider, not passing a different argument per call.

ready async 

ready() -> None

Verify the bound rerank model is reachable and serving.

rerank async 

rerank(
    query: str,
    documents: Sequence[str],
    *,
    top_k: int | None = None,
    config: RerankRuntimeConfig | None = None
) -> RerankResponse

Score documents against query, sorted by relevance.

Parameters:

Name Type Description Default
query str

The query string the documents are scored against.

 required
documents Sequence[str]

The candidate documents. Always a list, even for a single-document caller (wrap as a one-element list). Not mutated by the implementation.

 required
top_k int | None

The maximum number of results to return. None means "all" (up to len(documents)).

 None
config RerankRuntimeConfig | None

Optional per-call request parameters.

 None

Returns a :class:RerankResponse whose results are sorted by relevance_score descending; each result's index keys back into the input documents list.

CohereEmbeddingProvider 

CohereEmbeddingProvider(
    *,
    base_url: str = "https://api.cohere.com",
    model: str,
    api_key: str | None = None,
    transport: AsyncBaseTransport | None = None,
    timeout: float = 60.0,
    genai_system: str = "cohere",
    populate_caller_metadata: bool = True
)

Cohere /v2/embed wire-shape embedding provider.

Construct with the bound embedding model and an optional API key + transport. base_url is the host root and defaults to the Cohere origin (https://api.cohere.com), overridable for a proxy / gateway. embed() posts to /v2/embed, chunk-and-stitching across Cohere's 96-input per-call cap when the input list is larger.

Cohere /v2/embed requires input_type on every request, so the mapping always sends a value: "query" becomes "search_query", "document" becomes "search_document", and an absent input_type becomes "search_document" (the bulk-indexing default). An unrecognized input_type is rejected before the request is sent. embedding_types always requests "float" (the mapping reads embeddings.float) and truncate: "NONE" is sent explicitly (an over-length input errors rather than being silently truncated); dimensions maps to Cohere's output_dimension when set. Other precisions (int8 / base64 / ...) ride the extras pass-through bag; Cohere's other input_type values (classification / ...) are outside OA's input_type value space and are not reachable through this mapping.

ready() verifies the bound model with a minimal one-input /v2/embed probe. The Cohere /v2/embed wire exposes no model-catalog probe, so there is a single universal probe.

aclose async 

aclose() -> None

Close the underlying HTTP client (releases the connection pool).

ready async 

ready() -> None

Verify the bound embedding model is reachable and serving.

embed async 

embed(
    input: Sequence[str],
    *,
    config: EmbeddingRuntimeConfig | None = None
) -> EmbeddingResponse

Embed input into one vector per string, in input order.

CohereRerankProvider 

CohereRerankProvider(
    *,
    base_url: str = "https://api.cohere.com",
    model: str,
    api_key: str | None = None,
    transport: AsyncBaseTransport | None = None,
    timeout: float = 60.0,
    genai_system: str = "cohere",
    populate_caller_metadata: bool = True
)

Cohere /v2/rerank wire-shape rerank provider.

Construct with the bound rerank model and an optional API key + transport. base_url is the host root and defaults to the Cohere origin (https://api.cohere.com), overridable for a proxy / gateway. rerank() posts to /v2/rerank.

ready() verifies the bound model with a minimal one-document /v2/rerank probe. The Cohere /v2/rerank wire exposes no model-catalog probe (unlike the OpenAI-compatible embedding surface), so there is a single universal probe.

aclose async 

aclose() -> None

Close the underlying HTTP client (releases the connection pool).

ready async 

ready() -> None

Verify the bound rerank model is reachable and serving.

rerank async 

rerank(
    query: str,
    documents: Sequence[str],
    *,
    top_k: int | None = None,
    config: RerankRuntimeConfig | None = None
) -> RerankResponse

Score documents against query, sorted by relevance.

JinaEmbeddingProvider 

JinaEmbeddingProvider(
    *,
    model: str,
    api_key: str,
    base_url: str = "https://api.jina.ai",
    transport: AsyncBaseTransport | None = None,
    timeout: float = 60.0,
    genai_system: str = "jina",
    populate_caller_metadata: bool = True
)

Jina /v1/embeddings wire-mapping embedding provider.

Construct with the bound embedding model and the required API key; the base_url defaults to the Jina endpoint (https://api.jina.ai, origin only -- override for a proxy / private gateway). embed() posts to /v1/embeddings.

ready() verifies the bound model with a minimal one-input /v1/embeddings probe.

aclose async 

aclose() -> None

Close the underlying HTTP client (releases the connection pool).

ready async 

ready() -> None

Verify the bound embedding model is reachable and serving.

embed async 

embed(
    input: Sequence[str],
    *,
    config: EmbeddingRuntimeConfig | None = None
) -> EmbeddingResponse

Embed input into one vector per string, in input order.

JinaRerankProvider 

JinaRerankProvider(
    *,
    model: str,
    api_key: str,
    base_url: str = "https://api.jina.ai",
    transport: AsyncBaseTransport | None = None,
    timeout: float = 60.0,
    genai_system: str = "jina",
    populate_caller_metadata: bool = True
)

Jina /v1/rerank wire-mapping rerank provider.

Construct with the bound rerank model and the required API key; the base_url defaults to the Jina endpoint (https://api.jina.ai, origin only -- override for a proxy / private gateway). rerank() posts to /v1/rerank (a single request -- Jina batches server-side, so there is no client-side chunk-and-stitch).

ready() verifies the bound model with a minimal one-document /v1/rerank probe.

aclose async 

aclose() -> None

Close the underlying HTTP client (releases the connection pool).

ready async 

ready() -> None

Verify the bound rerank model is reachable and serving.

rerank async 

rerank(
    query: str,
    documents: Sequence[str],
    *,
    top_k: int | None = None,
    config: RerankRuntimeConfig | None = None
) -> RerankResponse

Score documents against query, sorted by relevance.

OpenAIEmbeddingProvider 

OpenAIEmbeddingProvider(
    *,
    base_url: str = "https://api.openai.com",
    model: str,
    api_key: str | None = None,
    transport: AsyncBaseTransport | None = None,
    timeout: float = 60.0,
    genai_system: str = "openai",
    readiness_probe: Literal[
        "embed", "models", "both"
    ] = "embed",
    query_prefix: str | None = None,
    document_prefix: str | None = None,
    populate_caller_metadata: bool = True
)

OpenAI /v1/embeddings wire-compatible embedding provider.

Construct with the bound embedding model and an optional API key + transport. base_url is the host root and defaults to the OpenAI origin (https://api.openai.com), overridable for any OpenAI-compatible backend. embed() posts to /v1/embeddings.

The optional query_prefix / document_prefix bind the client-side asymmetric prefixes -- off by default (pure-symmetric OpenAI). When bound (for an asymmetric model served behind a compatible endpoint), input_type selects which prefix to prepend to each input before sending, since the OpenAI wire carries no query/document field.

ready() verifies the bound model per the readiness_probe argument:

  • "embed" (default): a one-input /v1/embeddings probe. Works against any OpenAI-compatible backend, including ones that do not serve the /v1/models catalog (e.g. TEI's OpenAI surface).
  • "models": a GET /v1/models catalog check. Cheaper (no embed billed), but requires the endpoint to serve the catalog.
  • "both": the catalog check, then the embed probe.

aclose async 

aclose() -> None

Close the underlying HTTP client (releases the connection pool).

ready async 

ready() -> None

Verify the bound embedding model is reachable and serving.

embed async 

embed(
    input: Sequence[str],
    *,
    config: EmbeddingRuntimeConfig | None = None
) -> EmbeddingResponse

Embed input into one vector per string, in input order.

TeiEmbeddingProvider 

TeiEmbeddingProvider(
    *,
    base_url: str,
    model: str,
    api_key: str | None = None,
    transport: AsyncBaseTransport | None = None,
    timeout: float = 60.0,
    genai_system: str = "tei",
    chunk_size: int = 32,
    input_type_prompt_map: Mapping[str, str] | None = None,
    query_prefix: str | None = None,
    document_prefix: str | None = None,
    populate_caller_metadata: bool = True
)

TEI /embed wire-mapping embedding provider.

Construct with a base URL (the TEI embedding instance root), the bound embedding model, and optionally an input_type_prompt_map binding input_type to TEI's native prompt_name (server-side prompts) and/or client-side query_prefix / document_prefix strings (the fallback for models without configured prompts). embed() posts to /embed.

ready() verifies the bound model with a minimal one-input /embed probe (TEI serves no model catalog).

aclose async 

aclose() -> None

Close the underlying HTTP client (releases the connection pool).

ready async 

ready() -> None

Verify the bound embedding model is reachable and serving.

embed async 

embed(
    input: Sequence[str],
    *,
    config: EmbeddingRuntimeConfig | None = None
) -> EmbeddingResponse

Embed input into one vector per string, in input order.

TeiRerankProvider 

TeiRerankProvider(
    *,
    base_url: str,
    model: str,
    api_key: str | None = None,
    transport: AsyncBaseTransport | None = None,
    timeout: float = 60.0,
    genai_system: str = "tei",
    chunk_size: int = 32,
    populate_caller_metadata: bool = True
)

TEI /rerank wire-mapping rerank provider.

Construct with a base URL (the TEI reranker instance root), the bound rerank model, and chunk_size (TEI's max-client-batch-size, default 32). rerank() posts to /rerank, chunk-and-stitching across chunk_size when the candidate pool is larger.

ready() verifies the bound model with a minimal one-document /rerank probe (TEI serves no model catalog).

aclose async 

aclose() -> None

Close the underlying HTTP client (releases the connection pool).

ready async 

ready() -> None

Verify the bound rerank model is reachable and serving.

rerank async 

rerank(
    query: str,
    documents: Sequence[str],
    *,
    top_k: int | None = None,
    config: RerankRuntimeConfig | None = None
) -> RerankResponse

Score documents against query, sorted by relevance.

Chunk-and-stitches across chunk_size: one /rerank request per consecutive <= chunk_size slice, absolute-position re-basing, a global re-sort by score descending, then top_k.

EmbeddingResponse

Bases: BaseModel

The result of an EmbeddingProvider.embed() call.

Attributes:

Name Type Description
vectors list[list[float]]

One vector (a list of floats) per input string, in the order the inputs were supplied. len(vectors) equals the input length.
model str

The model identifier the provider returned; may be more specific than the bound identifier.
usage EmbeddingUsage | None

The token record, or None when the provider reports no usage object (e.g. TEI /embed). Never a fabricated record.
response_id str | None

The provider-returned response id when present; None otherwise.
dimensions int

The output vector dimensionality; equals the length of each inner vector.
raw dict[str, Any] | list[Any]

The parsed provider response verbatim -- a dict or a list, matching the provider's top-level JSON shape -- populated on every successful return (a chunked call carries the list of per-request responses in request order).

EmbeddingRuntimeConfig

Bases: BaseModel

Per-call embedding request parameters.

from_partial classmethod 

from_partial(**kwargs: Any) -> EmbeddingRuntimeConfig

Construct a config, dropping kwargs whose value is None.

EmbeddingUsage

Bases: BaseModel

Token-accounting record for an embedding call.

Carries input_tokens only; an embedding call has no output tokens (vectors are not tokens).

RerankResponse

Bases: BaseModel

The result of a RerankProvider.rerank() call.

Attributes:

Name Type Description
results list[ScoredDocument]

The scored documents sorted by relevance_score descending (most relevant first). Each entry's index keys back to the original input documents list.
model str

The model identifier the provider returned; may be more specific than the bound identifier.
usage RerankUsage | None

The usage record, or None when the provider reports no usage object.
response_id str | None

The provider-returned response id when present; None otherwise.
raw dict[str, Any] | list[Any]

The parsed provider response verbatim -- a dict or a list, matching the provider's top-level JSON shape -- populated on every successful return (a chunked call carries the list of per-request responses in request order).

RerankRuntimeConfig

Bases: BaseModel

Per-call rerank request parameters.

from_partial classmethod 

from_partial(**kwargs: Any) -> RerankRuntimeConfig

Construct a config, dropping kwargs whose value is None.

RerankUsage

Bases: BaseModel

Token-accounting record for a rerank call.

Both fields default to None and are individually nullable: a provider may surface one figure and not the other (Cohere reports search_units but no token count; Voyage AI reports input_tokens).

ScoredDocument

Bases: BaseModel

A single scored result entry in a RerankResponse.

Attributes:

Name Type Description
index int

The 0-based position of this document in the original input documents list. Load-bearing for caller-side lookup: documents[result.index] maps a result back to its input.
relevance_score float

The provider-assigned relevance score; higher = more relevant. Provider-specific scale (not normalized here).
document str | None

The echoed document text when the provider returns it; None otherwise. Never fabricated from the input documents.

validate_embedding_input 

validate_embedding_input(input: Sequence[str]) -> None

Validate the input list before sending.

Raises :class:ProviderInvalidRequest when the input list is empty.

validate_embedding_response 

validate_embedding_response(
    vectors: Sequence[Sequence[float]], input_count: int
) -> int

Validate the response invariants and return the dimensionality.

Raises :class:ProviderInvalidResponse when the vector count does not match the input count, when the response carries no vectors, or when the vectors are not all the same length. Returns the (consistent) dimensionality on success.

validate_rerank_input 

validate_rerank_input(
    query: str, documents: Sequence[str], top_k: int | None
) -> None

Validate the rerank inputs before sending.

Raises :class:ProviderInvalidRequest when the query is empty, the documents list is empty, or top_k is supplied and not positive. top_k may exceed len(documents); that is allowed.

validate_rerank_response 

validate_rerank_response(
    results: Sequence[ScoredDocument],
    document_count: int,
    top_k: int | None,
) -> None

Validate the rerank response invariants.

Raises :class:ProviderInvalidResponse when a result's index is out of range for the input documents, when an index appears twice, or when top_k is supplied and the provider returned more results than requested.