From e9713af0f6922902c44c5f9ba25a0e3066108ea6 Mon Sep 17 00:00:00 2001 From: Caleb Kiragu Date: Wed, 12 Aug 2026 13:50:49 +0300 Subject: [PATCH] docs: add agent memory Context API example Provide a runnable SDK walkthrough that mirrors the pgContext SQL demo, including schema seeding, collection preflight config, mocked dry-run checks, and application-side context-pack assembly. Co-authored-by: Cursor --- README.md | 7 + examples/agent_memory/README.md | 164 +++++++++++++++++++++ examples/agent_memory/collection.json | 29 ++++ examples/agent_memory/demo.py | 144 +++++++++++++++++++ examples/agent_memory/dry_run.py | 189 +++++++++++++++++++++++++ examples/agent_memory/requirements.txt | 1 + examples/agent_memory/schema.sql | 78 ++++++++++ 7 files changed, 612 insertions(+) create mode 100644 examples/agent_memory/README.md create mode 100644 examples/agent_memory/collection.json create mode 100644 examples/agent_memory/demo.py create mode 100644 examples/agent_memory/dry_run.py create mode 100644 examples/agent_memory/requirements.txt create mode 100644 examples/agent_memory/schema.sql diff --git a/README.md b/README.md index 6f28117..02db033 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,13 @@ print(readiness.graph, readiness.vector, readiness.hybrid) Use the Runtime API URL with the SDK. Do not use a direct or pooled PostgreSQL connection string. +## Agent memory example + +For a runnable end-to-end pattern over the Context API, see +[examples/agent_memory/README.md](examples/agent_memory/README.md). It mirrors +the pgContext SQL walkthrough for durable decision memory, hybrid retrieval, +and LLM context-pack assembly. + ## Choose a retrieval method | Need | Method | diff --git a/examples/agent_memory/README.md b/examples/agent_memory/README.md new file mode 100644 index 0000000..5f239c1 --- /dev/null +++ b/examples/agent_memory/README.md @@ -0,0 +1,164 @@ +# Agent Memory Example (Polygres SDK) + +This example shows how an application retrieves **durable agent memory** through +the Polygres Runtime **Context API** (`project.context`). It mirrors the SQL +walkthrough in [Evokoa/pgContext](https://github.com/Evokoa/pgContext): +`examples/sql/06_agent_memory.sql`. + +## What it demonstrates + +- Scoped memory retrieval for one tenant and user +- Filtered dense search over a Context collection +- Hybrid dense + full-text retrieval with `text_hybrid` +- Application-side context-pack assembly for LLM prompts + +The demo uses fixture 4-dimensional embeddings so you can run it without calling +an external embedding model. Replace the vectors in your application with outputs +from your production embedding model. + +## Prerequisites + +- Python 3.10+ +- A [Polygres](https://polygres.com) project with Context enabled +- [polygres-cli](https://github.com/Evokoa/polygres-cli) for one-time setup + +```bash +pip install polygres-cli polygres-sdk +polygres login +polygres projects use +``` + +Create a Project API key from the project **Connect** page and export: + +```bash +export POLYGRES_API_KEY=poly_live_... +export POLYGRES_RUNTIME_URL=https://... +``` + +## One-time project setup + +### 1. Load schema and seed data + +From this repository root: + +```bash +polygres db psql < examples/agent_memory/schema.sql +``` + +This creates `agent_users`, `agent_sessions`, `agent_messages`, and +`agent_decisions` with the same fixture rows as the pgContext SQL example. + +### 2. Preflight and create the Context collection + +Inspect capabilities: + +```bash +polygres --json context capabilities +``` + +Preflight the collection definition: + +```bash +polygres --json context sources preflight --file examples/agent_memory/collection.json +``` + +Create the collection (requires explicit approval in agent-driven workflows): + +```bash +polygres context collections create agent_memory_decisions \ + --source existing \ + --schema public \ + --table agent_decisions \ + --source-key-column id \ + --vector-column embedding \ + --dimensions 4 \ + --metric cosine \ + --text-column body \ + --result-column summary \ + --result-column body \ + --result-column category \ + --result-column decided_at \ + --result-column user_id \ + --result-column tenant_id \ + --result-column session_id \ + --filter-column tenant_id \ + --filter-column user_id \ + --filter-column category +``` + +Synchronize catalog points after loading seed rows. Use the collection UUID +returned by the create command (not the collection name): + +```bash +COLLECTION_ID= + +polygres context points upsert "$COLLECTION_ID" \ + d-billing-refund d-arch-postgres d-onboarding-delay d-other-billing +``` + +Or reconcile every row in the source table: + +```bash +polygres context points reconcile "$COLLECTION_ID" +``` + +Exact CLI flag names may vary slightly by `polygres-cli` version. Use +`polygres context collections create --help` if a flag changed. + +```bash +python examples/agent_memory/dry_run.py +``` + +This validates `collection.json` against the Context request schema and exercises +`demo.py` against mocked Runtime responses (no API key required). + +## Run the demo + +```bash +python -m venv .venv +source .venv/bin/activate +pip install -r examples/agent_memory/requirements.txt +pip install -e . + +python examples/agent_memory/demo.py +``` + +Optional override: + +```bash +export POLYGRES_AGENT_MEMORY_COLLECTION=agent_memory_decisions +``` + +## Expected output + +When setup succeeded, the script prints: + +1. Context capability flags (`dense_search`, `text_hybrid`) +2. Filtered dense search returning `d-billing-refund` for tenant `acme` and user `u-alice` +3. Hybrid search ranking the billing decision highest for `billing refund` +4. A scoped context pack with summary and body fields + +## Review checklist + +- [ ] `schema.sql` loads without errors on the project database +- [ ] Collection creation completes and verification passes +- [ ] `demo.py` returns billing memory for the scoped user +- [ ] No API keys committed to git + +## Related material + +- pgContext SQL example: `Evokoa/pgContext/examples/sql/06_agent_memory.sql` +- SDK Context reference: `docs/reference-v1.md` +- Agent skill guidance: `Evokoa/polygres-skills` → `polygres-sdk/references/context.md` + +## Troubleshooting + +| Symptom | Likely cause | +|---|---| +| `Dense search unavailable` | Context collection missing or not verified | +| Empty results | Points not upserted after seed load | +| Dimension mismatch | Collection expects 4-D cosine vectors | +| `401` / auth errors | Invalid or expired API key | +| Hybrid skipped | `text_hybrid` capability blocked on project | + +Never commit `POLYGRES_API_KEY` or paste it into logs. diff --git a/examples/agent_memory/collection.json b/examples/agent_memory/collection.json new file mode 100644 index 0000000..de9f91e --- /dev/null +++ b/examples/agent_memory/collection.json @@ -0,0 +1,29 @@ +{ + "name": "agent_memory_decisions", + "source": { + "mode": "existing", + "schema_name": "public", + "table_name": "agent_decisions", + "source_key_column": "id" + }, + "vector": { + "column_name": "embedding", + "dimensions": 4, + "metric": "cosine" + }, + "text_column": "body", + "result_columns": [ + "summary", + "body", + "category", + "decided_at", + "user_id", + "tenant_id", + "session_id" + ], + "filter_columns": [ + "tenant_id", + "user_id", + "category" + ] +} diff --git a/examples/agent_memory/demo.py b/examples/agent_memory/demo.py new file mode 100644 index 0000000..96d75d2 --- /dev/null +++ b/examples/agent_memory/demo.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +"""Retrieve durable agent memory through the Polygres Context API. + +Companion to the pgContext SQL example in Evokoa/pgContext: + examples/sql/06_agent_memory.sql + +Requires: + - POLYGRES_API_KEY + - POLYGRES_RUNTIME_URL + - A Context collection named agent_memory_decisions (see README.md) +""" + +from __future__ import annotations + +import os +import sys +from typing import Any + +from polygres import Polygres + +COLLECTION = os.environ.get("POLYGRES_AGENT_MEMORY_COLLECTION", "agent_memory_decisions") +QUERY_EMBEDDING = [0.95, 0.05, 0.0, 0.0] +HYBRID_QUERY = "billing refund" +TENANT_ID = "acme" +USER_ID = "u-alice" +CATEGORY = "billing" + + +def _require_env(name: str) -> str: + value = os.environ.get(name) + if not value: + print(f"Missing required environment variable: {name}", file=sys.stderr) + sys.exit(2) + return value + + +def _source_key(result: Any) -> str: + source = getattr(result, "source", None) + if source is not None and getattr(source, "id", None): + return str(source.id) + properties = getattr(result, "properties", {}) or {} + for key in ("id", "source_key"): + if key in properties: + return str(properties[key]) + return str(getattr(result, "point_id", "unknown")) + + +def _build_context_pack(results: list[Any], *, tenant_id: str, user_id: str) -> list[dict[str, Any]]: + pack: list[dict[str, Any]] = [] + for result in results: + properties = dict(getattr(result, "properties", {}) or {}) + if properties.get("tenant_id") not in (None, tenant_id): + continue + if properties.get("user_id") not in (None, user_id): + continue + pack.append( + { + "rank": getattr(result, "rank", len(pack) + 1), + "score": getattr(result, "score", None), + "decision_id": _source_key(result), + "summary": properties.get("summary"), + "body": properties.get("body"), + "category": properties.get("category"), + "decided_at": properties.get("decided_at"), + "session_id": properties.get("session_id"), + } + ) + return pack + + +def main() -> None: + api_key = _require_env("POLYGRES_API_KEY") + runtime_url = _require_env("POLYGRES_RUNTIME_URL") + + client = Polygres(api_key=api_key, runtime_url=runtime_url) + context = client.project().context + + capabilities = context.get_capabilities() + print("Context capabilities:") + print(f" setup={capabilities.setup}") + print(f" dense_search={capabilities.dense_search}") + print(f" text_hybrid={capabilities.text_hybrid}") + + if not capabilities.dense_search: + blocker = capabilities.dense_search_blocker or "unknown" + message = capabilities.dense_search_blocker_message or "" + print(f"Dense search unavailable ({blocker}): {message}", file=sys.stderr) + sys.exit(1) + + memory_filter = { + "must": [ + {"key": "tenant_id", "match": TENANT_ID}, + {"key": "user_id", "match": USER_ID}, + {"key": "category", "match": CATEGORY}, + ] + } + + print("\nFiltered dense memory search:") + dense = context.search( + COLLECTION, + QUERY_EMBEDDING, + filter=memory_filter, + limit=5, + ) + for result in dense.results: + print(f" {_source_key(result):<18} score={result.score:.6f}") + + hybrid_results = dense.results + if capabilities.text_hybrid: + print("\nHybrid memory search (dense + full-text):") + hybrid = context.text_hybrid( + COLLECTION, + QUERY_EMBEDDING, + query=HYBRID_QUERY, + limit=5, + ) + hybrid_results = hybrid.results + for result in hybrid.results: + print(f" {_source_key(result):<18} score={result.score:.6f}") + else: + blocker = capabilities.text_hybrid_blocker or "unknown" + print(f"\nSkipping hybrid search ({blocker}).") + + print("\nContext pack for downstream LLM prompt assembly:") + context_pack = _build_context_pack( + list(hybrid_results), + tenant_id=TENANT_ID, + user_id=USER_ID, + ) + if not context_pack: + print(" No scoped memory rows returned.", file=sys.stderr) + sys.exit(1) + + for item in context_pack: + print( + f" #{item['rank']} {item['decision_id']} " + f"({item['category']}) score={item['score']:.6f}" + ) + print(f" summary: {item['summary']}") + print(f" body: {item['body']}") + + +if __name__ == "__main__": + main() diff --git a/examples/agent_memory/dry_run.py b/examples/agent_memory/dry_run.py new file mode 100644 index 0000000..cd5bf3d --- /dev/null +++ b/examples/agent_memory/dry_run.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Dry-run checks for the agent memory SDK example without live credentials.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from copy import deepcopy +from pathlib import Path + +import httpx + +ROOT = Path(__file__).resolve().parents[2] +if str(ROOT / "src") not in sys.path: + sys.path.insert(0, str(ROOT / "src")) + +from polygres import Polygres +from polygres._vendor.polygres_lib.context.models import CollectionCreateRequest + +EXAMPLE_DIR = Path(__file__).resolve().parent +DEMO_PATH = EXAMPLE_DIR / "demo.py" +COLLECTION_PATH = EXAMPLE_DIR / "collection.json" +FIXTURES_PATH = ROOT / "tests" / "fixtures" / "context" / "contract-fixtures.json" +API_KEY = "poly_live_0123456789abcdef0123456789abcdef" +RUNTIME_URL = "https://runtime.example.test/v1" +COLLECTION_ID = "00000000-0000-0000-0000-000000000010" + +FIXTURES = json.loads(FIXTURES_PATH.read_text()) + + +def _load_demo_module(): + spec = importlib.util.spec_from_file_location("agent_memory_demo", DEMO_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def _validate_collection_json() -> None: + payload = json.loads(COLLECTION_PATH.read_text()) + request = CollectionCreateRequest.model_validate(payload) + assert request.name == "agent_memory_decisions" + assert request.source.table_name == "agent_decisions" + assert request.vector.dimensions == 4 + assert request.text_column == "body" + assert "body" in request.result_columns + assert set(request.filter_columns) == {"tenant_id", "user_id", "category"} + + +def _memory_properties(*, tenant_id: str, user_id: str) -> dict[str, str]: + return { + "summary": "Approved partial enterprise refund", + "body": "Approved a partial refund for the enterprise billing dispute after policy review.", + "category": "billing", + "tenant_id": tenant_id, + "user_id": user_id, + "session_id": "s-billing-1", + "decided_at": "2026-07-01T10:00:00Z", + } + + +def _mock_handler(request: httpx.Request) -> httpx.Response: + path = request.url.path.removeprefix("/v1") + if path == "/context/capabilities": + payload = deepcopy(FIXTURES["responses"]["CapabilitiesResponse"]) + payload.update( + { + "setup": True, + "setup_blocker": None, + "setup_blocker_message": None, + "dense_search": True, + "dense_search_blocker": None, + "dense_search_blocker_message": None, + "text_hybrid": True, + "text_hybrid_blocker": None, + "text_hybrid_blocker_message": None, + } + ) + return httpx.Response(200, json=payload) + if path == "/context/search": + payload = deepcopy(FIXTURES["responses"]["RankedResponse"]) + payload["mode"] = "dense" + payload["results"] = [ + { + "point_id": 1, + "source": { + "schema": "public", + "table": "agent_decisions", + "id": "d-billing-refund", + }, + "rank": 1, + "score": 0.001382, + "score_kind": "context_metric", + "metric": "cosine", + "properties": _memory_properties(tenant_id="acme", user_id="u-alice"), + } + ] + return httpx.Response(200, json=payload) + if path == "/context/hybrid/text": + payload = deepcopy(FIXTURES["responses"]["RankedResponse"]) + payload["mode"] = "text_hybrid" + payload["results"] = [ + { + "point_id": 1, + "source": { + "schema": "public", + "table": "agent_decisions", + "id": "d-billing-refund", + }, + "rank": 1, + "score": 0.032787, + "score_kind": "rrf", + "rrf_k": 60, + "properties": _memory_properties(tenant_id="acme", user_id="u-alice"), + }, + { + "point_id": 2, + "source": { + "schema": "public", + "table": "agent_decisions", + "id": "d-other-billing", + }, + "rank": 2, + "score": 0.032258, + "score_kind": "rrf", + "rrf_k": 60, + "properties": { + "summary": "Denied refund for other tenant", + "body": "Denied refund request for a billing dispute in another tenant workspace.", + "category": "billing", + "tenant_id": "other", + "user_id": "u-bob", + "session_id": "s-other-1", + "decided_at": "2026-07-02T12:00:00Z", + }, + }, + ] + return httpx.Response(200, json=payload) + return httpx.Response(404, json={"error": {"code": "NOT_FOUND", "message": path}}) + + +def _run_demo_dry_run() -> None: + demo = _load_demo_module() + client = Polygres(api_key=API_KEY, runtime_url=RUNTIME_URL) + client._client.close() + client._client = httpx.Client(transport=httpx.MockTransport(_mock_handler)) + + context = client.project().context + capabilities = context.get_capabilities() + assert capabilities.dense_search + assert capabilities.text_hybrid + + memory_filter = { + "must": [ + {"key": "tenant_id", "match": demo.TENANT_ID}, + {"key": "user_id", "match": demo.USER_ID}, + {"key": "category", "match": demo.CATEGORY}, + ] + } + dense = context.search(demo.COLLECTION, demo.QUERY_EMBEDDING, filter=memory_filter, limit=5) + assert len(dense.results) == 1 + assert demo._source_key(dense.results[0]) == "d-billing-refund" + + hybrid = context.text_hybrid( + demo.COLLECTION, + demo.QUERY_EMBEDDING, + query=demo.HYBRID_QUERY, + limit=5, + ) + pack = demo._build_context_pack( + list(hybrid.results), + tenant_id=demo.TENANT_ID, + user_id=demo.USER_ID, + ) + assert len(pack) == 1 + assert pack[0]["decision_id"] == "d-billing-refund" + assert pack[0]["summary"] == "Approved partial enterprise refund" + assert "partial refund" in pack[0]["body"] + + +def main() -> None: + _validate_collection_json() + _run_demo_dry_run() + print("agent memory example dry run passed") + + +if __name__ == "__main__": + main() diff --git a/examples/agent_memory/requirements.txt b/examples/agent_memory/requirements.txt new file mode 100644 index 0000000..bd5e826 --- /dev/null +++ b/examples/agent_memory/requirements.txt @@ -0,0 +1 @@ +polygres-sdk>=0.2.0 diff --git a/examples/agent_memory/schema.sql b/examples/agent_memory/schema.sql new file mode 100644 index 0000000..bde97ae --- /dev/null +++ b/examples/agent_memory/schema.sql @@ -0,0 +1,78 @@ +-- Agent memory schema for Polygres projects. +-- Mirrors examples/sql/06_agent_memory.sql in Evokoa/pgContext. +-- +-- Load with: +-- polygres db psql < examples/agent_memory/schema.sql + +CREATE TABLE IF NOT EXISTS agent_users ( + id text PRIMARY KEY, + tenant_id text NOT NULL, + display_name text NOT NULL +); + +CREATE TABLE IF NOT EXISTS agent_sessions ( + id text PRIMARY KEY, + user_id text NOT NULL REFERENCES agent_users(id), + topic text NOT NULL, + started_at timestamptz NOT NULL DEFAULT now() +); + +CREATE TABLE IF NOT EXISTS agent_messages ( + id text PRIMARY KEY, + session_id text NOT NULL REFERENCES agent_sessions(id), + user_id text NOT NULL REFERENCES agent_users(id), + tenant_id text NOT NULL, + role text NOT NULL, + body text NOT NULL, + embedding pgcontext.vector(4) NOT NULL, + created_at timestamptz NOT NULL DEFAULT now() +); + +CREATE TABLE IF NOT EXISTS agent_decisions ( + id text PRIMARY KEY, + user_id text NOT NULL REFERENCES agent_users(id), + session_id text REFERENCES agent_sessions(id), + tenant_id text NOT NULL, + category text NOT NULL, + summary text NOT NULL, + body text NOT NULL, + decided_at timestamptz NOT NULL, + embedding pgcontext.vector(4) NOT NULL +); + +TRUNCATE agent_messages, agent_decisions, agent_sessions, agent_users CASCADE; + +INSERT INTO agent_users (id, tenant_id, display_name) VALUES + ('u-alice', 'acme', 'Alice Okello'), + ('u-bob', 'other', 'Bob Mwangi'); + +INSERT INTO agent_sessions (id, user_id, topic, started_at) VALUES + ('s-billing-1', 'u-alice', 'Enterprise billing review', '2026-07-01 09:00:00+00'), + ('s-arch-1', 'u-alice', 'Retrieval architecture', '2026-07-08 14:00:00+00'), + ('s-other-1', 'u-bob', 'Billing escalation', '2026-07-02 11:00:00+00'); + +INSERT INTO agent_messages (id, session_id, user_id, tenant_id, role, body, embedding, created_at) VALUES + ('m-1', 's-billing-1', 'u-alice', 'acme', 'user', + 'Can we approve a partial refund for the billing dispute?', + '[0.9,0.1,0,0]'::pgcontext.vector, '2026-07-01 09:05:00+00'), + ('m-2', 's-billing-1', 'u-alice', 'acme', 'assistant', + 'Reviewing prior billing decisions and refund policy.', + '[0.85,0.15,0,0]'::pgcontext.vector, '2026-07-01 09:06:00+00'); + +INSERT INTO agent_decisions (id, user_id, session_id, tenant_id, category, summary, body, decided_at, embedding) VALUES + ('d-billing-refund', 'u-alice', 's-billing-1', 'acme', 'billing', + 'Approved partial enterprise refund', + 'Approved a partial refund for the enterprise billing dispute after policy review.', + '2026-07-01 10:00:00+00', '[1,0,0,0]'::pgcontext.vector), + ('d-arch-postgres', 'u-alice', 's-arch-1', 'acme', 'architecture', + 'Selected Postgres-native retrieval', + 'Selected Postgres-native hybrid retrieval instead of operating a separate vector database.', + '2026-07-08 15:00:00+00', '[0,1,0,0]'::pgcontext.vector), + ('d-onboarding-delay', 'u-alice', 's-arch-1', 'acme', 'onboarding', + 'Delayed multilingual rollout', + 'Delayed multilingual dataset onboarding to Q3 to finish annotation QA.', + '2026-07-08 16:00:00+00', '[0,0,1,0]'::pgcontext.vector), + ('d-other-billing', 'u-bob', 's-other-1', 'other', 'billing', + 'Denied refund for other tenant', + 'Denied refund request for a billing dispute in another tenant workspace.', + '2026-07-02 12:00:00+00', '[1,0,0,0]'::pgcontext.vector);