resolving comments

Author: Aayush Kataria

Docs/azure_testing.md

@@ -314,9 +314,9 @@ for i in range(10):
 
 # Wait for the change-feed processor to catch up, then read derived memories.
 import time; time.sleep(15)
-print(memory.get_memories(user_id="user-1", thread_id="thread-1", memory_types=["summary"]))
+print(memory.get_thread_summary(user_id="user-1", thread_id="thread-1"))
 print(memory.get_memories(user_id="user-1", memory_types=["fact"]))
-print(memory.get_memories(user_id="user-1", memory_types=["user_summary"]))
+print(memory.get_user_summary(user_id="user-1"))
 ```
 
 ### Change feed auto-processing
@@ -339,7 +339,7 @@ for i in range(3):
 # Wait a few seconds for the change feed to trigger, then check:
 import time
 time.sleep(10)
-results = memory.get_memories(user_id="user-1", thread_id=thread_id, memory_types=["summary"])
+results = memory.get_thread_summary(user_id="user-1", thread_id=thread_id)
 print(results)  # Should contain an auto-generated summary
 ```
 
@@ -348,9 +348,9 @@ Check the Function App logs to confirm the `on_memory_change` trigger fired and
 ### Verify stored results
 
 ```python
-print(memory.get_memories(user_id="user-1", memory_types=["summary"]))
+print(memory.get_thread_summary(user_id="user-1", thread_id="thread-1"))
 print(memory.get_memories(user_id="user-1", memory_types=["fact"]))
-print(memory.get_memories(user_id="user-1", memory_types=["user_summary"]))
+print(memory.get_user_summary(user_id="user-1"))
 ```
 
 ---

Docs/local_testing.md

@@ -115,9 +115,9 @@ No Azure resources are required for local in-memory operations.
 
 ```python
 import uuid
-from agent_memory_toolkit import AgentMemory
+from agent_memory_toolkit import CosmosMemoryClient
 
-memory = AgentMemory(use_default_credential=False)
+memory = CosmosMemoryClient(use_default_credential=False)
 
 THREAD_ID = str(uuid.uuid4())
 
@@ -133,7 +133,7 @@ memory.delete_local(mem_id)
 print(f"Remaining: {len(memory.get_local())}")
 ```
 
-`AsyncAgentMemory` works the same way for local operations (local methods are synchronous).
+`AsyncCosmosMemoryClient` works the same way for local operations (local methods are synchronous).
 
 ---
 
@@ -153,11 +153,11 @@ Then run a minimal smoke test:
 import os, uuid
 from dotenv import load_dotenv
 from azure.identity import DefaultAzureCredential
-from agent_memory_toolkit import AgentMemory
+from agent_memory_toolkit import CosmosMemoryClient
 
 load_dotenv()
 
-memory = AgentMemory(
+memory = CosmosMemoryClient(
     cosmos_endpoint=os.getenv("COSMOS_DB_ENDPOINT"),
     cosmos_database=os.getenv("COSMOS_DB_DATABASE"),
     cosmos_container=os.getenv("COSMOS_DB_MEMORIES_CONTAINER"),
@@ -196,11 +196,11 @@ for r in results:
 import os, uuid
 from dotenv import load_dotenv
 from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
-from agent_memory_toolkit.aio import AsyncAgentMemory
+from agent_memory_toolkit.aio import AsyncCosmosMemoryClient
 
 load_dotenv()
 
-memory = AsyncAgentMemory(
+memory = AsyncCosmosMemoryClient(
     cosmos_endpoint=os.getenv("COSMOS_DB_ENDPOINT"),
     cosmos_database=os.getenv("COSMOS_DB_DATABASE"),
     cosmos_container=os.getenv("COSMOS_DB_MEMORIES_CONTAINER"),

Docs/public_api.md

@@ -28,10 +28,11 @@
 - `delete_local(memory_id) -> None` — remove a local buffered memory.
 - `add_cosmos(user_id, role, content, memory_type='turn', metadata=None, thread_id=None, tags=None, ttl=None, salience=None, embedding=None, embed=None) -> str` — upsert one memory to Cosmos and return its id.
 - `push_to_cosmos(batch_size=25) -> None` — flush local buffered memories to Cosmos.
-- `get_memories(memory_id=None, user_id=None, thread_id=None, role=None, memory_types=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` — retrieve memories with filters.
-- `update_cosmos(memory_id, content=None, role=None, memory_type=None, metadata=None) -> None` — update a Cosmos memory.
-- `delete_cosmos(memory_id, thread_id, user_id) -> None` — delete a Cosmos memory.
-- `get_thread(thread_id, user_id=None, memory_types=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, created_after=None, created_before=None) -> list[dict]` — retrieve a thread oldest-first.
+- `get_memories(memory_id=None, user_id=None, thread_id=None, role=None, memory_types=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` — retrieve memories from the MEMORIES container. `memory_types` defaults to `["fact", "episodic", "procedural"]` and must be a subset of those three.
+- `update_cosmos(memory_id, *, user_id, thread_id, memory_type, content=None, role=None, metadata=None) -> None` — point-update a memory in the container that holds `memory_type`. The `type` field itself is never mutated.
+- `delete_cosmos(memory_id, *, user_id, thread_id, memory_type) -> None` — delete a memory from the container that holds `memory_type`.
+- `get_thread(thread_id, user_id=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, created_after=None, created_before=None) -> list[dict]` — retrieve turns from the TURNS container oldest-first.
+- `get_thread_summary(user_id, thread_id, recent_k=None) -> list[dict]` — retrieve thread summary documents from the SUMMARIES container for a single `(user_id, thread_id)` partition.
 - `get_user_summary(user_id) -> Optional[dict]` — retrieve the active user-summary document.
 
 ### Retrieval
@@ -56,8 +57,8 @@
 
 ### Tagging
 
-- `add_tags(memory_id, user_id, thread_id, tags) -> None` — add tags to a memory.
-- `remove_tags(memory_id, user_id, thread_id, tags) -> None` — remove tags from a memory.
+- `add_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` — add tags to a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`.
+- `remove_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` — remove tags from a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`.
 - `list_tags(user_id, *, thread_id=None, prefix=None, include_sys=False) -> list[str]` — list sorted, deduped tags for a user; omits `sys:*` by default.
 
 ## AsyncCosmosMemoryClient
@@ -80,10 +81,11 @@ Local-buffer methods remain synchronous in-memory operations; Cosmos, retrieval,
 - `delete_local(memory_id) -> None` — remove a local buffered memory.
 - `async add_cosmos(user_id, role, content, memory_type='turn', metadata=None, thread_id=None, tags=None, ttl=None, salience=None, embedding=None, embed=None) -> str` — upsert one memory to Cosmos and return its id.
 - `async push_to_cosmos(batch_size=25) -> None` — flush local buffered memories to Cosmos.
-- `async get_memories(memory_id=None, user_id=None, thread_id=None, role=None, memory_types=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` — retrieve memories with filters.
-- `async update_cosmos(memory_id, content=None, role=None, memory_type=None, metadata=None) -> None` — update a Cosmos memory.
-- `async delete_cosmos(memory_id, thread_id, user_id) -> None` — delete a Cosmos memory.
-- `async get_thread(thread_id, user_id=None, memory_types=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, created_after=None, created_before=None) -> list[dict]` — retrieve a thread oldest-first.
+- `async get_memories(memory_id=None, user_id=None, thread_id=None, role=None, memory_types=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` — retrieve memories from the MEMORIES container. `memory_types` defaults to `["fact", "episodic", "procedural"]` and must be a subset of those three.
+- `async update_cosmos(memory_id, *, user_id, thread_id, memory_type, content=None, role=None, metadata=None) -> None` — point-update a memory in the container that holds `memory_type`. The `type` field itself is never mutated.
+- `async delete_cosmos(memory_id, *, user_id, thread_id, memory_type) -> None` — delete a memory from the container that holds `memory_type`.
+- `async get_thread(thread_id, user_id=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, created_after=None, created_before=None) -> list[dict]` — retrieve turns from the TURNS container oldest-first.
+- `async get_thread_summary(user_id, thread_id, recent_k=None) -> list[dict]` — retrieve thread summary documents from the SUMMARIES container for a single `(user_id, thread_id)` partition.
 - `async get_user_summary(user_id) -> Optional[dict]` — retrieve the active user-summary document.
 
 ### Retrieval
@@ -108,8 +110,8 @@ Local-buffer methods remain synchronous in-memory operations; Cosmos, retrieval,
 
 ### Tagging
 
-- `async add_tags(memory_id, user_id, thread_id, tags) -> None` — add tags to a memory.
-- `async remove_tags(memory_id, user_id, thread_id, tags) -> None` — remove tags from a memory.
+- `async add_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` — add tags to a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`.
+- `async remove_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` — remove tags from a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`.
 - `async list_tags(user_id, *, thread_id=None, prefix=None, include_sys=False) -> list[str]` — list sorted, deduped tags for a user; omits `sys:*` by default.
 
 ## Topology validation

Samples/Advanced/advanced_memory_lifecycle.py

@@ -92,8 +92,8 @@ def main() -> None:
 
     _header(5, "Archive: delete raw turns, keep derived memories")
     deleted = 0
-    for m in mem.get_memories(user_id=user_id, thread_id=thread_id, memory_types=["turn"]):
-        mem.delete_cosmos(memory_id=m["id"], thread_id=thread_id, user_id=user_id)
+    for m in mem.get_thread(thread_id=thread_id, user_id=user_id):
+        mem.delete_cosmos(memory_id=m["id"], user_id=user_id, thread_id=thread_id, memory_type="turn")
         deleted += 1
     print(f"  deleted {deleted} raw turn(s)")
 

Samples/Advanced/advanced_search_patterns.py

@@ -10,13 +10,14 @@
 """
 
 import os
+import uuid
 
 from dotenv import load_dotenv
-load_dotenv()
-import uuid
 
 from agent_memory_toolkit import CosmosMemoryClient
 
+load_dotenv()
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------

Samples/Notebooks/Demo.ipynb

@@ -50,12 +50,14 @@
    },
    "outputs": [],
    "source": [
-    "import os, json\n",
-    "from dotenv import load_dotenv\n",
-    "from azure.identity import DefaultAzureCredential\n",
+    "import json\n",
+    "import os\n",
     "\n",
     "# Add repository root to path so we can import the package easily\n",
     "import sys\n",
+    "\n",
+    "from dotenv import load_dotenv\n",
+    "\n",
     "sys.path.insert(0, os.path.abspath(\"../..\"))\n",
     "from agent_memory_toolkit import CosmosMemoryClient\n",
     "\n",
@@ -143,6 +145,7 @@
    "outputs": [],
    "source": [
     "import uuid\n",
+    "\n",
     "THREAD_ID = str(uuid.uuid4())\n",
     "# Use a unique user_id per demo run so we get a clean extraction without\n",
     "# inheriting facts from prior runs that would dedup new content away.\n",
@@ -499,8 +502,8 @@
     "# Verify the memories were added directly to Cosmos DB (not in local memory)\n",
     "print(f\"Local memory count (should be unchanged): {len(memory.local_memory)}\\n\")\n",
     "\n",
-    "cosmos_results = memory.get_memories(user_id=\"user-002\", thread_id=new_thread_id)\n",
-    "print(f\"Memories in Cosmos DB for new thread: {len(cosmos_results)}\")\n",
+    "cosmos_results = memory.get_thread(thread_id=new_thread_id, user_id=\"user-002\")\n",
+    "print(f\"Turns in Cosmos DB for new thread: {len(cosmos_results)}\")\n",
     "for r in cosmos_results:\n",
     "    print(f\"  [{r['thread_id'][:8]}...] [{r['id'][:8]}...] role={r['role']:<6} {r['content'][:60]}\")"
    ]
@@ -510,9 +513,9 @@
    "id": "40d603f1",
    "metadata": {},
    "source": [
-    "### 3c. Retrieve memories from Cosmos DB with `get_memories`\n",
+    "### 3c. Retrieve memories from Cosmos DB\n",
     "\n",
-    "Supports the same filters as `get_local`: `memory_id`, `user_id`, `role`, `memory_type`."
+    "`get_memories(...)` returns derived memories (`fact`, `episodic`, `procedural`) from the memories container. For other containers use `get_thread(thread_id=...)` for turns, `get_thread_summary(user_id, thread_id)` for thread summaries, and `get_user_summary(user_id)` for the cross-thread profile."
    ]
   },
   {
@@ -533,16 +536,17 @@
    },
    "outputs": [],
    "source": [
-    "# Get all memories for user-001\n",
+    "# Derived memories (fact/episodic/procedural) for user-001. Empty until section 5 runs extract_memories.\n",
     "results = memory.get_memories(user_id=USER_ID)\n",
-    "print(f\"Memories for user-001: {len(results)}\\n\")\n",
+    "print(f\"Derived memories for {USER_ID}: {len(results)}\\n\")\n",
     "for r in results:\n",
-    "    print(f\"  [{r['thread_id'][:8]}...] [{r['id'][:8]}...] role={r['role']:<6} type={r['type']:<8} {r['content'][:50]}\")\n",
+    "    print(f\"  [{r['thread_id'][:8]}...] [{r['id'][:8]}...] type={r['type']:<10} {r['content'][:50]}\")\n",
     "\n",
-    "# Get only agent memories\n",
-    "agent_results = memory.get_memories(role=\"agent\")\n",
-    "print(f\"\\nAgent memories: {len(agent_results)}\")\n",
-    "for r in agent_results:\n",
+    "# Turns live in a separate container — fetch via get_thread and filter in-process.\n",
+    "turns = memory.get_thread(thread_id=THREAD_ID, user_id=USER_ID)\n",
+    "agent_turns = [t for t in turns if t.get(\"role\") == \"agent\"]\n",
+    "print(f\"\\nAgent turns in seed thread: {len(agent_turns)}\")\n",
+    "for r in agent_turns:\n",
     "    print(f\"  [{r['thread_id'][:8]}...] [{r['id'][:8]}...] role={r['role']:<6} type={r['type']:<8} {r['content'][:50]}\")"
    ]
   },
@@ -574,18 +578,21 @@
    },
    "outputs": [],
    "source": [
-    "# Update the user's budget message to add a hotel budget constraint\n",
-    "user_msgs = memory.get_memories(user_id=USER_ID, role=\"user\")\n",
-    "target = [m for m in user_msgs if \"Pike Place\" in m[\"content\"]][0]\n",
+    "# Update the user's budget turn (in the seed thread) to add a hotel budget constraint\n",
+    "turns = memory.get_thread(thread_id=THREAD_ID, user_id=USER_ID)\n",
+    "target = [m for m in turns if m.get(\"role\") == \"user\" and \"Pike Place\" in m[\"content\"]][0]\n",
     "print(f\"Before: {target['content']}\\n\")\n",
     "\n",
     "memory.update_cosmos(\n",
     "    memory_id=target[\"id\"],\n",
+    "    user_id=target[\"user_id\"],\n",
+    "    thread_id=target[\"thread_id\"],\n",
+    "    memory_type=target[\"type\"],\n",
     "    content=\"Something near Pike Place Market would be great. Keep flights under $300 round trip and hotels under $200/night.\",\n",
     "    metadata={\"edited\": True, \"reason\": \"user clarified hotel budget\"},\n",
     ")\n",
     "\n",
-    "updated = memory.get_memories(memory_id=target[\"id\"])[0]\n",
+    "updated = [m for m in memory.get_thread(thread_id=THREAD_ID, user_id=USER_ID) if m[\"id\"] == target[\"id\"]][0]\n",
     "print(f\"After:  {updated['content']}\")"
    ]
   },
@@ -607,20 +614,21 @@
    },
    "outputs": [],
    "source": [
-    "# Delete the tool memory from Cosmos\n",
-    "tool_mems = memory.get_memories(user_id=\"user-002\", role=\"tool\")\n",
+    "# Delete the tool turn from the new thread in Cosmos\n",
+    "tool_mems = [t for t in memory.get_thread(thread_id=new_thread_id, user_id=\"user-002\") if t.get(\"role\") == \"tool\"]\n",
     "print(tool_mems[0])\n",
     "if tool_mems:\n",
     "    memory.delete_cosmos(\n",
     "        tool_mems[0][\"id\"],\n",
-    "        thread_id=tool_mems[0][\"thread_id\"],\n",
     "        user_id=tool_mems[0][\"user_id\"],\n",
+    "        thread_id=tool_mems[0][\"thread_id\"],\n",
+    "        memory_type=tool_mems[0][\"type\"],\n",
     "    )\n",
     "    print(f\"Deleted tool memory {tool_mems[0]['id'][:8]}...\")\n",
     "\n",
-    "# Verify\n",
-    "remaining = memory.get_memories()\n",
-    "print(f\"\\nRemaining memories in Cosmos DB: {len(remaining)}\")\n",
+    "# Verify the remaining turns in the new thread\n",
+    "remaining = memory.get_thread(thread_id=new_thread_id, user_id=\"user-002\")\n",
+    "print(f\"\\nRemaining turns in new thread: {len(remaining)}\")\n",
     "for r in remaining:\n",
     "    print(f\"  [{r['thread_id'][:8]}...] [{r['id'][:8]}...] role={r['role']:<6} type={r['type']:<8} {r['content'][:50]}\")"
    ]
@@ -662,14 +670,14 @@
     "print(f\"Using thread_id: {thread_id}\\n\")\n",
     "\n",
     "# Get all documents in the thread\n",
-    "thread_all = memory.get_thread(thread_id=thread_id)\n",
+    "thread_all = memory.get_thread(thread_id=thread_id, user_id=USER_ID)\n",
     "print(f\"All memories in thread: {len(thread_all)}\")\n",
     "for m in thread_all:\n",
     "    print(f\"  [{m['thread_id'][:8]}...] [{m['id'][:8]}...] role={m.get('role','?'):6} type={m['type']:8}  {m['content'][:60]}\")\n",
     "\n",
     "# Most recent 2\n",
-    "recent = memory.get_thread(thread_id=thread_id, recent_k=2)\n",
-    "print(f\"\\nMost recent 2 memories:\")\n",
+    "recent = memory.get_thread(thread_id=thread_id, user_id=USER_ID, recent_k=2)\n",
+    "print(\"\\nMost recent 2 memories:\")\n",
     "for m in recent:\n",
     "    print(f\"  [{m['thread_id'][:8]}...] [{m['id'][:8]}...] role={m.get('role','?'):6} type={m['type']:8}  {m['content'][:60]}\")\n",
     "\n",