code improvements

Author: Aayush Kataria

Samples/Demo.ipynb

@@ -2879,7 +2879,7 @@
     "stored = memory.get_user_summary(user_id=user_id)\n",
     "if stored:\n",
     "    print(\"User Summary for\", user_id)\n",
-    "    print(stored[0][\"content\"])\n",
+    "    print(stored[\"content\"])\n",
     "else:\n",
     "    print(\"No user summary found — run the generate_user_summary cell first.\")\n"
    ]

Samples/Demo_async.ipynb

@@ -2765,7 +2765,7 @@
     "stored = await memory.get_user_summary(user_id=user_id)\n",
     "if stored:\n",
     "    print(\"User Summary for\", user_id)\n",
-    "    print(stored[0][\"content\"])\n",
+    "    print(stored[\"content\"])\n",
     "else:\n",
     "    print(\"No user summary found — run the generate_user_summary cell first.\")\n"
    ]

Samples/processing_user_profile.py

@@ -101,8 +101,11 @@ def main() -> None:
     print(json.dumps(doc.get("metadata", {}).get("structured_summary"), indent=2))
 
     _banner("Read profile back from Cosmos")
-    for s in mem.get_user_summary(user_id):
-        print(f"  • {s['content'][:200]}…")
+    summary = mem.get_user_summary(user_id)
+    if summary:
+        print(f"  • {summary['content'][:200]}…")
+    else:
+        print("  (no summary stored yet)")
 
     print("\nDone.")
 

Samples/scenario_customer_support.py

@@ -80,7 +80,7 @@ def _run_ticket_3_greeting(mem: CosmosMemoryClient, user_id: str, ticket: str) -
     facts = mem.get_memories(user_id=user_id, memory_types=["fact"])
     print("Profile content (truncated):")
     if profile:
-        print(f"  {profile[0]['content'][:200]}…")
+        print(f"  {profile['content'][:200]}…")
     print("\nKnown facts:")
     for f in facts[:8]:
         print(f"  • {f['content']}")

Samples/scenario_rag_with_memory.py

@@ -182,9 +182,8 @@ def run_demo() -> None:
 
     summaries = mem.get_user_summary(user_id=user_id)
     if summaries:
-        print(f"  Found {len(summaries)} summary document(s):")
-        for s in summaries:
-            print(f"    • {s.get('content', '')[:120]}")
+        print("  Found user summary:")
+        print(f"    • {summaries.get('content', '')[:120]}")
     else:
         print("  No summary available yet (generate one via generate_thread_summary).")
 
