Python: Fix in-memory UUID key lookup and deletion

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,

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()