Skip to content

Batched embeddings drop types absent from the first response #796

Description

@OllieinCanada

Problem

merge_embed_responses() determines the embedding fields to merge by inspecting only the first EmbeddingsByTypeEmbedResponse.

If the first batch has float embeddings while a later batch also contains int8, the later int8 values are silently omitted from the merged response.

Reproduction

first = EmbeddingsByTypeEmbedResponse(
    response_type="embeddings_by_type",
    id="1",
    embeddings=EmbedByTypeResponseEmbeddings(float_=[[1.0, 2.0]]),
)
second = EmbeddingsByTypeEmbedResponse(
    response_type="embeddings_by_type",
    id="2",
    embeddings=EmbedByTypeResponseEmbeddings(
        float_=[[3.0, 4.0]],
        int8=[[3, 4]],
    ),
)

merged = merge_embed_responses([first, second])
assert merged.embeddings.int8 == [[3, 4]]  # currently None

Expected behavior

The merger should preserve every embedding type returned by any batch while continuing to skip None values.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions