github-featured/microsoft--agent-framework
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Python: Fix Cosmos memory provider calling renamed add_cosmos toolkit API (#7635)

Browse Source 
* Python: Fix Cosmos memory provider calling renamed add_cosmos toolkit API

The Agent Memory Toolkit renamed AsyncCosmosMemoryClient.add_cosmos to
upsert_memory with an identical signature. The provider declares
azure-cosmos-agent-memory>=0.2.0b3 with no upper bound, so a resolved
install can expose either name. after_run swallows write errors and only
logs a warning, so on a post-rename toolkit the agent turn still looks
successful while long-term memory silently stops receiving turns.

Resolve the write method once per after_run, preferring upsert_memory and
falling back to add_cosmos, so both ends of the declared range keep working.
Same treatment for the emulator test's direct seed call.

Fixes #7633

* Ponytail comment erased

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* Clarify TODO comment regarding memory method rename

Updated TODO comment to include author and clarify context , to resolve linting error

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Atharva Vichare
2026-08-13 21:25:38 +00:00
committed by GitHub
parent 9a06fa3f42
commit 9645d33cde
3 changed files with 39 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py
+7 -3
View File
@@ -435,13 +435,17 @@ class CosmosMemoryContextProvider(ContextProvider):
user_id = self._resolve_user_id(state, session)
thread_id = state.get("thread_id") or session.session_id or "default"
# TODO(atty57): The toolkit renamed add_cosmos -> upsert_memory (same kwargs); accept either
# until the declared azure-cosmos-agent-memory floor is past the rename, then inline it.
write_turn = getattr(self.memory_client, "upsert_memory", None) or self.memory_client.add_cosmos
try:
# Store input messages (skip empty/whitespace-only content to avoid junk turns)
for msg in context.input_messages:
if hasattr(msg, "role") and hasattr(msg, "text") and msg.text and msg.text.strip():
role_value = getattr(msg.role, "value", None) or str(msg.role)
if role_value in {"user", "assistant", "system"}:
await self.memory_client.add_cosmos(
await write_turn(
user_id=user_id,
thread_id=thread_id,
role=self._ROLE_MAP.get(role_value, role_value),
@@ -454,7 +458,7 @@ class CosmosMemoryContextProvider(ContextProvider):
if hasattr(msg, "role") and hasattr(msg, "text") and msg.text and msg.text.strip():
role_value = getattr(msg.role, "value", None) or str(msg.role)
if role_value in {"user", "assistant", "system"}:
await self.memory_client.add_cosmos(
await write_turn(
user_id=user_id,
thread_id=thread_id,
role=self._ROLE_MAP.get(role_value, role_value),
@@ -462,7 +466,7 @@ class CosmosMemoryContextProvider(ContextProvider):
)
# Auto-extraction and processing:
# When auto_extract is True (default), add_cosmos() schedules cadence-aware background
# When auto_extract is True (default), the turn write schedules cadence-aware background
# processing (fact extraction, summaries, reconciliation) based on the configured
# thresholds (FACT_EXTRACTION_EVERY_N, DEDUP_EVERY_N, etc.), so no explicit
# process_now() call is needed. When auto_extract is False, those thresholds were
python/packages/azure-cosmos-memory/tests/test_context_provider.py
+30 -10
View File
@@ -52,7 +52,7 @@ def mock_memory_client() -> AsyncMock:
mock_client = AsyncMock()
mock_client.search_cosmos = AsyncMock(return_value=[])
mock_client.get_user_summary = AsyncMock(return_value=None)
mock_client.add_cosmos = AsyncMock()
mock_client.upsert_memory = AsyncMock()
mock_client.create_memory_store = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
@@ -468,8 +468,8 @@ class TestAfterRun:
agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert mock_memory_client.add_cosmos.await_count == 2
calls = mock_memory_client.add_cosmos.await_args_list
assert mock_memory_client.upsert_memory.await_count == 2
calls = mock_memory_client.upsert_memory.await_args_list
# Check input message stored
assert calls[0].kwargs["role"] == "user"
@@ -499,7 +499,7 @@ class TestAfterRun:
agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
stored_roles = [c.kwargs["role"] for c in mock_memory_client.add_cosmos.await_args_list]
stored_roles = [c.kwargs["role"] for c in mock_memory_client.upsert_memory.await_args_list]
assert stored_roles == ["user", "agent"]
# No raw "assistant" role should ever be sent to the toolkit.
assert "assistant" not in stored_roles
@@ -520,7 +520,7 @@ class TestAfterRun:
agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_memory_client.add_cosmos.await_args_list[0].kwargs
call_kwargs = mock_memory_client.upsert_memory.await_args_list[0].kwargs
assert call_kwargs["user_id"] == "user-456"
assert call_kwargs["thread_id"] == "thread-789"
@@ -541,8 +541,8 @@ class TestAfterRun:
)
# Only one message should be stored
assert mock_memory_client.add_cosmos.await_count == 1
call_kwargs = mock_memory_client.add_cosmos.await_args_list[0].kwargs
assert mock_memory_client.upsert_memory.await_count == 1
call_kwargs = mock_memory_client.upsert_memory.await_args_list[0].kwargs
assert call_kwargs["content"] == "Valid message"
async def test_skips_whitespace_only_messages(self, mock_memory_client: AsyncMock) -> None:
@@ -563,15 +563,35 @@ class TestAfterRun:
)
# Whitespace-only input and the whitespace-only response are both skipped.
assert mock_memory_client.add_cosmos.await_count == 1
call_kwargs = mock_memory_client.add_cosmos.await_args_list[0].kwargs
assert mock_memory_client.upsert_memory.await_count == 1
call_kwargs = mock_memory_client.upsert_memory.await_args_list[0].kwargs
assert call_kwargs["content"] == "Trimmed message"
async def test_falls_back_to_add_cosmos_on_older_toolkit(self) -> None:
"""Toolkit versions predating the upsert_memory rename still receive turns.
The declared azure-cosmos-agent-memory range spans both names, so a resolved
install can expose either one; picking neither would silently drop every turn.
"""
legacy_client = AsyncMock(spec=["add_cosmos", "search_cosmos", "get_user_summary"])
legacy_client.add_cosmos = AsyncMock()
provider = CosmosMemoryContextProvider(memory_client=legacy_client)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", contents=["Hello"])], session_id="s1")
await provider.after_run(
agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert legacy_client.add_cosmos.await_count == 1
assert legacy_client.add_cosmos.await_args_list[0].kwargs["content"] == "Hello"
async def test_storage_failure_logs_warning(
self, mock_memory_client: AsyncMock, caplog: pytest.LogCaptureFixture
) -> None:
"""Storage failures are logged but don't raise."""
mock_memory_client.add_cosmos.side_effect = Exception("Storage failed")
mock_memory_client.upsert_memory.side_effect = Exception("Storage failed")
provider = CosmosMemoryContextProvider(memory_client=mock_memory_client)
session = AgentSession(session_id="test-session")
python/packages/azure-cosmos-memory/tests/test_emulator.py
+2 -1
View File
@@ -183,7 +183,8 @@ class TestEmulatorVectorSearch:
# embeddings client). This lands in the memories container under the quantizedFlat
# vector index, without needing LLM extraction.
assert provider.memory_client is not None
await provider.memory_client.add_cosmos(
seed = getattr(provider.memory_client, "upsert_memory", None) or provider.memory_client.add_cosmos
await seed(
user_id=user_id,
thread_id=thread_id,
role="user",