Merge pull request #2 from aayush3011/users/akataria/asyncFixes

Code refactoring, async code fixes and error handling

Author: jcodella

.github/workflows/ci.yml

@@ -0,0 +1,49 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+permissions:
+  contents: read
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: pip install ruff
+      - name: Ruff check
+        run: ruff check agent_memory_toolkit/ tests/
+      - name: Ruff format check
+        run: ruff format --check agent_memory_toolkit/ tests/
+
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install package with dev dependencies
+        run: pip install -e ".[dev]"
+      - name: Run unit tests with coverage
+        run: pytest tests/unit/ --cov=agent_memory_toolkit --cov-report=xml --cov-report=term-missing -v
+      - name: Upload coverage
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-report-${{ matrix.python-version }}
+          path: coverage.xml

Docs/azure_testing.md

@@ -22,12 +22,12 @@ You need:
 
 - an Azure subscription
 - `az login`
-- Python 3.10+
+- Python 3.11+
 - Azure Functions Core Tools v4
 - dependencies installed:
 
 ```bash
-pip install -r requirements.txt
+pip install -e ".[dev]"
 pip install -r azure_functions/requirements.txt
 ```
 
@@ -190,7 +190,7 @@ memory.connect_cosmos()
 import os
 from dotenv import load_dotenv
 from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
-from agent_memory_toolkit import AsyncAgentMemory
+from agent_memory_toolkit.aio import AsyncAgentMemory
 
 load_dotenv()
 

Docs/design_patterns.md

@@ -12,7 +12,7 @@ Write a turn memory every time a user or agent message is produced. If the appli
 
 ```python
 from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
-from agent_memory_toolkit import AsyncAgentMemory
+from agent_memory_toolkit.aio import AsyncAgentMemory
 
 mem = AsyncAgentMemory(
     cosmos_endpoint=COSMOS_ENDPOINT,

Docs/local_testing.md

@@ -10,7 +10,7 @@ This guide covers the shortest path to running the library locally, then validat
 
 | Tool | Install command | Purpose |
 |------|-----------------|---------|
-| Python 3.10+ | `brew install python@3.12` | Runtime |
+| Python 3.11+ | `brew install python@3.13` | Runtime |
 | Azure CLI | `brew install azure-cli` | `az login` for `DefaultAzureCredential` |
 | Azure Functions Core Tools v4 | `brew install azure-functions-core-tools@4` | Run Functions locally |
 | Azurite | `npm install -g azurite` | Local storage emulator for Functions |
@@ -19,7 +19,7 @@ This guide covers the shortest path to running the library locally, then validat
 ### Python packages
 
 ```bash
-pip install -r requirements.txt
+pip install -e ".[dev]"
 pip install -r azure_functions/requirements.txt
 ```
 
@@ -168,7 +168,7 @@ for r in results:
 import os, uuid
 from dotenv import load_dotenv
 from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
-from agent_memory_toolkit import AsyncAgentMemory
+from agent_memory_toolkit.aio import AsyncAgentMemory
 
 load_dotenv()
 
@@ -217,7 +217,7 @@ azurite --silent --location /tmp/azurite --debug /tmp/azurite/debug.log
 
 ```bash
 cd azure_functions
-pip install -r requirements.txt
+pip install -r azure_functions/requirements.txt
 func start
 ```
 
@@ -324,7 +324,7 @@ User summaries also update incrementally when one already exists.
 
 | Problem | Fix |
 |---------|-----|
-| `ImportError: azure.identity` | Run `pip install -r requirements.txt` |
+| `ImportError: azure.identity` | Run `pip install -e ".[dev]"` |
 | `DefaultAzureCredential` fails | Run `az login` and confirm the active subscription |
 | Cosmos 403 | Check Cosmos DB RBAC and wait for propagation |
 | `func: command not found` | Install Azure Functions Core Tools v4 |
@@ -334,6 +334,3 @@ User summaries also update incrementally when one already exists.
 | Function 401 in Azure | Set `ADF_KEY` or pass `?code=<key>` |
 
 For full cloud deployment and validation, see `Docs/azure_testing.md`.
-| Functions host can't connect to storage | Make sure Azurite is running before starting `func start` |
-| Embeddings 401/403 | Confirm the "Cognitive Services OpenAI User" role is assigned to your identity on the AI resource |
-| Functions 401 Unauthorized | The HTTP trigger requires a function key. Pass it as `?code=<key>` in the URL or set `ADF_KEY` in `.env`. See "Get the function key" above. |

README.md

@@ -1,7 +1,7 @@
 # Azure Cosmos DB Agent Memory Toolkit
 
 [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
-[![Python 3.12+](https://img.shields.io/badge/python-3.12+-blue.svg)](https://www.python.org/downloads/)
+[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
 [![Azure Cosmos DB](https://img.shields.io/badge/Azure-Cosmos%20DB-0078D4?logo=microsoft-azure)](https://azure.microsoft.com/en-us/products/cosmos-db/)
 [![Follow on X](https://img.shields.io/twitter/follow/AzureCosmosDB?style=social)](https://twitter.com/AzureCosmosDB)
 [![LinkedIn](https://img.shields.io/badge/LinkedIn-Azure%20Cosmos%20DB-0077B5?logo=linkedin)](https://www.linkedin.com/showcase/azure-cosmos-db/)
@@ -71,11 +71,24 @@ Agent Memory Toolkit is a Python library and Azure-backed reference implementati
 ## Project Structure
 
 ```