@@ -204,7 +203,7 @@ def run_demo() -> None:
 
     # 4c. User summary context
     summary_context = (
-        summaries[0].get("content", "") if summaries else "No summary available."
+        summaries.get("content", "") if summaries else "No summary available."
     )
 
     augmented_prompt = textwrap.dedent(f"""\

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -1136,24 +1136,21 @@ async def get_thread(
         items.reverse()
         return items
 
-    async def get_user_summary(self, user_id: str) -> list[dict[str, Any]]:
-        """Retrieve user summary documents from Cosmos DB, newest first."""
-        await self._require_cosmos()
+    async def get_user_summary(self, user_id: str) -> Optional[dict[str, Any]]:
+        """Retrieve the user's summary document from Cosmos DB, or ``None`` if absent."""
+        from azure.cosmos.exceptions import CosmosResourceNotFoundError
 
-        query = (
-            "SELECT c.id, c.user_id, c.thread_id, c.role, c.type, "
-            "c.content, c.metadata, c.created_at "
-            "FROM c WHERE c.user_id = @user_id AND c.type = 'user_summary' "
-            "ORDER BY c.created_at DESC"
-        )
-        parameters = [{"name": "@user_id", "value": user_id}]
-        logger.debug("async get_user_summary query: %s", query)
+        await self._require_cosmos()
 
         try:
-            items_iter = self._container_client.query_items(query=query, parameters=parameters)
-            return [item async for item in items_iter]
+            return await self._container_client.read_item(
+                item=f"user_summary_{user_id}",
+                partition_key=[user_id, "__user_summary__"],
+            )
+        except CosmosResourceNotFoundError:
+            return None
         except Exception as exc:
-            raise CosmosOperationError(f"async get_user_summary query failed: {exc}") from exc
+            raise CosmosOperationError(f"async get_user_summary read failed: {exc}") from exc
 
     # ------------------------------------------------------------------
     # Tag operations

agent_memory_toolkit/cosmos_memory_client.py

@@ -1386,29 +1386,21 @@ def get_thread(
         items.reverse()
         return items
 
-    def get_user_summary(self, user_id: str) -> list[dict[str, Any]]:
-        """Retrieve user summary documents from Cosmos DB, newest first."""
-        self._require_cosmos()
+    def get_user_summary(self, user_id: str) -> Optional[dict[str, Any]]:
+        """Retrieve the user's summary document from Cosmos DB, or ``None`` if absent."""
+        from azure.cosmos.exceptions import CosmosResourceNotFoundError
 
-        query = (
-            "SELECT c.id, c.user_id, c.thread_id, c.role, c.type, "
-            "c.content, c.metadata, c.created_at "
-            "FROM c WHERE c.user_id = @user_id AND c.type = 'user_summary' "
-            "ORDER BY c.created_at DESC"
-        )
-        parameters = [{"name": "@user_id", "value": user_id}]
-        logger.debug("get_user_summary query: %s", query)
+        self._require_cosmos()
 
         try:
-            return list(
-                self._container_client.query_items(
-                    query=query,
-                    parameters=parameters,
-                    enable_cross_partition_query=True,
-                )
+            return self._container_client.read_item(
+                item=f"user_summary_{user_id}",
+                partition_key=[user_id, "__user_summary__"],
             )
+        except CosmosResourceNotFoundError:
+            return None
         except Exception as exc:
-            raise CosmosOperationError(f"get_user_summary query failed: {exc}") from exc
+            raise CosmosOperationError(f"get_user_summary read failed: {exc}") from exc
 
     # ------------------------------------------------------------------
     # Tag operations

tests/integration/test_full_pipeline.py

@@ -212,9 +212,9 @@ def test_multi_thread_user_summary(self, agent_memory, unique_user_id):
             assert doc.get("id"), f"Expected user_summary doc with id, got {doc}"
             assert doc.get("type") == "user_summary"
 
-            summaries = agent_memory.get_user_summary(unique_user_id)
-            assert len(summaries) >= 1
-            combined = " ".join(s.get("content", "") for s in summaries).lower()
+            summary = agent_memory.get_user_summary(unique_user_id)
+            assert summary is not None, "Expected user_summary doc to be readable after generation"
+            combined = summary.get("content", "").lower()
             assert any(t in combined for t in ("pasta", "italian", "cooking", "fettuccine"))
             assert any(t in combined for t in ("running", "run", "morning"))
         finally:

tests/unit/aio/test_cosmos_memory_client.py

@@ -502,16 +502,27 @@ async def test_not_found(self):
 
 
 class TestGetUserSummary:
-    async def test_filters_by_type(self):
+    async def test_returns_doc_when_present(self):
         mem, container = _connected_client()
-        doc = _make_doc(type="user_summary")
-        container.query_items = MagicMock(return_value=AsyncIterator([doc]))
+        doc = _make_doc(type="user_summary", id="user_summary_u1")
+        container.read_item = AsyncMock(return_value=doc)
 
         result = await mem.get_user_summary(user_id="u1")
 
-        call_kwargs = container.query_items.call_args.kwargs
-        assert "user_summary" in call_kwargs["query"]
-        assert result == [doc]
+        call_kwargs = container.read_item.call_args.kwargs
+        assert call_kwargs["item"] == "user_summary_u1"
+        assert call_kwargs["partition_key"] == ["u1", "__user_summary__"]
+        assert result == doc
+
+    async def test_returns_none_when_absent(self):
+        from azure.cosmos.exceptions import CosmosResourceNotFoundError
+
+        mem, container = _connected_client()
+        container.read_item = AsyncMock(side_effect=CosmosResourceNotFoundError(message="404"))
+
+        result = await mem.get_user_summary(user_id="u1")
+
+        assert result is None
 
 
 # ===================================================================

tests/unit/test_cosmos_memory_client.py

@@ -581,16 +581,27 @@ def test_not_found(self):
 
 
 class TestGetUserSummary:
-    def test_filters_by_type(self):
+    def test_returns_doc_when_present(self):
         mem, container = _connected_client()
-        doc = _make_doc(type="user_summary")
-        container.query_items.return_value = [doc]
+        doc = _make_doc(type="user_summary", id="user_summary_u1")
+        container.read_item.return_value = doc
 
         result = mem.get_user_summary(user_id="u1")
 
-        call_kwargs = container.query_items.call_args.kwargs
-        assert "user_summary" in call_kwargs["query"]
-        assert result == [doc]
+        call_kwargs = container.read_item.call_args.kwargs
+        assert call_kwargs["item"] == "user_summary_u1"
+        assert call_kwargs["partition_key"] == ["u1", "__user_summary__"]
+        assert result == doc
+
+    def test_returns_none_when_absent(self):
+        from azure.cosmos.exceptions import CosmosResourceNotFoundError
+
+        mem, container = _connected_client()
+        container.read_item.side_effect = CosmosResourceNotFoundError(message="404")
+
+        result = mem.get_user_summary(user_id="u1")
+
+        assert result is None
 
 
 # ===================================================================