crewai-memory-cosmosdb 0.1.0__tar.gz

crewai_memory_cosmosdb-0.1.0/.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.pytest_cache/
+.venv/
+dist/
+build/
+uv.lock
+.env

crewai_memory_cosmosdb-0.1.0/PKG-INFO

@@ -0,0 +1,53 @@
+Metadata-Version: 2.5
+Name: crewai-memory-cosmosdb
+Version: 0.1.0
+Summary: Azure Cosmos DB StorageBackend for CrewAI unified Memory — hierarchical scopes, categories, metadata filters and native vector search (VectorDistance)
+Project-URL: Homepage, https://github.com/skamalj/crewai-memory
+Project-URL: Repository, https://github.com/skamalj/crewai-memory.git
+Project-URL: Documentation, https://skamalj.github.io/agentstate-reducer/
+Author-email: Kamal <skamalj@gmail.com>
+Keywords: agent-memory,azure,cosmosdb,crewai,long-term-memory,memory,storage-backend,vector-search
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: azure-cosmos>=4.7
+Requires-Dist: crewai-memory-core>=0.1.0
+Description-Content-Type: text/markdown
+
+# crewai-memory-cosmosdb
+
+An **Azure Cosmos DB (NoSQL)** `StorageBackend` for [CrewAI](https://docs.crewai.com/en/concepts/memory)'s unified `Memory` — hierarchical scopes, categories, metadata filters, importance/recency, and **native vector search** via `VectorDistance` (diskANN index).
+
+```bash
+pip install crewai-memory-cosmosdb
+```
+
+```python
+from crewai import Crew
+from crewai.memory import Memory
+from crewai_memory_cosmosdb import CosmosDBMemoryBackend
+
+backend = CosmosDBMemoryBackend(
+    endpoint="https://<acct>.documents.azure.com:443/", key="<key>",
+    database_name="crewai", container_name="memory", dimensions=3072,   # match your embedder
+)
+memory = Memory(storage=backend)
+crew = Crew(agents=[...], tasks=[...], memory=memory)
+```
+
+Or once for every `Crew(memory=True)`: `set_memory_storage_factory(lambda spec: CosmosDBMemoryBackend(...))`.
+
+## How it works
+
+- A container partitioned by `/scope`, one document per record (`id` = record id), created with a vector embedding policy on `/embedding` (cosine, `dimensions`) and a `diskANN` vector index (`vector_index_type="quantizedFlat"|"flat"` to change). The policy is fixed at container creation — use a new container to change `dimensions`.
+- `search` is one SQL query: `ORDER BY VectorDistance(c.embedding, @q)` with the scope subtree in `WHERE`; the score is the cosine similarity. `categories` / `metadata_filter` / `min_score` are applied on the candidates.
+- Scope tree, category counts, listing, `delete(older_than=...)`, `reset` are cross-partition SQL queries.
+- The account needs **Vector Search for NoSQL**: `az cosmosdb update -n <acct> -g <rg> --capabilities EnableServerless EnableNoSQLVectorSearch` (list all existing capabilities; propagation can take a few minutes).
+
+Docs: <https://skamalj.github.io/agentstate-reducer/> · part of [crewai-memory](https://github.com/skamalj/crewai-memory)
+
+## License
+
+MIT

crewai_memory_cosmosdb-0.1.0/README.md

@@ -0,0 +1,35 @@
+# crewai-memory-cosmosdb
+
+An **Azure Cosmos DB (NoSQL)** `StorageBackend` for [CrewAI](https://docs.crewai.com/en/concepts/memory)'s unified `Memory` — hierarchical scopes, categories, metadata filters, importance/recency, and **native vector search** via `VectorDistance` (diskANN index).
+
+```bash
+pip install crewai-memory-cosmosdb
+```
+
+```python
+from crewai import Crew
+from crewai.memory import Memory
+from crewai_memory_cosmosdb import CosmosDBMemoryBackend
+
+backend = CosmosDBMemoryBackend(
+    endpoint="https://<acct>.documents.azure.com:443/", key="<key>",
+    database_name="crewai", container_name="memory", dimensions=3072,   # match your embedder
+)
+memory = Memory(storage=backend)
+crew = Crew(agents=[...], tasks=[...], memory=memory)
+```
+
+Or once for every `Crew(memory=True)`: `set_memory_storage_factory(lambda spec: CosmosDBMemoryBackend(...))`.
+
+## How it works
+
+- A container partitioned by `/scope`, one document per record (`id` = record id), created with a vector embedding policy on `/embedding` (cosine, `dimensions`) and a `diskANN` vector index (`vector_index_type="quantizedFlat"|"flat"` to change). The policy is fixed at container creation — use a new container to change `dimensions`.
+- `search` is one SQL query: `ORDER BY VectorDistance(c.embedding, @q)` with the scope subtree in `WHERE`; the score is the cosine similarity. `categories` / `metadata_filter` / `min_score` are applied on the candidates.
+- Scope tree, category counts, listing, `delete(older_than=...)`, `reset` are cross-partition SQL queries.
+- The account needs **Vector Search for NoSQL**: `az cosmosdb update -n <acct> -g <rg> --capabilities EnableServerless EnableNoSQLVectorSearch` (list all existing capabilities; propagation can take a few minutes).
+
+Docs: <https://skamalj.github.io/agentstate-reducer/> · part of [crewai-memory](https://github.com/skamalj/crewai-memory)
+
+## License
+
+MIT

crewai_memory_cosmosdb-0.1.0/pyproject.toml

@@ -0,0 +1,36 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "crewai-memory-cosmosdb"
+version = "0.1.0"
+description = "Azure Cosmos DB StorageBackend for CrewAI unified Memory — hierarchical scopes, categories, metadata filters and native vector search (VectorDistance)"
+authors = [{name = "Kamal", email = "skamalj@gmail.com"}]
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "crewai-memory-core>=0.1.0",
+    "azure-cosmos>=4.7",
+]
+keywords = ["crewai", "memory", "storage-backend", "cosmosdb", "azure", "long-term-memory", "agent-memory", "vector-search"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.urls]
+Homepage = "https://github.com/skamalj/crewai-memory"
+Repository = "https://github.com/skamalj/crewai-memory.git"
+Documentation = "https://skamalj.github.io/agentstate-reducer/"
+
+[dependency-groups]
+dev = ["pytest>=7.0", "pytest-asyncio"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/crewai_memory_cosmosdb"]

crewai_memory_cosmosdb-0.1.0/src/crewai_memory_cosmosdb/__init__.py

@@ -0,0 +1,141 @@
+"""Azure Cosmos DB (NoSQL) ``StorageBackend`` for CrewAI unified ``Memory``.
+
+A container partitioned by ``/scope`` holding one document per ``MemoryRecord``
+(``id`` = record id) with a **vector embedding policy** on ``/embedding``
+(cosine) and a ``diskANN`` vector index. ``search`` is a single SQL query:
+``ORDER BY VectorDistance(c.embedding, @q)`` with the scope subtree filtered in
+``WHERE``. Built on ``crewai-memory-core``.
+
+The account needs the ``EnableNoSQLVectorSearch`` capability; the vector policy
+can only be set at container creation.
+"""
+
+from __future__ import annotations
+
+from typing import Any, List, Optional, Tuple
+
+from azure.cosmos import CosmosClient, PartitionKey
+from azure.cosmos.exceptions import CosmosResourceNotFoundError
+
+from crewai_memory_core import MemoryBackend, ROW_FIELDS, norm_scope
+
+__all__ = ["CosmosDBMemoryBackend"]
+
+_PROJ = ", ".join(f'c["{f}"]' for f in ROW_FIELDS if f != "embedding")
+
+
+class CosmosDBMemoryBackend(MemoryBackend):
+    """CrewAI ``StorageBackend`` on Azure Cosmos DB with native vector search.
+
+    Example:
+        ```python
+        from crewai.memory import Memory
+        from crewai_memory_cosmosdb import CosmosDBMemoryBackend
+
+        backend = CosmosDBMemoryBackend(endpoint=..., key=..., database_name="crewai",
+                                        container_name="memory", dimensions=3072)
+        memory = Memory(storage=backend)
+        ```
+    """
+
+    supports_native_vector_search = True
+
+    def __init__(
+        self,
+        *,
+        endpoint: str,
+        key: str,
+        database_name: str,
+        container_name: str,
+        dimensions: int = 3072,
+        vector_index_type: str = "diskANN",
+    ) -> None:
+        self.dimensions = dimensions
+        client = CosmosClient(endpoint, credential=key)
+        db = client.create_database_if_not_exists(database_name)
+        self._container = db.create_container_if_not_exists(
+            id=container_name,
+            partition_key=PartitionKey(path="/scope"),
+            vector_embedding_policy={
+                "vectorEmbeddings": [
+                    {"path": "/embedding", "dataType": "float32", "distanceFunction": "cosine", "dimensions": dimensions}
+                ]
+            },
+            indexing_policy={
+                "automatic": True,
+                "indexingMode": "consistent",
+                "includedPaths": [{"path": "/*"}],
+                "excludedPaths": [{"path": "/embedding/*"}, {"path": '/"_etag"/?'}],
+                "vectorIndexes": [{"path": "/embedding", "type": vector_index_type}],
+            },
+        )
+
+    # ------------------------------------------------------------------ helpers
+    @staticmethod
+    def _row(doc: dict) -> dict:
+        row = {f: doc.get(f) for f in ROW_FIELDS}
+        row["scope"] = norm_scope(doc.get("scope"))
+        return row
+
+    @staticmethod
+    def _scope_where(scope_prefix: Optional[str]) -> Tuple[str, list]:
+        p = norm_scope(scope_prefix)
+        if p == "/":
+            return "true", []
+        return '(c["scope"] = @p OR STARTSWITH(c["scope"], @pc))', [
+            {"name": "@p", "value": p}, {"name": "@pc", "value": p + "/"},
+        ]
+
+    def _query(self, sql: str, params: list) -> List[dict]:
+        return list(self._container.query_items(query=sql, parameters=params, enable_cross_partition_query=True))
+
+    # ------------------------------------------------------------------ primitives
+    def _put(self, rows: List[dict]) -> None:
+        for row in rows:
+            existing = self._get(row["id"])
+            if existing and existing["scope"] != norm_scope(row["scope"]):
+                self._delete_doc(row["id"], existing["scope"])   # scope moved: partition changes
+            doc = {f: row.get(f) for f in ROW_FIELDS}
+            doc["scope"] = norm_scope(row["scope"])
+            if not doc.get("embedding"):
+                doc.pop("embedding", None)
+            self._container.upsert_item(doc)
+
+    def _get(self, record_id: str) -> Optional[dict]:
+        docs = self._query('SELECT TOP 1 * FROM c WHERE c["id"] = @id', [{"name": "@id", "value": record_id}])
+        return self._row(docs[0]) if docs else None
+
+    def _delete_doc(self, record_id: str, scope: str) -> bool:
+        try:
+            self._container.delete_item(item=record_id, partition_key=scope)
+            return True
+        except CosmosResourceNotFoundError:
+            return False
+
+    def _delete_ids(self, ids: List[str]) -> int:
+        n = 0
+        for rid in ids:
+            row = self._get(rid)
+            if row and self._delete_doc(rid, row["scope"]):
+                n += 1
+        return n
+
+    def _scan(self, scope_prefix: Optional[str]) -> List[dict]:
+        where, params = self._scope_where(scope_prefix)
+        return [self._row(d) for d in self._query(f"SELECT * FROM c WHERE {where}", params)]
+
+    def _vector_search(
+        self, vector: List[float], scope_prefix: Optional[str], limit: int
+    ) -> List[Tuple[dict, float]]:
+        where, params = self._scope_where(scope_prefix)
+        sql = (
+            f'SELECT TOP @k {_PROJ}, VectorDistance(c["embedding"], @q) AS score '
+            f'FROM c WHERE {where} AND IS_DEFINED(c["embedding"]) '
+            f'ORDER BY VectorDistance(c["embedding"], @q)'
+        )
+        params = params + [{"name": "@q", "value": vector}, {"name": "@k", "value": int(limit)}]
+        out: List[Tuple[dict, float]] = []
+        for d in self._query(sql, params):
+            row = self._row(d)
+            out.append((row, float(d["score"])))   # cosine similarity, higher is better
+        return out

crewai_memory_cosmosdb-0.1.0/tests/test_cosmosdb_backend.py

@@ -0,0 +1,46 @@
+"""Cosmos DB StorageBackend: the shared contract suite against a real Cosmos account."""
+import os
+import uuid
+
+import pytest
+
+from crewai_memory_core.contract import *  # noqa: F401,F403
+from crewai_memory_core.contract import DIMS
+from crewai_memory_cosmosdb import CosmosDBMemoryBackend
+
+EP = os.environ.get("COSMOS_ENDPOINT")
+KEY = os.environ.get("COSMOS_KEY")
+DB = os.environ.get("CREWAI_COSMOS_DB", "crewai_memory_test")
+
+pytestmark = pytest.mark.skipif(not (EP and KEY), reason="COSMOS_ENDPOINT/COSMOS_KEY not set")
+
+
+@pytest.fixture(scope="module")
+def _store():
+    cname = f"mem_{uuid.uuid4().hex[:8]}"
+    b = CosmosDBMemoryBackend(endpoint=EP, key=KEY, database_name=DB, container_name=cname, dimensions=DIMS)
+    yield b
+    from azure.cosmos import CosmosClient
+    CosmosClient(EP, credential=KEY).get_database_client(DB).delete_container(cname)
+
+
+@pytest.fixture()
+def backend(_store):
+    _store.reset()
+    yield _store
+    _store.reset()
+
+
+def test_native_vector_path_used(backend, monkeypatch):
+    from crewai_memory_core.contract import _seed, EMB
+    _seed(backend)
+    seen = []
+    orig = backend._container.query_items
+
+    def spy(query, **kw):
+        seen.append(query)
+        return orig(query=query, **kw)
+
+    monkeypatch.setattr(backend._container, "query_items", spy)
+    assert backend.search(EMB(["sushi"])[0], scope_prefix="/crew", limit=5)
+    assert any("VectorDistance" in q for q in seen)