diff --git a/README.md b/README.md index a4c7fee..a1fdf3c 100644 --- a/README.md +++ b/README.md @@ -150,7 +150,7 @@ from dakera import DakeraClient, RetryConfig client = DakeraClient(base_url="http://your-server:3000", api_key="your-key") # Cloud (early access) -client = DakeraClient(base_url="http://localhost:3000", api_key="your-key") +client = DakeraClient(base_url="http://:3000", api_key="your-key") # With custom retry config client = DakeraClient( diff --git a/pyproject.toml b/pyproject.toml index 5c29b72..a8b1371 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "dakera" -version = "0.11.56" +version = "0.11.57" description = "Python SDK for Dakera - AI memory platform" readme = "README.md" license = {text = "MIT"} diff --git a/src/dakera/__init__.py b/src/dakera/__init__.py index e10c534..f28dd30 100644 --- a/src/dakera/__init__.py +++ b/src/dakera/__init__.py @@ -48,6 +48,10 @@ BatchMemoryFilter, BatchRecallRequest, BatchRecallResponse, + BatchStoredMemory, + BatchStoreMemoryItem, + BatchStoreMemoryRequest, + BatchStoreMemoryResponse, # CE-12 CompressResponse, ConfigureNamespaceRequest, @@ -175,12 +179,16 @@ "RecalledMemory", "RecallResponse", "ConsolidateResponse", - # Batch memory operations (CE-2) + # Batch memory operations (CE-2, DAK-5508) "BatchMemoryFilter", "BatchRecallRequest", "BatchRecallResponse", "BatchForgetRequest", "BatchForgetResponse", + "BatchStoreMemoryItem", + "BatchStoreMemoryRequest", + "BatchStoreMemoryResponse", + "BatchStoredMemory", # Rate-limit headers (OPS-1) "RateLimitHeaders", # Session types diff --git a/src/dakera/client.py b/src/dakera/client.py index 50ff471..3bd6f52 100644 --- a/src/dakera/client.py +++ b/src/dakera/client.py @@ -35,6 +35,8 @@ BatchForgetResponse, BatchRecallRequest, BatchRecallResponse, + BatchStoreMemoryRequest, + BatchStoreMemoryResponse, BatchTextQueryResponse, # CE-12 CompressResponse, @@ -1260,6 +1262,31 @@ def batch_forget(self, request: BatchForgetRequest) -> BatchForgetResponse: result = self._request("DELETE", "/v1/memories/forget/batch", data=request.to_dict()) return BatchForgetResponse.from_dict(result) + def store_memories_batch(self, request: BatchStoreMemoryRequest) -> BatchStoreMemoryResponse: + """Store multiple memories in a single request (DAK-5508). + + Uses ``POST /v1/memories/store/batch``. The server embeds all contents + in a single ONNX inference pass, yielding ≥100× throughput vs. N + sequential single-store calls. Accepts up to 1 000 memories per call. + + Args: + request: Batch store request containing ``agent_id`` and list of + :class:`BatchStoreMemoryItem` (1–1000 items). + + Returns: + :class:`BatchStoreMemoryResponse` with stored memories and timing. + + Example: + >>> items = [ + ... BatchStoreMemoryItem("The user prefers dark mode", importance=0.8), + ... BatchStoreMemoryItem("The user is based in Berlin", importance=0.7), + ... ] + >>> resp = client.store_memories_batch(BatchStoreMemoryRequest("agent-1", items)) + >>> print(f"Stored {resp.stored_count} memories") + """ + result = self._request("POST", "/v1/memories/store/batch", data=request.to_dict()) + return BatchStoreMemoryResponse.from_dict(result) + def search_memories( self, agent_id: str, diff --git a/src/dakera/models.py b/src/dakera/models.py index dc4b448..2b8a2c7 100644 --- a/src/dakera/models.py +++ b/src/dakera/models.py @@ -1565,6 +1565,117 @@ def from_dict(cls, data: dict[str, Any]) -> "BatchForgetResponse": return cls(deleted_count=data.get("deleted_count", 0)) +@dataclass +class BatchStoreMemoryItem: + """A single memory entry within a :class:`BatchStoreMemoryRequest` (DAK-5508). + + Mirrors :class:`StoreMemoryRequest` but omits ``agent_id`` — supplied at batch level. + """ + + content: str + """Memory content text (required, max 100 000 chars).""" + memory_type: str = "episodic" + """One of ``"episodic"``, ``"semantic"``, ``"procedural"``, or ``"working"``.""" + importance: float = 0.5 + """Importance score 0.0–1.0 (default: 0.5).""" + tags: list[str] | None = None + """Optional tags to associate with the memory.""" + session_id: str | None = None + """Optional session ID to associate with.""" + metadata: dict[str, Any] | None = None + """Arbitrary metadata dictionary.""" + ttl_seconds: int | None = None + """Optional TTL in seconds.""" + expires_at: int | None = None + """Optional explicit expiry as a Unix timestamp (seconds).""" + id: str | None = None + """Optional custom ID. Auto-generated if not provided.""" + + def to_dict(self) -> dict[str, Any]: + d: dict[str, Any] = { + "content": self.content, + "memory_type": self.memory_type, + "importance": self.importance, + } + if self.tags is not None: + d["tags"] = self.tags + if self.session_id is not None: + d["session_id"] = self.session_id + if self.metadata is not None: + d["metadata"] = self.metadata + if self.ttl_seconds is not None: + d["ttl_seconds"] = self.ttl_seconds + if self.expires_at is not None: + d["expires_at"] = self.expires_at + if self.id is not None: + d["id"] = self.id + return d + + +@dataclass +class BatchStoreMemoryRequest: + """Request body for ``POST /v1/memories/store/batch`` (DAK-5508). + + Accepts up to 1 000 memories per call. The server embeds all contents in a + single ONNX inference pass and upserts them in one write, yielding ≥100× + throughput vs. N sequential single-store calls. + """ + + agent_id: str + """Agent namespace to store the memories in.""" + memories: list[BatchStoreMemoryItem] + """Memories to store (1–1000 items).""" + + def to_dict(self) -> dict[str, Any]: + return { + "agent_id": self.agent_id, + "memories": [m.to_dict() for m in self.memories], + } + + +@dataclass +class BatchStoredMemory: + """A single stored memory returned in a :class:`BatchStoreMemoryResponse`.""" + + id: str + content: str + agent_id: str + tags: list[str] + importance: float + created_at: int + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "BatchStoredMemory": + return cls( + id=data["id"], + content=data["content"], + agent_id=data["agent_id"], + tags=data.get("tags", []), + importance=data.get("importance", 0.5), + created_at=data.get("created_at", 0), + ) + + +@dataclass +class BatchStoreMemoryResponse: + """Response from ``POST /v1/memories/store/batch``.""" + + stored: list[BatchStoredMemory] + """Stored memories in the same order as the request items.""" + stored_count: int + """Number of memories successfully stored.""" + total_embedding_time_ms: int + """Time spent on ONNX embedding for the entire batch (milliseconds).""" + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "BatchStoreMemoryResponse": + return cls( + stored=[BatchStoredMemory.from_dict(m) for m in data.get("stored", [])], + stored_count=data.get("stored_count", 0), + total_embedding_time_ms=data.get("total_embedding_time_ms", 0), + ) + + # ============================================================================ # Memory Knowledge Graph Types (CE-5 / SDK-9) # ============================================================================ diff --git a/tests/test_client.py b/tests/test_client.py index 4dfd786..070f05f 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -829,6 +829,81 @@ def test_batch_recall_response_from_dict(self): assert resp.filtered == 0 assert resp.memories == [] + def test_store_memories_batch_sends_correct_request(self, client, mock_responses): + """store_memories_batch() POSTs to /v1/memories/store/batch.""" + from dakera import BatchStoreMemoryItem, BatchStoreMemoryRequest, BatchStoreMemoryResponse + + mock_responses.add( + responses.POST, + "http://localhost:3000/v1/memories/store/batch", + json={ + "stored": [ + { + "id": "mem-1", "content": "Dark mode", "agent_id": "agent-1", + "tags": [], "importance": 0.8, "created_at": 1700000000, + }, + { + "id": "mem-2", "content": "Berlin user", "agent_id": "agent-1", + "tags": [], "importance": 0.7, "created_at": 1700000001, + }, + ], + "stored_count": 2, + "total_embedding_time_ms": 42, + }, + status=200, + ) + + items = [ + BatchStoreMemoryItem("Dark mode", importance=0.8), + BatchStoreMemoryItem("Berlin user", importance=0.7), + ] + req = BatchStoreMemoryRequest("agent-1", items) + resp = client.store_memories_batch(req) + + assert isinstance(resp, BatchStoreMemoryResponse) + assert resp.stored_count == 2 + assert len(resp.stored) == 2 + assert resp.stored[0].id == "mem-1" + assert resp.total_embedding_time_ms == 42 + assert mock_responses.calls[0].request.method == "POST" + assert "/v1/memories/store/batch" in mock_responses.calls[0].request.url + + def test_batch_store_memory_item_to_dict(self): + """BatchStoreMemoryItem.to_dict() serializes all fields correctly.""" + from dakera import BatchStoreMemoryItem + + item = BatchStoreMemoryItem( + content="Test memory", + importance=0.9, + tags=["test", "qa"], + memory_type="semantic", + session_id="sess-1", + ) + d = item.to_dict() + assert d["content"] == "Test memory" + assert d["importance"] == 0.9 + assert d["tags"] == ["test", "qa"] + assert d["memory_type"] == "semantic" + assert d["session_id"] == "sess-1" + + def test_batch_store_memory_response_from_dict(self): + """BatchStoreMemoryResponse.from_dict() parses stored_count and memories.""" + from dakera import BatchStoreMemoryResponse + + resp = BatchStoreMemoryResponse.from_dict({ + "stored": [ + { + "id": "x", "content": "c", "agent_id": "a", + "tags": [], "importance": 0.5, "created_at": 0, + }, + ], + "stored_count": 1, + "total_embedding_time_ms": 10, + }) + assert resp.stored_count == 1 + assert resp.total_embedding_time_ms == 10 + assert resp.stored[0].id == "x" + class TestRateLimitHeaders: """Tests for OPS-1 RateLimitHeaders (v0.7.0)."""