Python: Bug: ChromaCollection mishandles storage_name on the key and vector fields, depending on whether the record is a dict, dataclass or pydantic model #14755

Description

ChromaCollection handles a field with a storage_name in four different ways, depending on the record type and the field:

• A dict or dataclass record with a storage_name on the key field fails to upsert: VectorStoreModelSerializationException: Error serializing records: 'id'.
• A pydantic record with a storage_name on the vector field fails to upsert: Error serializing records: 'doc_vector'.
• A dict or dataclass record with a storage_name on the vector field upserts, but get and search return vector=None even with include_vectors=True.
• A pydantic record is stored with its metadata under the declared field names. Filters only accept storage names, so a filter on a field with a storage_name finds nothing.

The base class hands the connector a dict keyed by storage name for dict and dataclass records, and keyed by field name for pydantic models. The Chroma serializer and deserializer assume one shape for some fields and the other shape for the rest.

Reproduction

I ran this on current main (df1cb33), Python 3.13, chromadb 1.5.9 with chromadb.EphemeralClient(), so no server is involved.

import asyncio
import uuid
import warnings
from dataclasses import dataclass

warnings.filterwarnings("ignore")
import chromadb
from pydantic import BaseModel

from semantic_kernel.connectors.chroma import ChromaCollection
from semantic_kernel.data.vector import VectorStoreCollectionDefinition, VectorStoreField

CLIENT = chromadb.EphemeralClient()


def definition(key_alias=None, vector_alias=None):
    return VectorStoreCollectionDefinition(
        fields=[
            VectorStoreField("key", name="id", type="str", storage_name=key_alias),
            VectorStoreField("data", name="text", type="str", storage_name="doc_text", is_indexed=True),
            VectorStoreField("vector", name="vector", dimensions=2, type="float", storage_name=vector_alias),
        ]
    )


@dataclass
class DC:
    id: str
    text: str
    vector: list[float] | None = None


class PM(BaseModel):
    id: str
    text: str
    vector: list[float] | None = None


def record(kind):
    values = {"id": "k1", "text": "a", "vector": [1.0, 0.0]}
    return values if kind is dict else kind(**values)


def show(rec):
    values = rec if isinstance(rec, dict) else vars(rec) if isinstance(rec, DC) else rec.model_dump()
    return {k: ([float(x) for x in v] if k == "vector" and v is not None else v) for k, v in values.items()}


async def roundtrip(kind, defn, label):
    collection = ChromaCollection(
        record_type=kind, definition=defn, collection_name=f"t-{uuid.uuid4().hex}", client=CLIENT
    )
    await collection.ensure_collection_exists()
    try:
        await collection.upsert(record(kind))
        print(f"{kind.__name__:5s} {label}: get -> {show(await collection.get('k1', include_vectors=True))}")
    except Exception as ex:
        print(f"{kind.__name__:5s} {label}: {type(ex).__name__}: {ex}")


async def filtered(kind):
    collection = ChromaCollection(
        record_type=kind, definition=definition(), collection_name=f"t-{uuid.uuid4().hex}", client=CLIENT
    )
    await collection.ensure_collection_exists()
    await collection.upsert(record(kind))
    results = await collection.search(vector=[1.0, 0.0], filter=lambda x: x.doc_text == "a")
    print(f"{kind.__name__:5s} filter on doc_text: {len([r async for r in results.results])} result(s)")


async def main():
    for kind in (dict, DC, PM):
        await roundtrip(kind, definition(vector_alias="doc_vector"), "vector storage_name")
    for kind in (dict, DC, PM):
        await roundtrip(kind, definition(key_alias="point_id"), "key storage_name   ")
    for kind in (dict, DC, PM):
        await filtered(kind)


asyncio.run(main())

Output on main:

dict  vector storage_name: get -> {'id': 'k1', 'text': 'a', 'vector': None}
DC    vector storage_name: get -> {'id': 'k1', 'text': 'a', 'vector': None}
PM    vector storage_name: VectorStoreModelSerializationException: Error serializing records: 'doc_vector'
dict  key storage_name   : VectorStoreModelSerializationException: Error serializing records: 'id'
DC    key storage_name   : VectorStoreModelSerializationException: Error serializing records: 'id'
PM    key storage_name   : get -> {'id': 'k1', 'text': 'a', 'vector': [1.0, 0.0]}
dict  filter on doc_text: 1 result(s)
DC    filter on doc_text: 1 result(s)
PM    filter on doc_text: 0 result(s)

The data field text has a storage_name in every run. That alone works for dict and dataclass records.

Cause

In python/semantic_kernel/connectors/chroma.py:

• _serialize_dicts_to_store_models reads the key with record[self.definition.key_name] (declared name) and the vector with record[vector_field.storage_name or vector_field.name] (storage name). Each is right for only one of the two record shapes. Everything else stays in the metadata under the key it arrived with, so pydantic metadata keeps the declared names.
• _deserialize_store_models_to_dicts puts the key and the vector back under the declared names. For dict and dataclass records the data model step reads record.get(field.storage_name or field.name), so a vector with a storage_name is lost.

The same kind of mismatch exists in the Qdrant connector, see #14754.

Expected

Upsert, get and search behave the same for dict, dataclass and pydantic records, with or without storage_name. The metadata uses the storage names, which is what filters use.

Possible fix

I have a branch ready: https://github.com/VANDRANKI/semantic-kernel/tree/fix/chroma-storage-name-record-keys

• _serialize_dicts_to_store_models renames declared names to storage names in a copy of the record (only when the storage name is not already there), and then reads the key and the vector by storage name.
• _deserialize_store_models_to_dicts returns the key and the vector under their storage names. The data model step maps them back for pydantic records.

It adds 6 tests to tests/unit/connectors/memory/test_chroma.py (dict, dataclass and pydantic, each with and without a key storage_name). They upsert, get, search with vectors and filter on a storage name, on an ephemeral Chroma client. All 6 fail on main. The 26 tests in that file pass on the branch, and ruff check and format are clean.

I can't open a pull request from this account right now, so the compare link is main...VANDRANKI:semantic-kernel:fix/chroma-storage-name-record-keys