-agent_memory_toolkit/          Python library — AgentMemory (sync) + AsyncAgentMemory (async)
+agent_memory_toolkit/          Python library — sync API
+  memory.py                    AgentMemory orchestrator
+  cosmos_memory_client.py      CosmosMemoryStore — Cosmos DB CRUD + vector search
+  embeddings.py                EmbeddingsClient — Azure OpenAI embeddings
+  processing.py                ProcessingClient — Durable Functions polling
+  models.py                    Pydantic data models (MemoryRecord, enums)
+  exceptions.py                Custom exception hierarchy
+  _query_builder.py            Shared query builder (private)
+  aio/                         Async API (mirrors azure.cosmos.aio convention)
+    memory.py                  AsyncAgentMemory
+    cosmos_memory_client.py    AsyncCosmosMemoryStore
+    embeddings.py              AsyncEmbeddingsClient
+    processing.py              AsyncProcessingClient
 azure_functions/               Durable Functions — orchestrator, activities, HTTP trigger
   prompts/                     LLM system prompts — summarize, facts, user_summary + update variants
 Samples/                       Demo notebooks — sync (Demo.ipynb) + async (Demo_async.ipynb)
 Docs/                          Documentation — concepts, local testing, Azure deployment
+tests/                         Unit tests (pytest) — 184 tests, 87% coverage
 ```
 
 ---
@@ -85,7 +98,10 @@ Docs/                          Documentation — concepts, local testing, Azure
 ### 1. Install
 
 ```bash
-pip install -r requirements.txt
+pip install .
+
+# With dev/test dependencies
+pip install ".[dev]"
 ```
 
 ### 2. Local-only (no Azure)
@@ -156,7 +172,11 @@ result = memory.generate_user_summary(user_id="user-001")
 summary = memory.get_user_summary(user_id="user-001")
 ```
 
-> The async API (`AsyncAgentMemory`) is identical — just `await` each call.
+> The async API (`AsyncAgentMemory`) is identical — just `await` each call. Import from the `aio` subpackage:
+>
+> ```python
+> from agent_memory_toolkit.aio import AsyncAgentMemory
+> ```
 
 ---
 

Samples/Demo_async.ipynb

@@ -43,7 +43,7 @@
     "sys.path.insert(0, os.path.abspath(\"..\"))\n",
     "\n",
     "from dotenv import load_dotenv\n",
-    "from agent_memory_toolkit import AsyncAgentMemory\n",
+    "from agent_memory_toolkit.aio import AsyncAgentMemory\n",
     "\n",
     "# Load environment variables from .env in the repo root\n",
     "load_dotenv(os.path.join(\"..\", \".env\"))\n",

agent_memory_toolkit/__init__.py

@@ -1,6 +1,25 @@
 """Agent Memory Toolkit – local and cloud agent memory management."""
 
+from agent_memory_toolkit.aio import AsyncAgentMemory
+from agent_memory_toolkit.exceptions import (
+    AgentMemoryError,
+    AuthenticationError,
+    ConfigurationError,
+    CosmosNotConnectedError,
+    CosmosOperationError,
+    EmbeddingError,
+    MemoryNotFoundError,
+    OrchestrationTimeoutError,
+    ProcessingError,
+    ValidationError,
+)
 from agent_memory_toolkit.memory import AgentMemory
-from agent_memory_toolkit.async_memory import AsyncAgentMemory
+from agent_memory_toolkit.models import MemoryRecord, MemoryRole, MemoryType, SearchResult
 
-__all__ = ["AgentMemory", "AsyncAgentMemory"]
+__all__ = [
+    "AgentMemory", "AsyncAgentMemory",
+    "MemoryRecord", "MemoryRole", "MemoryType", "SearchResult",
+    "AgentMemoryError", "ConfigurationError", "ValidationError",
+    "CosmosNotConnectedError", "CosmosOperationError", "MemoryNotFoundError",
+    "EmbeddingError", "ProcessingError", "OrchestrationTimeoutError", "AuthenticationError",
+]

agent_memory_toolkit/_query_builder.py

@@ -0,0 +1,43 @@
+"""Reusable query-builder for parameterized Cosmos DB queries.
+
+The :class:`_QueryBuilder` helper eliminates duplicated
+condition/parameter-building patterns across the sync and async clients.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class _QueryBuilder:
+    """Accumulates optional WHERE conditions and their parameterized values.
+
+    Usage::
+
+        qb = _QueryBuilder()
+        qb.add_filter("c.user_id", "@user_id", some_user_id)
+        qb.add_filter("c.role", "@role", some_role)
+        where = qb.build_where()        # " WHERE c.user_id = @user_id AND c.role = @role"
+        params = qb.get_parameters()     # [{"name": "@user_id", "value": ...}, ...]
+    """
+
+    def __init__(self) -> None:
+        self._conditions: list[str] = []
+        self._parameters: list[dict[str, Any]] = []
+
+    def add_filter(self, field: str, param_name: str, value: Any) -> None:
+        """Add a filter only when *value* is not ``None``."""
+        if value is None:
+            return
+        self._conditions.append(f"{field} = {param_name}")
+        self._parameters.append({"name": param_name, "value": value})
+
+    def build_where(self) -> str:
+        """Return the ``WHERE …`` clause (or empty string if no filters)."""
+        if not self._conditions:
+            return ""
+        return " WHERE " + " AND ".join(self._conditions)
+
+    def get_parameters(self) -> list[dict[str, Any]]:
+        """Return a *copy* of the accumulated parameters list."""
+        return list(self._parameters)