From 5b233c5e4d7f8b09522e796183959d94a73dd38f Mon Sep 17 00:00:00 2001 From: Zhibo Lin <147509942+LE0-Lin@users.noreply.github.com> Date: Mon, 5 Oct 2026 23:19:30 +0800 Subject: [PATCH] Python: Fix in-memory UUID key lookup and deletion Signed-off-by: Zhibo Lin <147509942+LE0-Lin@users.noreply.github.com> --- .../core/agent_framework/_in_memory.py | 6 ++-- .../core/tests/core/test_in_memory.py | 35 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_in_memory.py b/python/packages/core/agent_framework/_in_memory.py index cc8f8046cb7..4e0da04541e 100644 --- a/python/packages/core/agent_framework/_in_memory.py +++ b/python/packages/core/agent_framework/_in_memory.py @@ -303,7 +303,9 @@ async def _inner_get( ) -> Sequence[Any] | None: self._require_collection() if keys is not None: - return [deepcopy(self._state.records[key]) for key in keys if key in self._state.records] + # Shared record serialization stores UUID keys as strings. + normalized_keys = (str(key) if isinstance(key, UUID) else key for key in keys) + return [deepcopy(self._state.records[key]) for key in normalized_keys if key in self._state.records] records = list(self._state.records.values()) if filter is not None: _validate_in_memory_filter(filter, self.definition) @@ -333,7 +335,7 @@ async def _inner_delete( ) -> None: self._require_collection() for key in keys: - self._state.records.pop(key, None) + self._state.records.pop(str(key) if isinstance(key, UUID) else key, None) async def _inner_search( self, diff --git a/python/packages/core/tests/core/test_in_memory.py b/python/packages/core/tests/core/test_in_memory.py index 1253a189519..624c722840f 100644 --- a/python/packages/core/tests/core/test_in_memory.py +++ b/python/packages/core/tests/core/test_in_memory.py @@ -9,6 +9,7 @@ from collections.abc import Iterator, Sequence from dataclasses import dataclass from typing import Annotated, Any, cast +from uuid import UUID import pytest @@ -52,6 +53,12 @@ class Document: VectorStoreField("vector", dimensions=2, distance_function="cosine_similarity"), ] = None + @vectorstoremodel(collection_name="uuid_documents") + @dataclass + class UUIDDocument: + id: Annotated[UUID, VectorStoreField("key")] + text: Annotated[str, VectorStoreField("data")] + DOCUMENTS = ( Document("one", "Luxury hotel", "travel", 5, ["wifi", "pool"], "featured", [1.0, 0.0]), @@ -126,6 +133,34 @@ async def test_in_memory_collection_requires_creation() -> None: await collection.get() +@pytest.mark.parametrize("as_strings", [False, True]) +async def test_in_memory_get_uuid_keys(as_strings: bool) -> None: + collection: InMemoryCollection[UUID | str, UUIDDocument] = InMemoryCollection(UUIDDocument) + records = [UUIDDocument(UUID(int=1), "first"), UUIDDocument(UUID(int=2), "second")] + await collection.ensure_collection_exists() + await collection.upsert(records, generate_vectors=False) + requested: list[UUID | str] = [records[1].id, UUID(int=3), records[0].id] + if as_strings: + requested = [str(key) for key in requested] + + assert await collection.get(requested) == [records[1], records[0]] + + +@pytest.mark.parametrize("as_strings", [False, True]) +async def test_in_memory_delete_uuid_keys(as_strings: bool) -> None: + collection: InMemoryCollection[UUID | str, UUIDDocument] = InMemoryCollection(UUIDDocument) + records = [UUIDDocument(UUID(int=1), "first"), UUIDDocument(UUID(int=2), "second")] + await collection.ensure_collection_exists() + await collection.upsert(records, generate_vectors=False) + requested: list[UUID | str] = [records[0].id, UUID(int=3)] + if as_strings: + requested = [str(key) for key in requested] + + await collection.delete(requested) + + assert await collection.get() == [records[1]] + + async def test_in_memory_crud_listing_ordering_and_defensive_copies() -> None: collection = await _create_collection()